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

# Offline Reinforcement Learning for Plasma Control in Nuclear Fusion: Codebase and Benchmark
Source: [https://arxiv.org/html/2606.07550](https://arxiv.org/html/2606.07550)
Yang Fu1,†\\daggerHaomin Bao2,†\\daggerRohit Sonker3Xiaoyan Hu3 Aravind Venugopal3Jeff Schneider3Jiayu Chen4,‡\\ddagger 1Central South University2Chongqing University3Carnegie Mellon University4The University of Hong Kong†\\daggerEqual contribution\.‡\\ddaggerCorresponding author\.

###### Abstract

https://github\.com/LucasCJYSDL/Offline\-RL\-Kit\-for\-Nuclear\-FusionJiayu Chen \(jiayuc@hku\.hk\) Offline reinforcement learning \(RL\) offers a promising route for developing plasma controllers from historical tokamak data, since online trial\-and\-error on real devices is costly and risky\. However, progress in this direction remains difficult to measure due to the lack of a standardized offline RL benchmark for realistic multi\-actuator, long\-horizon plasma control problems in nuclear fusion\. We introduce RL4F, an Offline Reinforcement Learning Benchmark for Plasma Control in Nuclear Fusion, providing closed\-loop evaluation environments and baseline comparisons across four full\-profile tracking tasks: rotation, density, temperature, and pressure\. The dynamics function underlying the evaluation environment is built from historical discharge data from DIII\-D, a real\-world Tokamak\. We evaluate a broad set of imitation learning and offline RL baselines under a unified protocol\. We find that offline model\-based RL methods obtain the best average performance on most objectives, although no single method dominates all tasks, highlighting the importance of dynamics modeling in complex, long\-horizon plasma control tasks\. To foster further research, we open\-source the codebase, datasets, and evaluation framework, providing a benchmark not only for the fusion community but also for algorithm development in offline RL\.

## 1Introduction

Nuclear fusion offers a potential route toward abundant, low\-carbon energy by harnessing the reactions that power stars\(Giet al\.,[2020](https://arxiv.org/html/2606.07550#bib.bib34)\)\. The tokamak is one of the most promising confinement devices for achieving controllable nuclear fusion, but its operation requires real\-time control of hot, unstable, and strongly coupled plasmas\. Recent work has shown that reinforcement learning \(RL\) can be used to train such controllers\. In particular,Degraveet al\.\([2022](https://arxiv.org/html/2606.07550#bib.bib5)\)demonstrated deep RL\-based magnetic control on the TCV tokamak, and subsequent studies have extended RL\-based plasma control toward tearing\-mode avoidance, profile tracking, and ramp\-down control\(Traceyet al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib6); Seoet al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib7); Charet al\.,[2023](https://arxiv.org/html/2606.07550#bib.bib9); Wanget al\.,[2025](https://arxiv.org/html/2606.07550#bib.bib8)\)\. These advances suggest that RL can complement conventional plasma\-control design by directly optimizing feedback policies in high\-dimensional tokamak control problems\.

Developing RL controllers directly on real tokamaks is difficult to scale: tokamak operation is expensive, time\-limited, and safety\-critical\. A natural way is to train candidate policies and evaluate their closed\-loop behavior in simulation before possible deployment\. Physics\-based simulators, such as RAPTOR\(Feliciet al\.,[2011](https://arxiv.org/html/2606.07550#bib.bib32)\), Forward Grad\-Shafranov Static \(FGE\) simulator\(Carpanese,[2021](https://arxiv.org/html/2606.07550#bib.bib14)\), and the more recent TORAX\(Citrinet al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib33)\), provide reliable tools for forward modeling, trajectory optimization, and controller development\. Previous work has demonstrated promising control performance for RL policies trained on these simulators\(Degraveet al\.,[2022](https://arxiv.org/html/2606.07550#bib.bib5); Traceyet al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib6)\)\. However, physics\-based simulators are significantly more computationally expensive than traditional RL simulators, especially when RL agents explore the actuator space stochastically during training, which often results in slower convergence for the iterative solver\. Moreover, adapting these simulators to a specific tokamak can require substantial modeling choices, parameter identification, and calibration\. An alternative approach is to learn control\-oriented dynamics models from historical experimental data, yielding simulation environments tied more directly to a particular device\(Charet al\.,[2023](https://arxiv.org/html/2606.07550#bib.bib9); Sonkeret al\.,[2026](https://arxiv.org/html/2606.07550#bib.bib41)\)\. Despite recent progress, RL\-based profile control remains challenging: plasma profiles are high\-dimensional spatial quantities, their dynamics are nonlinear and coupled across multiple actuators, and practical RL controllers still face difficulties in reward specification, steady\-state tracking bias, and sample efficiency\(Traceyet al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib6)\)\. These challenges motivate a standardized benchmark for data\-driven training and offline evaluation of RL algorithms for plasma control in tokamaks\.

Our primary contribution is RL4F, a unified benchmark for offline RL in tokamak profile control\. A profile refers to the radial spatial variation of a plasma quantity from the core, i\.e\., the innermost region of the plasma, to the edge, i\.e\., the outermost region\. It is characterized by four key features\. 1\)Realistic pre\-deployment workflow\.We first train a reference dynamics model from historical DIII\-D experimental discharges, and then use this model to generate trajectories for offline policy learning\. Candidate algorithms learn only from the model\-generated datasets and are evaluated in closed loop on the reference dynamics model, mirroring the practical setting in which controllers must be developed from past experimental data before any real\-machine test\. 2\)Multi\-task profile tracking\.The benchmark covers four full\-profile tracking tasks, including rotation, density, temperature, and pressure, which expose different control difficulties and plasma response channels\. 3\)Scenario\-relevant actuator space\.Policies act through a shared action space consisting of neutral\-beam power, neutral\-beam torque, gas puffing, and electron\-cyclotron heating, covering the main heating, momentum\-input, and fueling channels relevant to profile control for DIII\-D, a tokamak device located in San Diego\. 4\)Large\-scale data\.The benchmark contains 5,882 shots and 945,828 transitions after filtering, with fixed training, validation, and test splits\. To our knowledge, this is the first benchmark specifically designed for offline RL in fusion plasma control\.

We evaluate representative baselines spanning imitation learning, model\-free offline RL, and model\-based offline RL under the same closed\-loop protocol and profile\-level tracking metrics\. Our evaluation shows that model\-based methods generally outperform model\-free baselines, while no single algorithm dominates across all profile objectives, highlighting tokamak profile control as a challenging offline RL benchmark\.

## 2Related Work

Fusion Plasma Control\.Nuclear fusion is a leading candidate for sustainable power generation\. A central challenge is the control of plasma profiles to achieve stable, high\-performance operation\. Conventional plasma control typically relies on precomputed feedforward coil current trajectories\(Walker and Humphreys,[2006](https://arxiv.org/html/2606.07550#bib.bib1)\)together with feedback loops for individual target quantities\. Such profile control systems have been implemented on several tokamaks, including JET for safetyq\-profile control\(Moreauet al\.,[2003](https://arxiv.org/html/2606.07550#bib.bib2)\), TCV forq\-profile and electron\-temperature control\(Bartonet al\.,[2015](https://arxiv.org/html/2606.07550#bib.bib3)\), and EAST forq\-profile control\(Wanget al\.,[2021](https://arxiv.org/html/2606.07550#bib.bib4)\)\. These approaches have demonstrated promising performance in a wide range of discharges, but design can be challenging and time\-consuming, especially in plasma scenarios where the control quantities are high\-dimensional or strongly coupled\. More recently, reinforcement learning \(RL\) has emerged as a new framework for designing feedback controllers in fusion systems\.Degraveet al\.\([2022](https://arxiv.org/html/2606.07550#bib.bib5)\)train an RL\-based controller under the Maximum\-a\-Posteriori \(MPO\)Abdolmalekiet al\.\([2018](https://arxiv.org/html/2606.07550#bib.bib15)\)framework to track the location, current, and shape across a diverse set of plasma configurations\. This line of work is further extended byTraceyet al\.\([2024](https://arxiv.org/html/2606.07550#bib.bib6)\)\.Charet al\.\([2023](https://arxiv.org/html/2606.07550#bib.bib9)\)develop an offline RL framework for trackingβN\\beta\_\{N\}and plasma rotation quantities and train the controller using Proximal Policy Optimization \(PPO\), which is later utilized to design a Bayesian optimization \(BO\)\-style controller\(Sonkeret al\.,[2025](https://arxiv.org/html/2606.07550#bib.bib12)\)for mitigating tearing instabilities\.Seoet al\.\([2024](https://arxiv.org/html/2606.07550#bib.bib7)\)apply the Deep Deterministic Policy Gradient \(DDPG\)Lillicrapet al\.\([2016](https://arxiv.org/html/2606.07550#bib.bib16)\)approach to maintain high\-pressure plasma at the H\-mode performance while avoiding tearing instabilities\.Wanget al\.\([2025](https://arxiv.org/html/2606.07550#bib.bib8)\)trains RL policies to avoid disruptions during the ramp\-down phase\.

Benchmarks for Offline RL\.The well\-known D4RL dataset\(Fuet al\.,[2021a](https://arxiv.org/html/2606.07550#bib.bib10)\)benchmarks offline RL in challenging robotic control scenarios with biased data distributions\. The Dope benchmark\(Fuet al\.,[2021b](https://arxiv.org/html/2606.07550#bib.bib17)\), which builds upon D4RL and RL Unplugged\(Gulcehreet al\.,[2020](https://arxiv.org/html/2606.07550#bib.bib11)\), focuses on off\-policy evaluation\.Qinet al\.\([2022](https://arxiv.org/html/2606.07550#bib.bib13)\)propose NeoRL to mitigate the gap between earlier offline RL benchmarks and real\-world scenarios\.Liuet al\.\([2023](https://arxiv.org/html/2606.07550#bib.bib19)\)present a benchmarking suite that facilitates the development and evaluation of offline safe RL algorithms in both the training and deployment phases\.Parket al\.\([2025](https://arxiv.org/html/2606.07550#bib.bib18)\)propose OGBench for offline goal\-conditioned RL\. Compared with these existing benchmarks, fusion plasma control tasks present unique challenges, including highly nonlinear and stochastic dynamics, partial observability, and safety\-critical operating constraints\. To the best of our knowledge, RL4F is the first offline RL benchmark tailored to fusion plasma control, a critical real\-world operational scenario\.

## 3Simulator

Problem Setup\.We formulate tokamak profile control as a finite\-horizon Markov decision process \(MDP\)ℳ=\(𝒮,𝒜,P,r,ρ0,γ,H\)\\mathcal\{M\}=\(\\mathcal\{S\},\\mathcal\{A\},P,r,\\rho\_\{0\},\\gamma,H\), where𝒮\\mathcal\{S\}is the state space,𝒜\\mathcal\{A\}is the action space,P​\(s′\|s,a\)P\(s^\{\\prime\}\|s,a\)is the transition distribution,r​\(s,a\)r\(s,a\)is the reward function,ρ0\\rho\_\{0\}is the initial\-state distribution,γ\\gammais the discount factor, andHHis the episode horizon\. A policyπ​\(a\|s\)\\pi\(a\|s\)aims to maximize the expected discounted return

J​\(π\)=𝔼π,P,ρ0​\[∑t=0H−1γt​r​\(st,at\)\]\.J\(\\pi\)=\\mathbb\{E\}\_\{\\pi,P,\\rho\_\{0\}\}\\left\[\\sum\_\{t=0\}^\{H\-1\}\\gamma^\{t\}r\(s\_\{t\},a\_\{t\}\)\\right\]\.\(1\)
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/dynamic-training.png)Figure 1:The reference RPNN dynamics ensemble is trained from historical DIII\-D operational data and used as the closed\-loop environment for evaluating trained policies\.In offline reinforcement learning, the agent cannot interact with the real environment during training\(Levineet al\.,[2020](https://arxiv.org/html/2606.07550#bib.bib42)\)and is instead given a fixed dataset of trajectories𝒟=\{τi\}i=1N\\mathcal\{D\}=\\\{\\tau\_\{i\}\\\}\_\{i=1\}^\{N\}, where each trajectory consists of transitions\(st,at,rt,st\+1\)\(s\_\{t\},a\_\{t\},r\_\{t\},s\_\{t\+1\}\)collected by an unknown behavior policy\. This setting is well matched to tokamak control, where online trial\-and\-error is expensive and risky\.

We use a recurrent probabilistic neural network \(RPNN\) to learn plasma dynamics from trajectory data, following prior data\-driven tokamak dynamics modeling and offline model\-based control work\(Charet al\.,[2023](https://arxiv.org/html/2606.07550#bib.bib9); Sonkeret al\.,[2026](https://arxiv.org/html/2606.07550#bib.bib41)\)\. Given the current plasma statests\_\{t\}, the previous actuator settingat−1a\_\{t\-1\}, and the actuator incrementΔ​at=at−at−1\\Delta a\_\{t\}=a\_\{t\}\-a\_\{t\-1\}, the model predicts a distribution over the next state change,Δ​st=st\+1−st\\Delta s\_\{t\}=s\_\{t\+1\}\-s\_\{t\}\. Specifically, the RPNN outputs the parameters of a Gaussian distribution,μt,log⁡σt2=fθ​\(st,at−1,Δ​at\)\\mu\_\{t\},\\log\\sigma\_\{t\}^\{2\}=f\_\{\\theta\}\(s\_\{t\},a\_\{t\-1\},\\Delta a\_\{t\}\), and the next state is advanced autoregressively by adding the predicted state change to the current state\. This recurrent probabilistic formulation allows the model to capture history\-dependent plasma evolution while providing uncertainty estimates for long\-horizon rollouts\. A schematic of the dynamics\-model training workflow is shown in Figure[1](https://arxiv.org/html/2606.07550#S3.F1)\.

Dynamics Modeling\.Given the limited accessibility of real tokamak devices, we train a reference dynamics model as a digital twin\. The reference dynamics model is trained from historical DIII\-D experimental discharges and is used to generate offline training data and to provide the closed\-loop evaluation environment\. As shown in Figure[2](https://arxiv.org/html/2606.07550#S3.F2), we adopt a two\-stage RPNN training procedure\(Sonkeret al\.,[2026](https://arxiv.org/html/2606.07550#bib.bib41)\)\.

The first stage optimizes the predictive mean using a mean\-squared error objective\. The second stage initializes from the first\-stage checkpoint, freezes the predictive backbone, and trains the log\-variance head using a negative log\-likelihood objective, thereby calibrating uncertainty without changing the learned mean dynamics\. Following prior model\-based control practice\(Chenet al\.,[2025](https://arxiv.org/html/2606.07550#bib.bib44)\), we train a bootstrapped ensemble of RPNNs; the predicted log variance captures aleatoric uncertainty, while disagreement across ensemble members provides an estimate of epistemic uncertainty\. More training details are given in Appendix[B](https://arxiv.org/html/2606.07550#A2)\.

![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/2stage.png)Figure 2:Two\-stage training procedure for the dynamics model\.We train a 25\-member RPNN ensemble on roughly 18,000 historical DIII\-D experimental discharges, spanning nearly a decade of data collection\. Each shot contains approximately four seconds of data sampled at 20 ms\. For profile quantities, including electron temperature, ion temperature, density, pressure, rotation, and the safety\-factorqqprofile, we use ZipFIT reconstructions\(Loganet al\.,[2018](https://arxiv.org/html/2606.07550#bib.bib43)\), which provide smooth, physics\-constrained profile estimates\. Following prior tokamak dynamics modeling work\(Charet al\.,[2023](https://arxiv.org/html/2606.07550#bib.bib9)\), we reduce the dimension of profile quantities with PCA before dynamics\-model training\. The held\-out predictive fidelity of the trained ensemble is reported in Appendix[B\.1](https://arxiv.org/html/2606.07550#A2.SS1)\. To simulate a tokamak experiment, i\.e\., a “shot”, we use the ensemble of dynamics models\{fθi\}i=125\\\{f\_\{\\theta\_\{i\}\}\\\}\_\{i=1\}^\{25\}\. Each ensemble member represents a plausible version of the device dynamics, and the ensemble captures uncertainty in the true transition model\. Notably, RL policies trained on the same ensemble of dynamics models have been verified to be effective for profile control on DIII\-D\(Sonkeret al\.,[2026](https://arxiv.org/html/2606.07550#bib.bib41)\)\.

Synthetic Dataset for Offline RL\.The offline dataset is generated by initializing rollouts from real flat\-top shot states after a warm\-up period, replaying the real actuator sequence, and autoregressively rolling out the reference dynamics model\. At timesteptt, the measured next state is replaced by the ensemble prediction

Δ​s~t\(m\)∼𝒩​\(μt\(m\),diag⁡\(\(σt\(m\)\)2\)\),s^t\+1=s^t\+1M​∑m=1MΔ​s~t\(m\),M=25,\\tilde\{\\Delta s\}\_\{t\}^\{\(m\)\}\\sim\\mathcal\{N\}\\\!\\bigl\(\\mu\_\{t\}^\{\(m\)\},\\operatorname\{diag\}\(\(\\sigma\_\{t\}^\{\(m\)\}\)^\{2\}\)\\bigr\),\\qquad\\hat\{s\}\_\{t\+1\}=\\hat\{s\}\_\{t\}\+\\frac\{1\}\{M\}\\sum\_\{m=1\}^\{M\}\\tilde\{\\Delta s\}\_\{t\}^\{\(m\)\},\\quad M=25,\(2\)whereμt\(m\)\\mu\_\{t\}^\{\(m\)\}andσt\(m\)\\sigma\_\{t\}^\{\(m\)\}denote the mean and standard deviation predicted by themm\-th RPNN ensemble member at time steptt, andΔ​s~t\(m\)\\tilde\{\\Delta s\}\_\{t\}^\{\(m\)\}denotes a sampled state increment drawn from the corresponding Gaussian predictive distribution\. The generated states^t\+1\\hat\{s\}\_\{t\+1\}is then fed back as the input for the next timestep until the end of the shot segment\. This procedure preserves the actuator schedules and shot segmentation of the real experiments, while replacing the measured plasma evolution with model\-generated dynamics\.

We restrict the synthesized data to the flat\-top phase of each discharge and remove outlier values, resulting in 5,882 shots and 945,828 timesteps\. From these filtered shots, we randomly sample 300 shots for validation and 300 shots for testing, and use the remaining 5,282 shots for training\. This yields 849,977 training timesteps, 48,010 validation timesteps, and 47,841 test timesteps, with a maximum shot length of 180 timesteps after filtering\. Offline RL and imitation learning algorithms are trained only on the generated training dataset and evaluated on the reference dynamics model\.

## 4Benchmarking Tasks

We benchmark algorithms on full\-profile tracking tasks in tokamak plasma control\. Unlike scalar tracking, profile tracking requires regulating high\-dimensional spatial quantities across the normalized radial coordinate\. This is challenging for offline RL because plasma profiles evolve over long horizons under nonlinear, stochastic dynamics, while policies must infer corrective actions from a fixed dataset\.

In this work we address profile tracking across four plasma quantities \- rotation, density, electron temperature, and pressure\.Rotation profile controlis beneficial as differential rotation control impacts suppression of tearing instabilities\(Richneret al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib36)\)\. Moreover, edge rotation affects penetration of neutral gas, leading to asymmetric fueling\(Emdeeet al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib37); Wilkieet al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib38)\), a process crucial for modern high\-performance fusion scenarios because it helps manage plasma exhaust and maintain the high density required for fusion without damaging the reactor walls\.Density profile controlregulates the core fuel inventory and fusion reaction rate while avoiding excessive edge density, radiation cooling, confinement degradation, and density\-limit disruptions\. Experiments and density\-limit studies show that peaked or optimized density profiles can extend the operational range and support high\-density, high\-confinement tokamak plasmas\(Greenwald,[2002](https://arxiv.org/html/2606.07550#bib.bib39)\)\.Electron temperature profile controlis important because the profile gradient directly affects heat transport and regulation of this profile helps maintain core thermal energy and confinement while keeping the plasma within favorable regimes\(Chapman\-Oplopoiouet al\.,[2025](https://arxiv.org/html/2606.07550#bib.bib35)\)\. Finally,Pressure profile controlis important because the achievable normalized pressure and MHD stability limit depend strongly on the shape of the pressure and current\-density profiles\. A poorly placed pressure gradient can drive instabilities, while an optimized pressure profile helps sustain high stored energy without crossing stability limits\(Strait,[1994](https://arxiv.org/html/2606.07550#bib.bib40)\)\. The profile tracking tasks can be modeled as MDPs\. We define their observation space, action space, and reward function as follows, the effectiveness of which has been testified in real\-world fusion control experiments\(Sonkeret al\.,[2026](https://arxiv.org/html/2606.07550#bib.bib41)\)\.

Observation Space\.We include the present profile, the current target profile, and a future target profile shifted 10 timesteps ahead\. Following the PID\-inspired observation design used in practical tokamak RL control\(Traceyet al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib6)\), we also include proportional error terms \(P\-terms\) for both the current and future targets\. These terms are defined as the difference between the target profile and the present profile, and provide the policy with explicit feedback on the current tracking error\.

Target Sampling\.Randomly set targets may be infeasible within a given regime\. To ensure feasibility, we construct targets from reference trajectories by sampling two time points and forming a step function with three segments \(first–second–first\)\. This strategy ensures that all generated targets are consistent with the system dynamics while introducing sufficient variation to enrich the training distribution\.

Action Space\.The policy controls NBI power, NBI torque, gas\-puffing voltage, and total ECH power\. NBI torque directly influences the plasma rotation profile at the beam deposition location, which is typically situated near the plasma core\. Both torque and power originate from the same neutral beam source, playing a central role in sustaining the plasma\. We follow a counter\-beam configuration of neutral beams, which allows independent control of injected power and torque\.

Reward Function\.The reward function is the mean squared tracking error on the whole profile \(33 dimensions per timestep\), which is reconstructed from corresponding PCA components\.

rt=−‖ptarget​\(t\)−p​\(t\)‖22r\_\{t\}=\-\\left\\\|p\_\{\\mathrm\{target\}\}\(t\)\-p\(t\)\\right\\\|\_\{2\}^\{2\}\(3\)whereptarget​\(t\)p\_\{\\mathrm\{target\}\}\(t\)is the target profile at timett, andp​\(t\)p\(t\)is the actual profile\.

## 5Benchmarking Results

We benchmark diverse baselines on four profile\-tracking tasks\. All policies are trained on a static dataset synthesized by the reference dynamics model and evaluated in closed loop on the same reference model, as illustrated in Figure[1](https://arxiv.org/html/2606.07550#S3.F1)\. FollowingSonkeret al\.\([2026](https://arxiv.org/html/2606.07550#bib.bib41)\), evaluation rollouts use the mean prediction of the reference ensemble as the next\-state estimate, and performance is measured by profile\-tracking error on held\-out test shots\.

#### Baselines

The evaluated methods include model\-free offline RL baselines –TD3\+Behavior Cloning \(TD3BC\)\(Fujimoto and Gu,[2021](https://arxiv.org/html/2606.07550#bib.bib31)\), Conservative Q\-Learning \(CQL\)\(Kumaret al\.,[2020](https://arxiv.org/html/2606.07550#bib.bib26)\), Implicit Q\-Learning \(IQL\)\(Kostrikovet al\.,[2022](https://arxiv.org/html/2606.07550#bib.bib24)\), Ensemble\-Diversified Actor\-Critic \(EDAC\)\(Anet al\.,[2021](https://arxiv.org/html/2606.07550#bib.bib27)\), and Mildly Conservative Q\-Learning \(MCQ\)\(Lyuet al\.,[2022](https://arxiv.org/html/2606.07550#bib.bib25)\); model\-based offline RL baselines – PPO\(Schulmanet al\.,[2017](https://arxiv.org/html/2606.07550#bib.bib30)\), Conservative Offline Model\-Based Policy Optimization \(COMBO\)\(Yuet al\.,[2021](https://arxiv.org/html/2606.07550#bib.bib22)\), Model\-based Offline Policy Optimization \(MOPO\)\(Yuet al\.,[2020](https://arxiv.org/html/2606.07550#bib.bib29)\), Model\-Bellman Inconsistency \(MOBILE\)\(Sunet al\.,[2023](https://arxiv.org/html/2606.07550#bib.bib28)\), Robust Adversarial Model\-Based Offline Reinforcement Learning \(RAMBO\)\(Rigteret al\.,[2022](https://arxiv.org/html/2606.07550#bib.bib23)\), and Bayes Adaptive Monte Carlo Tree Search \(BAMCTS\)\(Chenet al\.,[2026](https://arxiv.org/html/2606.07550#bib.bib21)\); and Goal\-Conditioned Imitation Learning \(GCIL\)\(Dinget al\.,[2019](https://arxiv.org/html/2606.07550#bib.bib20)\)\. Full experiment details in Appendix[C](https://arxiv.org/html/2606.07550#A3)

Model\-based offline RL baselines train their own dynamics models from the synthetic trajectories generated by the reference model\. We use the same RPNN ensemble architecture and two\-stage training pipeline described in Section 3, again training a 25\-member ensemble\. The held\-out predictive fidelity of the learned dynamics models \(for offline RL\) is summarized in Appendix[B\.1](https://arxiv.org/html/2606.07550#A2.SS1)\. During model\-based RL training, for each rollout, we sample one ensemble member and then generate the trajectory autoregressively by sampling from that member’s predictive distribution, parameterized by its mean and log\-variance\. A new ensemble member is selected whenever a new rollout trajectory is sampled\. FollowingChuaet al\.\([2018](https://arxiv.org/html/2606.07550#bib.bib45)\), this procedure trains policies to perform across the range of plausible learned dynamics\.

#### Evaluation Metrics

We evaluate the baselines on each profile tracking task using closed\-loop simulation over 300 held\-out test shots, with 10 random seeds per shot\. At each timestep, the environment provides the target profile and the agent\-achieved profile\.

We compute the tracking error at each normalized radial locationψN\\psi\_\{N\}, whereNNranges uniformly from 0 \(core\) to 1 \(edge\)\. There are 33 radial locations in total\. The per\-location tracking metric is RMSE, computed from the squared deviation between the target and achieved profiles over the rollout horizon\. We report the mean RMSE and its standard error \(SE\) across all rollout instances at six selected radial locations\. Lower values indicate better tracking performance\.

RMSE¯​\(ψN\)=1N​S​∑i=1N∑j=1SRMSE\(i,j\)​\(ψN\),\\overline\{\\mathrm\{RMSE\}\}\(\\psi\_\{N\}\)=\\frac\{1\}\{NS\}\\sum\_\{i=1\}^\{N\}\\sum\_\{j=1\}^\{S\}\\mathrm\{RMSE\}^\{\(i,j\)\}\(\\psi\_\{N\}\),\(4\)whereNNis the number of test shots andSSis the number of random seeds per shot\.RMSE\(i,j\)​\(ψN\)\\mathrm\{RMSE\}^\{\(i,j\)\}\(\\psi\_\{N\}\)denotes the per\-rollout RMSE at radial locationψN\\psi\_\{N\}for shotiiand seedjj\.

#### Main Results

Tables[1](https://arxiv.org/html/2606.07550#S5.T1)–[4](https://arxiv.org/html/2606.07550#S5.T4)summarize the tracking performance of all baselines across the four tasks\. Across all tasks, offline RL methods achieve substantially lower tracking errors than the goal\-conditioned imitation learning baseline, indicating that purely imitation\-based learning is insufficient for these long\-horizon profile\-tracking problems\.

For rotation tracking, RAMBO achieves the lowest average RMSE \(8\.03\), narrowly outperforming COMBO and MOPO\. Model\-free baselines perform substantially worse\. Rotation profiles have pronounced radial variation from the plasma core to the edge and respond primarily through torque actuation, requiring sustained correction over long horizons after each target switch\. In this setting, conservative or behavior\-regularized model\-free methods are more restricted in their ability to deviate from the behavior distribution within the static training dataset when active correction is needed\.

Table 1:Rotation profile tracking error \(RMSE↓\\downarrow\)\. Values are reported as mean±\\pmstandard error over test shots\.For density tracking, COMBO achieves the lowest average RMSE of 0\.691, with MOPO \(0\.727\) and PPO \(0\.733\) close behind\. The relatively small gap between model\-based and model\-free methods suggests that the offline dataset provides sufficient coverage for learning effective density\-control behavior\. By contrast, RAMBO performs noticeably worse, with an average RMSE of 1\.583, indicating that robustness\-oriented model\-based training does not consistently improve performance across all profile control tasks\.

Table 2:Density profile tracking error \(RMSE↓\\downarrow\)\. Values are reported as mean±\\pmstandard error over test shots\.For temperature tracking, MOPO achieves the best average RMSE \(0\.193\), outperforming PPO \(0\.240\), MCQ \(0\.244\), and COMBO \(0\.264\)\. Temperature dynamics involve delayed responses and accumulated transport effects, making long\-horizon accuracy particularly important\. Methods that penalize uncertain model predictions \(MOPO\) are well suited to this regime, as they avoid compounding errors from unreliable model regions\.

Table 3:Temperature profile tracking error \(RMSE↓\\downarrow\)\. Values are reported as mean±\\pmstandard error over test shots\.Pressure tracking is likely the most strongly coupled task\. The pressure profile depends directly on both density and temperature, and it is also influenced by the current profile and actuator dynamics\. In this task, MOPO achieves the lowest average RMSE of 5198\.6, followed by RAMBO at 5358\.2\. Among model\-free methods, EDAC performs best with an average RMSE of 6016\.8\. These results suggest that uncertainty\-aware methods may be better suited to pressure tracking\.

Table 4:Pressure profile tracking error \(RMSE↓\\downarrow\)\. Values are reported as mean±\\pmstandard error over test shots\.Methods without such stabilization mechanisms tend to obtain higher errors\.

Across all the tasks, MOPO is the most robust method; COMBO is strongest on density; RAMBO excels on rotation and pressure but fails on density and temperature; PPO is a strong baseline on rotation, density, and temperature, but falls short on the more coupled pressure profile tracking task\. Thus, no single algorithm is universally beneficial across plasma profile control objectives, calling for stronger offline RL algorithms\.

![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/rot/187108_profile_actuator_combined.png)\(a\)Higher rotation target
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/rot/186906_profile_actuator_combined.png)\(b\)Lower rotation target

Figure 3:Simulated results of RAMBO applied to Rotation profile tracking for shot 187108 and 186906 using the dynamics\-model environment\. Two target patterns are tested: \(a\) increasing the profile and returning, and \(b\) decreasing the profile and returning\. Left plots show the rotation profile at different normalized flux values \(ψn\\psi\_\{\\mathrm\{n\}\}\)\. Right plots show the RL\-controlled actuator signals\. Both cases demonstrate strong tracking performance in the absence of the sim\-to\-real gap\.A second consistent trend is that tracking near the plasma core remains more difficult than tracking near the edge\. Errors at smaller normalized flux values are systematically higher regardless of the algorithm or task, because both the profile magnitude and temporal variation are substantially larger near the core\. This tendency is also qualitatively reflected in Appendix[A](https://arxiv.org/html/2606.07550#A1)\.

![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/dens/160579_profile_actuator_combined.png)\(a\)Higher density target
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/temp/186501_profile_actuator_combined.png)\(b\)Higher temperature target
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/pres/162961_profile_actuator_combined.png)\(c\)Higher pressure target

Figure 4:Simulated results for the higher\-target cases of the remaining three tasks in the dynamics\-model environment\. From top to bottom, the panels show COMBO on density profile tracking \(shot 160579\), MOPO on temperature profile tracking \(shot 186501\), and MOPO on pressure profile tracking \(shot 162961\)\.The small standard errors across all tables indicate that performance is consistent across test shots and that the observed gaps reflect systematic differences between methods rather than noise or a few favorable trajectories\.

The main text shows higher\- and lower\-target rotation rollouts \(Fig\.[3](https://arxiv.org/html/2606.07550#S5.F3)\), as well as higher\-target rollouts for density, temperature, and pressure \(Fig\.[4](https://arxiv.org/html/2606.07550#S5.F4)\)\. The corresponding lower\-target cases for density, temperature, and pressure are provided in Appendix[A](https://arxiv.org/html/2606.07550#A1)\. Across all tasks, the policies exhibit a consistent closed\-loop pattern: each target switch triggers an immediate coordinated response across all four actuators, after which the profile converges smoothly to the new setpoint\. The dominant actuation pathway differs by task: rotation is driven primarily by torque, density by gas puffing, temperature by co\-dominant ECH and beam power with active gas puffing suppression, and pressure by balanced coordination of all four actuators\. Overall, although there is a consistent offset in the tracking values, the controller responded appropriately to each step change, coordinating the actuators in a physically meaningful way\.

## 6Conclusion and Discussion

We introduced RL4F, an open\-source benchmark for offline RL in nuclear\-fusion plasma profile control, releasing reference dynamics models trained on real\-world fusion operational data, synthetic datasets for offline RL training, and the full codebase to support reproducibility and community adoption\. The benchmark provides standardized task definitions and a unified evaluation protocol across four profile\-tracking objectives: rotation, density, temperature, and pressure\. Physics\-based simulators such as RAPTOR\(Feliciet al\.,[2011](https://arxiv.org/html/2606.07550#bib.bib32)\)and TORAX\(Citrinet al\.,[2024](https://arxiv.org/html/2606.07550#bib.bib33)\)can struggle to capture complex fusion dynamics, and calibrating them to a specific tokamak can be nontrivial\. In contrast, our benchmark is built directly from experimental discharge data and designed specifically for offline RL evaluation\. To the best of our knowledge, it is the first offline RL benchmark for nuclear\-fusion plasma control, providing a common testbed for comparing algorithms under fixed datasets and standardized evaluation conditions\.

A current limitation is that RL4F is constructed from DIII\-D data alone\. As a result, it remains unclear how well the learned dynamics model, or the relative performance of different algorithms, transfers to other tokamak devices\. Addressing this limitation will require broader collaboration across the fusion community and the inclusion of data from additional fusion devices\. Despite these limitations, we believe this benchmark is a useful step toward accelerating progress in both nuclear fusion and RL algorithm development\.

## References

- Maximum a posteriori policy optimisation\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=S1ANxQW0b)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- G\. An, S\. Moon, J\. Kim, and H\. O\. Song \(2021\)Uncertainty\-based offline reinforcement learning with diversified Q\-ensemble\.InAdvances in Neural Information Processing Systems,Vol\.34,pp\. 7436–7447\.External Links:[Link](https://proceedings.neurips.cc/paper/2021/hash/3d3d286a8d153a4a58156d0e02d8570c-Abstract.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- J\. E\. Barton, W\. P\. Wehner, E\. Schuster, F\. Felici, and O\. Sauter \(2015\)Simultaneous closed\-loop control of the current profile and the electron temperature profile in the tcv tokamak\.In2015 American Control Conference \(ACC\),pp\. 3316–3321\.External Links:[Document](https://dx.doi.org/10.1109/ACC.2015.7171844)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- F\. Carpanese \(2021\)Development of free\-boundary equilibrium and transport solvers for simulation and real\-time interpretation of tokamak experiments\.Ph\.D\. Thesis,EPFL,Lausanne, \(en\)\.External Links:[Link](https://infoscience.epfl.ch/handle/20.500.14299/175817),[Document](https://dx.doi.org/10.5075/epfl-thesis-7914)Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p2.1)\.
- B\. Chapman\-Oplopoiou, J\. Walker, D\. R\. Hatch, T\. Görler, and J\. contributors \(2025\)Composition of electron temperature gradient driven plasma turbulence in jet\-ilw tokamak plasmas\.Phys\. Rev\. Res\.7,pp\. L012004\.External Links:[Document](https://dx.doi.org/10.1103/PhysRevResearch.7.L012004),[Link](https://link.aps.org/doi/10.1103/PhysRevResearch.7.L012004)Cited by:[§4](https://arxiv.org/html/2606.07550#S4.p2.1)\.
- I\. Char, J\. Abbate, L\. Bardoczi, M\. Boyer, Y\. Chung, R\. Conlin, K\. Erickson, V\. Mehta, N\. Richner, E\. Kolemen, and J\. Schneider \(2023\)Offline model\-based reinforcement learning for tokamak control\.InProceedings of The 5th Annual Learning for Dynamics and Control Conference,N\. Matni, M\. Morari, and G\. J\. Pappas \(Eds\.\),Proceedings of Machine Learning Research, Vol\.211,pp\. 1357–1372\.External Links:[Link](https://proceedings.mlr.press/v211/char23a.html)Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p1.1),[§1](https://arxiv.org/html/2606.07550#S1.p2.1),[§2](https://arxiv.org/html/2606.07550#S2.p1.1),[§3](https://arxiv.org/html/2606.07550#S3.p3.5),[§3](https://arxiv.org/html/2606.07550#S3.p6.2)\.
- J\. Chen, L\. Xu, W\. Chen, and J\. Schneider \(2026\)Bayes adaptive monte carlo tree search for offline model\-based reinforcement learning\.InInternational Conference on Learning Representations,Note:PosterExternal Links:[Link](https://openreview.net/forum?id=UmCzGf4qNh)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- J\. Chen, L\. Xu, A\. Venugopal, and J\. Schneider \(2025\)Policy\-driven world model adaptation for robust offline model\-based reinforcement learning\.External Links:2505\.13709,[Link](https://arxiv.org/abs/2505.13709)Cited by:[§3](https://arxiv.org/html/2606.07550#S3.p5.1)\.
- K\. Chua, R\. Calandra, R\. McAllister, and S\. Levine \(2018\)Deep reinforcement learning in a handful of trials using probabilistic dynamics models\.Vol\.31\.External Links:[Link](https://proceedings.neurips.cc/paper/2018/hash/3de568f8597b94bda53149c7d7f5958c-Abstract.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p2.1)\.
- J\. Citrin, I\. Goodfellow, A\. Raju, J\. Chen, J\. Degrave, C\. Donner, F\. Felici, P\. Hamel, A\. Huber, D\. Nikulin, D\. Pfau, B\. Tracey, M\. Riedmiller, and P\. Kohli \(2024\)TORAX: a fast and differentiable tokamak transport simulator in JAX\.External Links:2406\.06718,[Link](https://arxiv.org/abs/2406.06718)Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p2.1),[§6](https://arxiv.org/html/2606.07550#S6.p1.1)\.
- J\. Degrave, F\. Felici, J\. Buchli, M\. Neunert, B\. Tracey, F\. Carpanese, T\. Ewalds, R\. Hafner, A\. Abdolmaleki, D\. de las Casas, C\. Donner, L\. Fritz, C\. Galperti, A\. Huber, J\. Keeling, M\. Tsimpoukelli, J\. Kay, A\. Merle, J\. Moret, S\. Noury, F\. Pesamosca, D\. Pfau, O\. Sauter, C\. Sommariva, S\. Coda, B\. Duval, A\. Fasoli, P\. Kohli, K\. Kavukcuoglu, D\. Hassabis, and M\. Riedmiller \(2022\)Magnetic control of tokamak plasmas through deep reinforcement learning\.Nature602\(7897\),pp\. 414–419\.Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p1.1),[§1](https://arxiv.org/html/2606.07550#S1.p2.1),[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- Y\. Ding, C\. Florensa, P\. Abbeel, and M\. Phielipp \(2019\)Goal\-conditioned imitation learning\.InAdvances in Neural Information Processing Systems,Vol\.32,pp\. 15298–15309\.External Links:[Link](https://proceedings.neurips.cc/paper/2019/hash/c8d3a760ebab631565f8509d84b3b3f1-Abstract.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- E\. Emdee, L\. Horvath, A\. Bortolon, and G\. Wilkie \(2024\)The influence of rotation and sol drifts on poloidal asymmetries of pedestal fueling\.InAPS Division of Plasma Physics Meeting Abstracts,Vol\.2024,pp\. GO06–014\.Cited by:[§4](https://arxiv.org/html/2606.07550#S4.p2.1)\.
- F\. Felici, O\. Sauter, S\. Coda, B\. P\. Duval, T\. P\. Goodman, J\.\-M\. Moret, and J\. I\. Paley \(2011\)Real\-time physics\-model\-based simulation of the current density profile in tokamak plasmas\.Nuclear Fusion51\(8\),pp\. 083052\.External Links:[Document](https://dx.doi.org/10.1088/0029-5515/51/8/083052)Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p2.1),[§6](https://arxiv.org/html/2606.07550#S6.p1.1)\.
- J\. Fu, A\. Kumar, O\. Nachum, G\. Tucker, and S\. Levine \(2021a\)D4RL: datasets for deep data\-driven reinforcement learning\.External Links:2004\.07219,[Link](https://arxiv.org/abs/2004.07219)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p2.1)\.
- J\. Fu, M\. Norouzi, O\. Nachum, G\. Tucker, Z\. Wang, A\. Novikov, M\. Yang, M\. R\. Zhang, Y\. Chen, A\. Kumar, C\. Paduraru, S\. Levine, and T\. L\. Paine \(2021b\)Benchmarks for deep off\-policy evaluation\.External Links:2103\.16596,[Link](https://arxiv.org/abs/2103.16596)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p2.1)\.
- S\. Fujimoto and S\. S\. Gu \(2021\)A minimalist approach to offline reinforcement learning\.InAdvances in Neural Information Processing Systems,Vol\.34,pp\. 20132–20145\.External Links:[Link](https://proceedings.neurips.cc/paper/2021/hash/a8166da05c5a094f7dc03724b41886e5-Abstract.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- K\. Gi, F\. Sano, K\. Akimoto, R\. Hiwatari, and K\. Tobita \(2020\)Potential contribution of fusion power generation to low\-carbon development under the paris agreement and associated uncertainties\.Energy Strategy Reviews27,pp\. 100432\.External Links:[Document](https://dx.doi.org/10.1016/j.esr.2019.100432),[Link](https://doi.org/10.1016/j.esr.2019.100432)Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p1.1)\.
- M\. Greenwald \(2002\)Density limits in toroidal plasmas\.Plasma Physics and Controlled Fusion44\(8\),pp\. R27–R53\.Cited by:[§4](https://arxiv.org/html/2606.07550#S4.p2.1)\.
- C\. Gulcehre, Z\. Wang, A\. Novikov, T\. Paine, S\. Gómez, K\. Zolna, R\. Agarwal, J\. S\. Merel, D\. J\. Mankowitz, C\. Paduraru, G\. Dulac\-Arnold, J\. Li, M\. Norouzi, M\. Hoffman, N\. Heess, and N\. de Freitas \(2020\)RL unplugged: a suite of benchmarks for offline reinforcement learning\.InAdvances in Neural Information Processing Systems,H\. Larochelle, M\. Ranzato, R\. Hadsell, M\.F\. Balcan, and H\. Lin \(Eds\.\),Vol\.33,pp\. 7248–7259\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2020/file/51200d29d1fc15f5a71c1dab4bb54f7c-Paper.pdf)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p2.1)\.
- I\. Kostrikov, A\. Nair, and S\. Levine \(2022\)Offline reinforcement learning with implicit Q\-learning\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=68n2s9ZJWF8)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- A\. Kumar, A\. Zhou, G\. Tucker, and S\. Levine \(2020\)Conservative Q\-learning for offline reinforcement learning\.InAdvances in Neural Information Processing Systems,Vol\.33,pp\. 1179–1191\.External Links:[Link](https://proceedings.neurips.cc/paper/2020/hash/0d2b2061826a5df3221116a5085a6052-Abstract.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- S\. Levine, A\. Kumar, G\. Tucker, and J\. Fu \(2020\)Offline reinforcement learning: tutorial, review, and perspectives on open problems\.External Links:2005\.01643,[Link](https://arxiv.org/abs/2005.01643)Cited by:[§3](https://arxiv.org/html/2606.07550#S3.p2.2)\.
- T\. P\. Lillicrap, J\. J\. Hunt, A\. Pritzel, N\. Heess, T\. Erez, Y\. Tassa, D\. Silver, and D\. Wierstra \(2016\)Continuous control with deep reinforcement learning\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/1509.02971)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- Z\. Liu, Z\. Guo, H\. Lin, Y\. Yao, J\. Zhu, Z\. Cen, H\. Hu, W\. Yu, T\. Zhang, J\. Tan, and D\. Zhao \(2023\)Datasets and benchmarks for offline safe reinforcement learning\.External Links:2306\.09303,[Link](https://arxiv.org/abs/2306.09303)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p2.1)\.
- N\. C\. Logan, B\. A\. Grierson, S\. R\. Haskey, S\. P\. Smith, O\. Meneghini, and D\. Eldon \(2018\)OMFIT tokamak profile data fitting and physics analysis\.Fusion Science and Technology74\(1\-2\),pp\. 125–134\.Cited by:[§3](https://arxiv.org/html/2606.07550#S3.p6.2)\.
- J\. Lyu, X\. Ma, X\. Li, and Z\. Lu \(2022\)Mildly conservative Q\-learning for offline reinforcement learning\.InAdvances in Neural Information Processing Systems,Vol\.35,pp\. 1711–1724\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2022/hash/0b5669c3b07bb8429af19a7919376ff5-Abstract-Conference.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- D\. Moreau, F\. Crisanti, X\. Litaudon, D\. Mazon, P\. D\. Vries, R\. Felton, E\. Joffrin, L\. Laborde, M\. Lennholm, A\. Murari, V\. Pericoli\-Ridolfini, M\. Riva, T\. Tala, G\. Tresset, L\. Zabeo, K\.D\. Zastrow, and contributors to the EFDA\-JET Workprogramme \(2003\)Real\-time control of the q\-profile in jet for steady state advanced tokamak operation\.Nuclear Fusion43\(9\),pp\. 870\.External Links:[Document](https://dx.doi.org/10.1088/0029-5515/43/9/311),[Link](https://doi.org/10.1088/0029-5515/43/9/311)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- S\. Park, K\. Frans, B\. Eysenbach, and S\. Levine \(2025\)OGBench: benchmarking offline goal\-conditioned RL\.InThe Thirteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=M992mjgKzI)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p2.1)\.
- R\. Qin, X\. Zhang, S\. Gao, X\. Chen, Z\. Li, W\. Zhang, and Y\. Yu \(2022\)NeoRL: a near real\-world benchmark for offline reinforcement learning\.InAdvances in Neural Information Processing Systems,S\. Koyejo, S\. Mohamed, A\. Agarwal, D\. Belgrave, K\. Cho, and A\. Oh \(Eds\.\),Vol\.35,pp\. 24753–24765\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2022/file/9cd828eb8dc81a84fb6bf89a94263e1b-Paper-Datasets_and_Benchmarks.pdf)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p2.1)\.
- N\. Richner, L\. Bardóczi, J\. Callen, R\. La Haye, N\. Logan, and E\. Strait \(2024\)Use of differential plasma rotation to prevent disruptive tearing mode onset from 3\-wave coupling\.Nuclear Fusion64\(10\),pp\. 106036\.Cited by:[§4](https://arxiv.org/html/2606.07550#S4.p2.1)\.
- M\. Rigter, B\. Lacerda, and N\. Hawes \(2022\)RAMBO\-RL: robust adversarial model\-based offline reinforcement learning\.InAdvances in Neural Information Processing Systems,Vol\.35,pp\. 16082–16097\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2022/hash/6691c5e4a199b72dffd9c90acb63bcd6-Abstract-Conference.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. Klimov \(2017\)Proximal policy optimization algorithms\.External Links:1707\.06347,[Link](https://arxiv.org/abs/1707.06347)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- J\. Seo, S\. Kim, A\. Jalalvand, R\. Conlin, A\. Rothstein, J\. Abbate, K\. Erickson, J\. Wai, R\. Shousha, and E\. Kolemen \(2024\)Avoiding fusion plasma tearing instability with deep reinforcement learning\.Nature626\(8000\),pp\. 746–751\.Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p1.1),[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- R\. Sonker, A\. Capone, A\. Rothstein, H\. J\. F\. Kaga, E\. Kolemen, and J\. Schneider \(2025\)Multi\-timescale dynamics model bayesian optimization for plasma stabilization in tokamaks\.InForty\-second International Conference on Machine Learning,External Links:[Link](https://openreview.net/forum?id=d7v7RVXbNH)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- R\. Sonker, H\. J\. F\. Kaga, J\. Chen, A\. Rothstein, I\. Char, R\. Shousha, E\. Kolemen, and J\. Schneider \(2026\)Offline reinforcement learning for rotation profile control in tokamaks\.Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p2.1),[§3](https://arxiv.org/html/2606.07550#S3.p3.5),[§3](https://arxiv.org/html/2606.07550#S3.p4.1),[§3](https://arxiv.org/html/2606.07550#S3.p6.2),[§4](https://arxiv.org/html/2606.07550#S4.p2.1),[§5](https://arxiv.org/html/2606.07550#S5.p1.1)\.
- E\. Strait \(1994\)Stability of high beta tokamak plasmas\.Physics of Plasmas1\(5\),pp\. 1415–1431\.Cited by:[§4](https://arxiv.org/html/2606.07550#S4.p2.1)\.
- Y\. Sun, J\. Zhang, C\. Jia, H\. Lin, J\. Ye, and Y\. Yu \(2023\)Model\-Bellman inconsistency for model\-based offline reinforcement learning\.InProceedings of the 40th International Conference on Machine Learning,A\. Krause, E\. Brunskill, K\. Cho, B\. Engelhardt, S\. Sabato, and J\. Scarlett \(Eds\.\),Proceedings of Machine Learning Research, Vol\.202,pp\. 33177–33194\.External Links:[Link](https://proceedings.mlr.press/v202/sun23q.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- B\. D\. Tracey, A\. Michi, Y\. Chervonyi, I\. Davies, C\. Paduraru, N\. Lazic, F\. Felici, T\. Ewalds, C\. Donner, C\. Galperti, J\. Buchli, M\. Neunert, A\. Huber, J\. Evens, P\. Kurylowicz, D\. J\. Mankowitz, M\. Riedmiller, and The TCV Team \(2024\)Towards practical reinforcement learning for tokamak magnetic control\.Fusion Engineering and Design200,pp\. 114161\.External Links:[Document](https://dx.doi.org/10.1016/j.fusengdes.2024.114161),[Link](https://doi.org/10.1016/j.fusengdes.2024.114161)Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p1.1),[§1](https://arxiv.org/html/2606.07550#S1.p2.1),[§2](https://arxiv.org/html/2606.07550#S2.p1.1),[§4](https://arxiv.org/html/2606.07550#S4.p3.1)\.
- M\. L\. Walker and D\. A\. Humphreys \(2006\)Valid coordinate systems for linearized plasma shape response models in tokamaks\.Fusion Science and Technology50\(4\),pp\. 473–489\.External Links:[Document](https://dx.doi.org/10.13182/FST06-A1271),[Link](https://doi.org/10.13182/FST06-A1271),https://doi\.org/10\.13182/FST06\-A1271Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- A\. M\. Wang, C\. Rea, O\. So, C\. Dawson, D\. T\. Garnier, and C\. Fan \(2025\)Active ramp\-down control and trajectory design for tokamaks with neural differential equations and reinforcement learning\.Communications Physics8\(1\),pp\. 231\.External Links:[Document](https://dx.doi.org/10.1038/s42005-025-02146-6),[Link](https://doi.org/10.1038/s42005-025-02146-6)Cited by:[§1](https://arxiv.org/html/2606.07550#S1.p1.1),[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- Z\. Wang, H\. Wang, E\. Schuster, Z\. Luo, Y\. Huang, Q\. Yuan, B\. Xiao, and D\. Humphreys \(2021\)Optimal shaping of the safety factor profile in the east tokamak\.In2021 IEEE Conference on Control Technology and Applications \(CCTA\),pp\. 63–68\.External Links:[Document](https://dx.doi.org/10.1109/CCTA48906.2021.9658937)Cited by:[§2](https://arxiv.org/html/2606.07550#S2.p1.1)\.
- G\. Wilkie, F\. Laggner, R\. Hager, A\. Rosenthal, S\. Ku, R\. M\. Churchill, L\. Horvath, C\. S\. Chang, and A\. Bortolon \(2024\)Reconstruction and interpretation of ionization asymmetry in magnetic confinement via synthetic diagnostics\.Nuclear Fusion64\(8\),pp\. 086028\.Cited by:[§4](https://arxiv.org/html/2606.07550#S4.p2.1)\.
- T\. Yu, A\. Kumar, R\. Rafailov, A\. Rajeswaran, S\. Levine, and C\. Finn \(2021\)COMBO: conservative offline model\-based policy optimization\.InAdvances in Neural Information Processing Systems,Vol\.34,pp\. 28954–28967\.External Links:[Link](https://proceedings.neurips.cc/paper/2021/hash/f29a179746902e331572c483c45e5086-Abstract.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.
- T\. Yu, G\. Thomas, L\. Yu, S\. Ermon, J\. Y\. Zou, S\. Levine, C\. Finn, and T\. Ma \(2020\)MOPO: model\-based offline policy optimization\.InAdvances in Neural Information Processing Systems,Vol\.33,pp\. 14129–14142\.External Links:[Link](https://proceedings.neurips.cc/paper/2020/hash/a322852ce0df73e204b7e67cbbef0d0a-Abstract.html)Cited by:[§5](https://arxiv.org/html/2606.07550#S5.SS0.SSS0.Px1.p1.1)\.

## Appendix AAdditional Results

We provide additional visualizations of closed\-loop profile tracking trajectories for each of the four benchmark tasks\. Figures[6](https://arxiv.org/html/2606.07550#A1.F6)\-[9](https://arxiv.org/html/2606.07550#A1.F9)show the temporal evolution of the full radial profiles at selected time instances for two representative test shots per task, corresponding to a higher and a lower target profile, respectively\. In each figure, solid lines denote the model\-predicted profile and dashed lines indicate the target profile\.

![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/dens/176849_profile_actuator_combined.png)\(a\)Lower density target
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/temp/190783_profile_actuator_combined.png)\(b\)Lower temperature target
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/pres/171975_profile_actuator_combined.png)\(c\)Lower pressure target

Figure 5:Simulated results for the lower\-target cases of the remaining three tasks in the dynamics\-model environment\. From top to bottom, the panels show COMBO on density profile tracking \(shot 176849\), MOPO on temperature profile tracking \(shot 190783\), and MOPO on pressure profile tracking \(shot 171975\)\.![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/rot/187108_tracking_quantities_with_units.png)\(a\)Higher rotation target
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/rot/186906_tracking_quantities_with_units.png)\(b\)Lower rotation target

Figure 6:Temporal evolution of full rotation profiles at selected time instances for two representative shots using RAMBO\. \(a\) Higher rotation target and \(b\) Lower rotation target\. Solid lines represent the present \(predicted\) profiles, while dashed lines indicate the target profiles\.![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/dens/160579_tracking_quantities_with_units.png)\(a\)Higher density target
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/dens/176849_tracking_quantities_with_units.png)\(b\)Lower density target

Figure 7:Temporal evolution of full density profiles at selected time instances for two representative shots using COMBO\. \(a\) Higher density target and \(b\) Lower density target\. Solid lines represent the present \(predicted\) profiles, while dashed lines indicate the target profiles\.![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/temp/186501_tracking_quantities_with_units.png)\(a\)Higher temperature target
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/temp/190783_tracking_quantities_with_units.png)\(b\)Lower temperature target

Figure 8:Temporal evolution of full temperature profiles at selected time instances for two representative shots using MOPO\. \(a\) Higher temperature target and \(b\) Lower temperature target\. Solid lines represent the present \(predicted\) profiles, while dashed lines indicate the target profiles\.![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/pres/162961_tracking_quantities_with_units.png)\(a\)Higher pressure target
![Refer to caption](https://arxiv.org/html/2606.07550v1/Figure/pres/171975_tracking_quantities_with_units.png)\(b\)Lower pressure target

Figure 9:Temporal evolution of full pressure profiles at selected time instances for two representative shots using MOPO\. \(a\) Higher pressure target and \(b\) Lower pressure target\. Solid lines represent the present \(predicted\) profiles, while dashed lines indicate the target profiles\.
## Appendix BRPNN

### B\.1RPNN Training

The benchmark uses two learned dynamics models: a reference dynamics model trained on historical DIII\-D experimental trajectories, and a separate dynamics model trained on the synthesized benchmark trajectories for model\-based baselines\. Both models use the same RPNN ensemble architecture described in Appendix[B\.2](https://arxiv.org/html/2606.07550#A2.SS2); they differ only in the data used for training\. Each model predicts transitions over the full dynamics state–actuator space, which is larger than the policy observation and action space\. The corresponding signal usage is summarized in Table[7](https://arxiv.org/html/2606.07550#A2.T7): variables marked as Actuators and States are used for dynamics modeling, while only a task\-dependent subset of the state variables is exposed to the policy as observations\.

We report held\-out predictive fidelity using per\-variable explained variance \(EV\)\. The reference dynamics model, trained on historical experimental data and used as the closed\-loop evaluation environment, is summarized in Table[5](https://arxiv.org/html/2606.07550#A2.T5)\. The dynamics model trained on synthesized benchmark trajectories, used by model\-based baselines for learned rollouts, is summarized in Table[6](https://arxiv.org/html/2606.07550#A2.T6)\.

Table 5:Per\-variable explained variance \(EV\) of the reference dynamics model trained on historical experimental data\. Higher values indicate better agreement with held\-out real trajectories\.Table 6:Per\-variable explained variance \(EV\) of the dynamics model trained on synthesized benchmark trajectories, evaluated on held\-out data\. Higher values indicate better predictive accuracy for model\-based baselines\.Table 7:Plasma signals and how they are used as state and actuator variables for dynamics modeling\. Policy observation and action space variables are also shown\. Profile dimensionality depends on the task: rotation, density, and temperature use 4 PCA components, while pressure uses 2 PCA components\.Signal GroupSignalsActuatorStateActionObservationScalar StatesβN\\beta\_\{N\}\(Normalized Plasma Pressure\)✗✓✗✗lil\_\{i\}\(Internal Inductance\)✗✓✗✗Line Averaged Density✗✓✗✗Loop Voltage✗✓✗✗MHD Stored Energy✗✓✗✗Profile StatesRotation✗✓✗✓\(rot task\)Density✗✓✗✓\(dens task\)Ion Temperature✗✓✗✗Electron Temperature✗✓✗✓\(temp task\)Pressure✗✓✗✓\(pressure task\)Safety Factorqq✗✓✗✗Shape VariablesElongation✓✗✗✗Upper TriangularityBottom Triangularityaminora\_\{\\text\{minor\}\}Radial and vertical positions of magnetic axisNeutral Beam VariablesPower Injected✓✗✓✗Torque Injected✓✗✓✗Gas PuffingGasA voltage✓✗✓✗Electron Cyclotron HeatingECH Total Power✓✗✓✗Other ActuatorsCurrent Target, Toroidal Field✓✗✗✗TargetsRotation Target\(t\)\(t\),Rotation Target\(t\+10\)\(t\+10\),✗✗✗✓Error Terms\(t\)\(t\)Error Terms\(t\+10\)\(t\+10\)Total Dimensions12D25D4D20D/10D\(pres task\)
### B\.2Network Architecture and Training Details

#### Network Architecture

- •Encoder: - –Fully Connected \(FC\) layer:input\_dim×512\\texttt\{input\\\_dim\}\\times 512 - –FC layer:512×512512\\times 512
- •Memory Unit: - –Gated Recurrent Unit \(GRU\) block:512×256512\\times 256
- •Decoder\(with residual connections between FC layers\): - –FC layer:256×512256\\times 512 - –FC layers:512×512512\\times 512\(repeated 8 times\) - –FC layer:512×128512\\times 128
- •Output Heads: - –Mean head:128×output\_dim128\\times\\texttt\{output\\\_dim\} - –Log\-variance head:128×output\_dim128\\times\\texttt\{output\\\_dim\}

## Appendix CImplementation Details

### C\.1Hyperparameter

We separate hyperparameter selection from final evaluation\. For each task and algorithm, hyperparameters were selected using the validation split, while the held\-out test split was used only for the final closed\-loop evaluation reported in the main results\. This protocol avoids tuning directly on the test shots and provides a less biased estimate of generalization performance\. Because the four profile\-control tasks differ in their response channels and long\-horizon dynamics, we allow a small set of task\-dependent hyperparameters, but keep all other settings shared across tasks or fixed by the released configuration files\.

Table[8](https://arxiv.org/html/2606.07550#A3.T8)summarizes the common training, evaluation, network, and model\-rollout settings used across the benchmark\. These settings define the default experimental budget and architecture choices\. Method\-specific deviations are limited to the hyperparameters listed in Table[9](https://arxiv.org/html/2606.07550#A3.T9)\.

For hyperparameter tuning, we used an Optuna\-based tuner implemented in the released codebase\. The tuner reads each algorithm’s search space from the configuration file, launches training trials with sampled hyperparameters, and maximizes the validation reward extracted from the training logs\. The default sampler is Optuna’s TPE sampler, with support for grid search when specified in the configuration\. Trials can be run in parallel across multiple GPUs, and the tuner records the best trial, selected parameters, log directory, and detailed trial results\.

Table[10](https://arxiv.org/html/2606.07550#A3.T10)reports the task\-dependent hyperparameters selected by this validation procedure\. These values were fixed before evaluating on the test split\. Parameters not shown in the table are either shared across all four tasks, listed in Table[8](https://arxiv.org/html/2606.07550#A3.T8), or specified in the released per\-method configuration files\.

Table 8:Common experimental settings used in the benchmark\. We report the main shared settings here and provide complete per\-method configurations in the released codebase\.CategorySettingValueNotesTrainingOptimization budgetGradient updates1000epochs,1000steps/epochUsed by most offline RL baselines\.Mini\-batch sizeDefault batch size256GCIL and PPO use method\-specific batch sizes\.Discount factorgamma0\.99Used by most offline RL baselines; PPO use separate settings\.Target updatetau0\.005Used by actor\-critic offline RL baselines\.EvaluationTest setHeld\-out shots300shotsFixed test split shared across all tasks\.RandomizationSeeds per shot10Used for closed\-loop policy evaluation\.MetricTracking errorRMSE↓\\downarrowReported with standard error across rollout instances\.Evaluation episodesDefault value5PPO uses3evaluation episodes\.Policy and value networksDefault MLPHidden dimensions\[256,256\]Used by several actor\-critic and model\-based baselines\.Larger MLPHidden dimensions\[256,256,256\]Used by COMBO, CQL, EDAC, GCIL, and MPPI\.PPO architecturePolicy/value networks\[256,256\]PPO uses separate policy and value networks\.Model\-based rolloutsRollout batch sizerollout\_batch\_size50000Used by model\-based offline RL baselines\.Model retentionmodel\_retain\_epochs5Controls the retained model\-generated replay buffer\.Rollout schedulerollout\_freq1000Default rollout frequency for model\-based methods\.Dynamics ensembleNumber of models25Used for learned dynamics\-model rollouts and closed\-loop evaluation\.Table 9:Algorithms evaluated in the benchmark and the hyperparameters tuned for each method\. BAMCTS is implemented asbambrlin the codebase\.CategoryAlgorithmTuned hyperparametersImitation LearningGCILbatch sizeModel\-freeoffline RLTD3BCα\\alphaCQLCQL weight, temperatureIQLexpectile, temperatureEDACnumber of critics,η\\etaMCQλ\\lambda, sampled actionsModel\-basedoffline RLPPOclip rangeCOMBOrollout length, CQL weightMOPOrollout length, penalty coef\.MOBILErollout length, penalty coef\.RAMBOrollout length, adv\. weightBAMCTSrollout length, penalty coef\.,use\_ba,search\_alpha∈\{0\.5,0\.8\}\\in\\\{0\.5,0\.8\\\}Table 10:Task\-dependent hyperparameters\. All other hyperparameters are shared across the four benchmark tasks or specified in the released configuration files\.
### C\.2Experiments compute resources

All experiments were conducted on an Ubuntu 22\.04 Linux server equipped with two Intel Xeon Platinum 8457C processors, providing 96 physical CPU cores \(192 threads\) in total, 503 GiB RAM, and 10 NVIDIA GeForce RTX 4090 GPUs\. Each GPU provides approximately 46\-49 GiB of memory\. Unless otherwise specified, each training run used a single GPU\. Training time varies across algorithm classes\. Model\-based methods generally require longer training due to learned\-model rollouts, with representative runs ranging from several to about 30 GPU\-hours \(MCQ: 3\.7h, MOPO: 11h, BAMCTS: 12\.0h, MOBILE: 17h, RAMBO: 24\.5h, COMBO: 1\.2 days\), whereas model\-free and imitation\-learning methods are comparatively faster, typically completing within minutes to 18 GPU\-hours \(GCIL: 43 min, IQL: 1\.8h, CQL: 18h, EDAC: 8h, TD3\-BC: 6h\)\.

## Appendix DBroader Impacts

RL4F is an open\-source benchmark for offline reinforcement learning in nuclear\-fusion plasma profile control, enabling researchers to develop data\-driven control algorithms and compare them under standardized tasks, datasets, dynamics\-model environments, and evaluation protocols\. By supporting reproducible offline evaluation, RL4F can help reduce the need for costly and safety\-critical trial\-and\-error on real tokamak devices during early\-stage controller development\. Progress on this problem may contribute to more effective regulation of plasma profiles related to confinement, fueling, heating, and overall plasma performance, which is an important step toward stable and efficient fusion operation\. More broadly, the benchmark provides a challenging real\-world testbed for offline RL research, with nonlinear dynamics, long\-horizon control, partial observability, and limited data coverage\. As with any simulation\-based benchmark in a safety\-critical domain, performance in RL4F should be interpreted as evidence from a controlled offline evaluation setting rather than as a substitute for real\-device validation; any potential deployment on physical tokamaks would require additional safety constraints, expert review, and device\-specific testing\.