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}
}

ECCV 2026 · Causal World Models

Ziyi Ding, Xianxin Lai, Weiyu Chen, Xiao-Ping Zhang, Jiayu Chen

Tsinghua Shenzhen International Graduate School · The University of Hong Kong · INFIFORCE Intelligent Technology

A world model should not only say what happens next. It should also know what would happen if something changed.

Latent world models compress images into compact states and roll them forward under actions. That recipe works well for factual prediction: given what actually happened, retrieve or generate the most likely future. But the same latent can be unreliable under an intervention, because features that are useful for prediction are not necessarily separated into causes and effects.

A counterfactual query is more demanding. If we move one object, change its velocity, or alter a mechanism, the model must update the variables downstream of that change while leaving unrelated factors alone. In an entangled latent space, the intervention often becomes a diffuse perturbation: too much changes, too little changes, or the effect travels along the wrong path.

Our question is simple: can we add explicit causal structure to an existing world model without redesigning the predictive backbone?

Attach causal structure where the world model already reasons

We introduce a CausalVAE-inspired structural branch directly in latent space. The original encoder still maps an observation to a latent state, and the original action-conditioned transition still predicts the next latent. Between them, our plug-in organizes latent factors through a learned DAG:

The key is modularity. The same structural branch can sit on AE, VAE, modular, or graph-network dynamics backbones. A differentiable acyclicity penalty keeps the learned adjacency close to a DAG, while a lightweight alignment head anchors latent coordinates to simulator state during training. Neither simulator state nor the alignment head is needed at inference.

A three-stage training strategy

This staged design lets the backbone first acquire predictive competence, gives the structural branch a stable representation to organize, and finally adapts the transition to the causally structured latent without retraining the entire system from scratch.

Intervene, re-simulate, then retrieve the right future

We evaluate four domains—Physics 3-body, Chemistry, 2D Shapes, and 3D Cubes—using state-, object-, and mechanism-level interventions. Each factual transition is paired with a re-simulated counterfactual target.

The evaluation asks a direct question: among candidate futures, does the model retrieve the one produced by the intervention? Factual H@1 and MRR track ordinary prediction, while CF-H@1 and CF-MRR isolate counterfactual retrieval.

The largest gains appear where ordinary dynamics are intervention-fragile

Across four benchmarks and eight paired backbones, the plug-in generally keeps factual retrieval competitive and improves CF-H@1 on 21 of 32 backbone–benchmark pairs. Physics shows the clearest effect: GNN–NLL rises from 11.0 to 41.0 CF-H@1, a gain of 30.0 points, while AE–NLL rises from 10.5 to 41.0, a gain of 30.5 points.

In the matched Physics comparison, the plug-in adds approximately 1M parameters.

The learned structure is also interpretable at the mechanism-trend level. On Physics, its strongest edges align with a first-order template derived from the locally linearized dynamics. This is evidence of meaningful physical interaction trends—not a claim that a local DAG exactly recovers the full nonlinear law.

Where the plug-in breaks: the host latent still matters

The gains are not universal. Chemistry exposes a capacity bottleneck: some backbones compress 5^5 = 3,125 global color configurations into a five-dimensional code. Passing that code through another structured encode–decode path can erase class-discriminative directions. In this setting, causal regularization has too little representational room to help.

For AE–Contrastive on Chemistry, increasing the latent dimension changes the CF-H@1 gain monotonically:

The practical lesson is specific: use the plug-in when the backbone latent can still carry the relevant causal factors, and validate it per domain. Long-horizon drift and strongly nonlinear regimes remain open problems.

The takeaway

Read the ECCV 2026 paper

Accepted to IEEE Robotics and Automation Letters, June 2026.

Wenqi Ge*, Junde Guo*, Zhen Fu, Shunpeng Yang, Jiayu Chen, Hua Chen

LimX Dynamics · The University of Hong Kong · SUSTech · HKUST · Zhejiang University

_{*Equal contribution}

Project page · arXiv:2606.27676 · Video (YouTube) · Video (Bilibili)

CWI on a full-size LimX Oli: squatting to pick up toys and drop them in a basket (a-1), standing up and walking off carrying it (a-2), opening a door (b), hanging cups on a rack (c), striking a suspended drum (d), assembling a small part (e), twisting a bottle cap (f).

A humanoid that picks a toy off the floor and carries it across the room is solving two problems at once. The legs have to hold a commanded velocity and height while the arms swing the base around. The arms have to reach wherever the operator points, with motion that stays inside what a human body would actually produce. Push hard on one objective and the other tends to give.

Human motion capture is the obvious source of priors for both halves. That is also where the trouble starts.

Where full-body MoCap breaks down

AMASS and similar corpora are not evenly useful across the body. Upper-body trajectories are rich: people reach, gesture, and handle objects in every direction, which is close to what a manipulation policy needs. Lower-body trajectories are another matter.

AMASS lower-body statistics against our training command range (blue). Pelvis height piles up near standing, 0.85 to 0.90 m, while we command down to 0.17 m. Base velocities cluster at the origin, while we command up to 1.0 m/s forward and 1.2 rad/s in yaw.

Two things go wrong. The clips do not cover the corners of the command box we need for deployment, so a policy trained to track them never learns to hold 0.3 m/s at a 0.4 m squat. And many of the clips that do move are too aggressive for a robot that is carrying something, so they have to be filtered out.

Both standard responses to this cost something. Filtering the corpus for lower-body stability throws away the upper-body motion in those same clips, which was perfectly good. Skipping MoCap for the legs and shaping rewards or sampling joint commands instead leaves the upper body without human statistics, and the arm motion drifts away from anything a person would produce.

Decouple the data, keep the policy whole

Several existing loco-manipulation controllers decouple the architecture: one policy for the legs and another for the arms, or RL for the legs and PD control on top for the arms. CWI keeps a single whole-body policy over all 31 DoF and decouples only what each half learns from.

The upper body trains on the full AMASS corpus with no filtering. The lower body trains on roughly ten curated expert-quality clips per class, one class for omnidirectional walking and one for squatting.

The CWI framework. Large-scale upper-body data drives manipulation, compact curated clips drive locomotion style through two discriminators, a multi-critic PPO stage trains the unified policy, and distillation produces a student driven by velocity/height commands plus bimanual hand poses.

The reason the upper-body corpus needs no filtering is a change of frame. We express every upper-body reference relative to the robot’s base. Once root motion is subtracted, a clip of someone spinning or jumping still contains a clean, human-feasible arm trajectory. The instability lives in the root, and the root is exactly what we drop.

Two details in how the body is partitioned turn out to matter later. The waist (yaw, roll, pitch) goes into the upper-body group, so torso bending is supervised by the same AMASS clip as the arms rather than being treated as part of the locomotion problem. The neck and head come out of the reference entirely, since AMASS contains very few head motions; instead the head gets a 2-DoF command sampled within its mechanical range each episode and a separate quadratic tracking term, which lets a VR headset drive it directly at deployment.

The upper-body reference combines joint targets q_u^ref ∈ R^17 with position, orientation, and keypoint-based pose for the shoulder, elbow, wrist, and hand links. The lower-body reference for the discriminators carries base velocity and height, the 12 leg joint positions and velocities, and hip/knee/ankle link states.

Two discriminators, one per gait class

The lower body gets its supervision through Adversarial Motion Priors rather than frame-by-frame tracking. Each discriminator scores a four-frame window of lower-body state transitions, with upper-body signals excluded so the prior stays purely about the legs, and trains with a least-squares GAN objective plus a gradient penalty.

Which discriminator is active depends on the command: D_walk when the velocity command is nonzero, D_squat when it is zero and the policy is tracking height. Both feed a continuous style reward that pays off states the discriminator classifies as expert-like.

The important consequence is that the curated clips act as a style prior, not a tracking target. Locomotion commands are sampled uniformly across the entire deployment range during training, far outside anything the ten clips demonstrate. The discriminator answers “does this look like expert walking”, not “does this match frame 47 of clip 3”, so command coverage and motion quality stop competing for the same data.

Walking and squatting are distinct enough that pooling them into one discriminator averages the two modes and blurs both, which the ablations below confirm.

Three reward groups, three critics

The full reward splits into locomotion tracking, upper-body tracking, and style:

Standard regularizers (projected-gravity alignment, action smoothness, torque penalties) are assigned to whichever group they primarily affect.

A single critic has to explain all three signals with one value function, and the adversarial style reward is high-variance enough to drag that estimate around. We give each group its own value head, compute group-wise GAE advantages, normalize each within the minibatch, and combine them with weights before the standard PPO clipped surrogate.

The multi-critic variant reaches a higher final reward and converges faster on both velocity and hand-position error.

Distilling down to two hand poses

Single-stage RL for whole-body tracking tends to come out jerky, so we train a teacher with privileged state and the full upper-body reference, then distill.

The student sees only a history of proprioceptive observations plus a reduced command: velocity and height for the legs, and an 18-dimensional bimanual hand pose (a 9D keypoint feature per hand) for the arms. Distillation combines PPO on the teacher’s reward with an action-matching BC term, querying the teacher on-policy over the student’s own state distribution and annealing the imitation weight.

Reducing the command to hand poses is a deployment decision. A VR operator supplies end-effector targets, not joint angles, and reconstructing joint targets through online IK risks discontinuities near singularities. Distilling directly on hand poses sidesteps that and keeps the student pinned to the AMASS manifold. The cost is a small open-loop accuracy gap against the joint-command interface.

How it compares

We evaluate on a 31-DoF LimX Oli (1.65 m, 50 kg), trained in IsaacLab with commands sampled over v_x ∈ [−0.5, 1.0] m/s, ω_yaw ∈ [−1.2, 1.2] rad/s, h ∈ [0.17, 0.9] m, and the usual domain randomization over mass, friction, sensor noise, and latency. The AMASS upper-body corpus is split 85/15, and evaluation uses sequences the policy never saw.

Because the original works target different robots, we re-implemented every baseline on Oli (marked with *) and unified observation composition, reward terms and weights, and PD gains across all methods, so the comparison isolates framework design and data usage.

Lower is better on every row except success rate. Position and orientation errors aggregate over all upper-body keypoints.

FALCON* carries higher arm tracking error, consistent with its separate upper and lower policies. HOVER* loses robustness once upper-body motions are paired with base commands outside its training distribution. HOMIE* shows larger joint-level deviation, which fits a PD-controlled upper body asked to follow non-smooth references.

The one metric where CWI does not come first is success rate, where FALCON* leads by about a tenth of a percentage point. Our policy holds strict upper-body tracking even at the boundary of the command box, and that occasionally costs a little stability.

Which pieces are doing the work

The ablations use the deployment-interface student, so the manipulation columns report end-effector error over the two hands rather than the full-keypoint aggregate above. We add a relaxed DTW distance d_dtw against reference clips as a naturalness measure, and the 95th-percentile joint acceleration q̈_p95 as a smoothness proxy.

_{† The single-stage variant fails upper-body tracking outright, so its locomotion numbers are reported for transparency but not ranked.}

Removing distillation is the sharpest failure. A single-stage policy cannot learn manipulation from the sparse hand-pose command alone, and end-effector error goes from 42.91 mm to 173.2 mm. The teacher has to learn whole-body control with dense references first; the student inherits it.

Dropping AMP barely touches command tracking but collapses gait quality: d_dtw jumps from 0.452 to 1.413 rad, with frequent stance violations and foot slips. Pooling the two discriminators into one is the same failure in milder form, with d_dtw at 0.615 and joint acceleration up from 28.8 to 35.1 rad/s². End-effector accuracy stays near 42.9 mm in both cases, so the dual-discriminator design buys lower-body style without charging manipulation for it.

The result we find most telling is the AMASS upper-body ablation. Swapping human references for randomly sampled joint targets degrades manipulation as expected (42.91 to 62.32 mm, joint acceleration 28.8 to 48.2 rad/s²), but it also degrades every locomotion metric: velocity error 0.100 to 0.131 m/s, yaw rate 0.183 to 0.249 rad/s, height 19.65 to 24.20 mm. The architecture supplies the structural decoupling. Only the data supplies upper-body motion whose statistics the legs can absorb, and a jerky upper body is a disturbance the legs then have to reject.

On the robot

Deployment runs off a Meta Quest headset and its hand controllers, with no MoCap suit and no exoskeleton cockpit. The operator’s hand poses and a velocity/height command are the entire interface.

The tasks in the opening figure cover precision work (twisting a bottle cap, assembling a small part), whole-body coordination (opening a door, striking a suspended drum), and long-horizon sequences (picking up toys, placing them, then carrying the basket away). The policy holds stable locomotion throughout.

Four phases of a teleoperated box lift: approaching and squatting (A), grasping and starting the lift (B), raising with both hands (C), standing up with the box overhead (D). Below, left and right hand distance from the base against waist pitch over time.

Watch the waist pitch. It rises as the hands extend forward and down, then straightens as the robot comes upright with the box. The operator never sent a torso command; there is no torso command in the interface. The coupling between hand reach and torso posture is in the AMASS clips because that is how people reach, and putting the waist in the upper-body group meant one clip supervised both. The composite objective keeps that coupling intact while the legs stay stable underneath it.

Limits and what comes next

Behavior coverage is bounded by the motion libraries and, more tightly, by the command interface. Bimanual hand poses plus velocity and height cannot express pressing a pedal with a foot or bracing against a surface with an elbow.

That same compactness is useful going forward. A portable teleoperation rig that works in ordinary rooms is a data collection instrument, and a small, task-oriented command space is structured supervision for higher-level autonomous policies, vision-language-action models included. Collecting multimodal demonstrations in everyday scenes is where we are taking this next.

Citation

@article{ge2026cwi,
  author  = {Ge, Wenqi and Guo, Junde and Fu, Zhen and Yang, Shunpeng
             and Chen, Jiayu and Chen, Hua},
  title   = {{CWI}: Composite Humanoid Whole-Body Imitation System
             for Loco-Manipulation},
  journal = {IEEE Robotics and Automation Letters},
  year    = {2026},
  doi     = {10.1109/LRA.2026.3710366}
}

• Paper (arXiv): arxiv.org/abs/2505.13709
• OpenReview: openreview.net/forum?id=eB5i7caric
• Project page: agentic-intelligence-lab.github.io/ROMBRL
• Code: github.com/Agentic-Intelligence-Lab/ROMBRL

Our ICML 2026 poster — click for the full-resolution PDF.

TL;DR

Offline model-based RL (MBRL) usually learns a world model and a policy in two separate stages — first fit the model to maximize data likelihood, then optimize the policy against the fixed model. This mismatch leaves policies brittle: even small amounts of deployment-time noise can wreck performance. We propose ROMBRL, which jointly adapts the world model with the policy under a single constrained maximin objective, solved via Stackelberg learning dynamics with the policy as the leader and the world model as an adversarial follower. ROMBRL comes with a formal suboptimality bound and achieves state-of-the-art robustness on D4RL MuJoCo and stochastic Tokamak Control benchmarks — with almost no cost to clean-environment performance.

Why Do Offline RL Policies Fall Apart Under a Little Noise?

Offline MBRL is appealing because it lets us train a policy from a static dataset by rolling it out inside a learned simulator, rather than the real environment — no risky exploration, no expensive interaction. But nearly all existing methods follow the same two-stage recipe: first fit a world model $P_\phi$ to maximize the likelihood of observed transitions, then freeze it (or only mildly adapt it) and optimize the policy $\pi_\theta$ against it.

That split has a quiet cost. The model is trained to explain the data, not to support good policy learning — an objective mismatch that’s easy to miss when you only evaluate in a clean, deterministic simulator. So we asked a more adversarial question: what happens to these policies once the environment doesn’t behave exactly like the simulator? We took a stable of strong offline RL baselines — CQL, EDAC, COMBO, MOBILE, and even RAMBO, which is specifically designed with an adversarial world model — and added just 5% Gaussian noise to state transitions at deployment time.

Average scores on nine D4RL MuJoCo tasks, before and after injecting 5% measurement noise at deployment. Every method — including the model-based robust baseline RAMBO — loses a substantial fraction of its performance.

Robustness, it turns out, is not something these methods get for free from good average returns.

The Trick: Turn Model Adaptation Into an Adversary’s Move

If the world model is going to be wrong somewhere, we’d rather it be wrong in a way our policy has already prepared for. That’s the intuition behind ROMBRL: instead of treating the world model as a fixed, passive stand-in for the environment, we make it move against the policy during training, so that whatever the policy learns has already survived a worst-case opponent.

Formally, we cast offline MBRL as a constrained maximin problem:

$$\max_\theta J(\theta, \phi’) \quad \text{s.t.} \quad \phi’ \in \arg\min_{\phi \in \Phi} J(\theta, \phi)$$

where $\Phi$ is an uncertainty set of world models $\phi \in \mathcal{M}$ anchored to the maximum-likelihood estimate $\hat\phi$ — a KL trust region around what the offline data actually supports:

$$\mathbb{E}_{(s,a)\sim\mathcal{D}}\left[ \mathrm{KL}\left(P_{\hat\phi}(\cdot \vert s,a) \,\Vert\, P_\phi(\cdot \vert s,a)\right) \right] \le \epsilon$$

(conservatism still matters offline, we’re not letting the model run wild). The policy maximizes its return against the worst model in that set; the model adversarially minimizes it.

We solve this as a Stackelberg game: the policy is the leader, the world model is the follower that best-responds to the policy — the mirror image of online MBRL, where the model is usually adapted to help the leader rather than oppose it. Using implicit differentiation through the follower’s best response, we derive primal-dual Stackelberg update rules for $(\theta, \phi, \lambda)$, and make them practical at scale with a few tricks:

• The Fisher Information Matrix stands in for expensive second-order terms, using $\mathbb{E}[\nabla^2 \log P_\phi] = -\mathbb{E}[\nabla \log P_\phi \nabla \log P_\phi^\top]$ to avoid an explicit Hessian.
• The Woodbury matrix identity inverts the resulting low-rank matrix in time linear (not cubic) in the parameter count.
• A gradient-mask mechanism keeps off-policy replay data from going stale as the model shifts under the policy — a subtlety that RAMBO’s simpler alternating-update scheme overlooks.

The full algorithm, ROMBRL, is summarized in Appendix J of the paper.

It’s not just an engineering trick

Theorem 3.1 gives a formal bound on the resulting policy’s suboptimality gap, assuming the true environment lies in the uncertainty set $\Phi$ with probability at least $1 - \delta/2$:

$$J(\theta^{\ast}, \phi^{\ast}) - J(\hat\theta, \phi^{\ast}) \;\le\; \frac{\sqrt{C}}{(1-\gamma)^2} \sqrt{4\epsilon + c\left(\sqrt{\frac{\log(2\vert\Phi\vert/\delta)}{N}} + \frac{\log(2\vert\Phi\vert/\delta)}{N}\right)}$$

where $N$ is the offline dataset size, $\vert\Phi\vert$ the covering number of the uncertainty set, $C$ a concentrability coefficient, and $\epsilon$ the uncertainty-set radius — so the gap shrinks as the dataset grows, giving a formal robustness guarantee rather than a purely heuristic one. Theorems 3.2 and 3.3 instantiate this bound for tabular MDPs and for continuous MDPs with Gaussian world models respectively, characterizing how $\epsilon$ should scale with dataset size and dimensionality.

From MuJoCo to a Real Fusion Reactor

Gaussian noise on MuJoCo is a convenient stress test, but it’s still synthetic. To find out whether any of this matters outside a benchmark, we went looking for a control problem where the “noise” isn’t something we inject — it’s baked into the physics.

An RL controller trained on a surrogate dynamics model drives plasma profiles toward a target using actuators such as power, torque, and ECH.

Tokamak plasma control was a natural fit: confining and shaping a superheated plasma is a genuinely stochastic process. We evaluated on tracking tasks (ion rotation, electron density, and $\beta_N$, a key economic indicator for fusion efficiency) built on an ensemble of recurrent dynamics models calibrated to real operational data from DIII-D, a working tokamak in San Diego.

This benchmark is considerably less forgiving than D4RL MuJoCo — several strong model-based baselines (RAMBO, MOBILE, and a Bayes-adaptive MCTS baseline, BAMCTS) that look competitive on MuJoCo degrade sharply once the dynamics get genuinely messy:

(negative tracking error, higher/less-negative is better; * marks the second-best result for each row)

ROMBRL ranked first on 2 of the 3 tracking targets, second on the third, and had by far the lowest variance across seeds of any method we tested. That gap is the clearest evidence we have that adversarial world-model adaptation isn’t just a MuJoCo party trick — it transfers to the kind of high-stakes, physically-grounded control problem where robustness actually matters.

Full Results

D4RL MuJoCo, under 5% measurement noise at deployment. ROMBRL ranks first on 7/12 tasks and second on 4, with the best average score by a wide margin (Cohen’s d ≥ 2 over every baseline indicates a large, statistically significant effect):

Robustness vs. clean-environment trade-off (OfflineRL-Kit protocol). ROMBRL matches the best clean-environment performance in the field while losing almost nothing once noise is introduced — where every other baseline gives up 11–27% of its performance:

RWRL-style deployment perturbations (dropped/stuck observations, body mass, friction, joint damping — applied only at evaluation, trained on clean hc-med-rep data). ROMBRL achieves the best score on every one of the 5 perturbation types and the smallest overall performance drop, beating dedicated robust baselines RORL and RFQI by a wide margin:

Ablation. Our full constrained Stackelberg update substantially outperforms both a naive alternating update (à la RAMBO) and an unconstrained Stackelberg update on walker2d-medium:

Robustness comes specifically from anticipating how the uncertainty-set boundary (via the dual variable λ) shifts with the policy — not just from anticipating the model itself.

Citation

The official PMLR proceedings for ICML 2026 (volume 306) have not been posted yet, so please cite the arXiv version for now — we’ll update this with the final proceedings entry once it’s available:

@article{chen2025rombrl,
  title   = {Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning},
  author  = {Chen, Jiayu and Xu, Le and Venugopal, Aravind and Schneider, Jeff},
  journal = {arXiv preprint arXiv:2505.13709},
  year    = {2025}
}

Acknowledgements

This work was funded in part by the Department of Energy Fusion Energy Sciences under grant DE-SC0024544.

Research Question

Can we establish a global convergence guarantee for online RLHF formulated as a bilevel optimization problem?

Proceedings · arXiv · Poster · MuJoCo Code · ** Code**

TL;DR

Our answer is conditional but positive. We first show that the original Self-Improving Alignment (SAIL) objective is guaranteed to have favorable curvature only near its initialization. We then introduce SAIL-RevKL, which adds

$$ \gamma\,\mathbb{E}_x!\left[ D_{\mathrm{KL}}!\left( \pi_{\mathrm{ref}}(\cdot\mid x)\,\Vert\,\pi_\theta(\cdot\mid x) \right) \right] $$

as a penalty. For a fixed reference and a log-linear policy, the Hessian of this reverse KL is the Fisher information matrix. With enough regularization, that positive curvature dominates the adverse curvature of SAIL throughout the full bounded parameter region. This yields global strong concavity and a Polyak-Lojasiewicz (PL) condition for the regularized surrogate, together with near-linear dependence on the inverse target accuracy.

Why Bilevel-Formulated Online RLHF Is Hard

Offline preference optimization trains on a fixed dataset. Online RLHF instead couples two learning problems: a reward model is fitted from preferences over responses generated by the current policy, and the policy is updated against that learned reward. This interaction is naturally formulated as a bilevel optimization problem because the solution of one level determines the objective and data distribution of the other.

SAIL uses the reward-policy equivalence to reduce this computationally expensive bilevel problem to an efficient single-level surrogate. If $\mathcal{D}_\theta$ denotes the policy-dependent response and preference distribution, that surrogate has the form

$$ J(\theta) = \mathbb{E}_{(x,y_w,y_l)\sim\mathcal{D}_\theta} !\left[ \log\sigma!\left( \beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\mathrm{SFT}}(y_w\mid x)} - \beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\mathrm{SFT}}(y_l\mid x)} \right) \right]. $$

This dependence matters. Standard performance-difference arguments assume that the reward or evaluation measure is independent of the policy appearing inside the advantage. That assumption no longer holds here, so convergence cannot be inherited directly from ordinary policy-gradient analysis. We instead analyze the geometry of $J$ itself.

The negative Hessian of $J$ decomposes into four terms. Some are stabilizing, while interaction terms can become unfavorable as $\theta$ moves away from the initial SFT parameters $\theta_0$. The proof must therefore control the objective’s curvature throughout the region the algorithm may visit, not only at its starting point.

Assumptions and Scope

The global theorem is explicit about its scope. It applies to the following setting:

Two construction choices are also essential. The reference policy $\pi_{\mathrm{ref}}$ is fixed during the analyzed update, and projected gradient ascent keeps every iterate inside $\Theta$, the radius-$B_\theta$ Euclidean ball centered at $\theta_0$.

Scope of the result. The word “global” below means global within this full bounded feasible set. The LoRA experiments test whether the method remains useful beyond these assumptions; the last-layer-only experiments are the closest empirical match to A1.

Core Theorem: Global Convergence of SAIL-RevKL

The original SAIL surrogate is guaranteed to have favorable geometry only near its initialization. We add a reverse-KL penalty and optimize

$$ J_\gamma(\theta) = J(\theta) - \gamma\,\mathbb{E}_x!\left[ D_{\mathrm{KL}}!\left( \pi_{\mathrm{ref}}(\cdot\mid x)\,\Vert\,\pi_\theta(\cdot\mid x) \right) \right]. $$

For a fixed reference and a log-linear policy, the reverse-KL Hessian is the Fisher information matrix, so

$$ -\nabla^2J_\gamma(\theta) = -\nabla^2J(\theta)+\gamma F(\theta). $$

Theorem (Global convergence on the feasible set)

Let $\Theta$ be the radius-$B_\theta$ Euclidean ball centered at $\theta_0$. Suppose A1-A4 hold, $\pi_{\mathrm{ref}}$ remains fixed during optimization, and projected updates keep every iterate in $\Theta$. Choose $\gamma$ sufficiently large to satisfy the explicit curvature threshold in Theorem 2 of the paper, so that the Fisher term dominates the adverse curvature of the original SAIL objective throughout $\Theta$.

Then $J_\gamma$ is $\mu$-strongly concave on all of $\Theta$ and satisfies the PL inequality

$$ \Vert\nabla J_\gamma(\theta)\Vert_2^2 \ge 2\mu\left(J_\gamma(\theta_\gamma^{\ast})-J_\gamma(\theta)\right), \qquad \theta\in\Theta, $$

where $\theta_\gamma^{\ast}=\arg\max_{\theta\in\Theta}J_\gamma(\theta)$ is the unique maximizer. Moreover, for projected mini-batch stochastic gradient ascent with a valid constant stepsize, there exist constants $q\in(0,1)$ and $C>0$ such that

$$ \mathbb{E}!\left[ J_\gamma(\theta_\gamma^{\ast})-J_\gamma(\theta_T) \right] \le q^T\left(J_\gamma(\theta_\gamma^{\ast})-J_\gamma(\theta_0)\right) + \frac{C\sigma_g^2}{B_s}. $$

Thus the optimization error decreases geometrically until it reaches the mini-batch noise floor. Taking $B_s=\widetilde{\mathcal{O}}(\varepsilon^{-1})$ and $T=\widetilde{\mathcal{O}}(\log\varepsilon^{-1})$ gives an $\varepsilon$ expected function-value gap with total sample complexity

$$ \widetilde{\mathcal{O}}!\left( \varepsilon^{-1}\log\varepsilon^{-1} \right). $$

The theorem is proved in a deliberately controlled log-linear setting. We next test whether the same regularized objective is useful when training real language models.

LLM Alignment Results

The main experiments use LoRA fine-tuning and GPT-4 as an LLM judge. PKU-SafeRLHF evaluates the balance between helpfulness and harmlessness, while UltraFeedback focuses on general instruction following. Pairwise win rate is the primary comparison below; tie rate and mean GPT score difference are reported separately because they capture different aspects of the judge outcomes.

PKU-SafeRLHF, Qwen 0.5B

UltraFeedback

The headline comparison is the pairwise win rate against vanilla SAIL: 43% vs. 29% on Qwen with PKU-SafeRLHF, 30% vs. 15% on Phi-3 with UltraFeedback, and 34% vs. 23% on LLaMA-3 with UltraFeedback.

A theory-aligned test

LoRA changes nonlinear features and is not a direct instantiation of the log-linear policy in A1. We therefore also freeze each backbone and train only its final linear layer on UltraFeedback U10. This isolates the regime that most closely matches the theory.

SAIL-RevKL improves the best win rate over vanilla SAIL on all three backbones. The preferred coefficient is scale-dependent: $10^{-2}$ for Qwen 0.5B and $10^{-1}$ for Phi-3 3.8B and LLaMA-3 8B. This is consistent with the theory’s curvature-bias trade-off and shows why $\gamma$ should be tuned rather than treated as a universal constant.

Citation

@inproceedings{wu2026convergence,
  title     = {On the Convergence of Self-Improving Online {LLM} Alignment},
  author    = {Wu, Xudong and Liu, Pangpang and Aggarwal, Vaneet and Chen, Jiayu},
  booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence},
  pages     = {7433--7467},
  year      = {2026},
  volume    = {337},
  series    = {Proceedings of Machine Learning Research}
}

Acknowledgements

This work was supported in part by the Seed Fund for PI Research, Basic Research, from The University of Hong Kong under Project Code 2502251784.

Source Code

Experiment Logs

Verlog is a multi-turn reinforcement learning framework built for long-horizon LLM-agentic tasks with highly variable episode lengths. Extending VeRL and BALROG while following the proven design principles of pytorch-a2c-ppo-acktr-gail, it introduces specialized optimizations for stable and efficient training when episodes span from short interactions to hundreds of turns. Whereas prior frameworks like VeRL and RAGEN effectively handle tasks with ~10 turns, and verl-agent scales up to 50 turns, Verlog is designed to operate in environments with over 400 turns, making it uniquely suited for complex, long-term decision-making. This capability has been validated across challenging domains such as BabyAI, BabaIsAI, and Crafter, where it consistently achieves strong performance out of the box. In Crafter, for instance, episode lengths range from 70 to 400 steps with an average of about 190.

Please kindly check https://wentsechen.github.io/Verlog_blogpost/ for details.

Jiayu Chen (陈佳玉) is an incoming Tenure-Track Assistant Professor at the Department of Data and Systems Engineering, The University of Hong Kong. Currently, he is a Postdoctoral Fellow at the School of Computer Science, Carnegie Mellon University, working under the supervision of Professor Jeff Schneider on data-driven profile control for nuclear fusion, as part of a joint research program with the Plasma Control Group, Princeton University.

Dr. Chen earned his Ph.D. in Industrial Engineering and Operations Research from the Edwardson School of Industrial Engineering, Purdue University and a Bachelor of Engineering from the College of Engineering, Peking University. During his Ph.D., he was co-advised by Professor Vaneet Aggarwal (Purdue University), an expert in theoretical reinforcement learning, and Professor Tian Lan (George Washington University), an expert in multi-agent systems.

He is the director of the Agentic Intelligence Lab, where his research focuses on Learning for Sequential Decision-Making and Continuous Control. Dr. Chen has published extensively in top venues such as NeurIPS, ICML, ICRA, ICAPS, AAMAS, TMLR, and IEEE TNNLS as the first author. His work has been recognized with prestigious awards, including the Oracle Research Award and the Purdue Research Grant.

About the University/Department

Founded in 1911, the University of Hong Kong (HKU) is the oldest tertiary education institution in Hong Kong, renowned globally for its excellence in teaching, research, and innovation. According to the QS World University Rankings, HKU is currently ranked 2nd in Asia and 17th worldwide. Studying at HKU offers unique advantages: its location in the vibrant, internationally connected city of Hong Kong provides unmatched access to global business networks, a multicultural environment, and close proximity to the thriving economies of the Greater Bay Area.

In the midst of the artificial intelligence and big data revolution, the Department of Data and Systems Engineering (DASE) upholds the philosophy of “computing-in-practice”. The vision is to bridge the cyber and physical worlds by integrating data science methodologies, such as artificial intelligence, with human tools like robots to create greater value for humanity. To achieve this goal, DASE focuses on three key areas: Data Engineering and Analytics, Human-Machine Cooperation, and Robotics and Automation.

Research Areas

The Agentic Intelligence Lab conducts broad research across Machine Learning, Robotics, and Operations Research. At present, our work focuses on advancing the following key areas:

• Theoretical foundations, algorithm design, and applications of (Multi-agent) Reinforcement Learning;
• Developing a unified theoretical framework for Reinforcement Learning, Optimal Control, and Stochastic Optimization;
• Scaling Laws in sequential decision-making, with a special emphasis on efficient (Multi-agent) Monte Carlo Tree Search algorithms and advanced Reasoning Models;
• Robotic Foundation Models;
• End-to-end multi-modal learning for Humanoid Robotic Control and its industrial applications;
• Post-training techniques for Large Language Model agents and their applications in sequential decision-making;
• Artificial General Intelligence grounded in the principles of The Alberta Plan.

Available Positions

We have the following openings:

• Fully funded Ph.D. positions starting in Fall 2025 have been taken.
• Openings for self-financed Ph.D. students and research assistants are available year-round.
• ❗❗❗ We have an available postdoc position in reinforcement learning and robotic learning.
• ❗❗❗ We need a remote research intern for a project to benchmark offline RL and control methods for plasma control in nuclear fusion. Please check this repo for details and reach out to Dr. Jiayu Chen if interested.

Ph.D. and postdoctoral applicants are expected to have strong proficiency in mathematics or programming, along with demonstrated self-motivation and resilience. Candidates with expertise in the following areas are highly preferred: Humanoid Robots, Large Language/Reasoning Models, Control Theory, Optimization Theory, Statistical Machine Learning.

We offer competitive remuneration. For Ph.D. students, the basic level of scholarship is 18,760 HK dollars per month. For postdocal researchers, the basic salary is no less than 30,000 HK dollars per month. Outstanding candidates are encouraged to apply for the HKU Presidential PhD Scholar Programme (HKUPS), Hong Kong PhD Fellowship Scheme (HKPFS), and RGC Postdoctoral Fellowship Scheme (PDFS) to get financial support of approximately 430,000 HK dollars annaly.

We also provide opportunities for visiting experiences at top U.S. universities such as UC Berkeley, UIUC, CMU, Purdue, and GWU. For research assistants, we offer strong recommendation letters and ample opportunities to publish papers in top machine learning and robotics venues.

For detailed requirements and additional information, please refer to the respective categories in openings. If interested, please contact Dr. Jiayu Chen directly through jiayuc@hku.hk. The Ph.D. application system at HKU accepts submissions year-round.

Deeply sorry that I am unable to respond to every inquiry unless it is to send an interview invitation.