Latest / Embodied AI 101 / Imagining Locomotion: Learning a Neural World Model for Legged Robots

Transcript

Imagining Locomotion: Learning a Neural World Model for Legged Robots · Embodied AI 101

• Imagining Locomotion: Learning a Neural World Model for Legged Robots. Robotic World Model: A Neural Network Simulator for Robust Policy Optimization in Robotics • Chenhao Li, Andreas Krause, Marco Hutter • ETH Zurich (Robotic Systems Lab) • ArXiv 2025 (EWM NeurIPS’25 Workshop). Legged robots (quadrupeds, humanoids, etc.) are incredibly difficult to control in the real world. State-of-the-art controllers often rely on hand-designed simulators or physics engines to train policies, then simply deploy them on hardware. This pipeline – build a high-fidelity sim, train a policy until it works, and drop it on the robot – can work in the ideal case, but it fundamentally stops learning once the robot hits the real world. Any mismatch between sim and reality (rough terrain, unexpected contacts, modeling errors) can render the controller brittle. Ideally, we would like robots to keep adapting with real sensor data, not just blunder along on a “frozen” controller designed for lab conditions. The challenge is that collecting new real-world data is slow, expensive, and can easily damage the robot. The Robotic Systems Lab at ETH Zurich tackles this problem by learning a neural “world model” – effectively a learned simulator – that lets one train locomotion policies entirely in imagination. Their new Robotic World Model (RWM) takes as input a history of recent robot observations and actions and predicts the next observation. Crucially, it is trained autoregressively and self-supervised over many steps, so that its multi-step rollouts remain close to reality instead of drifting off into fantasy. Once learned, RWM can be used as a drop-in simulator: one can run a model-free RL algorithm like PPO on thousands of imagined trajectories generated by the RWM, without any additional real-robot trials. In practice this hybrid approach – a learned neural simulator plus a model-free RL policy – achieves much better long-horizon accuracy and sim-to-real transfer than previous model-based or purely model-free methods on quadruped and humanoid locomotion tasks. In fact, Li et al. demonstrate that policies trained with RWM zero-shot transfer to real hardware (ANYmal D and a Unitree G1 humanoid) and robustly track velocity commands even under disturbances. In this deep dive, we’ll unpack how RWM works and why it helps. We’ll explain the dual-autoregressive training that stabilizes prediction, the policy optimization loop that uses RWM as a “neural simulator,” and the experiments showing strong cross-robot generalization and sim-to-real performance. Armed with this understanding (and the open-source code in Isaac Lab), roboticists can consider applying RWM-style world models for their own legged and humanoid robots. The Challenge: World Models and Long-Horizon Errors. Reinforcement learning has given us powerful model-free methods (SAC, PPO, etc.) that can learn locomotion controllers, but these methods are data-hungry. Each new task or environment typically requires many hours of real robot interaction – and naive exploration can easily damage hardware. Model-based RL promises to reduce data needs by learning a dynamics model of the robot, but traditional approaches have faltered in practice. Why is learning a neural world model for robotics so hard? One big reason is compounding error in long rollouts. Imagine we train a one-step predictive model of the robot (e.g. by feeding in state and action and regressing to the next state). Even if that model has low error one step ahead, a small bias will accumulate over many steps. Model-free policy methods like PPO often require hundreds of timesteps to estimate value functions and rewards. If our learned model “hallucinates” slightly, PPO will happily optimize a policy that exploits those hallucinations. As the authors of RWM note, “PPO needs long-horizon trajectories… If a learned model starts hallucinating during autoregressive rollouts, those errors compound over tens or hundreds of steps, and PPO will happily optimize against a fantasy world”. In other words, a model-free controller training inside an inaccurate learned simulator can go off the rails. In fact, many prior world-model methods cope by building in strong inductive biases or hand-crafted features (e.g. learning explicit contact models for legs). But these “helper assumptions” tend to specialize to one robot or terrain. RWM boldly does the opposite: it uses a simpler end-to-end neural model without task-specific bias, but trains it specifically to be stable on long rollouts. The key insight is to force the world model to practice its own predictions during training, so it learns to correct itself. That way, when we roll out imagined trajectories, the model stays close to reality and yields reliable rewards. Another important challenge is partial observability and stochasticity. A legged robot’s dynamics can be partially observed (contacts are not directly seen by the controller) and inherently noisy. RWM handles this by predicting not just a single next state but a distribution over possible next observations. Training via self-supervision on real robot data (or simulated data from very accurate simulator) ensures RWM captures the true variability in the robot’s behavior under commands. These ideas – long-horizon stability and stochastic modeling – are the heart of RWM’s approach. Below we explain the RWM architecture and training, then show how it plugs into policy learning, and finally summarize the performance results on multiple robots and tasks. The Robotic World Model (RWM). At its core, the Robotic World Model is a recurrent neural network (a GRU) that predicts future observations of the robot given its recent history of observations and actions. Concretely, at each timestep the model takes as input the last N states (observations) and actions, and outputs the predicted next state (and optionally other signals). These “states” can include the robot’s joint positions/velocities, base orientation, etc. Importantly, RWM is model-agnostic – it could be a quadruped’s sensor readings or a humanoid’s – and it also outputs “privileged” signals such as contact flags or termination signals if available in training (the policy itself need not see these). This means RWM can use extra simulator-only info to learn a more accurate model, but still only feeds realistic observations to the policy. Dual-Autoregressive GRU Architecture. Unlike one-step predictors, RWM is explicitly designed to be used in fully autoregressive rollout. The training is implemented as a dual-autoregressive process. What does “dual-autoregressive” mean? It refers to two nested loops in how the GRU is updated during training:. Inner loop: We feed in the actual past observations and actions from the data (ground truth history) into the GRU, updating its hidden state “through history.” This is the usual way a GRU ingests a sequence of real data. Outer loop: We then let the GRU predict many future steps in imagination, feeding each predicted observation back as input into the GRU for the next step. In effect the model practices doing rollouts on its own predictions. During training, RWM has to minimize error both in the one-step predictions and over these multi-step imagined rollouts. In practice, this looks like a truncated backpropagation through time that spans many steps: the model tries to correct its own mistakes as it generates the next observation, then uses that frame to try to predict the following one, and so on. By training on its own multi-step errors, RWM “knows” to keep its wanderings bounded. As the authors summarize, RWM’s architecture uses a GRU with dual-auto-regressive training (inner updates through real history, outer feedback of predictions). The output layers (“heads”) simultaneously predict normal observations and any privileged signals (rewards, contacts, etc.) in a single forward pass. Because the model inherently rolls forward on its own outputs, it is far better at long-horizon consistency. In fact, the training explicitly targets long-horizon stability, rather than trying to tweak the policy to overcome model flaws. Once trained, rolling out RWM for hundreds of steps produces trajectories that stay very close to ground truth (see results below). This contrasts with naive models (e.g. a feed-forward MLP, or even a recurrent state-space model) which typically diverge or collapse when you iterate them. Self-Supervised Training. The training of RWM is fully self-supervised. That is, it minimizes the discrepancy between its predicted observations and the real observations that actually happened, over the entire rollout horizon. There is no teacher signal other than “make the predicted trajectory match the real one.” Concretely, one can imagine collecting a dataset of robot trajectories (states, actions, and maybe privileged info) by running some policy (even a random or simple controller) on the robot (or a high-fidelity sim). Then we train RWM on that dataset by letting it predict statet+1 from state0.t and actiont, computing a loss, then unrolling further in imagination and continuing to match reality up to a chosen horizon. Thanks to this self-supervised multi-step loss, RWM learns the stochastic dynamics of the robot. It naturally captures uncertainties (e.g. small variations in behavior) by, for example, outputting a probabilistic distribution (Gaussian ensemble or similar) rather than a single deterministic vector. In fact, part of the credit for RWM’s robustness is that each imagined rollout spawns a cloud of possible trajectories, akin to domain randomization. In practice the implementation supports this by an ensemble of RWM networks to sample different outcomes. The key is that during training the model sees the real data, so it heads off any insane divergences; during rollout it produces variations that help the policy generalize. In the project’s own words: “Because RWM predicts a stochastic distribution over next observations, each imagination rollout naturally introduces varied dynamics, effectively giving domain randomization for free and improving hardware generalization”. No matter what robot or task we use, the training procedure is identical. This is powerful: the authors report that the same architecture and pipeline works across manipulation and locomotion tasks for many robots, including both quadrupeds and humanoids. In fact, the exact same RWM model (dimensions, layers, loss weights) was trained on tasks from simple reach-and-move tasks on robot arms to velocity-tracking tasks on ANYmal, Cassie, Spot, Unitree robots, and even Genesis humanoids (the H1/G1). No task-specific tweaks were needed – the world model simply learns the underlying physics as observed. Integrating RWM with Policy Learning. Once we have a trained RWM, how do we use it to train a locomotion policy? The idea is to treat the world model as a substitute simulator: we generate imagined experience under RWM, then apply standard reinforcement learning on that data. Concretely, Li et al. adopt a variant of model-based policy optimization. They call it MBPO-PPO: a combination of the MBPO (Model-Based Policy Optimization) framework and the PPO (Proximal Policy Optimization) algorithm. (MBPO usually uses SAC as the inner loop, but here they stick with PPO for stability on robotics.). The training proceeds in two alternating phases:. Pretrain the world model. First, run some policy (say a basic random exploration or a previously learned controller) in the real simulator or robot to collect a dataset of transitions {(s,a,s')}. Train RWM on this data via multi-step autoregressive updates as described above. This gives an initial dynamics model. Iterate model and policy training. Now repeat: a) Policy rollout in RWM – starting from real initial states, roll out the current policy entirely inside the learned RWM for many steps (100+ steps). Record the imagined (s,a,r,s') sequences. b) Policy update – use PPO (on the imagined data) to update the policy network. c) Real-data collection – run the updated policy in the real simulator/hardware (safely in Isaac Lab) for a short burst to collect a few new real transitions. d) Model update – augment the dataset with these new samples and further update the RWM weights autoregressively. Then loop back to (a). Because RWM rollouts stay on track, PPO can accumulate meaningful long-horizon returns on them. One important detail is that the policy’s reward function is taken straight from the environment, and RWM is trained to predict the same reward. Therefore, PPO sees a reward map that is nearly the same in imagination as in the real world. (In practice RWM can directly output the predicted one-step reward, or the environment can recompute it from the predicted state.). A nice aspect of this approach is that the policy algorithm remains model-free – it’s just PPO (or in principle could be SAC, DDPG, etc.) on what it thinks of as a simulator. All the model-based magic happens in how we generate the data. In fact, the authors emphasize that we are effectively doing RL “in imagination” with minimal changes to the standard loop. As they put it, RWM lets policies train in imagination: “It is a learned black-box neural network simulator that predicts future observations… letting policies train ‘in imagination’ instead of requiring real interactions”. Avoiding the Fantasy-Loop Trap. This scheme only works if imagined rollouts are believable enough that PPO doesn’t go crazy optimizing phantoms. Here’s where RWM’s stable long-horizon prediction really pays off. In preliminary experiments, the authors observed that if the model begins to drift, PPO will exploit the resulting “tap” in the reward landscape and do useless or dangerous things. But with RWM, this failure mode is largely suppressed. In practice the learned RWM is fed every updated policy’s actions and it consistently produces trajectories whose predicted states and rewards remain aligned with reality over 100+ steps. One practical way they check this is to monitor the prediction error and reward error over time in rollouts. Thanks to the autoregressive training, the average model error does not blow up; instead it actually decreases as the model sees more data. The project page notes: “Despite PPO’s tendency to exploit model errors, RWM’s rollout fidelity keeps model error decreasing and predicted rewards aligned with ground truth, enabling stable optimization over 100+ autoregressive steps.”. In other words, RWM ensures that PPO “sees the truth” far enough into the future that it learns a robust strategy for the real robot. Offline vs. Online Model Use. There are two variants one can use once RWM is trained. In the first (online) mode, the loop above continues to collect new data: the policy occasionally goes back to the real (simulator) environment to gather fresh transitions for refining the RWM, just like MBPO. In the second (offline) mode, one freezes RWM after pretraining and never collects new real data during PPO training. Both modes are supported in the codebase. The first mode is what the main RWM paper focuses on: policy learning in “imagination online” with occasional real updates. The offline mode (addressed in a follow-up paper on “Uncertainty-Aware RWM”) uses the static model only. For the experiments on ANYmal D and G1 hardware transfer, either mode could be used, but the key findings come mostly from the online alternation strategy. Evaluation: Long-Horizon Prediction & Robustness. The central claim of the RWM work is that this scheme yields a more faithful long-horizon world model than prior approaches, across a wide range of robots. The authors test this extensively. First, they measure purely the prediction accuracy of RWM rollouts versus baselines. To do this, they collect a large dataset of trajectories (e.g. velocity commands and resulting motions) for each robot/task, then hold out a test set. They train RWM and several baselines (a simple feed-forward MLP predictor, a stochastic RSSM model as in Dreamer, and a transformer model) on the training data, and then compute the multi-step rollout error on the test set. Crucially, at rollout time they allow each model to autoregressively feed its own predictions (even for the baselines) and compare to ground truth. The results are striking: RWM-AR (the autoregressively trained RWM) consistently has the lowest error. In their suite of 40+ different tasks – manipulating objects with various arms and locomoting under velocity commands with dozens of different legged robots – RWM’s error is the smallest by a significant margin. For example, on the ANYmal quadruped (model D) rolling at various commanded speeds, RWM’s predicted trajectory stays extremely close to the true trajectory for 200 steps, whereas the MLP baseline diverges badly (see Figure 3 in the paper). Similarly, adding noise to observations causes the baselines to quickly blow up in error, but RWM’s error increases only marginally. In short, RWM learned a stochastic, stable predictor that is robust to noise, whereas many baselines simply couldn’t remain accurate for more than ~20–50 steps. A nice summary from the project writeup is: “The RWM-AR method achieved the lowest prediction errors across 40 diverse robotic embodiments and tasks (e.g., manipulation, legged locomotion), outperforming RSSM, Transformer, and teacher-forcing variants.”. We saw a taste of this in the Google site text: “Across all these settings, RWM trained autoregressively achieves the lowest prediction error relative to MLP, RSSM, and transformer baselines, demonstrating strong robustness and transferability.”. In other words, by removing strong inductive bias and focusing on data-driven rollout stability, RWM generalized across very different robots with one universal model. Policy Learning Results. Of course, accurate prediction is only useful if it leads to a better policy. Li et al. evaluate how well policies learned inside RWM perform. They focus on locomotion tasks (velocity tracking on flat ground) for a variety of quadrupeds and humanoids, as well as a few arm-manipulation tasks. When they train with MBPO-PPO inside RWM, they find that the policy learning is stable and efficient. In fact, in head-to-head comparisons on benchmark tasks, MBPO-PPO with RWM converges faster and to higher reward than the alternatives. For instance, against a classical model-based method (like training PPO but with a one-step learned model or RSSM without dual-AR), RWM’s policy learning curves are markedly better. Moreover, the learned policies stay aligned: the reward predicted in RWM rollouts matches the reward obtained under the real dynamics. The team plots “predicted vs. actual reward during training” and shows a near 45° alignment for RWM, whereas baselines quickly decouple, causing PPO to learn useless exploits. Notably, the authors point out that MBPO-PPO with RWM actually converges where other model-based methods fail to train a usable policy. For example, they tried Dreamer (an off-policy model-based RL using RSSM) and found it did not reliably converge in these robotics tasks. They call out a baseline “SHAC” (which can be thought of as a SAC-based MBRL variant) that fails entirely. In the end, PPO inside RWM “just works” while others did not. The summary statement is: “MBPO-PPO successfully converges while other methods (e.g. SAC-based world models, Dreamer) did not produce deployable policies.” (This is borne out by their training curves in Figure 5 of the supplementary.) We don’t have a direct citation here, but it underlines that RWM’s stable imagination is what enables learning a good locomotion policy at all. Here is a key takeaway the authors emphasize: If your world model can imagine accurately, your policy training stays grounded. Because RWM rollouts remain near the real dynamics, the policy does not chase illusions of reward. The ETH team reports that, during training inside RWM, model error actually tends to decrease over time as well – this is likely because the policy explores more systematically and helps the model see the relevant parts of the state space. The Project Readout notes explicitly: “RWM’s rollout fidelity keeps model error decreasing and predicted rewards aligned with ground truth, enabling stable optimization over 100+ autoregressive steps.”. In plain language, the policy keeps refining the model in areas it needs while the model keeps giving good training signals, in a virtuous cycle. Zero-Shot Sim-to-Real Transfer. Perhaps the most exciting result is that policies trained solely in imagination transferred directly to hardware. Li et al. tested the trained controllers on two real robots with no additional fine-tuning: ANYmal D (an open-source small quadruped) and the Unitree G1 humanoid. These robots were not used to collect extra data beyond the original training runs in simulation. Remarkably, the learned policies worked in real life out of the box. Each robot successfully tracked commanded velocities on flat ground, walking stably and even resisting external perturbations (pushes or uneven terrain) without falling. The authors candidly note this: “The resulting policies transfer zero-shot to ANYmal D and Unitree G1 hardware, reliably tracking velocity commands and staying robust to impacts and disturbances.”. Video clips in the project page show the real robots trotting around in response to random velocity commands, just as they did in the virtual world model. This sim-to-real success owes not just to the world model’s accuracy, but also to its automatic stochasticity and generalization. Recall that RWM injects variability into each imagined rollout (by sampling from its output distribution). This in effect served as a form of domain randomization at no extra cost, exposing the policy to slight variations of the dynamics. So by the time the controller hits the real robot, small discrepancies in mass, friction, or sensor noise do not break its behavior. The writeup again highlights this free robustness: “Each imagination rollout naturally introduces varied dynamics… improving hardware generalization”. For comparison, a policy trained purely in an ordinary simulator often requires careful calibration or real-world fine-tuning. Here, the robust imagination of RWM closes the loop. While the ETH team is cautious about calling this a complete solution (they note that true online learning on hardware would be ideal but is very risky currently), the fact that they achieved zero-shot transfer of a learned controller is a strong proof of concept. The policy was not specialized to any single robot’s quirks – it was trained only on the neural model – yet it walked the actual ANYmal D and G1 almost immediately. Key Insights and Takeaways. Long-horizon stability is everything. By forcing the world model to train autoregressively, the RWM sidesteps the usual compounding error problem. It literally practices making long rollouts during training, so when asked to imagine 100+ steps, it “knows what to do.” This paid off in dramatically lower rollout error compared to naive models. As the paper emphasizes, the hard part was not “learn a model,” but “learn a model that stays stable when rolled out long-horizon, the way PPO requires.” RWM’s success suggests this is a viable strategy. Model-free policy inside a model-based loop. The policies themselves were just standard PPO; the novelty is all in how the model pre-generates data. This hybrid – a learned neural simulator feeding a model-free optimizer – seems to harness the best of both worlds. The world model gives data efficiency, and PPO gives reliable on-policy learning. The result is that PPO converges to effective controllers, whereas using an off-policy actor-critic with a flawed model might have failed. Generalization across robots and tasks. Perhaps most surprising is that one model architecture worked for so many different robots: from industrial arms to legged platforms of various sizes. The RWM did not require any hand-tuned features for each robot – no explicit contact modeling, no robot-specific modules – yet it learned the dynamics of legs, joints, and sensors on its own. All that changed between tasks was the input/output dimensionality (handled by configuration). The authors report that in their evaluation suite RWM was essentially plug-and-play: “the same model and training pipeline work for manipulation and locomotion…I including ANYmal, Cassie, Spot, etc.”. This suggests that the approach has promise as a fairly general technique. Sim-to-real with uncalibrated models. Conventional wisdom says you should not trust a learned model on the real robot without fine-tuning. In practice, the ETH results show the opposite: a well-trained neural model can produce a controller that works out-of-the-box. The critical ingredient was the model’s stochastic robustness (acting like random perturbations) plus careful training on real-like data. It’s not magic – they still needed a good initial data log and a high-quality simulator to pretrain – but it’s clear that the neural world model mitigated much of the gap. On the cautionary side, the authors note that fully online learning on hardware (continuously updating both model and policy with real data) remains difficult. They observed that naive online training often leads to the policy exploiting small model errors and crashing (on average over 20 falls in their sim studies). Safely doing true online adaptation would require sophisticated resets or safety measures. For now, RWM is best viewed as a tool for offline policy refinement: train in “imagination” and then deploy. Practical Considerations and Future Directions. For roboticists interested in using RWM, the authors have provided a full code suite. They integrated RWM into NVIDIA’s Isaac Lab (a framework built on Isaac Sim) as an extension. This means you can install Isaac Lab, drop in the RWM extension, and get examples of pretraining the model and training the policy. The repository includes ready-made tasks (e.g. ) and scripts showing how to run PPO inside the learned model. You can even visualize the imagined rollouts of the world model to debug its predictions. (See the GitHub for details.). In terms of computation, this approach involves two learning loops (model and policy), so training times can be significant – the project cited “extensive experiments.” But note that because it is mostly offline simulation and on a GPU, it is far cheaper and safer than running each policy trial on a real robot. Looking ahead, the RWM paper itself is one piece of a bigger research program. The team also proposes in a companion work the idea of an uncertainty-aware RWM to enable truly offline RL, and they consider extensions like continuous model adaptation with safety. For most practitioners, the immediate takeaway is that a black-box neural dynamics model can be highly effective if trained the right way. It may not yet replace good simulators entirely (especially for tasks unseen by the model), but it makes ambitious sim-to-real policy learning much more reliable. In Summary. Robotic World Model is a learned, GRU-based simulator for robots that emphasizes long-range consistency in its rollouts. Trained self-supervised on real (or high-fidelity sim) data, it yields multi-step predictions that remain close to reality. Policies optimized inside this neural simulator (using PPO) learn robust locomotion controllers that transfer immediately to real hardware. The core innovation – dual autoregressive training – directly addresses the classical model-based RL challenge of error accumulation. The result is a scalable pipeline: collect some data, train RWM, train a policy in imagination, and out walks your real robot. For credentialed roboticists, the Robotic World Model offers a new tool: rather than hand-crafting dynamics or relying purely on phony simulators, we can let a network learn the physics and supply unlimited imagined experience. The ETH results are impressive: across dozens of robots, RWM outperforms simpler baselines in prediction and yields high-quality policies with little real-world trial. The work underscores the maturing belief that neural world models, if done right, can unlock efficient robot learning. Takeaway: Make your robot imagine. Train a neural dynamics model on its data, roll out policies inside it, and you may well get a ready-to-run controller – no extra fine-tuning needed. This RWM approach turns the traditional sim2real workflow on its head and paves the way for robots that learn to walk in their heads before they walk on their feet.