simillusion/dreamerv2-dmc-cartpole-swingup

DreamerV2 on DMC cartpole swing-up

A demo checkpoint from Simulusion, an open-source PyTorch library that puts world models behind one API. It implements the paper Mastering Atari with Discrete World Models (Hafner et al., ICLR 2021).

DreamerV2 replaces Dreamer's Gaussian latents with discrete categorical ones and uses symlog two-hot reward and value heads with a balanced KL loss.

Results

Trained from scratch for 60k environment steps on a 4 GB RTX 3050 laptop GPU, at batch 32 (the most that fit beside MuJoCo's renderer). Evaluation return rose from about 75 to between 165 and 250.

This is a short demo run sized for a laptop GPU, meant to show what the model does. It is not a reproduction of the paper's results. More clips and learning curves are in the gallery.

Usage

pip install simulusion huggingface_hub

from huggingface_hub import snapshot_download
from simulusion.world_model import load_world_model

path = snapshot_download("simillusion/dreamerv2-dmc-cartpole-swingup")
wm = load_world_model(path, "dreamer")  # encode / step / rollout / decode

The same interface works for every world model in Simulusion. To get the full agent:

from simulusion.models.dreamer import DreamerAgent

# Builds the DeepMind Control environment too: pip install "simulusion[dmc]"
agent = DreamerAgent.from_pretrained(path)

Checkpoints are plain state dicts and load with torch.load(..., weights_only=True). They include optimizer state, so you can continue training from them.

Training

Reproduce this checkpoint from a clone of the repository:

python demos/dreamer_demo.py train --algo dreamer-v2 --batch-size 32 --run demos/runs/dreamer_v2

Citation

If you use this model, please cite the original paper:

@inproceedings{hafner2021dreamerv2,
  title     = {Mastering Atari with Discrete World Models},
  author    = {Hafner, Danijar and Lillicrap, Timothy and Norouzi, Mohammad and Ba, Jimmy},
  booktitle = {International Conference on Learning Representations},
  year      = {2021},
  url       = {https://arxiv.org/abs/2010.02193},
}

and Simulusion:

@software{simulusion,
  title  = {Simulusion: a PyTorch library for world models},
  author = {Thakkar, Param},
  year   = {2026},
  url    = {https://github.com/ParamThakkar123/simulusion},
}