Package Overview

Simulusion is organized into focused modules so you can use only the pieces you need.

Quick Import (Public API)

For applications and examples, import the simulusion package (installed from PyPI as simulusion). It mirrors the Simulusion implementation package and exposes the same lazy public API without importing optional training backends until you use them.

import simulusion

print(simulusion.list_models())
agent = simulusion.create_model("dreamer", env="walker-walk", total_steps=1_000_000)
env = simulusion.make_env("CartPole-v1", backend="gym")

Use simulusion for direct component imports as well as factory helpers:

from simulusion import DreamerAgent, DreamerConfig

cfg = DreamerConfig()
cfg.env = "walker-walk"
agent = DreamerAgent(cfg)

Available Exports

Category	Exports
Friendly factories	create_config, create_model, make_env, list_models, list_env_backends, list_envs
Models / Agents	Dreamer, DreamerV1, DreamerV2, DreamerAgent, Planet, JEPAAgent, IRISAgent, Genie, create_genie, DiT, create_dit
State-space models	RSSM, RecurrentStateSpaceModel, DreamerRSSM, ModularRSSM, create_modular_rssm
Vision	ConvEncoder, CNNEncoder, IRISEncoder, ConvDecoder, CNNDecoder, DenseDecoder, ActionDecoder, IRISDecoder, VideoTokenizer, create_video_tokenizer
Quantization	VectorQuantizer, VectorQuantizerEMA
Configs	DreamerConfig, JEPAConfig, DiTConfig, DiamondConfig, IRISConfig, GenieConfig, GenieSmallConfig, STTransformerConfig, VideoTokenizerConfig, LatentActionModelConfig, DynamicsModelConfig
Environments	make_atari_env, make_gym_env, make_mujoco_env, make_robotics_env, make_brax_env, make_procgen_env, GymImageEnv, ProcgenImageEnv, DeepMindControlEnv, DMLabEnv, make_dmlab_env, UnityMLAgentsEnv, TimeLimit, ActionRepeat, wrappers, etc.
Memory	ReplayBuffer, Memory, Episode, IRISReplayBuffer, IRISOnPolicyBuffer
Reward / Value	RewardModel, ValueModel, DreamerRewardModel, DreamerValueModel
Controllers	RSSMPolicy, RolloutGenerator, IRISPolicy, IRISActor, IRISCritic, CNNFeatureExtractor
Transformer blocks	STTransformer, MultiHeadSelfAttention, MultiHeadAttention, AdaLNNormalization, RMSNorm
Diffusion	DiT, DDPM, PatchEmbed, PatchUnEmbed, ActorCriticNetwork, RewardTerminationModel
Genie subcomponents	LatentActionModel, DynamicsModel, create_latent_action_model, create_dynamics_model
Export	export_any, export_model, ExportableAgentMixin
Registry / plugins	register_world_model, deregister_world_model, register_env_backend, deregister_env_backend
Deprecation	deprecated, deprecated_class, deprecated_function
Utilities	Logger, FreezeParameters, compute_return, preprocess_obs, __version__

Example usage:

import simulusion

# Training
agent = simulusion.create_model("dreamer", env="walker-walk", total_steps=1_000_000)
agent.train()

Core Modules

The module paths below are the public simulusion.* surface (from simulusion.models import Dreamer).

• simulusion.models: High-level models and agents (Dreamer, DreamerAgent, Planet, JEPAAgent)

• simulusion.configs: Configuration containers for Dreamer, JEPA, and diffusion runs

• simulusion.training: Script-style training entrypoints for world models (VAE, MDNRNN, Controller, Planet, RSSM, JEPA)

Environment Integration

• simulusion.envs: DMC, Gym/Gymnasium, Atari, MuJoCo, Unity ML-Agents adapters

• simulusion.envs.wrappers: Action repeat, normalization, time limits

World Model Building Blocks

• simulusion.models.dreamer_rssm: Recurrent state-space model used by Dreamer

• simulusion.models.modular_rssm: Modular RSSM with swappable encoder/decoder/backbone for research experiments

• simulusion.vision: Encoders/decoders and action heads for latent dynamics models

• simulusion.reward: Reward and value prediction heads

• simulusion.observations: Symbolic and visual observation reconstruction modules

Representation Learning and Diffusion

• simulusion.models.vit: Vision Transformer and JEPA predictor components

• simulusion.models.diffusion: DDPM scheduler and DiT model implementation

• simulusion.masks: Mask collators for JEPA-style context/target masking

Data and Memory

• simulusion.datasets: CIFAR-10, ImageNet-1K, and generic ImageFolder dataset loaders

• simulusion.memory: Replay buffers for Dreamer and episode-based memory for PlaNet/RSSM

Utilities

• simulusion.utils: Logging, parameter freezing, transforms

• simulusion.transforms: Data augmentation pipelines

• simulusion.benchmarks: CLI and reporting utilities

Which API Should I Use?

• End-to-end Dreamer training: DreamerAgent

• End-to-end JEPA training: JEPAAgent

• World model training scripts: simulusion.training modules (e.g., train_world_model for VAE+MDNRNN+Controller pipeline)

• Low-level model experimentation: Dreamer, RSSM, decoder/encoder modules

• Custom world model architectures: ModularRSSM with swappable encoder/decoder/backbone

• Custom data pipelines: make_cifar10, make_imagenet1k, make_imagefolder