One Trainer for Every Model

Every family in Simulusion trains on the same loop, simulusion.training.Trainer, from the same kind of file: a YAML run config that names a task and fills its sections. The trainer owns what is common to all of them: the device, one precision setting, seeding, the step budget, logging, evaluation, early stopping, checkpoints and resuming, and exporting the trained world model.

Train from the command line

# A shipped config by name, with key=value overrides for any field.
simulusion train --config dreamer model.env=cheetah-run trainer.precision=bf16

# Or name the task and set only what you need; the rest keeps its defaults.
simulusion train genie data.path=data/clips trainer.max_steps=20000

# See the fully resolved config without training.
simulusion train --config iris --print-config

Nested keys use dots (model.lr=3e-4, trainer.eval_every=1000). A key that does not exist is an error that names its full path, so a typo fails before the run starts instead of being ignored. simulusion train --legacy MODEL still reaches each family’s older entrypoint.

Train from Python

import simulusion

result = simulusion.run_training(
    "simulusion/configs/runs/genie.yaml",
    ["data.path=data/clips", "trainer.max_steps=500"],
)
model = simulusion.load_world_model(result.world_model_path)

or build the pieces yourself:

from simulusion.training import Trainer, TrainerConfig
from simulusion.training.tasks import GenieTask
from simulusion.training.tasks.genie import GenieTaskConfig

config = GenieTaskConfig()
config.data.path = "data/clips"
result = Trainer(TrainerConfig(max_steps=500, precision="bf16")).fit(GenieTask(config))

Tasks

A task is one family’s training: its networks, optimizers, where batches come from and what one unit of progress is. The trainer’s budget (trainer.max_steps) counts that unit.

Task	Trains	Data	Progress
dreamer, dreamer-v1, dreamer-v2	world model + actor + critic	environment	env steps
dreamer-wm	Dreamer’s world model alone	episode folder	optimizer steps
planet	RSSM, acting by CEM planning	environment	epochs
iris	tokenizer + transformer + actor-critic (paper settings)	environment	epochs
iris-wm	IRIS’s world model alone	episode folder	optimizer steps
diamond	denoiser + reward model + actor-critic	environment	epochs
diamond-wm	DIAMOND’s world model alone	episode folder	optimizer steps
genie	tokenizer + latent actions + dynamics	episode folder / videos	optimizer steps
vae, mdrnn	World Models’ V, then M	episode folder	optimizer steps
controller	World Models’ C, by CMA-ES (real env or dream)	world model	generations
dit	DiT image diffusion	image dataset	optimizer steps
jepa	I-JEPA image representations	image dataset	optimizer steps

The -wm tasks train a world model without its agent: no environment and no policy, only recorded episodes (your own, or converted from LeRobot; see Your Own Data: Episodes and LeRobot). What comes out is a WorldModel, usable by any agent that speaks the protocol, such as simulusion.agents.PlanningAgent:

from simulusion.agents import CEMPlanner, PlanningAgent, run_episode

agent = PlanningAgent(
    simulusion.load_world_model("runs/dreamer-wm_20261020-120000/world_model"),
    CEMPlanner(horizon=12, candidates=500),
)
episode_return = run_episode(env, agent)

Run configs

task: genie
trainer:              # TrainerConfig: shared by every task
  max_steps: 20000
  precision: auto     # fp32 | bf16 | fp16 | auto
  eval_every: 1000
  checkpoint_every: 2000
  early_stopping_metric: val_recon_loss
model:                # the family's own config (GenieConfig here)
  num_frames: 16
data:                 # where batches come from
  path: data/clips
  val_fraction: 0.05

Shipped configs for every task live in simulusion/configs/runs/; --config <name> finds them by name.

Trainer settings

Field	Meaning
max_steps	Budget in the task’s progress unit; the task’s default when unset.
seed, deterministic	Seeds Python, NumPy and PyTorch; deterministic asks for deterministic kernels.
precision	fp32, bf16, fp16 (loss-scaled) or auto (bf16 where supported, fp16 on older GPUs, fp32 on CPU). Replaces the per-family use_amp flags for runs on the trainer.
device	null (or auto) picks CUDA, then MPS, else the CPU; cpu, cuda, cuda:1 or mps choose one. Training runs on one device.
grad_accumulation	Batches per optimizer step (dataset-driven tasks).
log_every, eval_every, checkpoint_every	Intervals in progress units. Progress may jump (env steps), so each fires when a multiple is crossed.
keep_checkpoints	Periodic checkpoints kept; the final one always is.
out_dir, run_name	Run directory: out_dir/run_name, timestamped when run_name is unset.
resume	A checkpoint, a run directory, or auto (the newest checkpoint of out_dir/run_name).
early_stopping_metric, early_stopping_mode, patience, min_delta	Stop when an evaluation metric stops improving.
console	Also print metrics to the console (they always go to metrics.jsonl).
export_world_model	Write the world model to <run>/world_model at the end.

What a run leaves behind

runs/genie_20261020-120000/
    config.yaml            # the fully resolved config: `--config` it to repeat the run
    run.json               # git commit and dirty flag, package versions, GPUs, command line
    metrics.jsonl          # every logged metric (train/..., eval/...)
    checkpoints/
        step_00002000/     # resumable: trainer.resume=<this dir> or resume=auto
    world_model/           # load_world_model(...) reads it
    model/                 # tasks that are not world models (dit, jepa, vae) export here

Checkpoints are written to a staging directory and renamed when complete, so an interrupted write never looks like a checkpoint. They hold the task’s weights, optimizer and loss-scaler states, the loop position and RNG states (resuming is bit-exact), and for online tasks the replay buffer, so a resumed run trains on the same data. Ctrl+C writes one before exiting.

Writing a task

from dataclasses import dataclass, field

import torch

from simulusion.training import TrainingTask, register_task
from simulusion.training.tasks.common import SequenceDataConfig, sequence_loaders


@dataclass
class MyConfig:
    lr: float = 3e-4
    data: SequenceDataConfig = field(default_factory=SequenceDataConfig)


@register_task("my-model")
class MyTask(TrainingTask):
    config_cls = MyConfig  # its fields are the run config's sections

    def setup(self, trainer):
        self.model = trainer.prepare_module(MyModel())
        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config.lr)
        self.update = trainer.gradient_update(
            optimizer, self.model.parameters(), max_grad_norm=1.0
        )
        self.batches, self.val, _ = sequence_loaders(
            self.config.data, trainer, sequence_length=8, image_size=64
        )

    def train_step(self):
        batch = next(self.batches)
        with self.trainer.precision.autocast():
            loss = self.model(batch["frames"])
        self.update(loss)
        return {"loss": loss.item()}

    def state_dict(self):
        return {"model": self.model.state_dict()}

    def load_state_dict(self, state):
        self.model.load_state_dict(state["model"])

Optional hooks: progress() (a budget in env steps or epochs), evaluate(), final_evaluation(), world_model() (export through the protocol), export(directory) (other models), save_checkpoint() / load_checkpoint() (state too large for one file, such as a replay buffer) and close().

The older entrypoints

The per-family entrypoints keep their signatures and still write the files they always wrote, but each now runs on the shared Trainer: it builds its family’s task around the object it was called on and runs it with simulusion.training.legacy.fit_legacy. There is one training loop, and the old calls gain resumable checkpoints and a record of every run.

Entrypoint	Runs as	Still writes
DreamerAgent.train(total_steps)	dreamer task	<logdir>/ckpts/<step>_ckpt.pt; metrics and eval videos through the agent’s own logger
IRISTrainer.train(total_epochs, eval_interval, save_dir)	iris task	best_<game>.pt, final_<game>.pt, checkpoint_<epoch>.pt; returns the per-epoch metrics, sets final_eval_return
DiamondAgent.train()	diamond task	checkpoint_<epoch>.pt in checkpoint_dir (with its replay)
Planet.train(...), train_planet.main	planet task	ckpt_<epoch>.pth, vid_<epoch>, the TensorBoard summary
GenieTrainer.train(loaders, ...)	genie task	genie_step_<N>.pt in checkpoint_dir
DiT.fit(...), train_dit	dit task	dit_model.pth, dit_model_epoch<N>.pth, generated_samples.png
JEPAAgent.train()	jepa task	<write_tag>-latest.pth.tar, <write_tag>-ep<N>.pth.tar
train_convae, train_mdn_rnn	vae, mdrnn tasks	checkpoint.tar, best.tar, *_final.pth (and V’s samples)
train_controller	controller task	<logdir>/ctrl/best.tar

What they now all do the same way:

• The run directory is the folder the entrypoint always used. The trainer adds checkpoints/, run.json, run_config.yaml and train_metrics.jsonl there, under names that never overwrite a family’s own config.yaml or metrics.jsonl.

• The object keeps its own device and precision (its use_amp).

• Evaluation, early stopping and “best” checkpoints are the trainer’s eval_every, early_stopping_metric and best_metric, set from the entrypoint’s own arguments.

Small differences from the loops they replace:

• Periodic checkpoints are taken when progress crosses an interval and named by the step they were written at. Dreamer no longer writes one before any training, and the last one is written after the final evaluation.

• IRIS evaluates after epochs eval_interval, 2 * eval_interval, ... (not also after epoch 0), and its best checkpoint can come from the final evaluation too. DIAMOND likewise.

• train_mdn_rnn samples latents from the VAE as batches are drawn instead of precomputing them once (use_precomputed_latents is accepted and ignored). V and M read rollouts as episodes: windows of seq_len + 1 frames, resized to 64x64.

• train_controller scores candidates in one process (max_workers and display no longer apply) and takes the controller’s sizes from the trained models.

• JEPAAgent.train with enable_sweep (W&B sweeps) still runs train_jepa.main’s own loop.