Reinforcement-learning control of turbulence transition in the modified Hasegawa-Wakatani system

Problem Formulation

The paper addresses the challenge of controlling plasma turbulence, a key mechanism for anomalous transport in magnetically confined fusion. While experimental control often relies on static, empirical signals, the authors propose using Reinforcement Learning (RL) to find optimal, non-trivial control trajectories for nonlinear, chaotic plasma dynamics.

To make this computationally tractable, the study utilizes a reduced-order model: the Modified Hasegawa–Wakatani (MHW) system, which describes electrostatic drift-wave turbulence and zonal-flow dynamics. The authors identify two critical limitations in the standard MHW model for control applications:

1. Lack of Repeatability: The standard model supports effectively absorbing, long-lived zonal-flow states that persist indefinitely once formed, making bidirectional transitions (turbulence $\leftrightarrow$ zonal flow) impossible within a finite control episode.
2. Unconstrained Actuation: Real-world actuators (e.g., biasing electrodes, RF heating) have finite energy budgets.

To resolve these, the authors introduce two modifications:

• Zonal Drag: A weak linear drag term ($-\gamma_{ZF} \bar{\varpi}$) is added to the vorticity equation. This damps zonal structures, giving them a finite lifetime and restoring a drive-damping balance necessary for repeatable transitions.
• Budget-Constrained Actuation: Control is applied via a spatially distributed Gaussian source field ($\phi_{ext}$) composed of four independent actuators. A time-weighted cost is applied to the actuation, enforcing a fixed total budget ($B_{total}$) per episode.

The study investigates two complementary control tasks:

1. Turbulence Suppression: Driving a drift-wave turbulence-dominated state ($\alpha=0.1$) into a quiescent zonal-flow state.
2. Zonal-Break (Turbulence Enhancement): Disrupting a self-sustained zonal-flow state ($\alpha=0.5$) to restore cross-field turbulent transport.

Methodology

The authors employ a model-free, online RL framework interacting with a custom, GPU-native solver.

• Solver: The MHW equations are solved using JAX, a Python library optimized for GPU acceleration. The solver uses the Arakawa scheme for Poisson brackets, spectral methods for field inversion, and supports both leapfrog and Runge-Kutta time integration. It runs approximately an order of magnitude faster than legacy Fortran codes (GDB), enabling the high-throughput rollouts required for online training.
• Algorithms: Two actor-critic algorithms are used:
  • Soft Actor-Critic (SAC): Used for the turbulence-suppression task (continuous action space).
  • Twin-Delayed Deterministic Policy Gradient (TD3): Used for the zonal-break task (discrete action space), chosen for its ability to converge to boundary optima.
• Observation Space: The agent observes a 5-channel, $256 \times 256$ field consisting of the electrostatic potential ($\phi$), remaining budget, zonal kinetic energy, its running integral, and the fraction of the episode remaining.
• Reward Functions:
  • Suppression: A physics-based reward minimizing the time-integrated turbulent flux ($\langle n \partial_y \phi \rangle$), penalized by a time-weighted action cost and an overspend penalty.
  • Zonal-Break: A scale-invariant normalized flux reward to prevent "reward hacking" (where the agent artificially inflates potential amplitude without disrupting zonal structure).
• Initialization: A physics-informed warm-buffer strategy is employed. For the zonal-break task, where the optimal solution lies on a narrow manifold, the replay buffer is pre-filled with episodes containing specific antisymmetric actuation patterns to guide exploration.

Key Results

1. Turbulence Suppression Task

• Performance: The learned RL policy outperformed two heuristic baselines (constant drive and linear ramp-down) across five unseen initial conditions. It achieved the lowest time-integrated turbulent flux for the same consumed budget.
• Strategy: The agent learned a "front-loaded" schedule: applying strong forcing early to trigger the transition to a zonal state, then tapering along a non-trivial, nonlinear profile to maintain the state against zonal drag. This combined the rapid transition of the linear ramp with the stability of the constant drive.
• Generalization: The policy converged to a near-open-loop schedule that was robust across different initial conditions, effectively discarding initial-condition-specific details while maintaining the optimal temporal trajectory.

2. Zonal-Break Task

• Challenge: The optimal solution required an up-down antisymmetric actuation pattern (top actuators positive, bottom negative) to induce radial $E \times B$ convection. This solution occupied a vanishingly small manifold in the action space.
• Reward Hacking: Using a raw flux reward, the agent failed, converging to a uniform (same-sign) pattern that artificially inflated potential amplitude without disrupting the zonal structure.
• Solution: By switching to a scale-invariant normalized reward and using a physics-informed warm buffer (seeded with antisymmetric patterns), the agent successfully rediscovered the optimal configuration.
• Outcome: The learned policy applied a constant antisymmetric pattern $[-9, -9, +9, +9]$, successfully breaking the coherent zonal bands and enhancing cross-field transport by a factor of 1.8 compared to unactuated states.

Significance and Claims

The paper positions itself as an exploratory study demonstrating the viability of RL as a direct optimizer for nonlinear plasma dynamics. Its primary contributions and claims are:

• Practical Trajectory Optimization: RL can optimize control policies directly against the nonlinear turbulent transport itself, rather than relying on proxies like linear growth rates or background gradients.
• Necessity of Physics Guidance: The study highlights that for turbulent systems, the reward landscape is often flat or deceptive away from the optimum. Pure random exploration is insufficient; physics-informed initialization (warm buffers) and physics-based reward design (scale invariance) are critical to finding optimal policies.
• Modeling Framework: The introduction of the weak zonal drag and budget constraints creates a well-posed control problem that mimics the finite lifetime of zonal states and experimental energy limits, allowing for repeatable bidirectional transitions.
• Scalability: The use of a GPU-native JAX solver makes online training on reduced-order plasma models feasible, bridging the gap between fluid dynamics control and plasma physics.

The authors explicitly state that while the results are promising, the study is limited to a reduced-order model with specific parameters. They do not claim immediate applicability to full-scale fusion devices but argue that the geometric constraints and the need for physics-informed learning are likely generic features of turbulence control problems.