API reference

This page lists every public class, function, and constant exported from the simulusion top-level package, grouped by category. For signatures and docstrings generated from source, see api_reference.

Factory helpers

Name	Description
`create_config(model, **overrides)`	Return a default config dict for a model family, with optional overrides applied. Models: `dreamer`, `jepa`, `iris`, `dit`, `genie`, `planet`.
`create_model(model, config=None, **overrides)`	Instantiate a model or high-level agent by canonical name (see Public API).
`make_env(env_id, backend="auto", **kwargs)`	Create a Simulusion environment through a named backend. Backends: `dmc`, `gym`, `atari`, `mujoco`, `robotics`, `procgen`, `brax`, `bsuite`, `unity`.
`list_models()`	Return canonical model names accepted by `create_model`.
`list_env_backends()`	Return backend names accepted by `make_env`.
`list_envs(model=None)`	Return known environment IDs, optionally filtered by model family.
`get_model_spec(name)`	Return metadata (`ModelSpec`) for a model name or alias.
`get_env_backend_spec(name)`	Return metadata (`EnvBackendSpec`) for an environment backend.

Data classes

Name	Description
`ModelSpec`	Named tuple describing a registered model (name, import_path, config_path, aliases, description).
`EnvBackendSpec`	Named tuple describing a registered environment backend (name, factory_path, aliases, description).

World model agents

Name	Source	Description
`Dreamer`	`simulusion.models.dreamer`	Base Dreamer world model (RSSM-based, V1-style).
`DreamerV1`	`simulusion.models.dreamer_v1`	DreamerV1 (alias for base Dreamer).
`DreamerV2`	`simulusion.models.dreamer_v2`	DreamerV2 (symlog two-hot heads, balanced KL).
`DreamerV3`	`simulusion.models.dreamer_v3`	DreamerV3 core model (categorical latents, free bits, return normalization).
`DreamerAgent`	`simulusion.models.dreamer`	High-level Dreamer agent with train/evaluate helpers.
`DreamerV1Agent`	`simulusion.models.dreamer_v1`	`DreamerAgent` that trains DreamerV1 (`create_model("dreamer-v1")`).
`DreamerV2Agent`	`simulusion.models.dreamer_v2`	`DreamerAgent` that trains DreamerV2 (`create_model("dreamer-v2")`).
`DreamerV3Agent`	`simulusion.models.dreamer_v3`	`DreamerAgent` that trains DreamerV3 (`create_model("dreamer-v3")`).
`Planet`	`simulusion.models.planet`	PlaNet: Deep Planning Network.
`JEPAAgent`	`simulusion.models.jepa_agent`	I-JEPA agent for self-supervised visual representation learning.
`IRISAgent`	`simulusion.models.iris_agent`	IRIS agent for sample-efficient RL with Transformers.
`Genie`	`simulusion.models.genie`	Genie generative interactive environment.
`create_genie`	`simulusion.models.genie`	Create a Genie model with specified parameters.
`create_genie_small`	`simulusion.models.genie`	Create Genie-small variant (~50M params).
`create_genie_large`	`simulusion.models.genie`	Create Genie-large variant (~11B params).

Genie subcomponents

Name	Source	Description
`LatentActionModel`	`simulusion.models.latent_action_model`	Learns latent actions from pairs of video frames.
`DynamicsModel`	`simulusion.models.dynamics_model`	Transformer-based dynamics for future token prediction.
`create_latent_action_model`	`simulusion.models.latent_action_model`	Factory for LatentActionModel.
`create_dynamics_model`	`simulusion.models.dynamics_model`	Factory for DynamicsModel.

State-space models

Name	Source	Description
`RSSM`	`simulusion.models.rssm`	Recurrent State-Space Model (standalone).
`CategoricalRSSM`	`simulusion.models.categorical_rssm`	DreamerV3 RSSM with straight-through categorical latents.
`RecurrentStateSpaceModel`	`simulusion.models.rssm`	Alias for RSSM.
`DreamerRSSM`	`simulusion.models.dreamer_rssm`	RSSM variant used in Dreamer training loop.
`ModularRSSM`	`simulusion.models.modular_rssm`	Modular RSSM with swappable encoder/backbone/decoder.
`create_modular_rssm`	`simulusion.models.modular_rssm`	Factory for ModularRSSM.

Diffusion models

Name	Source	Description
`DiT`	`simulusion.models.diffusion`	Diffusion Transformer model.
`create_dit`	`simulusion.models.diffusion`	Factory for DiT.
`PatchEmbed`	`simulusion.models.diffusion`	Image-to-patch embedding layer.
`PatchUnEmbed`	`simulusion.models.diffusion`	Patch-to-image un-embedding layer.
`DDPM`	`simulusion.models.diffusion`	Denoising Diffusion Probabilistic Model.
`ActorCriticNetwork`	`simulusion.models.diffusion`	Actor-critic head for DIAMOND-style RL.
`RewardTerminationModel`	`simulusion.models.diffusion`	Reward + termination predictor for DIAMOND.
`sinusoidal_time_embedding`	`simulusion.models.diffusion`	Time-step embedding for diffusion.

Vision components

Name	Source	Description
`ConvEncoder`	`simulusion.vision.dreamer_encoder`	Dreamer convolutional encoder (image → embedding).
`DreamerV3Encoder`	`simulusion.vision.dreamer_v3_nets`	DreamerV3 encoder (CNN for images, MLP for vectors).
`DreamerV3Decoder`	`simulusion.vision.dreamer_v3_nets`	DreamerV3 observation decoder.
`DreamerV3Head`	`simulusion.vision.dreamer_v3_nets`	DreamerV3 MLP head (two-hot reward/critic, Bernoulli continue).
`DreamerV3Actor`	`simulusion.vision.dreamer_v3_nets`	DreamerV3 actor for discrete and continuous actions.
`CNNEncoder`	`simulusion.vision.planet_encoder`	PlaNet CNN encoder (image → embedding).
`IRISEncoder`	`simulusion.vision.iris_encoder`	IRIS encoder (image → discrete tokens).
`ConvDecoder`	`simulusion.vision.dreamer_decoder`	Dreamer convolutional decoder (latent → image distribution).
`CNNDecoder`	`simulusion.vision.planet_decoder`	PlaNet CNN decoder.
`DenseDecoder`	`simulusion.vision.dreamer_decoder`	MLP decoder for reward/value/discount.
`ActionDecoder`	`simulusion.vision.dreamer_decoder`	Dreamer policy head (latent → tanh-squashed action).
`IRISDecoder`	`simulusion.vision.iris_decoder`	IRIS decoder (tokens → image).
`VideoTokenizer`	`simulusion.vision.video_tokenizer`	Genie VQ-VAE video tokenizer.
`create_video_tokenizer`	`simulusion.vision.video_tokenizer`	Factory for VideoTokenizer.
`VectorQuantizer`	`simulusion.vision.vq_layer`	VQ-VAE vector quantization layer.
`VectorQuantizerEMA`	`simulusion.vision.vq_layer`	VQ-VAE with EMA codebook updates.
`TanhBijector`	`simulusion.vision.dreamer_decoder`	Tanh bijection for action squashing.
`SampleDist`	`simulusion.vision.dreamer_decoder`	MC-sampled distribution statistics.

Config classes

Name	Source	Description
`DreamerConfig`	`simulusion.configs.dreamer_config`	Dreamer hyperparameter config.
`DreamerV3Config`	`simulusion.configs.dreamer_v3_config`	DreamerV3 hyperparameters with 12M-400M model-size presets.
`JEPAConfig`	`simulusion.configs.jepa_config`	JEPA hyperparameter config.
`DiTConfig`	`simulusion.configs.dit_config`	DiT hyperparameter config.
`get_dit_config`	`simulusion.configs.dit_config`	Factory for DiTConfig with presets.
`DiamondConfig`	`simulusion.configs.diamond_config`	DIAMOND hyperparameter config.
`IRISConfig`	`simulusion.configs.iris_config`	IRIS hyperparameter config.
`GenieConfig`	`simulusion.configs.genie_config`	Genie hyperparameter config.
`GenieSmallConfig`	`simulusion.configs.genie_config`	Genie-small preset config.
`STTransformerConfig`	`simulusion.configs.st_transformer_config`	ST-Transformer config.
`VideoTokenizerConfig`	`simulusion.configs.video_tokenizer_config`	Video tokenizer config.
`LatentActionModelConfig`	`simulusion.configs.lam_config`	Latent action model config.
`DynamicsModelConfig`	`simulusion.configs.dynamics_config`	Dynamics model config.

Constants

Name	Description
`MODEL_SPECS`	Dict of all built-in model specs (name → `ModelSpec`).
`ENV_BACKEND_SPECS`	Dict of all built-in environment backend specs.
`ATARI_100K_GAMES`	List of Atari 100K benchmark game names.
`HUMAN_SCORES`	Dict of human baseline scores for Atari 100K.
`RANDOM_SCORES`	Dict of random baseline scores for Atari 100K.

Memory / replay buffers

Name	Source	Description
`ReplayBuffer`	`simulusion.memory.dreamer_memory`	Dreamer ring buffer (transitions → sequences).
`DreamerV3ReplayBuffer`	`simulusion.memory.dreamer_v3_memory`	DreamerV3 ring buffer with an online queue mixed into each batch.
`Memory`	`simulusion.memory.planet_memory`	Episode-based memory for PlaNet.
`Episode`	`simulusion.memory.planet_memory`	Single episode recording.
`IRISReplayBuffer`	`simulusion.memory.iris_memory`	Ring buffer for IRIS (uint8 images).
`IRISOnPolicyBuffer`	`simulusion.memory.iris_memory`	On-policy buffer for episode collection.

Environments and wrappers

Name	Source	Description
`DeepMindControlEnv`	`simulusion.envs.dmc_env`	DeepMind Control Suite adapter.
`DMLabEnv`	`simulusion.envs.dmlab_env`	DeepMind Lab adapter.
`make_dmlab_env`	`simulusion.envs.dmlab_env`	Factory for DeepMind Lab.
`DMLAB_LEVELS`	`simulusion.envs.dmlab_env`	Available DMLab level names.
`GymImageEnv`	`simulusion.envs.gym_env`	Gymnasium image adapter.
`make_gym_env`	`simulusion.envs.gym_env`	Factory for GymImageEnv.
`MuJoCoImageEnv`	`simulusion.envs.mujoco_env`	MuJoCo image adapter.
`make_mujoco_env`	`simulusion.envs.mujoco_env`	Factory for MuJoCo environments.
`MujocoEnv`	`simulusion.envs.mujoco_env`	Alias of `MuJoCoImageEnv`.
`BraxImageEnv`	`simulusion.envs.brax_env`	Brax image adapter.
`make_brax_env`	`simulusion.envs.brax_env`	Factory for BraxImageEnv.
`BSuiteImageEnv`	`simulusion.envs.bsuite_env`	BSuite image adapter.
`make_bsuite_env`	`simulusion.envs.bsuite_env`	Factory for BSuiteImageEnv.
`list_available_bsuite_ids`	`simulusion.envs.bsuite_env`	List BSuite environment IDs.
`make_atari_env`	`simulusion.envs.atari_env`	Factory for Atari ALE environments.
`list_available_atari_envs`	`simulusion.envs.atari_env`	List available Atari game IDs.
`make_atari_vector_env`	`simulusion.envs.atari_env`	Factory for vectorized Atari.
`make_diamond_atari_env`	`simulusion.envs.diamond_atari`	DIAMOND-style Atari preprocessing.
`make_procgen_env`	`simulusion.envs.procgen_env`	Factory for Procgen environments.
`make_robotics_env`	`simulusion.envs.robotics_env`	Factory for Gymnasium Robotics.
`register_gymnasium_robotics_envs`	`simulusion.envs.robotics_env`	Register Robotics envs.
`list_gymnasium_robotics_envs`	`simulusion.envs.robotics_env`	List installed Robotics envs.
`UnityMLAgentsEnv`	`simulusion.envs.unity_env`	Unity ML-Agents adapter.
`make_unity_mlagents_env`	`simulusion.envs.unity_env`	Factory for UnityMLAgentsEnv.
`WorldModelEnv`	`simulusion.envs.world_model_env`	Environment inside a learned world model.
`make_world_model_env`	`simulusion.envs.world_model_env`	Factory for WorldModelEnv.
`TimeLimit`	`simulusion.envs.wrappers`	Episode time limit wrapper.
`ActionRepeat`	`simulusion.envs.wrappers`	Action repeat wrapper.
`NormalizeActions`	`simulusion.envs.wrappers`	Action normalization to [-1, 1].
`ObsDict`	`simulusion.envs.wrappers`	Observation-to-dict conversion.
`OneHotAction`	`simulusion.envs.wrappers`	Discrete to one-hot action conversion.
`RewardObs`	`simulusion.envs.wrappers`	Reward observation injection.
`ResizeImage`	`simulusion.envs.wrappers`	Image resizing wrapper.
`RenderImage`	`simulusion.envs.wrappers`	Render-based image observation.
`SelectAction`	`simulusion.envs.wrappers`	Action selection wrapper.

Controllers and policies

Name	Source	Description
`RSSMPolicy`	`simulusion.controller`	RSSM-based policy for Dreamer.
`RolloutGenerator`	`simulusion.controller`	Policy rollouts in the environment.
`IRISActor`	`simulusion.controller`	IRIS actor head.
`IRISCritic`	`simulusion.controller`	IRIS critic head.
`IRISPolicy`	`simulusion.controller`	IRIS combined actor-critic policy.
`CNNFeatureExtractor`	`simulusion.controller`	CNN feature extractor for policy inputs.

Export

Name	Source	Description
`export_any(obj, path, format, ...)`	`simulusion.export`	Export a model or agent (`exported_program`, `aoti`, ONNX, TensorRT, TorchScript).
`export_model(module, path, format, ...)`	`simulusion.export`	Export a raw nn.Module.
`load_exported(path, format=None)`	`simulusion.export`	Load any exported artifact as a callable.
`verify_export(module, exported, example_inputs)`	`simulusion.export`	Check an artifact against its eager module; returns max abs error.
`ExportableAgentMixin`	`simulusion.export`	Mixin that adds `.export()` to custom agents.

Inference and deployment

See Deployment and export.

Name	Source	Description
`optimize_for_inference(module, ...)`	`simulusion.inference`	Wrap a module with a precision policy and optional `torch.compile` / CUDA graphs.
`InferenceModel`	`simulusion.inference`	The wrapper returned by `optimize_for_inference`.
`inference_context(device, precision)`	`simulusion.inference`	`inference_mode` plus autocast at one resolved precision.
`make_stepper(agent)`	`simulusion.inference`	Uniform `init_state` / `observe` / `act` / `imagine` step interface.
`DreamerStepper`, `IRISStepper`	`simulusion.inference`	Steppers for Dreamer and IRIS.
`DreamerStepModule`	`simulusion.inference`	One Dreamer observe+act step as a pure, exportable module.
`benchmark_step(fn, *args)`	`simulusion.inference`	Latency percentiles, throughput and peak memory.
`rollout_drift(reference, candidate, state, inputs)`	`simulusion.inference`	Closed-loop divergence of an optimized step from a reference.
`quantize_weights(module)`	`simulusion.inference`	Weight-only int8 quantization of `nn.Linear` layers.
`save_bundle(dir, module, ...)` / `load_bundle(dir)`	`simulusion.inference`	Deployment bundles: weights, config, artifacts, manifest.
`STKVCache`	`simulusion.blocks.st_transformer`	Temporal KV cache for incremental Genie / ST-transformer generation.

World model protocol

See One interface for every world model.

Training

See One Trainer for Every Model.

Name	Source	Description
`Trainer`	`simulusion.training.trainer`	Runs any training task to its budget on one device: precision, logging, evaluation, early stopping, checkpoints and resume, world model export.
`TrainerConfig`	`simulusion.training.trainer`	Settings shared by every run (the `trainer:` section of a run config).
`TrainResult`	`simulusion.training.trainer`	What `Trainer.fit` returns: run directory, progress, final metrics, checkpoint and export paths.
`TrainingTask`	`simulusion.training.task`	Base class of a family's training (networks, optimizers, data, one step).
`register_task` / `list_tasks`	`simulusion.training.task`	Register a task under a name / list the registered ones.
`run_training(config, overrides)`	`simulusion.training.run`	Train from a YAML run config (or mapping) plus `key=value` overrides.
`MixedPrecision`	`simulusion.training.precision`	Autocast dtype and loss scaling for one precision setting.
`GradientUpdate`	`simulusion.training.precision`	One optimizer's step: backward, unscale, clip, step, with accumulation and gradient all-reduce.

Data

See Episodes and datasets.

Name	Source	Description
`EpisodeFolder`	`simulusion.data`	A folder of episodes (`.npz`, videos, `.npy`), the dataset every task reads.
`EpisodeWriter`	`simulusion.data`	Append episodes to a folder and record its `meta.json`.
`LeRobotEpisodes`	`simulusion.data`	Local LeRobot datasets (v2.x and v3.0) as episodes.

Agents

Name	Source	Description
`CEMPlanner`	`simulusion.agents`	Cross-entropy-method planning over any world model's predicted returns.
`PlanningAgent`	`simulusion.agents`	Acts in a real environment by planning in a world model each step.

Name	Source	Description
`as_world_model(model)`	`simulusion.world_model`	Wrap a trained Dreamer, PlaNet, IRIS, DIAMOND, Genie or (ConvVAE, MDN-RNN) model in the common protocol.
`load_world_model(path, family=None)`	`simulusion.world_model`	Load a saved world model, or a family's own training checkpoint.
`WorldModel`	`simulusion.world_model`	The protocol: `encode` / `step` / `rollout` / `decode` / `save`.
`WorldModelSpec`	`simulusion.world_model`	Observation shape, action space, context length and which heads a model has.
`list_world_model_families()`	`simulusion.world_model`	The family names `load_world_model` accepts.

Reward and value models

Name	Source	Description
`RewardModel`	`simulusion.reward`	Base reward model.
`ValueModel`	`simulusion.reward`	Base value model.
`DreamerRewardModel`	`simulusion.reward`	Dreamer reward predictor.
`DreamerValueModel`	`simulusion.reward`	Dreamer value function.

Transformer blocks

Name	Source	Description
`STTransformer`	`simulusion.blocks`	Spatiotemporal Transformer (Genie).
`MultiHeadSelfAttention`	`simulusion.blocks`	Multi-head self-attention.
`MultiHeadAttention`	`simulusion.blocks`	Multi-head cross-attention.
`AdaLNNormalization`	`simulusion.blocks`	Adaptive layer norm for diffusion.
`RMSNorm`	`simulusion.blocks`	Root mean square layer norm.

Plugin registry

Name	Source	Description
`register_world_model(name, import_path, ...)`	`simulusion.registry`	Register a custom world model architecture.
`deregister_world_model(name)`	`simulusion.registry`	Remove a registered model.
`get_registered_model_spec(name)`	`simulusion.registry`	Look up a registered model spec.
`list_registered_models()`	`simulusion.registry`	List all externally registered model names.
`register_env_backend(name, factory_path, ...)`	`simulusion.registry`	Register a custom environment backend.
`deregister_env_backend(name)`	`simulusion.registry`	Remove a registered env backend.
`list_registered_env_backends()`	`simulusion.registry`	List all registered env backends.

Deprecation utilities

Name	Source	Description
`deprecated(version, reason)`	`simulusion.utils.deprecation`	Decorator to mark functions/classes as deprecated.
`deprecated_class(version, alternative)`	`simulusion.utils.deprecation`	Shortcut for deprecating a class.
`deprecated_function(version, alternative)`	`simulusion.utils.deprecation`	Shortcut for deprecating a function.

General utilities

Name	Source	Description
`Logger`	`simulusion.utils`	Logging utility.
`FreezeParameters`	`simulusion.utils`	Context manager to freeze model parameters.
`compute_return(rewards, values, gamma, lambda_)`	`simulusion.utils`	Compute GAE or λ-return.
`preprocess_obs(obs)`	`simulusion.utils`	Observation preprocessing (resize, normalize).

Version

Name	Description
`__version__`	Package version string (semver).

Edit this page on GitHub

Updated Nov 1, 2026