Robotic World Model Extension for Isaac Lab

Overview

This repository extends Isaac Lab with environments and training pipelines for Robotic World Model (RWM) and related model-based reinforcement learning methods.
It enables:

• joint training of policies and neural dynamics models in Isaac Lab,
• evaluation of model-based vs. model-free policies,
• visualization of autoregressive “imagination” rollouts from learned dynamics.

Paper: Robotic World Model: A Neural Network Simulator for Robust Policy Optimization in Robotics
Project Page: https://sites.google.com/view/roboticworldmodel

Authors: Chenhao Li, Andreas Krause, Marco Hutter
Affiliation: ETH AI Center, Learning & Adaptive Systems Group and Robotic Systems Lab, ETH Zurich

---

Installation

1. Install Isaac Lab

Follow the official installation guide. We recommend using the Conda installation as it simplifies calling Python scripts from the terminal.

2. Install model-based RSL RL

Follow the official installation guide of model-based RSL RL for model-based reinforcement learning to replace the rsl_rl_lib that comes with Isaac Lab.

3. Clone this repository (outside your Isaac Lab directory)

git clone git@github.com:leggedrobotics/robotic_world_model.git

4. Install the extension using the Python environment where Isaac Lab is installed

python -m pip install -e source/mbrl

5. Verify the installation

python scripts/reinforcement_learning/rsl_rl/train.py --task=Template-Isaac-Velocity-Flat-Anymal-D-Init-v0 --headless

---

World Model Pretraining & Evaluation

Robotic World Model is a model-based reinforcement learning algorithm that learns a dynamics model and a policy concurrently.

Configure model inputs/outputs

You can configure the model inputs and outputs under ObservationsCfg_PRETRAIN in AnymalDFlatEnvCfg_PRETRAIN.

Available components:

• SystemStateCfg: state input and output head
• SystemActionCfg: action input
• SystemExtensionCfg: continuous privileged output head (e.g. rewards etc.)
• SystemContactCfg: binary privileged output head (e.g. contacts)
• SystemTerminationCfg: binary privileged output head (e.g. terminations)

Run dynamics model pretraining:

python scripts/reinforcement_learning/rsl_rl/train.py \
  --task=Template-Isaac-Velocity-Flat-Anymal-D-Pretrain-v0 \
  --headless

It trains a PPO policy from scratch, while the induced experience during training is used to train the dynamics model.

Visualize autoregressive predictions

python scripts/reinforcement_learning/rsl_rl/visualize.py \
  --task=Template-Isaac-Velocity-Flat-Anymal-D-Visualize-v0 \
  --headless \
  --load_run <run_name> \
  --system_dynamics_load_path <dynamics_model_path>

---

Model-Based Policy Training & Evaluation

Once a dynamics model is pretrained, you can train a model-based policy purely from imagined rollouts generated by the learned dynamics.

Train policy in imagination

python scripts/reinforcement_learning/rsl_rl/train.py --task=Template-Isaac-Velocity-Flat-Anymal-D-Finetune-v0 --headless --load_run <run_name> --system_dynamics_load_path <dynamics_model_path>

You can either start the policy from pretrained checkpoints or from scratch by simply omitting the --load_run argument.

Play the learned model-based policy

python scripts/reinforcement_learning/rsl_rl/play.py --task=Template-Isaac-Velocity-Flat-Anymal-D-Play-v0 --headless --load_run <run_name>

Code Structure

We provide a reference pipeline that enables RWM on ANYmal D.

Key files:

• Environment configurations + dynamics model setup flat_env_cfg.py.
• Algorithm configuration + training parameters anymal_d_manager_based_mbrl_env.py.
• Imagination rollout logic (constructs policy observations & rewards from model outputs) anymal_d_manager_based_mbrl_env.
• Visualization environment + rollout reset anymal_d_manager_based_visualize_env.py.

Citation

If you find this repository useful for your research, please consider citing:

@article{li2025robotic,
  title={Robotic world model: A neural network simulator for robust policy optimization in robotics},
  author={Li, Chenhao and Krause, Andreas and Hutter, Marco},
  journal={arXiv preprint arXiv:2501.10100},
  year={2025}
}