Getting Started#

This page takes you from installation to a trained world model you can watch, then shows where to go for each model family.

1. Install#

pip install "simulusion[gym]"

The package installs as simulusion and imports as simulusion. The gym extra brings Gymnasium with its classic-control and Atari environments. Other extras add DeepMind Control (dmc), logging (ml), MuJoCo, Brax and more; see Installation.

Simulusion does not pin a PyTorch build. For GPU training, install the CUDA wheel that matches your driver from pytorch.org first, and check it with:

python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

2. Train your first world model#

Dreamer learns a model of the environment from raw pixels and trains its policy entirely inside that model:

import simulusion

agent = simulusion.create_model(
    "dreamer",
    env="Pendulum-v1",
    env_backend="gym",
    total_steps=20_000,
)
agent.train()

Training prints a line per iteration and writes everything to one run folder, runs/Pendulum-v1_Dreamerv1_<name>_<timestamp>/:

File	Contents
`config.yaml`	the full configuration of the run
`metrics.jsonl`	one JSON line per iteration: returns, losses, throughput
`ckpts/<step>_ckpt.pt`	checkpoints, the final one always written

create_model accepts any field of the model’s config as a keyword argument. To see them all, build the config first:

cfg = simulusion.create_config("dreamer", env="Pendulum-v1", env_backend="gym")
print(cfg)

3. Watch what it learned#

The demo scripts turn a run into video. For Dreamer, record shows the real environment beside the world model’s prediction of it, plus the policy acting:

python demos/dreamer_demo.py record --run runs/Pendulum-v1_Dreamerv1_<name>_<timestamp>

What record produces, here for a Dreamer trained on cartpole swing-up. Left: the real environment. Right: the world model, observing 5 frames and then imagining the rest.

For a live session instead of a recording, simulusion play --model dreamer -c <checkpoint> opens a window where you can take over the controls (needs a desktop session).

4. Choose a model#

Every model is built with simulusion.create_model(name, ...). The names are those returned by simulusion.list_models():

Model	`create_model` name	What it does	Demo
Dreamer	`dreamer`, `dreamer-v1`, `dreamer-v2`	Learns latent dynamics from pixels; trains an actor-critic in imagination	video
PlaNet	`planet`	Plans each action with CEM inside a learned latent model	video
DIAMOND	`diamond`	Diffusion world model of Atari; agent trained on generated frames	video
IRIS	`iris`	Discrete-token autoencoder + Transformer world model	video
Genie	`genie`, `genie-small`, `genie-large`	Playable world model learned from unlabelled video	video
DiT	`dit`	Diffusion Transformer for image generation	video
I-JEPA	`jepa`	Self-supervised image representations, predicted in latent space	images
ModularRSSM	`modular-rssm`	A world model assembled from swappable encoder, backbone and decoder	video

To run any of them on a different environment, pass env= and env_backend=. The same backend names work with simulusion.make_env():

Backend	Environments
`gym`	Gymnasium IDs (classic control, Box2D, Atari), or an existing env instance
`dmc`	DeepMind Control Suite tasks, written `domain-task` (e.g. `walker-walk`)
`mujoco`	Gymnasium MuJoCo tasks or native MJCF/MJB models
`robotics`	IDs registered by Gymnasium Robotics
`brax`	JAX/Brax continuous-control environments
`procgen`	Procgen games such as `coinrun` (manual install)
`dmlab`	DeepMind Lab 3D navigation tasks
`unity_mlagents`	Unity ML-Agents executables

5. Log to Weights & Biases or TensorBoard#

Both are off by default. Install the ml extra, then set them on the config:

cfg = simulusion.create_config("dreamer", env="Pendulum-v1", env_backend="gym")
cfg.enable_tensorboard = True  # then: tensorboard --logdir runs
cfg.enable_wandb = True  # needs WANDB_API_KEY in the environment
cfg.wandb_project = "simulusion"
agent = simulusion.create_model("dreamer", config=cfg)

Next steps#

The Gallery shows every model trained on a laptop GPU, with the commands to reproduce each clip.
Training Guide covers budgets, checkpointing, resuming and early stopping.
Inference Guide and Efficient Inference and Deployment cover fast inference and exporting to ONNX, TorchScript and TensorRT.
World Models Study Guide explains the ideas the models share.