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.
- 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
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:latestHydroGym 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 |
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)
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.
HydroGym includes comprehensive examples for each solver backend (internet connection required). We highly recommend using our provided docker containers:
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 trainSee 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 trainSee 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.shSee 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 5HydroGym works with standard RL libraries. Example with Stable-Baselines3:
from hydrogym import FlowEnv
import hydrogym.firedrake as hgym
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
# Create environment
def make_env():
env_config = {
'flow': hgym.Cylinder,
'flow_config': {'mesh': 'medium', 'Re': 100},
'solver': hgym.SemiImplicitBDF,
'solver_config': {'dt': 1e-2},
'actuation_config': {'num_substeps': 2},
}
return FlowEnv(env_config)
# Vectorize and normalize
env = DummyVecEnv([make_env])
env = VecNormalize(env, norm_obs=True, norm_reward=True)
# Train
model = PPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=100000)See also provided examples/ for more details how to leverage individual solver backends for training.
- Checkpoint management: Automatic loading from Hugging Face Hub
- Custom observations: Force sensors, velocity/pressure/vorticity probes, stress sensors
- Callbacks: Checkpointing, logging, Paraview export
- Stability analysis: Global stability analysis via SLEPc (Firedrake)
- Modal decompositions: DMD, POD via modred (Firedrake)
- Multi-agent RL: PettingZoo interface for distributed actuation (NEK5000)
- Getting Started Guides: See
examples/[backend]/getting_started/README.md - API Documentation: https://hydrogym.readthedocs.io
- Flow Configurations: docs/FlowConfigurations.md
- Paper: arXiv:2512.17534
If you use HydroGym in your research, please cite the following two papers:
@inproceedings{lagemann2025hydrogym_a,
title={HydroGym: A Reinforcement Learning Platform for Fluid Dynamics},
author={Lagemann, Christian and Paehler, Ludger and Callaham, Jared and Mokbel, Sajeda and Ahnert, Samuel and Lagemann, Kai and Lagemann, Esther and Adams, Nikolaus and Brunton, Steven},
booktitle={7th Annual Learning for Dynamics$\backslash$\& Control Conference},
pages={497--512},
year={2025},
organization={PMLR}
}
@article{lagemann2025hydrogym_b,
title={Hydrogym: A reinforcement learning platform for fluid dynamics},
author={Lagemann, Christian and Mokbel, Sajeda and Gondrum, Miro and R{\"u}ttgers, Mario and Callaham, Jared and Paehler, Ludger and Ahnert, Samuel and Zolman, Nicholas and Lagemann, Kai and Adams, Nikolaus and others},
journal={arXiv preprint arXiv:2512.17534},
year={2025}
}HydroGym is released under the MIT License. See LICENSE for details.