RL4F — Offline Reinforcement Learning for Plasma Control in Nuclear Fusion

RL4F is an open-source benchmark for learning plasma controllers from offline data. Built around a dynamics ensemble trained on historical DIII-D tokamak discharges, it provides datasets, closed-loop evaluation environments, and baselines for rotation, density, electron temperature, and pressure profile tracking.

Our paper, “Offline Reinforcement Learning for Plasma Control in Nuclear Fusion: Codebase and Benchmark,” introduces a common testbed for comparing offline RL algorithms on a demanding physical control problem.

Yang Fu*, Haomin Bao*, Rohit Sonker, Xiaoyan Hu, Aravind Venugopal, Jeff Schneider, Jiayu Chen

_{*Equal contribution}

Paper · Project page · Code · Dataset · Dynamics models

Why plasma control needs an offline benchmark

A tokamak controller must regulate a plasma whose behavior is nonlinear, coupled, and sensitive to its operating conditions. Changing heating, torque, or fueling affects several quantities at once, and the response unfolds over time.

The goal is also spatial: a plasma profile describes how a quantity varies from the core to the edge. Holding one scalar near its target does not necessarily produce the desired profile throughout the plasma.

Developing a controller through trial and error on a real device is expensive and risky. Historical experiments offer a starting point, but algorithm development also needs shared datasets and evaluation conditions. RL4F supplies that common ground, so researchers can compare how their methods handle the same tracking tasks before considering real-machine experiments.

From experimental discharges to policy evaluation

RL4F starts with a 25-member ensemble of recurrent probabilistic neural networks (RPNNs) trained on roughly 18,000 historical DIII-D discharges, with measurements sampled every 20 ms. Each model predicts a distribution over the next plasma-state change from the current state and actuator inputs.

Training proceeds in two stages: first learn the predictive mean, then freeze the mean-prediction backbone and fit the variance head. The ensemble represents different plausible dynamics, while the variance head models uncertainty in each member’s predictions.

The released offline RL trajectories are model-generated. We initialize the reference model from experimental shot states, replay the recorded actuator schedules, and roll the model forward. This preserves the experimental action sequences while replacing the measured plasma evolution with simulated trajectories.

After filtering, the dataset contains 5,882 shots and 945,828 transitions, with fixed splits:

Model-based baselines fit their own dynamics models from these synthetic training trajectories. Trained policies are then evaluated in closed loop against the reference ensemble, separating the models used for policy learning from the benchmark environment.

Four tasks, one shared actuator space

The paper studies four full-profile tracking objectives:

Every task uses the same four actuation channels: neutral-beam power, neutral-beam torque, gas puffing, and electron-cyclotron heating power. The policy receives the current profile, current and future targets, and explicit tracking-error terms.

Targets change during each rollout, requiring the controller to respond to a new setpoint and maintain it over time. Tracking is measured across 33 radial locations, rather than judged only from a scalar summary.

What the benchmark reveals

We compare goal-conditioned imitation learning, model-free offline RL methods such as TD3BC, CQL, IQL, EDAC, and MCQ, and model-based methods including PPO, COMBO, MOPO, MOBILE, RAMBO, and BAMCTS.

Evaluation uses 300 held-out test shots with 10 random seeds per shot. The paper’s main results report the following best task-average tracking errors:

Lower is better. Each task has its own physical scale; these raw RMSE values should be compared within a task.

The strongest method changes with the control objective. RAMBO performs well on rotation but struggles on density and temperature. COMBO leads density tracking, while MOPO achieves the lowest average error on temperature and pressure. Model-based approaches obtain the best averages across these tasks, but using a learned model alone does not guarantee strong performance.

Try RL4F

The repository provides environment setup instructions, pre-trained dynamics models, algorithm entry points, and a unified visualization script.

1. Follow the README’s Quick Start to install the Conda environment and dependencies.
2. Download the processed datasets and dynamics models, then place them in the directories described in the README.
3. Run a temperature-control baseline from the source checkout:

python scripts/run_ppo.py --env profile_control_new --task temp

Training scripts load the matching task and algorithm defaults from offlinerlkit/configs/benchmark_configs.json. Command-line arguments override these defaults, and policy checkpoints and return plots are saved under logs/.

To inspect a trained policy, replace /path/to/your_run_dir with its run directory:

python -m visualization.profile_tracking_save --actor_path /path/to/your_run_dir --test

The script saves actuator trajectories, profile-tracking outputs, and tracking-error summaries under results/<task>/<algo>/<val|test>/<run>/<seed>/.

What remains open

RL4F measures control performance in a learned environment derived from DIII-D data. Its results do not establish how these policies will behave on a real tokamak or transfer to another device. Model error and the simulation-to-reality gap remain important questions for future experiments.

By releasing the datasets, reference models, code, and evaluation protocol, we aim to make those questions easier to study and give offline RL researchers a reproducible setting for developing better plasma controllers.

Citation

@misc{fu2026offline,
  title={Offline Reinforcement Learning for Plasma Control in Nuclear Fusion: Codebase and Benchmark},
  author={Yang Fu and Haomin Bao and Rohit Sonker and Xiaoyan Hu and Aravind Venugopal and Jeff Schneider and Jiayu Chen},
  year={2026},
  eprint={2606.07550},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2606.07550}
}