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.

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 Episodes and datasets). What comes out is a simulusion.world_model.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

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.

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.