Each ensemble member of RWM-U independently predicts a Gaussian distribution over the next observation. The variance within each prediction captures aleatoric uncertainty, while the variance across ensemble means estimates epistemic uncertainty arising from limited or biased offline data. 

During policy training in imagination, MBPO-PPO penalizes high-risk imagined transitions using this uncertainty signal, balancing task performance against model confidence.

RWM-U builds on the same general, domain-agnostic neural simulator as RWM, but augments it with ensemble-based epistemic uncertainty estimation that remains consistent over long autoregressive rollouts.

In our evaluation suite, RWM-U is tested on both manipulation and locomotion tasks in fully offline settings, including Reach-Franka manipulation and velocity-tracking locomotion for ANYmal D and Unitree G1, with policies trained entirely from fixed datasets and no online interaction.

Across these environments, RWM-U produces uncertainty estimates that closely track long-horizon prediction error and enables uncertainty-penalized policy optimization to outperform uncertainty-unaware model-based and model-free baselines, demonstrating reliable long-horizon modeling and safe generalization under severe distribution shift.

Much of modern offline reinforcement learning is built around off-policy, value-function methods like CQL, IQL, and COMBO. On standard benchmarks these approaches work remarkably well, but robotics, especially legged locomotion, plays by different rules. In practice, the most reliable controllers for real robots today are trained with on-policy methods, and PPO in particular has emerged as the empirical workhorse for high-dimensional control. It scales cleanly with massively parallel simulation, optimizes stably over long horizons, and produces policies that reliably transfer to hardware. The underlying reasons for this gap are still not fully understood, and existing studies consistently show how difficult it is to tune SAC-based methods to achieve PPO-level stability in large-scale locomotion training. Despite continued research, there is still no reliable recipe for deploying SAC-style approaches in this domain.

If offline RL is to work on real robots, a different alignment of tools is required. Policies must be trained on-policy using long-horizon rollouts, and they must be able to leverage synthetic experience because no new real-world data is available. This makes model-based reinforcement learning not just attractive, but essential. A learned world model can expand a limited dataset into millions of imagined trajectories, enabling long-horizon credit assignment that would otherwise be impossible.

However, learned models are imperfect, and in fully offline settings there is no corrective feedback. Errors compound during long autoregressive rollouts, and PPO will optimize against hallucinated dynamics unless it is explicitly informed about what is unreliable. Prior offline model-based approaches typically avoid this issue by restricting rollouts to short horizons, relying on single-step predictions, or operating in simplified simulated domains. These design choices do not scale to real locomotion.

RWM-U is designed around this exact failure mode. By building on an autoregressive world model and explicitly estimating epistemic uncertainty, it tracks where predictions degrade as rollouts extend over tens or hundreds of steps. That uncertainty is propagated into policy optimization, allowing PPO to improve policies while remaining grounded in regions supported by data. The result is a fully offline, long-horizon, uncertainty-aware model-based RL pipeline that remains stable and deployable on real quadruped and humanoid robots.