One interface for every world model#

Each world model family in Simulusion predicts the future its own way: Dreamer and PlaNet roll a recurrent latent forward, IRIS and Genie generate tokens with a Transformer, DIAMOND samples pixels with a diffusion model, and World Models pairs a VAE with an MDN-RNN. The WorldModel protocol puts one interface over all of them, so code that evaluates, plans with, or plays a world model works with any family.

import torch
from simulusion.world_model import as_world_model, load_world_model

model = as_world_model(agent)  # a model you trained in this process
# or: model = load_world_model("runs/dreamer/ckpts/61000_ckpt.pt", "dreamer")

state = model.encode(frames, actions)  # real context -> model state
out = model.step(state, action)  # one imagined step
out.reward, out.terminated  # from the model's own heads
image = model.decode(out.state)  # (B, C, H, W) in [0, 1]

torch.manual_seed(0)
traj = model.rollout(state, planned_actions)  # many steps
traj.frames, traj.rewards, traj.terminated  # (B, H, ...)

The methods#

Method	Does
`encode(frames, actions=None)`	Turns real context frames, and the actions taken between them, into a model state.
`step(state, action)`	Imagines one step. Returns a `Transition`: the next state, plus `reward` and `terminated` where the model has those heads.
`decode(state)`	The current frame of a state, float `(B, C, H, W)` in `[0, 1]`.
`rollout(state, actions, decode=True)`	`H` steps at once; returns a `Rollout` of frames, rewards and terminations, each `(B, H, ...)`.
`save(path)`	Writes the model to a directory that `load_world_model` reads back.
`spec`	A `WorldModelSpec`: what the model reads and predicts (below).

Conventions#

The same for every family, so you never need to know a model’s internals:

Frames are RGB, channel-first and batch-first: (B, T, C, H, W), either float in [0, 1] or uint8 in [0, 255]. A 4-D input is one sequence and a 3-D input one frame. Channel-last input and other sizes are converted to the model’s own (spec.obs_shape).
Actions between context frames: actions[:, t] is the action taken at frame t, which led to frame t + 1. T frames come with T - 1 actions. Discrete actions are integers (B, ...), continuous ones float (B, ..., A). A single action passed to step is used for the whole batch.
State is a dict[str, Tensor] with the batch on dim 0. It is plain data: keep a copy to branch a dream from the same point.
Randomness comes from PyTorch’s global RNG. torch.manual_seed before encode and rollout makes a dream repeatable.
Inference only: methods run without autograd and put the model in eval mode. Training stays with each family’s trainer.

What each family provides#

model.spec says what a model can do, so generic code can check instead of special-casing families:

Family	Actions	Context	Reward	Termination	Notes
`dreamer` (V1, V2)	continuous in `[-1, 1]`	1+ frames, recurrent	yes	if trained with the discount head
`planet`	continuous, unbounded	1+ frames, recurrent	yes	no	frames are quantised to the trained bit depth
`iris`	discrete	up to `transformer_timesteps` frames	yes, on the trained (sign) scale	yes	`sample=False` makes steps deterministic
`diamond`	discrete	exactly `num_conditioning_frames`	yes, `{-1, 0, 1}`	yes	`sampling_steps=` trades quality for speed
`genie`	discrete, latent (`spec.action.latent`)	up to `num_frames`	no	no	without `actions`, Genie labels the context itself
`world-models`	continuous, unbounded	1+ frames, recurrent	yes	yes	`temperature=` as in the paper

JEPA and DiT have no adapter: I-JEPA learns image representations without actions and DiT generates images from a class label, so neither predicts what an action does. as_world_model says so if you pass one.

Loading checkpoints#

load_world_model(path) reads a directory written by save, which records its family in world_model.json. It also reads each family’s own training checkpoints when you name the family:

load_world_model("runs/dreamer/ckpts/61000_ckpt.pt", "dreamer")  # V1 or V2
load_world_model("runs/planet/rssm.pth", "planet")  # bare RSSM weights
load_world_model("runs/iris/iris.pt", "iris")
load_world_model("runs/diamond/checkpoint_10.pt", "diamond")  # no emulator needed
load_world_model("runs/genie/genie_sonic_final.pt", "genie")

World Models trains its two halves separately, so it loads from both files:

from simulusion.world_model import MDRNNWorldModel

model = MDRNNWorldModel.from_checkpoints("vae/best.tar", "mdrnn/best.tar")

Constructor options such as a sampling temperature are saved with the model and can be overridden on load: load_world_model(path, temperature=0.5).

Models and agents#

The protocol covers the model: what happens to the world under an action. The policy that chooses actions stays with each family’s agent, and steppers remain the interface for acting in a real environment. A DIAMOND world model, for example, loads without the Atari emulator or the replay buffer its agent needs for training.

Your own model#

Implement the methods above (and spec and device) and any code written against WorldModel accepts it; isinstance(obj, WorldModel) checks the shape at runtime. Subclassing simulusion.world_model.WorldModelBase gives you input conversion, rollout and the save format for free: implement _encode, _step, _decode, _save_weights and _load_weights.