Robotic World Model: A Neural Network Simulator for Robust Policy Optimization in Robotics

Chenhao Li; Andreas Krause; Marco Hutter

Source-linked AI summary

Robotic World Model: A Neural Network Simulator for Robust Policy Optimization in Robotics

Chenhao Li, Andreas Krause, Marco Hutter

arXiv:2501.10100v5Jan 17, 2025cs.ROcs.AIcs.LG

TL;DR

Robotic control needs world models that can learn complex, partially observable dynamics efficiently and generalize without domain-specific assumptions. The paper introduces RWM with dual-autoregressive, self-supervised training and combines it with MBPO-PPO for imagined policy optimization and hardware deployment. RWM reportedly improves autoregressive prediction across diverse environments, while MBPO-PPO enables robust zero-shot transfer, although current training relies on simulation pretraining and safe online hardware learning remains challenging.

Problem
Robotic world models must handle nonlinear, stochastic, partially observable dynamics while supporting efficient, generalizable control without domain-specific assumptions.
Method
RWM uses dual-autoregressive, self-supervised training for long-horizon prediction, and MBPO-PPO optimizes policies in learned simulations for deployment.
Results
RWM consistently outperforms RSSM and transformer-based architectures in autoregressive prediction accuracy across diverse robotic environments and supports robust hardware transfer with MBPO-PPO.
Takeaways & Limitations
The framework supports scalable model-based reinforcement learning for robotic tasks spanning simulation and real-world hardware.
Takeaways & Limitations
The method currently relies on simulation pretraining, while safe and effective online learning directly on hardware remains challenging.

Abstract

from arXiv · show

Learning robust and generalizable world models is crucial for enabling efficient and scalable robotic control in real-world environments. In this work, we introduce a novel framework for learning world models that accurately capture complex, partially observable, and stochastic dynamics. The proposed method employs a dual-autoregressive mechanism and self-supervised training to achieve reliable long-horizon predictions without relying on domain-specific inductive biases, ensuring adaptability across diverse robotic tasks. We further propose a policy optimization framework that leverages world models for efficient training in imagined environments and seamless deployment in real-world systems. This work advances model-based reinforcement learning by addressing the challenges of long-horizon prediction, error accumulation, and sim-to-real transfer. By providing a scalable and robust framework, the introduced methods pave the way for adaptive and efficient robotic systems in real-world applications.

1 Introduction

The paper targets robust, generalizable robotic world models that support long-horizon prediction and policy optimization without domain-specific assumptions. It combines a learned simulator with policy optimization for simulated and real-world robotic deployment.

World models simulate environment dynamics for planning and policy optimization, but reliable generalization is difficult under nonlinear, stochastic, and partially observable dynamics.
The framework learns robust world models for accurate long-horizon predictions without handcrafted representations or specialized architectural biases.
Its network and training framework enables reliable long autoregressive rollouts, which are important for downstream planning and control.
The proposed policy optimization framework uses learned world models for continuous control and generalizes to simulated and real-world robotic scenarios.
The framework aims to improve the scalability, adaptability, and robustness of robotic systems for real-world model-based reinforcement learning.

2 Related work

Prior work uses world models for planning and control, while many approaches add physics or structural inductive biases to improve generalization. This paper extends model-based policy optimization by combining MBPO with PPO over extensive autoregressive rollouts.

World models capture system dynamics and enable efficient planning and control through simulated trajectories in robotic tasks.
Many black-box neural world models incorporate physics principles or state structures to generalize beyond their training distributions.
Model-based reinforcement learning improves policy-learning efficiency and safety by simulating environment interactions with a learned model.
The framework extends MBPO by integrating it with PPO over extensive autoregressive rollouts for complex robotic control tasks.

3 Approach

RWM learns world models with self-supervised dual-autoregressive training and uses them to optimize policies through imagined rollouts. The framework targets partially observable, stochastic robotic environments and supports extended autoregressive policy optimization.

3.1 Reinforcement Learning and World Models: RWM combines historical context and autoregressive training to produce reliable long-horizon predictions in partially observable and dynamically complex environments.The approach is designed to address error accumulation and discontinuous dynamics without relying on simplified settings.
3.2 Self-supervised Autoregressive Training: The world model predicts future observations from historical observation-action sequences, appending each prediction to the history for N-step autoregressive rollout.The input spans M historical steps, and predicted observations are combined with subsequent actions for later predictions.
3.2 Self-supervised Autoregressive Training: Training minimizes multi-step errors for observations and privileged information, encouraging hidden states that support accurate long-horizon predictions.The objective uses observation and privileged-information discrepancies with a decay factor α.
3.2 Self-supervised Autoregressive Training: RWM uses a GRU-based network with a dual-autoregressive mechanism to maintain historical context and predict the next observation distribution.The architecture is suited to low-dimensional inputs and predicts Gaussian means and standard deviations.
3.3 Policy Optimization on Learned World Models: Policy optimization recursively generates imagined actions and observations with the learned model, then updates the policy using PPO or another reinforcement-learning algorithm.Imagination agents are initialized from replay-buffer observations and simulate trajectories for T steps.
3.3 Policy Optimization on Learned World Models: RWM successfully optimizes policies over more than one hundred autoregressive steps with MBPO-PPO, exceeding the stated rollout capabilities of existing frameworks.The extended-rollout result is presented as evidence of training robustness and stability for hardware-deployable policies.

4 Experiments

Experiments evaluate RWM’s long-horizon prediction accuracy, robustness to noise, cross-environment generalization, and policy-learning utility. Across these settings, autoregressive training supports stable predictions and MBPO-PPO enables hardware deployment.

4.1 Autoregressive Trajectory Prediction: RWM aligns predicted and ground-truth trajectories over extended rollouts, mitigating compounding errors through its dual-autoregressive mechanism.The trajectory evaluation uses ANYmal D data with history horizon M = 32 and forecast horizon N = 8.
4.1 Autoregressive Trajectory Prediction: Extending history and forecast horizons improves accuracy, but computational cost requires careful tuning of M and N.This result comes from the reported horizon ablation study.
4.2 Robustness under Noise: RWM maintains lower prediction errors than an autoregressively trained MLP across noise levels, with less error growth as forecast steps increase.The comparison evaluates Gaussian perturbations applied to observations and actions.
4.3 Generality across Robotic Environments: RWM-AR consistently achieves the lowest autoregressive prediction errors across diverse robotic environments, outperforming MLP, RSSM, transformer architectures, and RWM-TF.The advantage is especially pronounced in complex dynamic tasks such as legged-robot velocity tracking.
4.4 Policy Learning and Hardware Transfer: MBPO-PPO reduces model error during training and achieves higher policy rewards than SHAC and Dreamer in ANYmal D and Unitree G1 velocity tracking.Predicted rewards initially overshoot ground truth but progressively align with it as training proceeds.
4.4 Policy Learning and Hardware Transfer: MBPO-PPO transfers zero-shot to ANYmal D and Unitree G1 hardware, maintaining velocity tracking and stability under impacts and terrain disturbances.SHAC and Dreamer fail to produce deployable policies after training collapse.

5 Limitations

RWM and MBPO-PPO outperform existing model-based reinforcement-learning methods in robustness and generalization, but remain behind well-tuned model-free methods using high-fidelity simulators. The framework also relies on simulation pretraining and faces unresolved challenges for safe online hardware learning.

RWM and MBPO-PPO surpass existing MBRL methods in robustness and generalization but trail well-tuned model-free RL trained with high-fidelity simulators.The comparison reflects different strengths: model-free RL benefits from near-perfect simulator interaction, whereas MBRL is advantageous when accurate or efficient simulation is infeasible.
Pretraining the world model on simulation data reduces training instability, while training from scratch can suffer from policy exploitation of model inaccuracies.Additional environment interaction is needed to fine-tune the model, creating efficiency and stability trade-offs.
Safe and effective online learning directly on hardware remains challenging because deployment requires safety constraints and robust uncertainty estimates.Current simulation-based training avoids potential hardware damage but does not resolve lifelong real-world learning.

6 Conclusion

The conclusion presents RWM as a robust, scalable world-model framework for complex robotic tasks and MBPO-PPO as its policy-optimization extension. Reported results include accurate long-horizon prediction, strong simulation performance, and zero-shot transfer to two hardware platforms.

RWM outperforms RSSM and transformer-based architectures in autoregressive prediction accuracy across diverse robotic environments.The framework uses dual autoregression, historical context, and self-supervised long-horizon training without domain-specific inductive biases.
MBPO-PPO leverages long world-model rollouts to train policies that perform strongly in simulation and transfer zero-shot to ANYmal D and Unitree G1.The reported hardware demonstrations provide evidence of deployment beyond simulation.
RWM is presented as a generalizable, efficient, and scalable contribution to model-based reinforcement learning for complex robotic tasks.

A.1.1 Observation and action spaces

The appendix defines the observation, privileged-information, and action spaces used by the world model and policy. Robot observations include motion, orientation, joint, command, and torque-related variables, while privileged information includes contacts.

The world-model observation space contains base linear and angular velocities, gravity measurements, joint positions, velocities, and torques.These quantities are defined in the robot frame where specified.
Privileged information consists of knee and foot contacts and supplies an additional learning objective for accurate long-term prediction.
The action space uses joint position targets.
The ANYmal velocity-tracking policy observes base velocities, gravity, velocity commands, and joint positions and velocities.

A.1.2 Reward functions

The reward is a weighted sum of tracking, stability, smoothness, contact, clearance, and posture terms. Several terms penalize undesirable motion or contacts, while velocity-tracking terms compare commanded and current base motion.

The total reward sums weighted terms detailed in Table S6.
Linear and angular velocity tracking reward agreement between commanded and current base velocities using temperature factors.The linear term covers x and y motion, while the angular term tracks yaw velocity.
Additional terms account for vertical velocity, roll and pitch velocity, joint torques, and joint acceleration.
The reward includes penalties or incentives for action changes, undesired contacts, foot-air time, swing-foot clearance, projected gravity, and deviation from default joint positions.

A.2 Network Architecture

RWM uses a GRU-based world model with MLP prediction heads, while MBPO-PPO combines learned-model rollouts with real-environment data for policy optimization. Experiments examine architecture, horizon choices, imagination rollouts, pretraining, and real-world deployment constraints.

Network architecture: RWM combines a GRU base with MLP heads that predict next observations, privileged information, means, and standard deviations.The architecture is summarized in Table S7.
Policy optimization: MBPO-PPO stores observation-action pairs in a replay buffer, trains the world model, and uses imagined trajectories for PPO policy updates.Figure S7 depicts the interaction between real-environment data, simulated rollouts, and policy optimization.
Horizon ablation: Longer history horizons reduce autoregressive prediction error until improvements plateau, while longer forecast horizons improve long-term prediction accuracy.Figure S8 presents these effects as a heatmap over history horizon M and forecast horizon N.
Horizon ablation: M = 32 and N = 8 provide a practical trade-off between strong autoregressive performance and manageable training time.Teacher-forcing with N = 1 minimizes training time but produces poor autoregressive performance because it does not expose the model to compounding errors.
Deployment considerations: RWM uses simulation pretraining, predicts unsafe-contact terminations, and supports imagination rollouts across diverse robotic systems, but real-hardware policy training remains unevaluated.Pretraining is especially necessary for locomotion, while online learning can exploit model errors and cause collisions.