HydroGym: Reinforcement Learning for Fluid Dynamics

61+ environments | 6 solver backends | 2D & 3D | Ready for RL training

HydroGym is a comprehensive platform for applying reinforcement learning to fluid dynamics and flow control. With environments ranging from canonical benchmarks to turbulent flows, HydroGym provides a standardized Gymnasium-compatible interface for training RL agents on challenging CFD problems.

Key Features

• Diverse Environments: 61+ pre-configured environments across 6 CFD solvers
• Standard RL Interface: Gymnasium-compatible API works with Stable-Baselines3, RLlib, and other RL libraries
• Compute Efficient: Highly optimized GPU & CPU backends for efficient RL deployment ranging from local workstations to exascale HPC systems
• Scalable: MPI-parallelized solvers with distributed RL training support
• Multiple Backends: Finite Element (Firedrake), Lattice Boltzmann (MAIA LBM), Finite Volume (MAIA FV), Spectral Element (NEK5000), Fully Differentiable solvers (JAX-Fluids)
• 2D & 3D: From simple 2D benchmarks to complex 3D turbulent flows (Re up to 400,000)
• Research-Ready: Managed by a complementary HuggingFace repository

Quick Start with Docker (Recommended)

We strongly recommend using our pre-configured Docker containers for hassle-free setup:

# For NVIDIA GPUs (CUDA)
docker pull clagemann/hydrogym-nvhpc-26.1_cuda-12.9_hopper_blackwell:latest
# or
docker pull clagemann/hydrogym-nvhpc-26.1_cuda-12.9_turing_ampere:latest

# For AMD GPUs (ROCm)
docker pull clagemann/hydrogym-rocm-6.3.3:latest

# Run container
docker run -it --gpus all clagemann/hydrogym-nvhpc-26.1_cuda-12.9_turing_ampere:latest

Available Environments

HydroGym provides 61 environments across 6 solver backends:

Solver Backend	Count	Description	Dimensions
Firedrake (FEM)	20	Canonical flow control benchmarks	2D
MAIA LBM	55	Lattice Boltzmann method environments	2D, 3D
MAIA Structured FV	8	High-Reynolds turbulent boundary layers	3D
NEK5000	2	Spectral element turbulent channel flow	3D
JAX	2	Differentiable fluid dynamics	2D, 3D
JAX-Fluids	2	Compressible shock vector control	2D, 3D

Environment Categories

Canonical Benchmarks (Low-Mid Re):

• Cylinder wake (Re=100-3900, 2D/3D)
• Rotating cylinder (Re=100-3900, 2D/3D)
• Pinball (Re=30-150, 2D/3D)
• Cavity flow (Re=4140-7500, 2D/3D)
• Backward-facing step (Re=600, 2D)
• Square cylinder (Re=200-3900, 2D/3D)
• Sphere (Re=300-3700, 3D)
• Cube (Re=300-3700, 3D)
• Turbulent channel flow (Re_tau=206, 3D)

Airfoil Control:

• NACA0012 steady (Re=100-50000, AOA=12-40°, 2D/3D)
• NACA0012 with gust disturbance (Re=100-50000, 2D/3D)

High Reynolds Number Flows:

• Zero-pressure-gradient turbulent boundary layer with jet/surface wave actuation (Re_Tau=180-2200, 3D)
• DRA2303 airfoil with jet/surface wave actuation (Re=400000, Ma=0.2-0.7, 3D)
• NACA0012 airfoil with jet actuation (Re=200000, 3D)

Fully Differentiable Flows:

• Shock-Vector Control in single divergent nozzle (SVC, Ma>1.0, 2D/3D)
• Turbulent channel flow (Re_tau=180, 3D)
• Kolmogorov flow (up to Re=1000, 2D)

Environment Checkpoints

All required environment checkpoints are available via HuggingFace and are downloaded on the fly when an environment is first created (internet connection required). If no internet connection is available at runtime — e.g. on compute nodes in HPC clusters — you can pre-download the environment files as outlined in examples/maia/README.md.

Examples

HydroGym includes comprehensive examples for each solver backend (internet connection required). We highly recommend using our provided docker containers:

Firedrake Examples

See examples/firedrake/getting_started/ for detailed documentation.

cd examples/firedrake/getting_started

# Test environment interactively
./run_example_docker.sh

# Train with Stable-Baselines3
./run_example_docker.sh train

MAIA Examples

See examples/maia/getting_started/ for MPMD coupling details.

cd examples/maia/getting_started

# Prepare workspace (downloads from Hugging Face Hub) and 
# Run with MPMD execution (1 Python + 1 MAIA process on GPU)
./run_example_docker.sh

# Prepare workspace and train with Stable-Baselines3
./run_example_docker.sh train

NEK5000 Examples

See examples/nek/getting_started/ for interface patterns.

cd examples/nek/getting_started

# Test single-agent environment
cd 1_nekenv_single
./run_nekenv_docker.sh

# ... or train with pettinzoo wrapper and SB3
cd 3_pettingzoo
./run_pettingzoo_docker.sh train

# ... or run zero-shot transfer learning
cd 6_zeroshot_wing_demo
./run_pettingzoo_docker.sh

JAX Examples

See examples/jax/getting_started/ for detailed documentation.

cd examples/jax/getting_started

# Test Kolmogorov flow environment
cd 1_kolmogorov
./run_nekenv_docker.sh

# ... or test channel flow environment
cd 2_channel
./run_channel_docker.sh strong_actuation

# ... or run zero-shot transfer learning
cd 3_ppo
./run_ppo_docker.sh --env channel --num-envs 1 --num-steps 10 --num-minibatches 5

Training RL Agents

HydroGym works with standard RL libraries. Example with Stable-Baselines3:

This is the beginning of the README. Read the rest at the original. Read the full README ↗