🤖 machine learning

Offline Reinforcement Learning for Plasma Control in Nuclear Fusion: Codebase and Benchmark

This paper introduces RL4F, a standardized offline reinforcement learning benchmark for nuclear fusion plasma control built on real DIII-D tokamak data, which evaluates various algorithms across four full-profile tracking tasks and demonstrates that offline model-based methods generally achieve superior performance in these complex, long-horizon problems.

Original authors: Yang Fu, Haomin Bao, Rohit Sonker, Xiaoyan Hu, Aravind Venugopal, Jeff Schneider, Jiayu Chen

Published 2026-06-09

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yang Fu, Haomin Bao, Rohit Sonker, Xiaoyan Hu, Aravind Venugopal, Jeff Schneider, Jiayu Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to fly a very complicated, unstable airplane. The problem is that the plane is so expensive and dangerous that you can't let the robot crash it a thousand times to learn how to fly. Every time it crashes, it costs millions of dollars and could hurt people.

This is exactly the situation scientists face with nuclear fusion, the technology that aims to replicate the power of the sun to create clean energy. Inside a fusion machine (called a tokamak), there is a super-hot, swirling soup of gas called plasma. This plasma is incredibly unstable. If you don't control it perfectly, it cools down instantly or damages the machine.

For a long time, scientists have tried to use Reinforcement Learning (RL)—a type of AI that learns by trial and error—to control this plasma. But, just like the robot pilot, they can't let the AI "play" on the real machine to learn.

The Problem: No "Flight Simulator"

Usually, to train an AI, you build a perfect video game simulator. But for nuclear fusion, the physics are so complex that building a perfect simulator is like trying to write a video game that perfectly predicts the weather, the ocean currents, and the movement of every single atom at the same time. It's too slow and too hard to get right.

The Solution: "RL4F" (The Fusion Flight Simulator)

The authors of this paper created a new tool called RL4F. Think of it as a high-tech flight simulator built not from physics equations, but from historical flight logs.

The Data: They took thousands of hours of real data from a real fusion machine (the DIII-D tokamak).
The "Digital Twin": They trained an AI to look at those old logs and learn: "When the operators did X, the plasma usually reacted like Y."
The Result: This AI became a "Digital Twin" of the real machine. Now, researchers can train their new control algorithms inside this Digital Twin. The AI can crash, fail, and try again a million times without ever touching the real, dangerous machine.

The Challenge: The "Profile" Puzzle

Controlling fusion isn't just about keeping the temperature high. It's about shaping the entire profile of the plasma. Imagine the plasma as a loaf of bread. You don't just want the whole loaf to be warm; you want the crust to be a certain crispness, the middle to be soft, and the center to be perfectly baked, all at the same time.

The paper tested the AI on four specific "loaves" (tasks):

Rotation: How fast the plasma spins.
Density: How crowded the particles are.
Temperature: How hot it is.
Pressure: How much force it exerts.

The Race: Who Learned Best?

The authors invited many different AI "students" (algorithms) to take a test in this simulator. They wanted to see which one could best keep the plasma "loaf" shaped correctly.

Here is what they found:

The "Copycat" Students: Some AIs just tried to copy what the human operators did in the old logs. They were okay, but they couldn't handle new situations well.
The "Model-Free" Students: These AIs tried to learn purely by guessing and checking within the data. They did better than the copycats but still struggled with the long, complex chains of cause-and-effect in the plasma.
The "Model-Based" Students: These AIs built their own internal understanding of how the plasma works (a "world model") and then planned their moves based on that. These students won the race.

Specifically, the MOPO and COMBO algorithms (the "Model-Based" students) were the best at keeping the temperature and density profiles on track. However, no single student was perfect at everything. For example, one was great at controlling rotation, while another was better at pressure.

Why This Matters

The paper concludes that model-based learning (learning the rules of the game first) is the most promising path for controlling fusion.

They have made all their code, data, and the "Digital Twin" simulator free for everyone to use. This is like giving every researcher in the world the same flight simulator and the same test track. Now, instead of everyone building their own confusing simulators, they can all compete fairly to see who can build the best AI pilot for the future of clean energy.

In short: They built a safe, virtual playground using real-world data so AI can learn to control the sun without blowing up the lab. And they found that the AI that learns the "rules of physics" first performs the best.

Technical Summary: RL4F – Offline Reinforcement Learning for Plasma Control in Nuclear Fusion

Problem Statement

Developing controllers for nuclear fusion tokamaks via online reinforcement learning (RL) is prohibitively expensive and risky due to the high cost of operation, time limitations, and safety-critical nature of real devices. While physics-based simulators (e.g., RAPTOR, TORAX) exist, they are computationally expensive, require extensive calibration, and often struggle to capture complex, nonlinear plasma dynamics. Conversely, learning directly from historical experimental data without online interaction (offline RL) offers a safer alternative but lacks a standardized benchmark. Current challenges include the high dimensionality of plasma profiles, nonlinear and stochastic dynamics, partial observability, and the difficulty of specifying rewards for long-horizon control tasks.

Methodology

The paper introduces RL4F, a unified benchmark for offline RL in tokamak profile control. The methodology is structured around four core components:

1. Dynamics Modeling (The Environment)

Instead of using physics-based simulators, the authors construct a data-driven "digital twin" using a Recurrent Probabilistic Neural Network (RPNN) ensemble.

Training Data: The model is trained on 18,000 historical experimental discharges from the DIII-D tokamak, spanning nearly a decade.
Architecture: A 25-member ensemble of RPNNs predicts the distribution of state changes ( $\Delta s_t$ ) given the current state ( $s_t$ ), previous actuator setting ( $a_{t-1}$ ), and actuator increment ( $\Delta a_t$ ).
Uncertainty Quantification: The model outputs parameters for a Gaussian distribution ( $\mu_t, \sigma_t$ ). The ensemble captures epistemic uncertainty (via member disagreement), while the log-variance head captures aleatoric uncertainty.
Two-Stage Training: The first stage optimizes the predictive mean (MSE), and the second stage freezes the backbone to calibrate uncertainty via negative log-likelihood.

2. Dataset Construction

A synthetic offline dataset is generated by initializing rollouts from real "flat-top" shot states and autoregressively rolling out the reference dynamics model.

Scale: The final dataset comprises 5,882 shots and 945,828 transitions.
Split: Fixed splits are used: 5,282 shots for training, 300 for validation, and 300 for testing.
Procedure: Measured next states are replaced by ensemble predictions to simulate closed-loop behavior while preserving real actuator schedules.

3. Benchmark Tasks

The benchmark evaluates algorithms on four full-profile tracking tasks, where the goal is to regulate high-dimensional spatial quantities across the normalized radial coordinate ( $\psi_N \in [0, 1]$ ):

Rotation: Critical for tearing mode suppression and fueling asymmetry.
Density: Regulates core fuel inventory and fusion rates.
Electron Temperature: Affects heat transport and confinement.
Pressure: Determines MHD stability limits and stored energy.

Control Setup:

Action Space: Neutral Beam Injection (NBI) power, NBI torque, gas puffing voltage, and Electron Cyclotron Heating (ECH) power.
Observation Space: Current profile, target profile, future target (shifted 10 steps), and proportional error terms (P-terms).
Reward: Negative mean squared error (MSE) between the target and achieved profile.

4. Evaluation Protocol

Algorithms are trained on the static synthetic dataset and evaluated in a closed-loop setting on the reference dynamics model.

Baselines: The study evaluates a wide range of methods including Imitation Learning (GCIL), Model-Free Offline RL (TD3BC, CQL, IQL, EDAC, MCQ), and Model-Based Offline RL (PPO, COMBO, MOPO, MOBILE, RAMBO, BAMCTS).
Metric: Root Mean Squared Error (RMSE) across 33 radial locations, averaged over 300 test shots and 10 random seeds.

Key Contributions

RL4F Benchmark: The first standardized offline RL benchmark specifically designed for fusion plasma control, providing a common testbed for algorithm comparison.
Realistic Data-Driven Environment: A high-fidelity dynamics model trained on real DIII-D data, avoiding the calibration overhead of physics-based simulators while capturing device-specific nonlinearities.
Comprehensive Baseline Evaluation: A unified evaluation of diverse offline RL and imitation learning approaches across four distinct, long-horizon profile tracking tasks.
Open-Source Release: The codebase, datasets, and evaluation framework are open-sourced to foster reproducibility and community adoption.

Results

The evaluation reveals that model-based offline RL methods generally outperform model-free baselines, though no single algorithm dominates all tasks:

Rotation: RAMBO achieved the lowest average RMSE (8.03), followed closely by COMBO and MOPO. Model-free methods performed significantly worse, likely due to the difficulty of deviating from the behavior distribution in a static dataset for long-horizon torque correction.
Density: COMBO performed best (RMSE 0.691), with MOPO and PPO close behind. The smaller gap between model-based and model-free methods suggests the dataset provided sufficient coverage for density control.
Temperature: MOPO achieved the best performance (RMSE 0.193). Its ability to penalize uncertain model predictions proved advantageous for temperature dynamics, which involve delayed responses and accumulated transport effects.
Pressure: MOPO again led (RMSE 5198.6), followed by RAMBO. The strong coupling of pressure to density and temperature makes this the most challenging task, favoring uncertainty-aware methods.

General Trends:

Tracking near the plasma core ( $\psi_N \approx 0$ ) is consistently more difficult than at the edge due to larger profile magnitudes and temporal variations.
Pure imitation learning (GCIL) was insufficient for these long-horizon problems, yielding significantly higher errors than offline RL methods.
No single algorithm is universally optimal; performance varies based on the specific physical dynamics and coupling of the target profile.

Significance and Claims

The paper claims that RL4F provides a critical step toward accelerating progress in both nuclear fusion and offline RL algorithm development. By offering a standardized, data-driven environment, it enables researchers to develop and compare controllers without the immediate risks and costs of real-device experimentation.

The authors modestly note that while the benchmark is built on real DIII-D data, its generalizability to other tokamaks remains an open question requiring broader collaboration. They emphasize that RL4F serves as a controlled offline evaluation setting; performance here is evidence of algorithmic potential but does not substitute for safety-critical validation on physical devices. The benchmark addresses the specific challenges of nonlinear dynamics, long-horizon control, and partial observability inherent to fusion, positioning it as a rigorous testbed for the broader offline RL community.