HydroGym: Reinforcement Learning Platform for Fluid Dynamics

HydroGym models fluid-flow control as a discrete-time Markov decision process (MDP), represented by the tuple \(({\mathcal{S}},{\mathcal{A}},{\mathcal{P}},{\mathcal{R}})\). At each control step t, a reinforcement learning (RL) agent receives an observation state \({s}_{t}\in{\mathcal{S}}\) containing flow measurements, selects an action \({a}_{t}\in{\mathcal{A}}\) representing actuator inputs, and receives a scalar reward \({r}_{t}\in{\mathcal{R}}\) that measures the control objective, such as drag reduction or flow stabilization. The transition probability \({\mathcal{P}}\) is determined implicitly by the Navier–Stokes equations solved by the selected computational fluid dynamics (CFD) backend.

To connect reinforcement learning with computational fluid dynamics, HydroGym uses the Farama Foundation Gymnasium interface⁴⁴. This standard interface hides the complexity of CFD execution and provides a consistent workflow across different solvers. The following example shows how to initialize and control an open-cavity-flow environment:

# Select a CFD backend by changing one import

import hydrogym.firedrake as hgym

# import hydrogym.maia as hgym

# import hydrogym.nek as hgym

# import hydrogym.jax as hgym

# import hydrogym.jaxfluids as hgym

# Create the flow-control environment

env = hgym.FlowEnv(‘Cavity_2D_Re7500’, **kwargs)

# Reset the environment

obs, info = env.reset()

# Run the control loop

for i in range(num_interactions): 

obs, reward, terminated, truncated, info = env.step(action)

This solver-independent design allows machine learning and control engineers to develop RL algorithms without specialist CFD knowledge, while giving fluid-dynamics researchers access to modern reinforcement learning methods. Because the architecture is independent of any individual solver, new CFD capabilities can be incorporated without changing established training workflows.

A central feature of HydroGym is its temporal coupling between the physical solver and the RL agent. Turbulent flows evolve continuously, whereas RL policies generally select actions at discrete intervals. HydroGym therefore updates control inputs at fixed intervals Δt_ctrl. Each control interval contains multiple physical CFD time steps selected to resolve the characteristic timescale of the target instability, such as vortex shedding or shear-layer oscillation.

HydroGym also applies temporal smoothing to control actions to reduce numerical instabilities caused by abrupt boundary-condition changes. The instantaneous actuation magnitude A(t) is interpolated smoothly between the previous action A^old and the new target action A^new across a predefined number of sub-iterations using a hyperbolic-tangent or exponential ramp. In addition, observation features and action limits are normalized to the [−1, 1] interval by the environment wrapper, helping standardize RL training across different flow configurations.

Computational fluid dynamics backends

HydroGym uses a solver-independent CFD architecture that supports applications ranging from rapid prototyping to large-scale high-performance computing (HPC). Each computational environment has been validated against established benchmarks to confirm physical accuracy and high-fidelity flow resolution. Detailed validation results and grid-convergence studies for the individual solvers and flow cases are provided in Supplementary Sections 3 and 5.

Lattice Boltzmann solver (m-AIA LB)

HydroGym uses the lattice Boltzmann method (LBM) within the m-AIA solver framework⁴⁵ for most weakly compressible, large-scale two- and three-dimensional direct numerical simulations (DNS). Developed at RWTH Aachen University for more than 20 years, the solver advances the discrete particle distribution function f_i on a D3Q27 velocity lattice using a two-step collision–streaming update.

At low and moderate Reynolds numbers, the Bhatnagar–Gross–Krook (BGK) collision operator relaxes the distributions toward a local Maxwellian equilibrium. Its single relaxation frequency, ω_BGK, is linked to the kinematic viscosity. For higher-Reynolds-number simulations, HydroGym uses a cumulant-based collision operator that performs relaxation in cumulant space with individual rates ω_α (ref. ⁴⁶). Local grid refinement follows the method of Eitel-Amor and colleagues⁴⁷; ω_BGK is adjusted across hierarchical grid levels to maintain a constant kinematic viscosity. No-slip walls are implemented using a second-order interpolated bounce-back scheme.

The m-AIA lattice Boltzmann solver uses hybrid parallelization based on the message passing interface (MPI) and shared-memory computing. Shared-memory execution is supported through OpenMP or parallel algorithms in the C++ standard, including the NVIDIA HPC SDK and AMD ROCm HIPSTDPAR backends. This approach enables hardware-independent GPU acceleration.

Finite-volume solver (m-AIA FV)

In addition to its lattice Boltzmann implementation, m-AIA provides a compressible Navier–Stokes solver based on the finite-volume method. It is designed for DNS and large-eddy simulation (LES) of wall-bounded turbulent flows at higher Mach numbers. The governing equations are written in an arbitrary Lagrangian–Eulerian formulation and discretized on structured, body-conforming curvilinear grids using a cell-centred scheme. This arrangement supports anisotropic wall-normal refinement for flat-plate and airfoil simulations.

Inviscid fluxes are computed with the advection upstream splitting method and Monotonic Upstream-centred Scheme for Conservation Laws reconstruction, providing second-order spatial accuracy. Viscous fluxes use a modified cell-vertex formulation. Time integration is performed with a five-stage, low-storage Runge–Kutta method. For moving or deforming grids under surface actuation, the geometric conservation law is enforced at every stage.

Under-resolved LES is treated with a monotonically integrated LES strategy that uses the inherent numerical dissipation of the upwind-biased advection upstream splitting method. The finite-volume solver shares the lattice Boltzmann solver’s MPI, OpenMP and GPU parallelization framework.

For reinforcement learning, the solver supports synthetic-jet actuation and travelling transverse surface-wave boundary conditions^{14,48,49,50,51}. The agent updates the wave amplitude A, wavelength λ and phase speed c at each control step. A C¹-continuous cosine cross-fade transition ensures smooth parameter changes and numerical stability.

Spectral-element solver (Nek5000)

For incompressible, high-Reynolds-number flows that require high-order accuracy, HydroGym provides a Nek5000⁵² backend based on the spectral-element method. Velocity and pressure are approximated with high-order Lagrange interpolants on hexahedral elements using the \({{\mathbb{P}}}_{N}{{\mathbb{P}}}_{N-2}\) formulation⁵³. Velocity is collocated at N³ Gauss–Lobatto–Legendre points per element, while pressure is defined on a staggered grid containing (N − 2)³ Gauss–Legendre points.

Temporal discretization combines a third-order explicit extrapolation method for nonlinear convection with a third-order implicit backward-differentiation method for viscous terms. Overintegration, using an oversampling factor of 3/2 in each spatial direction, reduces aliasing errors. Flow actuation is prescribed as a time-dependent wall-normal Dirichlet velocity v_n; its spatial mean is removed to satisfy the zero-net-mass-flux (ZNMF) constraint. Hybrid OpenMPI distributed- and shared-memory parallelization is used, and RL coupling is implemented through dynamic MPI communication.

Differentiable incompressible solver

HydroGym includes differentiable JAX⁵⁴ solvers for the incompressible Navier–Stokes equations, enabling gradient-enhanced reinforcement learning²⁸. The Kolmogorov-flow environment uses a pseudo-spectral method with two-thirds de-aliasing on a doubly periodic domain. Its vorticity–streamfunction formulation is advanced with a fourth-order Runge–Kutta scheme.

The three-dimensional turbulent-channel environment uses a fully differentiable finite-difference DNS solver. Spatial derivatives are calculated with sparse differentiation matrices, while the pressure Poisson equation is solved using a biconjugate gradient stabilized method. Both environments implement the Gymnax interface⁵⁵ and use RL components from PureJaxRL⁵⁶, creating a synchronous training pipeline compatible with JAX automatic differentiation.

Using jax.grad, HydroGym computes exact analytical gradients of the objective with respect to control parameters through the complete simulation trajectory. Multiple random seeds are processed in parallel with vmap, and just-in-time compilation with jit accelerates GPU execution.

Differentiable compressible solver (JAX-Fluids)

For gradient-based flow control in compressible and multiphase regimes, HydroGym integrates JAX-Fluids^57,58, a fully differentiable, high-order CFD solver written in Python with JAX. The compressible Navier–Stokes equations are solved on structured Cartesian grids supporting arbitrary one-dimensional mesh stretching.

Convective fluxes use either fifth-order weighted essentially non-oscillatory WENO5-Z reconstruction or sixth-order targeted essentially non-oscillatory TENO6-A reconstruction, together with an approximate Harten–Lax–van Leer–Contact Riemann solver. Diffusive fluxes are calculated with high-order central differences, and time integration uses explicit total-variation-diminishing Runge–Kutta schemes.

JAX-Fluids supports multiphase simulations through a sharp-interface level-set method and a five-equation diffuse-interface model. Positivity-preserving limiters improve stability near strong shocks and large density ratios. The level-set formulation also acts as an immersed-boundary method for complex geometries. This is useful for active flow control because blowing and suction can be incorporated naturally through interface exchange terms.

Because the solver follows JAX’s functional programming model, the complete simulation pipeline remains end-to-end differentiable. Exact gradients of scalar objectives, including time-averaged drag and RL rewards, can be backpropagated through the full trajectory with jax.grad or value_and_grad. Gradient checkpointing with jax.checkpoint reduces memory requirements during long rollouts.

For large-scale HPC applications, JAX-Fluids uses homogeneous domain decomposition across multiple accelerator devices with jax.pmap. Inter-block halo exchanges are performed with jax.lax.ppermute, preserving the automatic-differentiation graph across distributed multi-node systems.

Finite-element solver (Firedrake)

HydroGym provides a Firedrake⁵⁹ backend for transparent code development and rapid prototyping of two-dimensional flow-control problems. Built on the Portable, Extensible Toolkit for Scientific Computation, Firedrake automatically generates code for variational formulations.

The incompressible Navier–Stokes equations are discretized with Taylor–Hood elements: second-order continuous Galerkin elements for velocity and first-order elements for pressure. This choice provides inf–sup stability. Fully implicit time-integration schemes are used, with solver parameters selected automatically according to the Reynolds number and grid resolution.

The implementation is divided into three modular layers: PDEBase defines the physical problem, TransientSolver advances the solution in time, and FlowEnv connects the CFD solver to the Gymnasium RL interface. Firedrake’s distributed meshes and PETSc parallel sparse linear-algebra routines provide scalable execution.

Reinforcement learning and multi-agent infrastructure

HydroGym emphasizes robust observation and action normalization rather than extensive hyperparameter tuning for every flow configuration. This strategy improves generalizability and allows standard, off-the-shelf RL algorithms to provide strong baseline results.

The platform integrates with StableBaselines3⁶⁰, TorchRL⁶¹ and CleanRL⁶² for model-free continuous control, and with PureJaxRL for differentiable environments. For standard model-free environments, HydroGym evaluates Proximal Policy Optimization (PPO), Deep Deterministic Policy Gradient (DDPG) and Twin Delayed DDPG (TD3).

Gradient-enhanced PPO

For differentiable environments, HydroGym implements gradient-enhanced PPO (GPPO). Conventional PPO estimates policy gradients with likelihood-ratio methods or generalized advantage estimation. GPPO additionally computes the analytical gradient of the reward trajectory with respect to policy parameters:

\({\nabla }_{\theta }{{\mathbb{E}}}_{\tau \sim {\pi }_{\theta }}[\sum R({s}_{t},{a}_{t})]\)

This gradient is obtained by backpropagating through the deterministic fluid-dynamics solver. Incorporating the analytical gradient into PPO’s clipped surrogate objective reduces gradient variance and improves sample efficiency.

Multi-agent reinforcement learning

HydroGym uses decentralized multi-agent reinforcement learning (MARL) to address the high dimensionality of spatially distributed three-dimensional flows, including cylinder, channel and airfoil configurations. A global control domain is divided into locally invariant pseudo-environments. Multiple agents operate within these partitions and may share a common control policy π(a∣s).

To avoid biased samples in the shared replay buffer, overlapping mesh nodes are excluded from pseudo-environment definitions. Identical actuation is enforced at partition boundaries to maintain control continuity. An MPI-based interface enables efficient communication between parallel CFD domains and the centralized RL policy.

Additional information about the reinforcement learning agents is provided in Supplementary Section 4.

Environment setup for the reported results

The following sections summarize the flow environments discussed in the main text. Full specifications, including physical parameters, boundary conditions, observation and action normalization, and reward definitions, are provided in Supplementary Section 5.

Circular cylinder flow (Re = 3,900)

The subcritical three-dimensional cylinder case produces a highly chaotic wake. The computational domain measures [51.2D × 48D] in two dimensions and 32D × 16D × 4D in three dimensions, with periodic boundaries in the spanwise direction.

Actuation is applied through zero-net-mass-flux synthetic jets located at the top and bottom of the cylinder. The reward promotes drag reduction while penalizing lift oscillations:

r = −∣C_D∣ − ω∣C_L∣.

For MARL experiments, the spanwise domain is divided into independent pseudo-environments, with each agent controlling a local pair of jets.

Fluidic pinball flow (Re = 100–150)

The fluidic-pinball environment contains three circular cylinders arranged in an equilateral triangle. It exposes RL agents to several bifurcation regimes, including symmetry-breaking pitchfork bifurcations and chaotic flow. Control is provided by independently rotating the surfaces of all three cylinders.

The reward targets collective drag reduction:

\(r=-{\sum }_{i=1}^{3}| {C}_{D,i}| \,-\) \(\omega {\sum }_{i=1}^{3}| {C}_{L,i}| \)

Here, the scaling factor ω discourages policies from exploiting asymmetric lift generation.

Open-cavity flow (Re = 4,200–7,500)

The open-cavity environment is designed to stabilize complex feedback loops and Kelvin–Helmholtz instabilities in the shear layer spanning the cavity opening. The reward penalizes deviations between measured flow quantities and a target reference state:

\(r=-{\sum }_{i}{\left(\frac{{{\rm{obs}}}_{i}-{\bar{o}}_{i}}{{{\sigma }}_{i}}\right)}^{2}\)

In this expression, obs_i denotes pressure and velocity measurements from shear-layer probes, \({\bar{o}}_{i}\) is the target mean estimated over 1,000 instability cycles, and σ_i provides the normalization scale.

Localized jet actuators are positioned at the upstream cavity edge and, optionally, inside the cavity. This arrangement enables multi-point flow manipulation designed to disrupt resonant shear-layer interactions.

Transverse-gust mitigation for NACA0012 (Re = 1,000)

This environment represents a severe aerodynamic operating condition in which a strongly disturbed inflow interacts with a NACA0012 airfoil. A 1-cosine transverse gust with a gust ratio of G = 2.0 encounters the airfoil at a high angle of attack, α = 20°, creating a risk of dynamic stall and significant aerodynamic-load fluctuations.

Control is provided by three independent jet actuators distributed along the leading edge. Each actuator covers 3% of the airfoil chord. The reward is designed to reduce gust-induced force variance while maintaining baseline aerodynamic efficiency.