simillusion/planet-dmc-cartpole-swingup

PlaNet 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 Learning Latent Dynamics for Planning from Pixels (Hafner et al., ICML 2019).

PlaNet has no policy network. At every step it searches over action sequences with the cross-entropy method inside its learned latent model, and executes the first action of the best plan.

Results

Trained from scratch for 75 minutes on a 4 GB RTX 3050 laptop GPU (action repeat 8). Evaluation return rose from 76 to 435; a recorded planner episode scores 388.

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/planet-dmc-cartpole-swingup")
wm = load_world_model(path, "planet")  # encode / step / rollout / decode

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

from simulusion.agents import CEMPlanner, PlanningAgent

# PlaNet acts by planning in the world model (wm, loaded above)
agent = PlanningAgent(wm, CEMPlanner())

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/planet_demo.py train --minutes 75

Citation

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

@inproceedings{hafner2019planet,
  title     = {Learning Latent Dynamics for Planning from Pixels},
  author    = {Hafner, Danijar and Lillicrap, Timothy and Fischer, Ian and Villegas, Ruben and Ha, David and Lee, Honglak and Davidson, James},
  booktitle = {International Conference on Machine Learning},
  year      = {2019},
  url       = {https://arxiv.org/abs/1811.04551},
}

and Simulusion:

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