simillusion/dreamer-dmc-cartpole-swingup

Dreamer (V1) 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 Dream to Control: Learning Behaviors by Latent Imagination (Hafner et al., ICLR 2020).

Dreamer learns a recurrent state-space model of the environment from pixels, then trains an actor-critic purely on imagined trajectories.

Results

Trained from scratch for 60k environment steps (about 95 minutes on a 4 GB RTX 3050 laptop GPU). Evaluation return rose from about 75 to 493, out of a maximum of roughly 880: the policy swings the pole up but cannot yet balance it. The world model tracks the real cart for about ten imagined steps before it drifts.

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/dreamer-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 --env cartpole-swingup --steps 60000

Citation

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

@inproceedings{hafner2020dreamer,
  title     = {Dream to Control: Learning Behaviors by Latent Imagination},
  author    = {Hafner, Danijar and Lillicrap, Timothy and Ba, Jimmy and Norouzi, Mohammad},
  booktitle = {International Conference on Learning Representations},
  year      = {2020},
  url       = {https://arxiv.org/abs/1912.01603},
}

and Simulusion:

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