Official implementation of "Diffusion World Model: Future Modeling Beyond Step-by-Step Rollout for Offline Reinforcement Learning" (arXiv:2402.03570).

DWM is a conditional diffusion model that predicts multi-step future states and rewards in a single forward pass, conditioned on the current state, action and return-to-go. It is integrated into a Dyna-style model-based offline RL pipeline via Diffusion Model Based Value Expansion (Diffusion-MVE): the world model generates short future rollouts that bootstrap the critic, sidestepping the compounding error of recursively-queried one-step dynamics models.

s_t, a_t, g_t  ──►  DWM  ──►  r_t, s_{t+1}, r_{t+1}, ..., s_{t+H-1}, r_{t+H-1}

Repository structure

Runnable Hydra entry scripts live at the repository root; all importable library code lives in the dwm/ package.

# Entry scripts (run from the repo root, e.g. `python train_world_model.py ...`)
train_world_model.py                 Stage 1: train the diffusion world model (DWM, or DD with add_condition=False)
cotrain_diffusion_actor_critic.py    Stage 2: model-based actor-critic co-training (Diffusion-MVE)
eval_cotrain_diffusion_policy.py     Evaluate Diffusion-QL policies (actor + critic)
eval_cotrain_policy.py               Evaluate any co-trained actor (TD3+BC / IQL / Diffusion-QL)
render_dwm_prediction.py             DWM multi-step prediction-error analysis
render_fdm_prediction.py             One-step model prediction-error analysis
train_fdm.py                         Baseline: one-step forward dynamics model
train_idm.py                         Inverse dynamics model (for the DD+IDM baseline)
eval_dd_idm.py                       Baseline: Decision Diffuser + inverse dynamics model
train_policy.py                      Baselines: model-free TD3+BC / IQL
train_decision_transformer.py        Baseline: Transformer world model (Appendix E.3)
cotrain_transformer_actor_critic.py  Baseline: Transformer-world-model co-training
# Library package
dwm/
  data.py            Dataset loading + PyTorch Dataset wrappers
  utils.py           Env specs, seeding, checkpoint I/O, result aggregation
  utils_dt.py        Decision-Transformer helpers (optimizer/loss/eval envs)
  logger.py          TensorBoard / CSV logging
  lamb.py            LAMB optimizer (used by the Transformer baseline)
  model/             Diffusion model, temporal U-Net, actor-critic networks, GPT-2 backbone
  trainer/           Training loops (world-model pre-training and actor-critic co-training)
  common/            Normalizers, evaluation/rendering, forward/inverse dynamics training
# Other
config/        Hydra configs (top-level + model/ + actor_critic_model/ groups)
dataset/       D4RL download scripts
scripts/       Portable sweep launcher (local or SLURM)

Installation

DWM is benchmarked on the D4RL MuJoCo locomotion datasets, which require MuJoCo and mujoco-py.

1. Install MuJoCo 2.1.0 (mujoco210) under ~/.mujoco/ and add it to your library path, e.g.:

   export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:~/.mujoco/mujoco210/bin

2. Create the environment and install dependencies:

   conda env create -f environment.yml && conda activate dwm
   # or, in an existing environment:
   bash install.sh        # pip install -r requirements.txt

   d4rl is installed from source; gym==0.18 and mujoco-py==2.1 are pinned for D4RL compatibility.

Configuration: paths via environment variables

All file paths live in the Hydra configs and are resolved from two environment variables so the repo stays portable:

export DWM_DATA_DIR=/abs/path/to/dataset
export DWM_EXP_DIR=/abs/path/to/experiments

Any value can also be overridden directly on the command line, e.g. data.dataset_folder=/my/data.

Download the datasets

# trajectory-level data -> $DWM_DATA_DIR/<env>.pkl       (world-model training)
python dataset/download_d4rl_gym_datasets.py
# transition-level data -> $DWM_DATA_DIR/raw/<env>.pkl   (model-free / co-training)
python dataset/download_d4rl_gym_datasets_raw.py

By default these download the 9 locomotion tasks ({hopper,walker2d,halfcheetah}-{medium,medium-replay,medium-expert}-v2); edit the ENVS / DATASET_TYPES lists in the scripts for other tasks.

Usage

The pipeline has two stages: train the world model, then co-train the policy on top of it. All entry points use Hydra, so any config value can be overridden as key=value.

Stage 1 — train the Diffusion World Model

python train_world_model.py env=hopper-medium-v2 data.reward_scale=400 \
    data.H=8 model.n_diffusion_steps=5

This writes a snapshot.pt to the run directory under $DWM_EXP_DIR.

Stage 2 — model-based actor-critic co-training (Diffusion-MVE)

Point diffusion_model_dir at the Stage-1 run directory:

python cotrain_diffusion_actor_critic.py env=hopper-medium-v2 \
    actor_critic_model=diffusion_ql diffusion_model_dir=$DWM_EXP_DIR/<dwm_run_dir> \
    eval_rtg=0.7 lookahead_steps=5

actor_critic_model selects the policy learner: diffusion_ql, td3bc, iql or pql (Pessimistic Q-learning). actor_*.pt / critic_*.pt checkpoints are saved periodically.

Evaluate

python eval_cotrain_diffusion_policy.py env=hopper-medium-v2 seed=1 \
    actor=DiffusionActor model_dir=$DWM_EXP_DIR/<cotrain_run_dir>

Use eval_cotrain_policy.py with actor=Actor (TD3+BC) or actor=TanhNormActor (IQL) for the other policy heads. (seed is required by the eval configs.)

Sweeping over all 9 tasks

scripts/launch_sweep.py loops over the locomotion tasks with the paper's per-environment hyper-parameters:

python scripts/launch_sweep.py train_dwm --local            # Stage 1, all tasks
python scripts/launch_sweep.py cotrain_dwm_policy --local \
    diffusion_model_dir=$DWM_EXP_DIR/<dwm_run_dir>          # Stage 2, all tasks
python scripts/launch_sweep.py eval_dwm_policy --local \
    model_dir=$DWM_EXP_DIR/<cotrain_run_dir>

Drop --local to submit to a SLURM cluster through Hydra's submitit launcher.

Baselines

Most baselines mirror the two-stage DWM pipeline (train a world/dynamics model, then learn a policy on top of it); the model-free ones are single-stage. Set DWM_DATA_DIR / DWM_EXP_DIR and download the datasets first (see above). Run directories are created under $DWM_EXP_DIR; substitute the generated paths for the <...> placeholders below. Use the per-task data.reward_scale from the table at the end (hopper 400 / walker2d 550 / halfcheetah 1200).

One-step dynamics model (FDM)

Replaces the diffusion world model with a one-step model p(s_{t+1}, r_t | s_t, a_t) that must be queried recursively at planning time (the main comparison of Section 5.1).

1. Train the one-step forward dynamics model — writes forward_model.pt:

   python train_fdm.py env=hopper-medium-v2 data.reward_scale=400

2. Co-train the policy on top of it (model=fdm, point fdm_dir at step 1; this uses the transition-level raw dataset):

   python cotrain_diffusion_actor_critic.py env=hopper-medium-v2 \
       model=fdm fdm_dir=$DWM_EXP_DIR/<fdm_run_dir> \
       actor_critic_model=diffusion_ql eval_rtg=0.7 lookahead_steps=5

3. Evaluate (identical to the DWM policy eval):

   python eval_cotrain_diffusion_policy.py env=hopper-medium-v2 seed=1 \
       actor=DiffusionActor model_dir=$DWM_EXP_DIR/<cotrain_run_dir>

Decision Diffuser + inverse dynamics model (DD+IDM)

Predicts a state-only future trajectory with a diffusion model and recovers the action to execute with an inverse dynamics model (Section 5.2).

1. Train the Decision-Diffuser world model. This is DWM Stage 1 but with add_condition=False, so the model is conditioned on the state and return-to-go only (no initial action), matching what eval_dd_idm.py expects:

   python train_world_model.py env=hopper-medium-v2 add_condition=False \
       data.reward_scale=400 data.H=8 model.n_diffusion_steps=5

2. Train the inverse dynamics model — writes inverse_model.pt:

   python train_idm.py env=hopper-medium-v2

3. Evaluate by planning with the world model and inferring actions with the IDM:

   python eval_dd_idm.py env=hopper-medium-v2 seed=1 \
       model_dir=$DWM_EXP_DIR/<dd_run_dir> \
       idm_dir=$DWM_EXP_DIR/<idm_run_dir> eval_rtg=0.7

Model-free TD3+BC / IQL

Single-stage model-free offline RL on transition-level data — the model-free counterparts DWM is compared against (Section 5.3). The policy is evaluated in the environment after every epoch and logged to results.csv.

1. Train (also evaluates during training; saves actor_*.pt):

   python train_policy.py env=hopper-medium-v2 model=td3bc   # or model=iql

2. Re-evaluate a saved checkpoint (optional):

   python eval_cotrain_policy.py env=hopper-medium-v2 seed=1 \
       actor=Actor model_dir=$DWM_EXP_DIR/<td3bc_run_dir>      # actor=TanhNormActor for IQL

Transformer world model (Appendix E.3)

Replaces the diffusion model with an autoregressive GPT-2-style Transformer world model; because it rolls out step by step it is more prone to compounding error.

1. Train the Transformer world model — writes snapshot.pt:

   python train_decision_transformer.py env=hopper-medium-v2 \
       data.reward_scale=400 data.H=8

2. Co-train the policy (transformer_model_dir points at step 1):

   python cotrain_transformer_actor_critic.py env=hopper-medium-v2 \
       transformer_model_dir=$DWM_EXP_DIR/<twm_run_dir> \
       actor_critic_model=iql eval_rtg=0.7 lookahead_steps=5

3. Evaluate (use the actor head matching actor_critic_model: TanhNormActor for iql, Actor for td3bc):

   python eval_cotrain_policy.py env=hopper-medium-v2 seed=1 \
       actor=TanhNormActor model_dir=$DWM_EXP_DIR/<cotrain_run_dir>

Prediction-error analysis

python render_dwm_prediction.py env=hopper-medium-v2 model_dir=... save_path=... eval_rtg='[0.6,0.7,0.8]'
python render_fdm_prediction.py env=hopper-medium-v2 fdm_dir=...   save_path=...

Per-environment hyper-parameters

The paper uses reward scaling and evaluation return-to-go (RTG) that depend on the task (also encoded in scripts/launch_sweep.py):

Key world-model settings used in the paper (applied via the launchers / command- line overrides): sequence length H=8, n_diffusion_steps=5, num_inference_steps=3, value-expansion horizon lookahead_steps=5.

Citation

@article{ding2024diffusion,
  title={Diffusion World Model: Future Modeling Beyond Step-by-Step Rollout for Offline Reinforcement Learning},
  author={Ding, Zihan and Zhang, Amy and Tian, Yuandong and Zheng, Qinqing},
  journal={arXiv preprint arXiv:2402.03570},
  year={2024}
}