API Reference #

classmethod from_pretrained(pretrained_model_name_or_path, *, config=None, checkpoint_filename=None, config_filename='config.yaml', repo_type=None, revision=None, map_location=None, **overrides)[source]

Load a Dreamer checkpoint from a local path/directory or the HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
config (DreamerConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
map_location (str | device | None)
overrides (Any)

Return type:

parameter_count(trainable_only=False)[source]

Return the total number of parameters owned by the Dreamer modules.

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]

Return a compact parameter-count summary for the Dreamer modules.

Return type:: dict[str, Any]

set_precision(precision)[source]

Train at precision from now on (fresh loss scalers).

Parameters:: precision (MixedPrecision)
Return type:: None

update_world_model(obs, acs, rews, nonterms, continues=None)[source]

One world-model optimizer step; returns (loss, grad_norm).

Inputs are time-major, as the replay buffer samples them: obs (L, B, C, H, W) uint8, acs (L, B, A), rews, nonterms and continues (L, B, 1), where row t holds the action taken at obs[t] and what followed it.

Parameters:

obs (Tensor)
acs (Tensor)
rews (Tensor)
nonterms (Tensor)
continues (Tensor | None)

Return type:

tuple[Tensor, Tensor | None]

world_model_loss(obs, acs, rews, nonterms, continues=None)[source]

World-model loss.

nonterms masks recurrent state at every episode boundary. continues is the discount-head target, which is 0 only on true terminations; it defaults to nonterms when not supplied.

Parameters:

obs (Tensor)
acs (Tensor)
rews (Tensor)
nonterms (Tensor)
continues (Tensor | None)

Return type:

Tensor

actor_loss()[source]

Return type:: Tensor

value_loss()[source]

Return type:: Tensor

train_one_batch()[source]

Return type:: list[float]

act_with_world_model(obs, prev_state, prev_action, explore=False)[source]

Parameters:

obs (Any)
prev_state (Any)
prev_action (Tensor)
explore (bool)

Return type:

tuple

act_and_collect_data(env, collect_steps)[source]

Parameters:

env (Any)
collect_steps (int)

Return type:

ndarray

evaluate(env, eval_episodes, render=False)[source]

Parameters:

env (Any)
eval_episodes (int)
render (bool)

Return type:

tuple

collect_random_episodes(env, seed_steps)[source]

Parameters:

env (Any)
seed_steps (int)

Return type:

ndarray

save(save_path)[source]

Parameters:: save_path (str)
Return type:: None

restore_checkpoint(ckpt_path, map_location=None)[source]

Parameters:

ckpt_path (str | Path)
map_location (Any)

Return type:

None

class simulusion.DreamerV1(args, obs_shape, action_size, device, restore=False)[source]

Bases: Dreamer

Parameters:

args (Any)
obs_shape (Any)
action_size (int)
device (device | str)
restore (bool)

class simulusion.DreamerV2(args, obs_shape, action_size, device, restore=False)[source]

Bases: Dreamer

Parameters:

args (Any)
obs_shape (Any)
action_size (int)
device (device | str)
restore (bool)

class simulusion.DreamerAgent(config=None, **kwargs)[source]

High-level user API for running Dreamer experiments end to end.

It builds environments from config, initializes seeds and logging, instantiates Dreamer, and exposes simple train() / evaluate() methods.

Parameters:

config (Any)
kwargs (Any)

core_cls: alias of Dreamer

algo_name: str | None = None: Run-folder name used when the config leaves algo at its default.

classmethod from_config(config=None, **overrides)[source]

Build a high-level Dreamer agent from a config object, dict, or YAML file.

Parameters:

config (DreamerConfig | dict[str, Any] | str | Path | None)
overrides (Any)

Return type:

Create a Dreamer agent and restore weights from a local path or HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
config (DreamerConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
map_location (str | device | None)
overrides (Any)

Return type:

parameter_count(trainable_only=False)[source]

Return the total number of Dreamer parameters.

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]

Return a compact parameter-count summary for the wrapped Dreamer model.

Return type:: dict[str, Any]

train(total_steps=None)[source]

Train for total_steps environment steps (args.total_steps by default).

Runs on the shared Trainer through the dreamer task, with this agent’s environments, model and logger, so metrics and evaluation videos go where they always went. Evaluation runs every test_interval steps, early stopping on its average reward when early_stopping is set, and checkpoints land in <logdir>/ckpts/<step>_ckpt.pt every checkpoint_interval steps and at the end, as before. The trainer adds a resumable copy of each under <logdir>/checkpoints/, plus run.json, run_config.yaml and train_metrics.jsonl.

Parameters:: total_steps (int | None)
Return type:: None

evaluate()[source]

Return type:: tuple

class simulusion.DreamerV1Agent(config=None, **kwargs)[source]

Bases: DreamerAgent

High-level agent that trains DreamerV1 (create_model("dreamer-v1")).

Parameters:

config (Any)
kwargs (Any)

core_cls: alias of DreamerV1

algo_name: str | None = 'Dreamerv1': Run-folder name used when the config leaves algo at its default.

class simulusion.DreamerV2Agent(config=None, **kwargs)[source]

Bases: DreamerAgent

High-level agent that trains DreamerV2 (create_model("dreamer-v2")).

Parameters:

config (Any)
kwargs (Any)

core_cls: alias of DreamerV2

algo_name: str | None = 'Dreamerv2': Run-folder name used when the config leaves algo at its default.

class simulusion.Planet(env, bit_depth=5, device=None, state_size=200, latent_size=30, embedding_size=1024, memory_size=100, policy_cfg=None, headless=False, max_episode_steps=None, action_repeats=1, results_dir=None)[source]

High-level Planet wrapper.

Usage example:: from simulusion.models.planet import Planet p = Planet(env=’CartPole-v1’, bit_depth=5) p.train(epochs=50)

Parameters:

env (Any)
bit_depth (int)
device (device | None)
state_size (int)
latent_size (int)
embedding_size (int)
memory_size (int)
policy_cfg (dict | None)
headless (bool)
max_episode_steps (int | None)
action_repeats (int)
results_dir (str | None)

warmup(n_episodes=1, random_policy=True)[source]

Collect n_episodes of rollouts into memory (used as warmup).

Parameters:

n_episodes (int)
random_policy (bool)

Return type:

None

train(epochs=100, steps_per_epoch=150, batch_size=32, H=50, beta=1.0, save_every=25, record_grads=False, results_dir=None, scheduler_type='step', scheduler_kwargs=None)[source]

Train for epochs epochs; returns the results directory.

Runs on the shared Trainer through the planet task, with this model, planner and replay. Each epoch is steps_per_epoch gradient steps on batch_size chunks of H steps, then one exploration and one evaluation episode. As before, metrics go to the TensorBoard summary in results_dir, every evaluation episode is saved as vid_<epoch>, and the RSSM is written to ckpt_<epoch>.pth every save_every epochs.

Parameters:

scheduler_type (str) – Type of scheduler to use (“step”, “cosine”, “exponential”, “plateau”, None)
scheduler_kwargs (dict) – Additional arguments for the scheduler
epochs (int)
steps_per_epoch (int)
batch_size (int)
H (int)
beta (float)
save_every (int)
record_grads (bool)
results_dir (str | None)

Return type:

str

class simulusion.JEPAAgent(config=None, **kwargs)[source]

Convenience interface for configuring and launching JEPA training runs.

Accepts a JEPAConfig plus keyword overrides, prepares output folders, and delegates execution to the JEPA training entrypoint.

Parameters:

config (JEPAConfig | None)
kwargs (Any)

classmethod from_config(config=None, **overrides)[source]

Build a JEPA agent from a config object, dict, YAML file, or YAML string.

Parameters:

config (JEPAConfig | dict[str, Any] | str | Path | None)
overrides (Any)

Return type:

classmethod from_pretrained(pretrained_model_name_or_path, *, config=None, checkpoint_filename=None, config_filename='config.yaml', repo_type=None, revision=None, **overrides)[source]

Create a JEPA agent from local/HF Hub config and checkpoint metadata.

Parameters:

pretrained_model_name_or_path (str | Path)
config (JEPAConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
overrides (Any)

Return type:

parameter_count(trainable_only=False)[source]

JEPA models are constructed inside training, so no parameters are resident.

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]

Return the configured JEPA run metadata.

Return type:: dict[str, Any]

train()[source]

Train for cfg.epochs epochs.

Runs on the shared Trainer through the jepa task, in cfg.folder: per-epoch checkpoints in train_jepa’s format (<write_tag>-latest.pth.tar, and <write_tag>-ep<N>.pth.tar every 50 epochs), resuming from read_checkpoint when load_checkpoint is set, bf16 with use_bfloat16, and per-epoch held-out early stopping with early_stopping. W&B sweeps (enable_sweep) still run through train_jepa.

Return type:: None

class simulusion.IRISAgent(config, action_size, device)[source]

Bases: ExportableAgentMixin, Module

Complete IRIS Agent with world model and policy.

Combines: - Discrete autoencoder (encoder + decoder) - Transformer world model - Actor-Critic for policy and value learning

Parameters:

config (IRISConfig)
action_size (int)
device (device)

use_cuda_graphs: Replay imagination’s single-position decode steps as a CUDA graph. Same computation, far fewer kernel launches; only used on CUDA.

reconstruct_chunk_size: Frames per autoencoder pass in reconstruct().

classmethod from_config(config=None, *, action_size, device=None, **overrides)[source]

Build an IRIS agent from a config object, dict, YAML file, or YAML string.

Parameters:

config (IRISConfig | dict[str, Any] | str | Path | None)
action_size (int)
device (device | str | None)
overrides (Any)

Return type:

classmethod from_pretrained(pretrained_model_name_or_path, *, action_size=None, device=None, config=None, checkpoint_filename=None, config_filename='config.yaml', repo_type=None, revision=None, **overrides)[source]

Load an IRIS agent checkpoint from a local path/directory or HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
action_size (int | None)
device (device | str | None)
config (IRISConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
overrides (Any)

Return type:

parameter_count(trainable_only=False)[source]

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]

Return type:: dict[str, Any]

set_precision(precision)[source]

Train at precision from now on (fresh loss scalers).

Parameters:: precision (MixedPrecision)
Return type:: None

forward_actor_critic(frames, hidden=None)[source]

Forward pass through actor-critic.

Parameters:

frames (Tensor) – Input frames (B, T, C, H, W)
hidden (Tuple[Tensor, Tensor] | None) – Optional LSTM hidden state

Returns:

(B, T, action_size) values: (B, T) hidden_state: (h, c)

Return type:

action_logits

act(frame, epsilon=0.0, temperature=1.0, hidden=None, return_hidden=False)[source]

Sample action from policy.

The policy is recurrent (paper A.3: CNN -> LSTM). Callers stepping an episode must thread hidden from one call to the next and reset it on episode boundaries; dropping it makes the policy effectively memoryless, which for games like Pong removes any way to infer the ball’s direction from a single frame.

Parameters:

frame (Tensor) – Single frame (B, C, H, W)
epsilon (float) – Random action probability
temperature (float) – Action distribution temperature
hidden (Tuple[Tensor, Tensor] | None) – LSTM state from the previous step, or None to start fresh
return_hidden (bool) – If True, also return the updated LSTM state

Returns:

Selected actions (B,), and the updated LSTM state when return_hidden is set.

Return type:

burn_in(frames)[source]

Initialise the LSTM state by replaying preceding frames.

Paper A.3: “Before starting the imagination procedure from a given frame, we burn-in the 20 previous frames to initialize the hidden state” (Kapturowski et al., 2019). Without this the rollout begins from a zero state that the policy never sees at collection time.

Parameters:: frames (Tensor) – Preceding observations (B, T_burn, C, H, W), already in the reconstruction domain the policy is trained on.
Returns:: The LSTM state after the burn-in, or None if no frames were given.
Return type:: Tuple[Tensor, Tensor] | None

transform_reward(rewards)[source]

Apply the configured reward transform.

Atari rewards are unbounded integers – in the thousands for games like Krull or UpNDown – and feeding them raw into the value function and lambda-return makes the critic’s target scale game-dependent. The standard Atari convention, which IRIS follows, is to take the sign.

Parameters:: rewards (Tensor)
Return type:: Tensor

reconstruct(frames)[source]

Pass frames through the discrete autoencoder: D(E(x)).

Paper A.1: “during experience collection in the real environment, frames still go through the autoencoder to keep the input distribution of the policy unchanged”. The policy only ever learns from reconstructions during imagination, so feeding it raw frames in the real environment is a distribution shift.

The encoder is forced into eval mode for the duration. This is not cosmetic: the quantizer’s dead-code revival only runs in training mode, and this method is called on single frames during experience collection and evaluation, when the agent is still in training mode. Leaving it there re-seeds most of the codebook from the 16 encoder outputs of one frame on every environment step, destroying the vocabulary the world model is being trained against.

Parameters:: frames (Tensor) – Real observations (B, C, H, W) or (B, T, C, H, W) in [0, 1].
Returns:: Reconstructions with the same shape, clamped to [0, 1].
Return type:: Tensor

imagine_rollout(initial_frame, horizon=20, burn_in_frames=None, sample_tokens=True, temperature=1.0, stop_on_termination=True)[source]

Generate imagined trajectories using the world model.

Implements the imagination MDP of paper 2.3 / Figure 1: starting from a real observation, the policy acts on decoded frames while the Transformer rolls the dynamics forward, conditioned on the full imagined history.

Parameters:

initial_frame (Tensor) – Starting frame x_0 (B, C, H, W), in [0, 1].
horizon (int) – Number of steps H to imagine.
burn_in_frames (Tensor | None) – The (B, T_burn, C, H, W) reconstructed observations preceding initial_frame, used to initialise the policy’s LSTM state (paper A.3). None starts from a zero state.
sample_tokens (bool) – Sample next-frame tokens instead of taking the argmax.
temperature (float) – Sampling temperature for token generation.
stop_on_termination (bool) – Cut the rollout short once every trajectory in the batch has hit a predicted episode end (paper 2.3). Set False to always return exactly horizon steps; correctness does not depend on it, since continues already zeroes the discount past a termination.

Returns:

dict with frames (B, T+1, C, H, W), actions (B, T), rewards (B, T) and continues (B, T). T may be shorter than horizon if every rollout predicted an episode end.

Return type:

trajectory

update_autoencoder(frames)[source]

Update discrete autoencoder.

Parameters:: frames (Tensor) – Training frames (B, C, H, W)
Returns:: Dictionary of loss values
Return type:: losses

update_transformer(frames, actions, rewards, terminals)[source]

Update transformer world model.

Parameters:

frames (Tensor) – Frame sequence
actions (Tensor) – Actions taken
rewards (Tensor) – Rewards received
terminals (Tensor) – Terminal flags

Returns:

Dictionary of loss values

Return type:

losses

update_actor_critic(imagined_trajectory)[source]

Update actor-critic in imagination.

Parameters:: imagined_trajectory (dict) – Dictionary from imagine_rollout
Returns:: Dictionary of loss values
Return type:: losses

CHECKPOINT_FORMAT = 5

save(path)[source]

Save agent state.

Parameters:: path (str)
Return type:: None

load(path)[source]

Load agent state.

Parameters:: path (str)
Return type:: None

class simulusion.IRISTransformer(vocab_size=512, tokens_per_frame=16, action_size=18, embed_dim=256, num_layers=10, num_heads=4, dropout=0.1, gradient_checkpointing=False, reward_classes=3)[source]

Bases: ExportableAgentMixin, Module

GPT-like autoregressive Transformer for world modeling.

Models the dynamics of the environment by predicting, autoregressively over an interleaved sequence of frame tokens and actions:

Next frame tokens (transition model), one token at a time
Rewards
Episode termination

The sequence layout for S frames and S - 1 actions is:

z_0^1 ... z_0^K, a_0, z_1^1 ... z_1^K, a_1, ..., z_{S-2}^1 ... z_{S-2}^K,
a_{S-2}, z_{S-1}^1 ... z_{S-1}^K

A causal (lower-triangular) attention mask is always applied, so every position only attends to itself and preceding positions. The tokens of frame t + 1 are predicted starting from the action position a_t (which sees the whole of frame t and the action), then autoregressively from each previously predicted token of frame t + 1. This matches the paper’s

z_{t+1}^k ~ p(. | z_{<=t}, a_{<=t}, z_{t+1}^{<k})

Parameters:

vocab_size (int)
tokens_per_frame (int)
action_size (int)
embed_dim (int)
num_layers (int)
num_heads (int)
dropout (float)
gradient_checkpointing (bool)
reward_classes (int)

expected_reward(action_hidden)[source]

Scalar reward prediction from an action-position hidden state.

With a categorical head this is the expectation under the predicted distribution over {-1, 0, +1} rather than an argmax, so the imagined return reflects the model’s uncertainty instead of committing to the modal class.

Parameters:: action_hidden (Tensor) – (…, embed_dim) hidden states at action positions.
Returns:: (…) scalar reward predictions.
Return type:: Tensor

gpt_blocks()[source]

The transformer blocks, typed (nn.ModuleList erases the element type).

Return type:: List[GPTBlock]

forward(tokens, actions)[source]

Teacher-forced forward pass through the Transformer world model.

Parameters:

tokens (Tensor) – Frame tokens (B, S, K) for S consecutive frames.
actions (Tensor) – Actions (B, S-1); actions[:, t] is taken after frame t.

Returns:

Predictions of frames 1..S-1 (B, S-1, K, vocab_size). rewards: Predicted rewards r_0..r_{S-2} (B, S-1). terminations: Predicted terminations d_0..d_{S-2} (B, S-1, 2).

Return type:

token_logits

init_cache(batch_size, device, dtype=None, max_len=None)[source]

Allocate an empty KV cache.

It holds max_len positions, by default this model’s full context capacity, and never more than it.

Parameters:

batch_size (int)
device (device)
dtype (dtype | None)
max_len (int | None)

Return type:

KVCache

prime_cache(tokens, actions, cache, start_pos=0)[source]

Fill a cache with an interleaved (frame, action) history in one pass.

This is the “conditioning” phase of imagination: the real frames (and the actions taken between them) are pushed through the model so that subsequent single-position steps attend to the full history, exactly as the teacher-forced training sequence does.

Parameters:

tokens (Tensor) – Frame token grids (B, Tc, K).
actions (Tensor | None) – Actions following each frame. Pass Tc actions to end the primed sequence on an action (ready to generate the next frame), or Tc - 1 / None to end on the last frame’s tokens.
cache (KVCache) – Cache to fill; must be empty or positioned at start_pos.
start_pos (int) – Absolute position of the first embedded token.

Returns:

The absolute position just past the primed sequence.

Return type:

int

generate_frame_cached(action, cache, start_pos, sample=True, temperature=1.0, decode_step=None)[source]

Generate one frame’s K tokens incrementally, one position at a time.

The cache must already hold the history up to and including the current frame’s tokens. This appends the action, reads the reward/termination hidden state from that position, then autoregressively appends the K - 1 predicted tokens – K single-position forwards in total, rather than K full-sequence passes.

Returns:

Per-token logits (B, K, vocab). generated: Token indices (B, K). action_hidden: Hidden state at the action position (B, E). next_pos: Absolute position after the generated frame’s tokens.

Return type:

step_logits

Parameters:

action (Tensor)
cache (KVCache)
start_pos (int)
sample (bool)
temperature (float)
decode_step (GraphedDecodeStep | None)

predict_next_tokens(tokens, actions)[source]

Greedily predict the next frame tokens autoregressively.

Parameters:

tokens (Tensor) – Current frame tokens (B, K) or (B, H, W).
actions (Tensor) – Actions taken (B,).

Returns:

Next frame token logits (B, K, vocab_size). Their argmax: equals the greedily generated tokens.

action_hidden: Hidden states for reward prediction (B, embed_dim).

Return type:

token_logits

imagine_step(tokens, actions, sample=True, temperature=1.0)[source]

Advance imagination one step: next frame tokens + the action hidden state.

Unlike predict_next_tokens() this can sample rather than take the argmax, and it returns the generated tokens directly so a rollout does not have to re-derive them from logits.

Parameters:

tokens (Tensor) – Current frame tokens (B, K) or (B, H, W).
actions (Tensor) – Actions taken (B,).
sample (bool) – If True, sample tokens from the predicted distribution.
temperature (float) – Sampling temperature (ignored when sample is False).

Returns:

Generated token indices (B, K). action_hidden: Hidden state at the action position (B, embed_dim),

the input expected by reward_head / termination_head.

Return type:

next_tokens

sample_next_tokens(tokens, actions, temperature=1.0)[source]

Sample next tokens autoregressively from the distribution.

Parameters:

tokens (Tensor) – Current frame tokens (B, K) or (B, H, W).
actions (Tensor) – Actions taken (B,).
temperature (float) – Sampling temperature (higher = more random).

Returns:

Sampled token indices (B, K). log_probs: Log probabilities of sampled tokens (B, K).

Return type:

sampled_tokens

class simulusion.IRISWorldModel(encoder, decoder, transformer)[source]

Bases: ExportableAgentMixin, Module

Complete IRIS World Model combining autoencoder and transformer.

This is the core component that learns environment dynamics entirely in the “imaginary” latent space.

Parameters:

encoder (Module)
decoder (Module)
transformer (IRISTransformer)

decode_tokens(indices)[source]

Decode token indices (B, K) or (B, H, W) to images.

The lookup goes through the encoder’s quantizer codebook, which is the table the reconstruction and commitment losses actually train. IRISDecoder.index_to_embedding is a separate, never-optimised embedding table, so decoding through it returns noise.

Parameters:: indices (Tensor)
Return type:: Tensor

forward(observations, actions)[source]

Full world model forward pass.

Parameters:

observations (Tensor) – Image sequence (B, T+1, C, H, W)
actions (Tensor) – Actions (B, T)

Returns:

Dictionary with predicted tokens, rewards, terminations losses: Dictionary with loss components

Return type:

predictions

imagine(initial_tokens, policy, horizon=20, temperature=1.0)[source]

Generate imagined trajectories.

Parameters:

initial_tokens (Tensor) – Initial frame tokens (B, K)
policy (Module) – Policy network to sample actions
horizon (int) – Number of steps to imagine
temperature (float) – Sampling temperature for token prediction

Returns:

Dictionary with imagined trajectories

Return type:

imagined

class simulusion.LPIPSPerceptualLoss(num_blocks=5, use_batch_norm=False, linear_weights=None)[source]

Bases: ExportableAgentMixin, Module

LPIPS-structured perceptual distance over VGG16 features.

Pipeline, following Zhang et al. (2018) as used by VQGAN:

Map inputs from [0, 1] to [-1, 1] and apply LPIPS’s fixed per-channel shift/scale (this is not ImageNet mean/std – LPIPS uses its own calibration).
Extract features at the end of each VGG16 conv block.
Unit-normalise each feature vector across channels, so no single block’s activation magnitude dominates the sum.
Square the difference, weight per channel with a 1x1 linear layer, then average spatially and sum over blocks.

Parameters:

num_blocks (int) – How many VGG16 conv blocks to compare (1-5). LPIPS uses 5.
use_batch_norm (bool) – Load vgg16_bn instead of vgg16.
linear_weights (str | Path | None) – Path to LPIPS learned linear weights (the lin*.model /vgg.pth state dict from the reference implementation). When None or unloadable, uniform weights are used and has_learned_weights stays False.

Variables:

has_learned_weights – Whether the calibrated LPIPS linear weights were loaded. False means this is the unweighted VGG-feature variant.

train(mode=True)[source]

Keep the frozen backbone in eval mode regardless of the parent’s mode.

Parameters:: mode (bool)
Return type:: LPIPSPerceptualLoss

forward(x, y)[source]

Perceptual distance between two batches of images in [0, 1].

Parameters:

x (Tensor) – Images (B, 3, H, W), typically the ground-truth frames.
y (Tensor) – Images (B, 3, H, W), typically the reconstructions.

Returns:

Scalar loss, averaged over the batch and summed over blocks.

Return type:

Tensor

simulusion.build_perceptual_loss(enabled=True, num_blocks=5, linear_weights=None, download_weights=True)[source]

Build the perceptual loss, returning None if it cannot be constructed.

Loading VGG16 requires torchvision and, on first use, a weight download. When either is unavailable this returns None and logs a warning rather than failing training – the autoencoder then falls back to the L1 + commitment objective, which trains but reconstructs small sprites poorly.

vgg16 is tried first (the variant LPIPS is calibrated on), then vgg16_bn as a fallback since it is more often cached locally.

When linear_weights is not given, the reference LPIPS weights are looked up in the torch hub cache and fetched if missing, so the loss is calibrated LPIPS by default. Set download_weights=False to stay strictly offline.

Parameters:

enabled (bool)
num_blocks (int)
linear_weights (str | Path | None)
download_weights (bool)

Return type:

Module | None

simulusion.compute_lambda_return(rewards, values, discounts, lambda_coef=0.95)[source]

Compute λ-return target for value function training.

Parameters:

rewards (Tensor) – Rewards (B, T)
values (Tensor) – Value estimates (B, T+1)
discounts (Tensor) – Discount factors (B, T)
lambda_coef (float) – Lambda parameter for bootstrapping

Returns:

λ-return targets (B, T)

Return type:

lambda_returns

class simulusion.VisionTransformer(img_size=[224], patch_size=16, in_chans=3, embed_dim=768, predictor_embed_dim=384, depth=12, predictor_depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>, init_std=0.02, **kwargs)[source]

Bases: ExportableAgentMixin, Module

Vision Transformer

Parameters:

img_size (list[int])
patch_size (int)
in_chans (int)
embed_dim (int)
predictor_embed_dim (int)
depth (int)
predictor_depth (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
qk_scale (float | None)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)
norm_layer (type[Module])
init_std (float)
kwargs (Any)

fix_init_weight()[source]

Return type:: None

forward(x, masks=None)[source]

Parameters:

x (Tensor)
masks (Tensor | list[Tensor] | None)

Return type:

Tensor

get_intermediate_layers(x, n=1)[source]

Return the patch tokens output by each of the last n blocks.

I-JEPA’s linear-evaluation protocol (Appendix A.2) reports the better of the average-pooled last layer and the concatenation of the average-pooled last four layers, which this method makes available. Every returned tensor has the final norm applied.

Parameters:

x (Tensor)
n (int)

Return type:

list[Tensor]

interpolate_pos_encoding(x, pos_embed)[source]

Parameters:

x (Tensor)
pos_embed (Tensor)

Return type:

Tensor

class simulusion.ModularRSSM(encoder, decoder, backbone, reward_decoder=None)[source]

Bases: ExportableAgentMixin, Module

Modular RSSM with swappable encoder, decoder, and backbone.

This class allows researchers to easily experiment with different: - Encoders: Conv, MLP, ViT - Decoders: Conv, MLP - Backbones: GRU, LSTM, Transformer

Example

>>> encoder = ConvEncoder((3, 64, 64), embed_size=1024)
>>> decoder = ConvDecoder(32, 200, (3, 64, 64))
>>> backbone = GRUBackbone(action_size=6, stoch_size=32, deter_size=200, hidden_size=200, embed_size=1024)
>>> rssm = ModularRSSM(encoder, decoder, backbone)

Parameters:

encoder (EncoderBase)
decoder (DecoderBase)
backbone (BackboneBase)
reward_decoder (DecoderBase | None)

property stoch_size: int

property deter_size: int

property embed_size: int

init_state(batch_size, device)[source]

Parameters:

batch_size (int)
device (device)

Return type:

Dict[str, Tensor]

get_dist(mean, std)[source]

Parameters:

mean (Tensor)
std (Tensor)

Return type:

Distribution

observe_step(prev_state, prev_action, obs, nonterm=1.0)[source]

Encode obs and advance one step.

Returns:

the core Dreamer simulusion.models.dreamer_rssm.RSSM.observe_step() returns (posterior, prior). Both states have identical keys, so mixing them up raises no error; the posterior is the one that has seen obs and is what should be acted on and carried forward.

Return type:

(prior, posterior). Note the order

Parameters:

prev_state (Dict[str, Tensor])
prev_action (Tensor)
obs (Tensor)
nonterm (Any)

imagine_step(prev_state, prev_action, nonterm=1.0)[source]

Parameters:

prev_state (Dict[str, Tensor])
prev_action (Tensor)
nonterm (Any)

Return type:

Dict[str, Tensor]

observe_rollout(obs, actions, nonterms, prev_state, horizon)[source]

Parameters:

obs (Tensor)
actions (Tensor)
nonterms (Tensor)
prev_state (Dict[str, Tensor])
horizon (int)

Return type:

Tuple[Dict[str, Tensor], Dict[str, Tensor]]

imagine_rollout(actor, prev_state, horizon)[source]

Parameters:

actor (Module)
prev_state (Dict[str, Tensor])
horizon (int)

Return type:

Dict[str, Tensor]

decode_observation(features)[source]

Parameters:: features (Tensor)
Return type:: Tensor

decode_reward(features)[source]

Parameters:: features (Tensor)
Return type:: Tensor

detach_state(state)[source]

Parameters:: state (Dict[str, Tensor])
Return type:: Dict[str, Tensor]

seq_to_batch(state)[source]

Parameters:: state (Dict[str, Tensor])
Return type:: Dict[str, Tensor]

simulusion.create_modular_rssm(encoder_type='conv', decoder_type='conv', backbone_type='gru', obs_shape=(3, 64, 64), action_size=6, stoch_size=32, deter_size=200, embed_size=1024, hidden_size=200, activation='elu', **kwargs)[source]

Factory function to create a modular RSSM with specified components.

Parameters:

encoder_type (str) – Type of encoder (“conv”, “mlp”, “vit”)
decoder_type (str) – Type of decoder (“conv”, “mlp”)
backbone_type (str) – Type of backbone (“gru”, “lstm”, “transformer”)
obs_shape (Tuple[int, int, int] | Tuple[int]) – Shape of observations (C, H, W) for images or (D,) for state
action_size (int) – Action space dimension
stoch_size (int) – Stochastic latent dimension
deter_size (int) – Deterministic hidden dimension
embed_size (int) – Encoder embedding dimension
hidden_size (int) – Hidden layer dimension
activation (str) – Activation function name
kwargs (Any)

Returns:

Configured ModularRSSM instance

Return type:

ModularRSSM

class simulusion.Genie(num_frames=16, image_size=64, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, tokenizer_encoder_dim=512, tokenizer_decoder_dim=1024, action_vocab_size=8, action_embedding_dim=32, action_encoder_dim=1024, action_decoder_dim=1024, dynamics_dim=5120, dynamics_depth=48, dynamics_num_heads=40, encoder_depth=12, decoder_depth=20, latent_action_depth=20, use_bfloat16=False, action_pooling='mean', window_attention_heads=1, tokenizer_num_heads=16, action_num_heads=16)[source]

Bases: ExportableAgentMixin, Module

Genie: Generative Interactive Environment.

A generative model trained from video-only data that can be used as an interactive environment. Contains three key components: 1. Video Tokenizer: Converts raw video frames into discrete tokens 2. Latent Action Model (LAM): Infers latent actions between frames 3. Dynamics Model: Predicts future frames given past frames and latent actions

Based on “Genie: Generative Interactive Environments” paper (arXiv:2402.15391).

Training follows two phases as per paper: 1. Train video tokenizer first (on video tokens) 2. Co-train LAM (from pixels) and dynamics model (on video tokens)

The LAM uses VQ-VAE training with: - Encoder: Takes x1:t and x_{t+1} → outputs latent actions - Decoder: Takes x1:t-1 (masked) + actions → reconstructs x_t - Auxiliary variance loss to prevent action collapse

At inference, latent actions are stopgrad’d when passed to dynamics model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
tokenizer_encoder_dim (int)
tokenizer_decoder_dim (int)
action_vocab_size (int)
action_embedding_dim (int)
action_encoder_dim (int)
action_decoder_dim (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
encoder_depth (int)
decoder_depth (int)
latent_action_depth (int)
use_bfloat16 (bool)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)
tokenizer_num_heads (int)
action_num_heads (int)

classmethod from_config(config=None, **overrides)[source]

Build Genie from a config object, dict, YAML file, or YAML string.

Parameters:

config (GenieConfig | GenieSmallConfig | dict[str, Any] | str | Path | None)
overrides (Any)

Return type:

Load Genie weights from a local path/directory or HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
config (GenieConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
map_location (str | device | None)
overrides (Any)

Return type:

CHECKPOINT_FORMAT = 2

save_pretrained(path)[source]

Save Genie weights and config in a from_pretrained-compatible format.

Parameters:: path (str | Path)
Return type:: None

parameter_count(trainable_only=False)[source]

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]

Return type:: dict[str, Any]

forward(video, mask_prob=None, training_phase='all')[source]

Full forward pass through all components.

Parameters:

video (Tensor) – (B, C, T, H, W) input video
mask_prob (float | None) – Dynamics masking rate. None (default) lets the dynamics model draw it from U(0.5, 1) per step, per paper 2.1.
training_phase (str) – “all”, “tokenizer”, or “lam_dynamics”

Returns:

Dictionary containing losses and predictions

Return type:

Dict[str, Tensor]

training_step(video, mask_prob=None, training_phase='all')[source]

Single training step computing all losses.

Parameters:

video (Tensor) – (B, C, T, H, W) input video
mask_prob (float | None) – Dynamics masking rate; None draws it per step.
training_phase (str) – “all”, “tokenizer”, or “lam_dynamics”

Returns:

Dictionary containing all losses for backpropagation

Return type:

Dict[str, Tensor]

encode_video(video)[source]

Encode video to discrete tokens.

Parameters:: video (Tensor) – (B, C, T, H, W)
Returns:: (B, T, H*W)
Return type:: video_tokens

infer_actions(frames)[source]

Infer latent actions from a sequence of frames.

Parameters:: frames (Tensor) – (B, C, T, H, W) video frames
Returns:: (B, T-1) inferred latent action indices
Return type:: latent_actions

generate(prompt_frame, num_frames=16, actions=None, use_maskgit=True, use_cache=False)[source]

Generate video frames given a prompt frame and actions.

Parameters:

prompt_frame (Tensor) – (B, C, H, W) initial frame
num_frames (int) – Total number of frames to generate
actions (Tensor | None) – (B, num_frames-1) latent action indices, or None for random
use_maskgit (bool) – Whether to use MaskGIT sampling
use_cache (bool) – Generate each frame from a temporal KV cache instead of re-running the dynamics model over the whole prefix; O(T) rather than O(T^2) in the number of frames. Same logits up to float rounding.

Returns:

(B, C, num_frames, H, W)

Return type:

generated_video

play(current_frame, action, current_frames=None)[source]

Play step - generate next frame given current frame and action.

Parameters:

current_frame (Tensor) – (B, C, H, W) current frame
action (Tensor) – (B,) latent action indices
current_frames (Tensor | None) – (B, C, T, H, W) history frames, or None for first frame

Returns:

(B, C, H, W)

Return type:

next_frame

get_num_parameters()[source]

Return total number of parameters.

Return type:: int

class simulusion.LatentActionModel(num_frames=16, image_size=64, in_channels=3, encoder_dim=256, decoder_dim=512, encoder_depth=4, decoder_depth=4, num_heads=8, patch_size=16, vocab_size=8, embedding_dim=32, commitment_weight=1.0, action_pooling='mean', window_attention_heads=1)[source]

Bases: ExportableAgentMixin, Module

Latent Action Model (LAM) for unsupervised action learning.

Learns discrete latent actions from unlabeled video frames using a VQ-VAE based objective. The model infers latent actions between frames that encode the most meaningful changes for future frame prediction.

Based on Genie paper - learns actions without action labels from Internet videos.

Components: - Encoder: Takes all previous frames x1:t and next frame x_t+1 → outputs latent actions - Decoder: Takes previous frames x1:t-1 and latent actions a1:t-1 → predicts next frame x_t

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
commitment_weight (float)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

encode(x_prev, x_next)[source]

Encode frames to latent actions.

Parameters:

x_prev (Tensor) – Previous frames (B, C, T, H, W)
x_next (Tensor) – Next frame (B, C, H, W)

Returns:

Discrete latent action indices (B, T) z_q: Quantized embeddings (B, T, embedding_dim)

Return type:

latent_actions

decode(x_prev, z_q)[source]

Decode latent actions to predict next frame.

Parameters:

x_prev (Tensor) – Previous frames (B, C, T, H, W) - will mask all but first
z_q (Tensor) – Quantized action embeddings (B, T-1, embedding_dim)

Returns:

(B, C, H, W)

Return type:

predicted_next_frame

forward(x_prev, x_next)[source]

Full forward pass: encode to get actions, decode to reconstruct.

Parameters:

x_prev (Tensor) – Previous frames (B, C, T, H, W)
x_next (Tensor) – Next frame (B, C, H, W)

Returns:

Dictionary with losses and outputs

Return type:

Dict[str, Tensor]

class simulusion.DynamicsModel(num_frames=16, image_size=64, vocab_size=1024, embedding_dim=32, action_vocab_size=8, dim=5120, depth=48, num_heads=40, patch_size=4, mlp_ratio=4.0, qkv_bias=True, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0, gradient_checkpointing=True, mask_prob_min=0.5, mask_prob_max=1.0)[source]

Bases: ExportableAgentMixin, Module

Dynamics Model for action-controllable video generation.

A decoder-only transformer that predicts future frame tokens given past frame tokens and latent actions. Uses MaskGIT for training and sampling.

Based on Genie paper - uses cross-entropy loss with random masking during training, and MaskGIT iterative refinement at inference.

Parameters:

num_frames (int)
image_size (int)
vocab_size (int)
embedding_dim (int)
action_vocab_size (int)
dim (int)
depth (int)
num_heads (int)
patch_size (int)
mlp_ratio (float)
qkv_bias (bool)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)
gradient_checkpointing (bool)
mask_prob_min (float)
mask_prob_max (float)

forward(video_tokens, actions, mask_prob=None)[source]

Forward pass for training.

Parameters:

video_tokens (Tensor) – (B, T, H*W) - token indices for frames 1 to T
actions (Tensor) – (B, T) - latent action indices for frames 1 to T
mask_prob (float | None) – Masking rate. None (the default) draws it from U(mask_prob_min, mask_prob_max) per call, as the paper specifies. Pass 0.0 to disable masking, e.g. at inference.

Returns:

(B, T, H*W, vocab_size)

Return type:

logits

init_cache(batch_size, device=None, dtype=None)[source]

Allocate a temporal KV cache for forward_cached().

Parameters:

batch_size (int)
device (device | None)
dtype (dtype | None)

Return type:

forward_cached(video_tokens, actions, cache, commit=True)[source]

Inference forward over only the frames after those in cache.

Equivalent to the matching frames of forward(all_tokens, all_actions, mask_prob=0.0) in eval mode, without recomputing the cached prefix.

Parameters:

video_tokens (Tensor) – (B, T_new, N) tokens of the new frames. T_new may be larger than 1 only while the cache is empty (the prompt).
actions (Tensor) – (B, T_new) latent actions for those frames.
cache (STKVCache) – From init_cache(); its length is the index of the first new frame.
commit (bool) – Advance the cache past these frames. False evaluates a candidate frame without keeping it.

Returns:

(B, T_new, N, vocab_size)

Return type:

logits

sample(prompt_tokens, prompt_actions, num_frames, sampler=None)[source]

Sample future frames using MaskGIT.

Parameters:

prompt_tokens (Tensor) – (B, T_prompt, N) - starting frame tokens
prompt_actions (Tensor) – (B, T_prompt) - actions for prompt frames
num_frames (int) – Total number of frames to generate
sampler (MaskGITSampler | None) – MaskGIT sampler instance

Returns:

(B, num_frames, N)

Return type:

generated_tokens

autoregressive_sample(prompt_tokens, actions, num_frames, temperature=1.0, use_cache=False)[source]

Simple autoregressive sampling (frame by frame).

Parameters:

prompt_tokens (Tensor) – (B, T_prompt, N) - starting frame tokens
actions (Tensor) – (B, num_frames - 1) - latent actions for each transition. The action driving frame t -> t+1 is actions[:, t-1]. If fewer actions are supplied than transitions, the remainder are sampled at random.
num_frames (int) – Total number of frames to generate
temperature (float) – Sampling temperature
use_cache (bool) – Reuse the temporal keys/values of already generated frames (STKVCache) instead of re-running the transformer over the whole prefix for every frame. The logits match the uncached path up to float rounding. One difference: when actions must be padded at random, the uncached path redraws the padding for every earlier frame at every step, while the cached path draws each frame’s padding once.

Returns:

(B, num_frames, N)

Return type:

generated_tokens

simulusion.create_genie(num_frames=16, image_size=64, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, action_vocab_size=8, action_embedding_dim=32, dynamics_dim=5120, dynamics_depth=48, dynamics_num_heads=40, use_bfloat16=False, action_pooling='mean', window_attention_heads=1)[source]

Factory function to create a Genie model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
action_vocab_size (int)
action_embedding_dim (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
use_bfloat16 (bool)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

Return type:

simulusion.create_genie_small(num_frames=16, image_size=64, use_bfloat16=False, action_pooling='mean', window_attention_heads=1)[source]

Create a smaller Genie model for development/testing.

Parameters:

num_frames (int)
image_size (int)
use_bfloat16 (bool)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

Return type:

simulusion.create_genie_large(num_frames=16, image_size=64, use_bfloat16=True, action_pooling='mean', window_attention_heads=1)[source]

Create the full 11B parameter Genie model (approximate).

Parameters:

num_frames (int)
image_size (int)
use_bfloat16 (bool)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

Return type:

simulusion.create_latent_action_model(num_frames=16, image_size=64, in_channels=3, encoder_dim=256, decoder_dim=512, encoder_depth=4, decoder_depth=4, num_heads=8, patch_size=16, vocab_size=8, embedding_dim=32, action_pooling='mean', window_attention_heads=1)[source]

Factory function to create a Latent Action Model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

Return type:

LatentActionModel

simulusion.create_dynamics_model(num_frames=16, image_size=64, vocab_size=1024, embedding_dim=32, action_vocab_size=8, dim=5120, depth=48, num_heads=40, patch_size=4)[source]

Factory function to create a Dynamics Model.

Parameters:

num_frames (int)
image_size (int)
vocab_size (int)
embedding_dim (int)
action_vocab_size (int)
dim (int)
depth (int)
num_heads (int)
patch_size (int)

Return type:

DynamicsModel

class simulusion.RSSM(action_size, stoch_size, deter_size, hidden_size, obs_embed_size, activation)[source]

Bases: ExportableAgentMixin, Module

Recurrent State-Space Model used by Dreamer for latent dynamics learning.

The RSSM is the core world model component that learns compact representations of environment dynamics. It maintains a hybrid state consisting of:

Deterministic State (h) – A recurrent hidden state updated by a GRU, capturing sequential/temporal information and deterministic transitions.
Stochastic State (s) – A latent variable representing stochastic, multi-modal uncertainty in the environment (e.g., ambiguous observations).

The model operates in two modes:

Observe Mode – Updates states using actual observations from the environment. Uses the representation model: p(s_t | h_t, obs_t)
Imagine Mode – Predicts future states without observations. Uses the transition/prior model: p(s_t | h_t)

Architecture

Input: Previous state (h_{t-1}, s_{t-1}) and action a_{t-1}
Process: GRU updates deterministic state, MLP computes stochastic prior/posterior
Output: Updated state (h_t, s_t) and distributions

State Representation

deter (h): GRU hidden state, captures sequential context
stoch (s): Stochastic latent, multi-modal uncertainty
mean/std: Parameters of the stochastic distribution

Usage with DreamerAgent:

rssm = RSSM(
    action_size=action_dim,
    stoch_size=30,      # Stochastic state dimension
    deter_size=200,     # Deterministic (GRU) state dimension
    hidden_size=200,    # MLP hidden layer size
    obs_embed_size=256,  # Observation embedding from encoder
    activation='elu'
)

# Observe with actual observation
posterior, prior = rssm.observe_step(prev_state, prev_action, obs_embed)

# Imagine future without observation
prior = rssm.imagine_step(current_state, action)

Training

The RSSM is trained by maximizing the ELBO (Evidence Lower Bound):

KL divergence between prior and posterior encourages the prior to capture environment dynamics
Reconstruction loss from decoder ensures state captures observation info

Reference:: Dreamer: Scalable Reinforcement Learning Using World Models Hafner et al., 2020 - https://arxiv.org/abs/1912.01603

Parameters:

action_size (int)
stoch_size (int)
deter_size (int)
hidden_size (int)
obs_embed_size (int)
activation (str)

init_state(batch_size, device)[source]

Initialize RSSM state with zeros.

Parameters:

batch_size (int) – Number of parallel sequences
device (device) – torch device for tensors

Returns:

mean, std: Stochastic distribution parameters
stoch: Stochastic state sample
deter: Deterministic GRU hidden state

Return type:

Dictionary containing zero-initialized state components

get_dist(mean, std)[source]

Create an Independent Normal distribution from mean and std.

Parameters:

mean (Tensor) – Location parameter
std (Tensor) – Scale parameter

Returns:

Independent Normal distribution with given parameters

Return type:

Independent

observe_step(prev_state, prev_action, obs_embed, nonterm=tensor(1.), noise=None)[source]

Update state using actual observation (observe mode).

In observe mode, the RSSM first computes a transition prior from the previous state and action, then refines the stochastic state using the actual observation embedding to form the posterior.

Parameters:

prev_state (dict) – Dictionary with ‘deter’ (h_{t-1}) and ‘stoch’ (s_{t-1})
prev_action (Tensor) – Previous action a_{t-1}, shape (B, action_size)
obs_embed (Tensor) – Observation embedding from encoder, shape (B, obs_embed_size)
nonterm (Tensor) – Termination mask (1.0 = continue, 0.0 = terminal)
noise (Tuple[Tensor, Tensor] | None) – Optional (prior_noise, posterior_noise) standard-normal tensors of shape (B, stoch_size) used for the two samples instead of drawing them internally. Passing noise makes the step a pure function of its inputs, which exported graphs and eager-vs-exported parity checks need.

Returns:

A tuple (posterior, prior) of state dictionaries. The posterior incorporates observation information; the prior is the transition prediction before observation. Both share the same deterministic state because the GRU is only advanced once per timestep.

The order differs from observe_rollout() and from ModularRSSM.observe_step, which both return (prior, posterior). The dictionaries have identical keys, so a swapped unpacking fails silently: acting on the prior discards the current observation. It is kept for backward compatibility.

Return type:

Tuple[dict, dict]

imagine_step(prev_state, prev_action, nonterm=tensor(1.), noise=None)[source]

Predict next state without observation (imagine mode).

In imagine mode, the RSSM predicts future states using only the prior distribution. This is used for planning and policy learning where actual observations are not available.

Parameters:

prev_state (dict) – Dictionary with ‘deter’ (h_{t-1}) and ‘stoch’ (s_{t-1})
prev_action (Tensor) – Previous action a_{t-1}, shape (B, action_size)
nonterm (Tensor) – Termination mask (1.0 = continue, 0.0 = terminal)
noise (Tensor | None) – Optional standard-normal tensor (B, stoch_size) used for the stochastic sample instead of drawing one internally.

Returns:

deter: Predicted deterministic state
mean, std, stoch: Prior stochastic state distribution

Return type:

Dictionary with predicted state containing

get_prior(prev_state, prev_action, nonterm=tensor(1.))[source]

Parameters:

prev_state (dict)
prev_action (Tensor)
nonterm (Tensor)

Return type:

dict

get_posterior(prev_state, prev_action, obs_embed, nonterm=tensor(1.))[source]

Compute posterior distribution over stochastic state.

The posterior incorporates observation information to produce a more accurate state estimate.

Parameters:

prev_state (dict) – Previous state dictionary
prev_action (Tensor) – Previous action
obs_embed (Tensor) – Observation embedding
nonterm (Tensor) – Termination mask

Returns:

Dictionary with posterior state (observation-informed). Note that the previous-state shape (B, ...) is preserved; the batch dimension is not flattened.

Return type:

dict

detach_state(state)[source]

Detach state tensors from computation graph.

Used during DreamerV2 training to prevent gradient flow through the observation/update pathway.

Parameters:: state (dict) – State dictionary with tensor values
Returns:: Detached state dictionary
Return type:: dict

seq_to_batch(state_dict)[source]

Convert sequence state to batch format.

Parameters:: state_dict (dict) – Dictionary with sequence-dimension tensors (T, B, …)
Returns:: Dictionary with batch-dimension tensors (B*T, …)
Return type:: dict

observe_rollout(obs_embed, actions, nonterms, init_state, seq_len)[source]

Process a sequence of observations (observe mode rollout).

At each timestep we run observe_step once to obtain the transition prior (the prediction given the previous state and action) and the observation-informed posterior. The posterior is then used as the previous state for the next step, matching the standard Dreamer inference pattern.

Parameters:

obs_embed (Tensor) – Observation embeddings, shape (T+1, B, obs_embed_size)
actions (Tensor) – Actions, shape (T, B, action_size)
nonterms (Tensor) – Non-termination flags, shape (T, B, 1)
init_state (dict) – Initial state dictionary
seq_len (int) – Sequence length T

Returns:

Dictionary with prior states stacked along the time axis posterior: Dictionary with posterior states stacked along the time axis

Return type:

prior

imagine_rollout(policy, init_state, horizon)[source]

Generate imagined trajectory using policy (imagine mode rollout).

Parameters:

policy (Module) – Actor network that outputs actions from state features
init_state (dict) – Initial state dictionary
horizon (int) – Number of steps to imagine

Returns:

Dictionary with imagined states for each step

Return type:

dict

forward(x, u)[source]

Forward pass for training (computes sequence of states).

Parameters:

x (Tensor) – Observations, shape (B, T+1, C, H, W)
u (Tensor) – Actions, shape (B, T, action_size)

Returns:

List of state dictionaries for each timestep priors: List of prior distributions (tuples of mean, std) posteriors: List of posterior distributions (tuples of mean, std)

Return type:

states

class simulusion.RecurrentStateSpaceModel(action_size, state_size=200, latent_size=30, hidden_size=200, embed_size=1024, activation_function='relu')[source]

Bases: ExportableAgentMixin, Module

A Recurrent State Space Model (RSSM) for modeling latent dynamics in sequential data.

Parameters:

action_size (int)
state_size (int)
latent_size (int)
hidden_size (int)
embed_size (int)
activation_function (str)

get_init_state(enc, h_t=None, s_t=None, a_t=None, mean=None, *, sample=None)[source]

Returns the next deterministic state and the posterior latent.

By default the latent is the posterior mean. Pass sample=True to draw it from the posterior instead. mean is the old, inverted name for that flag (mean=True used to sample); it is still accepted with its old meaning so existing callers keep their behaviour.

Parameters:

enc (Tensor)
h_t (Tensor | None)
s_t (Tensor | None)
a_t (Tensor | None)
mean (bool | None)
sample (bool | None)

Return type:

tuple[Tensor, Tensor]

deterministic_state_fwd(h_t, s_t, a_t)[source]

Deterministic transition update.

Ensures a_t is 2D and matches batch dimension of h_t before concatenation. Accepts a_t shaped [B, action_size], [action_size] (expanded to [B, action_size]), or [B]/scalar (reshaped appropriately).

Parameters:

h_t (Tensor)
s_t (Tensor)
a_t (Tensor)

Return type:

Tensor

state_prior(h_t, sample=False)[source]

Returns the prior distribution over the latent state given the deterministic state

Parameters:

h_t (Tensor)
sample (bool)

Return type:

tuple[Tensor, Tensor] | Tensor

state_posterior(h_t, e_t, sample=False)[source]

Returns the state prior given the deterministic state and obs

Parameters:

h_t (Tensor)
e_t (Tensor)
sample (bool)

Return type:

tuple[Tensor, Tensor] | Tensor

pred_reward(h_t, s_t)[source]

Parameters:

h_t (Tensor)
s_t (Tensor)

Return type:

Tensor

rollout_prior(act, h_t, s_t)[source]

Parameters:

act (Tensor)
h_t (Tensor)
s_t (Tensor)

Return type:

tuple[Tensor, Tensor]

forward(x, u)[source]

Forward through the RSSM for a batch of sequences.

Parameters:

x (Tensor) – Tensor [B, T+1, C, H, W] (observations including initial frame)
u (Tensor) – Tensor [B, T, action_size] (actions for T steps)

Returns:

list[T] of tensors [B, state_size] priors: list[T] of tuples (mean, std) each [B, latent_size] posteriors: list[T] of tuples (mean, std) each [B, latent_size]

Return type:

states

simulusion.DreamerRSSM: alias of RSSM

class simulusion.ConvEncoder(input_shape, embed_size, activation, depth=32)[source]

Bases: ExportableAgentMixin, Module

Convolutional observation encoder used by Dreamer world models.

This encoder transforms raw image observations (typically RGB frames from environments like Atari or DeepMind Control) into compact latent embeddings that can be processed by the RSSM (Recurrent State-Space Model).

Input: (B, C, H, W) raw images, values in [-0.5, 0.5]
Process: 4 convolutional layers with stride 2, halving spatial dimensions
Output: (B, embed_size) compact representation

The encoder uses a depth doubling pattern: 32 -> 64 -> 128 -> 256 channels. After convolutions, a fully connected layer projects from 1024 features to the desired embedding size.

Usage with Dreamer:

encoder = ConvEncoder(
    input_shape=(3, 64, 64),  # RGB 64x64 images
    embed_size=256,           # RSSM observation embedding size
    activation='relu'         # Activation function
)
obs_embedding = encoder(observation)  # (B, 256)

Parameters:

input_shape (tuple) – Tuple (C, H, W) for input images, typically (3, 64, 64)
embed_size (int) – Output embedding dimension, typically 256 or 1024
activation (str) – Activation function name (‘relu’, ‘elu’, ‘tanh’, etc.)
depth (int) – Base channel depth for first layer (default 32)

forward(inputs)[source]

Parameters:: inputs (Tensor)
Return type:: Tensor

class simulusion.CNNEncoder(embedding_size, activation_function='relu')[source]

Bases: ExportableAgentMixin, Module

A Convolutional Neural Network (CNN) encoder for processing image inputs.

Parameters:

embedding_size (int)
activation_function (str)

forward(observation)[source]

Parameters:: observation (Tensor)
Return type:: Tensor

class simulusion.ConvDecoder(stoch_size, deter_size, output_shape, activation, depth=32)[source]

Bases: ExportableAgentMixin, Module

Convolutional decoder for reconstructing observations from latent states.

Part of Dreamer’s world model, this decoder reconstructs image observations from the combined stochastic (s) and deterministic (h) RSSM states.

Input: Concatenated [stoch_state, deter_state], shape (B, stoch+deter)
Process: Dense projection + 4 transposed convolutions (upsampling 2x each)
Output: Independent Normal distribution over observation pixels

The decoder mirrors the ConvEncoder’s structure but in reverse (transposed convs instead of regular convs). This creates a symmetric autoencoder where the encoder and decoder can be trained jointly to learn compressed representations.

Returns torch.distributions.Independent(Normal(mean, std), len(shape)) allowing log_prob(observation) computation for reconstruction loss.

Usage in Dreamer world model:

decoder = ConvDecoder(
    stoch_size=30,
    deter_size=200,
    output_shape=(3, 64, 64),  # RGB images
    activation='relu'
)
obs_dist = decoder(latent_features)  # Returns distribution
log_prob = obs_dist.log_prob(target_observation)

The reconstruction loss is -log_prob(observation), which encourages the RSSM to learn states that capture observation information.

Parameters:

stoch_size (int)
deter_size (int)
output_shape (tuple[int, ...])
activation (str)
depth (int)

forward(features)[source]

Parameters:: features (Tensor)
Return type:: Independent

class simulusion.CNNDecoder(state_size, latent_size, embedding_size, activation_function='relu')[source]

Bases: ExportableAgentMixin, Module

A Convolutional Neural Network (CNN) decoder for reconstructing image outputs.

Parameters:

state_size (int)
latent_size (int)
embedding_size (int)
activation_function (str)

forward(latent, state)[source]

Parameters:

latent (Tensor)
state (Tensor)

Return type:

Tensor

class simulusion.DenseDecoder(stoch_size, deter_size, output_shape, n_layers, units, activation, dist, num_buckets=255, symlog_range=10.0)[source]

Bases: ExportableAgentMixin, Module

MLP decoder for reward/value/discount prediction from latent features.

Part of Dreamer’s world model, this decoder predicts scalar quantities (rewards, values, discount factors) from RSSM latent states.

Input: [stoch_state, deter_state] concatenated, shape (B, stoch+deter)
Process: MLP with configurable layers and hidden units
Output: Predicted quantity with distribution (normal, binary, or raw)

Supports three output types: - 'normal': Gaussian distribution for regression (rewards, values) - 'binary': Bernoulli distribution for binary classification (discount) - 'none': Raw tensor for non-probabilistic outputs

Usage:

reward_decoder = DenseDecoder(
    stoch_size=30,
    deter_size=200,
    output_shape=(1,),
    n_layers=2,
    units=400,
    activation='elu',
    dist='normal'
)
reward_dist = reward_decoder(latent_features)
reward_loss = -reward_dist.log_prob(target_reward)

For discount prediction (binary):

discount_decoder = DenseDecoder(
    stoch_size=30,
    deter_size=200,
    output_shape=(1,),
    n_layers=2,
    units=400,
    activation='elu',
    dist='binary'
)

Parameters:

stoch_size (int)
deter_size (int)
output_shape (tuple[int, ...])
n_layers (int)
units (int)
activation (str)
dist (str)
num_buckets (int)
symlog_range (float)

forward(features)[source]

Parameters:: features (Tensor)
Return type:: Any

class simulusion.ActionDecoder(action_size, stoch_size, deter_size, n_layers, units, activation, min_std=0.0001, init_std=5, mean_scale=5)[source]

Bases: ExportableAgentMixin, Module

Dreamer actor head producing squashed continuous actions from latent features.

Outputs a transformed Gaussian policy with optional deterministic mode and utility for additive exploration noise.

Parameters:

action_size (int)
stoch_size (int)
deter_size (int)
n_layers (int)
units (int)
activation (str)
min_std (float)
init_std (float)
mean_scale (float)

forward(features, deter=False)[source]

Parameters:

features (Tensor)
deter (bool)

Return type:

Tensor

mean_action(features)[source]

Deterministic action tanh(mean) of the squashed Gaussian policy.

forward(deter=True) estimates the mode by drawing 100 samples and keeping the most likely, so it is random and ~100x the cost of one sample. This is the usual deployment choice instead: one pass, no randomness, and traceable by every exporter. It is not bit-identical to the Monte Carlo mode, so evaluation numbers can differ slightly.

Parameters:: features (Tensor)
Return type:: Tensor

add_exploration(action, action_noise=0.3)[source]

Parameters:

action (Tensor)
action_noise (float)

Return type:

Tensor

class simulusion.TanhBijector[source]

Bases: Transform

Bijective tanh transform for squashing Gaussian distributions to [-1, 1].

This transformation is essential for Dreamer’s action policy. Raw neural network outputs are Gaussian distributions over R^n, but actions in continuous control environments are typically bounded in [-1, 1]. The tanh bijector provides:

Bijective mapping: tanh is invertible (with atanh as inverse)
Stable log-det Jacobian: Computable for gradient-based training
Clipped actions: During inference, actions are naturally bounded

Forward: y = tanh(x)
Inverse: x = atanh(y) = 0.5 * log((1+y)/(1-y))
Log-det: log|dy/dx| = 2*(log(2) - x - softplus(-2x))

Usage with Dreamer ActionDecoder:

dist = TransformedDistribution(
    Normal(mean, std),
    TanhBijector()
)
action = dist.sample()  # Bounded to [-1, 1]

Reference:: Building a Scalable Deep RL Library by Learning from Mistakes, Haarnoja et al.

property sign: int

atanh(x)[source]

Parameters:: x (Tensor)
Return type:: Tensor

log_abs_det_jacobian(x, y)[source]

Parameters:

x (Tensor)
y (Tensor)

Return type:

Tensor

class simulusion.SampleDist(dist, samples=100)[source]

Bases: object

Distribution wrapper that estimates statistics via Monte Carlo sampling.

Provides approximated mean, mode, and entropy helpers for transformed distributions where analytic forms may be inconvenient.

Parameters:

dist (Any)
samples (int)

property name: str

mean()[source]

Return type:: Tensor

mode()[source]

Return type:: Tensor

entropy()[source]

Return type:: Tensor

sample()[source]

Return type:: Tensor

class simulusion.IRISEncoder(vocab_size=512, tokens_per_frame=16, embedding_dim=512, in_channels=3, base_channels=64, num_layers=4, num_residual_blocks=2, frame_shape=(3, 64, 64), commitment_weight=1.0, quantizer='gradient')[source]

Bases: ExportableAgentMixin, Module

CNN Encoder for IRIS discrete autoencoder.

Encodes image observations into latent features, which are then quantized into discrete tokens using the VectorQuantizer.

Architecture:

4 convolutional layers with residual blocks
Self-attention at 8x8 and 16x16 resolutions
Vector quantization to produce discrete tokens

Parameters:

vocab_size (int)
tokens_per_frame (int)
embedding_dim (int)
in_channels (int)
base_channels (int)
num_layers (int)
num_residual_blocks (int)
frame_shape (Tuple[int, int, int])
commitment_weight (float)
quantizer (str)

forward(x)[source]

Encode images to discrete tokens.

Parameters:: x (Tensor) – Input images (B, C, H, W) - should be 64x64
Returns:: Quantized tokens (B, C, H’, W’) indices: Token indices (B, H’, W’) vq_loss: Dictionary with VQ loss components
Return type:: z_q

encode_to_indices(x)[source]

Encode directly to token indices (for world model).

Parameters:: x (Tensor)
Return type:: Tensor

decode_from_indices(indices)[source]

Decode token indices to embeddings (for decoder).

Parameters:: indices (Tensor)
Return type:: Tensor

class simulusion.IRISDecoder(vocab_size=512, embedding_dim=512, base_channels=64, out_channels=3, frame_shape=(3, 64, 64), num_residual_blocks=2)[source]

Bases: ExportableAgentMixin, Module

CNN Decoder for IRIS discrete autoencoder.

Decodes discrete tokens back into image observations. Uses transposed convolutions to upsample from 4x4 to 64x64.

Parameters:

vocab_size (int)
embedding_dim (int)
base_channels (int)
out_channels (int)
frame_shape (Tuple[int, int, int])
num_residual_blocks (int)

forward(z)[source]

Decode tokens to images.

Parameters:: z (Tensor) – Token embeddings (B, C, H, W) - e.g., (B, 512, 4, 4)
Returns:: Reconstructed images (B, C, H, W) - e.g., (B, 3, 64, 64)
Return type:: reconstructed

decode_from_embeddings(z_flat)[source]

Decode flattened token embeddings to images.

Parameters:: z_flat (Tensor) – Flattened tokens (B, H*W, C) or (B, C, H, W)
Returns:: Reconstructed images
Return type:: Tensor

decode_from_indices(indices, codebook)[source]

Decode discrete token indices into images.

The codebook must be passed in explicitly – it is the quantizer’s table (IRISEncoder.quantizer.codebook), the only one the commitment and reconstruction losses train. This decoder previously owned a private index_to_embedding table that no objective ever touched, so decoding through it returned noise that looked plausible enough to go unnoticed.

Parameters:

indices (Tensor) – Tensor of shape (B, H, W) or (B, H*W) containing integer token indices in the range [0, vocab_size).
codebook (Embedding) – The encoder’s quantizer codebook.

Returns:

Reconstructed images (B, C, H, W)

Return type:

Tensor

class simulusion.VideoTokenizer(num_frames=16, image_size=64, in_channels=3, encoder_dim=512, decoder_dim=1024, encoder_depth=12, decoder_depth=20, num_heads=16, patch_size=4, vocab_size=1024, embedding_dim=32, commitment_weight=0.25, use_ema=False, ema_decay=0.99)[source]

Bases: ExportableAgentMixin, Module

Video Tokenizer using VQ-VAE with Spatiotemporal Transformer.

This is a core component of Genie (Google DeepMind, 2024), used to compress raw video frames into discrete latent tokens that can be processed by downstream models like the LatentActionModel and DynamicsModel.

The tokenizer uses Vector Quantized Variational Autoencoder (VQ-VAE) objective to learn a discrete codebook of video representations. Unlike standard VQ-VAE, this uses a Spatiotemporal (ST) Transformer in both encoder and decoder to better capture temporal dynamics in videos.

Architecture

Patch Embedding: Convert (B, C, T, H, W) video to patch tokens
Encoder ST-Transformer: Process spatial-temporal patches
Vector Quantization: Discretize continuous embeddings to codebook entries
Decoder ST-Transformer: Reconstruct video from quantized tokens
Patch Unembedding: Convert tokens back to video frames

Key Features

Causal processing: Each frame’s encoding only uses previous frames
Discrete tokens: Enables autoregressive prediction with latent actions
Memory efficient: Uses ST-Transformer instead of full ViT to reduce complexity

Usage with Genie:

tokenizer = VideoTokenizer(
    num_frames=16,
    image_size=64,
    patch_size=4,
    vocab_size=1024,
    embedding_dim=32
)
reconstructed, indices, loss_dict = tokenizer(video_frames)

# For discrete token input to dynamics model:
token_embeddings = tokenizer.decode_indices(indices)

The tokenizer is trained with VQ-VAE objective: - Reconstruction loss: MSE between input and reconstructed video - VQ loss: Commit to codebook embeddings - Commitment loss: Penalizes encoder outputs drifting from codebook

Reference:: Genie: Generative Interactive Environments Bruce et al., Google DeepMind, 2024 - https://arxiv.org/abs/2402.15391

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
commitment_weight (float)
use_ema (bool)
ema_decay (float)

encode(x)[source]

Encode video to discrete tokens.

Parameters:: x (Tensor) – Video tensor (B, C, T, H, W)
Returns:: Quantized embeddings (B, T, H’, W’, embedding_dim) indices: Token indices (B, T, H’, W’) vq_loss: Dictionary with VQ loss components
Return type:: z_q

decode_indices(indices)[source]

Decode token indices to embeddings for video frames.

Parameters:: indices (Tensor) – Token indices (B, T, H’, W’) or (B, T, N) where N = H’ x W’
Returns:: Quantized embeddings (B, T, H’, W’, embedding_dim)
Return type:: z_q

decode(z_q)[source]

Decode discrete tokens to video frames.

Parameters:: z_q (Tensor) – Quantized embeddings (B, T, H’, W’, embedding_dim)
Returns:: Reconstructed video (B, C, T, H, W)
Return type:: Tensor

forward(x)[source]

Full forward pass with VQ-VAE objective.

Parameters:: x (Tensor) – Video tensor (B, C, T, H, W)
Returns:: Reconstructed video (B, C, T, H, W) indices: Token indices (B, T, H’, W’) loss_dict: Dictionary containing loss components
Return type:: reconstructed

simulusion.create_video_tokenizer(num_frames=16, image_size=64, in_channels=3, encoder_dim=512, decoder_dim=1024, encoder_depth=12, decoder_depth=20, num_heads=16, patch_size=4, vocab_size=1024, embedding_dim=32, use_ema=False)[source]

Factory function to create a Video Tokenizer.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
use_ema (bool)

Return type:

VideoTokenizer

class simulusion.VectorQuantizer(vocab_size=512, embedding_dim=512, commitment_weight=1.0, restart_dead_codes_after=0.01, usage_decay=0.99)[source]

Bases: ExportableAgentMixin, Module

Vector Quantizer for discrete autoencoder.

Implements the VQ-VAE quantization from: “Neural Discrete Representation Learning” (Van Den Oord et al., 2017)

Uses exponential moving averages for codebook updates and straight-through estimator for gradient flow.

Parameters:

vocab_size (int)
embedding_dim (int)
commitment_weight (float)
restart_dead_codes_after (float)
usage_decay (float)

forward(z)[source]

Quantize the input latents.

Parameters:: z (Tensor) – Input tensor of shape (B, C, H, W) or (B, C)
Returns:: Quantized tensor (same shape as input) indices: Token indices for each position (B, H, W) or (B,) loss: Dictionary containing VQ loss components
Return type:: z_q

decode_indices(indices)[source]

Decode token indices back to embeddings.

Parameters:: indices (Tensor) – Token indices (B, H, W) or (B,)
Returns:: Embeddings (B, C, H, W) or (B, C)
Return type:: Tensor

class simulusion.VectorQuantizerEMA(vocab_size=512, embedding_dim=512, commitment_weight=1.0, ema_decay=0.99, epsilon=1e-05, restart_dead_codes_after=0.01)[source]

Bases: ExportableAgentMixin, Module

Vector Quantizer with Exponential Moving Average updates.

Uses EMA updates for the codebook instead of gradient-based updates, which leads to more stable training.

Parameters:

vocab_size (int)
embedding_dim (int)
commitment_weight (float)
ema_decay (float)
epsilon (float)
restart_dead_codes_after (float)

forward(z)[source]

Quantize with EMA updates.

Parameters:: z (Tensor)
Return type:: tuple[Tensor, Tensor, dict[str, Tensor]]

decode_indices(indices)[source]

Decode token indices to embeddings.

Parameters:: indices (Tensor) – Token indices (B, H, W) or (B,)
Returns:: Embeddings (B, C, H, W) or (B, C)
Return type:: Tensor

class simulusion.ReplayBuffer(size, obs_shape, action_size, seq_len, batch_size)[source]

Bases: object

Fixed-size replay buffer for Dreamer with image observations and transitions.

Stores (observation, action, reward, terminal) tuples in a ring buffer and supports sampling contiguous sequences for world-model training.

Key Features

Ring buffer with fixed capacity (FIFO eviction when full)
Stores raw uint8 images to save memory
Samples sequences (not single transitions) for temporal modeling
Validates sampled sequences don’t span episode boundaries

Memory Layout

observations: (capacity, C, H, W) uint8 images
actions: (capacity, action_dim) float32
rewards: (capacity,) float32
terminals: (capacity,) float32 (1.0 = terminal, 0.0 = continue)

Sampling Process

Random start index (avoiding episode boundaries)
Collect sequence of length seq_len with wraparound
Validate no terminal in middle of sequence
Return batch of sequences

Usage with Dreamer:

buffer = ReplayBuffer(
    size=100000,           # Max transitions to store
    obs_shape=(3, 64, 64), # RGB images
    action_size=6,         # Continuous action dim
    seq_len=50,            # Sequence length for training
    batch_size=50          # Parallel sequences per batch
)

# Add transitions during interaction
buffer.add(obs, action, reward, done)

# Sample batch for world model training
obs_batch, action_batch, reward_batch, term_batch = buffer.sample()

Memory Efficiency

Uses uint8 for images (1 byte per pixel vs 4 for float32)
Sequences share observations (overlapping windows)
Configurable capacity based on available system memory

Note

The buffer stores observations as {“image”: …} dicts but returns just the image arrays for training efficiency.

Parameters:

size (int)
obs_shape (Tuple[int, ...])
action_size (int)
seq_len (int)
batch_size (int)

add(obs, ac, rew, done, terminated=None)[source]

Add a transition to the buffer.

Parameters:

obs (dict) – Observation dict with ‘image’ key containing the observation
ac (ndarray) – Action taken, shape (action_size,)
rew (float) – Reward received, scalar
done (float) – Episode-boundary flag, 1.0 if the episode ended for any reason (termination or truncation), 0.0 otherwise
terminated (float | None) – 1.0 only if the environment genuinely terminated. Defaults to done when the caller cannot tell the two apart.

Return type:

None

sample(include_terminated=False)[source]

Sample a batch of sequences for training.

Parameters:

include_terminated (bool) – Also return the true-termination flags as a fifth element.

Returns:

(observations, actions, rewards, terminals[, terminated])

observations: (seq_len, batch, C, H, W)
actions: (seq_len, batch, action_dim)
rewards: (seq_len, batch)
terminals: (seq_len, batch) episode boundaries
terminated: (seq_len, batch) true terminations only

Return type:

tuple

class simulusion.Memory(size=None)[source]

Bases: deque

Episode-based replay memory for PlaNet/RSSM training.

Stores episodes as variable-length trajectories and supports sampling sub-sequences for training. Implements a ring-buffer style eviction when capacity is reached.

Stores complete episodes as lists of transitions
Samples contiguous sub-sequences for sequence models
Supports time-major formatting (time-first) for RNN input
Memory usage estimation to prevent OOM errors

Parameters:

size (int, optional) – Maximum number of episodes to store. If None, deque grows without limit (useful for unpickling).

Variables:

episodes (deque) – Collection of Episode objects.
eps_lengths (deque) – Length of each episode.
size (property) – Total number of transitions across all episodes.

Example:

memory = Memory(size=100)
memory.append([episode1, episode2])
batch, lengths = memory.sample(batch_size=32, tracelen=50)

property size: int

append(episodes)[source]

Parameters:: episodes (list[Episode])
Return type:: None

sample(batch_size, tracelen=1, time_first=False)[source]

Sample random sub-sequences from stored episodes.

Randomly selects episodes and starting positions to create batches of contiguous sequences for training sequence models.

Parameters:

batch_size (int) – Number of sequences to sample.
tracelen (int) – Length of each sequence (default: 1).
time_first (bool) – If True, returns tensors with time dimension first (T, B, …) instead of batch first (B, T, …).

Returns:

(observations, actions, rewards, terminals, lengths)

observations: (batch, tracelen+1, *obs_shape) or (tracelen+1, batch, …)
actions: (batch, tracelen, action_dim) or (tracelen, batch, …)
rewards: (batch, tracelen) or (tracelen, batch)
terminals: (batch, tracelen) or (tracelen, batch)
lengths: (batch,) original episode lengths for each sample

Return type:

tuple

Raises:

ValueError – If memory is empty or no episodes meet minimum length.
MemoryError – If estimated memory usage exceeds 200 MiB threshold.

class simulusion.Episode(postprocess_fn=None)[source]

Bases: object

Records the agent’s interaction with the environment for a single episode.

Stores observations, actions, rewards, and terminal flags during a single trajectory. At termination, converts all lists to numpy arrays for efficient batch processing.

Variables:

x (list or np.ndarray) – Observations collected during the episode.
u (list or np.ndarray) – Actions taken.
r (list or np.ndarray) – Rewards received.
t (list or np.ndarray) – Terminal flags (0.0 = continue, 1.0 = terminal).
info (dict) – Additional episode metadata.

Parameters:

postprocess_fn (callable, optional) – Function to apply to observations before storing (e.g., normalization). Default: identity function.

Example:

episode = Episode()
episode.append(obs, action, reward, False)
episode.append(obs, action, reward, True)
episode.terminate(final_obs)
print(episode.x.shape)  # Now a numpy array

property size: int

append(obs, act, reward, terminal)[source]

Parameters:

obs (Any)
act (Any)
reward (Any)
terminal (Any)

Return type:

None

terminate(obs)[source]

Parameters:: obs (Any)
Return type:: None

class simulusion.IRISReplayBuffer(size, obs_shape, action_size, seq_len=20, batch_size=64)[source]

Bases: object

Replay buffer for IRIS (Imagination with auto-Regression over an Inner Speech) training.

Stores (observation, action, reward, terminal) tuples in a ring buffer and supports sampling contiguous sequences for world model training.

Features:

Ring buffer with fixed capacity (FIFO eviction when full)
Stores uint8 images for memory efficiency
Samples sequences with validation to avoid episode boundaries
Supports sequence sampling for temporal learning

Memory Layout:

observations: (capacity, C, H, W) uint8
actions: (capacity, action_size) float32
rewards: (capacity,) float32
terminals: (capacity,) float32

Parameters:

size (int) – Maximum number of transitions to store.
obs_shape (tuple) – Shape of observations as (C, H, W).
action_size (int) – Dimension of actions.
seq_len (int) – Length of sequences to sample (default: 20).
batch_size (int) – Number of sequences per batch (default: 64).

Variables:

size (int) – Buffer capacity.
obs_shape (tuple) – Observation shape.
action_size (int) – Action dimension.
seq_len (int) – Sequence length.
batch_size (int) – Batch size.
steps (int) – Total transitions added.
episodes (int) – Number of episode terminations observed.

add(obs, action, reward, terminal)[source]

Add a transition to the buffer.

Parameters:

obs (ndarray) – Observation array with shape (C, H, W).
action (ndarray) – Action array with shape (action_size,).
reward (float) – Scalar reward value.
terminal (bool) – Boolean indicating if episode terminated.

Return type:

None

sample_sequence(seq_len=None)[source]

Sample a batch of sequences for world model training.

Returns:: (batch_size, seq_len+1, C, H, W) actions: (batch_size, seq_len, action_size) rewards: (batch_size, seq_len) terminals: (batch_size, seq_len)
Return type:: observations
Parameters:: seq_len (int | None)

sample_with_burn_in(batch_size, burn_in)[source]

Sample start frames together with the frames that precede them.

IRIS burns in the previous frames to initialise the actor-critic’s LSTM state before imagining from a given frame (paper A.3). This returns both halves so the caller does not have to reason about buffer indexing.

Episode boundaries are respected: any burn-in frame at or before a terminal is replaced by a repeat of the oldest valid frame, so context never bleeds across episodes.

Parameters:

batch_size (int) – Number of start frames to draw.
burn_in (int) – Number of preceding frames to return per start frame.

Returns:

(batch_size, C, H, W) uint8 frames to imagine from. burn_in_obs: (batch_size, burn_in, C, H, W) uint8 preceding frames.

Empty along axis 1 when burn_in is 0.

Return type:

start_obs

sample_single()[source]

Sample a single transition for online updates.

Return type:: Tuple[ndarray, ndarray, float, float]

property buffer_capacity: int: Returns the total capacity of the buffer.

class simulusion.IRISOnPolicyBuffer(max_steps=1000)[source]

Bases: object

On-policy buffer for collecting trajectories during environment interaction.

Used to store the current episode data before adding to the main replay buffer. Unlike the main replay buffer, this collects trajectories in a list-based structure that’s cleared after each episode.

Useful for:

Collecting complete episode trajectories
Storing data before batch processing
Temporary storage during environment interaction

Parameters:

max_steps (int) – Maximum number of steps to store (default: 1000).

Variables:

max_steps (int) – Maximum buffer capacity.
observations (list) – List of observations.
actions (list) – List of actions.
rewards (list) – List of rewards.
terminals (list) – List of terminal flags.

add(obs, action, reward, terminal)[source]

Parameters:

obs (ndarray)
action (ndarray)
reward (float)
terminal (bool)

Return type:

None

clear()[source]

Return type:: None

get_arrays()[source]

Return type:: tuple[ndarray, ndarray, ndarray, ndarray]

class simulusion.DiT(img_size, patch_size, in_channels, d_model, depth, heads, drop=0.0, t_dim=256, num_classes=0, class_dropout_prob=0.1, learn_sigma=True, mlp_ratio=4.0)[source]

Bases: ExportableAgentMixin, Module

Diffusion Transformer model for image denoising and generation.

The module maps noisy images and timesteps to predicted noise residuals and also provides a classmethod training entrypoint for common datasets.

Parameters:

img_size (int)
patch_size (int)
in_channels (int)
d_model (int)
depth (int)
heads (int)
drop (float)
t_dim (int)
num_classes (int)
class_dropout_prob (float)
learn_sigma (bool)
mlp_ratio (float)

unpatchify_tokens(x)[source]

Rearrange decoded tokens (B, T, p*p*C_out) back to (B, C_out, H, W).

Parameters:: x (Tensor)
Return type:: Tensor

forward(x, t, y=None)[source]

Predict noise (and covariance when learn_sigma).

Parameters:

x (Tensor) – Noised input (B, C, H, W).
t (Tensor) – Diffusion timesteps (B,).
y (Tensor | None) – Class labels (B,). Required when the model is class-conditional.

Returns:

(B, out_channels, H, W); when learn_sigma the first C channels are the predicted noise and the rest the covariance.

Return type:

Tensor

forward_with_cfg(x, t, y, cfg_scale, guided_channels=None)[source]

Forward pass with classifier-free guidance (paper 3.1).

Computes eps = eps(x, null) + s * (eps(x, y) - eps(x, null)) by running the conditional and unconditional branches as one batch.

Parameters:

x (Tensor) – Noised input (B, C, H, W).
t (Tensor) – Timesteps (B,).
y (Tensor) – Class labels (B,).
cfg_scale (float) – Guidance scale s; 1.0 recovers standard sampling.
guided_channels (int | None) – Apply guidance to only the first N channels. The paper’s appendix guides 3 of the 4 latent channels; None guides all of them, for which the equivalent scale is roughly 1 + 0.75 * (s - 1).

Returns:

Guided model output, same shape as a plain forward pass.

Return type:

Tensor

classmethod from_config(config=None, **overrides)[source]

Build DiT from a config object, dict, YAML file, or YAML string.

Parameters:

config (DiTConfig | dict[str, Any] | str | Path | None)
overrides (Any)

Return type:

Load DiT weights from a local path/directory or HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
config (DiTConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
map_location (str | device | None)
overrides (Any)

Return type:

save_pretrained(path)[source]

Save DiT weights and config in a from_pretrained-compatible format.

Parameters:: path (str | Path)
Return type:: None

parameter_count(trainable_only=False)[source]

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]

Return type:: dict[str, Any]

train(mode=True)[source]

Set training mode – the standard torch.nn.Module.train().

This used to be shadowed by the training-loop classmethod, which made model.eval() raise TypeError: missing 1 required positional argument: 'dataset' and forced callers to reach for nn.Module.train(model, False). The training loop now lives in fit().

Parameters:: mode (bool)
Return type:: DiT

classmethod fit(epochs, dataset, batch_size=256, lr=0.0001, img_size=32, channels=3, patch=4, width=384, depth=12, heads=6, drop=0.0, timesteps=1000, beta_start=0.0001, beta_end=0.02, ema=True, ema_decay=0.9999, num_classes=0, class_dropout_prob=0.1, learn_sigma=True, workdir='./dit_demo', root_path='./data', image_folder=None, crop_size=None, num_workers=4, download=True, copy_data=False, subset_file=None, val_split=None, early_stopping=False, patience=10, min_delta=0.0001, checkpoint_every=0)[source]

Parameters:

epochs (int)
dataset (Any)
batch_size (int)
lr (float)
img_size (int)
channels (int)
patch (int)
width (int)
depth (int)
heads (int)
drop (float)
timesteps (int)
beta_start (float)
beta_end (float)
ema (bool)
ema_decay (float)
num_classes (int)
class_dropout_prob (float)
learn_sigma (bool)
workdir (str)
root_path (str)
image_folder (str | None)
crop_size (int | None)
num_workers (int)
download (bool)
copy_data (bool)
subset_file (str | None)
val_split (float | None)
early_stopping (bool)
patience (int)
min_delta (float)
checkpoint_every (int)

Return type:

None

simulusion.create_dit(config=None, **overrides)[source]

Create a DiT from a DiTConfig or keyword overrides.

The public factory API works with config objects, while DiT itself has a compact constructor. This adapter keeps the lower-level model constructor unchanged and maps the public config fields onto the expected arguments.

Parameters:

config (Any)
overrides (Any)

Return type:

class simulusion.PatchEmbed(img_size, patch_size, in_channels, embed_dim, learnable_pos=False)[source]

Bases: ExportableAgentMixin, Module

Patchify an image into a sequence of learnable patch tokens.

Used in Vision Transformers (ViT) and DiT to convert 2D images into sequences of token embeddings that can be processed by transformers.

Process:

Conv2d with kernel_size=stride=patch_size extracts non-overlapping patches
Each patch is projected to embed_dim via linear layer (Conv2d)
Learnable positional embeddings are added for spatial information

Input: (B, C, H, W) images Output: (B, N, embed_dim) where N = (H/patch_size) * (W/patch_size)

Parameters:

img_size (int) – Image size (assumes square), e.g., 32 for CIFAR
patch_size (int) – Size of each patch (typically 4, 8, or 16)
in_channels (int) – Number of input channels (3 for RGB)
embed_dim (int) – Output dimension for each patch token
learnable_pos (bool)

Usage with DiT:: patch_embed = PatchEmbed(img_size=32, patch_size=4, in_channels=3, embed_dim=256) tokens = patch_embed(images) # (B, 64, 256) for 32x32 image with patch_size=4

forward(x)[source]

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.PatchUnEmbed(img_size, patch_size, embed_dim, out_channels)[source]

Bases: ExportableAgentMixin, Module

Reconstruct image-like tensors from patch-token sequences.

The inverse of PatchEmbed, this module reshapes token sequences into grids and uses transposed convolution to decode spatial outputs.

Parameters:

img_size (int)
patch_size (int)
embed_dim (int)
out_channels (int)

forward(x)[source]

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.DDPM(timesteps, beta_start, beta_end)[source]

Bases: ExportableAgentMixin, Module

Utility module implementing forward and reverse DDPM diffusion steps.

Precomputes diffusion schedule terms and exposes helpers for noising training inputs (q_sample) and iterative denoising sampling (sample).

Parameters:

timesteps (int)
beta_start (float)
beta_end (float)

q_sample(x_start, t, noise=None)[source]

Parameters:

x_start (Tensor)
t (Tensor)
noise (Tensor | None)

Return type:

Tensor

p_sample(model, x_t, t, y=None)[source]

Parameters:

model (Module)
x_t (Tensor)
t (Tensor)
y (Tensor | None)

Return type:

Tensor

sample(model, n, img_size, channels, y=None)[source]

Parameters:

model (Module)
n (int)
img_size (int)
channels (int)
y (Tensor | None)

Return type:

Tensor

class simulusion.ActorCriticNetwork(obs_channels=3, action_dim=18, channels=(32, 32, 64, 64), lstm_dim=512, res_blocks=1, frame_size=64)[source]

Bases: ExportableAgentMixin, Module

Actor-Critic network for DIAMOND RL training. Shared CNN-LSTM trunk with separate policy and value heads.

Parameters:

obs_channels (int)
action_dim (int)
channels (Tuple[int, ...])
lstm_dim (int)
res_blocks (int)
frame_size (int)

forward(obs, hidden_state=None)[source]

Forward pass of actor-critic network.

Parameters:

obs (Tensor) – Observations [B, T, C, H, W]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) hidden states

Returns:

[B, T, action_dim] values: [B, T, 1] hidden_state: (h, c)

Return type:

policy_logits

get_action(obs, hidden_state=None, deterministic=False)[source]

Get action from a single observation.

Parameters:

obs (Tensor) – Single observation [B, C, H, W]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) hidden states
deterministic (bool) – If True, take argmax; else sample

Returns:

Selected action [B] hidden_state: (h, c)

Return type:

get_actions(obs, hidden_state=None, deterministic=False)[source]

Batched version of get_action.

Parameters:

obs (Tensor) – Tensor of shape [B, C, H, W]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional LSTM hidden state tuple matching batch size
deterministic (bool) – If True, take argmax; else sample from policy

Returns:

LongTensor of shape [B] hidden_state: updated LSTM hidden state tuple

Return type:

get_value(obs, hidden_state=None)[source]

Get value for a single observation.

Parameters:

obs (Tensor)
hidden_state (Tuple[Tensor, Tensor] | None)

Return type:

Tuple[Tensor, Tuple[Tensor, Tensor] | None]

init_hidden(batch_size, device)[source]

Initialize LSTM hidden states.

Parameters:

batch_size (int)
device (device)

Return type:

Tuple[Tensor, Tensor]

get_hidden_size()[source]

Get LSTM hidden size.

Return type:: int

class simulusion.RewardTerminationModel(obs_channels=3, action_dim=18, channels=(32, 32, 32, 32), lstm_dim=512, cond_dim=128, res_blocks=2, frame_size=64)[source]

Bases: ExportableAgentMixin, Module

Reward and termination prediction model. CNN + LSTM architecture following DIAMOND paper specifications.

Parameters:

obs_channels (int) – Number of observation channels (3 for RGB)
action_dim (int) – Number of possible actions
channels (Tuple[int, ...]) – List of channel sizes for conv blocks
lstm_dim (int) – LSTM hidden dimension
cond_dim (int) – Conditioning dimension for adaptive norm
res_blocks (int)
frame_size (int)

forward(obs, actions, hidden_state=None)[source]

Forward pass of reward/termination model.

Parameters:

obs (Tensor) – Observations [B, T, C, H, W]
actions (Tensor) – Actions [B, T]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) hidden states

Returns:

Reward predictions [B, T, 3] (for -1, 0, 1) termination_logits: Termination predictions [B, T, 2] hidden_state: Updated (h, c) hidden states

Return type:

reward_logits

predict(obs, actions, hidden_state=None)[source]

Predict reward and termination for a single step.

Parameters:

obs (Tensor) – Single observation [B, C, H, W]
actions (Tensor) – Single action [B]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) hidden states

Returns:

Predicted reward classes as tensor (values -1,0,1) terminated: Predicted termination tensor (bool tensor) hidden_state: Updated (h, c) hidden states

Return type:

reward

init_hidden(batch_size, device)[source]

Initialize LSTM hidden states.

Parameters:

batch_size (int)
device (device)

Return type:

Tuple[Tensor, Tensor]

simulusion.sinusoidal_time_embedding(timesteps, dim, max_period=10000.0)[source]

Create sinusoidal timestep embeddings for diffusion conditioning.

Math:: embedding[t] = [sin(t / P^(2i/d)), cos(t / P^(2i/d))] for i in [0, d/2)

Note the division by increasing powers of max_period: frequencies decay from 1 down to 1/max_period. Building them the other way round (ascending to max_period) makes the sine argument reach ~1e7 radians at t=999, which aliases so badly that adjacent timesteps receive near-orthogonal embeddings – the model then has to memorise every noise level instead of interpolating between them.

Parameters:

timesteps (Tensor) – Tensor of timesteps, shape (B,) or (B, 1)
dim (int) – Embedding dimension
max_period (float) – Longest sinusoid period; 10000 follows DDPM/ADM.

Returns:

Tensor of shape (B, dim) with sinusoidal embeddings.

Return type:

Tensor

class simulusion.STTransformer(num_frames=16, num_patches_per_frame=256, dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>, gradient_checkpointing=False)[source]

Bases: ExportableAgentMixin, Module

Spatiotemporal Transformer for video modeling.

Contains L spatiotemporal blocks with interleaved spatial and temporal attention.

Parameters:

num_frames (int)
num_patches_per_frame (int)
dim (int)
depth (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
qk_scale (float | None)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)
norm_layer (type[Module])
gradient_checkpointing (bool)

init_cache(batch_size, max_frames=None, device=None, dtype=None)[source]

Allocate an empty temporal KV cache for incremental generation.

Parameters:

batch_size (int)
max_frames (int | None)
device (device | None)
dtype (dtype | None)

Return type:

forward(x, cache=None, commit=True)[source]

Parameters:

x (Tensor) – (B, T*N, C) where T is num_frames, N is num_patches_per_frame
cache (STKVCache | None) – Optional temporal KV cache. When given, x holds only the frames after the cached ones, and the result equals the corresponding frames of a full-prefix forward (eval mode).
commit (bool) – Advance cache past these frames. Pass False to evaluate a candidate frame that will be recomputed (e.g. MaskGIT steps).

Returns:

(B, T*N, C)

Return type:

Tensor

class simulusion.MultiHeadSelfAttention(d, n_heads=2)[source]

Bases: ExportableAgentMixin, Module

Multi-head scaled dot-product self-attention over sequence tokens.

This module projects the input sequence into query/key/value heads, performs attention independently per head, and merges the heads back into the original feature dimension. It is used as a lightweight transformer attention block.

Parameters:

d (int)
n_heads (int)

forward(x)[source]

Parameters:: x (Tensor)
Return type:: Tensor

simulusion.MultiHeadAttention: alias of MultiHeadSelfAttention

class simulusion.AdaLNNormalization(d_model, t_dim)[source]

Bases: ExportableAgentMixin, Module

Adaptive layer normalization conditioned on an external embedding.

The module applies RMS normalization and predicts per-channel scale/shift from a conditioning vector (for example diffusion timestep embeddings).

Parameters:

d_model (int)
t_dim (int)

forward(x, t_emb)[source]

Parameters:

x (Tensor)
t_emb (Tensor)

Return type:

Tensor

class simulusion.RMSNorm(dim, eps=1e-06)[source]

Bases: ExportableAgentMixin, Module

Root Mean Square Layer Normalization with a learned gain parameter.

RMSNorm rescales activations using their RMS magnitude without centering, providing a lightweight normalization alternative to LayerNorm.

Parameters:

dim (int)
eps (float)

forward(x)[source]

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.RSSMPolicy(model, planning_horizon, num_candidates, num_iterations, top_candidates, device, compile_rollout=False)[source]

Bases: object

Model-predictive controller using Cross-Entropy Method (CEM) with RSSM.

Plans actions by optimizing a sequence of future actions in the RSSM’s latent space. Uses Cross-Entropy Method to refine action sequences based on predicted returns.

The policy uses a Cross-Entropy Method style loop: it samples candidate action sequences, rolls them forward in latent space, scores predicted returns, and refits a Gaussian proposal to top-performing candidates.

Algorithm:

Initialize Gaussian distribution over action sequences
Sample N candidate action sequences
Rollout each sequence in RSSM latent space
Score by predicted cumulative rewards
Keep top K candidates, fit Gaussian to them
Repeat for T iterations
Execute first action from best sequence

Variables:

rssm – The RSSM world model.
N – Number of candidate action sequences to sample.
K – Number of top candidates to use for updating the proposal.
T – Number of CEM iterations per planning step.
H – Planning horizon (number of future steps to consider).
d – Action dimensionality.
device – Device to run computations on.
state_size – Hidden state dimensionality.
latent_size – Latent state dimensionality.

Parameters:

model (Any)
planning_horizon (int)
num_candidates (int)
num_iterations (int)
top_candidates (int)
device (device | str)
compile_rollout (bool)

Example

>>> policy = RSSMPolicy(
...     model=rssm,
...     planning_horizon=12,
...     num_candidates=1000,
...     num_iterations=5,
...     top_candidates=100,
...     device='cuda'
... )
>>> policy.reset()
>>> action = policy.poll(observation)

reset()[source]

Reset the policy state.

Initializes the hidden state, latent state, and action to zeros. Should be called at the beginning of each episode.

Return type:: None

poll(observation, explore=False)[source]

Get action for given observation.

Parameters:

observation (Tensor) – Current observation tensor of shape (channels, height, width).
explore (bool) – If True, add exploration noise to the selected action.

Returns:

Action tensor of shape (1, action_size).

Return type:

Tensor

class simulusion.RolloutGenerator(env, device, policy=None, max_episode_steps=None, episode_gen=None, name='', enable_streaming_video=False, streaming_video_path=None, streaming_video_fps=20, streaming_video_format='mp4')[source]

Bases: object

Generator for collecting environment rollouts.

This class handles environment interactions and rollout collection, supporting both random and policy-based action selection.

Variables:

env – The environment to interact with.
device – Device to run computations on.
policy – The policy to use for action selection (optional).
episode_gen – Factory for creating episode objects.
name – Name identifier for the generator.
max_episode_steps – Maximum steps per episode.

Parameters:

env (Any)
device (device | str)
policy (Any)
max_episode_steps (int | None)
episode_gen (Any)
name (str)
enable_streaming_video (bool)
streaming_video_path (str | None)
streaming_video_fps (int)
streaming_video_format (str)

Example

>>> generator = RolloutGenerator(
...     env=env,
...     device='cuda',
...     policy=policy,
...     max_episode_steps=1000
... )
>>> episode = generator.rollout_once()

rollout_once(random_policy=False, explore=False)[source]

Perform a single rollout of the environment.

Parameters:

random_policy (bool) – If True, use random actions instead of policy.
explore (bool) – If True, add exploration noise to policy actions.

Returns:

Episode object containing the rollout experience.

Return type:

rollout_n(n=1, random_policy=False)[source]

Perform multiple rollouts.

Parameters:

n (int) – Number of rollouts to perform.
random_policy (bool) – If True, use random actions.

Returns:

List of Episode objects.

Return type:

list

rollout_eval_n(n)[source]

Perform multiple evaluation rollouts with metrics.

Parameters:: n (int) – Number of evaluation rollouts.
Returns:: Tuple of (episodes, frames, metrics).
Return type:: tuple

rollout_eval(collect_latents=False)[source]

Parameters:: collect_latents (bool)
Return type:: tuple

class simulusion.IRISActor(action_size, hidden_size=512, num_layers=4, frame_shape=(3, 64, 64))[source]

Bases: ExportableAgentMixin, Module

Actor network for the IRIS (Imagination with auto-Regression over an Inner Speech) policy.

Takes reconstructed frames as input and outputs action logits for policy control. Uses a CNN feature extractor followed by an LSTM for temporal processing. Supports a burn-in mechanism for initializing the hidden state with context frames.

This standalone actor owns its own CNN and LSTM. (When actor and critic share a backbone, as in the paper, that sharing is done at the IRISAgent level, which builds a single CNN + LSTM feeding separate actor/critic heads.)

Architecture:

CNN: Extracts features from input frames (3x64x64 -> 512)
LSTM: Processes temporal sequences with configurable layers
Linear: Maps hidden states to action logits

Parameters:

action_size (int) – Number of discrete actions.
hidden_size (int) – LSTM hidden state size (default: 512).
num_layers (int) – Number of LSTM layers (default: 4).
frame_shape (tuple) – Shape of input frames as (C, H, W) (default: (3, 64, 64)).

Variables:

action_size (int) – Number of discrete actions.
hidden_size (int) – LSTM hidden state size.
num_layers (int) – Number of LSTM layers.
frame_shape (tuple) – Input frame shape.

forward(frames, hidden_state=None, burn_in_frames=None)[source]

Forward pass through actor.

Parameters:

frames (Tensor) – Input frames (B, T, C, H, W) or (B, C, H, W)
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) tuple for LSTM state
burn_in_frames (Tensor | None) – Frames to use for initializing hidden state

Returns:

Action logits (B, T, action_size) or (B, action_size) hidden_state: Updated (h, c) tuple

Return type:

action_logits

init_hidden_state(batch_size, device)[source]

Initialize LSTM hidden state.

Parameters:

batch_size (int)
device (device)

Return type:

Tuple[Tensor, Tensor]

get_action(frame, temperature=1.0, deterministic=False)[source]

Get action from a single frame.

Parameters:

frame (Tensor) – Single frame (B, C, H, W)
temperature (float) – Softmax temperature (higher = more random)
deterministic (bool) – If True, return argmax; else sample

Returns:

Selected action indices (B,)

Return type:

GymImageEnv | MuJoCoImageEnv

class simulusion.IRISCritic(hidden_size=512, num_layers=4, frame_shape=(3, 64, 64))[source]

Bases: ExportableAgentMixin, Module

Critic network for IRIS value estimation.

Estimates the value function for given frame sequences. It uses the same architecture as the actor (CNN feature extractor + LSTM) and a value head that predicts expected cumulative rewards. This standalone critic instantiates its own CNN and LSTM; backbone sharing between actor and critic is handled at the IRISAgent level, not here.

Architecture:

CNN: Feature extractor with the same architecture as the actor (3x64x64 -> 512)
LSTM: Temporal processing with same architecture as actor
Linear: Maps hidden states to scalar values

Parameters:

hidden_size (int) – LSTM hidden state size (default: 512).
num_layers (int) – Number of LSTM layers (default: 4).
frame_shape (tuple) – Shape of input frames as (C, H, W) (default: (3, 64, 64)).

Variables:

hidden_size (int) – LSTM hidden state size.
num_layers (int) – Number of LSTM layers.
frame_shape (tuple) – Input frame shape.

Returns:

Value estimates with shape (B, T). hidden_state: Updated LSTM hidden state (h, c) tuple.

Return type:

values

forward(frames, hidden_state=None)[source]

Forward pass through critic.

Parameters:

frames (Tensor) – Input frames (B, T, C, H, W)
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) tuple

Returns:

Value estimates (B, T) hidden_state: Updated (h, c) tuple

Return type:

values

init_hidden_state(batch_size, device)[source]

Initialize LSTM hidden state.

Parameters:

batch_size (int)
device (device)

Return type:

Tuple[Tensor, Tensor]

class simulusion.IRISPolicy(action_size, hidden_size=512, num_layers=4, frame_shape=(3, 64, 64))[source]

Bases: ExportableAgentMixin, Module

Combined policy module for IRIS (Imagination with auto-Regression over an Inner Speech).

Provides a unified interface for actor-only or actor-critic policies. Used in the IRIS algorithm where the actor generates actions from reconstructed frames and the critic estimates value functions for training.

Parameters:

action_size (int) – Number of discrete actions.
hidden_size (int) – LSTM hidden state size (default: 512).
num_layers (int) – Number of LSTM layers (default: 4).
frame_shape (tuple) – Shape of input frames as (C, H, W) (default: (3, 64, 64)).

Variables:

actor (IRISActor) – The actor network for action selection.
hidden_size (int) – LSTM hidden state size.
num_layers (int) – Number of LSTM layers.
frame_shape (tuple) – Input frame shape.

Example

>>> policy = IRISPolicy(
...     action_size=18,
...     hidden_size=512,
...     num_layers=4,
...     frame_shape=(3, 64, 64)
... )
>>> action = policy.act(frame, temperature=1.0, deterministic=False)

forward(frames)[source]

Get action logits from frames.

Parameters:: frames (Tensor)
Return type:: Tensor

act(frame, temperature=1.0, deterministic=False)[source]

Sample action from policy.

Parameters:

frame (Tensor)
temperature (float)
deterministic (bool)

Return type:

Tensor

init_hidden(batch_size, device)[source]

Initialize hidden state.

Parameters:

batch_size (int)
device (device)

Return type:

tuple[Tensor, Tensor]

class simulusion.CNNFeatureExtractor(frame_shape=(3, 64, 64), output_size=512, channels=(32, 64, 128, 256))[source]

Bases: ExportableAgentMixin, Module

CNN feature extractor shared between actor and critic networks.

Reproduces the convolutional block of the IRIS actor-critic (paper A.3):

“The convolutional block consists of the same layer repeated four times: a 3x3 convolution with stride 1 and padding 1, a ReLU activation, and 2x2 max-pooling with stride 2.”

Downsampling is therefore done by max-pooling, not by strided convolution. Both reach 64 -> 4 spatially, but max-pooling keeps a full-resolution convolution before each reduction and selects the strongest activation in each window, which preserves small bright objects (the ball in Pong, a bullet) that a stride-2 convolution can skip over entirely.

The paper does not state the channel widths; 32 -> 64 -> 128 -> 256 is kept from the previous implementation.

Parameters:

frame_shape (tuple) – Shape of input frames as (C, H, W) (default: (3, 64, 64)).
output_size (int) – Size of output feature vector (default: 512).
channels (tuple) – Per-layer output channel counts.

Variables:

frame_shape (tuple) – Input frame shape.
output_size (int) – Output feature dimension.

Returns:

Feature vectors with shape (B, output_size).

Return type:

features

forward(x)[source]

Extract features from frames.

Parameters:: x (Tensor) – Frames (B, C, H, W)
Returns:: Feature vectors (B, output_size)
Return type:: features

class simulusion.DreamerConfig(env_backend='dmc', env='walker-walk', env_instance=None, image_size=(64, 64), gym_render_mode='rgb_array', frame_stack=1, dmlab_action_repeat=4, dmlab_action_set=None, dmlab_observations=None, dmlab_config=None, dmlab_renderer='hardware', procgen_distribution_mode='easy', procgen_num_levels=0, procgen_start_level=None, mujoco_xml_path=None, mujoco_xml_string=None, mujoco_binary_path=None, mujoco_camera=None, mujoco_frame_skip=1, mujoco_reset_noise_scale=0.0, brax_backend='generalized', brax_jit=True, brax_auto_reset=False, brax_suppress_warp_warnings=True, unity_file_name=None, unity_behavior_name=None, unity_worker_id=0, unity_base_port=5005, unity_no_graphics=True, unity_time_scale=20.0, unity_quality_level=1, algo='Dreamerv1', exp_name='lr1e-3', train=True, evaluate=False, seed=1, no_gpu=False, max_episode_length=1000, buffer_size=800000, time_limit=1000, cnn_activation_function='relu', dense_activation_function='elu', obs_embed_size=1024, num_units=400, deter_size=200, stoch_size=30, action_repeat=2, action_noise=0.3, total_steps=5000000, seed_steps=5000, update_steps=100, collect_steps=1000, batch_size=50, train_seq_len=50, imagine_horizon=15, use_disc_model=False, free_nats=3.0, discount=0.99, reward_scale=1.0, td_lambda=0.95, kl_loss_coeff=1.0, kl_alpha=0.8, disc_loss_coeff=10.0, num_buckets=255, symlog_range=10.0, model_learning_rate=0.0006, actor_learning_rate=8e-05, value_learning_rate=8e-05, adam_epsilon=1e-07, grad_clip_norm=100.0, use_amp=False, perf_defaults=True, tf32=True, test=False, test_interval=10000, test_episodes=10, scalar_freq=1000, log_video_freq=-1, max_videos_to_save=2, video_format='gif', video_fps=20, checkpoint_interval=10000, early_stopping=False, patience=10, min_delta=0.0001, checkpoint_path='', restore=False, experience_replay='', render=False, enable_wandb=False, wandb_project='simulusion', wandb_entity='', log_dir='runs', logdir=None, data_dir=None, log_level='INFO', log_file=None, enable_tensorboard=False, enable_console_metrics=True, enable_jsonl=True, jsonl_filename='metrics.jsonl', log_system_stats_freq=1000, detect_anomaly=False)[source]

Bases: SerializableConfigMixin

Configuration container for Dreamer training, evaluation, and environment setup.

This class centralizes environment backend selection (DMC/DMLab/Gym/MuJoCo/Robotics/Unity/Brax), model dimensions, replay and optimization settings, logging cadence, and checkpoint options consumed by DreamerAgent.

Parameters:

env_backend (str)
env (str)
env_instance (Any)
image_size (tuple[int, int])
gym_render_mode (str)
frame_stack (int)
dmlab_action_repeat (int)
dmlab_action_set (Any)
dmlab_observations (Any)
dmlab_config (Any)
dmlab_renderer (str)
procgen_distribution_mode (str)
procgen_num_levels (int)
procgen_start_level (Any)
mujoco_xml_path (Any)
mujoco_xml_string (Any)
mujoco_binary_path (Any)
mujoco_camera (Any)
mujoco_frame_skip (int)
mujoco_reset_noise_scale (float)
brax_backend (str)
brax_jit (bool)
brax_auto_reset (bool)
brax_suppress_warp_warnings (bool)
unity_file_name (Any)
unity_behavior_name (Any)
unity_worker_id (int)
unity_base_port (int)
unity_no_graphics (bool)
unity_time_scale (float)
unity_quality_level (int)
algo (str)
exp_name (str)
train (bool)
evaluate (bool)
seed (int)
no_gpu (bool)
max_episode_length (int)
buffer_size (int)
time_limit (int)
cnn_activation_function (str)
dense_activation_function (str)
obs_embed_size (int)
num_units (int)
deter_size (int)
stoch_size (int)
action_repeat (int)
action_noise (float)
total_steps (int)
seed_steps (int)
update_steps (int)
collect_steps (int)
batch_size (int)
train_seq_len (int)
imagine_horizon (int)
use_disc_model (bool)
free_nats (float)
discount (float)
reward_scale (float)
td_lambda (float)
kl_loss_coeff (float)
kl_alpha (float)
disc_loss_coeff (float)
num_buckets (int)
symlog_range (float)
model_learning_rate (float)
actor_learning_rate (float)
value_learning_rate (float)
adam_epsilon (float)
grad_clip_norm (float)
use_amp (bool)
perf_defaults (bool)
tf32 (bool)
test (bool)
test_interval (int)
test_episodes (int)
scalar_freq (int)
log_video_freq (int)
max_videos_to_save (int)
video_format (str)
video_fps (int)
checkpoint_interval (int)
early_stopping (bool)
patience (int)
min_delta (float)
checkpoint_path (str)
restore (bool)
experience_replay (str)
render (bool)
enable_wandb (bool)
wandb_project (str)
wandb_entity (str)
log_dir (str)
logdir (Any)
data_dir (Any)
log_level (str)
log_file (Any)
enable_tensorboard (bool)
enable_console_metrics (bool)
enable_jsonl (bool)
jsonl_filename (str)
log_system_stats_freq (int)
detect_anomaly (bool)

env_backend: str = 'dmc'

env: str = 'walker-walk'

env_instance: Any = None

image_size: tuple[int, int] = (64, 64)

gym_render_mode: str = 'rgb_array'

frame_stack: int = 1

dmlab_action_repeat: int = 4

dmlab_action_set: Any = None

dmlab_observations: Any = None

dmlab_config: Any = None

dmlab_renderer: str = 'hardware'

procgen_distribution_mode: str = 'easy'

procgen_num_levels: int = 0

procgen_start_level: Any = None

mujoco_xml_path: Any = None

mujoco_xml_string: Any = None

mujoco_binary_path: Any = None

mujoco_camera: Any = None

mujoco_frame_skip: int = 1

mujoco_reset_noise_scale: float = 0.0

brax_backend: str = 'generalized'

brax_jit: bool = True

brax_auto_reset: bool = False

brax_suppress_warp_warnings: bool = True

unity_file_name: Any = None

unity_behavior_name: Any = None

unity_worker_id: int = 0

unity_base_port: int = 5005

unity_no_graphics: bool = True

unity_time_scale: float = 20.0

unity_quality_level: int = 1

algo: str = 'Dreamerv1'

exp_name: str = 'lr1e-3'

train: bool = True

evaluate: bool = False

seed: int = 1

no_gpu: bool = False

max_episode_length: int = 1000

buffer_size: int = 800000

time_limit: int = 1000

cnn_activation_function: str = 'relu'

dense_activation_function: str = 'elu'

obs_embed_size: int = 1024

num_units: int = 400

deter_size: int = 200

stoch_size: int = 30

action_repeat: int = 2

action_noise: float = 0.3

total_steps: int = 5000000

seed_steps: int = 5000

update_steps: int = 100

collect_steps: int = 1000

batch_size: int = 50

train_seq_len: int = 50

imagine_horizon: int = 15

use_disc_model: bool = False

free_nats: float = 3.0

discount: float = 0.99

reward_scale: float = 1.0

td_lambda: float = 0.95

kl_loss_coeff: float = 1.0

kl_alpha: float = 0.8

disc_loss_coeff: float = 10.0

num_buckets: int = 255

symlog_range: float = 10.0

model_learning_rate: float = 0.0006

actor_learning_rate: float = 8e-05

value_learning_rate: float = 8e-05

adam_epsilon: float = 1e-07

grad_clip_norm: float = 100.0

use_amp: bool = False

perf_defaults: bool = True

tf32: bool = True

test: bool = False

test_interval: int = 10000

test_episodes: int = 10

scalar_freq: int = 1000

log_video_freq: int = -1

max_videos_to_save: int = 2

video_format: str = 'gif'

video_fps: int = 20

checkpoint_interval: int = 10000

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

checkpoint_path: str = ''

restore: bool = False

experience_replay: str = ''

render: bool = False

enable_wandb: bool = False

wandb_project: str = 'simulusion'

wandb_entity: str = ''

log_dir: str = 'runs'

logdir: Any = None

data_dir: Any = None

log_level: str = 'INFO'

log_file: Any = None

enable_tensorboard: bool = False

enable_console_metrics: bool = True

enable_jsonl: bool = True

jsonl_filename: str = 'metrics.jsonl'

log_system_stats_freq: int = 1000

detect_anomaly: bool = False

class simulusion.JEPAConfig(use_bfloat16=False, model_name='vit_base', load_checkpoint=False, read_checkpoint=None, copy_data=False, pred_depth=None, pred_emb_dim=384, loss_type='l2', dataset='imagenet', val_split=None, use_gaussian_blur=False, use_horizontal_flip=False, use_color_distortion=False, color_jitter_strength=0.0, batch_size=2048, pin_mem=True, num_workers=8, root_path='/data/imagenet', image_folder='train', crop_size=224, crop_scale=(0.3, 1.0), download=False, allow_overlap=False, patch_size=16, num_enc_masks=1, min_keep=10, enc_mask_scale=(0.85, 1.0), num_pred_masks=4, pred_mask_scale=(0.15, 0.2), aspect_ratio=(0.75, 1.5), ema=(0.996, 1.0), ipe_scale=1.0, weight_decay=0.04, final_weight_decay=0.4, epochs=300, warmup=15, start_lr=0.0001, lr=0.001, final_lr=1e-06, lr_reference_batch_size=2048, early_stopping=False, patience=10, min_delta=0.0001, folder='results/jepa', write_tag='jepa_run', enable_wandb=False, wandb_project='simulusion', wandb_entity='', enable_sweep=False, sweep_config=<factory>)[source]

Bases: SerializableConfigMixin

Minimal configuration container for JEPA training. Converts to the nested dict expected by train_jepa.main.

Defaults reproduce I-JEPA as specified in the paper (Assran et al., CVPR 2023, Sec. 3 and Appendix A): no hand-crafted view augmentations, 4 target blocks of scale (0.15, 0.2), a single context block of scale (0.85, 1.0), and the paper’s AdamW/EMA schedules at the reference batch size of 2048. See configs/experiments/jepa_small_gpu.yaml for a single-GPU preset.

Parameters:

use_bfloat16 (bool)
model_name (str)
load_checkpoint (bool)
read_checkpoint (str | None)
copy_data (bool)
pred_depth (int | None)
pred_emb_dim (int)
loss_type (str)
dataset (str)
val_split (float | None)
use_gaussian_blur (bool)
use_horizontal_flip (bool)
use_color_distortion (bool)
color_jitter_strength (float)
batch_size (int)
pin_mem (bool)
num_workers (int)
root_path (str)
image_folder (str)
crop_size (int)
crop_scale (Tuple[float, float])
download (bool)
allow_overlap (bool)
patch_size (int)
num_enc_masks (int)
min_keep (int)
enc_mask_scale (Tuple[float, float])
num_pred_masks (int)
pred_mask_scale (Tuple[float, float])
aspect_ratio (Tuple[float, float])
ema (Tuple[float, float])
ipe_scale (float)
weight_decay (float)
final_weight_decay (float)
epochs (int)
warmup (int)
start_lr (float)
lr (float)
final_lr (float)
lr_reference_batch_size (int | None)
early_stopping (bool)
patience (int)
min_delta (float)
folder (str)
write_tag (str)
enable_wandb (bool)
wandb_project (str)
wandb_entity (str)
enable_sweep (bool)
sweep_config (Dict[str, Any])

use_bfloat16: bool = False

model_name: str = 'vit_base'

load_checkpoint: bool = False

read_checkpoint: str | None = None

copy_data: bool = False

pred_depth: int | None = None

pred_emb_dim: int = 384

loss_type: str = 'l2'

dataset: str = 'imagenet'

val_split: float | None = None

use_gaussian_blur: bool = False

use_horizontal_flip: bool = False

use_color_distortion: bool = False

color_jitter_strength: float = 0.0

batch_size: int = 2048

pin_mem: bool = True

num_workers: int = 8

root_path: str = '/data/imagenet'

image_folder: str = 'train'

crop_size: int = 224

crop_scale: Tuple[float, float] = (0.3, 1.0)

download: bool = False

allow_overlap: bool = False

patch_size: int = 16

num_enc_masks: int = 1

min_keep: int = 10

enc_mask_scale: Tuple[float, float] = (0.85, 1.0)

num_pred_masks: int = 4

pred_mask_scale: Tuple[float, float] = (0.15, 0.2)

aspect_ratio: Tuple[float, float] = (0.75, 1.5)

ema: Tuple[float, float] = (0.996, 1.0)

ipe_scale: float = 1.0

weight_decay: float = 0.04

final_weight_decay: float = 0.4

epochs: int = 300

warmup: int = 15

start_lr: float = 0.0001

lr: float = 0.001

final_lr: float = 1e-06

lr_reference_batch_size: int | None = 2048

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

folder: str = 'results/jepa'

write_tag: str = 'jepa_run'

enable_wandb: bool = False

wandb_project: str = 'simulusion'

wandb_entity: str = ''

enable_sweep: bool = False

sweep_config: Dict[str, Any]

to_dict()[source]

Return type:: Dict[str, Dict[str, Any]]

classmethod from_dict(values)[source]

Load flat field values or the nested trainer dictionary.

Parameters:: values (Dict[str, Any])
Return type:: JEPAConfig

to_train_dict()[source]

Return the nested dictionary expected by train_jepa.

Return type:: Dict[str, Dict[str, Any]]

to_nested_dict()[source]

Backward-compatible alias for the nested JEPA trainer dictionary.

Return type:: Dict[str, Dict[str, Any]]

class simulusion.DiTConfig(DATASET='CIFAR10', BATCH=256, EPOCHS=3, LR=0.0001, WEIGHT_DECAY=0.0, IMG_SIZE=32, CHANNELS=3, PATCH=4, WIDTH=384, DEPTH=12, HEADS=6, DROP=0.0, NUM_CLASSES=0, CLASS_DROPOUT_PROB=0.1, LEARN_SIGMA=True, BETA_START=0.0001, BETA_END=0.02, TIMESTEPS=1000, EMA=True, EMA_DECAY=0.9999, WORKDIR='./dit_demo', ROOT_PATH='./data', CHECKPOINT_EVERY=0, EARLY_STOPPING=False, PATIENCE=10, MIN_DELTA=0.0001, VAL_SPLIT=0.05, CROP_SIZE=None, NUM_WORKERS=4)[source]

Bases: SerializableConfigMixin

Default configuration values for Diffusion Transformer (DiT) training.

The fields define dataset selection, model architecture, diffusion schedule, optimization hyperparameters, and output paths used by the built-in training entrypoints.

Field names use UPPER_CASE for backward compatibility with the original DiT codebase. Snake-case aliases are accepted via __getattr__ and get_dit_config().

Parameters:

DATASET (str)
BATCH (int)
EPOCHS (int)
LR (float)
WEIGHT_DECAY (float)
IMG_SIZE (int)
CHANNELS (int)
PATCH (int)
WIDTH (int)
DEPTH (int)
HEADS (int)
DROP (float)
NUM_CLASSES (int)
CLASS_DROPOUT_PROB (float)
LEARN_SIGMA (bool)
BETA_START (float)
BETA_END (float)
TIMESTEPS (int)
EMA (bool)
EMA_DECAY (float)
WORKDIR (str)
ROOT_PATH (str)
CHECKPOINT_EVERY (int)
EARLY_STOPPING (bool)
PATIENCE (int)
MIN_DELTA (float)
VAL_SPLIT (float)
CROP_SIZE (Any)
NUM_WORKERS (int)

DATASET: str = 'CIFAR10'

BATCH: int = 256

EPOCHS: int = 3

LR: float = 0.0001

WEIGHT_DECAY: float = 0.0

IMG_SIZE: int = 32

CHANNELS: int = 3

PATCH: int = 4

WIDTH: int = 384

DEPTH: int = 12

HEADS: int = 6

DROP: float = 0.0

NUM_CLASSES: int = 0

CLASS_DROPOUT_PROB: float = 0.1

LEARN_SIGMA: bool = True

BETA_START: float = 0.0001

BETA_END: float = 0.02

TIMESTEPS: int = 1000

EMA: bool = True

EMA_DECAY: float = 0.9999

WORKDIR: str = './dit_demo'

ROOT_PATH: str = './data'

CHECKPOINT_EVERY: int = 0

EARLY_STOPPING: bool = False

PATIENCE: int = 10

MIN_DELTA: float = 0.0001

VAL_SPLIT: float = 0.05

CROP_SIZE: Any = None

NUM_WORKERS: int = 4

simulusion.dit_preset_config(name, patch_size, **overrides)[source]

Build a config for a named Table 1 model, e.g. dit_preset_config("DiT-XL", 2).

Parameters:

name (str) – One of DiT-S, DiT-B, DiT-L, DiT-XL (case-insensitive, and the DiT- prefix is optional).
patch_size (int) – Latent patch size; the paper explores 2, 4 and 8.
**overrides (Any) – Further config fields, UPPER_CASE or snake_case.

Returns:

A DiTConfig for that model. Defaults target latent diffusion of 256x256 ImageNet: a 32x32x4 latent with 1000 classes.

Return type:

DiTConfig

simulusion.list_dit_presets()[source]

Return the Table 1 model names.

Return type:: list[str]

simulusion.get_dit_config(**overrides)[source]

Returns a DiTConfig instance with default values overridden by the provided keyword arguments.

Both UPPER_CASE and snake_case override keys are accepted.

Example usage:: cfg = get_dit_config(BATCH=64, EPOCHS=10, LR=1e-3) cfg = get_dit_config(batch=64, epochs=10, lr=1e-3)

Parameters:: overrides (Any)
Return type:: DiTConfig

class simulusion.DiamondConfig(preset: str | None = None, game: str = 'Breakout-v5', seed: int = 0, obs_size: int = 64, frameskip: int = 4, max_noop: int = 30, terminate_on_life_loss: bool = True, reward_clip: List[int] = <factory>, max_episode_steps: int = 27000, num_conditioning_frames: int = 4, diffusion_channels: List[int] = <factory>, diffusion_res_blocks: int = 2, diffusion_cond_dim: int = 256, sigma_data: float = 0.5, sigma_min: float = 0.002, sigma_max: float = 80.0, rho: int = 7, p_mean: float = -0.4, p_std: float = 1.2, sampling_method: str = 'euler', num_sampling_steps: int = 3, reward_channels: List[int] = <factory>, reward_res_blocks: int = 2, reward_cond_dim: int = 128, reward_lstm_dim: int = 512, burn_in_length: int = 4, actor_channels: List[int] = <factory>, actor_res_blocks: int = 1, actor_lstm_dim: int = 512, num_epochs: int = 1000, training_steps_per_epoch: int = 400, batch_size: int = 32, environment_steps_per_epoch: int = 100, epsilon_greedy: float = 0.01, data_loader_num_workers: int = 4, pin_memory: bool = True, persistent_workers: bool = True, imagination_horizon: int = 15, discount_factor: float = 0.985, entropy_weight: float = 0.001, lambda_returns: float = 0.95, learning_rate: float = 0.0001, adam_epsilon: float = 1e-08, weight_decay_diffusion: float = 0.01, weight_decay_reward: float = 0.01, weight_decay_actor: float = 0.0, use_amp: bool = True, device: str = <factory>, log_interval: int = 10, eval_interval: int = 50, save_interval: int = 100, checkpoint_dir: str = 'checkpoints/diamond', early_stopping: bool = False, patience: int = 10, min_delta: float = 0.0001)[source]

Bases: SerializableConfigMixin

Parameters:

preset (str | None)
game (str)
seed (int)
obs_size (int)
frameskip (int)
max_noop (int)
terminate_on_life_loss (bool)
reward_clip (List[int])
max_episode_steps (int)
num_conditioning_frames (int)
diffusion_channels (List[int])
diffusion_res_blocks (int)
diffusion_cond_dim (int)
sigma_data (float)
sigma_min (float)
sigma_max (float)
rho (int)
p_mean (float)
p_std (float)
sampling_method (str)
num_sampling_steps (int)
reward_channels (List[int])
reward_res_blocks (int)
reward_cond_dim (int)
reward_lstm_dim (int)
burn_in_length (int)
actor_channels (List[int])
actor_res_blocks (int)
actor_lstm_dim (int)
num_epochs (int)
training_steps_per_epoch (int)
batch_size (int)
environment_steps_per_epoch (int)
epsilon_greedy (float)
data_loader_num_workers (int)
pin_memory (bool)
persistent_workers (bool)
imagination_horizon (int)
discount_factor (float)
entropy_weight (float)
lambda_returns (float)
learning_rate (float)
adam_epsilon (float)
weight_decay_diffusion (float)
weight_decay_reward (float)
weight_decay_actor (float)
use_amp (bool)
device (str)
log_interval (int)
eval_interval (int)
save_interval (int)
checkpoint_dir (str)
early_stopping (bool)
patience (int)
min_delta (float)

preset: str | None = None

game: str = 'Breakout-v5'

seed: int = 0

obs_size: int = 64

frameskip: int = 4

max_noop: int = 30

terminate_on_life_loss: bool = True

reward_clip: List[int]

max_episode_steps: int = 27000

num_conditioning_frames: int = 4

diffusion_channels: List[int]

diffusion_res_blocks: int = 2

diffusion_cond_dim: int = 256

sigma_data: float = 0.5

sigma_min: float = 0.002

sigma_max: float = 80.0

rho: int = 7

p_mean: float = -0.4

p_std: float = 1.2

sampling_method: str = 'euler'

num_sampling_steps: int = 3

reward_channels: List[int]

reward_res_blocks: int = 2

reward_cond_dim: int = 128

reward_lstm_dim: int = 512

burn_in_length: int = 4

actor_channels: List[int]

actor_res_blocks: int = 1

actor_lstm_dim: int = 512

num_epochs: int = 1000

training_steps_per_epoch: int = 400

batch_size: int = 32

environment_steps_per_epoch: int = 100

epsilon_greedy: float = 0.01

data_loader_num_workers: int = 4

pin_memory: bool = True

persistent_workers: bool = True

imagination_horizon: int = 15

discount_factor: float = 0.985

entropy_weight: float = 0.001

lambda_returns: float = 0.95

learning_rate: float = 0.0001

adam_epsilon: float = 1e-08

weight_decay_diffusion: float = 0.01

weight_decay_reward: float = 0.01

weight_decay_actor: float = 0.0

use_amp: bool = True

device: str

log_interval: int = 10

eval_interval: int = 50

save_interval: int = 100

checkpoint_dir: str = 'checkpoints/diamond'

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

class simulusion.IRISConfig(frame_height=64, frame_width=64, frame_channels=3, vocab_size=512, tokens_per_frame=16, token_embedding_dim=512, encoder_channels=64, encoder_layers=4, encoder_residual_blocks=2, decoder_depth=64, reconstruction_weight=1.0, commitment_weight=1.0, perceptual_weight=1.0, perceptual_blocks=5, perceptual_linear_weights='', quantizer='gradient', transformer_timesteps=20, transformer_embed_dim=256, transformer_layers=10, transformer_heads=4, transformer_dropout=0.1, reward_transform='sign', reward_loss='cross_entropy', imagination_horizon=20, burn_in_length=20, discount=0.995, td_lambda=0.95, entropy_coef=0.001, actor_hidden_size=512, actor_layers=1, total_epochs=600, collection_epochs=500, env_steps_per_epoch=200, training_steps_per_epoch=200, transformer_steps_per_epoch=200, actor_critic_steps_per_epoch=200, model_learning_rate=0.0001, actor_learning_rate=0.0001, value_learning_rate=0.0001, adam_beta1=0.9, adam_beta2=0.999, weight_decay=0.01, grad_clip_norm=10.0, use_amp=True, gradient_checkpointing=True, collect_epsilon=0.01, eval_temperature=0.5, collect_temperature=1.0, start_autoencoder_after=5, start_transformer_after=25, start_actor_critic_after=50, autoencoder_batch_size=256, transformer_batch_size=64, actor_critic_batch_size=64, atari_100k=True, max_env_steps=100000, env_backend='gym', env='ALE/Pong-v5', action_repeat=4, repeat_action_probability=0.0, max_episode_steps=27000, log_interval=1000, eval_episodes=100, checkpoint_interval=50, early_stopping=False, patience=10, min_delta=0.0001)[source]

Bases: SerializableConfigMixin

Configuration for IRIS (Imagination with auto-Regression over an Inner Speech)

Based on paper: “Transformers are Sample-Efficient World Models” Implements discrete autoencoder + autoregressive Transformer for sample-efficient RL.

Parameters:

frame_height (int)
frame_width (int)
frame_channels (int)
vocab_size (int)
tokens_per_frame (int)
token_embedding_dim (int)
encoder_channels (int)
encoder_layers (int)
encoder_residual_blocks (int)
decoder_depth (int)
reconstruction_weight (float)
commitment_weight (float)
perceptual_weight (float)
perceptual_blocks (int)
perceptual_linear_weights (str)
quantizer (str)
transformer_timesteps (int)
transformer_embed_dim (int)
transformer_layers (int)
transformer_heads (int)
transformer_dropout (float)
reward_transform (str)
reward_loss (str)
imagination_horizon (int)
burn_in_length (int)
discount (float)
td_lambda (float)
entropy_coef (float)
actor_hidden_size (int)
actor_layers (int)
total_epochs (int)
collection_epochs (int)
env_steps_per_epoch (int)
training_steps_per_epoch (int)
transformer_steps_per_epoch (int)
actor_critic_steps_per_epoch (int)
model_learning_rate (float)
actor_learning_rate (float)
value_learning_rate (float)
adam_beta1 (float)
adam_beta2 (float)
weight_decay (float)
grad_clip_norm (float)
use_amp (bool)
gradient_checkpointing (bool)
collect_epsilon (float)
eval_temperature (float)
collect_temperature (float)
start_autoencoder_after (int)
start_transformer_after (int)
start_actor_critic_after (int)
autoencoder_batch_size (int)
transformer_batch_size (int)
actor_critic_batch_size (int)
atari_100k (bool)
max_env_steps (int)
env_backend (str)
env (str)
action_repeat (int)
repeat_action_probability (float)
max_episode_steps (int)
log_interval (int)
eval_episodes (int)
checkpoint_interval (int)
early_stopping (bool)
patience (int)
min_delta (float)

frame_height: int = 64

frame_width: int = 64

frame_channels: int = 3

vocab_size: int = 512

tokens_per_frame: int = 16

token_embedding_dim: int = 512

encoder_channels: int = 64

encoder_layers: int = 4

encoder_residual_blocks: int = 2

decoder_depth: int = 64

reconstruction_weight: float = 1.0

commitment_weight: float = 1.0

perceptual_weight: float = 1.0

perceptual_blocks: int = 5

perceptual_linear_weights: str = ''

quantizer: str = 'gradient'

transformer_timesteps: int = 20

transformer_embed_dim: int = 256

transformer_layers: int = 10

transformer_heads: int = 4

transformer_dropout: float = 0.1

reward_transform: str = 'sign'

reward_loss: str = 'cross_entropy'

imagination_horizon: int = 20

burn_in_length: int = 20

discount: float = 0.995

td_lambda: float = 0.95

entropy_coef: float = 0.001

actor_hidden_size: int = 512

actor_layers: int = 1

total_epochs: int = 600

collection_epochs: int = 500

env_steps_per_epoch: int = 200

training_steps_per_epoch: int = 200

transformer_steps_per_epoch: int = 200

actor_critic_steps_per_epoch: int = 200

model_learning_rate: float = 0.0001

actor_learning_rate: float = 0.0001

value_learning_rate: float = 0.0001

adam_beta1: float = 0.9

adam_beta2: float = 0.999

weight_decay: float = 0.01

grad_clip_norm: float = 10.0

use_amp: bool = True

gradient_checkpointing: bool = True

collect_epsilon: float = 0.01

eval_temperature: float = 0.5

collect_temperature: float = 1.0

start_autoencoder_after: int = 5

start_transformer_after: int = 25

start_actor_critic_after: int = 50

autoencoder_batch_size: int = 256

transformer_batch_size: int = 64

actor_critic_batch_size: int = 64

atari_100k: bool = True

max_env_steps: int = 100000

env_backend: str = 'gym'

env: str = 'ALE/Pong-v5'

action_repeat: int = 4

repeat_action_probability: float = 0.0

max_episode_steps: int = 27000

log_interval: int = 1000

eval_episodes: int = 100

checkpoint_interval: int = 50

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

get_frame_shape()[source]

Return type:: tuple

get_autoencoder_config()[source]

Return type:: dict

get_transformer_config()[source]

Return type:: dict

get_rl_config()[source]

Return type:: dict

class simulusion.GenieConfig(num_frames=8, image_size=32, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, tokenizer_encoder_dim=256, tokenizer_decoder_dim=512, tokenizer_encoder_depth=4, tokenizer_decoder_depth=8, tokenizer_num_heads=16, action_vocab_size=8, action_embedding_dim=32, action_encoder_dim=256, action_decoder_dim=1024, action_encoder_depth=4, action_num_heads=16, action_pooling='mean', window_attention_heads=1, dynamics_dim=512, dynamics_depth=8, dynamics_num_heads=8, batch_size=4, learning_rate=3e-05, weight_decay=0.0001, warmup_steps=5000, max_steps=125000, mask_prob_min=0.5, mask_prob_max=1.0, sample_temperature=2.0, maskgit_steps=25, use_amp=False, early_stopping=False, patience=10, min_delta=0.0001, val_split=0.1)[source]

Bases: SerializableConfigMixin

Configuration for Genie model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
tokenizer_encoder_dim (int)
tokenizer_decoder_dim (int)
tokenizer_encoder_depth (int)
tokenizer_decoder_depth (int)
tokenizer_num_heads (int)
action_vocab_size (int)
action_embedding_dim (int)
action_encoder_dim (int)
action_decoder_dim (int)
action_encoder_depth (int)
action_num_heads (int)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
batch_size (int)
learning_rate (float)
weight_decay (float)
warmup_steps (int)
max_steps (int)
mask_prob_min (float)
mask_prob_max (float)
sample_temperature (float)
maskgit_steps (int)
use_amp (bool)
early_stopping (bool)
patience (int)
min_delta (float)
val_split (float)

num_frames: int = 8

image_size: int = 32

in_channels: int = 3

tokenizer_vocab_size: int = 1024

tokenizer_embedding_dim: int = 32

tokenizer_encoder_dim: int = 256

tokenizer_decoder_dim: int = 512

tokenizer_encoder_depth: int = 4

tokenizer_decoder_depth: int = 8

tokenizer_num_heads: int = 16

action_vocab_size: int = 8

action_embedding_dim: int = 32

action_encoder_dim: int = 256

action_decoder_dim: int = 1024

action_encoder_depth: int = 4

action_num_heads: int = 16

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'

window_attention_heads: int = 1

dynamics_dim: int = 512

dynamics_depth: int = 8

dynamics_num_heads: int = 8

batch_size: int = 4

learning_rate: float = 3e-05

weight_decay: float = 0.0001

warmup_steps: int = 5000

max_steps: int = 125000

mask_prob_min: float = 0.5

mask_prob_max: float = 1.0

sample_temperature: float = 2.0

maskgit_steps: int = 25

use_amp: bool = False

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

val_split: float = 0.1

class simulusion.GenieSmallConfig(num_frames=16, image_size=64, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, tokenizer_encoder_dim=256, tokenizer_decoder_dim=512, tokenizer_encoder_depth=4, tokenizer_decoder_depth=8, tokenizer_num_heads=16, action_vocab_size=8, action_embedding_dim=32, action_encoder_dim=512, action_decoder_dim=1024, action_encoder_depth=8, action_num_heads=16, action_pooling='mean', window_attention_heads=1, dynamics_dim=512, dynamics_depth=8, dynamics_num_heads=8, batch_size=2, learning_rate=0.0001, weight_decay=0.0001, warmup_steps=1000, max_steps=50000, mask_prob_min=0.5, mask_prob_max=1.0, sample_temperature=2.0, maskgit_steps=25, use_amp=False, early_stopping=False, patience=10, min_delta=0.0001, val_split=0.1)[source]

Bases: SerializableConfigMixin

Small configuration for development/testing.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
tokenizer_encoder_dim (int)
tokenizer_decoder_dim (int)
tokenizer_encoder_depth (int)
tokenizer_decoder_depth (int)
tokenizer_num_heads (int)
action_vocab_size (int)
action_embedding_dim (int)
action_encoder_dim (int)
action_decoder_dim (int)
action_encoder_depth (int)
action_num_heads (int)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
batch_size (int)
learning_rate (float)
weight_decay (float)
warmup_steps (int)
max_steps (int)
mask_prob_min (float)
mask_prob_max (float)
sample_temperature (float)
maskgit_steps (int)
use_amp (bool)
early_stopping (bool)
patience (int)
min_delta (float)
val_split (float)

num_frames: int = 16

image_size: int = 64

in_channels: int = 3

tokenizer_vocab_size: int = 1024

tokenizer_embedding_dim: int = 32

tokenizer_encoder_dim: int = 256

tokenizer_decoder_dim: int = 512

tokenizer_encoder_depth: int = 4

tokenizer_decoder_depth: int = 8

tokenizer_num_heads: int = 16

action_vocab_size: int = 8

action_embedding_dim: int = 32

action_encoder_dim: int = 512

action_decoder_dim: int = 1024

action_encoder_depth: int = 8

action_num_heads: int = 16

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'

window_attention_heads: int = 1

dynamics_dim: int = 512

dynamics_depth: int = 8

dynamics_num_heads: int = 8

batch_size: int = 2

learning_rate: float = 0.0001

weight_decay: float = 0.0001

warmup_steps: int = 1000

max_steps: int = 50000

mask_prob_min: float = 0.5

mask_prob_max: float = 1.0

sample_temperature: float = 2.0

maskgit_steps: int = 25

use_amp: bool = False

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

val_split: float = 0.1

class simulusion.STTransformerConfig(num_frames=16, num_patches_per_frame=256, dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0)[source]

Bases: SerializableConfigMixin

Configuration for Spatiotemporal Transformer.

Parameters:

num_frames (int)
num_patches_per_frame (int)
dim (int)
depth (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)

num_frames: int = 16

num_patches_per_frame: int = 256

dim: int = 768

depth: int = 12

num_heads: int = 12

mlp_ratio: float = 4.0

qkv_bias: bool = True

drop_rate: float = 0.0

attn_drop_rate: float = 0.0

drop_path_rate: float = 0.0

class simulusion.VideoTokenizerConfig(num_frames=16, image_size=64, in_channels=3, encoder_dim=512, decoder_dim=1024, encoder_depth=12, decoder_depth=20, num_heads=16, patch_size=4, vocab_size=1024, embedding_dim=32, use_ema=False, ema_decay=0.99, commitment_weight=0.25)[source]

Bases: SerializableConfigMixin

Configuration for Video Tokenizer.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
use_ema (bool)
ema_decay (float)
commitment_weight (float)

num_frames: int = 16

image_size: int = 64

in_channels: int = 3

encoder_dim: int = 512

decoder_dim: int = 1024

encoder_depth: int = 12

decoder_depth: int = 20

num_heads: int = 16

patch_size: int = 4

vocab_size: int = 1024

embedding_dim: int = 32

use_ema: bool = False

ema_decay: float = 0.99

commitment_weight: float = 0.25

class simulusion.LatentActionModelConfig(num_frames=16, image_size=64, in_channels=3, encoder_dim=1024, encoder_depth=20, num_heads=16, patch_size=16, vocab_size=8, embedding_dim=32, commitment_weight=1.0, action_pooling='mean', window_attention_heads=1)[source]

Bases: SerializableConfigMixin

Configuration for Latent Action Model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
encoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
commitment_weight (float)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

num_frames: int = 16

image_size: int = 64

in_channels: int = 3

encoder_dim: int = 1024

encoder_depth: int = 20

num_heads: int = 16

patch_size: int = 16

vocab_size: int = 8

embedding_dim: int = 32

commitment_weight: float = 1.0

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'

window_attention_heads: int = 1

class simulusion.DynamicsModelConfig(num_frames=16, image_size=64, vocab_size=1024, embedding_dim=32, action_vocab_size=8, dim=5120, depth=48, num_heads=40, patch_size=4, mlp_ratio=4.0, qkv_bias=True, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0)[source]

Bases: SerializableConfigMixin

Configuration for Dynamics Model.

Parameters:

num_frames (int)
image_size (int)
vocab_size (int)
embedding_dim (int)
action_vocab_size (int)
dim (int)
depth (int)
num_heads (int)
patch_size (int)
mlp_ratio (float)
qkv_bias (bool)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)

num_frames: int = 16

image_size: int = 64

vocab_size: int = 1024

embedding_dim: int = 32

action_vocab_size: int = 8

dim: int = 5120

depth: int = 48

num_heads: int = 40

patch_size: int = 4

mlp_ratio: float = 4.0

qkv_bias: bool = True

drop_rate: float = 0.0

attn_drop_rate: float = 0.0

drop_path_rate: float = 0.0

class simulusion.BSuiteImageEnv(bsuite_id, seed=0, size=(64, 64), env=None, include_state=False)[source]

Bases: object

Gym-like wrapper for DeepMind BSuite dm_env environments.

BSuite tasks expose compact dm_env observations and mostly discrete actions. This adapter presents a Gym/Gymnasium-style API with image observations under a dict with required "image" key so Simulusion’s pixel-based world models can train and evaluate on BSuite diagnostic tasks without requiring the base environment to implement rendering. When include_state=True, observations also include a flattened "state" key.

Parameters:

bsuite_id (str)
seed (int)
size (tuple[int, int])
env (Any | None)
include_state (bool)

property observation_space: Dict

property action_space: Space

property max_episode_steps: int

reset(seed=None)[source]

Parameters:: seed (int | None)
Return type:: dict[str, ndarray]

step(action)[source]

Parameters:: action (Any)
Return type:: tuple[dict[str, ndarray], float, bool, dict[str, Any]]

render(*args, **kwargs)[source]

Parameters:

args (Any)
kwargs (Any)

Return type:

ndarray

close()[source]

Return type:: None

simulusion.make_bsuite_env(bsuite_id, **kwargs)[source]

Create a Dreamer-compatible image wrapper around a BSuite task.

Parameters:

bsuite_id (str)
kwargs (Any)

Return type:

BSuiteImageEnv

simulusion.list_available_bsuite_ids()[source]

Return the installed BSuite sweep ids, or examples if BSuite is absent.

Return type:: list[str]

simulusion.make_diamond_atari_env(game, frameskip=4, max_noop=30, terminate_on_life_loss=True, reward_clip=True, resize=(64, 64), seed=None)[source]

Create a DIAMOND-compatible Atari environment.

Parameters:

game (str) – Atari game name. Accepts either the bare "Breakout-v5" or the namespaced "ALE/Breakout-v5"; the ALE/ prefix is added when missing. DIAMOND checkpoints store the bare form, so requiring the namespaced one here would make them unloadable.
frameskip (int) – Number of frames to skip between actions
max_noop (int) – Maximum number of noop actions at reset
terminate_on_life_loss (bool) – Whether to terminate on life loss
reward_clip (bool) – Whether to clip rewards to [-1, 0, 1]
resize (Tuple[int, int]) – Target size for observations
seed (int | None) – Random seed

Returns:

Configured Atari environment

Return type:

DiamondAtariWrapper

class simulusion.MuJoCoImageEnv(xml_path=None, *, xml_string=None, binary_path=None, assets=None, seed=0, size=(64, 64), camera=None, reward_fn=None, terminal_fn=None, frame_skip=1, reset_noise_scale=0.0, default_control_range=(-1.0, 1.0), include_state=False)[source]

Bases: object

Native MuJoCo environment adapter for pixel-based world-model training.

The adapter uses the low-level mujoco Python package directly: models are compiled from MJCF XML strings/files or MJB binaries via mujoco.MjModel; simulation state lives in mujoco.MjData; actions are written to data.ctrl; and images are produced with mujoco.Renderer. Observations follow Simulusion’s Dreamer-style contract: a dict with required key "image" containing uint8[C, H, W] frames.

Native MuJoCo models do not define task rewards or episode termination by themselves, so callers can supply reward_fn and terminal_fn callbacks. By default, rewards are 0.0 and episodes terminate only through external wrappers such as TimeLimit.

Parameters:

xml_path (str | Path | None)
xml_string (str | None)
binary_path (str | Path | None)
assets (dict[str, bytes] | None)
seed (int)
size (tuple[int, int])
camera (str | int | None)
reward_fn (RewardFn | None)
terminal_fn (TerminalFn | None)
frame_skip (int)
reset_noise_scale (float)
default_control_range (tuple[float, float])
include_state (bool)

property observation_space: Dict

property action_space: Box

reset(seed=None)[source]

Parameters:: seed (int | None)
Return type:: dict[str, ndarray]

step(action)[source]

Parameters:: action (Any)
Return type:: tuple[dict[str, ndarray], float, bool, dict[str, Any]]

render()[source]

Return type:: Any

close()[source]

Return type:: None

simulusion.make_mujoco_env(model=None, *, backend='auto', seed=0, size=(64, 64), render_mode='rgb_array', gym_kwargs=None, **kwargs)[source]

Create one MuJoCo image environment factory for tasks and MJCF/MJB models.

Parameters:

model (str | Path | None) – Either a Gymnasium MuJoCo task id such as "Humanoid-v4", an MJCF XML path/string, or an MJB binary path.
backend (str) – "auto" infers native vs Gymnasium task mode. Use "native" for MJCF/MJB, "gymnasium" for task ids, or "robotics" for Gymnasium Robotics registrations.
seed (int) – Seed forwarded to the image wrapper.
size (tuple[int, int]) – Target (height, width) image size.
render_mode (str) – Render mode used for Gymnasium MuJoCo task ids.
gym_kwargs (dict[str, Any] | None) – Optional keyword arguments forwarded to gymnasium.make in task-id mode. Extra **kwargs are also forwarded there.
**kwargs (Any) – Native MuJoCoImageEnv options for MJCF/MJB mode, or environment-constructor options for Gymnasium task-id mode.

Returns:

A Simulusion image environment returning {"image": uint8[C, H, W]}.

Return type:

simulusion.make_mujoco_env_from_config(args, size)[source]

Build a MuJoCo image environment from a DreamerConfig-like object.

Parameters:

args (Any)
size (tuple[int, int])

Return type:

Any

simulusion.list_gymnasium_robotics_envs()[source]

List all Gymnasium Robotics ids registered by the installed package.

Returns an empty list when the optional dependency is not installed. When it is installed, the list is derived from Gymnasium’s registry rather than a hand-maintained subset, so newly added Robotics environments are exposed automatically.

Return type:: list[str]

simulusion.make_robotics_env(env, *, seed=0, size=(64, 64), render_mode='rgb_array', gym_kwargs=None, **kwargs)[source]

Create a Simulusion image wrapper for a Gymnasium Robotics environment.

Parameters:

env (str) – Any environment id registered by gymnasium-robotics.
seed (int) – Seed forwarded to GymImageEnv.
size (tuple[int, int]) – Target (height, width) image size.
render_mode (str) – Render mode forwarded to gymnasium.make.
gym_kwargs (dict[str, Any] | None) – Optional keyword arguments forwarded to gymnasium.make.
**kwargs (Any) – Additional keyword arguments forwarded to gymnasium.make.

Returns:

A GymImageEnv that emits {"image": uint8[C, H, W]} observations.

Return type:

GymImageEnv

simulusion.register_gymnasium_robotics_envs()[source]

Import Gymnasium Robotics so its environments are registered with Gymnasium.

Gymnasium moved legacy MuJoCo v2/v3 task registrations into the external gymnasium-robotics package. Current Gymnasium Robotics versions register environments during import, while older plugin-style installations may rely on gymnasium.register_envs; this helper supports both paths.

Return type:: Any

class simulusion.GymImageEnv(env, seed=0, size=(64, 64), render_mode='rgb_array', include_state=False)[source]

Bases: object

Gym-like environment wrapper that always returns image observations.

This wrapper normalizes diverse environment interfaces to return consistent image-based observations suitable for pixel-based world models like Dreamer.

Features:

Supports environment IDs (string) and pre-built environment objects.
For vector observations, uses the environment’s own render() frame, and only synthesizes band images from the vector when rendering fails.
Exposes continuous action spaces mapped to [-1, 1] range.
Converts discrete actions to one-hot vectors.
Returns observations as dicts with required key "image" and optional key "state" when include_state=True and a vector observation is available.

Parameters:

env (Any) – Either a string environment ID (e.g., “Pendulum-v1”) or a pre-built gym environment instance.
seed (int) – Random seed for environment reset (default: 0).
size (tuple) – Target image size as (height, width) (default: (64, 64)).
render_mode (str) – Render mode for environment (default: “rgb_array”).
include_state (bool) – Include a flattened low-dimensional "state" key in observations when one can be derived from the underlying env.

Variables:

observation_space – Dict space with required "image" key and optional "state" key when enabled.
action_space – Box space with actions in [-1, 1] range.
max_episode_steps – Maximum steps per episode (default: 1000).

property observation_space: Dict

property action_space: Box

property max_episode_steps: int

reset(seed=None)[source]

Parameters:: seed (int | None)
Return type:: dict[str, Any]

step(action)[source]

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], float, bool, dict[str, Any]]

render(*args, **kwargs)[source]

Parameters:

args (Any)
kwargs (Any)

Return type:

ndarray

close()[source]

Return type:: None

simulusion.make_gym_env(env, **kwargs)[source]

Create a GymImageEnv wrapper for generic Gym/Gymnasium environments.

Parameters:

env (Any) – Either a string environment ID (e.g., “Pendulum-v1”) or a pre-built gym environment instance.
**kwargs (Any) – Additional keyword arguments passed to GymImageEnv, including: - seed (int): Random seed for environment (default: 0) - size (tuple): Target image size as (height, width) (default: (64, 64)) - render_mode (str): Render mode for environment (default: “rgb_array”)

Returns:

A wrapper that always returns image observations in the: format {“image”: (C, H, W)} suitable for pixel-based world models.

Return type:

GymImageEnv

class simulusion.WorldModelEnv(world_model, observation_space, action_space, *, initial_observation=None, initial_state=None, reset_fn=None, transition_fn=None, reward_fn=None, terminal_fn=None, render_fn=None, action_transform_fn=None, max_episode_steps=None, render_mode=None, device=None, torch_actions=True, seed=None)[source]

Bases: Env

Expose a trained world model through the Gymnasium Env API.

WorldModelEnv keeps the current latent/model state and advances it with a transition callable or with a compatible method on world_model. The wrapper returns Gymnasium-style (obs, info) from reset and (obs, reward, terminated, truncated, info) from step, making learned model rollouts pluggable into RL libraries such as Stable-Baselines3, TorchRL, and CleanRL.

Parameters:

world_model (Any) – Trained model or lightweight adapter object used for simulated dynamics.
observation_space (gym.Space) – Gymnasium observation space emitted by the wrapper.
action_space (gym.Space) – Gymnasium action space accepted by the wrapper.
initial_observation (Any | None) – Optional observation returned when no reset callable provides one. Defaults to observation_space.sample().
initial_state (Any | None) – Optional latent/model state used at reset.
reset_fn (ResetFn | None) – Optional callable for resetting model state. Accepted return forms are obs, (obs, info), (state, obs), (state, obs, info), or a mapping with state/observation.
transition_fn (TransitionFn | None) – Optional callable for one model step. If omitted, the wrapper tries common methods on world_model: env_step, step, predict_step, predict, imagine_step, transition, then __call__.
reward_fn (RewardFn | None) – Optional callable used when the transition output omits a reward.
terminal_fn (TerminalFn | None) – Optional callable used when the transition output omits a termination flag.
render_fn (RenderFn | None) – Optional callable used by render.
action_transform_fn (ActionTransformFn | None) – Optional callable that converts library actions into the format expected by the world model.
max_episode_steps (int | None) – Optional time limit. Reaching it sets truncated.
render_mode (str | None) – Optional Gymnasium render mode. rgb_array is supported by default when observations contain image-like data.
device (Any | None) – Device used for tensor actions when torch_actions=True.
torch_actions (bool) – Convert actions to torch.Tensor before model calls.
seed (int | None) – Optional RNG seed for observation/action spaces and NumPy.

metadata = {'render_fps': 30, 'render_modes': ['rgb_array']}

property state: Any: Current latent/model state tracked by the wrapper.

reset(*, seed=None, options=None)[source]

Reset the simulated rollout and return (observation, info).

Parameters:

seed (int | None)
options (dict[str, Any] | None)

Return type:

tuple[Any, dict[str, Any]]

step(action)[source]

Roll the learned model forward for one simulated environment step.

Parameters:: action (Any)
Return type:: tuple[Any, float, bool, bool, dict[str, Any]]

render()[source]

Render the latest simulated observation or delegate to render_fn.

Return type:: Any

close()[source]

Close the wrapped world model if it exposes close.

Return type:: None

simulusion.make_world_model_env(world_model, **kwargs)[source]

Create a WorldModelEnv from a trained model and spaces.

Parameters:

world_model (Any)
kwargs (Any)

Return type:

WorldModelEnv

class simulusion.BraxImageEnv(env, seed=0, size=(64, 64), backend=None, episode_length=None, auto_reset=False, jit=True, suppress_warp_warnings=True, include_state=False, **env_kwargs)[source]

Bases: object

Gym-like adapter for training Simulusion world models on Brax tasks.

Brax environments are functional JAX environments: reset consumes a PRNG key and returns a state, while step consumes the previous state plus an action and returns the next state. This adapter stores the Brax state between calls and converts state observations into image observations compatible with pixel-based Simulusion agents such as Dreamer.

If a Brax renderer is not available, vector observations are rendered as deterministic feature-band images so training code can still consume a pixel stream. The original vector observation is also exposed through info["vector_observation"] after step for diagnostics. When include_state=True, observations also expose a flattened "state" key.

Parameters:

env (str | Any)
seed (int)
size (tuple[int, int])
backend (str | None)
episode_length (int | None)
auto_reset (bool)
jit (bool)
suppress_warp_warnings (bool)
include_state (bool)
env_kwargs (Any)

property observation_space: Space

property action_space: Space

property max_episode_steps: int

reset(seed=None)[source]

Parameters:: seed (int | None)
Return type:: dict[str, Any]

step(action)[source]

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], float, bool, dict[str, Any]]

render(*args, **kwargs)[source]

Parameters:

args (Any)
kwargs (Any)

Return type:

ndarray

close()[source]

Return type:: None

simulusion.make_brax_env(env, **kwargs)[source]

Create a Simulusion image wrapper for Brax environments.

Parameters:

env (str | Any) – Brax environment name (for example, "ant") or a pre-built Brax environment object exposing reset(rng) and step(state, action).
**kwargs (Any) – Additional keyword arguments passed to BraxImageEnv.

Returns:

A Gym-like wrapper that returns {"image": (C, H, W)} observations and exposes continuous actions in the Brax [-1, 1] range.

Return type:

BraxImageEnv

class simulusion.DeepMindControlEnv(name, seed, size=(64, 64), camera=None)[source]

Bases: object

Gym-style adapter for DeepMind Control Suite tasks.

The wrapper exposes DMC observations and actions through Gym spaces and adds a rendered RGB image to each observation dict so image-based world model pipelines can train consistently across backends.

Features:

Parses domain-task names (e.g., “cheetah-run” -> domain=”cheetah”, task=”run”)
Automatically handles special cases like “cup” -> “ball_in_cup”
Renders RGB images at configurable resolution
Returns observations as dict with both state vectors and images

Parameters:

name (str) – Environment name in format “domain-task” (e.g., “cheetah-run”).
seed (int) – Random seed for environment initialization.
size (tuple) – Target image size as (height, width) (default: (64, 64)).
camera (int, optional) – Camera ID for rendering. Defaults to 0 for most domains, 2 for quadruped.

Variables:

observation_space (gym.spaces.Dict) – Dict space with state keys and “image”.
action_space (gym.spaces.Box) – Continuous action space from DMC spec.

Example

>>> env = DeepMindControlEnv("cheetah-run", seed=0, size=(64, 64))
>>> obs = env.reset()
>>> print(obs.keys())  # dict_keys(['position', 'velocity', 'image'])

property observation_space: Dict

property action_space: Box

step(action)[source]

Parameters:: action (ndarray)
Return type:: tuple[dict, float, bool, dict]

reset(seed=None)[source]

Parameters:: seed (int | None)
Return type:: dict

render(*args, **kwargs)[source]

Parameters:

args (Any)
kwargs (Any)

Return type:

ndarray

class simulusion.DMLabEnv(level, seed=0, size=(64, 64), action_repeat=4, action_set=None, observations=None, config=None, renderer='hardware', **lab_kwargs)[source]

Bases: object

Gym-style adapter for DeepMind Lab 3D environments.

The native deepmind_lab API exposes RGB observations as HWC arrays and expects a seven-element integer action vector. This adapter presents a Simulusion-friendly image observation dict and a Box action space containing a one-hot vector in [-1, 1] so it composes with Dreamer’s normalization wrappers.

Parameters:

level (str)
seed (int)
size (tuple[int, int])
action_repeat (int)
action_set (Sequence[Sequence[int]] | np.ndarray | None)
observations (Sequence[str] | None)
config (dict[str, Any] | None)
renderer (str)
lab_kwargs (Any)

property observation_space: Dict

property action_space: _OneHotActionSpace

property max_episode_steps: int

reset(seed=None)[source]

Parameters:: seed (int | None)
Return type:: dict[str, ndarray]

step(action)[source]

Parameters:: action (ndarray)
Return type:: tuple[dict[str, ndarray], float, bool, dict[str, Any]]

render(*args, **kwargs)[source]

Parameters:

args (Any)
kwargs (Any)

Return type:

ndarray

close()[source]

Return type:: None

simulusion.make_dmlab_env(level, **kwargs)[source]

Create a DeepMind Lab environment adapter for Simulusion.

Parameters:

level (str) – DeepMind Lab level name, for example "rooms_collect_good_objects_train".
**kwargs (Any) – Additional keyword arguments passed to DMLabEnv.

Returns:

A Gym-like wrapper returning {"image": (C, H, W)} uint8 observations and normalized one-hot discrete actions.

Return type:

DMLabEnv

class simulusion.UnityMLAgentsEnv(file_name, behavior_name=None, seed=0, size=(64, 64), worker_id=0, base_port=5005, no_graphics=True, time_scale=20.0, quality_level=1, max_episode_steps=1000, include_state=False)[source]

Bases: object

Gym-like wrapper for Unity ML-Agents environments.

Provides a unified interface for Unity-based environments, converting observations to image format compatible with pixel-based world models.

Features:

Supports single-agent control with continuous action spaces.
Returns observations as dicts with required key "image".
Normalizes actions to [-1, 1] range.
Exposes non-visual sensors in info["vector_observation"] for debugging.

Parameters:

file_name (str) – Path to the Unity environment binary.
behavior_name (str, optional) – Name of the behavior to use. If None, uses the first available behavior.
seed (int) – Random seed for environment (default: 0).
size (tuple) – Target image size as (height, width) (default: (64, 64)).
worker_id (int) – Worker ID for multi-environment setup (default: 0).
base_port (int) – Base port for Unity environment communication (default: 5005).
no_graphics (bool) – Disable graphics rendering for faster simulation (default: True).
time_scale (float) – Simulation time scale multiplier (default: 20.0).
quality_level (int) – Graphics quality level 0-5 (default: 1).
max_episode_steps (int) – Maximum steps per episode (default: 1000).
include_state (bool) – Include a flattened non-visual "state" key in observations when the Unity behavior exposes non-image sensors.

Variables:

observation_space – Dict space with required "image" key containing (3, H, W) uint8 frames.
action_space – Box space with actions in [-1, 1] range.
max_episode_steps – Maximum steps per episode.

Raises:

ValueError – If no behaviors found or action space is not continuous.
RuntimeError – If no agents available after reset.

property observation_space: Dict

property action_space: Box

property max_episode_steps: int

reset(seed=None)[source]

Parameters:: seed (int | None)
Return type:: dict[str, Any]

step(action)[source]

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], float, bool, dict[str, Any]]

render(*args, **kwargs)[source]

Parameters:

args (Any)
kwargs (Any)

Return type:

Any

close()[source]

Return type:: None

simulusion.make_unity_mlagents_env(env_id=None, **kwargs)[source]

Create a Unity ML-Agents environment wrapper.

Factory function that instantiates a UnityMLAgentsEnv with the provided keyword arguments. Suitable for integrating Unity-based environments with Dreamer-style world model pipelines.

Parameters:

**kwargs (Any) – Keyword arguments passed to UnityMLAgentsEnv, including: - file_name (str): Path to the Unity environment binary. - behavior_name (str, optional): Name of the behavior to use. - seed (int): Random seed (default: 0). - size (tuple): Image size as (height, width) (default: (64, 64)). - worker_id (int): Worker ID for multi-environment setup (default: 0). - base_port (int): Base port for communication (default: 5005). - no_graphics (bool): Disable graphics rendering (default: True). - time_scale (float): Simulation time scale (default: 20.0). - quality_level (int): Graphics quality level (default: 1). - max_episode_steps (int): Max steps per episode (default: 1000).
env_id (str | None)
**kwargs

Returns:

A Gym-compatible wrapper for Unity environments.

Return type:

UnityMLAgentsEnv

simulusion.MujocoEnv: alias of MuJoCoImageEnv

class simulusion.TimeLimit(env, duration)[source]

Bases: object

Terminate episodes after a fixed number of wrapper steps.

If the wrapped environment does not provide a discount flag at timeout, the wrapper injects a default discount of 1.0 for downstream learners.

Parameters:

env (Any)
duration (int)

step(action)[source]

Parameters:: action (Any)
Return type:: tuple[Any, Any, bool, dict[str, Any]]

reset(*, seed=None)[source]

Parameters:: seed (int | None)
Return type:: Any

class simulusion.ActionRepeat(env, amount)[source]

Bases: object

Repeat each action for a fixed number of environment steps.

Rewards are accumulated and the loop stops early if the environment terminates, mirroring common action-repeat behavior in world model papers.

Parameters:

env (Any)
amount (int)

step(action)[source]

Parameters:: action (Any)
Return type:: tuple[Any, float, bool, dict[str, Any]]

class simulusion.NormalizeActions(env)[source]

Bases: object

Expose a normalized [-1, 1] action space for bounded continuous controls.

Incoming normalized actions are mapped back to the wrapped environment action bounds before stepping the environment.

Parameters:: env (Any)

property action_space: Box

step(action)[source]

Parameters:: action (ndarray)
Return type:: tuple[Any, Any, bool, dict[str, Any]]

class simulusion.ObsDict(env, key='obs')[source]

Bases: object

Convert scalar/array observations into a dictionary observation format.

This harmonizes outputs for code paths that expect keyed observations (for example {“image”: …} style world model inputs).

Parameters:

env (Any)
key (str)

property observation_space: Dict

property action_space: Any

step(action)[source]

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], Any, bool, dict[str, Any]]

reset(*, seed=None)[source]

Parameters:: seed (int | None)
Return type:: dict[str, Any]

class simulusion.OneHotAction(env)[source]

Bases: object

Wrap discrete-action environments to accept one-hot action vectors.

The wrapper validates one-hot inputs and converts them to integer action indices before forwarding to the underlying environment.

Parameters:: env (Any)

property action_space: Box

step(action)[source]

Parameters:: action (ndarray)
Return type:: tuple[Any, Any, bool, dict[str, Any]]

reset(*, seed=None)[source]

Parameters:: seed (int | None)
Return type:: Any

class simulusion.RewardObs(env)[source]

Bases: object

Augment observations with the latest scalar reward under obs[“reward”].

Useful for agents that consume reward as part of the observation stream during model learning or recurrent policy inference.

Parameters:: env (Any)

property observation_space: Dict

step(action)[source]

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], Any, bool, dict[str, Any]]

reset(*, seed=None)[source]

Parameters:: seed (int | None)
Return type:: dict[str, Any]

class simulusion.ResizeImage(env, size=(64, 64))[source]

Bases: object

Resize image-like observation entries to a target spatial size.

The wrapper discovers image keys from env.obs_space, applies nearest neighbor resizing, and updates the advertised observation space shapes.

Parameters:

env (Any)
size (tuple[int, int])

property obs_space: dict[str, Any]

step(action)[source]

Parameters:: action (Any)
Return type:: Any

reset(*, seed=None)[source]

Parameters:: seed (int | None)
Return type:: Any

class simulusion.RenderImage(env, key='image')[source]

Bases: object

Inject RGB renders from env.render(“rgb_array”) into observations.

This is useful when the base environment returns non-image observations but a rendered camera view is needed for world-model training.

Parameters:

env (Any)
key (str)

property obs_space: dict[str, Any]

step(action)[source]

Parameters:: action (Any)
Return type:: Any

reset(*, seed=None)[source]

Parameters:: seed (int | None)
Return type:: Any

class simulusion.SelectAction(env, key)[source]

Bases: Wrapper

Gym wrapper for dictionary actions that forwards a selected key only.

This enables integration with policies that emit action dicts while the environment expects a single tensor/array action payload.

Parameters:

env (Any)
key (str)

step(action)[source]

Parameters:: action (dict[str, Any])
Return type:: Any

simulusion.jepa_linear_probe(checkpoint, root_path, dataset='imagenet', model_name='vit_base', patch_size=16, crop_size=224, weights='target_encoder', train_folder='train', val_folder='val', download=False, batch_size=256, num_workers=8, epochs=50, head_batch_size=16384, reference_lrs=(0.01, 0.05, 0.001), weight_decays=(0.0005, 0.0), representations=(1, 4), device=None, use_bfloat16=False)[source]

Run the paper’s linear evaluation and return every swept result.

Returns a dict with the best top-1 accuracy under "top1" and the full sweep under "sweep". representations lists how many trailing blocks to average-pool and concatenate – the paper tries 1 and 4.

Parameters:

checkpoint (str)
root_path (str)
dataset (str)
model_name (str)
patch_size (int)
crop_size (int)
weights (str)
train_folder (str)
val_folder (str)
download (bool)
batch_size (int)
num_workers (int)
epochs (int)
head_batch_size (int)
reference_lrs (Sequence[float])
weight_decays (Sequence[float])
representations (Sequence[int])
device (str | None)
use_bfloat16 (bool)

Return type:

dict[str, Any]

simulusion.load_jepa_encoder(checkpoint, device, model_name='vit_base', patch_size=16, crop_size=224, weights='target_encoder')[source]

Load a frozen I-JEPA encoder from a training checkpoint.

weights selects which set of encoder weights to evaluate; the paper uses the EMA target_encoder (“We use the target-encoder for evaluation and average pool its output”).

Parameters:

checkpoint (str)
device (device)
model_name (str)
patch_size (int)
crop_size (int)
weights (str)

Return type:

Module

class simulusion.RewardModel(belief_size, state_size, hidden_size, activation_function='relu')[source]

Bases: ExportableAgentMixin, Module

Predict scalar rewards from Dreamer latent belief and state vectors.

Implemented as an MLP used for model-based reward supervision and imagined rollout return estimation.

Parameters:

belief_size (int)
state_size (int)
hidden_size (int)
activation_function (str)

forward(belief, state)[source]

Parameters:

belief (Tensor)
state (Tensor)

Return type:

Tensor

class simulusion.ValueModel(belief_size, state_size, hidden_size, activation_function='relu')[source]

Bases: ExportableAgentMixin, Module

Estimate scalar value from Dreamer latent belief and state vectors.

This MLP is trained on imagined returns and used for actor/value updates.

Parameters:

belief_size (int)
state_size (int)
hidden_size (int)
activation_function (str)

forward(belief, state)[source]

Parameters:

belief (Tensor)
state (Tensor)

Return type:

Tensor

simulusion.DreamerRewardModel: alias of RewardModel

simulusion.DreamerValueModel: alias of ValueModel

simulusion.register_world_model(name, *, import_path, config_path=None, description='', aliases=(), override=False)[source]

Can be used as a decorator or called directly:

@register_world_model("my-model", import_path="my_pkg.model:MyModel")
class MyModel:
    ...

# or without a class:
register_world_model("other-model", import_path="other:factory")

Once registered, the model is available via create_model("my-model") and appears in list_models().

Parameters:

name (str)
import_path (str)
config_path (str | None)
description (str)
aliases (tuple[str, ...])
override (bool)

Return type:

Callable[[type], type] | None

simulusion.deregister_world_model(name)[source]

Remove a previously registered model from the plugin registry.

Parameters:: name (str)
Return type:: None

simulusion.get_registered_model_spec(name)[source]

Look up a registered (non-built-in) model spec by name or alias.

Parameters:: name (str)
Return type:: ModelSpec | None

simulusion.list_registered_models()[source]

Return names of all externally registered world models.

Return type:: list[str]

simulusion.register_env_backend(name, *, factory_path, description='', aliases=(), override=False)[source]

Parameters:

name (str)
factory_path (str)
description (str)
aliases (tuple[str, ...])
override (bool)

Return type:

None

simulusion.deregister_env_backend(name)[source]

Remove a previously registered environment backend.

Parameters:: name (str)
Return type:: None

simulusion.list_registered_env_backends()[source]

Return names of all externally registered env backends.

Return type:: list[str]

simulusion.deprecated(version, reason='', stacklevel=3, category=<class 'DeprecationWarning'>)[source]

Mark a function, method, or class as deprecated.

Usage:

@deprecated(version="0.5.0", reason="Use new_func instead")
def old_func():
    ...

The warning is emitted on every call for functions/methods, or on instantiation for classes.

Parameters:

version (str)
reason (str)
stacklevel (int)
category (type[Warning])

Return type:

Callable[[F], F]

simulusion.deprecated_class(version, reason='', alternative=None)[source]

Shortcut decorator for deprecating a class with an alternative.

Parameters:

version (str)
reason (str)
alternative (str | None)

Return type:

Callable[[type], type]

simulusion.deprecated_function(version, reason='', alternative=None)[source]

Shortcut decorator for deprecating a function with an alternative.

Parameters:

version (str)
reason (str)
alternative (str | None)

Return type:

Callable[[F], F]

class simulusion.Logger(log_dir, enable_wandb=False, wandb_project='simulusion', wandb_entity='', video_format='gif', video_fps=20, enable_tensorboard=False, enable_console=True, enable_jsonl=True, jsonl_filename='metrics.jsonl')[source]

Bases: object

Experiment logger for scalars and GIF rollouts using WandB.

Provides helpers to write scalar metrics, dump pickle snapshots, and save video previews during Dreamer training/evaluation.

Parameters:

log_dir (str)
enable_wandb (bool)
wandb_project (str)
wandb_entity (str)
video_format (str)
video_fps (int)
enable_tensorboard (bool)
enable_console (bool)
enable_jsonl (bool)
jsonl_filename (str)

log_scalar(scalar, name, step_)[source]

Parameters:

scalar (Any)
name (str)
step_ (int)

Return type:

None

log_scalars(scalar_dict, step)[source]

Parameters:

scalar_dict (dict[str, Any])
step (int)

Return type:

None

log_videos(videos, step, max_videos_to_save=1, fps=None, video_title='video')[source]

Parameters:

videos (Any)
step (int)
max_videos_to_save (int)
fps (int | None)
video_title (str)

Return type:

None

dump_scalars_to_pickle(metrics, step, log_title=None)[source]

Parameters:

metrics (dict[str, Any])
step (int)
log_title (str | None)

Return type:

None

flush()[source]

Return type:: None

class simulusion.FreezeParameters(modules)[source]

Bases: object

Context manager that temporarily disables gradients for given modules.

Useful during imagination or target-network forward passes where gradients through certain components should be blocked for speed and correctness.

Parameters:: modules (Iterable[Module])

simulusion.compute_return(rewards, values, discounts, td_lam, last_value)[source]

Compute TD(lambda) returns from imagined rewards, values, and discounts.

Implements backward recursion used by Dreamer actor/value objectives.

Parameters:

rewards (Tensor)
values (Tensor)
discounts (Tensor)
td_lam (float)
last_value (Tensor)

Return type:

Tensor

simulusion.preprocess_obs(obs)[source]

Convert raw uint8 image observations to Dreamer float input space.

Images are scaled from [0, 255] to roughly [-0.5, 0.5], matching the normalization expected by Dreamer encoders.

Parameters:: obs (Tensor)
Return type:: Tensor

class simulusion.ThroughputMeter(device=None, window=50)[source]

Bases: object

Steps per second and bytes-to-device, over a sliding window.

Parameters:

device (torch.device | str | None) – Device the work runs on. When it is CUDA the meter synchronises before reading the clock, because kernel launches are asynchronous and an unsynchronised timer measures queueing rather than execution.
window (int) – Number of recent steps the windowed rate is computed over.

record_transfer(*tensors)[source]

Attribute tensors’ bytes to the step currently in flight.

Parameters:: tensors (Any)
Return type:: None

step()[source]

Close out one step.

Return type:: None

reset()[source]

Discard history, e.g. after warmup and compilation have settled.

Return type:: None

property stats: Dict[str, float]

summary()[source]

One-line human-readable report.

Return type:: str

simulusion.measure_steps(step_fn, iterations, warmup=3, device=None)[source]

Time step_fn over iterations, discarding warmup steps first.

Warmup matters: the first calls pay for lazy CUDA context creation, cuDNN autotuning, and torch.compile tracing, none of which recur.

Parameters:

step_fn (Any)
iterations (int)
warmup (int)
device (Any)

Return type:

Dict[str, float]

simulusion.tensor_nbytes(*tensors)[source]

Total bytes backing tensors, skipping anything that is not a tensor.

element_size() * nelement() rather than untyped_storage().nbytes(): a sliced view would otherwise be charged for the whole base allocation.

Parameters:: tensors (Any)
Return type:: int

simulusion.enable_performance_defaults(*, tf32=True, cudnn_benchmark=True)[source]

Turn on the CUDA throughput settings Simulusion trainers expect.

Call this once at the start of a training run. It is deliberately explicit and not applied on import, because both settings trade something away:

cudnn_benchmark autotunes convolution algorithms on first sight of each input shape. That is a large win for fixed-shape training and a loss for workloads whose shapes keep changing, and the autotuner’s choice is not guaranteed stable run to run.

tf32 lets matmuls and convolutions use TensorFloat-32 on Ampere and later: same exponent range as float32, but a 10-bit mantissa. For model-based RL and vision training this is the standard setting and costs no measurable quality, but it is not bit-reproducible against a float32 baseline - leave it off for numerics regression work.

Both are no-ops without CUDA, so this is safe to call unconditionally.

Parameters:

tf32 (bool)
cudnn_benchmark (bool)

Return type:

None

simulusion.maybe_compile(module_or_fn, *, enabled=False, mode='reduce-overhead')[source]

Optionally wrap a callable in torch.compile.

RSSM rollouts and CEM planning are sequences of very small kernels driven by a Python loop, so they are launch-bound rather than FLOP-bound; compiling the step collapses that overhead. It is off by default because compilation costs seconds to minutes on first call, recompiles whenever an input shape changes, and is unavailable on some builds - none of which a short run or a test wants to pay for.

Falls back to the eager callable if compilation is unsupported, so callers never need to guard. torch.compile is lazy: a missing backend (for example no Triton on a Windows CUDA build) surfaces on the first call, not when wrapping, so the fallback also covers that first call. Once a compiled call has succeeded, later errors propagate unchanged - by then they are genuine errors, not a missing backend.

An nn.Module is compiled by swapping its instance forward, so the module object, its attributes and its state_dict keys are unchanged. (torch.compile(module) would return an OptimizedModule whose keys gain an _orig_mod. prefix, and checkpoints written from it would not load into an uncompiled model.)

Parameters:

module_or_fn (Any)
enabled (bool)
mode (str)

Return type:

Any

simulusion.to_channels_last(model)[source]

Convert a conv-heavy model to NHWC memory format.

This is a layout change only - values are untouched - but it lets cuDNN pick NHWC tensor-core kernels for convolutions instead of transposing on every call. Inputs must be converted to match; mixing layouts silently costs more than it saves.

Parameters:: model (Module)
Return type:: Module

class simulusion.InferenceModel(module, *, precision='fp32', compile=False, compile_mode='reduce-overhead', channels_last=False, cast_weights=False, output_dtype=torch.float32, clone_outputs=None)[source]

Bases: ExportableAgentMixin, Module

A module wrapped for fast, gradient-free inference.

Parameters:

module (nn.Module) – Module to wrap. It is put in eval mode and its parameters stop requiring grad; its weights are only modified if cast_weights.
precision (str | None) – "fp32" (default), "bf16", "fp16" or "auto". Reduced precision runs under autocast, so numerically sensitive ops (softmax, norms, reductions) stay in fp32.
compile (bool) – Wrap the forward in torch.compile. Falls back to eager, with a warning, if the backend is unavailable (e.g. no Triton).
compile_mode (str) – torch.compile mode. "reduce-overhead" (default) uses CUDA graphs, the right choice for small, fixed-shape steps. Use "max-autotune" for large batched steps.
channels_last (bool) – Convert the module (and 4D inputs) to NHWC, which lets cuDNN use tensor-core convolution kernels.
cast_weights (bool) – Store weights in the reduced dtype as well, halving their memory. Inputs are cast to match. Ignored at fp32.
output_dtype (torch.dtype | None) – Cast floating outputs to this dtype (default fp32) so the wrapper is a drop-in replacement. None returns them as computed.
clone_outputs (bool | None) – Clone outputs after each call. CUDA-graph replays reuse their output buffers, so an output kept across calls would otherwise be overwritten by the next step. Defaults to on exactly when compiling with a CUDA-graph mode.

property device: device

forward(*args, **kwargs)[source]

Parameters:

args (Any)
kwargs (Any)

Return type:

Any

simulusion.optimize_for_inference(module, **kwargs)[source]

Wrap module for inference. See InferenceModel for options.

Example:

policy = optimize_for_inference(agent.dreamer.actor, precision="auto",
                                compile=True)
action = policy(features)

Parameters:

module (Module)
kwargs (Any)

Return type:

InferenceModel

simulusion.inference_context(device, precision='fp32')[source]

torch.inference_mode plus autocast at precision.

fp32 enters no autocast region at all, so it is bit-identical to calling the model directly under inference_mode.

Parameters:

device (device | str)
precision (str | None)

Return type:

Iterator[None]

simulusion.make_stepper(model, **kwargs)[source]

Build the stepper matching model’s family.

Parameters:

model (Any)
kwargs (Any)

Return type:

WorldModelStepper

class simulusion.DreamerStepper(agent, *, action_mode='mode')[source]

Bases: object

Stepper for Dreamer (V1) agents.

Parameters:

agent (Any) – A DreamerAgent or the inner Dreamer.
action_mode (str) – How act(explore=False) picks an action. "mode" (default) calls the actor exactly as Dreamer.act_with_world_model does, i.e. a 100-sample Monte Carlo mode estimate. "mean" uses mean_action(): deterministic, one pass, and what step_module() exports.

supports_imagination = True

static features(state)[source]

Parameters:: state (Dict[str, Tensor])
Return type:: Tensor

init_state(batch_size=1)[source]

Parameters:: batch_size (int)
Return type:: Dict[str, Tensor]

observe(state, obs, prev_action, *, noise=None)[source]

Filter a raw [0, 255] image batch (B, C, H, W) into the state.

Parameters:

state (Dict[str, Tensor])
obs (Tensor)
prev_action (Tensor | None)
noise (tuple[Tensor, Tensor] | None)

Return type:

Dict[str, Tensor]

act(state, *, explore=False)[source]

Parameters:

state (Dict[str, Tensor])
explore (bool)

Return type:

Tensor

imagine(state, action, *, noise=None)[source]

Parameters:

state (Dict[str, Tensor])
action (Tensor)
noise (Tensor | None)

Return type:

ImagineOutput

decode(state)[source]

Reconstructed observation, in Dreamer’s [-0.5, 0.5] space.

Parameters:: state (Dict[str, Tensor])
Return type:: Tensor

step_module()[source]

A pure nn.Module for one observe+act step, ready to export.

Return type:: DreamerStepModule

class simulusion.DreamerStepModule(dreamer)[source]

Bases: ExportableAgentMixin, Module

One Dreamer observe+act step as a pure function of tensors.

forward(deter, stoch, prev_action, obs, prior_noise, posterior_noise) returns (deter, stoch, action). Sampling noise is an input, so the graph contains no RNG: it is exportable by torch.export/ONNX, gives identical results eager vs. exported, and the caller controls seeding (pass zeros for the posterior mean). The action is the deterministic tanh(mean) (see ActionDecoder.mean_action()).

Parameters:: dreamer (Any)

forward(deter, stoch, prev_action, obs, prior_noise, posterior_noise)[source]

Parameters:

deter (Tensor)
stoch (Tensor)
prev_action (Tensor)
obs (Tensor)
prior_noise (Tensor)
posterior_noise (Tensor)

Return type:

tuple[Tensor, Tensor, Tensor]

example_inputs(batch_size=1, obs_shape=None)[source]

Zero inputs of the right shapes, e.g. for export or benchmarking.

Parameters:

batch_size (int)
obs_shape (tuple[int, ...] | None)

Return type:

tuple[Tensor, …]

class simulusion.IRISStepper(agent, *, temperature=1.0)[source]

Bases: object

Stepper for IRISAgent acting.

The IRIS policy is a CNN -> LSTM over frames, so observe advances the LSTM on the new frame and caches the action logits, and act samples from them. Imagination in IRIS is a whole-rollout procedure with a KV cache and cache rebuilding (paper 2.3); use IRISAgent.imagine_rollout for it rather than a per-step API, so supports_imagination is False.

Parameters:

agent (Any) – The IRIS agent. It is switched to eval mode.
temperature (float) – Softmax temperature for action sampling; 0 with explore=False takes the argmax.

supports_imagination = False

init_state(batch_size=1)[source]

Parameters:: batch_size (int)
Return type:: Dict[str, Tensor]

observe(state, obs, prev_action)[source]

Advance the policy on a (B, C, H, W) frame batch in [0, 1].

Parameters:

state (Dict[str, Tensor])
obs (Tensor)
prev_action (Tensor | None)

Return type:

Dict[str, Tensor]

act(state, *, explore=False)[source]

Parameters:

state (Dict[str, Tensor])
explore (bool)

Return type:

Tensor

simulusion.benchmark_step(step, *args, warmup=10, iterations=100, batch_size=1, device=None, name='step', **kwargs)[source]

Time step(*args, **kwargs).

Each iteration is timed individually (with a device synchronisation on CUDA, because kernel launches are asynchronous) so tail latency is visible, which matters for acting in a real-time environment. warmup iterations run first and are discarded: they absorb torch.compile compilation, CUDA-graph capture and cuDNN autotuning.

Parameters:

step (Callable[[...], Any]) – Callable to time. For a stateful loop, pass a closure that advances its own state.
warmup (int) – Untimed iterations.
iterations (int) – Timed iterations.
batch_size (int) – Items processed per call, for items_per_sec.
device (device | str | None) – Device to synchronise and read peak memory on. Inferred from the first tensor argument when omitted.
name (str) – Label for the report.
args (Any)
kwargs (Any)

Return type:

LatencyReport

simulusion.rollout_drift(reference, candidate, initial_state, step_inputs)[source]

Run both steps closed-loop from initial_state and compare states.

Each step is called as step(state, *step_inputs[t]) and must return the next state (any pytree of tensors, e.g. a tuple or dict). Each rollout feeds back its own output, so the report shows compounding error, not only one-step error.

Stochastic models must be given their randomness through step_inputs (for example the noise argument of DreamerStepModule): compiled and exported graphs do not consume the global RNG the way eager code does, so seeding alone does not make the two rollouts comparable.

Parameters:

reference (Callable[[...], Any])
candidate (Callable[[...], Any])
initial_state (Any)
step_inputs (Sequence[Sequence[Any]])

Return type:

DriftReport

simulusion.quantize_weights(module, *, skip=('quantizer', 'codebook', 'vq', 'fc_state_prior', 'fc_state_posterior'), min_in_features=64, backend='native')[source]

Quantise nn.Linear weights of module to int8, in place.

Parameters:

module (Module) – Module to quantise. Put it in eval mode first; the result is for inference only.
skip (Iterable[str]) – Substrings of qualified module names to leave in floating point. A layer is skipped if any of its ancestors’ names matches too.
min_in_features (int) – Leave smaller nn.Linear layers alone.
backend (str) – "native" (default) uses Int8WeightOnlyLinear, which needs no extra dependency and exports with torch.export. "torchao" delegates to torchao’s int8 weight-only kernels, which are faster on recent GPUs; it applies the same filter.

Returns:

Qualified names of the layers that were quantised.

Return type:

list[str]

simulusion.save_bundle(directory, module, *, example_inputs=None, formats=('exported_program',), config=None, name='model', metadata=None, verify=True, export_kwargs=None)[source]

Write module as a deployment bundle.

Parameters:

directory (str | Path) – Output directory (created if missing).
module (Module) – Module to bundle - typically a pure step module such as DreamerStepModule.
example_inputs (Any) – Inputs used to export and to record the input spec. Required when formats is non-empty.
formats (Iterable[str]) – Export formats to include (see simulusion.export). Pass () for a weights-and-config bundle.
config (Any) – Model config (dataclass, object or dict), stored as JSON.
name (str) – Base file name of the artifacts.
metadata (dict[str, Any] | None) – Extra JSON-serialisable fields for the manifest.
verify (bool) – Check each artifact against the eager module with verify_export() and record the max error.
export_kwargs (dict[str, Any] | None) – Extra keyword arguments per format, e.g. {"exported_program": {"dynamic_shapes": ...}}.

Return type:

Bundle

simulusion.load_bundle(directory)[source]

Read a bundle written by save_bundle().

Parameters:: directory (str | Path)
Return type:: Bundle

class simulusion.STKVCache(num_layers, batch_size, num_heads, num_patches, head_dim, max_frames, device, dtype)[source]

Bases: object

Per-layer temporal key/value cache for frame-by-frame generation.

In an ST-transformer, spatial attention and the MLP act within a frame and temporal attention is causal, so a frame’s activations never depend on later frames. Generating frame t+1 therefore only needs the temporal keys/values of frames 0..t - it does not have to re-run the network over the whole prefix. With the cache, each new frame costs one frame’s worth of compute instead of t frames’, turning an O(T^2) rollout into O(T).

Storage is a pre-allocated (B, heads, N, max_frames, head_dim) buffer per layer. Writes go to length onwards; advance() commits them. A forward that is not committed can be repeated (e.g. MaskGIT refinement of the same frame) and simply overwrites the same slots.

Parameters:

num_layers (int)
batch_size (int)
num_heads (int)
num_patches (int)
head_dim (int)
max_frames (int)
device (device)
dtype (dtype)

append(layer, k, v)[source]

Write (B, heads, N, T_new, hd) keys/values; return the full prefix.

Parameters:

layer (int)
k (Tensor)
v (Tensor)

Return type:

Tuple[Tensor, Tensor]

advance(frames)[source]

Commit frames newly written frames.

Parameters:: frames (int)
Return type:: None

class simulusion.WorldModel(*args, **kwargs)[source]

Bases: Protocol

The interface every world model in Simulusion implements.

Wrap a trained model with as_world_model(), or load one with load_world_model(). See the module docstring for the conventions.

spec: WorldModelSpec

device: device

encode(frames, actions=None)[source]

Turn real context frames (and the actions between them) into a state.

Parameters:

frames (Any)
actions (Any)

Return type:

Dict[str, Tensor]

step(state, action)[source]

Imagine one step: the state after taking action.

Parameters:

state (Dict[str, Tensor])
action (Any)

Return type:

decode(state)[source]

The current frame of state: float (B, C, H, W) in [0, 1].

Parameters:: state (Dict[str, Tensor])
Return type:: Tensor

rollout(state, actions, *, decode=True)[source]

Imagine H steps; actions is (B, H) or (B, H, A).

Parameters:

state (Dict[str, Tensor])
actions (Any)
decode (bool)

Return type:

save(path)[source]

Write the model to directory path; returns it.

Parameters:: path (str | Path)
Return type:: Path

class simulusion.WorldModelSpec(family, obs_shape, action, context_frames=1, max_context=None, has_reward=True, has_termination=False, stochastic=True)[source]

Bases: object

What a world model consumes and produces.

Variables:

family (str) – "dreamer", "planet", "iris", "diamond", "genie" or "world-models".
obs_shape (tuple[int, int, int]) – (C, H, W) of the frames it reads and decodes.
action (simulusion.world_model.protocol.ActionSpec) – The actions WorldModel.step() accepts.
context_frames (int) – The fewest frames WorldModel.encode() needs.
max_context (int | None) – The most frames the model attends to; older ones are forgotten. None for recurrent models, which carry everything in their state.
has_reward (bool) – Transition carries a predicted reward.
has_termination (bool) – Transition carries a predicted episode end.
stochastic (bool) – WorldModel.step() samples, so repeated calls on the same state differ unless the RNG is reseeded.

Parameters:

family (str)
obs_shape (tuple[int, int, int])
action (ActionSpec)
context_frames (int)
max_context (int | None)
has_reward (bool)
has_termination (bool)
stochastic (bool)

family: str

obs_shape: tuple[int, int, int]

action: ActionSpec

context_frames: int = 1

max_context: int | None = None

has_reward: bool = True

has_termination: bool = False

stochastic: bool = True

to_dict()[source]

Return type:: dict[str, Any]

simulusion.as_world_model(model, *args, **kwargs)[source]

Wrap a trained model of any supported family as a WorldModel.

Accepts the agent or the bare model: a Dreamer/DreamerAgent, Planet or its RSSM, IRISAgent, DiamondAgent, Genie, or a (ConvVAE, MDRNN) pair. Something that already implements WorldModel is returned as is. Extra arguments go to the family’s constructor, e.g. temperature=0 for Genie.

Parameters:

model (Any)
args (Any)
kwargs (Any)

Return type:

simulusion.load_world_model(path, family=None, *, device=None, **options)[source]

Load a world model saved by WorldModel.save(), or a training checkpoint.

A directory written by save names its family in world_model.json. Anything else is read as family’s own checkpoint (a file, a run directory, or for most families a Hugging Face Hub repo id), so family is required then. options override the saved constructor options, e.g. temperature=0.5.

Parameters:

path (str | Path)
family (str | None)
device (str | device | None)
options (Any)

Return type:

simulusion.list_world_model_families()[source]

The family names load_world_model() and world_model_class() accept.

Return type:: list[str]

class simulusion.Trainer(config=None, *, run_config=None, config_filename='config.yaml', metrics_filename='metrics.jsonl')[source]

Bases: object

Runs a TrainingTask to its budget.

Parameters:

config (Optional[TrainerConfig]) – Run settings; defaults throughout when None.
run_config (Optional[Mapping[str, Any]]) – The full run config (all sections) to write as config.yaml. simulusion.training.run.run() passes it; without it only the trainer and task sections are written.
config_filename (str) – Name of the resolved run config in the run directory.
metrics_filename (str) – Name of the metrics log in the run directory. The older entrypoints (simulusion.training.legacy) rename both so they never overwrite a family’s own config.yaml or metrics.jsonl in the folder it has always written to.

prepare_module(module)[source]

module on the run’s device.

Parameters:: module (Module)
Return type:: Module

prepare_dataloader(loader)[source]

loader with every batch’s tensors moved to the run’s device.

Parameters:: loader (Iterable[Any])
Return type:: DeviceLoader

gradient_update(optimizer, parameters, *, max_grad_norm=None, accumulate=True)[source]

A GradientUpdate at this run’s precision and accumulation.

Parameters:

optimizer (Optimizer)
parameters (Iterable[Tensor])
max_grad_norm (float | None)
accumulate (bool)

Return type:

GradientUpdate

static infinite(loader)[source]

Cycle loader forever, re-shuffling each pass.

Parameters:: loader (Iterable[Any])
Return type:: Iterator[Any]

fit(task)[source]

Train task until its budget is spent or early stopping fires.

Parameters:: task (TrainingTask)
Return type:: TrainResult

save_checkpoint(*, final=False, name=None)[source]

Write a resumable checkpoint of the task and the loop; returns its dir.

name writes it to checkpoints/<name> instead of checkpoints/step_<progress>, outside the pruned periodic ones.

Parameters:

final (bool)
name (str | None)

Return type:

Path

class simulusion.TrainerConfig(max_steps=None, seed=0, deterministic=False, precision='fp32', device=None, grad_accumulation=1, log_every=100, eval_every=0, checkpoint_every=0, keep_checkpoints=2, out_dir='runs', run_name=None, resume=None, export_world_model=True, best_metric=None, best_mode='max', early_stopping_metric=None, early_stopping_mode='min', patience=10, min_delta=0.0001, console=True)[source]

Bases: SerializableConfigMixin

Settings shared by every training run, whatever the family.

Variables:

max_steps (int | None) – Budget in the task’s progress unit (optimizer steps, environment steps or epochs; see the task). None uses the task’s default.
seed (int) – Seeds Python, NumPy and PyTorch.
deterministic (bool) – Ask PyTorch for deterministic kernels. Slower; some ops only warn.
precision (str) – fp32, bf16, fp16 or auto.
device (str | None) – None (or "auto") picks CUDA, then MPS, else CPU; "cpu", "cuda", "cuda:1" or "mps" choose one. A device that is not available falls back with a warning.
grad_accumulation (int) – Batches per optimizer step, for tasks whose updates support it (dataset-driven ones).
log_every (int) – Log metrics every this many units of progress.
eval_every (int) – Evaluate every this many units; 0 only at the end.
checkpoint_every (int) – Checkpoint every this many units; 0 only at the end.
keep_checkpoints (int) – Periodic checkpoints kept on disk (oldest dropped). The final one is always kept.
out_dir (str) – Parent directory of run directories.
run_name (str | None) – Run directory name; a timestamped one when None.
resume (str | None) – A checkpoint directory to resume from, or "auto" for the newest complete checkpoint of out_dir/run_name.
export_world_model (bool) – Write the task’s world model to <run>/world_model when training ends.
best_metric (str | None) – Evaluation metric whose best value so far is kept as a checkpoint in checkpoints/best (rewritten on every improvement, never pruned); None keeps none.
best_mode (str) – "max" or "min": which direction is better.
early_stopping_metric (str | None) – Evaluation metric that stops the run once it stops improving; None disables early stopping.
early_stopping_mode (str) – "min" or "max".
patience (int) – Evaluations without improvement before stopping.
min_delta (float) – Relative improvement that counts.
console (bool) – Print metrics to the console.

Parameters:

max_steps (int | None)
seed (int)
deterministic (bool)
precision (str)
device (str | None)
grad_accumulation (int)
log_every (int)
eval_every (int)
checkpoint_every (int)
keep_checkpoints (int)
out_dir (str)
run_name (str | None)
resume (str | None)
export_world_model (bool)
best_metric (str | None)
best_mode (str)
early_stopping_metric (str | None)
early_stopping_mode (str)
patience (int)
min_delta (float)
console (bool)

max_steps: int | None = None

seed: int = 0

deterministic: bool = False

precision: str = 'fp32'

device: str | None = None

grad_accumulation: int = 1

log_every: int = 100

eval_every: int = 0

checkpoint_every: int = 0

keep_checkpoints: int = 2

out_dir: str = 'runs'

run_name: str | None = None

resume: str | None = None

export_world_model: bool = True

best_metric: str | None = None

best_mode: str = 'max'

early_stopping_metric: str | None = None

early_stopping_mode: str = 'min'

patience: int = 10

min_delta: float = 0.0001

console: bool = True

class simulusion.TrainResult(run_dir, progress, iterations, metrics=<factory>, checkpoint=None, world_model_path=None, model_path=None, best_checkpoint=None, best_value=None, stopped_early=False)[source]

Bases: object

What Trainer.fit() returns.

Parameters:

run_dir (Path)
progress (int)
iterations (int)
metrics (dict[str, float])
checkpoint (Path | None)
world_model_path (Path | None)
model_path (Path | None)
best_checkpoint (Path | None)
best_value (float | None)
stopped_early (bool)

run_dir: Path

progress: int

iterations: int

metrics: dict[str, float]

checkpoint: Path | None = None

world_model_path: Path | None = None

model_path: Path | None = None

best_checkpoint: Path | None = None

best_value: float | None = None

stopped_early: bool = False

class simulusion.TrainingTask(config=None)[source]

Bases: ABC

One family’s training, run by Trainer.

Subclasses set config_cls (a dataclass whose fields are the run config’s sections besides trainer) and implement setup() and train_step().

Variables:

name (ClassVar[str]) – Registry name, set by register_task().
config_cls (ClassVar[type]) – The task’s config dataclass.
progress_unit (ClassVar[str]) – What progress() counts, for logs: "steps" (calls to train_step()), "env_steps" or "epochs".
default_max_steps (ClassVar[int]) – Budget used when the trainer config gives none.

Parameters:

config (Any)

name: ClassVar[str] = ''

config_cls: ClassVar[type]

progress_unit: ClassVar[str] = 'steps'

default_max_steps: ClassVar[int] = 1000

property trainer: Trainer

attach(trainer)[source]

Parameters:: trainer (Trainer)
Return type:: None

max_steps()[source]

The budget, in progress_unit, when the trainer sets none.

Return type:: int

abstractmethod setup(trainer)[source]

Build networks, optimizers and data on trainer.device.

Parameters:: trainer (Trainer)
Return type:: None

abstractmethod train_step()[source]

One unit of training; returns scalar metrics.

Return type:: Mapping[str, float]

progress()[source]

How far training is, in progress_unit.

Defaults to the number of train_step() calls. Tasks whose budget is counted in environment steps or epochs override it.

Return type:: int

evaluate()[source]

Periodic evaluation metrics (validation loss, episode return).

Return type:: Mapping[str, float]

final_evaluation()[source]

Metrics computed once, after the last step. Defaults to evaluate().

Return type:: Mapping[str, float]

world_model()[source]

The trained world model behind the shared protocol, if the task has one.

Return type:: Optional[‘WorldModel’]

export(directory)[source]

Write the trained model in its own reusable form; return directory.

For models that are not world models (DiT, I-JEPA, a VAE), which world_model() cannot return. Called once, when training ends, with <run_dir>/model. The default writes nothing.

Parameters:: directory (Path)
Return type:: Path | None

state_dict()[source]

Everything needed to resume: weights, optimizer and scaler states.

Return type:: dict[str, Any]

load_state_dict(state)[source]

Parameters:: state (dict[str, Any])
Return type:: None

save_checkpoint(directory)[source]

Write the task’s resumable state into directory.

The default saves state_dict() to task.pt. Tasks with state too large or too odd for one file (a replay buffer) extend it.

Parameters:: directory (Path)
Return type:: None

load_checkpoint(directory)[source]

Parameters:: directory (Path)
Return type:: None

should_stop()[source]

True to end training now, before the budget (a search that converged).

Checked after every step. The final evaluation, checkpoint and export still run.

Return type:: bool

close()[source]

Release environments and other resources.

Return type:: None

simulusion.register_task(name, target=None)[source]

Use as a decorator on the class, or call with target="module:Class" to register lazily without importing it.

Parameters:

name (str)
target (str | None)

Return type:

Callable[[type[TrainingTask]], type[TrainingTask]]

simulusion.list_tasks()[source]

Return type:: list[str]

simulusion.run_training(config=None, overrides=(), *, task=None)

Train the task a run config describes; see the module docstring.

Parameters:

config (str | Path | Mapping[str, Any] | None)
overrides (Sequence[str] | Mapping[str, Any])
task (str | None)

Return type:

TrainResult

class simulusion.MixedPrecision(precision, device, *, scale_loss=None)[source]

Bases: object

Autocast dtype and loss scaling for one training run.

Parameters:

precision (str | None) – "fp32", "bf16", "fp16" or "auto" (bf16 on CUDA devices that support it, fp16 on older ones, fp32 elsewhere).
device (torch.device | str) – The device the run trains on.
scale_loss (Optional[bool])

fp32 enters no autocast region, so it is bit-identical to running without one. Only fp16 scales the loss: bf16 has fp32’s dynamic range.

classmethod from_use_amp(use_amp, device)[source]

The precision the per-family use_amp flags used to select.

use_amp meant autocast at the device’s default dtype: fp16 on CUDA (with a scaler), bf16 on CPU and fp16 on MPS, both without one: the GradScaler("cuda") they built disabled itself off CUDA.

Parameters:

use_amp (bool)
device (device | str)

Return type:

MixedPrecision

property enabled: bool

property scales_loss: bool

autocast(enabled=True)[source]

Autocast at this precision; enabled=False runs a block in fp32.

Parameters:: enabled (bool)
Return type:: Iterator[None]

grad_scaler()[source]

A fresh loss scaler, enabled only for fp16.

Return type:: GradScaler

class simulusion.GradientUpdate(optimizer, parameters, precision, *, max_grad_norm=None, accumulation_steps=1)[source]

Bases: object

One optimizer’s step: zero, backward, unscale, clip, step.

Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to step.
parameters (Iterable[torch.Tensor]) – The parameters to clip; normally the ones optimizer owns.
precision (MixedPrecision) – Decides whether the loss is scaled.
max_grad_norm (Optional[float]) – Clip threshold, or None not to clip. The norm is measured after unscaling, so the threshold means the same thing at every precision.
accumulation_steps (int) – Optimizer steps happen every this many calls; the calls in between only add gradients. Each loss is divided by it, so the step sees the mean over the accumulated batches.

Calling it returns the gradient norm before clipping (a tensor), or None on a call that only accumulated. A non-finite norm under fp16 means the scaler skipped the step.

state_dict()[source]

Return type:: dict[str, Any]

load_state_dict(state)[source]

Parameters:: state (dict[str, Any])
Return type:: None

class simulusion.EpisodeFolder(root, *, episodes=None)[source]

Bases: Sequence[Episode]

Read a folder of episodes; folder[i] loads episode i.

Parameters:

root (str | Path) – The folder.
episodes (Optional[Iterable[int]]) – Indices of the episodes to keep (after sorting by path), for train/validation splits.

property lengths: List[int]: Frames per episode, read without loading the frames where possible.

split(val_fraction, seed=0)[source]

Disjoint train / validation folders, split by whole episodes.

Parameters:

val_fraction (float)
seed (int)

Return type:

tuple[EpisodeFolder, EpisodeFolder]

class simulusion.EpisodeWriter(root, *, meta=None, compress=True)[source]

Bases: object

Append episodes to a folder, numbering after any already there.

with EpisodeWriter("data/pusht", meta=DatasetMeta(fps=10)) as writer:
    for episode in episodes:
        writer.write(episode)

Parameters:

root (str | Path)
meta (Optional[DatasetMeta])
compress (bool)

write(episode)[source]

Parameters:: episode (Episode)
Return type:: Path

close()[source]

Return type:: None

class simulusion.LeRobotEpisodes(source, *, camera=None, extra_cameras=(), episodes=None)[source]

Bases: Sequence[Episode]

The episodes of a LeRobot dataset; dataset[i] decodes episode i.

Parameters:

source (str | Path) – A local dataset folder, the one holding meta/info.json.
camera (Optional[str]) – Feature key of the view to use as frames ("observation.images.top", or just "top"). Defaults to the first camera in name order.
extra_cameras (str | Sequence[str]) – Further views to keep as cameras: "all", or a list of keys. () keeps none (less decoding).
episodes (Optional[Iterable[int]]) – Episode indices to read; all when None.

property lengths: List[int]

property meta: DatasetMeta: The episode-folder metadata this dataset converts to.

class simulusion.CEMPlanner(horizon=12, candidates=500, iterations=5, elites=50, discount=0.99, init_std=1.0, min_std=0.05, smoothing=0.0)[source]

Bases: object

Cross-entropy-method planning over a world model’s imagined returns.

Each of iterations rounds samples candidates action sequences of length horizon, imagines them all in one batched rollout, and refits the sampling distribution to the elites with the highest discounted predicted return. Continuous actions are Gaussian (clipped to the action bounds); discrete ones are a categorical per step. Imagined returns stop accumulating once the model predicts the episode ended.

Variables:

horizon (int) – Steps per imagined sequence.
candidates (int) – Sequences per round.
iterations (int) – Refitting rounds.
elites (int) – Best sequences each refit uses.
discount (float) – Per-step discount of imagined rewards.
init_std (float) – Initial standard deviation of continuous actions.
min_std (float) – Floor on the refitted standard deviation.
smoothing (float) – Weight of the old distribution in each refit (0 = none).

Parameters:

horizon (int)
candidates (int)
iterations (int)
elites (int)
discount (float)
init_std (float)
min_std (float)
smoothing (float)

horizon: int = 12

candidates: int = 500

iterations: int = 5

elites: int = 50

discount: float = 0.99

init_std: float = 1.0

min_std: float = 0.05

smoothing: float = 0.0

plan(model, state)[source]

The first action of the best plan from state (a batch of one).

Returns (A,) float for continuous actions, a scalar int64 for discrete ones.

Parameters:

model (WorldModel)
state (Dict[str, Tensor])

Return type:

Tensor

class simulusion.PlanningAgent(model, planner=None, *, context=None)[source]

Bases: object

An agent that keeps recent frames and plans each action in the world model.

Every act() re-encodes the last context frames and the actions between them (the protocol’s encode()), then plans from that state. Recurrent models (Dreamer, PlaNet, V + M) carry the whole history in their state, so a short context is only an approximation for them; models with a fixed window (IRIS, DIAMOND, Genie) see exactly what they were trained to see.

Parameters:

model (WorldModel) – Any world model with a reward head.
planner (Optional[CEMPlanner]) – How actions are chosen; CEMPlanner by default.
context (Optional[int]) – Frames kept; the model’s max_context (or its context_frames when that is unbounded), at least 1.

frames: Deque[ndarray]

actions: Deque[Any]

reset(observation)[source]

Start an episode from its first observation.

Parameters:: observation (Any)
Return type:: None

act()[source]

The action to take now.

Return type:: ndarray

observe(action, observation)[source]

Record the action just taken and the observation it led to.

Parameters:

action (Any)
observation (Any)

Return type:

None

User-facing convenience APIs for Simulusion.

The lower-level modules remain available for research workflows, but this module collects the common discovery and construction paths behind small, predictable factory functions.

class simulusion.api.EnvBackendSpec(name, factory_path, description='', aliases=())[source]#

Bases: NamedTuple

Metadata describing an environment backend available through make_env.

Parameters:

name (str)
factory_path (str)
description (str)
aliases (tuple[str, ...])

name: str#: Alias for field number 0

factory_path: str#: Alias for field number 1

description: str#: Alias for field number 2

aliases: tuple[str, ...]#: Alias for field number 3

class simulusion.api.ModelSpec(name, import_path, config_path=None, description='', aliases=())[source]#

Bases: NamedTuple

Metadata describing a model available through create_model().

Parameters:

name (str)
import_path (str)
config_path (str | None)
description (str)
aliases (tuple[str, ...])

name: str#: Alias for field number 0

import_path: str#: Alias for field number 1

config_path: str | None#: Alias for field number 2

description: str#: Alias for field number 3

aliases: tuple[str, ...]#: Alias for field number 4

simulusion.api.create_config(model, **overrides)[source]#

Create the default config object for model and apply overrides.

Examples

>>> cfg = create_config("dreamer", env="walker-walk", seed=7)
>>> cfg.env
'walker-walk'

Parameters:

model (str)
overrides (Any)

Return type:

Any

simulusion.api.create_model(model, config=None, **overrides)[source]#

Instantiate a model or agent from a simple string name.

config is optional for models that define a config class. Keyword overrides are applied to the config when possible, otherwise they are passed directly to the underlying constructor/factory.

Examples

>>> agent = create_model(
...     "dreamer", env="Pendulum-v1", env_backend="gym", total_steps=1000
... )
>>> genie = create_model("genie-small", image_size=32)

Parameters:

model (str)
config (Any | None)
overrides (Any)

Return type:

Any

simulusion.api.get_env_backend_spec(name)[source]#

Return metadata for an environment backend name or alias.

Parameters:: name (str)
Return type:: EnvBackendSpec

simulusion.api.get_model_spec(name)[source]#

Return metadata for a model name or alias.

Parameters:: name (str)
Return type:: ModelSpec

simulusion.api.list_env_backends()[source]#

Return canonical backend names accepted by make_env().

Return type:: list[str]

simulusion.api.list_envs(model=None)[source]#

List known environment ids, optionally filtered by model family.

Parameters:: model (str | None)
Return type:: list[str] | dict[str, list[str]]

simulusion.api.list_models()[source]#

Return canonical model names accepted by create_model().

Return type:: list[str]

simulusion.api.make_env(env_id, backend='auto', **kwargs)[source]#

Create an environment with a consistent Simulusion entry point.

Parameters:

env_id (str) – Environment id, XML path, Unity executable path, or backend-specific id.
backend (str) – One of list_env_backends(); "auto" tries Simulusion’s compatibility helper.
**kwargs (Any) – Backend-specific options.

Return type:

Any

Export utilities for production deployment.

The public entry points are export_model() / export_any() and the obj.export(path, format=...) method. Supported formats:

"exported_program" - torch.export graph saved as .pt2. The recommended format: loadable without the model’s source code, and the input to AOTInductor, ExecuTorch and TensorRT.
"aoti" - AOTInductor package (.pt2): an ahead-of-time compiled shared library runnable from Python or C++ without the model’s Python code.
"onnx" - for ONNX Runtime / TensorRT / edge runtimes.
"tensorrt" - Torch-TensorRT, compiled through the dynamo IR by default.
"torchscript" - legacy; TorchScript is in maintenance mode upstream.

load_exported() loads any of these back and verify_export() checks an artifact against its eager module.

Simulusion’s public models and agents get .export() from ExportableAgentMixin. Any other nn.Module goes through export_model(); Simulusion does not modify torch.nn.Module.

class simulusion.export.DreamerPolicyExport(actor)[source]#

Bases: Module

Traceable Dreamer policy head used by the generic export resolver.

Parameters:: actor (nn.Module)

forward(features)[source]#

Parameters:: features (Tensor)
Return type:: Tensor

class simulusion.export.ExportableAgentMixin[source]#

Bases: object

Adds .export() to an agent or nn.Module.

Delegates to export_any(), which resolves the deployable module (the object itself, or target / a preferred submodule such as policy for agent wrappers) and infers example inputs where it can.

export(path, format='onnx', *, example_inputs=None, target=None, input_names=None, output_names=None, dynamic_axes=None, opset_version=17, **kwargs)[source]#

Export this agent or one of its contained modules for deployment.

Parameters:

path (str | Path)
format (str)
example_inputs (Any | None)
target (str | None)
input_names (list[str] | None)
output_names (list[str] | None)
dynamic_axes (dict[str, dict[int, str]] | None)
opset_version (int)
kwargs (Any)

Return type:

Path

class simulusion.export.IRISActorCriticExport(agent)[source]#

Bases: Module

Traceable IRIS policy/value head used by the generic export resolver.

Parameters:: agent (Any)

forward(frames)[source]#

Parameters:: frames (Tensor)
Return type:: tuple[Tensor, Tensor]

simulusion.export.export_any(obj, path, format='onnx', *, example_inputs=None, target=None, input_names=None, output_names=None, dynamic_axes=None, opset_version=17, **kwargs)[source]#

Export any Simulusion model/agent or a target module contained by it.

Parameters:

obj (Any)
path (str | Path)
format (str)
example_inputs (Any | None)
target (str | None)
input_names (list[str] | None)
output_names (list[str] | None)
dynamic_axes (dict[str, dict[int, str]] | None)
opset_version (int)
kwargs (Any)

Return type:

Path

simulusion.export.export_model(module, path, format='onnx', *, example_inputs=None, input_names=None, output_names=None, dynamic_axes=None, opset_version=17, **kwargs)[source]#

Export a torch.nn.Module to ONNX, TorchScript, or TensorRT.

Parameters:

module (Module)
path (str | Path)
format (str)
example_inputs (Any | None)
input_names (list[str] | None)
output_names (list[str] | None)
dynamic_axes (dict[str, dict[int, str]] | None)
opset_version (int)
kwargs (Any)

Return type:

Path

simulusion.export.load_exported(path, format=None, *, device=None)[source]#

Load an exported artifact back as a callable.

Parameters:

path (str | Path) – Artifact written by export_model().
format (str | None) – Its export format. Inferred from the suffix when omitted, with .pt2 read as exported_program; pass "aoti" for AOTInductor packages, which share the suffix.
device (device | str | None) – Device to move an ExportedProgram or TorchScript module to.

Return type:

Callable[[…], Any]

ONNX artifacts run through ONNX Runtime (optional dependency); the returned callable takes and returns tensors like the other formats.

simulusion.export.verify_export(module, exported, example_inputs, *, format=None, atol=0.0001, rtol=0.0001)[source]#

Check that an exported artifact reproduces module on example_inputs.

Parameters:

module (Module)
exported (str | Path | Callable[[...], Any])
example_inputs (Any)
format (str | None)
atol (float)
rtol (float)

Return type:

float

In-process inference optimisation for any Simulusion module.

World-model inference is a loop of small steps (an RSSM transition, one token, one denoising pass). Each step does little arithmetic, so on a GPU its cost is dominated by Python and kernel-launch overhead rather than FLOPs. The two levers that matter most are therefore:

torch.compile(mode="reduce-overhead"), which fuses kernels and replays the step as a CUDA graph - one launch instead of hundreds;
reduced precision (bf16/fp16), which halves memory traffic.

optimize_for_inference() applies both behind one wrapper, and every setting is opt-in: with the defaults the wrapped module is numerically identical to calling it under torch.inference_mode.

class simulusion.inference.runtime.InferenceModel(module, *, precision='fp32', compile=False, compile_mode='reduce-overhead', channels_last=False, cast_weights=False, output_dtype=torch.float32, clone_outputs=None)[source]#

Bases: ExportableAgentMixin, Module

A module wrapped for fast, gradient-free inference.

Parameters:

module (nn.Module) – Module to wrap. It is put in eval mode and its parameters stop requiring grad; its weights are only modified if cast_weights.
precision (str | None) – "fp32" (default), "bf16", "fp16" or "auto". Reduced precision runs under autocast, so numerically sensitive ops (softmax, norms, reductions) stay in fp32.
compile (bool) – Wrap the forward in torch.compile. Falls back to eager, with a warning, if the backend is unavailable (e.g. no Triton).
compile_mode (str) – torch.compile mode. "reduce-overhead" (default) uses CUDA graphs, the right choice for small, fixed-shape steps. Use "max-autotune" for large batched steps.
channels_last (bool) – Convert the module (and 4D inputs) to NHWC, which lets cuDNN use tensor-core convolution kernels.
cast_weights (bool) – Store weights in the reduced dtype as well, halving their memory. Inputs are cast to match. Ignored at fp32.
output_dtype (torch.dtype | None) – Cast floating outputs to this dtype (default fp32) so the wrapper is a drop-in replacement. None returns them as computed.
clone_outputs (bool | None) – Clone outputs after each call. CUDA-graph replays reuse their output buffers, so an output kept across calls would otherwise be overwritten by the next step. Defaults to on exactly when compiling with a CUDA-graph mode.

property device: device#

forward(*args, **kwargs)[source]#

Parameters:

args (Any)
kwargs (Any)

Return type:

Any

simulusion.inference.runtime.optimize_for_inference(module, **kwargs)[source]#

Wrap module for inference. See InferenceModel for options.

Example:

policy = optimize_for_inference(agent.dreamer.actor, precision="auto",
                                compile=True)
action = policy(features)

Parameters:

module (Module)
kwargs (Any)

Return type:

InferenceModel

One precision policy for inference, instead of a flag per model.

Models in Simulusion grew their own switches (use_amp, use_bfloat16, …). At inference time the question is always the same - which dtype should matmuls and convolutions run in on this device - so it is answered once here.

simulusion.inference.precision.inference_context(device, precision='fp32')[source]#

torch.inference_mode plus autocast at precision.

fp32 enters no autocast region at all, so it is bit-identical to calling the model directly under inference_mode.

Parameters:

device (device | str)
precision (str | None)

Return type:

Iterator[None]

simulusion.inference.precision.precision_dtype(precision, device)[source]#

The torch.dtype that precision resolves to on device.

Parameters:

precision (str | None)
device (device | str)

Return type:

dtype

simulusion.inference.precision.resolve_precision(precision, device)[source]#

Normalise a precision name and resolve "auto" for device.

auto picks bf16 on CUDA devices that support it (Ampere and later), fp16 on older CUDA devices, and fp32 elsewhere. On CPU, reduced precision is usually slower unless the CPU has AMX/AVX512-BF16, so auto stays at fp32 and bf16 must be requested explicitly.

Parameters:

precision (str | None)
device (device | str)

Return type:

str

A uniform, stateful step interface over Simulusion world models.

Deployed world models run as a loop: filter the latest real observation into a latent state, pick an action from that state, and - for planning or dreaming - roll the state forward without observations. Each model family spells these steps differently. A stepper gives them one shape:

state = stepper.init_state(batch_size)
while running:
    state = stepper.observe(state, obs, prev_action)   # real observation
    action = stepper.act(state)
    step = stepper.imagine(state, action)               # optional
    ...

State is always a flat dict[str, Tensor] with the batch on dim 0. Keeping it explicit, rather than hidden inside the model, is what makes the loop batchable across environments, compilable (the step is a pure function of its inputs), and exportable: see DreamerStepper.step_module().

Steppers never change a model’s numerics: with default options every method calls the model’s own code path.

class simulusion.inference.steppers.DreamerStepModule(dreamer)[source]#

Bases: ExportableAgentMixin, Module

One Dreamer observe+act step as a pure function of tensors.

Parameters:: dreamer (Any)

forward(deter, stoch, prev_action, obs, prior_noise, posterior_noise)[source]#

Parameters:

deter (Tensor)
stoch (Tensor)
prev_action (Tensor)
obs (Tensor)
prior_noise (Tensor)
posterior_noise (Tensor)

Return type:

tuple[Tensor, Tensor, Tensor]

example_inputs(batch_size=1, obs_shape=None)[source]#

Zero inputs of the right shapes, e.g. for export or benchmarking.

Parameters:

batch_size (int)
obs_shape (tuple[int, ...] | None)

Return type:

tuple[Tensor, …]

class simulusion.inference.steppers.DreamerStepper(agent, *, action_mode='mode')[source]#

Bases: object

Stepper for Dreamer (V1) agents.

Parameters:

agent (Any) – A DreamerAgent or the inner Dreamer.
action_mode (str) – How act(explore=False) picks an action. "mode" (default) calls the actor exactly as Dreamer.act_with_world_model does, i.e. a 100-sample Monte Carlo mode estimate. "mean" uses mean_action(): deterministic, one pass, and what step_module() exports.

supports_imagination = True#

static features(state)[source]#

Parameters:: state (Dict[str, Tensor])
Return type:: Tensor

init_state(batch_size=1)[source]#

Parameters:: batch_size (int)
Return type:: Dict[str, Tensor]

observe(state, obs, prev_action, *, noise=None)[source]#

Filter a raw [0, 255] image batch (B, C, H, W) into the state.

Parameters:

state (Dict[str, Tensor])
obs (Tensor)
prev_action (Tensor | None)
noise (tuple[Tensor, Tensor] | None)

Return type:

Dict[str, Tensor]

act(state, *, explore=False)[source]#

Parameters:

state (Dict[str, Tensor])
explore (bool)

Return type:

Tensor

imagine(state, action, *, noise=None)[source]#

Parameters:

state (Dict[str, Tensor])
action (Tensor)
noise (Tensor | None)

Return type:

ImagineOutput

decode(state)[source]#

Reconstructed observation, in Dreamer’s [-0.5, 0.5] space.

Parameters:: state (Dict[str, Tensor])
Return type:: Tensor

step_module()[source]#

A pure nn.Module for one observe+act step, ready to export.

Return type:: DreamerStepModule

class simulusion.inference.steppers.IRISStepper(agent, *, temperature=1.0)[source]#

Bases: object

Stepper for IRISAgent acting.

Parameters:

agent (Any) – The IRIS agent. It is switched to eval mode.
temperature (float) – Softmax temperature for action sampling; 0 with explore=False takes the argmax.

supports_imagination = False#

init_state(batch_size=1)[source]#

Parameters:: batch_size (int)
Return type:: Dict[str, Tensor]

observe(state, obs, prev_action)[source]#

Advance the policy on a (B, C, H, W) frame batch in [0, 1].

Parameters:

state (Dict[str, Tensor])
obs (Tensor)
prev_action (Tensor | None)

Return type:

Dict[str, Tensor]

act(state, *, explore=False)[source]#

Parameters:

state (Dict[str, Tensor])
explore (bool)

Return type:

Tensor

class simulusion.inference.steppers.ImagineOutput(state, reward, continue_prob=None)[source]#

Bases: object

Result of one imagined (observation-free) step.

Parameters:

state (Dict[str, Tensor])
reward (Tensor)
continue_prob (Tensor | None)

state: Dict[str, Tensor]#

reward: Tensor#

continue_prob: Tensor | None = None#

class simulusion.inference.steppers.WorldModelStepper(*args, **kwargs)[source]#

Bases: Protocol

Protocol implemented by every stepper.

imagine and decode are optional; supports_imagination says whether imagine is available.

device: device#

supports_imagination: bool#

init_state(batch_size=1)[source]#

Parameters:: batch_size (int)
Return type:: Dict[str, Tensor]

observe(state, obs, prev_action)[source]#

Parameters:

state (Dict[str, Tensor])
obs (Tensor)
prev_action (Tensor | None)

Return type:

Dict[str, Tensor]

act(state, *, explore=False)[source]#

Parameters:

state (Dict[str, Tensor])
explore (bool)

Return type:

Tensor

simulusion.inference.steppers.make_stepper(model, **kwargs)[source]#

Build the stepper matching model’s family.

Parameters:

model (Any)
kwargs (Any)

Return type:

WorldModelStepper

Measure inference speed and what an optimisation costs in fidelity.

Every inference optimisation trades something. Compilation and CUDA graphs are exact up to kernel-level float reordering; bf16, quantisation and fewer sampling steps are not. For a world model a tiny per-step error is not the whole story, because the state is fed back into itself and errors compound over a rollout. So this module reports two things:

benchmark_step() - latency percentiles, throughput and peak memory of a step callable, with proper warm-up and CUDA synchronisation;
rollout_drift() - how far a candidate step (compiled, bf16, quantised, exported, …) drifts from a reference step over an N-step closed-loop rollout.

class simulusion.inference.benchmark.DriftReport(steps, max_abs_error=<factory>, rel_error=<factory>)[source]#

Bases: object

Per-step divergence of a candidate rollout from a reference rollout.

max_abs_error[t] is the largest absolute difference over every tensor in the state after step t; rel_error[t] is the L2 norm of the difference divided by the L2 norm of the reference.

Parameters:

steps (int)
max_abs_error (list[float])
rel_error (list[float])

steps: int#

max_abs_error: list[float]#

rel_error: list[float]#

property final_max_abs_error: float#

property worst_rel_error: float#

within(*, atol=inf, rtol=inf)[source]#

True if every step stays within atol absolute and rtol relative error.

Parameters:

atol (float)
rtol (float)

Return type:

bool

as_dict()[source]#

Return type:: dict[str, Any]

summary()[source]#

Return type:: str

class simulusion.inference.benchmark.LatencyReport(name, iterations, batch_size, mean_ms, p50_ms, p90_ms, p99_ms, min_ms, steps_per_sec, items_per_sec, peak_memory_mb=None)[source]#

Bases: object

Timing summary for one benchmarked callable. Times are milliseconds.

Parameters:

name (str)
iterations (int)
batch_size (int)
mean_ms (float)
p50_ms (float)
p90_ms (float)
p99_ms (float)
min_ms (float)
steps_per_sec (float)
items_per_sec (float)
peak_memory_mb (float | None)

name: str#

iterations: int#

batch_size: int#

mean_ms: float#

p50_ms: float#

p90_ms: float#

p99_ms: float#

min_ms: float#

steps_per_sec: float#

items_per_sec: float#

peak_memory_mb: float | None = None#

as_dict()[source]#

Return type:: dict[str, Any]

summary()[source]#

Return type:: str

simulusion.inference.benchmark.benchmark_step(step, *args, warmup=10, iterations=100, batch_size=1, device=None, name='step', **kwargs)[source]#

Time step(*args, **kwargs).

Parameters:

step (Callable[[...], Any]) – Callable to time. For a stateful loop, pass a closure that advances its own state.
warmup (int) – Untimed iterations.
iterations (int) – Timed iterations.
batch_size (int) – Items processed per call, for items_per_sec.
device (device | str | None) – Device to synchronise and read peak memory on. Inferred from the first tensor argument when omitted.
name (str) – Label for the report.
args (Any)
kwargs (Any)

Return type:

LatencyReport

simulusion.inference.benchmark.rollout_drift(reference, candidate, initial_state, step_inputs)[source]#

Run both steps closed-loop from initial_state and compare states.

Parameters:

reference (Callable[[...], Any])
candidate (Callable[[...], Any])
initial_state (Any)
step_inputs (Sequence[Sequence[Any]])

Return type:

DriftReport

Weight-only int8 quantisation for world-model inference.

Weight-only quantisation stores nn.Linear weights as int8 with one fp scale per output channel and dequantises on the fly. It cuts weight memory ~4x (vs fp32) and, for the memory-bound small-batch steps world models run, usually speeds them up - especially under torch.compile, which fuses the dequantise into the matmul. Activations stay in floating point, so no calibration data is needed.

What is not quantised by default, and why:

VQ codebooks / quantizers (IRIS, Genie tokenizers): token identity is decided by nearest-neighbour distances, so small weight error flips tokens.
Stochastic-state heads (the RSSM prior/posterior projections): they output distribution parameters that are sampled from and fed back every step.
Small layers (min_in_features): they save little memory and are where quantisation error is proportionally largest.

Embeddings, norms and convolutions are never touched. Always check the result with simulusion.inference.rollout_drift().

class simulusion.inference.quantize.Int8WeightOnlyLinear(weight_int8, scale, bias)[source]#

Bases: Module

nn.Linear with int8 weights and per-output-channel scales.

Parameters:

weight_int8 (Tensor)
scale (Tensor)
bias (torch.Tensor | None)

classmethod from_linear(linear)[source]#

Parameters:: linear (Linear)
Return type:: Int8WeightOnlyLinear

weight_int8: Tensor#

scale: Tensor#

dequantized_weight(dtype)[source]#

Parameters:: dtype (dtype)
Return type:: Tensor

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

extra_repr()[source]#

Return type:: str

simulusion.inference.quantize.quantize_weights(module, *, skip=('quantizer', 'codebook', 'vq', 'fc_state_prior', 'fc_state_posterior'), min_in_features=64, backend='native')[source]#

Quantise nn.Linear weights of module to int8, in place.

Parameters:

module (Module) – Module to quantise. Put it in eval mode first; the result is for inference only.
skip (Iterable[str]) – Substrings of qualified module names to leave in floating point. A layer is skipped if any of its ancestors’ names matches too.
min_in_features (int) – Leave smaller nn.Linear layers alone.
backend (str) – "native" (default) uses Int8WeightOnlyLinear, which needs no extra dependency and exports with torch.export. "torchao" delegates to torchao’s int8 weight-only kernels, which are faster on recent GPUs; it applies the same filter.

Returns:

Qualified names of the layers that were quantised.

Return type:

list[str]

simulusion.inference.quantize.weight_memory_bytes(module)[source]#

Bytes held by parameters and buffers, e.g. to compare before/after.

Parameters:: module (Module)
Return type:: int

Self-describing deployment bundles.

A checkpoint alone is not deployable: a server also needs the config the model was built with, the input shapes and dtypes it expects, and ideally a compiled artifact it can run without the training code. A bundle is a directory holding all of that, plus a manifest.json describing it:

policy_bundle/
    manifest.json          # versions, input/output spec, artifact index
    weights.safetensors    # or weights.pt without safetensors
    config.json            # optional: the model config
    model.pt2              # optional: exported artifacts, one per format
    model.onnx

load_bundle reads it back; Bundle.load_artifact() returns a callable for the deployment format and Bundle.load_weights(module) restores eager weights for the Python path.

class simulusion.inference.bundle.Bundle(path, manifest, config=None, artifacts=<factory>)[source]#

Bases: object

A loaded deployment bundle.

Parameters:

path (Path)
manifest (dict[str, Any])
config (dict[str, Any] | None)
artifacts (dict[str, str])

path: Path#

manifest: dict[str, Any]#

config: dict[str, Any] | None = None#

artifacts: dict[str, str]#

property formats: list[str]#

state_dict(map_location='cpu')[source]#

Parameters:: map_location (str | device)
Return type:: dict[str, Any]

load_weights(module, strict=True)[source]#

Load the bundled weights into an eagerly constructed module.

Parameters:

module (Module)
strict (bool)

Return type:

Module

load_artifact(format=None, *, device=None)[source]#

Load an exported artifact as a callable (first available by default).

Parameters:

format (str | None)
device (device | str | None)

Return type:

Callable[[…], Any]

simulusion.inference.bundle.load_bundle(directory)[source]#

Read a bundle written by save_bundle().

Parameters:: directory (str | Path)
Return type:: Bundle

simulusion.inference.bundle.save_bundle(directory, module, *, example_inputs=None, formats=('exported_program',), config=None, name='model', metadata=None, verify=True, export_kwargs=None)[source]#

Write module as a deployment bundle.

Parameters:

directory (str | Path) – Output directory (created if missing).
module (Module) – Module to bundle - typically a pure step module such as DreamerStepModule.
example_inputs (Any) – Inputs used to export and to record the input spec. Required when formats is non-empty.
formats (Iterable[str]) – Export formats to include (see simulusion.export). Pass () for a weights-and-config bundle.
config (Any) – Model config (dataclass, object or dict), stored as JSON.
name (str) – Base file name of the artifacts.
metadata (dict[str, Any] | None) – Extra JSON-serialisable fields for the manifest.
verify (bool) – Check each artifact against the eager module with verify_export() and record the max error.
export_kwargs (dict[str, Any] | None) – Extra keyword arguments per format, e.g. {"exported_program": {"dynamic_shapes": ...}}.

Return type:

Bundle

One interface for every world model in Simulusion.

import torch
from simulusion.world_model import as_world_model, load_world_model

model = as_world_model(trained_agent)          # or load_world_model(path)
state = model.encode(frames, actions)          # real context
traj = model.rollout(state, planned_actions)   # imagined future
traj.frames, traj.rewards, traj.terminated

model.save("runs/my_model/world_model")        # any family, same format
model = load_world_model("runs/my_model/world_model")

WorldModel is the protocol; see simulusion.world_model.protocol for the conventions (frame layout and range, action layout, state, seeding) shared by every family. Supported families:

`dreamer`	`DreamerWorldModel` (Dreamer V1 and V2)
`planet`	`PlanetWorldModel`
`iris`	`IRISTokenWorldModel`
`diamond`	`DiamondWorldModel`
`genie`	`GenieWorldModel` (latent actions)
`world-models`	`MDRNNWorldModel` (ConvVAE + MDN-RNN)

JEPA and DiT are not world models in this sense: I-JEPA learns image representations without actions, and DiT generates images from a class label, so neither predicts what an action does. They have no adapter.

The protocol covers the model, not the agent: policies, planners and trainers stay with each family. Steppers (simulusion.inference.steppers) remain the acting interface.

class simulusion.world_model.ActionSpec(kind, n, low=None, high=None, latent=False)[source]#

Bases: object

The actions a world model accepts.

Variables:

kind (Literal['discrete', 'continuous']) – "discrete" (integer actions in [0, n)) or "continuous" (float vectors of length n).
n (int) – Number of actions (discrete) or action dimensions (continuous).
low (float | None) – Lower bound of continuous actions, or None if unbounded.
high (float | None) – Upper bound of continuous actions, or None if unbounded.
latent (bool) – The actions were learned from video rather than taken from an environment (Genie), so what each one does is for the model to say.

Parameters:

kind (Literal['discrete', 'continuous'])
n (int)
low (float | None)
high (float | None)
latent (bool)

kind: Literal['discrete', 'continuous']#

n: int#

low: float | None = None#

high: float | None = None#

latent: bool = False#

as_tensor(action, batch_size, device)[source]#

One step’s actions as (B,) int64 or (B, n) float32.

Accepts a tensor, array, list or scalar. A single action is repeated across the batch. Continuous actions are clipped to [low, high].

Parameters:

action (Any)
batch_size (int)
device (device)

Return type:

Tensor

class simulusion.world_model.Rollout(state, frames, rewards, terminated)[source]#

Bases: object

Many imagined steps, as returned by WorldModel.rollout().

Variables:

state (Dict[str, torch.Tensor]) – The model state after the last step.
frames (torch.Tensor | None) – Decoded frames (B, H, C, H, W) in [0, 1], one per step, or None when the rollout was run with decode=False.
rewards (torch.Tensor | None) – (B, H), or None without a reward head.
terminated (torch.Tensor | None) – (B, H) bool, or None without a termination head.

Parameters:

state (Dict[str, Tensor])
frames (Tensor | None)
rewards (Tensor | None)
terminated (Tensor | None)

state: Dict[str, Tensor]#

frames: Tensor | None#

rewards: Tensor | None#

terminated: Tensor | None#

class simulusion.world_model.Transition(state, reward=None, terminated=None, info=<factory>)[source]#

Bases: object

One imagined step, as returned by WorldModel.step().

Variables:

state (Dict[str, torch.Tensor]) – The model state after the step.
reward (torch.Tensor | None) – Predicted reward (B,), or None without a reward head.
terminated (torch.Tensor | None) – Predicted episode end (B,) bool, or None.
info (Dict[str, torch.Tensor]) – Family-specific extras, such as Dreamer’s continue probability.

Parameters:

state (Dict[str, Tensor])
reward (Tensor | None)
terminated (Tensor | None)
info (Dict[str, Tensor])

state: Dict[str, Tensor]#

reward: Tensor | None = None#

terminated: Tensor | None = None#

info: Dict[str, Tensor]#

class simulusion.world_model.WorldModel(*args, **kwargs)[source]#

Bases: Protocol

The interface every world model in Simulusion implements.

Wrap a trained model with as_world_model(), or load one with load_world_model(). See the module docstring for the conventions.

spec: WorldModelSpec#

device: device#

encode(frames, actions=None)[source]#

Turn real context frames (and the actions between them) into a state.

Parameters:

frames (Any)
actions (Any)

Return type:

Dict[str, Tensor]

step(state, action)[source]#

Imagine one step: the state after taking action.

Parameters:

state (Dict[str, Tensor])
action (Any)

Return type:

decode(state)[source]#

The current frame of state: float (B, C, H, W) in [0, 1].

Parameters:: state (Dict[str, Tensor])
Return type:: Tensor

rollout(state, actions, *, decode=True)[source]#

Imagine H steps; actions is (B, H) or (B, H, A).

Parameters:

state (Dict[str, Tensor])
actions (Any)
decode (bool)

Return type:

save(path)[source]#

Write the model to directory path; returns it.

Parameters:: path (str | Path)
Return type:: Path

class simulusion.world_model.WorldModelBase[source]#

Bases: ABC

Shared plumbing for the built-in world models.

Subclasses set family, build spec and device in __init__, and implement _encode(), _step(), _decode(), _save_weights() and _load_weights(). Input normalisation, autograd, rollout() and the save format live here.

family: ClassVar[str]#: Registry key, also written to the manifest.

spec: WorldModelSpec#

device: device#

encode(frames, actions=None)[source]#

Parameters:

frames (Any)
actions (Any)

Return type:

Dict[str, Tensor]

step(state, action)[source]#

Parameters:

state (Dict[str, Tensor])
action (Any)

Return type:

decode(state)[source]#

Parameters:: state (Dict[str, Tensor])
Return type:: Tensor

rollout(state, actions, *, decode=True)[source]#

Parameters:

state (Dict[str, Tensor])
actions (Any)
decode (bool)

Return type:

save(path)[source]#

Write world_model.json and the weights into directory path.

Read it back with load_world_model().

Parameters:: path (str | Path)
Return type:: Path

classmethod load(path, *, device=None, **options)[source]#

Load a directory written by save(), or the family’s own checkpoint.

Options saved in the manifest (such as a sampling temperature) are restored; keyword options override them.

Parameters:

path (str | Path)
device (str | device | None)
options (Any)

Return type:

WorldModelBase

static batch_size(state)[source]#

Parameters:: state (Dict[str, Tensor])
Return type:: int

as_frames(frames)[source]#

Frames as float (B, T, C, H, W) in [0, 1] on device.

Parameters:: frames (Any)
Return type:: Tensor

class simulusion.world_model.WorldModelSpec(family, obs_shape, action, context_frames=1, max_context=None, has_reward=True, has_termination=False, stochastic=True)[source]#

Bases: object

What a world model consumes and produces.

Variables:

family (str) – "dreamer", "planet", "iris", "diamond", "genie" or "world-models".
obs_shape (tuple[int, int, int]) – (C, H, W) of the frames it reads and decodes.
action (simulusion.world_model.protocol.ActionSpec) – The actions WorldModel.step() accepts.
context_frames (int) – The fewest frames WorldModel.encode() needs.
max_context (int | None) – The most frames the model attends to; older ones are forgotten. None for recurrent models, which carry everything in their state.
has_reward (bool) – Transition carries a predicted reward.
has_termination (bool) – Transition carries a predicted episode end.
stochastic (bool) – WorldModel.step() samples, so repeated calls on the same state differ unless the RNG is reseeded.

Parameters:

family (str)
obs_shape (tuple[int, int, int])
action (ActionSpec)
context_frames (int)
max_context (int | None)
has_reward (bool)
has_termination (bool)
stochastic (bool)

family: str#

obs_shape: tuple[int, int, int]#

action: ActionSpec#

context_frames: int = 1#

max_context: int | None = None#

has_reward: bool = True#

has_termination: bool = False#

stochastic: bool = True#

to_dict()[source]#

Return type:: dict[str, Any]

simulusion.world_model.as_world_model(model, *args, **kwargs)[source]#

Wrap a trained model of any supported family as a WorldModel.

Parameters:

model (Any)
args (Any)
kwargs (Any)

Return type:

simulusion.world_model.list_world_model_families()[source]#

The family names load_world_model() and world_model_class() accept.

Return type:: list[str]

simulusion.world_model.load_world_model(path, family=None, *, device=None, **options)[source]#

Load a world model saved by WorldModel.save(), or a training checkpoint.

Parameters:

path (str | Path)
family (str | None)
device (str | device | None)
options (Any)

Return type:

simulusion.world_model.world_model_class(family)[source]#

The WorldModelBase subclass for family.

Parameters:: family (str)
Return type:: type[WorldModelBase]

The protocol every Simulusion world model implements.

A world model predicts how observations evolve under actions. Each family (Dreamer, PlaNet, IRIS, DIAMOND, Genie, the VAE + MDN-RNN World Models) does this differently inside: an RNN latent, transformer tokens, a diffusion sampler. WorldModel gives them one outside:

state = model.encode(frames, actions)       # real context -> model state
out = model.step(state, action)             # one imagined transition
image = model.decode(out.state)             # model state -> pixels
traj = model.rollout(state, action_seq)     # many steps at once

Conventions, the same for every family:

Frames are RGB, channel-first, batch-first: (B, T, C, H, W), float in [0, 1] or uint8 in [0, 255]. A 4-D input is one sequence (T, C, H, W) and a 3-D input one frame. Frames of another size are resized to WorldModelSpec.obs_shape. WorldModel.decode() always returns float (B, C, H, W) in [0, 1].
Actions between context frames: actions[:, t] is the action taken at frame t, which led to frame t + 1, so T frames come with T - 1 actions. Discrete actions are integer (B, ...), continuous ones float (B, ..., A).
State is a dict[str, Tensor] with the batch on dim 0. It is plain data: copy it to branch a dream, stack it to batch dreams.
Rewards and terminations come from the model’s own heads, (B,) per step, or None when the family has no such head (see WorldModelSpec).
Randomness comes from PyTorch’s global RNG, so torch.manual_seed before a rollout makes it repeatable.

Every method runs without autograd and puts the model in eval mode: this is the inference surface. Training stays with each family’s trainer.

class simulusion.world_model.protocol.ActionSpec(kind, n, low=None, high=None, latent=False)[source]#

Bases: object

The actions a world model accepts.

Variables:

kind (Literal['discrete', 'continuous']) – "discrete" (integer actions in [0, n)) or "continuous" (float vectors of length n).
n (int) – Number of actions (discrete) or action dimensions (continuous).
low (float | None) – Lower bound of continuous actions, or None if unbounded.
high (float | None) – Upper bound of continuous actions, or None if unbounded.
latent (bool) – The actions were learned from video rather than taken from an environment (Genie), so what each one does is for the model to say.

Parameters:

kind (Literal['discrete', 'continuous'])
n (int)
low (float | None)
high (float | None)
latent (bool)

kind: Literal['discrete', 'continuous']#

n: int#

low: float | None = None#

high: float | None = None#

latent: bool = False#

as_tensor(action, batch_size, device)[source]#

One step’s actions as (B,) int64 or (B, n) float32.

Accepts a tensor, array, list or scalar. A single action is repeated across the batch. Continuous actions are clipped to [low, high].

Parameters:

action (Any)
batch_size (int)
device (device)

Return type:

Tensor

simulusion.world_model.protocol.MANIFEST_NAME = 'world_model.json'#: Written next to the weights by WorldModelBase.save().

class simulusion.world_model.protocol.Rollout(state, frames, rewards, terminated)[source]#

Bases: object

Many imagined steps, as returned by WorldModel.rollout().

Variables:

state (Dict[str, torch.Tensor]) – The model state after the last step.
frames (torch.Tensor | None) – Decoded frames (B, H, C, H, W) in [0, 1], one per step, or None when the rollout was run with decode=False.
rewards (torch.Tensor | None) – (B, H), or None without a reward head.
terminated (torch.Tensor | None) – (B, H) bool, or None without a termination head.

Parameters:

state (Dict[str, Tensor])
frames (Tensor | None)
rewards (Tensor | None)
terminated (Tensor | None)

state: Dict[str, Tensor]#

frames: Tensor | None#

rewards: Tensor | None#

terminated: Tensor | None#

class simulusion.world_model.protocol.Transition(state, reward=None, terminated=None, info=<factory>)[source]#

Bases: object

One imagined step, as returned by WorldModel.step().

Variables:

state (Dict[str, torch.Tensor]) – The model state after the step.
reward (torch.Tensor | None) – Predicted reward (B,), or None without a reward head.
terminated (torch.Tensor | None) – Predicted episode end (B,) bool, or None.
info (Dict[str, torch.Tensor]) – Family-specific extras, such as Dreamer’s continue probability.

Parameters:

state (Dict[str, Tensor])
reward (Tensor | None)
terminated (Tensor | None)
info (Dict[str, Tensor])

state: Dict[str, Tensor]#

reward: Tensor | None = None#

terminated: Tensor | None = None#

info: Dict[str, Tensor]#

class simulusion.world_model.protocol.WorldModel(*args, **kwargs)[source]#

Bases: Protocol

The interface every world model in Simulusion implements.

Wrap a trained model with as_world_model(), or load one with load_world_model(). See the module docstring for the conventions.

spec: WorldModelSpec#

device: device#

encode(frames, actions=None)[source]#

Turn real context frames (and the actions between them) into a state.

Parameters:

frames (Any)
actions (Any)

Return type:

Dict[str, Tensor]

step(state, action)[source]#

Imagine one step: the state after taking action.

Parameters:

state (Dict[str, Tensor])
action (Any)

Return type:

decode(state)[source]#

The current frame of state: float (B, C, H, W) in [0, 1].

Parameters:: state (Dict[str, Tensor])
Return type:: Tensor

rollout(state, actions, *, decode=True)[source]#

Imagine H steps; actions is (B, H) or (B, H, A).

Parameters:

state (Dict[str, Tensor])
actions (Any)
decode (bool)

Return type:

save(path)[source]#

Write the model to directory path; returns it.

Parameters:: path (str | Path)
Return type:: Path

class simulusion.world_model.protocol.WorldModelBase[source]#

Bases: ABC

Shared plumbing for the built-in world models.

family: ClassVar[str]#: Registry key, also written to the manifest.

spec: WorldModelSpec#

device: device#

encode(frames, actions=None)[source]#

Parameters:

frames (Any)
actions (Any)

Return type:

Dict[str, Tensor]

step(state, action)[source]#

Parameters:

state (Dict[str, Tensor])
action (Any)

Return type:

decode(state)[source]#

Parameters:: state (Dict[str, Tensor])
Return type:: Tensor

rollout(state, actions, *, decode=True)[source]#

Parameters:

state (Dict[str, Tensor])
actions (Any)
decode (bool)

Return type:

save(path)[source]#

Write world_model.json and the weights into directory path.

Read it back with load_world_model().

Parameters:: path (str | Path)
Return type:: Path

classmethod load(path, *, device=None, **options)[source]#

Load a directory written by save(), or the family’s own checkpoint.

Options saved in the manifest (such as a sampling temperature) are restored; keyword options override them.

Parameters:

path (str | Path)
device (str | device | None)
options (Any)

Return type:

WorldModelBase

static batch_size(state)[source]#

Parameters:: state (Dict[str, Tensor])
Return type:: int

as_frames(frames)[source]#

Frames as float (B, T, C, H, W) in [0, 1] on device.

Parameters:: frames (Any)
Return type:: Tensor

class simulusion.world_model.protocol.WorldModelSpec(family, obs_shape, action, context_frames=1, max_context=None, has_reward=True, has_termination=False, stochastic=True)[source]#

Bases: object

What a world model consumes and produces.

Variables:

family (str) – "dreamer", "planet", "iris", "diamond", "genie" or "world-models".
obs_shape (tuple[int, int, int]) – (C, H, W) of the frames it reads and decodes.
action (simulusion.world_model.protocol.ActionSpec) – The actions WorldModel.step() accepts.
context_frames (int) – The fewest frames WorldModel.encode() needs.
max_context (int | None) – The most frames the model attends to; older ones are forgotten. None for recurrent models, which carry everything in their state.
has_reward (bool) – Transition carries a predicted reward.
has_termination (bool) – Transition carries a predicted episode end.
stochastic (bool) – WorldModel.step() samples, so repeated calls on the same state differ unless the RNG is reseeded.

Parameters:

family (str)
obs_shape (tuple[int, int, int])
action (ActionSpec)
context_frames (int)
max_context (int | None)
has_reward (bool)
has_termination (bool)
stochastic (bool)

family: str#

obs_shape: tuple[int, int, int]#

action: ActionSpec#

context_frames: int = 1#

max_context: int | None = None#

has_reward: bool = True#

has_termination: bool = False#

stochastic: bool = True#

to_dict()[source]#

Return type:: dict[str, Any]

simulusion.world_model.protocol.read_manifest(path)[source]#

The world_model.json in directory path, or None.

Parameters:: path (str | Path)
Return type:: dict[str, Any] | None

Dreamer (V1 and V2) as a WorldModel.

class simulusion.world_model.dreamer.DreamerWorldModel(model)[source]#

Bases: WorldModelBase

The RSSM of a Dreamer agent, with its decoder, reward and discount heads.

encode filters the context through the posterior, step samples the prior, exactly as DreamerStepper does (and so as Dreamer’s own imagination does). Actions are continuous in [-1, 1]. terminated is the discount head’s continue probability falling below 0.5; models trained without that head have none.

Parameters:: model (Any) – A Dreamer, DreamerV2 or the DreamerAgent wrapping one.

family: ClassVar[str] = 'dreamer'#: Registry key, also written to the manifest.

PlaNet as a WorldModel.

class simulusion.world_model.planet.PlanetWorldModel(model, *, bit_depth=None)[source]#

Bases: WorldModelBase

PlaNet’s recurrent state-space model.

encode filters the context with the posterior mean, the way PlaNet’s CEM planner (RSSMPolicy) tracks the real episode; step advances the deterministic state and samples the prior, as the planner’s candidate rollouts do. Frames are quantised to bit_depth bits first, matching PlaNet’s preprocessing without its training-time dequantisation noise. PlaNet has a reward head but no termination head; its continuous actions are unbounded.

Parameters:

model (Any) – A Planet or its RecurrentStateSpaceModel.
bit_depth (Optional[int]) – Bits per colour channel the model was trained on (PlaNet’s default is 5).

family: ClassVar[str] = 'planet'#: Registry key, also written to the manifest.

IRIS as a WorldModel.

class simulusion.world_model.iris.IRISTokenWorldModel(agent, *, temperature=1.0, sample=True)[source]#

Bases: WorldModelBase

IRIS’s discrete autoencoder and Transformer, one frame at a time.

The state is the token history: tokens (B, T, K) for the last T frames and actions (B, T - 1) between them. step primes a fresh key/value cache with that history, then generates the next frame’s K tokens and reads the reward and termination heads at the action position, through the same Transformer calls imagine_rollout() makes. The history keeps at most WorldModelSpec.max_context frames, the window imagine_rollout falls back to when its cache fills up.

Rewards come out on the scale IRIS was trained on (sign-transformed by default, see IRISConfig.reward_transform).

The agent is switched to eval mode: its quantizer revives dead codes in training mode, which would rewrite the codebook on every call.

Parameters:

agent (Any) – An IRISAgent.
temperature (float) – Softmax temperature for sampling next-frame tokens.
sample (bool) – Sample tokens (default, as IRIS imagines); False takes the most likely token, which makes step deterministic.

family: ClassVar[str] = 'iris'#: Registry key, also written to the manifest.

DIAMOND as a WorldModel.

simulusion.world_model.diamond.build_diamond_networks(config, action_dim, device)[source]#

DIAMOND’s world-model networks: (diffusion_model, sampler, reward_model).

The one place they are built, so a DiamondAgent and a DiamondWorldModel loaded from its checkpoint are the same networks.

Parameters:

config (Any)
action_dim (int)
device (Any)

Return type:

tuple

class simulusion.world_model.diamond.DiamondWorldModel(model, *, sampling_steps=None)[source]#

Bases: WorldModelBase

DIAMOND’s diffusion model and reward/termination model.

Each frame is sampled by the diffusion model conditioned on the last L = num_conditioning_frames frames and the actions taken at them, exactly as in DIAMOND’s imagination: the reward/termination model scores (x_i, a_i), then the sampler draws x_{i+1}. encode needs L frames; without actions the context is taken to have been reached with action 0 (NOOP in Atari). Rewards are DIAMOND’s clipped classes {-1, 0, 1}.

Parameters:

model (Any) – A DiamondAgent, or any object with its config, diffusion_model, sampler and reward_model (None for a model trained without one).
sampling_steps (Optional[int]) – Override the Euler denoising steps per frame. Fewer is faster and blurrier; the paper uses 3.

family: ClassVar[str] = 'diamond'#: Registry key, also written to the manifest.

Genie as a WorldModel.

class simulusion.world_model.genie.GenieWorldModel(model, *, temperature=None)[source]#

Bases: WorldModelBase

Genie: video tokens in, latent actions drive the next frame.

Genie learns its actions from unlabelled video, so the action space is action_vocab_size latent actions (spec.action.latent is set): what each one does is whatever the model discovered. Without actions, encode labels the context’s transitions with the actions Genie infers itself. Genie has no reward or termination head.

The dynamics model sees at most num_frames frames, so the state keeps a rolling window of that many: tokens (B, T, N) and actions (B, T - 1). decode decodes the whole window, because the tokenizer’s decoder attends over time, and returns its last frame.

Parameters:

model (Any) – A Genie.
temperature (Optional[float]) – Sampling temperature for next-frame tokens. None keeps the model’s (2.0 in the paper); 0 takes the most likely token.

family: ClassVar[str] = 'genie'#: Registry key, also written to the manifest.

Ha & Schmidhuber’s World Models (ConvVAE + MDN-RNN) as a WorldModel.

class simulusion.world_model.vae_mdrnn.MDRNNWorldModel(vae, mdrnn, *, temperature=1.0)[source]#

Bases: WorldModelBase

A ConvVAE for pixels and an MDN-RNN for dynamics (Ha & Schmidhuber, 2018).

encode maps each frame to its VAE latent mean and runs the MDN-RNN’s LSTM over the context. step predicts a Gaussian mixture over the next latent and samples it, at temperature as in the paper: the mixture weights are sharpened or flattened by it and the standard deviations scaled by its square root. The reward and done heads give reward and terminated (done logit above 0). Actions are continuous and unbounded.

Parameters:

vae (Any) – The ConvVAE.
mdrnn (Any) – The MDN-RNN, as trained (MDRNN) or as a single-step MDRNNCell.
temperature (float) – Sampling temperature tau; 0 takes the mean of the most likely mixture component.

family: ClassVar[str] = 'world-models'#: Registry key, also written to the manifest.

classmethod from_checkpoints(vae_checkpoint, mdrnn_checkpoint, *, device=None, temperature=1.0)[source]#

Load the files train_convvae and train_mdn_rnn write.

Parameters:

vae_checkpoint (str | Path)
mdrnn_checkpoint (str | Path)
device (str | device | None)
temperature (float)

Return type:

MDRNNWorldModel

Training#

The one training loop every Simulusion family runs on.

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

result = Trainer(TrainerConfig(max_steps=10_000, precision="bf16")).fit(
    GenieTask(GenieTaskConfig(data=VideoDataConfig(path="data/clips")))
)
result.world_model_path  # runs/<name>/world_model, loadable with load_world_model

A run uses one device: a CPU, one GPU, or Apple MPS. Precision is a run setting (fp32, bf16, fp16, auto) applied through MixedPrecision, rather than a flag per model.

Every run gets a directory holding everything needed to reproduce it:

runs/<run_name>/
    config.yaml          # the fully resolved run config: re-run with --config
    run.json             # git commit, package versions, host, command line
    metrics.jsonl        # every logged metric
    checkpoints/step_000500/   # resumable: trainer.resume=<this dir> or "auto"
    world_model/         # the trained world model (WorldModel protocol)

class simulusion.training.trainer.DeviceLoader(loader, device)[source]#

Bases: object

Iterate a loader, moving every tensor in each batch to device.

Batches may be tensors or nested dicts, lists and tuples of them; other values pass through. len() and other attributes (dataset, batch_size, …) are the wrapped loader’s.

Parameters:

loader (Iterable[Any])
device (torch.device)

class simulusion.training.trainer.TrainResult(run_dir, progress, iterations, metrics=<factory>, checkpoint=None, world_model_path=None, model_path=None, best_checkpoint=None, best_value=None, stopped_early=False)[source]#

Bases: object

What Trainer.fit() returns.

Parameters:

run_dir (Path)
progress (int)
iterations (int)
metrics (dict[str, float])
checkpoint (Path | None)
world_model_path (Path | None)
model_path (Path | None)
best_checkpoint (Path | None)
best_value (float | None)
stopped_early (bool)

run_dir: Path#

progress: int#

iterations: int#

metrics: dict[str, float]#

checkpoint: Path | None = None#

world_model_path: Path | None = None#

model_path: Path | None = None#

best_checkpoint: Path | None = None#

best_value: float | None = None#

stopped_early: bool = False#

class simulusion.training.trainer.Trainer(config=None, *, run_config=None, config_filename='config.yaml', metrics_filename='metrics.jsonl')[source]#

Bases: object

Runs a TrainingTask to its budget.

Parameters:

config (Optional[TrainerConfig]) – Run settings; defaults throughout when None.
run_config (Optional[Mapping[str, Any]]) – The full run config (all sections) to write as config.yaml. simulusion.training.run.run() passes it; without it only the trainer and task sections are written.
config_filename (str) – Name of the resolved run config in the run directory.
metrics_filename (str) – Name of the metrics log in the run directory. The older entrypoints (simulusion.training.legacy) rename both so they never overwrite a family’s own config.yaml or metrics.jsonl in the folder it has always written to.

device: device#

run_dir: Path#

task: TrainingTask | None#

prepare_module(module)[source]#

module on the run’s device.

Parameters:: module (Module)
Return type:: Module

prepare_dataloader(loader)[source]#

loader with every batch’s tensors moved to the run’s device.

Parameters:: loader (Iterable[Any])
Return type:: DeviceLoader

gradient_update(optimizer, parameters, *, max_grad_norm=None, accumulate=True)[source]#

A GradientUpdate at this run’s precision and accumulation.

Parameters:

optimizer (Optimizer)
parameters (Iterable[Tensor])
max_grad_norm (float | None)
accumulate (bool)

Return type:

GradientUpdate

static infinite(loader)[source]#

Cycle loader forever, re-shuffling each pass.

Parameters:: loader (Iterable[Any])
Return type:: Iterator[Any]

fit(task)[source]#

Train task until its budget is spent or early stopping fires.

Parameters:: task (TrainingTask)
Return type:: TrainResult

save_checkpoint(*, final=False, name=None)[source]#

Write a resumable checkpoint of the task and the loop; returns its dir.

name writes it to checkpoints/<name> instead of checkpoints/step_<progress>, outside the pruned periodic ones.

Parameters:

final (bool)
name (str | None)

Return type:

Path

class simulusion.training.trainer.TrainerConfig(max_steps=None, seed=0, deterministic=False, precision='fp32', device=None, grad_accumulation=1, log_every=100, eval_every=0, checkpoint_every=0, keep_checkpoints=2, out_dir='runs', run_name=None, resume=None, export_world_model=True, best_metric=None, best_mode='max', early_stopping_metric=None, early_stopping_mode='min', patience=10, min_delta=0.0001, console=True)[source]#

Bases: SerializableConfigMixin

Settings shared by every training run, whatever the family.

Variables:

max_steps (int | None) – Budget in the task’s progress unit (optimizer steps, environment steps or epochs; see the task). None uses the task’s default.
seed (int) – Seeds Python, NumPy and PyTorch.
deterministic (bool) – Ask PyTorch for deterministic kernels. Slower; some ops only warn.
precision (str) – fp32, bf16, fp16 or auto.
device (str | None) – None (or "auto") picks CUDA, then MPS, else CPU; "cpu", "cuda", "cuda:1" or "mps" choose one. A device that is not available falls back with a warning.
grad_accumulation (int) – Batches per optimizer step, for tasks whose updates support it (dataset-driven ones).
log_every (int) – Log metrics every this many units of progress.
eval_every (int) – Evaluate every this many units; 0 only at the end.
checkpoint_every (int) – Checkpoint every this many units; 0 only at the end.
keep_checkpoints (int) – Periodic checkpoints kept on disk (oldest dropped). The final one is always kept.
out_dir (str) – Parent directory of run directories.
run_name (str | None) – Run directory name; a timestamped one when None.
resume (str | None) – A checkpoint directory to resume from, or "auto" for the newest complete checkpoint of out_dir/run_name.
export_world_model (bool) – Write the task’s world model to <run>/world_model when training ends.
best_metric (str | None) – Evaluation metric whose best value so far is kept as a checkpoint in checkpoints/best (rewritten on every improvement, never pruned); None keeps none.
best_mode (str) – "max" or "min": which direction is better.
early_stopping_metric (str | None) – Evaluation metric that stops the run once it stops improving; None disables early stopping.
early_stopping_mode (str) – "min" or "max".
patience (int) – Evaluations without improvement before stopping.
min_delta (float) – Relative improvement that counts.
console (bool) – Print metrics to the console.

Parameters:

max_steps (int | None)
seed (int)
deterministic (bool)
precision (str)
device (str | None)
grad_accumulation (int)
log_every (int)
eval_every (int)
checkpoint_every (int)
keep_checkpoints (int)
out_dir (str)
run_name (str | None)
resume (str | None)
export_world_model (bool)
best_metric (str | None)
best_mode (str)
early_stopping_metric (str | None)
early_stopping_mode (str)
patience (int)
min_delta (float)
console (bool)

max_steps: int | None = None#

seed: int = 0#

deterministic: bool = False#

precision: str = 'fp32'#

device: str | None = None#

grad_accumulation: int = 1#

log_every: int = 100#

eval_every: int = 0#

checkpoint_every: int = 0#

keep_checkpoints: int = 2#

out_dir: str = 'runs'#

run_name: str | None = None#

resume: str | None = None#

export_world_model: bool = True#

best_metric: str | None = None#

best_mode: str = 'max'#

early_stopping_metric: str | None = None#

early_stopping_mode: str = 'min'#

patience: int = 10#

min_delta: float = 0.0001#

console: bool = True#

simulusion.training.trainer.latest_checkpoint(run_dir)[source]#

The newest complete checkpoint in a run directory, or None.

Parameters:: run_dir (str | Path)
Return type:: Path | None

What a family plugs into the shared Trainer.

A TrainingTask is everything that differs between families: which networks and optimizers exist, where batches come from (a dataset, a replay buffer filled by an environment), and what one unit of training does. The trainer owns everything that does not differ: the device, precision, seeding, the step budget, logging, evaluation cadence, early stopping, checkpoints and resuming, and exporting the world model at the end.

@register_task("my-model")
class MyTask(TrainingTask):
    config_cls = MyTaskConfig        # dataclass; its fields are YAML sections

    def setup(self, trainer):
        self.model = trainer.prepare_module(MyModel(**...))
        self.update = trainer.gradient_update(
            torch.optim.Adam(self.model.parameters()), self.model.parameters()
        )
        self.batches = trainer.infinite(trainer.prepare_dataloader(loader))

    def train_step(self):
        with self.trainer.precision.autocast():
            loss = self.model(next(self.batches).to(self.trainer.device))
        self.update(loss)
        return {"loss": loss.item()}

A task is the “agent” half of a family: it may hold a world model and a policy and decide how they train together. The world model itself stays a separate object, reachable through TrainingTask.world_model(), so it can be saved, evaluated or planned with on its own.

class simulusion.training.task.TrainingTask(config=None)[source]#

Bases: ABC

One family’s training, run by Trainer.

Subclasses set config_cls (a dataclass whose fields are the run config’s sections besides trainer) and implement setup() and train_step().

Variables:

name (ClassVar[str]) – Registry name, set by register_task().
config_cls (ClassVar[type]) – The task’s config dataclass.
progress_unit (ClassVar[str]) – What progress() counts, for logs: "steps" (calls to train_step()), "env_steps" or "epochs".
default_max_steps (ClassVar[int]) – Budget used when the trainer config gives none.

Parameters:

config (Any)

name: ClassVar[str] = ''#

config_cls: ClassVar[type]#

progress_unit: ClassVar[str] = 'steps'#

default_max_steps: ClassVar[int] = 1000#

property trainer: Trainer#

attach(trainer)[source]#

Parameters:: trainer (Trainer)
Return type:: None

max_steps()[source]#

The budget, in progress_unit, when the trainer sets none.

Return type:: int

abstractmethod setup(trainer)[source]#

Build networks, optimizers and data on trainer.device.

Parameters:: trainer (Trainer)
Return type:: None

abstractmethod train_step()[source]#

One unit of training; returns scalar metrics.

Return type:: Mapping[str, float]

progress()[source]#

How far training is, in progress_unit.

Defaults to the number of train_step() calls. Tasks whose budget is counted in environment steps or epochs override it.

Return type:: int

evaluate()[source]#

Periodic evaluation metrics (validation loss, episode return).

Return type:: Mapping[str, float]

final_evaluation()[source]#

Metrics computed once, after the last step. Defaults to evaluate().

Return type:: Mapping[str, float]

world_model()[source]#

The trained world model behind the shared protocol, if the task has one.

Return type:: Optional[‘WorldModel’]

export(directory)[source]#

Write the trained model in its own reusable form; return directory.

For models that are not world models (DiT, I-JEPA, a VAE), which world_model() cannot return. Called once, when training ends, with <run_dir>/model. The default writes nothing.

Parameters:: directory (Path)
Return type:: Path | None

state_dict()[source]#

Everything needed to resume: weights, optimizer and scaler states.

Return type:: dict[str, Any]

load_state_dict(state)[source]#

Parameters:: state (dict[str, Any])
Return type:: None

save_checkpoint(directory)[source]#

Write the task’s resumable state into directory.

The default saves state_dict() to task.pt. Tasks with state too large or too odd for one file (a replay buffer) extend it.

Parameters:: directory (Path)
Return type:: None

load_checkpoint(directory)[source]#

Parameters:: directory (Path)
Return type:: None

should_stop()[source]#

True to end training now, before the budget (a search that converged).

Checked after every step. The final evaluation, checkpoint and export still run.

Return type:: bool

close()[source]#

Release environments and other resources.

Return type:: None

simulusion.training.task.get_task(name)[source]#

The task class registered as name.

Parameters:: name (str)
Return type:: type[TrainingTask]

simulusion.training.task.list_tasks()[source]#

Return type:: list[str]

simulusion.training.task.register_task(name, target=None)[source]#

Use as a decorator on the class, or call with target="module:Class" to register lazily without importing it.

Parameters:

name (str)
target (str | None)

Return type:

Callable[[type[TrainingTask]], type[TrainingTask]]

Mixed precision and optimizer steps shared by every trainer.

Each family used to carry its own use_amp flag, a GradScaler("cuda") per optimizer and the same six lines around backward(). They live here once, so a run’s precision is one setting rather than a flag per model:

precision = MixedPrecision("bf16", device)
update = GradientUpdate(optimizer, model.parameters(), precision, max_grad_norm=1.0)

with precision.autocast():
    loss = model(batch)
grad_norm = update(loss)

GradientUpdate keeps a loss scaler per optimizer. Algorithms with several optimizers (a world model, an actor and a critic) step them one after another, and a scaler shared between them would let an overflow in one skip the others.

class simulusion.training.precision.GradientUpdate(optimizer, parameters, precision, *, max_grad_norm=None, accumulation_steps=1)[source]#

Bases: object

One optimizer’s step: zero, backward, unscale, clip, step.

Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to step.
parameters (Iterable[torch.Tensor]) – The parameters to clip; normally the ones optimizer owns.
precision (MixedPrecision) – Decides whether the loss is scaled.
max_grad_norm (Optional[float]) – Clip threshold, or None not to clip. The norm is measured after unscaling, so the threshold means the same thing at every precision.
accumulation_steps (int) – Optimizer steps happen every this many calls; the calls in between only add gradients. Each loss is divided by it, so the step sees the mean over the accumulated batches.

Calling it returns the gradient norm before clipping (a tensor), or None on a call that only accumulated. A non-finite norm under fp16 means the scaler skipped the step.

state_dict()[source]#

Return type:: dict[str, Any]

load_state_dict(state)[source]#

Parameters:: state (dict[str, Any])
Return type:: None

class simulusion.training.precision.MixedPrecision(precision, device, *, scale_loss=None)[source]#

Bases: object

Autocast dtype and loss scaling for one training run.

Parameters:

precision (str | None) – "fp32", "bf16", "fp16" or "auto" (bf16 on CUDA devices that support it, fp16 on older ones, fp32 elsewhere).
device (torch.device | str) – The device the run trains on.
scale_loss (Optional[bool])

fp32 enters no autocast region, so it is bit-identical to running without one. Only fp16 scales the loss: bf16 has fp32’s dynamic range.

classmethod from_use_amp(use_amp, device)[source]#

The precision the per-family use_amp flags used to select.

use_amp meant autocast at the device’s default dtype: fp16 on CUDA (with a scaler), bf16 on CPU and fp16 on MPS, both without one: the GradScaler("cuda") they built disabled itself off CUDA.

Parameters:

use_amp (bool)
device (device | str)

Return type:

MixedPrecision

property enabled: bool#

property scales_loss: bool#

autocast(enabled=True)[source]#

Autocast at this precision; enabled=False runs a block in fp32.

Parameters:: enabled (bool)
Return type:: Iterator[None]

grad_scaler()[source]#

A fresh loss scaler, enabled only for fp16.

Return type:: GradScaler

Run any training task from one YAML file.

A run config names a task and fills its sections:

task: genie
trainer:            # TrainerConfig: budget, precision, logging, checkpoints
  max_steps: 20000
  precision: bf16
model:              # the family's own config (GenieConfig here)
  num_frames: 16
data:               # where batches come from
  path: data/clips

Train it from the command line, with key=value overrides for any field (nested keys use dots):

simulusion train --config genie.yaml trainer.max_steps=500 model.dynamics_depth=4
simulusion train genie data.path=data/clips          # defaults for the rest

or from Python with run(). Every run writes its fully resolved config to <run_dir>/config.yaml; passing that file back with --config repeats the run. Example configs for each task live in simulusion/configs/runs/.

simulusion.training.run.RUN_CONFIG_DIR = PosixPath('/home/runner/work/simulusion/simulusion/simulusion/configs/runs')#: Shipped example run configs.

simulusion.training.run.run_training(config=None, overrides=(), *, task=None)#

simulusion.run_training: the same as run().

Parameters:

config (str | Path | Mapping[str, Any] | None)
overrides (Sequence[str] | Mapping[str, Any])
task (str | None)

Return type:

TrainResult

simulusion.training.run.dump_run_config(config)[source]#

Parameters:: config (Mapping[str, Any])
Return type:: str

simulusion.training.run.main(argv=None)[source]#

simulusion train [TASK] [--config FILE] [--print-config] [key=value ...].

Parameters:: argv (Sequence[str] | None)
Return type:: TrainResult | None

simulusion.training.run.resolve_run_config(config=None, overrides=(), *, task=None)[source]#

Compose defaults, a YAML file (or mapping) and overrides into a run.

Returns (task_class, trainer_config, task_config, resolved_dict). Precedence, lowest first: the task’s defaults, config, overrides. An explicit task wins over one named in config or overrides.

Parameters:

config (str | Path | Mapping[str, Any] | None)
overrides (Sequence[str] | Mapping[str, Any])
task (str | None)

Return type:

tuple[type[TrainingTask], TrainerConfig, Any, dict[str, Any]]

simulusion.training.run.run(config=None, overrides=(), *, task=None)[source]#

Train the task a run config describes; see the module docstring.

Parameters:

config (str | Path | Mapping[str, Any] | None)
overrides (Sequence[str] | Mapping[str, Any])
task (str | None)

Return type:

TrainResult

simulusion.training.run.run_config_dict(task, trainer, task_config)[source]#

The YAML-ready form of a run: task, trainer and the task’s sections.

Parameters:

task (str)
trainer (TrainerConfig)
task_config (Any)

Return type:

dict[str, Any]

Pieces several tasks share: where offline batches come from.

class simulusion.training.tasks.common.SequenceDataConfig(path=None, source='folder', camera=None, sequence_length=None, stride=1, image_size=None, batch_size=16, num_workers=0, val_fraction=0.0, val_batches=8, episodes=None, preload=False, seed=0)[source]#

Bases: object

Where an offline task reads its episodes, and how it cuts them.

Variables:

path (str | None) – An episode folder, or for source: lerobot a local LeRobot dataset folder.
source (str) – "folder" (Simulusion episodes, videos, .npy) or "lerobot".
camera (str | None) – The view to train on, for datasets with several.
sequence_length (int | None) – Frames per training window; the task’s natural length (its model’s context) when None.
stride (int) – Frames between window starts.
image_size (int | None) – Resize frames to this square size; the model’s size when None.
batch_size (int) – Windows per batch (per process under data parallelism).
num_workers (int) – DataLoader worker processes.
val_fraction (float) – Fraction of episodes held out for evaluation; 0 for none.
val_batches (int) – Batches scored per evaluation.
episodes (List[int] | None) – Use only these episode indices.
preload (bool) – Decode every episode into memory up front.
seed (int) – Seed for the train / validation split.

Parameters:

path (str | None)
source (str)
camera (str | None)
sequence_length (int | None)
stride (int)
image_size (int | None)
batch_size (int)
num_workers (int)
val_fraction (float)
val_batches (int)
episodes (List[int] | None)
preload (bool)
seed (int)

path: str | None = None#

source: str = 'folder'#

camera: str | None = None#

sequence_length: int | None = None#

stride: int = 1#

image_size: int | None = None#

batch_size: int = 16#

num_workers: int = 0#

val_fraction: float = 0.0#

val_batches: int = 8#

episodes: List[int] | None = None#

preload: bool = False#

seed: int = 0#

simulusion.training.tasks.common.load_replay(buffer, path)[source]#

Restore what save_replay() wrote into a buffer of the same size.

Parameters:

buffer (Any)
path (Path)

Return type:

None

simulusion.training.tasks.common.mean_metrics(rows)[source]#

Parameters:: rows (List[dict])
Return type:: dict

simulusion.training.tasks.common.open_episodes(config)[source]#

The episodes config points at, and their metadata if known.

Parameters:: config (SequenceDataConfig)
Return type:: Tuple[Any, DatasetMeta | None]

simulusion.training.tasks.common.save_replay(buffer, path)[source]#

Write the filled part of a ring replay buffer, plus its write position.

Parameters:

buffer (Any)
path (Path)

Return type:

None

simulusion.training.tasks.common.sequence_loaders(config, trainer, *, sequence_length, image_size, require_actions=False)[source]#

An endless iterator of training batches, a validation loader, and metadata.

Parameters:

config (SequenceDataConfig)
trainer (Any)
sequence_length (int)
image_size (int | Tuple[int, int] | None)
require_actions (bool)

Return type:

Tuple[Iterator[dict], DataLoader | None, DatasetMeta | None]

Data#

The episode: Simulusion’s one trajectory format, for video and for RL.

An Episode with T frames has T - 1 transitions, the same convention as the WorldModel protocol: actions[t] is the action taken at frames[t], which led to frames[t + 1] and earned rewards[t]; terminated[t] says frames[t + 1] ended the episode, truncated[t] that it was cut off (a time limit) without ending. Everything but frames is optional, so a plain video is an episode too.

On disk an episode is one .npz file:

Key	Contents	Required
`frames`	`(T, H, W, C)` uint8, the main camera	yes
`actions`	`(T-1,)` int (discrete) or `(T-1, A)`	no
`rewards`	`(T-1,)` float	no
`terminated`	`(T-1,)` bool	no
`truncated`	`(T-1,)` bool	no
`frames.<camera>`	`(T, H, W, C)` uint8, further cameras	no
`obs.<name>`	`(T, ...)` low-dimensional observations (robot joint state, proprioception)	no
`language`	0-d string, a task instruction	no

This is the layout RecordingWrapper writes, extended for robot data (several cameras, proprioception, language).

class simulusion.data.episode.Episode(frames, actions=None, rewards=None, terminated=None, truncated=None, cameras=<factory>, observations=<factory>, language=None)[source]#

Bases: object

One trajectory; see the module docstring for the conventions.

Variables:

frames (numpy.ndarray) – (T, H, W, C) uint8 frames from the main camera.
actions (numpy.ndarray | None) – (T - 1, ...) actions, or None for video.
rewards (numpy.ndarray | None) – (T - 1,) rewards, or None.
terminated (numpy.ndarray | None) – (T - 1,) bool episode ends, or None.
truncated (numpy.ndarray | None) – (T - 1,) bool time-limit cuts, or None.
cameras (Dict[str, numpy.ndarray]) – Further views, each (T, H, W, C) uint8.
observations (Dict[str, numpy.ndarray]) – Low-dimensional observations, each (T, ...).
language (str | None) – A natural-language task instruction, if any.

Parameters:

frames (ndarray)
actions (ndarray | None)
rewards (ndarray | None)
terminated (ndarray | None)
truncated (ndarray | None)
cameras (Dict[str, ndarray])
observations (Dict[str, ndarray])
language (str | None)

frames: ndarray#

actions: ndarray | None = None#

rewards: ndarray | None = None#

terminated: ndarray | None = None#

truncated: ndarray | None = None#

cameras: Dict[str, ndarray]#

observations: Dict[str, ndarray]#

language: str | None = None#

property num_transitions: int#

property has_actions: bool#

validate()[source]#

Raise ValueError if the arrays disagree on the episode’s length.

Return type:: None

slice(start, stop)[source]#

Frames start:stop and the transitions between them.

Parameters:

start (int)
stop (int)

Return type:

to_arrays()[source]#

Return type:: dict[str, ndarray]

save(path, *, compress=True)[source]#

Write the episode as .npz; returns the path written.

Parameters:

path (str | Path)
compress (bool)

Return type:

Path

classmethod from_arrays(arrays)[source]#

Parameters:: arrays (Any)
Return type:: Episode

classmethod load(path)[source]#

Parameters:: path (str | Path)
Return type:: Episode

simulusion.data.episode.as_uint8_frames(frames)[source]#

Frames as (T, H, W, C) uint8, from uint8 or [0, 1] floats, HWC or CHW.

Parameters:: frames (Any)
Return type:: ndarray

simulusion.data.episode.npz_length(path)[source]#

The number of frames in an episode file, read from its header only.

Parameters:: path (str | Path)
Return type:: int

A folder of episodes: the on-disk dataset every task reads.

my_dataset/
    meta.json                 # optional: fps, action space, cameras, source
    episode_000000.npz        # one Episode each (see simulusion.data.episode)
    episode_000001.npz
    ...

Folders without meta.json work too, as do loose videos (.mp4 etc., read as frame-only episodes) and .npy frame arrays, so a directory of clips is already a dataset. Sub-folders are searched recursively.

class simulusion.data.folder.DatasetMeta(action_kind=None, action_dim=None, action_low=None, action_high=None, fps=None, cameras=<factory>, observation_keys=<factory>, source=None, info=<factory>)[source]#

Bases: object

What meta.json records about a folder of episodes.

Variables:

action_kind (str | None) – "discrete", "continuous" or None (video).
action_dim (int | None) – Number of discrete actions, or the action vector length.
action_high (action_low /) – Bounds of continuous actions, if known.
fps (float | None) – Frame rate of the recording, if known.
cameras (List[str]) – Names of the extra cameras (frames.<name> arrays).
observation_keys (List[str]) – Names of the low-dimensional observations.
source (str | None) – Where the data came from ("lerobot:datasets/pusht").
info (dict[str, Any]) – Anything else worth keeping.

Parameters:

action_kind (str | None)
action_dim (int | None)
action_low (List[float] | None)
action_high (List[float] | None)
fps (float | None)
cameras (List[str])
observation_keys (List[str])
source (str | None)
info (dict[str, Any])

action_kind: str | None = None#

action_dim: int | None = None#

action_low: List[float] | None = None#

action_high: List[float] | None = None#

fps: float | None = None#

cameras: List[str]#

observation_keys: List[str]#

source: str | None = None#

info: dict[str, Any]#

to_json()[source]#

Return type:: dict[str, Any]

classmethod from_json(values)[source]#

Parameters:: values (dict[str, Any])
Return type:: DatasetMeta

classmethod infer(episode, **extra)[source]#

Fill the action space and keys from one episode.

Parameters:

episode (Episode)
extra (Any)

Return type:

DatasetMeta

class simulusion.data.folder.EpisodeFolder(root, *, episodes=None)[source]#

Bases: Sequence[Episode]

Read a folder of episodes; folder[i] loads episode i.

Parameters:

root (str | Path) – The folder.
episodes (Optional[Iterable[int]]) – Indices of the episodes to keep (after sorting by path), for train/validation splits.

files: List[Path]#

property lengths: List[int]#: Frames per episode, read without loading the frames where possible.

split(val_fraction, seed=0)[source]#

Disjoint train / validation folders, split by whole episodes.

Parameters:

val_fraction (float)
seed (int)

Return type:

tuple[EpisodeFolder, EpisodeFolder]

class simulusion.data.folder.EpisodeWriter(root, *, meta=None, compress=True)[source]#

Bases: object

Append episodes to a folder, numbering after any already there.

with EpisodeWriter("data/pusht", meta=DatasetMeta(fps=10)) as writer:
    for episode in episodes:
        writer.write(episode)

Parameters:

root (str | Path)
meta (Optional[DatasetMeta])
compress (bool)

write(episode)[source]#

Parameters:: episode (Episode)
Return type:: Path

close()[source]#

Return type:: None

simulusion.data.folder.load_episode(path)[source]#

One episode from an .npz, a .npy frame array or a video file.

Parameters:: path (str | Path)
Return type:: Episode

simulusion.data.folder.read_meta(root)[source]#

Parameters:: root (str | Path)
Return type:: DatasetMeta | None

simulusion.data.folder.write_meta(root, meta)[source]#

Parameters:

root (str | Path)
meta (DatasetMeta)

Return type:

Path

Fixed-length windows over episodes, as a PyTorch dataset.

World models train on short clips cut from longer episodes. A SequenceDataset indexes every window of sequence_length frames that fits inside an episode and returns it as a dict of tensors:

Frames stay uint8 so batches are a quarter of the size on the way to the GPU; tasks convert them with frames_to_float().

class simulusion.data.sequence.SequenceDataset(episodes, sequence_length, *, stride=1, image_size=None, cameras=(), main_camera=None, observation_keys=(), require_actions=False, cache_episodes=16, preload=False)[source]#

Bases: Dataset

Every sequence_length-frame window of a collection of episodes.

Parameters:

episodes (Sequence[Episode]) – An EpisodeFolder, or any sequence of Episode.
sequence_length (int) – Frames per window L.
stride (int) – Frames between consecutive window starts.
image_size (Optional[int | Tuple[int, int]]) – Resize frames to (size, size) or (H, W); None keeps them as stored.
cameras (Sequence[str]) – Extra cameras to return as frames.<cam>.
main_camera (Optional[str]) – Return this extra camera as frames instead of the episode’s main view.
observation_keys (Sequence[str]) – Low-dimensional observations to return.
require_actions (bool) – Raise if an episode has no actions, rather than returning windows without them.
cache_episodes (int) – Decoded episodes kept per worker. Raise it (or preload) when episodes are re-read often.
preload (bool) – Load every episode up front.

simulusion.data.sequence.collate_sequences(items)[source]#

Stack windows into a batch, keeping only the keys every window has.

Parameters:: items (List[Dict[str, Any]])
Return type:: Dict[str, Tensor]

simulusion.data.sequence.frames_to_float(frames)[source]#

uint8 frames to float in [0, 1] (floats pass through).

Parameters:: frames (Tensor)
Return type:: Tensor

simulusion.data.sequence.one_hot_actions(actions, n)[source]#

Integer actions (..., ) to float one-hot (..., n); floats pass through.

Parameters:

actions (Tensor)
n (int)

Return type:

Tensor

Read LeRobot datasets as Simulusion episodes.

LeRobot is the Hugging Face format most open robot datasets now ship in. LeRobotEpisodes reads one from a local folder, without the lerobot package: the format is Parquet for the low-dimensional signals plus MP4 per camera, described by meta/info.json. Both layouts are supported:

v2.x: one Parquet file and one MP4 per camera per episode, listed in meta/episodes.jsonl.
v3.0: many episodes per Parquet / MP4 file, located through the meta/episodes/ tables (row ranges and video timestamps).

from simulusion.data import LeRobotEpisodes, SequenceDataset

episodes = LeRobotEpisodes("datasets/pusht")         # a local dataset folder
episodes[0].frames.shape                               # (T, 96, 96, 3)
episodes[0].actions.shape                              # (T - 1, 2)
windows = SequenceDataset(episodes, sequence_length=16)

or convert once to an episode folder: simulusion datasets import lerobot datasets/pusht data/pusht. Needs pip install simulusion[lerobot] (PyArrow and PyAV).

How LeRobot maps onto Episode: the chosen camera becomes frames and the others cameras; action becomes actions minus its last row (it leads past the final frame); next.reward and next.done become rewards and terminated the same way, with the final next.done moved onto the last transition; other numeric observation.* features become observations (so observation.state is observations["state"]); the episode’s first task becomes language.

class simulusion.data.lerobot.LeRobotEpisodes(source, *, camera=None, extra_cameras=(), episodes=None)[source]#

Bases: Sequence[Episode]

The episodes of a LeRobot dataset; dataset[i] decodes episode i.

Parameters:

source (str | Path) – A local dataset folder, the one holding meta/info.json.
camera (Optional[str]) – Feature key of the view to use as frames ("observation.images.top", or just "top"). Defaults to the first camera in name order.
extra_cameras (str | Sequence[str]) – Further views to keep as cameras: "all", or a list of keys. () keeps none (less decoding).
episodes (Optional[Iterable[int]]) – Episode indices to read; all when None.

property lengths: List[int]#

property meta: DatasetMeta#: The episode-folder metadata this dataset converts to.

simulusion.data.lerobot.import_lerobot(source, out_dir, *, camera=None, extra_cameras=(), episodes=None)[source]#

Convert a LeRobot dataset into an episode folder; returns out_dir.

Parameters:

source (str | Path)
out_dir (str | Path)
camera (str | None)
extra_cameras (str | Sequence[str])
episodes (Iterable[int] | None)

Return type:

Path

Convert RLDS episodes into Simulusion episodes.

RLDS (the format of Open X-Embodiment and many other robot datasets) stores each episode as a sequence of steps, every step a dict with an observation (itself a dict of cameras and state), an action, a reward and the flags is_first / is_last / is_terminal.

episode_from_rlds() turns one such episode, given as plain Python and NumPy values, into an Episode; reading the dataset files themselves is left to whatever produced the steps.

RLDS puts the action in the step it is taken from, so a T-step episode has T actions; the last one leads past the final observation and is dropped, as are its reward and flags. is_terminal on the final step marks the last transition terminated; an episode that ends without it is truncated.

class simulusion.data.rlds.RLDSKeys(image='image', cameras=(), state=(), action=(), language='language_instruction')[source]#

Bases: object

Where the parts of an episode live inside an RLDS step.

Variables:

image (str) – Observation key of the main camera.
cameras (Sequence[str]) – Further observation keys to keep as cameras.
state (Sequence[str]) – Observation keys of low-dimensional state to keep.
action (Sequence[str]) – Action keys, in order, when action is a dict (the RT-1 format); they are flattened and concatenated. Empty when action is already an array.
language (str) – Step or observation key of the language instruction.

Parameters:

image (str)
cameras (Sequence[str])
state (Sequence[str])
action (Sequence[str])
language (str)

image: str = 'image'#

cameras: Sequence[str] = ()#

state: Sequence[str] = ()#

action: Sequence[str] = ()#

language: str = 'language_instruction'#

simulusion.data.rlds.episode_from_rlds(episode, keys=RLDSKeys(image='image', cameras=(), state=(), action=(), language='language_instruction'))[source]#

One RLDS episode ({"steps": [...]} or the steps themselves) as an Episode.

Parameters:

episode (Mapping[str, Any] | Iterable[Mapping[str, Any]])
keys (RLDSKeys)

Return type:

Agents#

Act through any world model by planning in it.

A WorldModel with a reward head is enough to choose actions: imagine many action sequences, score them by predicted return, keep the best. CEMPlanner does that with the cross-entropy method, the planner PlaNet uses, but through the shared protocol only, so it works on every family’s world model: one trained with its agent, or alone from a dataset (dreamer-wm, iris-wm, diamond-wm, mdrnn).

from simulusion.agents import CEMPlanner, PlanningAgent, run_episode
from simulusion.world_model import load_world_model

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

class simulusion.agents.planning.CEMPlanner(horizon=12, candidates=500, iterations=5, elites=50, discount=0.99, init_std=1.0, min_std=0.05, smoothing=0.0)[source]#

Bases: object

Cross-entropy-method planning over a world model’s imagined returns.

Variables:

horizon (int) – Steps per imagined sequence.
candidates (int) – Sequences per round.
iterations (int) – Refitting rounds.
elites (int) – Best sequences each refit uses.
discount (float) – Per-step discount of imagined rewards.
init_std (float) – Initial standard deviation of continuous actions.
min_std (float) – Floor on the refitted standard deviation.
smoothing (float) – Weight of the old distribution in each refit (0 = none).

Parameters:

horizon (int)
candidates (int)
iterations (int)
elites (int)
discount (float)
init_std (float)
min_std (float)
smoothing (float)

horizon: int = 12#

candidates: int = 500#

iterations: int = 5#

elites: int = 50#

discount: float = 0.99#

init_std: float = 1.0#

min_std: float = 0.05#

smoothing: float = 0.0#

plan(model, state)[source]#

The first action of the best plan from state (a batch of one).

Returns (A,) float for continuous actions, a scalar int64 for discrete ones.

Parameters:

model (WorldModel)
state (Dict[str, Tensor])

Return type:

Tensor

class simulusion.agents.planning.PlanningAgent(model, planner=None, *, context=None)[source]#

Bases: object

An agent that keeps recent frames and plans each action in the world model.

Parameters:

model (WorldModel) – Any world model with a reward head.
planner (Optional[CEMPlanner]) – How actions are chosen; CEMPlanner by default.
context (Optional[int]) – Frames kept; the model’s max_context (or its context_frames when that is unbounded), at least 1.

frames: Deque[ndarray]#

actions: Deque[Any]#

reset(observation)[source]#

Start an episode from its first observation.

Parameters:: observation (Any)
Return type:: None

act()[source]#

The action to take now.

Return type:: ndarray

observe(action, observation)[source]#

Record the action just taken and the observation it led to.

Parameters:

action (Any)
observation (Any)

Return type:

None

simulusion.agents.planning.expand_state(state, n)[source]#

Repeat a batch-of-one state n times along the batch dimension.

Parameters:

state (Dict[str, Tensor])
n (int)

Return type:

Dict[str, Tensor]

simulusion.agents.planning.run_episode(env, agent, *, max_steps=1000)[source]#

Play one episode of a Gymnasium-style env with agent; returns its return.

Parameters:

env (Any)
agent (PlanningAgent)
max_steps (int)

Return type:

float

Models sub-module - Core world model implementations.

Exported Components:

Agents (High-level training wrappers):

DreamerAgent: High-level Dreamer training API
JEPAAgent: JEPA agent for self-supervised learning
Planet: PlaNet planning agent
VisionTransformer: Vision Transformer for image encoding
ModularRSSM: Modular RSSM with swappable components
Genie: Generative Interactive Environment model

Core Models:

Dreamer: Core Dreamer implementation with RSSM, actor, critic
RSSM: Recurrent State-Space Model (Dreamer-style)
RecurrentStateSpaceModel: PlaNet-style RSSM
LatentActionModel: Latent action learning for Genie
DynamicsModel: Future frame prediction for Genie

Factory Functions:

create_genie, create_genie_small, create_genie_large
create_modular_rssm
create_latent_action_model, create_dynamics_model

Small, import-safe catalog of available environments and backends.

This module replaces the previous simulusion.ui.catalog and is safe to import from lightweight CLI tools and tests without pulling in any UI dependencies.

Model catalog#

Core model families#

Key classes: Dreamer, DreamerAgent, RSSM, RecurrentStateSpaceModel, Planet, ModularRSSM, JEPAAgent, VisionTransformer, IRISAgent, IRISTransformer, IRISWorldModel, Genie, LatentActionModel, and DynamicsModel.

simulusion.models.dreamer.get_available_memory()[source]#

Get available physical memory in bytes.

Return type:: int

simulusion.models.dreamer.make_env(args)[source]#

Construct a Dreamer-compatible environment from DreamerConfig options.

Supports DMC, DMLab, Gym/Gymnasium, MuJoCo, Gymnasium Robotics, Procgen, Brax, BSuite, and Unity ML-Agents backends and applies the standard wrapper stack: action repeat, action normalization, and time limit.

Parameters:: args (Any)
Return type:: Any

simulusion.models.dreamer.preprocess_obs(obs)[source]#

Convert raw uint8 image observations to Dreamer float input space.

Images are scaled from [0, 255] to roughly [-0.5, 0.5], matching the normalization expected by Dreamer encoders.

Parameters:: obs (Tensor)
Return type:: Tensor

class simulusion.models.dreamer.Dreamer(args, obs_shape, action_size, device, restore=False)[source]#

Bases: object

Core Dreamer training system combining world model, actor, and value nets.

This class owns model construction, replay sampling, imagination rollouts, loss computation, optimization steps, evaluation loops, and checkpoint I/O.

Parameters:

args (Any)
obs_shape (Any)
action_size (int)
device (device | str)
restore (bool)

metrics: dict[str, float]#

classmethod from_config(config=None, *, obs_shape=None, action_size=None, device=None, restore=None, **overrides)[source]#

Build a core Dreamer model from a config object, dict, or YAML file.

obs_shape and action_size may be supplied directly. When either is omitted, this method constructs a temporary environment from the config to infer the model shapes.

Parameters:

config (DreamerConfig | dict[str, Any] | str | Path | None)
obs_shape (tuple[int, ...] | None)
action_size (int | None)
device (str | device | None)
restore (bool | None)
overrides (Any)

Return type:

Load a Dreamer checkpoint from a local path/directory or the HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
config (DreamerConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
map_location (str | device | None)
overrides (Any)

Return type:

parameter_count(trainable_only=False)[source]#

Return the total number of parameters owned by the Dreamer modules.

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]#

Return a compact parameter-count summary for the Dreamer modules.

Return type:: dict[str, Any]

set_precision(precision)[source]#

Train at precision from now on (fresh loss scalers).

Parameters:: precision (MixedPrecision)
Return type:: None

update_world_model(obs, acs, rews, nonterms, continues=None)[source]#

One world-model optimizer step; returns (loss, grad_norm).

Parameters:

obs (Tensor)
acs (Tensor)
rews (Tensor)
nonterms (Tensor)
continues (Tensor | None)

Return type:

tuple[Tensor, Tensor | None]

world_model_loss(obs, acs, rews, nonterms, continues=None)[source]#

World-model loss.

nonterms masks recurrent state at every episode boundary. continues is the discount-head target, which is 0 only on true terminations; it defaults to nonterms when not supplied.

Parameters:

obs (Tensor)
acs (Tensor)
rews (Tensor)
nonterms (Tensor)
continues (Tensor | None)

Return type:

Tensor

actor_loss()[source]#

Return type:: Tensor

value_loss()[source]#

Return type:: Tensor

train_one_batch()[source]#

Return type:: list[float]

act_with_world_model(obs, prev_state, prev_action, explore=False)[source]#

Parameters:

obs (Any)
prev_state (Any)
prev_action (Tensor)
explore (bool)

Return type:

tuple

act_and_collect_data(env, collect_steps)[source]#

Parameters:

env (Any)
collect_steps (int)

Return type:

ndarray

evaluate(env, eval_episodes, render=False)[source]#

Parameters:

env (Any)
eval_episodes (int)
render (bool)

Return type:

tuple

collect_random_episodes(env, seed_steps)[source]#

Parameters:

env (Any)
seed_steps (int)

Return type:

ndarray

save(save_path)[source]#

Parameters:: save_path (str)
Return type:: None

restore_checkpoint(ckpt_path, map_location=None)[source]#

Parameters:

ckpt_path (str | Path)
map_location (Any)

Return type:

None

class simulusion.models.dreamer.DreamerAgent(config=None, **kwargs)[source]#

High-level user API for running Dreamer experiments end to end.

It builds environments from config, initializes seeds and logging, instantiates Dreamer, and exposes simple train() / evaluate() methods.

Parameters:

config (Any)
kwargs (Any)

core_cls#

The world model + actor-critic this agent trains. Versioned agents (DreamerV1Agent, DreamerV2Agent) override it.

alias of Dreamer

algo_name: str | None = None#: Run-folder name used when the config leaves algo at its default.

classmethod from_config(config=None, **overrides)[source]#

Build a high-level Dreamer agent from a config object, dict, or YAML file.

Parameters:

config (DreamerConfig | dict[str, Any] | str | Path | None)
overrides (Any)

Return type:

Create a Dreamer agent and restore weights from a local path or HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
config (DreamerConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
map_location (str | device | None)
overrides (Any)

Return type:

parameter_count(trainable_only=False)[source]#

Return the total number of Dreamer parameters.

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]#

Return a compact parameter-count summary for the wrapped Dreamer model.

Return type:: dict[str, Any]

train(total_steps=None)[source]#

Train for total_steps environment steps (args.total_steps by default).

Parameters:: total_steps (int | None)
Return type:: None

evaluate()[source]#

Return type:: tuple

class simulusion.models.dreamer_rssm.RSSM(action_size, stoch_size, deter_size, hidden_size, obs_embed_size, activation)[source]#

Bases: ExportableAgentMixin, Module

Recurrent State-Space Model used by Dreamer for latent dynamics learning.

The RSSM is the core world model component that learns compact representations of environment dynamics. It maintains a hybrid state consisting of:

Deterministic State (h) – A recurrent hidden state updated by a GRU, capturing sequential/temporal information and deterministic transitions.
Stochastic State (s) – A latent variable representing stochastic, multi-modal uncertainty in the environment (e.g., ambiguous observations).

The model operates in two modes:

Observe Mode – Updates states using actual observations from the environment. Uses the representation model: p(s_t | h_t, obs_t)
Imagine Mode – Predicts future states without observations. Uses the transition/prior model: p(s_t | h_t)

Architecture

Input: Previous state (h_{t-1}, s_{t-1}) and action a_{t-1}
Process: GRU updates deterministic state, MLP computes stochastic prior/posterior
Output: Updated state (h_t, s_t) and distributions

State Representation

deter (h): GRU hidden state, captures sequential context
stoch (s): Stochastic latent, multi-modal uncertainty
mean/std: Parameters of the stochastic distribution

Usage with DreamerAgent:

rssm = RSSM(
    action_size=action_dim,
    stoch_size=30,      # Stochastic state dimension
    deter_size=200,     # Deterministic (GRU) state dimension
    hidden_size=200,    # MLP hidden layer size
    obs_embed_size=256,  # Observation embedding from encoder
    activation='elu'
)

# Observe with actual observation
posterior, prior = rssm.observe_step(prev_state, prev_action, obs_embed)

# Imagine future without observation
prior = rssm.imagine_step(current_state, action)

Training

The RSSM is trained by maximizing the ELBO (Evidence Lower Bound):

KL divergence between prior and posterior encourages the prior to capture environment dynamics
Reconstruction loss from decoder ensures state captures observation info

Reference:: Dreamer: Scalable Reinforcement Learning Using World Models Hafner et al., 2020 - https://arxiv.org/abs/1912.01603

Parameters:

action_size (int)
stoch_size (int)
deter_size (int)
hidden_size (int)
obs_embed_size (int)
activation (str)

init_state(batch_size, device)[source]#

Initialize RSSM state with zeros.

Parameters:

batch_size (int) – Number of parallel sequences
device (device) – torch device for tensors

Returns:

mean, std: Stochastic distribution parameters
stoch: Stochastic state sample
deter: Deterministic GRU hidden state

Return type:

Dictionary containing zero-initialized state components

get_dist(mean, std)[source]#

Create an Independent Normal distribution from mean and std.

Parameters:

mean (Tensor) – Location parameter
std (Tensor) – Scale parameter

Returns:

Independent Normal distribution with given parameters

Return type:

Independent

observe_step(prev_state, prev_action, obs_embed, nonterm=tensor(1.), noise=None)[source]#

Update state using actual observation (observe mode).

In observe mode, the RSSM first computes a transition prior from the previous state and action, then refines the stochastic state using the actual observation embedding to form the posterior.

Parameters:

prev_state (dict) – Dictionary with ‘deter’ (h_{t-1}) and ‘stoch’ (s_{t-1})
prev_action (Tensor) – Previous action a_{t-1}, shape (B, action_size)
obs_embed (Tensor) – Observation embedding from encoder, shape (B, obs_embed_size)
nonterm (Tensor) – Termination mask (1.0 = continue, 0.0 = terminal)
noise (Tuple[Tensor, Tensor] | None) – Optional (prior_noise, posterior_noise) standard-normal tensors of shape (B, stoch_size) used for the two samples instead of drawing them internally. Passing noise makes the step a pure function of its inputs, which exported graphs and eager-vs-exported parity checks need.

Returns:

Return type:

Tuple[dict, dict]

imagine_step(prev_state, prev_action, nonterm=tensor(1.), noise=None)[source]#

Predict next state without observation (imagine mode).

In imagine mode, the RSSM predicts future states using only the prior distribution. This is used for planning and policy learning where actual observations are not available.

Parameters:

prev_state (dict) – Dictionary with ‘deter’ (h_{t-1}) and ‘stoch’ (s_{t-1})
prev_action (Tensor) – Previous action a_{t-1}, shape (B, action_size)
nonterm (Tensor) – Termination mask (1.0 = continue, 0.0 = terminal)
noise (Tensor | None) – Optional standard-normal tensor (B, stoch_size) used for the stochastic sample instead of drawing one internally.

Returns:

deter: Predicted deterministic state
mean, std, stoch: Prior stochastic state distribution

Return type:

Dictionary with predicted state containing

get_prior(prev_state, prev_action, nonterm=tensor(1.))[source]#

Parameters:

prev_state (dict)
prev_action (Tensor)
nonterm (Tensor)

Return type:

dict

get_posterior(prev_state, prev_action, obs_embed, nonterm=tensor(1.))[source]#

Compute posterior distribution over stochastic state.

The posterior incorporates observation information to produce a more accurate state estimate.

Parameters:

prev_state (dict) – Previous state dictionary
prev_action (Tensor) – Previous action
obs_embed (Tensor) – Observation embedding
nonterm (Tensor) – Termination mask

Returns:

Dictionary with posterior state (observation-informed). Note that the previous-state shape (B, ...) is preserved; the batch dimension is not flattened.

Return type:

dict

detach_state(state)[source]#

Detach state tensors from computation graph.

Used during DreamerV2 training to prevent gradient flow through the observation/update pathway.

Parameters:: state (dict) – State dictionary with tensor values
Returns:: Detached state dictionary
Return type:: dict

seq_to_batch(state_dict)[source]#

Convert sequence state to batch format.

Parameters:: state_dict (dict) – Dictionary with sequence-dimension tensors (T, B, …)
Returns:: Dictionary with batch-dimension tensors (B*T, …)
Return type:: dict

observe_rollout(obs_embed, actions, nonterms, init_state, seq_len)[source]#

Process a sequence of observations (observe mode rollout).

Parameters:

obs_embed (Tensor) – Observation embeddings, shape (T+1, B, obs_embed_size)
actions (Tensor) – Actions, shape (T, B, action_size)
nonterms (Tensor) – Non-termination flags, shape (T, B, 1)
init_state (dict) – Initial state dictionary
seq_len (int) – Sequence length T

Returns:

Dictionary with prior states stacked along the time axis posterior: Dictionary with posterior states stacked along the time axis

Return type:

prior

imagine_rollout(policy, init_state, horizon)[source]#

Generate imagined trajectory using policy (imagine mode rollout).

Parameters:

policy (Module) – Actor network that outputs actions from state features
init_state (dict) – Initial state dictionary
horizon (int) – Number of steps to imagine

Returns:

Dictionary with imagined states for each step

Return type:

dict

forward(x, u)[source]#

Forward pass for training (computes sequence of states).

Parameters:

x (Tensor) – Observations, shape (B, T+1, C, H, W)
u (Tensor) – Actions, shape (B, T, action_size)

Returns:

List of state dictionaries for each timestep priors: List of prior distributions (tuples of mean, std) posteriors: List of posterior distributions (tuples of mean, std)

Return type:

states

class simulusion.models.rssm.RecurrentStateSpaceModel(action_size, state_size=200, latent_size=30, hidden_size=200, embed_size=1024, activation_function='relu')[source]#

Bases: ExportableAgentMixin, Module

A Recurrent State Space Model (RSSM) for modeling latent dynamics in sequential data.

Parameters:

action_size (int)
state_size (int)
latent_size (int)
hidden_size (int)
embed_size (int)
activation_function (str)

get_init_state(enc, h_t=None, s_t=None, a_t=None, mean=None, *, sample=None)[source]#

Returns the next deterministic state and the posterior latent.

Parameters:

enc (Tensor)
h_t (Tensor | None)
s_t (Tensor | None)
a_t (Tensor | None)
mean (bool | None)
sample (bool | None)

Return type:

tuple[Tensor, Tensor]

deterministic_state_fwd(h_t, s_t, a_t)[source]#

Deterministic transition update.

Ensures a_t is 2D and matches batch dimension of h_t before concatenation. Accepts a_t shaped [B, action_size], [action_size] (expanded to [B, action_size]), or [B]/scalar (reshaped appropriately).

Parameters:

h_t (Tensor)
s_t (Tensor)
a_t (Tensor)

Return type:

Tensor

state_prior(h_t, sample=False)[source]#

Returns the prior distribution over the latent state given the deterministic state

Parameters:

h_t (Tensor)
sample (bool)

Return type:

tuple[Tensor, Tensor] | Tensor

state_posterior(h_t, e_t, sample=False)[source]#

Returns the state prior given the deterministic state and obs

Parameters:

h_t (Tensor)
e_t (Tensor)
sample (bool)

Return type:

tuple[Tensor, Tensor] | Tensor

pred_reward(h_t, s_t)[source]#

Parameters:

h_t (Tensor)
s_t (Tensor)

Return type:

Tensor

rollout_prior(act, h_t, s_t)[source]#

Parameters:

act (Tensor)
h_t (Tensor)
s_t (Tensor)

Return type:

tuple[Tensor, Tensor]

forward(x, u)[source]#

Forward through the RSSM for a batch of sequences.

Parameters:

x (Tensor) – Tensor [B, T+1, C, H, W] (observations including initial frame)
u (Tensor) – Tensor [B, T, action_size] (actions for T steps)

Returns:

list[T] of tensors [B, state_size] priors: list[T] of tuples (mean, std) each [B, latent_size] posteriors: list[T] of tuples (mean, std) each [B, latent_size]

Return type:

states

class simulusion.models.planet.Planet(env, bit_depth=5, device=None, state_size=200, latent_size=30, embedding_size=1024, memory_size=100, policy_cfg=None, headless=False, max_episode_steps=None, action_repeats=1, results_dir=None)[source]#

High-level Planet wrapper.

Usage example:: from simulusion.models.planet import Planet p = Planet(env=’CartPole-v1’, bit_depth=5) p.train(epochs=50)

Parameters:

env (Any)
bit_depth (int)
device (device | None)
state_size (int)
latent_size (int)
embedding_size (int)
memory_size (int)
policy_cfg (dict | None)
headless (bool)
max_episode_steps (int | None)
action_repeats (int)
results_dir (str | None)

warmup(n_episodes=1, random_policy=True)[source]#

Collect n_episodes of rollouts into memory (used as warmup).

Parameters:

n_episodes (int)
random_policy (bool)

Return type:

None

Train for epochs epochs; returns the results directory.

Parameters:

scheduler_type (str) – Type of scheduler to use (“step”, “cosine”, “exponential”, “plateau”, None)
scheduler_kwargs (dict) – Additional arguments for the scheduler
epochs (int)
steps_per_epoch (int)
batch_size (int)
H (int)
beta (float)
save_every (int)
record_grads (bool)
results_dir (str | None)

Return type:

str

simulusion.models.planet.make_planet_scheduler(optimizer, scheduler_type, scheduler_kwargs, epochs)[source]#

The per-epoch learning-rate scheduler Planet.train uses.

scheduler_type is "step" (halve every 50 epochs by default), "cosine", "exponential", "plateau" or None / "none".

Parameters:

optimizer (Optimizer)
scheduler_type (str | None)
scheduler_kwargs (dict | None)
epochs (int)

Return type:

Any

Mixture Density Recurrent Neural Network (MDRNN) model implementation.

This module provides implementations of MDRNN models for world modeling. The MDRNN is used to predict future latent states given current latent states and actions, using a Gaussian Mixture Model (GMM) for the output.

Reference:: Ha & Schmidhuber (2018). Recurrent World Models Facilitate Policy Evolution. https://arxiv.org/abs/1805.11111

class simulusion.models.mdrnn.MDRNN(latents, actions, hiddens, gaussians)[source]#

Bases: _MDRNNBase

MDRNN model for multi-step sequence prediction.

This model processes entire sequences of latent states and actions, predicting the next latent state using a Gaussian Mixture Model (GMM). It also predicts rewards and terminal states.

Parameters:

latents (int) – Dimensionality of latent space (input and output).
actions (int) – Dimensionality of action space.
hiddens (int) – Number of hidden units in LSTM.
gaussians (int) – Number of Gaussian components in GMM output.

Example

>>> mdrnn = MDRNN(latents=32, actions=3, hiddens=256, gaussians=5)
>>> actions = torch.randn(10, 4, 3)  # seq_len, batch, action
>>> latents = torch.randn(10, 4, 32)  # seq_len, batch, latent
>>> mus, sigmas, logpi, rs, ds = mdrnn(actions, latents)
>>> # mus.shape = (10, 4, 5, 32)

forward(actions, latents)[source]#

Multi-step forward pass through the MDRNN.

Parameters:

actions (Tensor) – (SEQ_LEN, BSIZE, ASIZE) Tensor of actions.
latents (Tensor) – (SEQ_LEN, BSIZE, LSIZE) Tensor of latent states.

Returns:

mus: (SEQ_LEN, BSIZE, N_GAUSS, LSIZE) GMM means
sigmas: (SEQ_LEN, BSIZE, N_GAUSS, LSIZE) GMM standard deviations
logpi: (SEQ_LEN, BSIZE, N_GAUSS) log GMM weights
rs: (SEQ_LEN, BSIZE) predicted rewards
ds: (SEQ_LEN, BSIZE) predicted terminal state logits

Return type:

Tuple of

get_init_hidden(batch_size=1)[source]#

Return initial hidden state for the LSTM.

Parameters:: batch_size (int) – Number of sequences in the batch.
Returns:: Tuple of (h, c) with shapes (batch_size, hiddens).
Return type:: tuple[Tensor, Tensor]

class simulusion.models.mdrnn.MDRNNCell(latents, actions, hiddens, gaussians)[source]#

Bases: _MDRNNBase

MDRNN model for single-step forward prediction.

This model processes a single step of latent state and action,

This model processes a single step of latent state and action, predicting the next latent state using a Gaussian Mixture Model (GMM). It also predicts rewards and terminal states. Useful for real-time inference.

Parameters:

latents (int) – Dimensionality of latent space (input and output).
actions (int) – Dimensionality of action space.
hiddens (int) – Number of hidden units in LSTMCell.
gaussians (int) – Number of Gaussian components in GMM output.

Example

>>> cell = MDRNNCell(latents=32, actions=3, hiddens=256, gaussians=5)
>>> action = torch.randn(4, 3)  # batch, action
>>> latent = torch.randn(4, 32)  # batch, latent
>>> hidden = (torch.randn(4, 256), torch.randn(4, 256))
>>> mus, sigmas, logpi, r, d, next_hidden = cell(action, latent, hidden)

forward(action, latent, hidden)[source]#

Single-step forward pass through the MDRNN cell.

Parameters:

action (Tensor) – (BSIZE, ASIZE) Tensor of actions for current batch.
latent (Tensor) – (BSIZE, LSIZE) Tensor of latent states for current batch.
hidden (Tuple[Tensor, Tensor]) – Tuple of (h, c) hidden states for LSTMCell.

Returns:

mus: (BSIZE, N_GAUSS, LSIZE) GMM means
sigmas: (BSIZE, N_GAUSS, LSIZE) GMM standard deviations
logpi: (BSIZE, N_GAUSS) log GMM weights
r: (BSIZE,) predicted rewards
d: (BSIZE,) predicted terminal state logits
next_hidden: Tuple of (h, c) next hidden states

Return type:

Tuple of

get_init_hidden(batch_size=1)[source]#

Parameters:: batch_size (int)
Return type:: Tuple[Tensor, Tensor]

Linear Controller for World Models.

This module provides a simple linear controller that maps latent states and recurrent hidden states to actions. The controller is trained using CMA-ES (Covariance Matrix Adaptation Evolution Strategy).

Reference:: Ha & Schmidhuber (2018). Recurrent World Models Facilitate Policy Evolution. https://arxiv.org/abs/1805.11111

class simulusion.models.controller.Controller(latent_size, hidden_size, action_size)[source]#

Bases: Module

Linear controller that maps latent + hidden state to actions.

This is a simple linear controller that takes the latent state and recurrent hidden state as input and outputs actions. It is trained separately from the world model using black-box optimization (CMA-ES).

Variables:

latent_size – Dimensionality of latent state from VAE.
hidden_size – Dimensionality of RSSM hidden state.
action_size – Dimensionality of action space.

Parameters:

latent_size (int)
hidden_size (int)
action_size (int)

Example

>>> controller = Controller(latent_size=32, hidden_size=200, action_size=3)
>>> state = torch.cat([latent, hidden], dim=-1)
>>> action = controller(state)

forward(state)[source]#

Compute actions from latent and hidden states.

Parameters:: state (Tensor) – Concatenated [latent, hidden] state tensor.
Returns:: Action tensor of shape (batch, action_size).
Return type:: Tensor

Modular RSSM with swappable encoder/decoder/backbone components.

This module provides a flexible architecture for world model research, allowing researchers to easily swap different encoder, decoder, and backbone implementations for ablations and experimentation.

class simulusion.models.modular_rssm.EncoderBase(*args, **kwargs)[source]#

Bases: Module, ABC

Abstract base class for observation encoders.

Parameters:

args (Any)
kwargs (Any)

abstractmethod forward(obs)[source]#

Encode observations to embeddings.

Parameters:: obs (Tensor)
Return type:: Tensor

embed_size: int#

get_embed_size()[source]#

Return the embedding size. Override in subclasses.

Return type:: int

class simulusion.models.modular_rssm.DecoderBase(*args, **kwargs)[source]#

Bases: Module, ABC

Abstract base class for observation decoders.

Parameters:

args (Any)
kwargs (Any)

abstractmethod forward(features)[source]#

Decode latent features to observation distributions.

Parameters:: features (Tensor)
Return type:: Any

class simulusion.models.modular_rssm.BackboneBase(*args, **kwargs)[source]#

Bases: Module, ABC

Abstract base class for recurrent dynamics backbones.

Parameters:

args (Any)
kwargs (Any)

abstractmethod forward(state, action, obs_embed=None, nonterm=1.0)[source]#

Process one step of dynamics. Returns (prior, posterior).

Parameters:

state (Dict[str, Tensor])
action (Tensor)
obs_embed (Tensor | None)
nonterm (float)

Return type:

Tuple[Dict[str, Tensor], Dict[str, Tensor]]

abstractmethod init_state(batch_size, device)[source]#

Initialize hidden state.

Parameters:

batch_size (int)
device (device)

Return type:

Dict[str, Tensor]

stoch_size: int#

deter_size: int#

class simulusion.models.modular_rssm.ConvEncoder(input_shape, embed_size, activation='elu', depth=32)[source]#

Bases: EncoderBase

Convolutional encoder from Dreamer (image observations).

Parameters:

input_shape (Tuple[int, int, int])
embed_size (int)
activation (str)
depth (int)

forward(obs)[source]#

Parameters:: obs (Tensor)
Return type:: Tensor

class simulusion.models.modular_rssm.MLPEncoder(input_dim, embed_size, hidden_sizes=[256, 256], activation='elu')[source]#

Bases: EncoderBase

MLP encoder for state-based observations.

Parameters:

input_dim (int)
embed_size (int)
hidden_sizes (List[int])
activation (str)

forward(obs)[source]#

Parameters:: obs (Tensor)
Return type:: Tensor

class simulusion.models.modular_rssm.ViTEncoder(input_shape, embed_size, patch_size=8, depth=6, num_heads=8, mlp_ratio=4.0, activation='gelu')[source]#

Bases: EncoderBase

Vision Transformer encoder for image observations.

Parameters:

input_shape (Tuple[int, int, int])
embed_size (int)
patch_size (int)
depth (int)
num_heads (int)
mlp_ratio (float)
activation (str)

forward(obs)[source]#

Parameters:: obs (Tensor)
Return type:: Tensor

class simulusion.models.modular_rssm.TransformerBlock(embed_size, num_heads, mlp_ratio, activation)[source]#

Bases: Module

Transformer block for ViT encoder.

Parameters:

embed_size (int)
num_heads (int)
mlp_ratio (float)
activation (str)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.models.modular_rssm.ConvDecoder(stoch_size, deter_size, output_shape, activation='elu', depth=32)[source]#

Bases: DecoderBase

Convolutional decoder for image observations.

Parameters:

stoch_size (int)
deter_size (int)
output_shape (Tuple[int, int, int])
activation (str)
depth (int)

forward(features)[source]#

Parameters:: features (Tensor)
Return type:: Any

class simulusion.models.modular_rssm.MLPDecoder(stoch_size, deter_size, output_dim, hidden_sizes=[256, 256], activation='elu', dist='normal')[source]#

Bases: DecoderBase

MLP decoder for state-based observations.

Parameters:

stoch_size (int)
deter_size (int)
output_dim (int)
hidden_sizes (List[int])
activation (str)
dist (str)

forward(features)[source]#

Parameters:: features (Tensor)
Return type:: Any

class simulusion.models.modular_rssm.GRUBackbone(action_size, stoch_size, deter_size, hidden_size, embed_size, activation='elu')[source]#

Bases: BackboneBase

GRU-based recurrent dynamics backbone (standard RSSM).

Parameters:

action_size (int)
stoch_size (int)
deter_size (int)
hidden_size (int)
embed_size (int)
activation (str)

property embedding_size: int#

init_state(batch_size, device)[source]#

Parameters:

batch_size (int)
device (device)

Return type:

Dict[str, Tensor]

forward(state, action, obs_embed=None, nonterm=1.0)[source]#

Parameters:

state (Dict[str, Tensor])
action (Tensor)
obs_embed (Tensor | None)
nonterm (float)

Return type:

Tuple[Dict[str, Tensor], Dict[str, Tensor]]

class simulusion.models.modular_rssm.LSTMBackbone(action_size, stoch_size, deter_size, hidden_size, embed_size, activation='elu')[source]#

Bases: BackboneBase

LSTM-based recurrent dynamics backbone.

Parameters:

action_size (int)
stoch_size (int)
deter_size (int)
hidden_size (int)
embed_size (int)
activation (str)

property embedding_size: int#

init_state(batch_size, device)[source]#

Parameters:

batch_size (int)
device (device)

Return type:

Dict[str, Tensor]

forward(state, action, obs_embed=None, nonterm=1.0)[source]#

Parameters:

state (Dict[str, Tensor])
action (Tensor)
obs_embed (Tensor | None)
nonterm (float)

Return type:

Tuple[Dict[str, Tensor], Dict[str, Tensor]]

class simulusion.models.modular_rssm.TransformerBackbone(action_size, stoch_size, deter_size, embed_size, num_heads=4, num_layers=2, activation='gelu')[source]#

Bases: BackboneBase

Transformer-based dynamics backbone for long-range dependencies.

Parameters:

action_size (int)
stoch_size (int)
deter_size (int)
embed_size (int)
num_heads (int)
num_layers (int)
activation (str)

property embedding_size: int#

init_state(batch_size, device)[source]#

Parameters:

batch_size (int)
device (device)

Return type:

Dict[str, Tensor]

forward(state, action, obs_embed=None, nonterm=1.0)[source]#

Parameters:

state (Dict[str, Tensor])
action (Tensor)
obs_embed (Tensor | None)
nonterm (float)

Return type:

Tuple[Dict[str, Tensor], Dict[str, Tensor]]

class simulusion.models.modular_rssm.ModularRSSM(encoder, decoder, backbone, reward_decoder=None)[source]#

Bases: ExportableAgentMixin, Module

Modular RSSM with swappable encoder, decoder, and backbone.

This class allows researchers to easily experiment with different: - Encoders: Conv, MLP, ViT - Decoders: Conv, MLP - Backbones: GRU, LSTM, Transformer

Example

>>> encoder = ConvEncoder((3, 64, 64), embed_size=1024)
>>> decoder = ConvDecoder(32, 200, (3, 64, 64))
>>> backbone = GRUBackbone(action_size=6, stoch_size=32, deter_size=200, hidden_size=200, embed_size=1024)
>>> rssm = ModularRSSM(encoder, decoder, backbone)

Parameters:

encoder (EncoderBase)
decoder (DecoderBase)
backbone (BackboneBase)
reward_decoder (DecoderBase | None)

property stoch_size: int#

property deter_size: int#

property embed_size: int#

init_state(batch_size, device)[source]#

Parameters:

batch_size (int)
device (device)

Return type:

Dict[str, Tensor]

get_dist(mean, std)[source]#

Parameters:

mean (Tensor)
std (Tensor)

Return type:

Distribution

observe_step(prev_state, prev_action, obs, nonterm=1.0)[source]#

Encode obs and advance one step.

Returns:

Return type:

(prior, posterior). Note the order

Parameters:

prev_state (Dict[str, Tensor])
prev_action (Tensor)
obs (Tensor)
nonterm (Any)

imagine_step(prev_state, prev_action, nonterm=1.0)[source]#

Parameters:

prev_state (Dict[str, Tensor])
prev_action (Tensor)
nonterm (Any)

Return type:

Dict[str, Tensor]

observe_rollout(obs, actions, nonterms, prev_state, horizon)[source]#

Parameters:

obs (Tensor)
actions (Tensor)
nonterms (Tensor)
prev_state (Dict[str, Tensor])
horizon (int)

Return type:

Tuple[Dict[str, Tensor], Dict[str, Tensor]]

imagine_rollout(actor, prev_state, horizon)[source]#

Parameters:

actor (Module)
prev_state (Dict[str, Tensor])
horizon (int)

Return type:

Dict[str, Tensor]

decode_observation(features)[source]#

Parameters:: features (Tensor)
Return type:: Tensor

decode_reward(features)[source]#

Parameters:: features (Tensor)
Return type:: Tensor

detach_state(state)[source]#

Parameters:: state (Dict[str, Tensor])
Return type:: Dict[str, Tensor]

seq_to_batch(state)[source]#

Parameters:: state (Dict[str, Tensor])
Return type:: Dict[str, Tensor]

simulusion.models.modular_rssm.create_modular_rssm(encoder_type='conv', decoder_type='conv', backbone_type='gru', obs_shape=(3, 64, 64), action_size=6, stoch_size=32, deter_size=200, embed_size=1024, hidden_size=200, activation='elu', **kwargs)[source]#

Factory function to create a modular RSSM with specified components.

Parameters:

encoder_type (str) – Type of encoder (“conv”, “mlp”, “vit”)
decoder_type (str) – Type of decoder (“conv”, “mlp”)
backbone_type (str) – Type of backbone (“gru”, “lstm”, “transformer”)
obs_shape (Tuple[int, int, int] | Tuple[int]) – Shape of observations (C, H, W) for images or (D,) for state
action_size (int) – Action space dimension
stoch_size (int) – Stochastic latent dimension
deter_size (int) – Deterministic hidden dimension
embed_size (int) – Encoder embedding dimension
hidden_size (int) – Hidden layer dimension
activation (str) – Activation function name
kwargs (Any)

Returns:

Configured ModularRSSM instance

Return type:

ModularRSSM

class simulusion.models.jepa_agent.JEPAAgent(config=None, **kwargs)[source]#

Convenience interface for configuring and launching JEPA training runs.

Accepts a JEPAConfig plus keyword overrides, prepares output folders, and delegates execution to the JEPA training entrypoint.

Parameters:

config (JEPAConfig | None)
kwargs (Any)

classmethod from_config(config=None, **overrides)[source]#

Build a JEPA agent from a config object, dict, YAML file, or YAML string.

Parameters:

config (JEPAConfig | dict[str, Any] | str | Path | None)
overrides (Any)

Return type:

classmethod from_pretrained(pretrained_model_name_or_path, *, config=None, checkpoint_filename=None, config_filename='config.yaml', repo_type=None, revision=None, **overrides)[source]#

Create a JEPA agent from local/HF Hub config and checkpoint metadata.

Parameters:

pretrained_model_name_or_path (str | Path)
config (JEPAConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
overrides (Any)

Return type:

parameter_count(trainable_only=False)[source]#

JEPA models are constructed inside training, so no parameters are resident.

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]#

Return the configured JEPA run metadata.

Return type:: dict[str, Any]

train()[source]#

Train for cfg.epochs epochs.

Return type:: None

simulusion.models.vit.get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False)[source]#

Generate fixed 2D sine/cosine positional embeddings on a square patch grid.

Returns NumPy embeddings used to initialize non-trainable transformer position encodings, with optional prepended class-token embedding.

Parameters:

embed_dim (int)
grid_size (int)
cls_token (bool)

Return type:

ndarray

simulusion.models.vit.get_2d_sincos_pos_embed_from_grid(embed_dim, grid)[source]#

Build 2D sine/cosine embeddings from precomputed meshgrid coordinates.

The final embedding concatenates independent encodings for vertical and horizontal coordinates.

Parameters:

embed_dim (int)
grid (ndarray)

Return type:

ndarray

simulusion.models.vit.get_1d_sincos_pos_embed(embed_dim, grid_size, cls_token=False)[source]#

Generate 1D sine/cosine positional embeddings for integer positions.

Useful for sequence-style positional encoding and as a building block for 2D embedding construction.

Parameters:

embed_dim (int)
grid_size (int)
cls_token (bool)

Return type:

ndarray

simulusion.models.vit.get_1d_sincos_pos_embed_from_grid(embed_dim, pos)[source]#

Generate 1D sine/cosine positional embeddings from explicit positions.

Positions are projected onto a log-frequency basis and encoded with sine and cosine components.

Parameters:

embed_dim (int)
pos (ndarray)

Return type:

ndarray

simulusion.models.vit.drop_path(x, drop_prob=0.0, training=False)[source]#

Apply stochastic depth (DropPath) regularization to residual branches.

Randomly drops entire residual paths per sample during training and scales the surviving activations to preserve expected magnitude.

Parameters:

x (Tensor)
drop_prob (float)
training (bool)

Return type:

Tensor

class simulusion.models.vit.DropPath(drop_prob=None)[source]#

Bases: Module

Module wrapper around the functional drop_path stochastic depth utility.

Parameters:: drop_prob (float | None)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.models.vit.MLP(in_features, hidden_features=None, out_features=None, act_layer=<class 'torch.nn.modules.activation.GELU'>, drop=0.0)[source]#

Bases: Module

Two-layer feed-forward network used inside transformer blocks.

Applies linear projection, activation, dropout, and output projection in the standard Vision Transformer MLP pattern.

Parameters:

in_features (int)
hidden_features (int | None)
out_features (int | None)
act_layer (type[Module])
drop (float)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.models.vit.Attention(dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0)[source]#

Bases: Module

Multi-head self-attention block for token sequences.

Computes QKV projections, scaled dot-product attention, and output projection with configurable dropout.

Parameters:

dim (int)
num_heads (int)
qkv_bias (bool)
qk_scale (float | None)
attn_drop (float)
proj_drop (float)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.models.vit.Block(dim, num_heads, mlp_ratio=4.0, qkv_bias=False, qk_scale=None, drop=0.0, attn_drop=0.0, drop_path=0.0, act_layer=<class 'torch.nn.modules.activation.GELU'>, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>)[source]#

Bases: Module

Transformer encoder block combining attention and MLP residual branches.

Each branch uses pre-normalization and optional stochastic depth.

Parameters:

dim (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
qk_scale (float | None)
drop (float)
attn_drop (float)
drop_path (float)
act_layer (type[Module])
norm_layer (type[Module])

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.models.vit.PatchEmbed(img_size=224, patch_size=16, in_chans=3, embed_dim=768)[source]#

Bases: Module

Image to Patch Embedding

Parameters:

img_size (int)
patch_size (int)
in_chans (int)
embed_dim (int)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.models.vit.ConvEmbed(channels, strides, img_size=224, in_chans=3, batch_norm=True)[source]#

Bases: Module

3x3 Convolution stems for ViT following ViTC models

Parameters:

channels (list[int])
strides (list[int])
img_size (int)
in_chans (int)
batch_norm (bool)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.models.vit.VisionTransformerPredictor(num_patches, embed_dim=768, predictor_embed_dim=384, depth=6, num_heads=12, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>, init_std=0.02, **kwargs)[source]#

Bases: Module

Vision Transformer

Parameters:

num_patches (int)
embed_dim (int)
predictor_embed_dim (int)
depth (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
qk_scale (float | None)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)
norm_layer (type[Module])
init_std (float)
kwargs (Any)

fix_init_weight()[source]#

Return type:: None

forward(x, masks_x, masks)[source]#

Parameters:

x (Tensor)
masks_x (Tensor | list[Tensor])
masks (Tensor | list[Tensor])

Return type:

Tensor

class simulusion.models.vit.VisionTransformer(img_size=[224], patch_size=16, in_chans=3, embed_dim=768, predictor_embed_dim=384, depth=12, predictor_depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>, init_std=0.02, **kwargs)[source]#

Bases: ExportableAgentMixin, Module

Vision Transformer

Parameters:

img_size (list[int])
patch_size (int)
in_chans (int)
embed_dim (int)
predictor_embed_dim (int)
depth (int)
predictor_depth (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
qk_scale (float | None)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)
norm_layer (type[Module])
init_std (float)
kwargs (Any)

fix_init_weight()[source]#

Return type:: None

forward(x, masks=None)[source]#

Parameters:

x (Tensor)
masks (Tensor | list[Tensor] | None)

Return type:

Tensor

get_intermediate_layers(x, n=1)[source]#

Return the patch tokens output by each of the last n blocks.

Parameters:

x (Tensor)
n (int)

Return type:

list[Tensor]

interpolate_pos_encoding(x, pos_embed)[source]#

Parameters:

x (Tensor)
pos_embed (Tensor)

Return type:

Tensor

simulusion.models.vit.vit_predictor(**kwargs)[source]#

Factory for a JEPA predictor transformer with sensible defaults.

Parameters:: kwargs (Any)
Return type:: VisionTransformerPredictor

simulusion.models.vit.vit_tiny(patch_size=16, **kwargs)[source]#

Factory for a tiny Vision Transformer encoder backbone.

Parameters:

patch_size (int)
kwargs (Any)

Return type:

Any

simulusion.models.vit.vit_small(patch_size=16, **kwargs)[source]#

Factory for a small Vision Transformer encoder backbone.

Parameters:

patch_size (int)
kwargs (Any)

Return type:

Any

simulusion.models.vit.vit_base(patch_size=16, **kwargs)[source]#

Factory for a base Vision Transformer encoder backbone.

Parameters:

patch_size (int)
kwargs (Any)

Return type:

Any

simulusion.models.vit.vit_large(patch_size=16, **kwargs)[source]#

Factory for a large Vision Transformer encoder backbone.

Parameters:

patch_size (int)
kwargs (Any)

Return type:

Any

simulusion.models.vit.vit_huge(patch_size=16, **kwargs)[source]#

Factory for a huge Vision Transformer encoder backbone.

Parameters:

patch_size (int)
kwargs (Any)

Return type:

Any

simulusion.models.vit.vit_giant(patch_size=16, **kwargs)[source]#

Factory for a giant Vision Transformer encoder backbone.

Parameters:

patch_size (int)
kwargs (Any)

Return type:

Any

simulusion.models.iris_agent.compute_lambda_return(rewards, values, discounts, lambda_coef=0.95)[source]#

Compute λ-return target for value function training.

Parameters:

rewards (Tensor) – Rewards (B, T)
values (Tensor) – Value estimates (B, T+1)
discounts (Tensor) – Discount factors (B, T)
lambda_coef (float) – Lambda parameter for bootstrapping

Returns:

λ-return targets (B, T)

Return type:

lambda_returns

class simulusion.models.iris_agent.IRISAgent(config, action_size, device)[source]#

Bases: ExportableAgentMixin, Module

Complete IRIS Agent with world model and policy.

Combines: - Discrete autoencoder (encoder + decoder) - Transformer world model - Actor-Critic for policy and value learning

Parameters:

config (IRISConfig)
action_size (int)
device (device)

use_cuda_graphs#: Replay imagination’s single-position decode steps as a CUDA graph. Same computation, far fewer kernel launches; only used on CUDA.

reconstruct_chunk_size#: Frames per autoencoder pass in reconstruct().

classmethod from_config(config=None, *, action_size, device=None, **overrides)[source]#

Build an IRIS agent from a config object, dict, YAML file, or YAML string.

Parameters:

config (IRISConfig | dict[str, Any] | str | Path | None)
action_size (int)
device (device | str | None)
overrides (Any)

Return type:

Load an IRIS agent checkpoint from a local path/directory or HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
action_size (int | None)
device (device | str | None)
config (IRISConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
overrides (Any)

Return type:

parameter_count(trainable_only=False)[source]#

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]#

Return type:: dict[str, Any]

set_precision(precision)[source]#

Train at precision from now on (fresh loss scalers).

Parameters:: precision (MixedPrecision)
Return type:: None

forward_actor_critic(frames, hidden=None)[source]#

Forward pass through actor-critic.

Parameters:

frames (Tensor) – Input frames (B, T, C, H, W)
hidden (Tuple[Tensor, Tensor] | None) – Optional LSTM hidden state

Returns:

(B, T, action_size) values: (B, T) hidden_state: (h, c)

Return type:

action_logits

act(frame, epsilon=0.0, temperature=1.0, hidden=None, return_hidden=False)[source]#

Sample action from policy.

Parameters:

frame (Tensor) – Single frame (B, C, H, W)
epsilon (float) – Random action probability
temperature (float) – Action distribution temperature
hidden (Tuple[Tensor, Tensor] | None) – LSTM state from the previous step, or None to start fresh
return_hidden (bool) – If True, also return the updated LSTM state

Returns:

Selected actions (B,), and the updated LSTM state when return_hidden is set.

Return type:

burn_in(frames)[source]#

Initialise the LSTM state by replaying preceding frames.

Parameters:: frames (Tensor) – Preceding observations (B, T_burn, C, H, W), already in the reconstruction domain the policy is trained on.
Returns:: The LSTM state after the burn-in, or None if no frames were given.
Return type:: Tuple[Tensor, Tensor] | None

transform_reward(rewards)[source]#

Apply the configured reward transform.

Parameters:: rewards (Tensor)
Return type:: Tensor

reconstruct(frames)[source]#

Pass frames through the discrete autoencoder: D(E(x)).

Parameters:: frames (Tensor) – Real observations (B, C, H, W) or (B, T, C, H, W) in [0, 1].
Returns:: Reconstructions with the same shape, clamped to [0, 1].
Return type:: Tensor

imagine_rollout(initial_frame, horizon=20, burn_in_frames=None, sample_tokens=True, temperature=1.0, stop_on_termination=True)[source]#

Generate imagined trajectories using the world model.

Parameters:

initial_frame (Tensor) – Starting frame x_0 (B, C, H, W), in [0, 1].
horizon (int) – Number of steps H to imagine.
burn_in_frames (Tensor | None) – The (B, T_burn, C, H, W) reconstructed observations preceding initial_frame, used to initialise the policy’s LSTM state (paper A.3). None starts from a zero state.
sample_tokens (bool) – Sample next-frame tokens instead of taking the argmax.
temperature (float) – Sampling temperature for token generation.
stop_on_termination (bool) – Cut the rollout short once every trajectory in the batch has hit a predicted episode end (paper 2.3). Set False to always return exactly horizon steps; correctness does not depend on it, since continues already zeroes the discount past a termination.

Returns:

dict with frames (B, T+1, C, H, W), actions (B, T), rewards (B, T) and continues (B, T). T may be shorter than horizon if every rollout predicted an episode end.

Return type:

trajectory

update_autoencoder(frames)[source]#

Update discrete autoencoder.

Parameters:: frames (Tensor) – Training frames (B, C, H, W)
Returns:: Dictionary of loss values
Return type:: losses

update_transformer(frames, actions, rewards, terminals)[source]#

Update transformer world model.

Parameters:

frames (Tensor) – Frame sequence
actions (Tensor) – Actions taken
rewards (Tensor) – Rewards received
terminals (Tensor) – Terminal flags

Returns:

Dictionary of loss values

Return type:

losses

update_actor_critic(imagined_trajectory)[source]#

Update actor-critic in imagination.

Parameters:: imagined_trajectory (dict) – Dictionary from imagine_rollout
Returns:: Dictionary of loss values
Return type:: losses

CHECKPOINT_FORMAT = 5#

save(path)[source]#

Save agent state.

Parameters:: path (str)
Return type:: None

load(path)[source]#

Load agent state.

Parameters:: path (str)
Return type:: None

class simulusion.models.iris_transformer.KVCache(num_layers, batch_size, num_heads, head_dim, max_len, device, dtype)[source]#

Bases: object

Per-layer key/value cache for incremental decoding.

Imagination generates one position at a time. Without a cache, producing the K tokens of a frame means re-running the whole Transformer over the entire prefix K times, which is O(K * L^2) per imagined step. Caching the keys and values of every position already processed reduces that to O(K * L).

Storage is a pre-allocated (B, num_heads, max_len, head_dim) buffer per layer, filled left to right; length marks the valid prefix.

Parameters:

num_layers (int)
batch_size (int)
num_heads (int)
head_dim (int)
max_len (int)
device (device)
dtype (dtype)

append(layer, k, v)[source]#

Append this step’s keys/values for layer and return the full prefix.

The write offset is self.length, which the caller advances once per forward pass (after all layers have been visited), so every layer writes to the same slice.

Parameters:

layer (int)
k (Tensor)
v (Tensor)

Return type:

Tuple[Tensor, Tensor]

advance(steps)[source]#

Commit steps newly written positions.

Parameters:: steps (int)
Return type:: None

reset()[source]#

Empty the cache while keeping its buffers.

Re-priming a reset cache gives the same result as priming a fresh one: positions past length are never read. Reusing the buffers keeps a GraphedDecodeStep captured against them valid.

Return type:: None

class simulusion.models.iris_transformer.CausalSelfAttention(embed_dim, num_heads, dropout=0.1)[source]#

Bases: Module

GPT-2 style multi-head causal self-attention with optional KV caching.

Parameters:

embed_dim (int)
num_heads (int)
dropout (float)

forward(x, cache=None, layer_idx=0)[source]#

Parameters:

x (Tensor)
cache (KVCache | None)
layer_idx (int)

Return type:

Tensor

class simulusion.models.iris_transformer.GPTBlock(embed_dim, num_heads, dropout=0.1)[source]#

Bases: Module

Pre-norm GPT-2 block: LN -> attention -> residual, LN -> MLP -> residual.

Parameters:

embed_dim (int)
num_heads (int)
dropout (float)

forward(x, cache=None, layer_idx=0)[source]#

Parameters:

x (Tensor)
cache (KVCache | None)
layer_idx (int)

Return type:

Tensor

class simulusion.models.iris_transformer.GraphedDecodeStep(model, cache, batch_size)[source]#

Bases: object

One single-position forward against a KVCache, as a CUDA graph.

Imagination decodes one position at a time, and each forward is about 160 small kernels. On a laptop GPU the CPU cost of launching them outweighs the GPU work several times over. Capturing the step once and replaying it launches it as one graph.

A graph needs fixed shapes, so instead of slicing the cache to its current length the step writes the new key/value at a position tensor and attends over the whole buffer with a mask hiding positions after it. Masked positions contribute exactly zero weight, so this computes the same function as IRISTransformer._forward_positions() for one position (inference only: dropout is not applied).

Capture writes to position 0 of cache, so build the step before priming the cache, and keep the same cache afterwards (KVCache.reset() instead of a new cache) so the captured buffers stay the ones in use.

Parameters:

model (IRISTransformer)
cache (KVCache)
batch_size (int)

class simulusion.models.iris_transformer.IRISTransformer(vocab_size=512, tokens_per_frame=16, action_size=18, embed_dim=256, num_layers=10, num_heads=4, dropout=0.1, gradient_checkpointing=False, reward_classes=3)[source]#

Bases: ExportableAgentMixin, Module

GPT-like autoregressive Transformer for world modeling.

Models the dynamics of the environment by predicting, autoregressively over an interleaved sequence of frame tokens and actions:

Next frame tokens (transition model), one token at a time
Rewards
Episode termination

The sequence layout for S frames and S - 1 actions is:

z_0^1 ... z_0^K, a_0, z_1^1 ... z_1^K, a_1, ..., z_{S-2}^1 ... z_{S-2}^K,
a_{S-2}, z_{S-1}^1 ... z_{S-1}^K

z_{t+1}^k ~ p(. | z_{<=t}, a_{<=t}, z_{t+1}^{<k})

Parameters:

vocab_size (int)
tokens_per_frame (int)
action_size (int)
embed_dim (int)
num_layers (int)
num_heads (int)
dropout (float)
gradient_checkpointing (bool)
reward_classes (int)

expected_reward(action_hidden)[source]#

Scalar reward prediction from an action-position hidden state.

Parameters:: action_hidden (Tensor) – (…, embed_dim) hidden states at action positions.
Returns:: (…) scalar reward predictions.
Return type:: Tensor

gpt_blocks()[source]#

The transformer blocks, typed (nn.ModuleList erases the element type).

Return type:: List[GPTBlock]

forward(tokens, actions)[source]#

Teacher-forced forward pass through the Transformer world model.

Parameters:

tokens (Tensor) – Frame tokens (B, S, K) for S consecutive frames.
actions (Tensor) – Actions (B, S-1); actions[:, t] is taken after frame t.

Returns:

Predictions of frames 1..S-1 (B, S-1, K, vocab_size). rewards: Predicted rewards r_0..r_{S-2} (B, S-1). terminations: Predicted terminations d_0..d_{S-2} (B, S-1, 2).

Return type:

token_logits

init_cache(batch_size, device, dtype=None, max_len=None)[source]#

Allocate an empty KV cache.

It holds max_len positions, by default this model’s full context capacity, and never more than it.

Parameters:

batch_size (int)
device (device)
dtype (dtype | None)
max_len (int | None)

Return type:

KVCache

prime_cache(tokens, actions, cache, start_pos=0)[source]#

Fill a cache with an interleaved (frame, action) history in one pass.

Parameters:

tokens (Tensor) – Frame token grids (B, Tc, K).
actions (Tensor | None) – Actions following each frame. Pass Tc actions to end the primed sequence on an action (ready to generate the next frame), or Tc - 1 / None to end on the last frame’s tokens.
cache (KVCache) – Cache to fill; must be empty or positioned at start_pos.
start_pos (int) – Absolute position of the first embedded token.

Returns:

The absolute position just past the primed sequence.

Return type:

int

generate_frame_cached(action, cache, start_pos, sample=True, temperature=1.0, decode_step=None)[source]#

Generate one frame’s K tokens incrementally, one position at a time.

Returns:

Per-token logits (B, K, vocab). generated: Token indices (B, K). action_hidden: Hidden state at the action position (B, E). next_pos: Absolute position after the generated frame’s tokens.

Return type:

step_logits

Parameters:

action (Tensor)
cache (KVCache)
start_pos (int)
sample (bool)
temperature (float)
decode_step (GraphedDecodeStep | None)

predict_next_tokens(tokens, actions)[source]#

Greedily predict the next frame tokens autoregressively.

Parameters:

tokens (Tensor) – Current frame tokens (B, K) or (B, H, W).
actions (Tensor) – Actions taken (B,).

Returns:

Next frame token logits (B, K, vocab_size). Their argmax: equals the greedily generated tokens.

action_hidden: Hidden states for reward prediction (B, embed_dim).

Return type:

token_logits

imagine_step(tokens, actions, sample=True, temperature=1.0)[source]#

Advance imagination one step: next frame tokens + the action hidden state.

Unlike predict_next_tokens() this can sample rather than take the argmax, and it returns the generated tokens directly so a rollout does not have to re-derive them from logits.

Parameters:

tokens (Tensor) – Current frame tokens (B, K) or (B, H, W).
actions (Tensor) – Actions taken (B,).
sample (bool) – If True, sample tokens from the predicted distribution.
temperature (float) – Sampling temperature (ignored when sample is False).

Returns:

Generated token indices (B, K). action_hidden: Hidden state at the action position (B, embed_dim),

the input expected by reward_head / termination_head.

Return type:

next_tokens

sample_next_tokens(tokens, actions, temperature=1.0)[source]#

Sample next tokens autoregressively from the distribution.

Parameters:

tokens (Tensor) – Current frame tokens (B, K) or (B, H, W).
actions (Tensor) – Actions taken (B,).
temperature (float) – Sampling temperature (higher = more random).

Returns:

Sampled token indices (B, K). log_probs: Log probabilities of sampled tokens (B, K).

Return type:

sampled_tokens

class simulusion.models.iris_transformer.IRISWorldModel(encoder, decoder, transformer)[source]#

Bases: ExportableAgentMixin, Module

Complete IRIS World Model combining autoencoder and transformer.

This is the core component that learns environment dynamics entirely in the “imaginary” latent space.

Parameters:

encoder (Module)
decoder (Module)
transformer (IRISTransformer)

decode_tokens(indices)[source]#

Decode token indices (B, K) or (B, H, W) to images.

Parameters:: indices (Tensor)
Return type:: Tensor

forward(observations, actions)[source]#

Full world model forward pass.

Parameters:

observations (Tensor) – Image sequence (B, T+1, C, H, W)
actions (Tensor) – Actions (B, T)

Returns:

Dictionary with predicted tokens, rewards, terminations losses: Dictionary with loss components

Return type:

predictions

imagine(initial_tokens, policy, horizon=20, temperature=1.0)[source]#

Generate imagined trajectories.

Parameters:

initial_tokens (Tensor) – Initial frame tokens (B, K)
policy (Module) – Policy network to sample actions
horizon (int) – Number of steps to imagine
temperature (float) – Sampling temperature for token prediction

Returns:

Dictionary with imagined trajectories

Return type:

imagined

class simulusion.models.genie.Genie(num_frames=16, image_size=64, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, tokenizer_encoder_dim=512, tokenizer_decoder_dim=1024, action_vocab_size=8, action_embedding_dim=32, action_encoder_dim=1024, action_decoder_dim=1024, dynamics_dim=5120, dynamics_depth=48, dynamics_num_heads=40, encoder_depth=12, decoder_depth=20, latent_action_depth=20, use_bfloat16=False, action_pooling='mean', window_attention_heads=1, tokenizer_num_heads=16, action_num_heads=16)[source]#

Bases: ExportableAgentMixin, Module

Genie: Generative Interactive Environment.

Based on “Genie: Generative Interactive Environments” paper (arXiv:2402.15391).

Training follows two phases as per paper: 1. Train video tokenizer first (on video tokens) 2. Co-train LAM (from pixels) and dynamics model (on video tokens)

At inference, latent actions are stopgrad’d when passed to dynamics model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
tokenizer_encoder_dim (int)
tokenizer_decoder_dim (int)
action_vocab_size (int)
action_embedding_dim (int)
action_encoder_dim (int)
action_decoder_dim (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
encoder_depth (int)
decoder_depth (int)
latent_action_depth (int)
use_bfloat16 (bool)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)
tokenizer_num_heads (int)
action_num_heads (int)

classmethod from_config(config=None, **overrides)[source]#

Build Genie from a config object, dict, YAML file, or YAML string.

Parameters:

config (GenieConfig | GenieSmallConfig | dict[str, Any] | str | Path | None)
overrides (Any)

Return type:

Load Genie weights from a local path/directory or HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
config (GenieConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
map_location (str | device | None)
overrides (Any)

Return type:

CHECKPOINT_FORMAT = 2#

save_pretrained(path)[source]#

Save Genie weights and config in a from_pretrained-compatible format.

Parameters:: path (str | Path)
Return type:: None

parameter_count(trainable_only=False)[source]#

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]#

Return type:: dict[str, Any]

forward(video, mask_prob=None, training_phase='all')[source]#

Full forward pass through all components.

Parameters:

video (Tensor) – (B, C, T, H, W) input video
mask_prob (float | None) – Dynamics masking rate. None (default) lets the dynamics model draw it from U(0.5, 1) per step, per paper 2.1.
training_phase (str) – “all”, “tokenizer”, or “lam_dynamics”

Returns:

Dictionary containing losses and predictions

Return type:

Dict[str, Tensor]

training_step(video, mask_prob=None, training_phase='all')[source]#

Single training step computing all losses.

Parameters:

video (Tensor) – (B, C, T, H, W) input video
mask_prob (float | None) – Dynamics masking rate; None draws it per step.
training_phase (str) – “all”, “tokenizer”, or “lam_dynamics”

Returns:

Dictionary containing all losses for backpropagation

Return type:

Dict[str, Tensor]

encode_video(video)[source]#

Encode video to discrete tokens.

Parameters:: video (Tensor) – (B, C, T, H, W)
Returns:: (B, T, H*W)
Return type:: video_tokens

infer_actions(frames)[source]#

Infer latent actions from a sequence of frames.

Parameters:: frames (Tensor) – (B, C, T, H, W) video frames
Returns:: (B, T-1) inferred latent action indices
Return type:: latent_actions

generate(prompt_frame, num_frames=16, actions=None, use_maskgit=True, use_cache=False)[source]#

Generate video frames given a prompt frame and actions.

Parameters:

prompt_frame (Tensor) – (B, C, H, W) initial frame
num_frames (int) – Total number of frames to generate
actions (Tensor | None) – (B, num_frames-1) latent action indices, or None for random
use_maskgit (bool) – Whether to use MaskGIT sampling
use_cache (bool) – Generate each frame from a temporal KV cache instead of re-running the dynamics model over the whole prefix; O(T) rather than O(T^2) in the number of frames. Same logits up to float rounding.

Returns:

(B, C, num_frames, H, W)

Return type:

generated_video

play(current_frame, action, current_frames=None)[source]#

Play step - generate next frame given current frame and action.

Parameters:

current_frame (Tensor) – (B, C, H, W) current frame
action (Tensor) – (B,) latent action indices
current_frames (Tensor | None) – (B, C, T, H, W) history frames, or None for first frame

Returns:

(B, C, H, W)

Return type:

next_frame

get_num_parameters()[source]#

Return total number of parameters.

Return type:: int

simulusion.models.genie.genie_kwargs_from_config(config)[source]#

Map a Genie config onto Genie constructor keyword arguments.

The config and constructor name several fields differently (for example tokenizer_encoder_depth -> encoder_depth), so filtering the config by parameter name silently drops them and builds those components at the constructor defaults. Every construction path goes through this mapping.

Parameters:: config (GenieConfig | GenieSmallConfig | dict[str, Any])
Return type:: dict[str, Any]

simulusion.models.genie.create_genie(num_frames=16, image_size=64, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, action_vocab_size=8, action_embedding_dim=32, dynamics_dim=5120, dynamics_depth=48, dynamics_num_heads=40, use_bfloat16=False, action_pooling='mean', window_attention_heads=1)[source]#

Factory function to create a Genie model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
action_vocab_size (int)
action_embedding_dim (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
use_bfloat16 (bool)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

Return type:

simulusion.models.genie.create_genie_small(num_frames=16, image_size=64, use_bfloat16=False, action_pooling='mean', window_attention_heads=1)[source]#

Create a smaller Genie model for development/testing.

Parameters:

num_frames (int)
image_size (int)
use_bfloat16 (bool)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

Return type:

simulusion.models.genie.create_genie_large(num_frames=16, image_size=64, use_bfloat16=True, action_pooling='mean', window_attention_heads=1)[source]#

Create the full 11B parameter Genie model (approximate).

Parameters:

num_frames (int)
image_size (int)
use_bfloat16 (bool)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

Return type:

class simulusion.models.latent_action_model.LatentActionModel(num_frames=16, image_size=64, in_channels=3, encoder_dim=256, decoder_dim=512, encoder_depth=4, decoder_depth=4, num_heads=8, patch_size=16, vocab_size=8, embedding_dim=32, commitment_weight=1.0, action_pooling='mean', window_attention_heads=1)[source]#

Bases: ExportableAgentMixin, Module

Latent Action Model (LAM) for unsupervised action learning.

Based on Genie paper - learns actions without action labels from Internet videos.

Components: - Encoder: Takes all previous frames x1:t and next frame x_t+1 → outputs latent actions - Decoder: Takes previous frames x1:t-1 and latent actions a1:t-1 → predicts next frame x_t

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
commitment_weight (float)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

encode(x_prev, x_next)[source]#

Encode frames to latent actions.

Parameters:

x_prev (Tensor) – Previous frames (B, C, T, H, W)
x_next (Tensor) – Next frame (B, C, H, W)

Returns:

Discrete latent action indices (B, T) z_q: Quantized embeddings (B, T, embedding_dim)

Return type:

latent_actions

decode(x_prev, z_q)[source]#

Decode latent actions to predict next frame.

Parameters:

x_prev (Tensor) – Previous frames (B, C, T, H, W) - will mask all but first
z_q (Tensor) – Quantized action embeddings (B, T-1, embedding_dim)

Returns:

(B, C, H, W)

Return type:

predicted_next_frame

forward(x_prev, x_next)[source]#

Full forward pass: encode to get actions, decode to reconstruct.

Parameters:

x_prev (Tensor) – Previous frames (B, C, T, H, W)
x_next (Tensor) – Next frame (B, C, H, W)

Returns:

Dictionary with losses and outputs

Return type:

Dict[str, Tensor]

simulusion.models.latent_action_model.create_latent_action_model(num_frames=16, image_size=64, in_channels=3, encoder_dim=256, decoder_dim=512, encoder_depth=4, decoder_depth=4, num_heads=8, patch_size=16, vocab_size=8, embedding_dim=32, action_pooling='mean', window_attention_heads=1)[source]#

Factory function to create a Latent Action Model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

Return type:

LatentActionModel

class simulusion.models.dynamics_model.MaskGITSampler(num_steps=25, temperature=2.0, mask_schedule='cosine')[source]#

Bases: object

MaskGIT sampling for token-based video generation.

Uses iterative refinement with a mask schedule to progressively reveal tokens during generation.

Parameters:

num_steps (int)
temperature (float)
mask_schedule (str)

get_mask_prob(step)[source]#

Get mask probability for given step.

Parameters:: step (int)
Return type:: float

sample_frame(logits)[source]#

Sample a full frame of tokens from per-token logits in one shot.

Genie samples each frame with a temperature (2.0 in the paper) using random sampling. This helper draws one categorical sample per spatial position.

Parameters:: logits (Tensor) – (B, N, vocab_size) next-frame token logits.
Returns:: (B, N) sampled token indices.
Return type:: tokens

sample(logits, tokens, mask, step)[source]#

One MaskGIT refinement step over a single frame’s tokens.

Parameters:

logits (Tensor) – (B, N, vocab_size) per-position logits for the frame.
tokens (Tensor) – (B, N) tokens committed so far (values at masked positions are placeholders and ignored).
mask (Tensor) – (B, N) - 1 for positions still to predict, 0 for committed.
step (int) – Current refinement step in [0, num_steps).

Returns:

(B, N) with newly revealed positions filled in. new_mask: (B, N) with newly revealed positions set to 0.

Return type:

new_tokens

class simulusion.models.dynamics_model.DynamicsModel(num_frames=16, image_size=64, vocab_size=1024, embedding_dim=32, action_vocab_size=8, dim=5120, depth=48, num_heads=40, patch_size=4, mlp_ratio=4.0, qkv_bias=True, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0, gradient_checkpointing=True, mask_prob_min=0.5, mask_prob_max=1.0)[source]#

Bases: ExportableAgentMixin, Module

Dynamics Model for action-controllable video generation.

A decoder-only transformer that predicts future frame tokens given past frame tokens and latent actions. Uses MaskGIT for training and sampling.

Based on Genie paper - uses cross-entropy loss with random masking during training, and MaskGIT iterative refinement at inference.

Parameters:

num_frames (int)
image_size (int)
vocab_size (int)
embedding_dim (int)
action_vocab_size (int)
dim (int)
depth (int)
num_heads (int)
patch_size (int)
mlp_ratio (float)
qkv_bias (bool)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)
gradient_checkpointing (bool)
mask_prob_min (float)
mask_prob_max (float)

forward(video_tokens, actions, mask_prob=None)[source]#

Forward pass for training.

Parameters:

video_tokens (Tensor) – (B, T, H*W) - token indices for frames 1 to T
actions (Tensor) – (B, T) - latent action indices for frames 1 to T
mask_prob (float | None) – Masking rate. None (the default) draws it from U(mask_prob_min, mask_prob_max) per call, as the paper specifies. Pass 0.0 to disable masking, e.g. at inference.

Returns:

(B, T, H*W, vocab_size)

Return type:

logits

init_cache(batch_size, device=None, dtype=None)[source]#

Allocate a temporal KV cache for forward_cached().

Parameters:

batch_size (int)
device (device | None)
dtype (dtype | None)

Return type:

forward_cached(video_tokens, actions, cache, commit=True)[source]#

Inference forward over only the frames after those in cache.

Equivalent to the matching frames of forward(all_tokens, all_actions, mask_prob=0.0) in eval mode, without recomputing the cached prefix.

Parameters:

video_tokens (Tensor) – (B, T_new, N) tokens of the new frames. T_new may be larger than 1 only while the cache is empty (the prompt).
actions (Tensor) – (B, T_new) latent actions for those frames.
cache (STKVCache) – From init_cache(); its length is the index of the first new frame.
commit (bool) – Advance the cache past these frames. False evaluates a candidate frame without keeping it.

Returns:

(B, T_new, N, vocab_size)

Return type:

logits

sample(prompt_tokens, prompt_actions, num_frames, sampler=None)[source]#

Sample future frames using MaskGIT.

Parameters:

prompt_tokens (Tensor) – (B, T_prompt, N) - starting frame tokens
prompt_actions (Tensor) – (B, T_prompt) - actions for prompt frames
num_frames (int) – Total number of frames to generate
sampler (MaskGITSampler | None) – MaskGIT sampler instance

Returns:

(B, num_frames, N)

Return type:

generated_tokens

autoregressive_sample(prompt_tokens, actions, num_frames, temperature=1.0, use_cache=False)[source]#

Simple autoregressive sampling (frame by frame).

Parameters:

prompt_tokens (Tensor) – (B, T_prompt, N) - starting frame tokens
actions (Tensor) – (B, num_frames - 1) - latent actions for each transition. The action driving frame t -> t+1 is actions[:, t-1]. If fewer actions are supplied than transitions, the remainder are sampled at random.
num_frames (int) – Total number of frames to generate
temperature (float) – Sampling temperature
use_cache (bool) – Reuse the temporal keys/values of already generated frames (STKVCache) instead of re-running the transformer over the whole prefix for every frame. The logits match the uncached path up to float rounding. One difference: when actions must be padded at random, the uncached path redraws the padding for every earlier frame at every step, while the cached path draws each frame’s padding once.

Returns:

(B, num_frames, N)

Return type:

generated_tokens

simulusion.models.dynamics_model.create_dynamics_model(num_frames=16, image_size=64, vocab_size=1024, embedding_dim=32, action_vocab_size=8, dim=5120, depth=48, num_heads=40, patch_size=4)[source]#

Factory function to create a Dynamics Model.

Parameters:

num_frames (int)
image_size (int)
vocab_size (int)
embedding_dim (int)
action_vocab_size (int)
dim (int)
depth (int)
num_heads (int)
patch_size (int)

Return type:

DynamicsModel

Diffusion and DIAMOND components#

Key classes: DiamondAgent, DDPM, DiT, DiffusionUNet, EDMPreconditioner, EulerSampler, RewardTerminationModel, and ActorCriticNetwork.

DIAMOND exposes DiamondAgent from simulusion.training.train_diamond; there is no separate DIAMONDAgent class name in the package.

Diffusion sub-module - Diffusion model components for world models.

Exported Components:

DiT: Diffusion Transformer model
PatchEmbed: Image patch embedding
PatchUnEmbed: Patch unembedding (decode tokens to image)
DDPM: Denoising Diffusion Probabilistic Model implementation
ActorCriticNetwork: DIAMOND actor-critic network
RewardTerminationModel: Reward/termination prediction model
sinusoidal_time_embedding: Time embedding for diffusion models

class simulusion.models.diffusion.DDPM.DDPM(timesteps, beta_start, beta_end)[source]#

Bases: ExportableAgentMixin, Module

Utility module implementing forward and reverse DDPM diffusion steps.

Precomputes diffusion schedule terms and exposes helpers for noising training inputs (q_sample) and iterative denoising sampling (sample).

Parameters:

timesteps (int)
beta_start (float)
beta_end (float)

q_sample(x_start, t, noise=None)[source]#

Parameters:

x_start (Tensor)
t (Tensor)
noise (Tensor | None)

Return type:

Tensor

p_sample(model, x_t, t, y=None)[source]#

Parameters:

model (Module)
x_t (Tensor)
t (Tensor)
y (Tensor | None)

Return type:

Tensor

sample(model, n, img_size, channels, y=None)[source]#

Parameters:

model (Module)
n (int)
img_size (int)
channels (int)
y (Tensor | None)

Return type:

Tensor

simulusion.models.diffusion.DiT.sinusoidal_time_embedding(timesteps, dim, max_period=10000.0)[source]#

Create sinusoidal timestep embeddings for diffusion conditioning.

Math:: embedding[t] = [sin(t / P^(2i/d)), cos(t / P^(2i/d))] for i in [0, d/2)

Parameters:

timesteps (Tensor) – Tensor of timesteps, shape (B,) or (B, 1)
dim (int) – Embedding dimension
max_period (float) – Longest sinusoid period; 10000 follows DDPM/ADM.

Returns:

Tensor of shape (B, dim) with sinusoidal embeddings.

Return type:

Tensor

simulusion.models.diffusion.DiT.get_2d_sincos_pos_embed(embed_dim, grid_size)[source]#

Fixed 2D sine-cosine positional embeddings (ViT/DiT convention).

The paper applies “standard ViT frequency-based positional embeddings (the sine-cosine version) to all input tokens” after patchify. These are constant, not learned, which is what lets a trained model be evaluated at a different token count without the embeddings being meaningless.

Parameters:

embed_dim (int) – Token dimension; must be divisible by 4.
grid_size (int) – Tokens per side, i.e. input_size // patch_size.

Returns:

(grid_size**2, embed_dim) positional embeddings.

Return type:

Tensor

class simulusion.models.diffusion.DiT.TimestepEmbedder(hidden_size, frequency_embedding_size=256)[source]#

Bases: Module

Embed diffusion timesteps into the transformer’s conditioning space.

Paper A: “a 256-dimensional frequency embedding followed by a two-layer MLP with dimensionality equal to the transformer’s hidden size and SiLU activations”.

Parameters:

hidden_size (int)
frequency_embedding_size (int)

forward(t)[source]#

Parameters:: t (Tensor)
Return type:: Tensor

class simulusion.models.diffusion.DiT.LabelEmbedder(num_classes, hidden_size, dropout_prob=0.1)[source]#

Bases: Module

Embed class labels, with dropout to a learned null token for guidance.

Classifier-free guidance (paper 3.1) needs the model to also score the unconditional distribution. That is obtained by randomly replacing the label with a learned “null” embedding during training, so the embedding table holds num_classes + 1 entries and index num_classes is the null token.

Parameters:

num_classes (int) – Number of real classes.
hidden_size (int) – Conditioning dimension.
dropout_prob (float) – Probability of dropping the label during training. 0 disables guidance support entirely.

token_drop(labels, force_drop_ids=None)[source]#

Replace a random subset of labels with the null class.

Parameters:

labels (Tensor)
force_drop_ids (Tensor | None)

Return type:

Tensor

forward(labels, train, force_drop_ids=None)[source]#

Parameters:

labels (Tensor)
train (bool)
force_drop_ids (Tensor | None)

Return type:

Tensor

simulusion.models.diffusion.DiT.modulate(x, shift, scale)[source]#

Apply adaptive-layer-norm scale and shift to a token sequence.

Parameters:

x (Tensor)
shift (Tensor)
scale (Tensor)

Return type:

Tensor

class simulusion.models.diffusion.DiT.PatchEmbed(img_size, patch_size, in_channels, embed_dim, learnable_pos=False)[source]#

Bases: ExportableAgentMixin, Module

Patchify an image into a sequence of learnable patch tokens.

Used in Vision Transformers (ViT) and DiT to convert 2D images into sequences of token embeddings that can be processed by transformers.

Process:

Conv2d with kernel_size=stride=patch_size extracts non-overlapping patches
Each patch is projected to embed_dim via linear layer (Conv2d)
Learnable positional embeddings are added for spatial information

Input: (B, C, H, W) images Output: (B, N, embed_dim) where N = (H/patch_size) * (W/patch_size)

Parameters:

img_size (int) – Image size (assumes square), e.g., 32 for CIFAR
patch_size (int) – Size of each patch (typically 4, 8, or 16)
in_channels (int) – Number of input channels (3 for RGB)
embed_dim (int) – Output dimension for each patch token
learnable_pos (bool)

Usage with DiT:: patch_embed = PatchEmbed(img_size=32, patch_size=4, in_channels=3, embed_dim=256) tokens = patch_embed(images) # (B, 64, 256) for 32x32 image with patch_size=4

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.models.diffusion.DiT.PatchUnEmbed(img_size, patch_size, embed_dim, out_channels)[source]#

Bases: ExportableAgentMixin, Module

Reconstruct image-like tensors from patch-token sequences.

The inverse of PatchEmbed, this module reshapes token sequences into grids and uses transposed convolution to decode spatial outputs.

Parameters:

img_size (int)
patch_size (int)
embed_dim (int)
out_channels (int)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.models.diffusion.DiT.TransformerBlock(d_model, n_heads, mlp_ratio, drop, t_dim=None)[source]#

Bases: Module

DiT block with adaLN-Zero conditioning (paper 3.2).

The conditioning vector – the sum of the timestep and class embeddings – is mapped by a single SiLU + Linear to six vectors per block: shift, scale and gate for each of the attention and MLP sub-layers. That is the 6x hidden output the paper specifies for adaLN-Zero (vanilla adaLN uses 4x because it has no gates).

The “-Zero” is the important half. The modulation layer is zero-initialised, so at step 0 every gate is 0 and the whole block is the identity function. Figure 5 shows this matters a lot: adaLN-Zero reaches roughly half the FID of in-context conditioning at 400K steps, and clearly beats vanilla adaLN, which is identical apart from the gates and their initialisation.

Normalisation is LayerNorm without affine parameters – the scale and shift come from the conditioning instead, so learned per-channel affines would be redundant.

Parameters:

d_model (int)
n_heads (int)
mlp_ratio (float)
drop (float)
t_dim (int | None)

forward(x, c)[source]#

Args: x (B, T, D) tokens; c (B, cond_dim) conditioning.

Parameters:

x (Tensor)
c (Tensor)

Return type:

Tensor

class simulusion.models.diffusion.DiT.FinalLayer(hidden_size, patch_size, out_channels)[source]#

Bases: Module

Adaptive-layer-norm + linear decode to per-patch outputs (paper 3.2).

“We apply the final layer norm (adaptive if using adaLN) and linearly decode each token into a p x p x 2C tensor”. Both the modulation and the linear decode are zero-initialised, so the model starts by predicting zero noise rather than an arbitrary field.

Parameters:

hidden_size (int)
patch_size (int)
out_channels (int)

forward(x, c)[source]#

Parameters:

x (Tensor)
c (Tensor)

Return type:

Tensor

class simulusion.models.diffusion.DiT.DiT(img_size, patch_size, in_channels, d_model, depth, heads, drop=0.0, t_dim=256, num_classes=0, class_dropout_prob=0.1, learn_sigma=True, mlp_ratio=4.0)[source]#

Bases: ExportableAgentMixin, Module

Diffusion Transformer model for image denoising and generation.

The module maps noisy images and timesteps to predicted noise residuals and also provides a classmethod training entrypoint for common datasets.

Parameters:

img_size (int)
patch_size (int)
in_channels (int)
d_model (int)
depth (int)
heads (int)
drop (float)
t_dim (int)
num_classes (int)
class_dropout_prob (float)
learn_sigma (bool)
mlp_ratio (float)

unpatchify_tokens(x)[source]#

Rearrange decoded tokens (B, T, p*p*C_out) back to (B, C_out, H, W).

Parameters:: x (Tensor)
Return type:: Tensor

forward(x, t, y=None)[source]#

Predict noise (and covariance when learn_sigma).

Parameters:

x (Tensor) – Noised input (B, C, H, W).
t (Tensor) – Diffusion timesteps (B,).
y (Tensor | None) – Class labels (B,). Required when the model is class-conditional.

Returns:

(B, out_channels, H, W); when learn_sigma the first C channels are the predicted noise and the rest the covariance.

Return type:

Tensor

forward_with_cfg(x, t, y, cfg_scale, guided_channels=None)[source]#

Forward pass with classifier-free guidance (paper 3.1).

Computes eps = eps(x, null) + s * (eps(x, y) - eps(x, null)) by running the conditional and unconditional branches as one batch.

Parameters:

x (Tensor) – Noised input (B, C, H, W).
t (Tensor) – Timesteps (B,).
y (Tensor) – Class labels (B,).
cfg_scale (float) – Guidance scale s; 1.0 recovers standard sampling.
guided_channels (int | None) – Apply guidance to only the first N channels. The paper’s appendix guides 3 of the 4 latent channels; None guides all of them, for which the equivalent scale is roughly 1 + 0.75 * (s - 1).

Returns:

Guided model output, same shape as a plain forward pass.

Return type:

Tensor

classmethod from_config(config=None, **overrides)[source]#

Build DiT from a config object, dict, YAML file, or YAML string.

Parameters:

config (DiTConfig | dict[str, Any] | str | Path | None)
overrides (Any)

Return type:

Load DiT weights from a local path/directory or HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
config (DiTConfig | dict[str, Any] | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
map_location (str | device | None)
overrides (Any)

Return type:

save_pretrained(path)[source]#

Save DiT weights and config in a from_pretrained-compatible format.

Parameters:: path (str | Path)
Return type:: None

parameter_count(trainable_only=False)[source]#

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]#

Return type:: dict[str, Any]

train(mode=True)[source]#

Set training mode – the standard torch.nn.Module.train().

Parameters:: mode (bool)
Return type:: DiT

Parameters:

epochs (int)
dataset (Any)
batch_size (int)
lr (float)
img_size (int)
channels (int)
patch (int)
width (int)
depth (int)
heads (int)
drop (float)
timesteps (int)
beta_start (float)
beta_end (float)
ema (bool)
ema_decay (float)
num_classes (int)
class_dropout_prob (float)
learn_sigma (bool)
workdir (str)
root_path (str)
image_folder (str | None)
crop_size (int | None)
num_workers (int)
download (bool)
copy_data (bool)
subset_file (str | None)
val_split (float | None)
early_stopping (bool)
patience (int)
min_delta (float)
checkpoint_every (int)

Return type:

None

simulusion.models.diffusion.DiT.dit_loaders(dataset, *, batch_size, img_size, root_path='./data', image_folder=None, crop_size=None, num_workers=4, download=True, copy_data=False, subset_file=None, val_split=None, with_val=False)[source]#

DiT’s (train_loader, val_loader); val_loader is None without with_val.

dataset is "cifar10" (downloads into root_path), "imagenet" or "imagefolder". Images are scaled to [-1, 1] for CIFAR-10.

Parameters:

dataset (str)
batch_size (int)
img_size (int)
root_path (str)
image_folder (str | None)
crop_size (int | None)
num_workers (int)
download (bool)
copy_data (bool)
subset_file (str | None)
val_split (float | None)
with_val (bool)

Return type:

tuple[Any, Any]

simulusion.models.diffusion.DiT.dit_ema_update(model, ema_model, decay)[source]#

ema = decay * ema + (1 - decay) * model, parameter by parameter.

Parameters:

model (Module)
ema_model (Module)
decay (float)

Return type:

None

simulusion.models.diffusion.DiT.dit_l_simple(model, ddpm, images, labels, timesteps)[source]#

DiT’s training objective, L_simple, on one batch.

With learn_sigma the model emits 2C channels: noise then the covariance parameterisation. Paper 3.1 trains eps with L_simple and Sigma with the full variational bound; only the L_simple half is applied here, so the covariance head is carried but not learned. Set learn_sigma=False for a strictly epsilon-only model.

Parameters:

model (Module)
ddpm (Any)
images (Tensor)
labels (Tensor | None)
timesteps (int)

Return type:

Tensor

simulusion.models.diffusion.DiT.create_dit(config=None, **overrides)[source]#

Create a DiT from a DiTConfig or keyword overrides.

Parameters:

config (Any)
overrides (Any)

Return type:

class simulusion.models.diffusion.diamond_diffusion.AdaptiveGroupNorm(num_groups, num_channels, cond_dim)[source]#

Bases: Module

Adaptive Group Normalization that conditions on actions and diffusion time.

Parameters:

num_groups (int)
num_channels (int)
cond_dim (int)

forward(x, cond)[source]#

Parameters:

x (Tensor) – Input tensor [B, C, H, W]
cond (Tensor) – Conditioning tensor [B, cond_dim]

Return type:

Tensor

class simulusion.models.diffusion.diamond_diffusion.ResBlock(in_channels, out_channels, cond_dim, dropout=0.0)[source]#

Bases: Module

Residual block with adaptive group normalization.

Parameters:

in_channels (int)
out_channels (int)
cond_dim (int)
dropout (float)

forward(x, cond)[source]#

Parameters:

x (Tensor)
cond (Tensor)

Return type:

Tensor

class simulusion.models.diffusion.diamond_diffusion.AttentionBlock(channels, cond_dim)[source]#

Bases: Module

Self-attention block for U-Net.

Parameters:

channels (int)
cond_dim (int)

forward(x, cond)[source]#

Parameters:

x (Tensor)
cond (Tensor)

Return type:

Tensor

class simulusion.models.diffusion.diamond_diffusion.TimestepEmbedding(dim, freq_dim=256, max_period=10000.0)[source]#

Bases: Module

Sinusoidal noise-level embedding followed by an MLP.

The input is EDM’s c_noise = log(sigma) / 4 (Appendix C, eq. 11), which ranges over several units either side of zero as sigma sweeps the schedule. A plain Linear(1, ...) on that scalar – the previous implementation, despite the docstring – gives the network a single monotone feature to condition on, so nearby noise levels are nearly indistinguishable after the first layer. The sinusoidal expansion used by every U-Net in this lineage spreads the scalar across freq_dim frequencies first, which is what makes the noise level a usable conditioning signal.

Parameters:

dim (int)
freq_dim (int)
max_period (float)

forward(t)[source]#

Parameters:: t (Tensor)
Return type:: Tensor

class simulusion.models.diffusion.diamond_diffusion.DownBlock(in_channels, out_channels, cond_dim, num_res_blocks=2, attention=False)[source]#

Bases: Module

Downsampling block for U-Net encoder.

Parameters:

in_channels (int)
out_channels (int)
cond_dim (int)
num_res_blocks (int)
attention (bool)

forward(x, cond)[source]#

Parameters:

x (Tensor)
cond (Tensor)

Return type:

Tensor

class simulusion.models.diffusion.diamond_diffusion.UpBlock(in_channels, out_channels, cond_dim, num_res_blocks=2, attention=False, skip_channels=None)[source]#

Bases: Module

Upsampling block for U-Net decoder with skip connections.

Parameters:

in_channels (int)
out_channels (int)
cond_dim (int)
num_res_blocks (int)
attention (bool)
skip_channels (int | None)

forward(x, cond, skip=None)[source]#

Parameters:

x (Tensor)
cond (Tensor)
skip (Tensor | None)

Return type:

Tensor

class simulusion.models.diffusion.diamond_diffusion.DiffusionUNet(obs_channels=3, num_conditioning_frames=4, base_channels=64, channel_multipliers=(1, 1, 1, 1), num_res_blocks=2, cond_dim=256, action_dim=18)[source]#

Bases: Module

U-Net architecture for EDM diffusion world model. Uses frame stacking for observation conditioning and adaptive group norm for action conditioning.

Parameters:

obs_channels (int)
num_conditioning_frames (int)
base_channels (int)
channel_multipliers (Tuple[int, ...])
num_res_blocks (int)
cond_dim (int)
action_dim (int)

forward(x, t, obs_history, actions)[source]#

Forward pass of the diffusion model.

Parameters:

x (Tensor) – Noised observation at timestep t [B, C, H, W]
t (Tensor) – Diffusion timestep [B]
obs_history (Tensor) – Past observations for conditioning [B, L, C, H, W]
actions (Tensor) – Past actions [B, L]

Returns:

Predicted clean observation [B, C, H, W]

Return type:

Tensor

class simulusion.models.diffusion.diamond_diffusion.EDMPreconditioner(sigma_data=0.5, p_mean=-0.4, p_std=1.2)[source]#

Bases: object

EDM preconditioner following Karras et al. (2022).

Parameters:

sigma_data (float)
p_mean (float)
p_std (float)

get_preconditioners(sigma)[source]#

Compute EDM preconditioners for given noise levels.

Returns:: Dictionary with c_skip, c_out, c_in, c_noise
Parameters:: sigma (Tensor)
Return type:: dict

sample_noise_level(batch_size, device)[source]#

Sample noise level from log-normal distribution.

Parameters:

batch_size (int)
device (device)

Return type:

Tensor

denoise(model, x, sigma, **kwargs)[source]#

Apply EDM denoising with preconditioners.

Parameters:

model (Module) – Diffusion model
x (Tensor) – Noised input [B, C, H, W]
sigma (Tensor) – Noise level [B]
**kwargs (Any) – Additional conditioning (obs_history, actions)

Returns:

Denoised prediction [B, C, H, W]

Return type:

Tensor

class simulusion.models.diffusion.diamond_diffusion.EulerSampler(sigma_min=0.002, sigma_max=80.0, rho=7, num_steps=3, edm_precond=None)[source]#

Bases: object

Euler method sampler for reverse diffusion.

Parameters:

sigma_min (float)
sigma_max (float)
rho (int)
num_steps (int)
edm_precond (EDMPreconditioner | None)

sample(model, shape, device, obs_history=None, actions=None)[source]#

Generate samples using Euler method.

Parameters:

model (Module) – Diffusion model
shape (Tuple[int, ...]) – Output shape [B, C, H, W]
device (device) – Device to run on
obs_history (Tensor | None) – Conditioning observations [B, L, C, H, W]
actions (Tensor | None) – Conditioning actions [B, L]

Returns:

Generated samples [B, C, H, W]

Return type:

Tensor

class simulusion.models.diffusion.reward_termination.ResidualBlock(in_channels, out_channels, cond_dim=None)[source]#

Bases: Module

Residual block following DIAMOND Appendix D.

The main path is GroupNorm -> SiLU -> 3x3 convolution (stride 1, padding 1), added to a (optionally projected) skip connection. When cond_dim is provided the group normalization is made adaptive, i.e. its scale/shift are predicted from a conditioning vector (the action embedding) as used by the reward/termination model. The actor-critic omits conditioning and uses a plain group normalization.

Parameters:

in_channels (int)
out_channels (int)
cond_dim (int | None)

forward(x, cond=None)[source]#

Parameters:

x (Tensor)
cond (Tensor | None)

Return type:

Tensor

class simulusion.models.diffusion.reward_termination.RewardTerminationModel(obs_channels=3, action_dim=18, channels=(32, 32, 32, 32), lstm_dim=512, cond_dim=128, res_blocks=2, frame_size=64)[source]#

Bases: ExportableAgentMixin, Module

Reward and termination prediction model. CNN + LSTM architecture following DIAMOND paper specifications.

Parameters:

obs_channels (int) – Number of observation channels (3 for RGB)
action_dim (int) – Number of possible actions
channels (Tuple[int, ...]) – List of channel sizes for conv blocks
lstm_dim (int) – LSTM hidden dimension
cond_dim (int) – Conditioning dimension for adaptive norm
res_blocks (int)
frame_size (int)

forward(obs, actions, hidden_state=None)[source]#

Forward pass of reward/termination model.

Parameters:

obs (Tensor) – Observations [B, T, C, H, W]
actions (Tensor) – Actions [B, T]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) hidden states

Returns:

Reward predictions [B, T, 3] (for -1, 0, 1) termination_logits: Termination predictions [B, T, 2] hidden_state: Updated (h, c) hidden states

Return type:

reward_logits

predict(obs, actions, hidden_state=None)[source]#

Predict reward and termination for a single step.

Parameters:

obs (Tensor) – Single observation [B, C, H, W]
actions (Tensor) – Single action [B]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) hidden states

Returns:

Predicted reward classes as tensor (values -1,0,1) terminated: Predicted termination tensor (bool tensor) hidden_state: Updated (h, c) hidden states

Return type:

reward

init_hidden(batch_size, device)[source]#

Initialize LSTM hidden states.

Parameters:

batch_size (int)
device (device)

Return type:

Tuple[Tensor, Tensor]

class simulusion.models.diffusion.reward_termination.RewardTerminationLoss[source]#

Bases: Module

Loss function for reward and termination prediction.

forward(reward_logits, termination_logits, rewards, terminated)[source]#

Compute loss for reward and termination predictions.

Parameters:

reward_logits (Tensor) – [B, T, 3]
termination_logits (Tensor) – [B, T, 2]
rewards (Tensor) – Rewards [B, T]. Mapped to class indices via sign(r) + 1, i.e. {-1, 0, +1} reward signs -> classes {0, 1, 2}.
terminated (Tensor) – Termination flags [B, T]

Returns:

total_loss, reward_loss, termination_loss

Return type:

Tuple[Tensor, Tensor, Tensor]

class simulusion.models.diffusion.actor_critic.ActorCriticNetwork(obs_channels=3, action_dim=18, channels=(32, 32, 64, 64), lstm_dim=512, res_blocks=1, frame_size=64)[source]#

Bases: ExportableAgentMixin, Module

Actor-Critic network for DIAMOND RL training. Shared CNN-LSTM trunk with separate policy and value heads.

Parameters:

obs_channels (int)
action_dim (int)
channels (Tuple[int, ...])
lstm_dim (int)
res_blocks (int)
frame_size (int)

forward(obs, hidden_state=None)[source]#

Forward pass of actor-critic network.

Parameters:

obs (Tensor) – Observations [B, T, C, H, W]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) hidden states

Returns:

[B, T, action_dim] values: [B, T, 1] hidden_state: (h, c)

Return type:

policy_logits

get_action(obs, hidden_state=None, deterministic=False)[source]#

Get action from a single observation.

Parameters:

obs (Tensor) – Single observation [B, C, H, W]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) hidden states
deterministic (bool) – If True, take argmax; else sample

Returns:

Selected action [B] hidden_state: (h, c)

Return type:

get_actions(obs, hidden_state=None, deterministic=False)[source]#

Batched version of get_action.

Parameters:

obs (Tensor) – Tensor of shape [B, C, H, W]
hidden_state (Tuple[Tensor, Tensor] | None) – Optional LSTM hidden state tuple matching batch size
deterministic (bool) – If True, take argmax; else sample from policy

Returns:

LongTensor of shape [B] hidden_state: updated LSTM hidden state tuple

Return type:

get_value(obs, hidden_state=None)[source]#

Get value for a single observation.

Parameters:

obs (Tensor)
hidden_state (Tuple[Tensor, Tensor] | None)

Return type:

Tuple[Tensor, Tuple[Tensor, Tensor] | None]

init_hidden(batch_size, device)[source]#

Initialize LSTM hidden states.

Parameters:

batch_size (int)
device (device)

Return type:

Tuple[Tensor, Tensor]

get_hidden_size()[source]#

Get LSTM hidden size.

Return type:: int

class simulusion.models.diffusion.actor_critic.RLLoss(discount_factor=0.985, lambda_returns=0.95, entropy_weight=0.001)[source]#

Bases: Module

RL loss functions for DIAMOND. Implements REINFORCE with value baseline and λ-returns.

Parameters:

discount_factor (float)
lambda_returns (float)
entropy_weight (float)

compute_lambda_returns(rewards, values, dones)[source]#

Compute λ-returns.

Parameters:

rewards (Tensor) – [B, T]
values (Tensor) – [B, T+1]
dones (Tensor) – [B, T]

Returns:

[B, T]

Return type:

lambda_returns

policy_loss(policy_logits, actions, lambda_returns, values)[source]#

Compute policy loss with REINFORCE and entropy regularization.

Parameters:

policy_logits (Tensor) – [B, T, A]
actions (Tensor) – [B, T]
lambda_returns (Tensor) – [B, T]
values (Tensor) – [B, T+1]

Returns:

scalar

Return type:

policy_loss

value_loss(values, lambda_returns)[source]#

Compute value loss (MSE between value and lambda returns).

Parameters:

values (Tensor)
lambda_returns (Tensor)

Return type:

Tensor

Vision, tokenization, and layers#

Key classes: ConvEncoder, ConvDecoder, DenseDecoder, ActionDecoder, CNNEncoder, CNNDecoder, IRISEncoder, IRISDecoder, DiscreteAutoencoder, VectorQuantizer, VectorQuantizerEMA, VideoTokenizer, MultiHeadSelfAttention, and STTransformer.

Convolutional Variational Autoencoder (ConvVAE) implementation.

This module provides the ConvVAE model architecture for encoding and decoding images in the World Models framework. The VAE uses a convolutional encoder and decoder with a variational latent space.

class simulusion.vision.VAE.ConvVAE.ConvVAEEncoder(img_channels, latent_size)[source]#

Bases: Module

Convolutional encoder for VAE.

This encoder takes images and produces the parameters (mean and log variance) of a Gaussian distribution in the latent space.

Variables:

latent_size – Dimensionality of the latent space.
img_channels – Number of input image channels.

Parameters:

img_channels (int)
latent_size (int)

Example

>>> encoder = ConvVAEEncoder(img_channels=3, latent_size=32)
>>> mu, logsigma = encoder(images)

forward(x)[source]#

Encode images to latent distribution parameters.

Parameters:

x (Tensor) – Input tensor of shape (batch, channels, height, width).

Returns:

mu: Mean of the latent distribution
logsigma: Log variance of the latent distribution

Return type:

Tuple of (mu, logsigma) where

class simulusion.vision.VAE.ConvVAE.ConvVAEDecoder(latent_size, img_channels)[source]#

Bases: Module

Convolutional decoder for VAE.

This decoder takes latent vectors and reconstructs images.

Variables:

latent_size – Dimensionality of the input latent space.
img_channels – Number of output image channels.

Parameters:

latent_size (int)
img_channels (int)

forward(z)[source]#

Decode latent vectors to images.

Parameters:: z (Tensor) – Latent vector of shape (batch, latent_size).
Returns:: Reconstructed image tensor of shape (batch, channels, height, width).
Return type:: Tensor

class simulusion.vision.VAE.ConvVAE.ConvVAE(img_channels, latent_size)[source]#

Bases: Module

Convolutional Variational Autoencoder.

The ConvVAE is a generative model that encodes images into a latent distribution and reconstructs them. It uses the reparameterization trick to enable backpropagation through the sampling process.

Variables:

encoder – ConvVAEEncoder that encodes images to latent parameters.
decoder – ConvVAEDecoder that decodes latent vectors to images.

Parameters:

img_channels (int)
latent_size (int)

Example

>>> vae = ConvVAE(img_channels=3, latent_size=32)
>>> recon_x, mu, logsigma = vae(images)
>>> # Training loss combines reconstruction and KL divergence

forward(x)[source]#

Encode and decode an image.

Parameters:

x (Tensor) – Input image tensor of shape (batch, channels, height, width).

Returns:

recon_x: Reconstructed image
mu: Mean of latent distribution
logsigma: Log variance of latent distribution

Return type:

Tuple of (recon_x, mu, logsigma)

class simulusion.vision.dreamer_encoder.ConvEncoder(input_shape, embed_size, activation, depth=32)[source]#

Bases: ExportableAgentMixin, Module

Convolutional observation encoder used by Dreamer world models.

Input: (B, C, H, W) raw images, values in [-0.5, 0.5]
Process: 4 convolutional layers with stride 2, halving spatial dimensions
Output: (B, embed_size) compact representation

The encoder uses a depth doubling pattern: 32 -> 64 -> 128 -> 256 channels. After convolutions, a fully connected layer projects from 1024 features to the desired embedding size.

Usage with Dreamer:

encoder = ConvEncoder(
    input_shape=(3, 64, 64),  # RGB 64x64 images
    embed_size=256,           # RSSM observation embedding size
    activation='relu'         # Activation function
)
obs_embedding = encoder(observation)  # (B, 256)

Parameters:

input_shape (tuple) – Tuple (C, H, W) for input images, typically (3, 64, 64)
embed_size (int) – Output embedding dimension, typically 256 or 1024
activation (str) – Activation function name (‘relu’, ‘elu’, ‘tanh’, etc.)
depth (int) – Base channel depth for first layer (default 32)

forward(inputs)[source]#

Parameters:: inputs (Tensor)
Return type:: Tensor

class simulusion.vision.dreamer_decoder.TanhBijector[source]#

Bases: Transform

Bijective tanh transform for squashing Gaussian distributions to [-1, 1].

Bijective mapping: tanh is invertible (with atanh as inverse)
Stable log-det Jacobian: Computable for gradient-based training
Clipped actions: During inference, actions are naturally bounded

Forward: y = tanh(x)
Inverse: x = atanh(y) = 0.5 * log((1+y)/(1-y))
Log-det: log|dy/dx| = 2*(log(2) - x - softplus(-2x))

Usage with Dreamer ActionDecoder:

dist = TransformedDistribution(
    Normal(mean, std),
    TanhBijector()
)
action = dist.sample()  # Bounded to [-1, 1]

Reference:: Building a Scalable Deep RL Library by Learning from Mistakes, Haarnoja et al.

property sign: int#

atanh(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

log_abs_det_jacobian(x, y)[source]#

Parameters:

x (Tensor)
y (Tensor)

Return type:

Tensor

class simulusion.vision.dreamer_decoder.ConvDecoder(stoch_size, deter_size, output_shape, activation, depth=32)[source]#

Bases: ExportableAgentMixin, Module

Convolutional decoder for reconstructing observations from latent states.

Part of Dreamer’s world model, this decoder reconstructs image observations from the combined stochastic (s) and deterministic (h) RSSM states.

Input: Concatenated [stoch_state, deter_state], shape (B, stoch+deter)
Process: Dense projection + 4 transposed convolutions (upsampling 2x each)
Output: Independent Normal distribution over observation pixels

Returns torch.distributions.Independent(Normal(mean, std), len(shape)) allowing log_prob(observation) computation for reconstruction loss.

Usage in Dreamer world model:

decoder = ConvDecoder(
    stoch_size=30,
    deter_size=200,
    output_shape=(3, 64, 64),  # RGB images
    activation='relu'
)
obs_dist = decoder(latent_features)  # Returns distribution
log_prob = obs_dist.log_prob(target_observation)

The reconstruction loss is -log_prob(observation), which encourages the RSSM to learn states that capture observation information.

Parameters:

stoch_size (int)
deter_size (int)
output_shape (tuple[int, ...])
activation (str)
depth (int)

forward(features)[source]#

Parameters:: features (Tensor)
Return type:: Independent

class simulusion.vision.dreamer_decoder.DenseDecoder(stoch_size, deter_size, output_shape, n_layers, units, activation, dist, num_buckets=255, symlog_range=10.0)[source]#

Bases: ExportableAgentMixin, Module

MLP decoder for reward/value/discount prediction from latent features.

Part of Dreamer’s world model, this decoder predicts scalar quantities (rewards, values, discount factors) from RSSM latent states.

Input: [stoch_state, deter_state] concatenated, shape (B, stoch+deter)
Process: MLP with configurable layers and hidden units
Output: Predicted quantity with distribution (normal, binary, or raw)

Usage:

reward_decoder = DenseDecoder(
    stoch_size=30,
    deter_size=200,
    output_shape=(1,),
    n_layers=2,
    units=400,
    activation='elu',
    dist='normal'
)
reward_dist = reward_decoder(latent_features)
reward_loss = -reward_dist.log_prob(target_reward)

For discount prediction (binary):

discount_decoder = DenseDecoder(
    stoch_size=30,
    deter_size=200,
    output_shape=(1,),
    n_layers=2,
    units=400,
    activation='elu',
    dist='binary'
)

Parameters:

stoch_size (int)
deter_size (int)
output_shape (tuple[int, ...])
n_layers (int)
units (int)
activation (str)
dist (str)
num_buckets (int)
symlog_range (float)

forward(features)[source]#

Parameters:: features (Tensor)
Return type:: Any

class simulusion.vision.dreamer_decoder.SampleDist(dist, samples=100)[source]#

Bases: object

Distribution wrapper that estimates statistics via Monte Carlo sampling.

Provides approximated mean, mode, and entropy helpers for transformed distributions where analytic forms may be inconvenient.

Parameters:

dist (Any)
samples (int)

property name: str#

mean()[source]#

Return type:: Tensor

mode()[source]#

Return type:: Tensor

entropy()[source]#

Return type:: Tensor

sample()[source]#

Return type:: Tensor

class simulusion.vision.dreamer_decoder.ActionDecoder(action_size, stoch_size, deter_size, n_layers, units, activation, min_std=0.0001, init_std=5, mean_scale=5)[source]#

Bases: ExportableAgentMixin, Module

Dreamer actor head producing squashed continuous actions from latent features.

Outputs a transformed Gaussian policy with optional deterministic mode and utility for additive exploration noise.

Parameters:

action_size (int)
stoch_size (int)
deter_size (int)
n_layers (int)
units (int)
activation (str)
min_std (float)
init_std (float)
mean_scale (float)

forward(features, deter=False)[source]#

Parameters:

features (Tensor)
deter (bool)

Return type:

Tensor

mean_action(features)[source]#

Deterministic action tanh(mean) of the squashed Gaussian policy.

Parameters:: features (Tensor)
Return type:: Tensor

add_exploration(action, action_noise=0.3)[source]#

Parameters:

action (Tensor)
action_noise (float)

Return type:

Tensor

class simulusion.vision.planet_encoder.CNNEncoder(embedding_size, activation_function='relu')[source]#

Bases: ExportableAgentMixin, Module

A Convolutional Neural Network (CNN) encoder for processing image inputs.

Parameters:

embedding_size (int)
activation_function (str)

forward(observation)[source]#

Parameters:: observation (Tensor)
Return type:: Tensor

class simulusion.vision.planet_decoder.CNNDecoder(state_size, latent_size, embedding_size, activation_function='relu')[source]#

Bases: ExportableAgentMixin, Module

A Convolutional Neural Network (CNN) decoder for reconstructing image outputs.

Parameters:

state_size (int)
latent_size (int)
embedding_size (int)
activation_function (str)

forward(latent, state)[source]#

Parameters:

latent (Tensor)
state (Tensor)

Return type:

Tensor

class simulusion.vision.iris_encoder.IRISEncoder(vocab_size=512, tokens_per_frame=16, embedding_dim=512, in_channels=3, base_channels=64, num_layers=4, num_residual_blocks=2, frame_shape=(3, 64, 64), commitment_weight=1.0, quantizer='gradient')[source]#

Bases: ExportableAgentMixin, Module

CNN Encoder for IRIS discrete autoencoder.

Encodes image observations into latent features, which are then quantized into discrete tokens using the VectorQuantizer.

Architecture:

4 convolutional layers with residual blocks
Self-attention at 8x8 and 16x16 resolutions
Vector quantization to produce discrete tokens

Parameters:

vocab_size (int)
tokens_per_frame (int)
embedding_dim (int)
in_channels (int)
base_channels (int)
num_layers (int)
num_residual_blocks (int)
frame_shape (Tuple[int, int, int])
commitment_weight (float)
quantizer (str)

forward(x)[source]#

Encode images to discrete tokens.

Parameters:: x (Tensor) – Input images (B, C, H, W) - should be 64x64
Returns:: Quantized tokens (B, C, H’, W’) indices: Token indices (B, H’, W’) vq_loss: Dictionary with VQ loss components
Return type:: z_q

encode_to_indices(x)[source]#

Encode directly to token indices (for world model).

Parameters:: x (Tensor)
Return type:: Tensor

decode_from_indices(indices)[source]#

Decode token indices to embeddings (for decoder).

Parameters:: indices (Tensor)
Return type:: Tensor

class simulusion.vision.iris_encoder.ResidualBlock(channels)[source]#

Bases: Module

Residual block for encoder.

Parameters:: channels (int)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.vision.iris_encoder.SelfAttentionBlock(channels)[source]#

Bases: Module

Self-attention block for encoder.

Applies spatial self-attention to capture long-range dependencies.

Parameters:: channels (int)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.vision.iris_decoder.IRISDecoder(vocab_size=512, embedding_dim=512, base_channels=64, out_channels=3, frame_shape=(3, 64, 64), num_residual_blocks=2)[source]#

Bases: ExportableAgentMixin, Module

CNN Decoder for IRIS discrete autoencoder.

Decodes discrete tokens back into image observations. Uses transposed convolutions to upsample from 4x4 to 64x64.

Parameters:

vocab_size (int)
embedding_dim (int)
base_channels (int)
out_channels (int)
frame_shape (Tuple[int, int, int])
num_residual_blocks (int)

forward(z)[source]#

Decode tokens to images.

Parameters:: z (Tensor) – Token embeddings (B, C, H, W) - e.g., (B, 512, 4, 4)
Returns:: Reconstructed images (B, C, H, W) - e.g., (B, 3, 64, 64)
Return type:: reconstructed

decode_from_embeddings(z_flat)[source]#

Decode flattened token embeddings to images.

Parameters:: z_flat (Tensor) – Flattened tokens (B, H*W, C) or (B, C, H, W)
Returns:: Reconstructed images
Return type:: Tensor

decode_from_indices(indices, codebook)[source]#

Decode discrete token indices into images.

Parameters:

indices (Tensor) – Tensor of shape (B, H, W) or (B, H*W) containing integer token indices in the range [0, vocab_size).
codebook (Embedding) – The encoder’s quantizer codebook.

Returns:

Reconstructed images (B, C, H, W)

Return type:

Tensor

class simulusion.vision.iris_decoder.UpsampleBlock(in_channels, mid_channels, out_channels)[source]#

Bases: Module

Upsampling block with optional residual connection.

Parameters:

in_channels (int)
mid_channels (int)
out_channels (int)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.vision.iris_decoder.ResidualBlock(channels)[source]#

Bases: Module

Residual block for decoder.

Parameters:: channels (int)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.vision.iris_decoder.DiscreteAutoencoder(vocab_size=512, tokens_per_frame=16, embedding_dim=512, base_channels=64, frame_shape=(3, 64, 64))[source]#

Bases: Module

Complete Discrete Autoencoder combining encoder and decoder.

Used for training the VQVAE component of IRIS.

Parameters:

vocab_size (int)
tokens_per_frame (int)
embedding_dim (int)
base_channels (int)
frame_shape (Tuple[int, int, int])

forward(x)[source]#

Full encode-decode forward pass.

Parameters:: x (Tensor) – Input images (B, C, H, W)
Returns:: Reconstructed images indices: Token indices (B, H’, W’) loss_dict: Dictionary with loss components
Return type:: reconstruction

encode(x)[source]#

Encode to token indices.

Parameters:: x (Tensor)
Return type:: Tensor

decode(indices)[source]#

Decode token indices to images.

Parameters:: indices (Tensor)
Return type:: Tensor

simulusion.vision.vq_layer.restart_dead_codes(codebook, usage, z_flat, threshold)[source]#

Re-seed codebook entries that have fallen out of use.

Nearest-neighbour quantizers are prone to codebook collapse: a code that stops winning any assignment receives no further update and can never come back, so the effective vocabulary shrinks (visible as the perplexity metric dropping toward 1). The standard remedy is to periodically reset unused entries onto randomly drawn encoder outputs, which puts them back in a region of space where they can win assignments again.

Parameters:

codebook (Embedding) – The embedding table to modify in place.
usage (Tensor) – (vocab_size,) EMA of how often each code was selected. Under an EMA with decay d, this converges to the code’s mean assignments per step, so the threshold is interpretable in those units.
z_flat (Tensor) – (N, C) encoder outputs from the current batch, used as the pool of candidate re-seed locations.
threshold (float) – Codes with usage below this are considered dead. Values near 1.0 are far too aggressive for a full codebook – with V codes and roughly V assignments per step the mean usage is ~1, so half the book would be restarted every step. See DEFAULT_DEAD_CODE_THRESHOLD.

Returns:

(num_restarted, restarted_mask) – a scalar tensor counting the codes that were re-seeded, and the boolean mask identifying them. Callers that keep their own accumulators for the codebook (the EMA quantizer’s ema_embed_avg) need the mask: they cannot recover it afterwards by comparing usage against the threshold, because the restart deliberately lifts usage above it.

Return type:

tuple[Tensor, Tensor]

class simulusion.vision.vq_layer.VectorQuantizer(vocab_size=512, embedding_dim=512, commitment_weight=1.0, restart_dead_codes_after=0.01, usage_decay=0.99)[source]#

Bases: ExportableAgentMixin, Module

Vector Quantizer for discrete autoencoder.

Implements the VQ-VAE quantization from: “Neural Discrete Representation Learning” (Van Den Oord et al., 2017)

Uses exponential moving averages for codebook updates and straight-through estimator for gradient flow.

Parameters:

vocab_size (int)
embedding_dim (int)
commitment_weight (float)
restart_dead_codes_after (float)
usage_decay (float)

forward(z)[source]#

Quantize the input latents.

Parameters:: z (Tensor) – Input tensor of shape (B, C, H, W) or (B, C)
Returns:: Quantized tensor (same shape as input) indices: Token indices for each position (B, H, W) or (B,) loss: Dictionary containing VQ loss components
Return type:: z_q

decode_indices(indices)[source]#

Decode token indices back to embeddings.

Parameters:: indices (Tensor) – Token indices (B, H, W) or (B,)
Returns:: Embeddings (B, C, H, W) or (B, C)
Return type:: Tensor

class simulusion.vision.vq_layer.VectorQuantizerEMA(vocab_size=512, embedding_dim=512, commitment_weight=1.0, ema_decay=0.99, epsilon=1e-05, restart_dead_codes_after=0.01)[source]#

Bases: ExportableAgentMixin, Module

Vector Quantizer with Exponential Moving Average updates.

Uses EMA updates for the codebook instead of gradient-based updates, which leads to more stable training.

Parameters:

vocab_size (int)
embedding_dim (int)
commitment_weight (float)
ema_decay (float)
epsilon (float)
restart_dead_codes_after (float)

forward(z)[source]#

Quantize with EMA updates.

Parameters:: z (Tensor)
Return type:: tuple[Tensor, Tensor, dict[str, Tensor]]

decode_indices(indices)[source]#

Decode token indices to embeddings.

Parameters:: indices (Tensor) – Token indices (B, H, W) or (B,)
Returns:: Embeddings (B, C, H, W) or (B, C)
Return type:: Tensor

class simulusion.vision.video_tokenizer.VideoTokenizer(num_frames=16, image_size=64, in_channels=3, encoder_dim=512, decoder_dim=1024, encoder_depth=12, decoder_depth=20, num_heads=16, patch_size=4, vocab_size=1024, embedding_dim=32, commitment_weight=0.25, use_ema=False, ema_decay=0.99)[source]#

Bases: ExportableAgentMixin, Module

Video Tokenizer using VQ-VAE with Spatiotemporal Transformer.

Architecture

Patch Embedding: Convert (B, C, T, H, W) video to patch tokens
Encoder ST-Transformer: Process spatial-temporal patches
Vector Quantization: Discretize continuous embeddings to codebook entries
Decoder ST-Transformer: Reconstruct video from quantized tokens
Patch Unembedding: Convert tokens back to video frames

Key Features

Causal processing: Each frame’s encoding only uses previous frames
Discrete tokens: Enables autoregressive prediction with latent actions
Memory efficient: Uses ST-Transformer instead of full ViT to reduce complexity

Usage with Genie:

tokenizer = VideoTokenizer(
    num_frames=16,
    image_size=64,
    patch_size=4,
    vocab_size=1024,
    embedding_dim=32
)
reconstructed, indices, loss_dict = tokenizer(video_frames)

# For discrete token input to dynamics model:
token_embeddings = tokenizer.decode_indices(indices)

Reference:: Genie: Generative Interactive Environments Bruce et al., Google DeepMind, 2024 - https://arxiv.org/abs/2402.15391

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
commitment_weight (float)
use_ema (bool)
ema_decay (float)

encode(x)[source]#

Encode video to discrete tokens.

Parameters:: x (Tensor) – Video tensor (B, C, T, H, W)
Returns:: Quantized embeddings (B, T, H’, W’, embedding_dim) indices: Token indices (B, T, H’, W’) vq_loss: Dictionary with VQ loss components
Return type:: z_q

decode_indices(indices)[source]#

Decode token indices to embeddings for video frames.

Parameters:: indices (Tensor) – Token indices (B, T, H’, W’) or (B, T, N) where N = H’ x W’
Returns:: Quantized embeddings (B, T, H’, W’, embedding_dim)
Return type:: z_q

decode(z_q)[source]#

Decode discrete tokens to video frames.

Parameters:: z_q (Tensor) – Quantized embeddings (B, T, H’, W’, embedding_dim)
Returns:: Reconstructed video (B, C, T, H, W)
Return type:: Tensor

forward(x)[source]#

Full forward pass with VQ-VAE objective.

Parameters:: x (Tensor) – Video tensor (B, C, T, H, W)
Returns:: Reconstructed video (B, C, T, H, W) indices: Token indices (B, T, H’, W’) loss_dict: Dictionary containing loss components
Return type:: reconstructed

simulusion.vision.video_tokenizer.create_video_tokenizer(num_frames=16, image_size=64, in_channels=3, encoder_dim=512, decoder_dim=1024, encoder_depth=12, decoder_depth=20, num_heads=16, patch_size=4, vocab_size=1024, embedding_dim=32, use_ema=False)[source]#

Factory function to create a Video Tokenizer.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
use_ema (bool)

Return type:

VideoTokenizer

Blocks sub-module - Transformer blocks and attention mechanisms.

Exported Components:

Transformers:

STTransformer: Spatiotemporal Transformer for video processing
STSpatialAttention: Spatial attention layer
STTemporalAttention: Temporal attention layer
STTransformerBlock: Combined spatiotemporal transformer block
STBlock: Backwards-compatible alias for STTransformerBlock

Attention:

MultiHeadSelfAttention: Multi-head self-attention
MultiHeadAttention: Backwards-compatible alias for MultiHeadSelfAttention
Attention: Attention mechanism

Normalization:

RMSNorm: Root Mean Square Layer Normalization
AdaLNNormalization: Adaptive Layer Normalization

class simulusion.blocks.mhsa.MultiHeadSelfAttention(d, n_heads=2)[source]#

Bases: ExportableAgentMixin, Module

Multi-head scaled dot-product self-attention over sequence tokens.

Parameters:

d (int)
n_heads (int)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.blocks.st_transformer.STKVCache(num_layers, batch_size, num_heads, num_patches, head_dim, max_frames, device, dtype)[source]#

Bases: object

Per-layer temporal key/value cache for frame-by-frame generation.

Parameters:

num_layers (int)
batch_size (int)
num_heads (int)
num_patches (int)
head_dim (int)
max_frames (int)
device (device)
dtype (dtype)

append(layer, k, v)[source]#

Write (B, heads, N, T_new, hd) keys/values; return the full prefix.

Parameters:

layer (int)
k (Tensor)
v (Tensor)

Return type:

Tuple[Tensor, Tensor]

advance(frames)[source]#

Commit frames newly written frames.

Parameters:: frames (int)
Return type:: None

class simulusion.blocks.st_transformer.STSpatialAttention(dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0)[source]#

Bases: Module

Spatial attention layer for spatiotemporal transformer.

Processes video tokens by attending over spatial positions (H*W) within each time step independently. Captures within-frame spatial relationships.

Input: (B, T, N, C) – B batches, T time steps, N spatial positions (H*W), C channels
Output: (B, T, N, C) – Same shape, spatially attended features

Architecture

QKV projection: Linear(dim, dim*3)
Reshape to multi-head attention format
Fused scaled dot-product attention (FlashAttention on supported GPUs)
Output projection

Applied to video tokens of shape (B, T, N, C) to capture within-frame spatial structure (e.g., object positions).

Parameters:

dim (int)
num_heads (int)
qkv_bias (bool)
qk_scale (float | None)
attn_drop (float)
proj_drop (float)

forward(x)[source]#

Parameters:: x (Tensor) – (B, T, N, C) where T is temporal dim, N is spatial dim (H*W)
Returns:: (B, T, N, C)
Return type:: Tensor

class simulusion.blocks.st_transformer.STTemporalAttention(dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0)[source]#

Bases: Module

Temporal attention layer with causal masking for spatiotemporal transformer.

Processes video tokens by attending over time steps (T) across all spatial positions. Uses causal masking to ensure each frame only attends to previous frames (important for autoregressive video generation).

Input: (B, T, N, C) – B batches, T time steps, N spatial positions, C channels
Output: (B, T, N, C) – Same shape, temporally attended features

Causal masking

Frame t can only attend to frames 0…t-1
Prevents information leakage from future frames
Essential for autoregressive video generation models

Applied after STSpatialAttention to model temporal dynamics in the Genie VideoTokenizer.

Parameters:

dim (int)
num_heads (int)
qkv_bias (bool)
qk_scale (float | None)
attn_drop (float)
proj_drop (float)

forward(x, causal=True, cache=None, layer_idx=0)[source]#

Parameters:

x (Tensor) – (B, T, N, C) where T is temporal dim, N is spatial dim (H*W)
causal (bool) – whether to apply causal masking
cache (STKVCache | None) – Optional temporal KV cache. x then holds only the new frames: either the whole prompt (cache empty) or one frame attending to every cached frame.
layer_idx (int) – This layer’s slot in cache.

Returns:

(B, T, N, C)

Return type:

Tensor

class simulusion.blocks.st_transformer.STMLP(in_features, hidden_features=None, out_features=None, act_layer=<class 'torch.nn.modules.activation.GELU'>, drop=0.0)[source]#

Bases: Module

MLP for ST-Transformer block.

Parameters:

in_features (int)
hidden_features (int | None)
out_features (int | None)
act_layer (type[Module])
drop (float)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.blocks.st_transformer.STTransformerBlock(dim, num_heads=8, mlp_ratio=4.0, qkv_bias=False, qk_scale=None, drop=0.0, attn_drop=0.0, drop_path=0.0, act_layer=<class 'torch.nn.modules.activation.GELU'>, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>)[source]#

Bases: Module

Combined spatiotemporal transformer block with interleaved attention.

A single block applies:

Spatial attention (within each time frame)
Temporal attention (across frames with causal mask)
MLP projection

The order is: x -> + SpatialAttn -> + TemporalAttn -> + MLP -> x

This interleaved design captures both spatial structure and temporal dynamics efficiently, used in Genie’s VideoTokenizer and DynamicsModel.

Parameters:

dim (int) – Feature dimension (must match patch embedding dimension)
num_heads (int) – Number of attention heads
mlp_ratio (float) – MLP hidden dim = dim * mlp_ratio
drop (float) – Dropout rates
attn_drop (float) – Dropout rates
drop_path (float) – Stochastic depth rate for drop path regularization
norm_layer (type[Module]) – Normalization layer class (default: nn.LayerNorm)
qkv_bias (bool)
qk_scale (float | None)
act_layer (type[Module])

Usage in Genie:

# VideoTokenizer encoder (12 layers)
encoder = STTransformer(
    num_frames=16,
    num_patches_per_frame=256,
    dim=512,
    depth=12,
    num_heads=16
)
encoded = encoder(tokens)  # (B, T*N, C)

# Dynamics model decoder (24 layers)
decoder = STTransformer(
    num_frames=16,
    num_patches_per_frame=256,
    dim=1024,
    depth=24,
    num_heads=16
)
decoded = decoder(tokens)

forward(x, cache=None, layer_idx=0)[source]#

Parameters:

x (Tensor) – (B, T, N, C) or (B, T*H*W, C)
cache (STKVCache | None) – Optional temporal KV cache (see STKVCache).
layer_idx (int) – This block’s slot in cache.

Returns:

Same shape as input

Return type:

Tensor

class simulusion.blocks.st_transformer.DropPath(drop_prob=0.0)[source]#

Bases: Module

Drop paths (Stochastic Depth) per sample.

Parameters:: drop_prob (float)

forward(x)[source]#

Parameters:: x (Tensor)
Return type:: Tensor

class simulusion.blocks.st_transformer.STTransformer(num_frames=16, num_patches_per_frame=256, dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>, gradient_checkpointing=False)[source]#

Bases: ExportableAgentMixin, Module

Spatiotemporal Transformer for video modeling.

Contains L spatiotemporal blocks with interleaved spatial and temporal attention.

Parameters:

num_frames (int)
num_patches_per_frame (int)
dim (int)
depth (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
qk_scale (float | None)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)
norm_layer (type[Module])
gradient_checkpointing (bool)

init_cache(batch_size, max_frames=None, device=None, dtype=None)[source]#

Allocate an empty temporal KV cache for incremental generation.

Parameters:

batch_size (int)
max_frames (int | None)
device (device | None)
dtype (dtype | None)

Return type:

forward(x, cache=None, commit=True)[source]#

Parameters:

x (Tensor) – (B, T*N, C) where T is num_frames, N is num_patches_per_frame
cache (STKVCache | None) – Optional temporal KV cache. When given, x holds only the frames after the cached ones, and the result equals the corresponding frames of a full-prefix forward (eval mode).
commit (bool) – Advance cache past these frames. Pass False to evaluate a candidate frame that will be recomputed (e.g. MaskGIT steps).

Returns:

(B, T*N, C)

Return type:

Tensor

simulusion.blocks.st_transformer.create_st_transformer(num_frames=16, patch_size=4, img_size=64, dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0, gradient_checkpointing=False)[source]#

Factory function to create an ST-Transformer.

Parameters:

num_frames (int)
patch_size (int)
img_size (int)
dim (int)
depth (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)
gradient_checkpointing (bool)

Return type:

STTransformer

Configuration objects#

Lazy config exports.

Configuration modules can have optional training dependencies, so the package initializer avoids importing every config eagerly.

class simulusion.configs.DreamerConfig(env_backend='dmc', env='walker-walk', env_instance=None, image_size=(64, 64), gym_render_mode='rgb_array', frame_stack=1, dmlab_action_repeat=4, dmlab_action_set=None, dmlab_observations=None, dmlab_config=None, dmlab_renderer='hardware', procgen_distribution_mode='easy', procgen_num_levels=0, procgen_start_level=None, mujoco_xml_path=None, mujoco_xml_string=None, mujoco_binary_path=None, mujoco_camera=None, mujoco_frame_skip=1, mujoco_reset_noise_scale=0.0, brax_backend='generalized', brax_jit=True, brax_auto_reset=False, brax_suppress_warp_warnings=True, unity_file_name=None, unity_behavior_name=None, unity_worker_id=0, unity_base_port=5005, unity_no_graphics=True, unity_time_scale=20.0, unity_quality_level=1, algo='Dreamerv1', exp_name='lr1e-3', train=True, evaluate=False, seed=1, no_gpu=False, max_episode_length=1000, buffer_size=800000, time_limit=1000, cnn_activation_function='relu', dense_activation_function='elu', obs_embed_size=1024, num_units=400, deter_size=200, stoch_size=30, action_repeat=2, action_noise=0.3, total_steps=5000000, seed_steps=5000, update_steps=100, collect_steps=1000, batch_size=50, train_seq_len=50, imagine_horizon=15, use_disc_model=False, free_nats=3.0, discount=0.99, reward_scale=1.0, td_lambda=0.95, kl_loss_coeff=1.0, kl_alpha=0.8, disc_loss_coeff=10.0, num_buckets=255, symlog_range=10.0, model_learning_rate=0.0006, actor_learning_rate=8e-05, value_learning_rate=8e-05, adam_epsilon=1e-07, grad_clip_norm=100.0, use_amp=False, perf_defaults=True, tf32=True, test=False, test_interval=10000, test_episodes=10, scalar_freq=1000, log_video_freq=-1, max_videos_to_save=2, video_format='gif', video_fps=20, checkpoint_interval=10000, early_stopping=False, patience=10, min_delta=0.0001, checkpoint_path='', restore=False, experience_replay='', render=False, enable_wandb=False, wandb_project='simulusion', wandb_entity='', log_dir='runs', logdir=None, data_dir=None, log_level='INFO', log_file=None, enable_tensorboard=False, enable_console_metrics=True, enable_jsonl=True, jsonl_filename='metrics.jsonl', log_system_stats_freq=1000, detect_anomaly=False)[source]

Bases: SerializableConfigMixin

Configuration container for Dreamer training, evaluation, and environment setup.

Parameters:

env_backend (str)
env (str)
env_instance (Any)
image_size (tuple[int, int])
gym_render_mode (str)
frame_stack (int)
dmlab_action_repeat (int)
dmlab_action_set (Any)
dmlab_observations (Any)
dmlab_config (Any)
dmlab_renderer (str)
procgen_distribution_mode (str)
procgen_num_levels (int)
procgen_start_level (Any)
mujoco_xml_path (Any)
mujoco_xml_string (Any)
mujoco_binary_path (Any)
mujoco_camera (Any)
mujoco_frame_skip (int)
mujoco_reset_noise_scale (float)
brax_backend (str)
brax_jit (bool)
brax_auto_reset (bool)
brax_suppress_warp_warnings (bool)
unity_file_name (Any)
unity_behavior_name (Any)
unity_worker_id (int)
unity_base_port (int)
unity_no_graphics (bool)
unity_time_scale (float)
unity_quality_level (int)
algo (str)
exp_name (str)
train (bool)
evaluate (bool)
seed (int)
no_gpu (bool)
max_episode_length (int)
buffer_size (int)
time_limit (int)
cnn_activation_function (str)
dense_activation_function (str)
obs_embed_size (int)
num_units (int)
deter_size (int)
stoch_size (int)
action_repeat (int)
action_noise (float)
total_steps (int)
seed_steps (int)
update_steps (int)
collect_steps (int)
batch_size (int)
train_seq_len (int)
imagine_horizon (int)
use_disc_model (bool)
free_nats (float)
discount (float)
reward_scale (float)
td_lambda (float)
kl_loss_coeff (float)
kl_alpha (float)
disc_loss_coeff (float)
num_buckets (int)
symlog_range (float)
model_learning_rate (float)
actor_learning_rate (float)
value_learning_rate (float)
adam_epsilon (float)
grad_clip_norm (float)
use_amp (bool)
perf_defaults (bool)
tf32 (bool)
test (bool)
test_interval (int)
test_episodes (int)
scalar_freq (int)
log_video_freq (int)
max_videos_to_save (int)
video_format (str)
video_fps (int)
checkpoint_interval (int)
early_stopping (bool)
patience (int)
min_delta (float)
checkpoint_path (str)
restore (bool)
experience_replay (str)
render (bool)
enable_wandb (bool)
wandb_project (str)
wandb_entity (str)
log_dir (str)
logdir (Any)
data_dir (Any)
log_level (str)
log_file (Any)
enable_tensorboard (bool)
enable_console_metrics (bool)
enable_jsonl (bool)
jsonl_filename (str)
log_system_stats_freq (int)
detect_anomaly (bool)

env_backend: str = 'dmc'

env: str = 'walker-walk'

env_instance: Any = None

image_size: tuple[int, int] = (64, 64)

gym_render_mode: str = 'rgb_array'

frame_stack: int = 1

dmlab_action_repeat: int = 4

dmlab_action_set: Any = None

dmlab_observations: Any = None

dmlab_config: Any = None

dmlab_renderer: str = 'hardware'

procgen_distribution_mode: str = 'easy'

procgen_num_levels: int = 0

procgen_start_level: Any = None

mujoco_xml_path: Any = None

mujoco_xml_string: Any = None

mujoco_binary_path: Any = None

mujoco_camera: Any = None

mujoco_frame_skip: int = 1

mujoco_reset_noise_scale: float = 0.0

brax_backend: str = 'generalized'

brax_jit: bool = True

brax_auto_reset: bool = False

brax_suppress_warp_warnings: bool = True

unity_file_name: Any = None

unity_behavior_name: Any = None

unity_worker_id: int = 0

unity_base_port: int = 5005

unity_no_graphics: bool = True

unity_time_scale: float = 20.0

unity_quality_level: int = 1

algo: str = 'Dreamerv1'

exp_name: str = 'lr1e-3'

train: bool = True

evaluate: bool = False

seed: int = 1

no_gpu: bool = False

max_episode_length: int = 1000

buffer_size: int = 800000

time_limit: int = 1000

cnn_activation_function: str = 'relu'

dense_activation_function: str = 'elu'

obs_embed_size: int = 1024

num_units: int = 400

deter_size: int = 200

stoch_size: int = 30

action_repeat: int = 2

action_noise: float = 0.3

total_steps: int = 5000000

seed_steps: int = 5000

update_steps: int = 100

collect_steps: int = 1000

batch_size: int = 50

train_seq_len: int = 50

imagine_horizon: int = 15

use_disc_model: bool = False

free_nats: float = 3.0

discount: float = 0.99

reward_scale: float = 1.0

td_lambda: float = 0.95

kl_loss_coeff: float = 1.0

kl_alpha: float = 0.8

disc_loss_coeff: float = 10.0

num_buckets: int = 255

symlog_range: float = 10.0

model_learning_rate: float = 0.0006

actor_learning_rate: float = 8e-05

value_learning_rate: float = 8e-05

adam_epsilon: float = 1e-07

grad_clip_norm: float = 100.0

use_amp: bool = False

perf_defaults: bool = True

tf32: bool = True

test: bool = False

test_interval: int = 10000

test_episodes: int = 10

scalar_freq: int = 1000

log_video_freq: int = -1

max_videos_to_save: int = 2

video_format: str = 'gif'

video_fps: int = 20

checkpoint_interval: int = 10000

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

checkpoint_path: str = ''

restore: bool = False

experience_replay: str = ''

render: bool = False

enable_wandb: bool = False

wandb_project: str = 'simulusion'

wandb_entity: str = ''

log_dir: str = 'runs'

logdir: Any = None

data_dir: Any = None

log_level: str = 'INFO'

log_file: Any = None

enable_tensorboard: bool = False

enable_console_metrics: bool = True

enable_jsonl: bool = True

jsonl_filename: str = 'metrics.jsonl'

log_system_stats_freq: int = 1000

detect_anomaly: bool = False

class simulusion.configs.JEPAConfig(use_bfloat16=False, model_name='vit_base', load_checkpoint=False, read_checkpoint=None, copy_data=False, pred_depth=None, pred_emb_dim=384, loss_type='l2', dataset='imagenet', val_split=None, use_gaussian_blur=False, use_horizontal_flip=False, use_color_distortion=False, color_jitter_strength=0.0, batch_size=2048, pin_mem=True, num_workers=8, root_path='/data/imagenet', image_folder='train', crop_size=224, crop_scale=(0.3, 1.0), download=False, allow_overlap=False, patch_size=16, num_enc_masks=1, min_keep=10, enc_mask_scale=(0.85, 1.0), num_pred_masks=4, pred_mask_scale=(0.15, 0.2), aspect_ratio=(0.75, 1.5), ema=(0.996, 1.0), ipe_scale=1.0, weight_decay=0.04, final_weight_decay=0.4, epochs=300, warmup=15, start_lr=0.0001, lr=0.001, final_lr=1e-06, lr_reference_batch_size=2048, early_stopping=False, patience=10, min_delta=0.0001, folder='results/jepa', write_tag='jepa_run', enable_wandb=False, wandb_project='simulusion', wandb_entity='', enable_sweep=False, sweep_config=<factory>)[source]

Bases: SerializableConfigMixin

Minimal configuration container for JEPA training. Converts to the nested dict expected by train_jepa.main.

Parameters:

use_bfloat16 (bool)
model_name (str)
load_checkpoint (bool)
read_checkpoint (str | None)
copy_data (bool)
pred_depth (int | None)
pred_emb_dim (int)
loss_type (str)
dataset (str)
val_split (float | None)
use_gaussian_blur (bool)
use_horizontal_flip (bool)
use_color_distortion (bool)
color_jitter_strength (float)
batch_size (int)
pin_mem (bool)
num_workers (int)
root_path (str)
image_folder (str)
crop_size (int)
crop_scale (Tuple[float, float])
download (bool)
allow_overlap (bool)
patch_size (int)
num_enc_masks (int)
min_keep (int)
enc_mask_scale (Tuple[float, float])
num_pred_masks (int)
pred_mask_scale (Tuple[float, float])
aspect_ratio (Tuple[float, float])
ema (Tuple[float, float])
ipe_scale (float)
weight_decay (float)
final_weight_decay (float)
epochs (int)
warmup (int)
start_lr (float)
lr (float)
final_lr (float)
lr_reference_batch_size (int | None)
early_stopping (bool)
patience (int)
min_delta (float)
folder (str)
write_tag (str)
enable_wandb (bool)
wandb_project (str)
wandb_entity (str)
enable_sweep (bool)
sweep_config (Dict[str, Any])

use_bfloat16: bool = False

model_name: str = 'vit_base'

load_checkpoint: bool = False

read_checkpoint: str | None = None

copy_data: bool = False

pred_depth: int | None = None

pred_emb_dim: int = 384

loss_type: str = 'l2'

dataset: str = 'imagenet'

val_split: float | None = None

use_gaussian_blur: bool = False

use_horizontal_flip: bool = False

use_color_distortion: bool = False

color_jitter_strength: float = 0.0

batch_size: int = 2048

pin_mem: bool = True

num_workers: int = 8

root_path: str = '/data/imagenet'

image_folder: str = 'train'

crop_size: int = 224

crop_scale: Tuple[float, float] = (0.3, 1.0)

download: bool = False

allow_overlap: bool = False

patch_size: int = 16

num_enc_masks: int = 1

min_keep: int = 10

enc_mask_scale: Tuple[float, float] = (0.85, 1.0)

num_pred_masks: int = 4

pred_mask_scale: Tuple[float, float] = (0.15, 0.2)

aspect_ratio: Tuple[float, float] = (0.75, 1.5)

ema: Tuple[float, float] = (0.996, 1.0)

ipe_scale: float = 1.0

weight_decay: float = 0.04

final_weight_decay: float = 0.4

epochs: int = 300

warmup: int = 15

start_lr: float = 0.0001

lr: float = 0.001

final_lr: float = 1e-06

lr_reference_batch_size: int | None = 2048

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

folder: str = 'results/jepa'

write_tag: str = 'jepa_run'

enable_wandb: bool = False

wandb_project: str = 'simulusion'

wandb_entity: str = ''

enable_sweep: bool = False

sweep_config: Dict[str, Any]

to_dict()[source]

Return type:: Dict[str, Dict[str, Any]]

classmethod from_dict(values)[source]

Load flat field values or the nested trainer dictionary.

Parameters:: values (Dict[str, Any])
Return type:: JEPAConfig

to_train_dict()[source]

Return the nested dictionary expected by train_jepa.

Return type:: Dict[str, Dict[str, Any]]

to_nested_dict()[source]

Backward-compatible alias for the nested JEPA trainer dictionary.

Return type:: Dict[str, Dict[str, Any]]

class simulusion.configs.DiTConfig(DATASET='CIFAR10', BATCH=256, EPOCHS=3, LR=0.0001, WEIGHT_DECAY=0.0, IMG_SIZE=32, CHANNELS=3, PATCH=4, WIDTH=384, DEPTH=12, HEADS=6, DROP=0.0, NUM_CLASSES=0, CLASS_DROPOUT_PROB=0.1, LEARN_SIGMA=True, BETA_START=0.0001, BETA_END=0.02, TIMESTEPS=1000, EMA=True, EMA_DECAY=0.9999, WORKDIR='./dit_demo', ROOT_PATH='./data', CHECKPOINT_EVERY=0, EARLY_STOPPING=False, PATIENCE=10, MIN_DELTA=0.0001, VAL_SPLIT=0.05, CROP_SIZE=None, NUM_WORKERS=4)[source]

Bases: SerializableConfigMixin

Default configuration values for Diffusion Transformer (DiT) training.

The fields define dataset selection, model architecture, diffusion schedule, optimization hyperparameters, and output paths used by the built-in training entrypoints.

Field names use UPPER_CASE for backward compatibility with the original DiT codebase. Snake-case aliases are accepted via __getattr__ and get_dit_config().

Parameters:

DATASET (str)
BATCH (int)
EPOCHS (int)
LR (float)
WEIGHT_DECAY (float)
IMG_SIZE (int)
CHANNELS (int)
PATCH (int)
WIDTH (int)
DEPTH (int)
HEADS (int)
DROP (float)
NUM_CLASSES (int)
CLASS_DROPOUT_PROB (float)
LEARN_SIGMA (bool)
BETA_START (float)
BETA_END (float)
TIMESTEPS (int)
EMA (bool)
EMA_DECAY (float)
WORKDIR (str)
ROOT_PATH (str)
CHECKPOINT_EVERY (int)
EARLY_STOPPING (bool)
PATIENCE (int)
MIN_DELTA (float)
VAL_SPLIT (float)
CROP_SIZE (Any)
NUM_WORKERS (int)

DATASET: str = 'CIFAR10'

BATCH: int = 256

EPOCHS: int = 3

LR: float = 0.0001

WEIGHT_DECAY: float = 0.0

IMG_SIZE: int = 32

CHANNELS: int = 3

PATCH: int = 4

WIDTH: int = 384

DEPTH: int = 12

HEADS: int = 6

DROP: float = 0.0

NUM_CLASSES: int = 0

CLASS_DROPOUT_PROB: float = 0.1

LEARN_SIGMA: bool = True

BETA_START: float = 0.0001

BETA_END: float = 0.02

TIMESTEPS: int = 1000

EMA: bool = True

EMA_DECAY: float = 0.9999

WORKDIR: str = './dit_demo'

ROOT_PATH: str = './data'

CHECKPOINT_EVERY: int = 0

EARLY_STOPPING: bool = False

PATIENCE: int = 10

MIN_DELTA: float = 0.0001

VAL_SPLIT: float = 0.05

CROP_SIZE: Any = None

NUM_WORKERS: int = 4

simulusion.configs.get_dit_config(**overrides)[source]

Returns a DiTConfig instance with default values overridden by the provided keyword arguments.

Both UPPER_CASE and snake_case override keys are accepted.

Example usage:: cfg = get_dit_config(BATCH=64, EPOCHS=10, LR=1e-3) cfg = get_dit_config(batch=64, epochs=10, lr=1e-3)

Parameters:: overrides (Any)
Return type:: DiTConfig

simulusion.configs.dit_preset_config(name, patch_size, **overrides)[source]

Build a config for a named Table 1 model, e.g. dit_preset_config("DiT-XL", 2).

Parameters:

name (str) – One of DiT-S, DiT-B, DiT-L, DiT-XL (case-insensitive, and the DiT- prefix is optional).
patch_size (int) – Latent patch size; the paper explores 2, 4 and 8.
**overrides (Any) – Further config fields, UPPER_CASE or snake_case.

Returns:

A DiTConfig for that model. Defaults target latent diffusion of 256x256 ImageNet: a 32x32x4 latent with 1000 classes.

Return type:

DiTConfig

simulusion.configs.list_dit_presets()[source]

Return the Table 1 model names.

Return type:: list[str]

class simulusion.configs.DiamondConfig(preset: str | None = None, game: str = 'Breakout-v5', seed: int = 0, obs_size: int = 64, frameskip: int = 4, max_noop: int = 30, terminate_on_life_loss: bool = True, reward_clip: List[int] = <factory>, max_episode_steps: int = 27000, num_conditioning_frames: int = 4, diffusion_channels: List[int] = <factory>, diffusion_res_blocks: int = 2, diffusion_cond_dim: int = 256, sigma_data: float = 0.5, sigma_min: float = 0.002, sigma_max: float = 80.0, rho: int = 7, p_mean: float = -0.4, p_std: float = 1.2, sampling_method: str = 'euler', num_sampling_steps: int = 3, reward_channels: List[int] = <factory>, reward_res_blocks: int = 2, reward_cond_dim: int = 128, reward_lstm_dim: int = 512, burn_in_length: int = 4, actor_channels: List[int] = <factory>, actor_res_blocks: int = 1, actor_lstm_dim: int = 512, num_epochs: int = 1000, training_steps_per_epoch: int = 400, batch_size: int = 32, environment_steps_per_epoch: int = 100, epsilon_greedy: float = 0.01, data_loader_num_workers: int = 4, pin_memory: bool = True, persistent_workers: bool = True, imagination_horizon: int = 15, discount_factor: float = 0.985, entropy_weight: float = 0.001, lambda_returns: float = 0.95, learning_rate: float = 0.0001, adam_epsilon: float = 1e-08, weight_decay_diffusion: float = 0.01, weight_decay_reward: float = 0.01, weight_decay_actor: float = 0.0, use_amp: bool = True, device: str = <factory>, log_interval: int = 10, eval_interval: int = 50, save_interval: int = 100, checkpoint_dir: str = 'checkpoints/diamond', early_stopping: bool = False, patience: int = 10, min_delta: float = 0.0001)[source]

Bases: SerializableConfigMixin

Parameters:

preset (str | None)
game (str)
seed (int)
obs_size (int)
frameskip (int)
max_noop (int)
terminate_on_life_loss (bool)
reward_clip (List[int])
max_episode_steps (int)
num_conditioning_frames (int)
diffusion_channels (List[int])
diffusion_res_blocks (int)
diffusion_cond_dim (int)
sigma_data (float)
sigma_min (float)
sigma_max (float)
rho (int)
p_mean (float)
p_std (float)
sampling_method (str)
num_sampling_steps (int)
reward_channels (List[int])
reward_res_blocks (int)
reward_cond_dim (int)
reward_lstm_dim (int)
burn_in_length (int)
actor_channels (List[int])
actor_res_blocks (int)
actor_lstm_dim (int)
num_epochs (int)
training_steps_per_epoch (int)
batch_size (int)
environment_steps_per_epoch (int)
epsilon_greedy (float)
data_loader_num_workers (int)
pin_memory (bool)
persistent_workers (bool)
imagination_horizon (int)
discount_factor (float)
entropy_weight (float)
lambda_returns (float)
learning_rate (float)
adam_epsilon (float)
weight_decay_diffusion (float)
weight_decay_reward (float)
weight_decay_actor (float)
use_amp (bool)
device (str)
log_interval (int)
eval_interval (int)
save_interval (int)
checkpoint_dir (str)
early_stopping (bool)
patience (int)
min_delta (float)

preset: str | None = None

game: str = 'Breakout-v5'

seed: int = 0

obs_size: int = 64

frameskip: int = 4

max_noop: int = 30

terminate_on_life_loss: bool = True

reward_clip: List[int]

max_episode_steps: int = 27000

num_conditioning_frames: int = 4

diffusion_channels: List[int]

diffusion_res_blocks: int = 2

diffusion_cond_dim: int = 256

sigma_data: float = 0.5

sigma_min: float = 0.002

sigma_max: float = 80.0

rho: int = 7

p_mean: float = -0.4

p_std: float = 1.2

sampling_method: str = 'euler'

num_sampling_steps: int = 3

reward_channels: List[int]

reward_res_blocks: int = 2

reward_cond_dim: int = 128

reward_lstm_dim: int = 512

burn_in_length: int = 4

actor_channels: List[int]

actor_res_blocks: int = 1

actor_lstm_dim: int = 512

num_epochs: int = 1000

training_steps_per_epoch: int = 400

batch_size: int = 32

environment_steps_per_epoch: int = 100

epsilon_greedy: float = 0.01

data_loader_num_workers: int = 4

pin_memory: bool = True

persistent_workers: bool = True

imagination_horizon: int = 15

discount_factor: float = 0.985

entropy_weight: float = 0.001

lambda_returns: float = 0.95

learning_rate: float = 0.0001

adam_epsilon: float = 1e-08

weight_decay_diffusion: float = 0.01

weight_decay_reward: float = 0.01

weight_decay_actor: float = 0.0

use_amp: bool = True

device: str

log_interval: int = 10

eval_interval: int = 50

save_interval: int = 100

checkpoint_dir: str = 'checkpoints/diamond'

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

class simulusion.configs.IRISConfig(frame_height=64, frame_width=64, frame_channels=3, vocab_size=512, tokens_per_frame=16, token_embedding_dim=512, encoder_channels=64, encoder_layers=4, encoder_residual_blocks=2, decoder_depth=64, reconstruction_weight=1.0, commitment_weight=1.0, perceptual_weight=1.0, perceptual_blocks=5, perceptual_linear_weights='', quantizer='gradient', transformer_timesteps=20, transformer_embed_dim=256, transformer_layers=10, transformer_heads=4, transformer_dropout=0.1, reward_transform='sign', reward_loss='cross_entropy', imagination_horizon=20, burn_in_length=20, discount=0.995, td_lambda=0.95, entropy_coef=0.001, actor_hidden_size=512, actor_layers=1, total_epochs=600, collection_epochs=500, env_steps_per_epoch=200, training_steps_per_epoch=200, transformer_steps_per_epoch=200, actor_critic_steps_per_epoch=200, model_learning_rate=0.0001, actor_learning_rate=0.0001, value_learning_rate=0.0001, adam_beta1=0.9, adam_beta2=0.999, weight_decay=0.01, grad_clip_norm=10.0, use_amp=True, gradient_checkpointing=True, collect_epsilon=0.01, eval_temperature=0.5, collect_temperature=1.0, start_autoencoder_after=5, start_transformer_after=25, start_actor_critic_after=50, autoencoder_batch_size=256, transformer_batch_size=64, actor_critic_batch_size=64, atari_100k=True, max_env_steps=100000, env_backend='gym', env='ALE/Pong-v5', action_repeat=4, repeat_action_probability=0.0, max_episode_steps=27000, log_interval=1000, eval_episodes=100, checkpoint_interval=50, early_stopping=False, patience=10, min_delta=0.0001)[source]

Bases: SerializableConfigMixin

Configuration for IRIS (Imagination with auto-Regression over an Inner Speech)

Based on paper: “Transformers are Sample-Efficient World Models” Implements discrete autoencoder + autoregressive Transformer for sample-efficient RL.

Parameters:

frame_height (int)
frame_width (int)
frame_channels (int)
vocab_size (int)
tokens_per_frame (int)
token_embedding_dim (int)
encoder_channels (int)
encoder_layers (int)
encoder_residual_blocks (int)
decoder_depth (int)
reconstruction_weight (float)
commitment_weight (float)
perceptual_weight (float)
perceptual_blocks (int)
perceptual_linear_weights (str)
quantizer (str)
transformer_timesteps (int)
transformer_embed_dim (int)
transformer_layers (int)
transformer_heads (int)
transformer_dropout (float)
reward_transform (str)
reward_loss (str)
imagination_horizon (int)
burn_in_length (int)
discount (float)
td_lambda (float)
entropy_coef (float)
actor_hidden_size (int)
actor_layers (int)
total_epochs (int)
collection_epochs (int)
env_steps_per_epoch (int)
training_steps_per_epoch (int)
transformer_steps_per_epoch (int)
actor_critic_steps_per_epoch (int)
model_learning_rate (float)
actor_learning_rate (float)
value_learning_rate (float)
adam_beta1 (float)
adam_beta2 (float)
weight_decay (float)
grad_clip_norm (float)
use_amp (bool)
gradient_checkpointing (bool)
collect_epsilon (float)
eval_temperature (float)
collect_temperature (float)
start_autoencoder_after (int)
start_transformer_after (int)
start_actor_critic_after (int)
autoencoder_batch_size (int)
transformer_batch_size (int)
actor_critic_batch_size (int)
atari_100k (bool)
max_env_steps (int)
env_backend (str)
env (str)
action_repeat (int)
repeat_action_probability (float)
max_episode_steps (int)
log_interval (int)
eval_episodes (int)
checkpoint_interval (int)
early_stopping (bool)
patience (int)
min_delta (float)

frame_height: int = 64

frame_width: int = 64

frame_channels: int = 3

vocab_size: int = 512

tokens_per_frame: int = 16

token_embedding_dim: int = 512

encoder_channels: int = 64

encoder_layers: int = 4

encoder_residual_blocks: int = 2

decoder_depth: int = 64

reconstruction_weight: float = 1.0

commitment_weight: float = 1.0

perceptual_weight: float = 1.0

perceptual_blocks: int = 5

perceptual_linear_weights: str = ''

quantizer: str = 'gradient'

transformer_timesteps: int = 20

transformer_embed_dim: int = 256

transformer_layers: int = 10

transformer_heads: int = 4

transformer_dropout: float = 0.1

reward_transform: str = 'sign'

reward_loss: str = 'cross_entropy'

imagination_horizon: int = 20

burn_in_length: int = 20

discount: float = 0.995

td_lambda: float = 0.95

entropy_coef: float = 0.001

actor_hidden_size: int = 512

actor_layers: int = 1

total_epochs: int = 600

collection_epochs: int = 500

env_steps_per_epoch: int = 200

training_steps_per_epoch: int = 200

transformer_steps_per_epoch: int = 200

actor_critic_steps_per_epoch: int = 200

model_learning_rate: float = 0.0001

actor_learning_rate: float = 0.0001

value_learning_rate: float = 0.0001

adam_beta1: float = 0.9

adam_beta2: float = 0.999

weight_decay: float = 0.01

grad_clip_norm: float = 10.0

use_amp: bool = True

gradient_checkpointing: bool = True

collect_epsilon: float = 0.01

eval_temperature: float = 0.5

collect_temperature: float = 1.0

start_autoencoder_after: int = 5

start_transformer_after: int = 25

start_actor_critic_after: int = 50

autoencoder_batch_size: int = 256

transformer_batch_size: int = 64

actor_critic_batch_size: int = 64

atari_100k: bool = True

max_env_steps: int = 100000

env_backend: str = 'gym'

env: str = 'ALE/Pong-v5'

action_repeat: int = 4

repeat_action_probability: float = 0.0

max_episode_steps: int = 27000

log_interval: int = 1000

eval_episodes: int = 100

checkpoint_interval: int = 50

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

get_frame_shape()[source]

Return type:: tuple

get_autoencoder_config()[source]

Return type:: dict

get_transformer_config()[source]

Return type:: dict

get_rl_config()[source]

Return type:: dict

class simulusion.configs.GenieConfig(num_frames=8, image_size=32, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, tokenizer_encoder_dim=256, tokenizer_decoder_dim=512, tokenizer_encoder_depth=4, tokenizer_decoder_depth=8, tokenizer_num_heads=16, action_vocab_size=8, action_embedding_dim=32, action_encoder_dim=256, action_decoder_dim=1024, action_encoder_depth=4, action_num_heads=16, action_pooling='mean', window_attention_heads=1, dynamics_dim=512, dynamics_depth=8, dynamics_num_heads=8, batch_size=4, learning_rate=3e-05, weight_decay=0.0001, warmup_steps=5000, max_steps=125000, mask_prob_min=0.5, mask_prob_max=1.0, sample_temperature=2.0, maskgit_steps=25, use_amp=False, early_stopping=False, patience=10, min_delta=0.0001, val_split=0.1)[source]

Bases: SerializableConfigMixin

Configuration for Genie model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
tokenizer_encoder_dim (int)
tokenizer_decoder_dim (int)
tokenizer_encoder_depth (int)
tokenizer_decoder_depth (int)
tokenizer_num_heads (int)
action_vocab_size (int)
action_embedding_dim (int)
action_encoder_dim (int)
action_decoder_dim (int)
action_encoder_depth (int)
action_num_heads (int)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
batch_size (int)
learning_rate (float)
weight_decay (float)
warmup_steps (int)
max_steps (int)
mask_prob_min (float)
mask_prob_max (float)
sample_temperature (float)
maskgit_steps (int)
use_amp (bool)
early_stopping (bool)
patience (int)
min_delta (float)
val_split (float)

num_frames: int = 8

image_size: int = 32

in_channels: int = 3

tokenizer_vocab_size: int = 1024

tokenizer_embedding_dim: int = 32

tokenizer_encoder_dim: int = 256

tokenizer_decoder_dim: int = 512

tokenizer_encoder_depth: int = 4

tokenizer_decoder_depth: int = 8

tokenizer_num_heads: int = 16

action_vocab_size: int = 8

action_embedding_dim: int = 32

action_encoder_dim: int = 256

action_decoder_dim: int = 1024

action_encoder_depth: int = 4

action_num_heads: int = 16

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'

window_attention_heads: int = 1

dynamics_dim: int = 512

dynamics_depth: int = 8

dynamics_num_heads: int = 8

batch_size: int = 4

learning_rate: float = 3e-05

weight_decay: float = 0.0001

warmup_steps: int = 5000

max_steps: int = 125000

mask_prob_min: float = 0.5

mask_prob_max: float = 1.0

sample_temperature: float = 2.0

maskgit_steps: int = 25

use_amp: bool = False

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

val_split: float = 0.1

class simulusion.configs.GenieSmallConfig(num_frames=16, image_size=64, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, tokenizer_encoder_dim=256, tokenizer_decoder_dim=512, tokenizer_encoder_depth=4, tokenizer_decoder_depth=8, tokenizer_num_heads=16, action_vocab_size=8, action_embedding_dim=32, action_encoder_dim=512, action_decoder_dim=1024, action_encoder_depth=8, action_num_heads=16, action_pooling='mean', window_attention_heads=1, dynamics_dim=512, dynamics_depth=8, dynamics_num_heads=8, batch_size=2, learning_rate=0.0001, weight_decay=0.0001, warmup_steps=1000, max_steps=50000, mask_prob_min=0.5, mask_prob_max=1.0, sample_temperature=2.0, maskgit_steps=25, use_amp=False, early_stopping=False, patience=10, min_delta=0.0001, val_split=0.1)[source]

Bases: SerializableConfigMixin

Small configuration for development/testing.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
tokenizer_encoder_dim (int)
tokenizer_decoder_dim (int)
tokenizer_encoder_depth (int)
tokenizer_decoder_depth (int)
tokenizer_num_heads (int)
action_vocab_size (int)
action_embedding_dim (int)
action_encoder_dim (int)
action_decoder_dim (int)
action_encoder_depth (int)
action_num_heads (int)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
batch_size (int)
learning_rate (float)
weight_decay (float)
warmup_steps (int)
max_steps (int)
mask_prob_min (float)
mask_prob_max (float)
sample_temperature (float)
maskgit_steps (int)
use_amp (bool)
early_stopping (bool)
patience (int)
min_delta (float)
val_split (float)

num_frames: int = 16

image_size: int = 64

in_channels: int = 3

tokenizer_vocab_size: int = 1024

tokenizer_embedding_dim: int = 32

tokenizer_encoder_dim: int = 256

tokenizer_decoder_dim: int = 512

tokenizer_encoder_depth: int = 4

tokenizer_decoder_depth: int = 8

tokenizer_num_heads: int = 16

action_vocab_size: int = 8

action_embedding_dim: int = 32

action_encoder_dim: int = 512

action_decoder_dim: int = 1024

action_encoder_depth: int = 8

action_num_heads: int = 16

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'

window_attention_heads: int = 1

dynamics_dim: int = 512

dynamics_depth: int = 8

dynamics_num_heads: int = 8

batch_size: int = 2

learning_rate: float = 0.0001

weight_decay: float = 0.0001

warmup_steps: int = 1000

max_steps: int = 50000

mask_prob_min: float = 0.5

mask_prob_max: float = 1.0

sample_temperature: float = 2.0

maskgit_steps: int = 25

use_amp: bool = False

early_stopping: bool = False

patience: int = 10

min_delta: float = 0.0001

val_split: float = 0.1

class simulusion.configs.STTransformerConfig(num_frames=16, num_patches_per_frame=256, dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0)[source]

Bases: SerializableConfigMixin

Configuration for Spatiotemporal Transformer.

Parameters:

num_frames (int)
num_patches_per_frame (int)
dim (int)
depth (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)

num_frames: int = 16

num_patches_per_frame: int = 256

dim: int = 768

depth: int = 12

num_heads: int = 12

mlp_ratio: float = 4.0

qkv_bias: bool = True

drop_rate: float = 0.0

attn_drop_rate: float = 0.0

drop_path_rate: float = 0.0

class simulusion.configs.VideoTokenizerConfig(num_frames=16, image_size=64, in_channels=3, encoder_dim=512, decoder_dim=1024, encoder_depth=12, decoder_depth=20, num_heads=16, patch_size=4, vocab_size=1024, embedding_dim=32, use_ema=False, ema_decay=0.99, commitment_weight=0.25)[source]

Bases: SerializableConfigMixin

Configuration for Video Tokenizer.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
use_ema (bool)
ema_decay (float)
commitment_weight (float)

num_frames: int = 16

image_size: int = 64

in_channels: int = 3

encoder_dim: int = 512

decoder_dim: int = 1024

encoder_depth: int = 12

decoder_depth: int = 20

num_heads: int = 16

patch_size: int = 4

vocab_size: int = 1024

embedding_dim: int = 32

use_ema: bool = False

ema_decay: float = 0.99

commitment_weight: float = 0.25

class simulusion.configs.LatentActionModelConfig(num_frames=16, image_size=64, in_channels=3, encoder_dim=1024, encoder_depth=20, num_heads=16, patch_size=16, vocab_size=8, embedding_dim=32, commitment_weight=1.0, action_pooling='mean', window_attention_heads=1)[source]

Bases: SerializableConfigMixin

Configuration for Latent Action Model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
encoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
commitment_weight (float)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

num_frames: int = 16

image_size: int = 64

in_channels: int = 3

encoder_dim: int = 1024

encoder_depth: int = 20

num_heads: int = 16

patch_size: int = 16

vocab_size: int = 8

embedding_dim: int = 32

commitment_weight: float = 1.0

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'

window_attention_heads: int = 1

class simulusion.configs.DynamicsModelConfig(num_frames=16, image_size=64, vocab_size=1024, embedding_dim=32, action_vocab_size=8, dim=5120, depth=48, num_heads=40, patch_size=4, mlp_ratio=4.0, qkv_bias=True, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0)[source]

Bases: SerializableConfigMixin

Configuration for Dynamics Model.

Parameters:

num_frames (int)
image_size (int)
vocab_size (int)
embedding_dim (int)
action_vocab_size (int)
dim (int)
depth (int)
num_heads (int)
patch_size (int)
mlp_ratio (float)
qkv_bias (bool)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)

num_frames: int = 16

image_size: int = 64

vocab_size: int = 1024

embedding_dim: int = 32

action_vocab_size: int = 8

dim: int = 5120

depth: int = 48

num_heads: int = 40

patch_size: int = 4

mlp_ratio: float = 4.0

qkv_bias: bool = True

drop_rate: float = 0.0

attn_drop_rate: float = 0.0

drop_path_rate: float = 0.0

Configuration classes for World Models training.

class simulusion.configs.wm_config.WMVAEConfig(height, width, device='cuda', train_batch_size=32, num_epochs=10, latent_size=32, data_dir='./data', learning_rate=0.001, logdir='results', noreload=False, nosamples=False, scheduler_patience=5, scheduler_factor=0.5, early_stopping_patience=30, sample_interval=5)[source]#

Bases: object

Configuration class for Variational Autoencoder (VAE) training.

Manages all hyperparameters and settings for training a ConvVAE model on observation data.

Variables:

height (int) – Height of input images (pixels).
width (int) – Width of input images (pixels).
device (str) – Device to train on (‘cpu’ or ‘cuda’).
train_batch_size (int) – Number of samples per training batch.
num_epochs (int) – Total number of training epochs.
latent_size (int) – Dimensionality of the VAE latent space.
data_dir (str) – Path to the dataset directory.
learning_rate (float) – Initial learning rate for optimizer.
logdir (str) – Directory for saving logs and checkpoints.
noreload (bool) – If True, skip loading existing checkpoints.
nosamples (bool) – If True, skip saving sample images during training.
scheduler_patience (int) – Epochs to wait before reducing learning rate.
scheduler_factor (float) – Multiplicative factor for learning rate reduction.
early_stopping_patience (int) – Epochs to wait before early stopping.
sample_interval (int) – Epoch interval for saving sample images.

Parameters:

height (int)
width (int)
device (str)
train_batch_size (int)
num_epochs (int)
latent_size (int)
data_dir (str)
learning_rate (float)
logdir (str)
noreload (bool)
nosamples (bool)
scheduler_patience (int)
scheduler_factor (float)
early_stopping_patience (int)
sample_interval (int)

height: int#

width: int#

device: str = 'cuda'#

train_batch_size: int = 32#

num_epochs: int = 10#

latent_size: int = 32#

data_dir: str = './data'#

learning_rate: float = 0.001#

logdir: str = 'results'#

noreload: bool = False#

nosamples: bool = False#

scheduler_patience: int = 5#

scheduler_factor: float = 0.5#

early_stopping_patience: int = 30#

sample_interval: int = 5#

validate()[source]#

Return type:: bool

to_dict()[source]#

Return type:: Dict[str, Any]

class simulusion.configs.wm_config.WMMDNRNNConfig(latent_size=32, action_size=3, hidden_size=256, gmm_components=5, device='cuda', batch_size=16, seq_len=32, num_epochs=30, data_dir='./data', learning_rate=0.001, logdir='results', noreload=False, include_reward=True, scheduler_patience=5, scheduler_factor=0.5, early_stopping_patience=30)[source]#

Bases: object

Configuration class for Mixture Density Recurrent Neural Network (MDRNN) training.

Manages all hyperparameters and settings for training an MDRNN model on sequence data.

Variables:

latent_size (int) – Dimensionality of the latent space from VAE.
action_size (int) – Dimensionality of action space.
hidden_size (int) – Number of hidden units in RNN.
gmm_components (int) – Number of Gaussian mixture components.
device (str) – Device to train on (‘cpu’ or ‘cuda’).
batch_size (int) – Number of sequences per batch.
seq_len (int) – Length of each sequence.
num_epochs (int) – Total number of training epochs.
data_dir (str) – Path to the dataset directory.
learning_rate (float) – Initial learning rate for optimizer.
logdir (str) – Directory for saving logs and checkpoints.
noreload (bool) – If True, skip loading existing checkpoints.
include_reward (bool) – If True, include reward prediction in loss.
scheduler_patience (int) – Epochs to wait before reducing learning rate.
scheduler_factor (float) – Multiplicative factor for learning rate reduction.
early_stopping_patience (int) – Epochs to wait before early stopping.

Parameters:

latent_size (int)
action_size (int)
hidden_size (int)
gmm_components (int)
device (str)
batch_size (int)
seq_len (int)
num_epochs (int)
data_dir (str)
learning_rate (float)
logdir (str)
noreload (bool)
include_reward (bool)
scheduler_patience (int)
scheduler_factor (float)
early_stopping_patience (int)

latent_size: int = 32#

action_size: int = 3#

hidden_size: int = 256#

gmm_components: int = 5#

device: str = 'cuda'#

batch_size: int = 16#

seq_len: int = 32#

num_epochs: int = 30#

data_dir: str = './data'#

learning_rate: float = 0.001#

logdir: str = 'results'#

noreload: bool = False#

include_reward: bool = True#

scheduler_patience: int = 5#

scheduler_factor: float = 0.5#

early_stopping_patience: int = 30#

validate()[source]#

Return type:: bool

to_dict()[source]#

Return type:: Dict[str, Any]

class simulusion.configs.wm_config.WMControllerConfig(latent_size=32, hidden_size=200, action_size=3, env_name='CarRacing-v3', logdir='results', n_samples=4, pop_size=10, target_return=950.0, max_workers=32, display=True, time_limit=1000)[source]#

Bases: object

Configuration class for Controller training with CMA-ES.

Manages hyperparameters for training a linear controller using Covariance Matrix Adaptation Evolution Strategy (CMA-ES).

Variables:

latent_size (int) – Dimensionality of latent state from VAE.
hidden_size (int) – Dimensionality of RSSM hidden state.
action_size (int) – Dimensionality of action space.
env_name (str) – Gym environment name.
logdir (str) – Directory for saving logs and checkpoints.
n_samples (int) – Number of samples used to obtain return estimate.
pop_size (int) – Population size for CMA-ES.
target_return (float) – Stop once the return gets above this threshold.
max_workers (int) – Maximum number of workers for parallel evaluation.
display (bool) – If True, show progress bars during training.
time_limit (int) – Maximum steps per episode.

Parameters:

latent_size (int)
hidden_size (int)
action_size (int)
env_name (str)
logdir (str)
n_samples (int)
pop_size (int)
target_return (float)
max_workers (int)
display (bool)
time_limit (int)

latent_size: int = 32#

hidden_size: int = 200#

action_size: int = 3#

env_name: str = 'CarRacing-v3'#

logdir: str = 'results'#

n_samples: int = 4#

pop_size: int = 10#

target_return: float = 950.0#

max_workers: int = 32#

display: bool = True#

time_limit: int = 1000#

validate()[source]#

Return type:: bool

to_dict()[source]#

Return type:: Dict[str, Any]

class simulusion.configs.dreamer_config.DreamerConfig(env_backend='dmc', env='walker-walk', env_instance=None, image_size=(64, 64), gym_render_mode='rgb_array', frame_stack=1, dmlab_action_repeat=4, dmlab_action_set=None, dmlab_observations=None, dmlab_config=None, dmlab_renderer='hardware', procgen_distribution_mode='easy', procgen_num_levels=0, procgen_start_level=None, mujoco_xml_path=None, mujoco_xml_string=None, mujoco_binary_path=None, mujoco_camera=None, mujoco_frame_skip=1, mujoco_reset_noise_scale=0.0, brax_backend='generalized', brax_jit=True, brax_auto_reset=False, brax_suppress_warp_warnings=True, unity_file_name=None, unity_behavior_name=None, unity_worker_id=0, unity_base_port=5005, unity_no_graphics=True, unity_time_scale=20.0, unity_quality_level=1, algo='Dreamerv1', exp_name='lr1e-3', train=True, evaluate=False, seed=1, no_gpu=False, max_episode_length=1000, buffer_size=800000, time_limit=1000, cnn_activation_function='relu', dense_activation_function='elu', obs_embed_size=1024, num_units=400, deter_size=200, stoch_size=30, action_repeat=2, action_noise=0.3, total_steps=5000000, seed_steps=5000, update_steps=100, collect_steps=1000, batch_size=50, train_seq_len=50, imagine_horizon=15, use_disc_model=False, free_nats=3.0, discount=0.99, reward_scale=1.0, td_lambda=0.95, kl_loss_coeff=1.0, kl_alpha=0.8, disc_loss_coeff=10.0, num_buckets=255, symlog_range=10.0, model_learning_rate=0.0006, actor_learning_rate=8e-05, value_learning_rate=8e-05, adam_epsilon=1e-07, grad_clip_norm=100.0, use_amp=False, perf_defaults=True, tf32=True, test=False, test_interval=10000, test_episodes=10, scalar_freq=1000, log_video_freq=-1, max_videos_to_save=2, video_format='gif', video_fps=20, checkpoint_interval=10000, early_stopping=False, patience=10, min_delta=0.0001, checkpoint_path='', restore=False, experience_replay='', render=False, enable_wandb=False, wandb_project='simulusion', wandb_entity='', log_dir='runs', logdir=None, data_dir=None, log_level='INFO', log_file=None, enable_tensorboard=False, enable_console_metrics=True, enable_jsonl=True, jsonl_filename='metrics.jsonl', log_system_stats_freq=1000, detect_anomaly=False)[source]#

Bases: SerializableConfigMixin

Configuration container for Dreamer training, evaluation, and environment setup.

Parameters:

env_backend (str)
env (str)
env_instance (Any)
image_size (tuple[int, int])
gym_render_mode (str)
frame_stack (int)
dmlab_action_repeat (int)
dmlab_action_set (Any)
dmlab_observations (Any)
dmlab_config (Any)
dmlab_renderer (str)
procgen_distribution_mode (str)
procgen_num_levels (int)
procgen_start_level (Any)
mujoco_xml_path (Any)
mujoco_xml_string (Any)
mujoco_binary_path (Any)
mujoco_camera (Any)
mujoco_frame_skip (int)
mujoco_reset_noise_scale (float)
brax_backend (str)
brax_jit (bool)
brax_auto_reset (bool)
brax_suppress_warp_warnings (bool)
unity_file_name (Any)
unity_behavior_name (Any)
unity_worker_id (int)
unity_base_port (int)
unity_no_graphics (bool)
unity_time_scale (float)
unity_quality_level (int)
algo (str)
exp_name (str)
train (bool)
evaluate (bool)
seed (int)
no_gpu (bool)
max_episode_length (int)
buffer_size (int)
time_limit (int)
cnn_activation_function (str)
dense_activation_function (str)
obs_embed_size (int)
num_units (int)
deter_size (int)
stoch_size (int)
action_repeat (int)
action_noise (float)
total_steps (int)
seed_steps (int)
update_steps (int)
collect_steps (int)
batch_size (int)
train_seq_len (int)
imagine_horizon (int)
use_disc_model (bool)
free_nats (float)
discount (float)
reward_scale (float)
td_lambda (float)
kl_loss_coeff (float)
kl_alpha (float)
disc_loss_coeff (float)
num_buckets (int)
symlog_range (float)
model_learning_rate (float)
actor_learning_rate (float)
value_learning_rate (float)
adam_epsilon (float)
grad_clip_norm (float)
use_amp (bool)
perf_defaults (bool)
tf32 (bool)
test (bool)
test_interval (int)
test_episodes (int)
scalar_freq (int)
log_video_freq (int)
max_videos_to_save (int)
video_format (str)
video_fps (int)
checkpoint_interval (int)
early_stopping (bool)
patience (int)
min_delta (float)
checkpoint_path (str)
restore (bool)
experience_replay (str)
render (bool)
enable_wandb (bool)
wandb_project (str)
wandb_entity (str)
log_dir (str)
logdir (Any)
data_dir (Any)
log_level (str)
log_file (Any)
enable_tensorboard (bool)
enable_console_metrics (bool)
enable_jsonl (bool)
jsonl_filename (str)
log_system_stats_freq (int)
detect_anomaly (bool)

env_backend: str = 'dmc'#

env: str = 'walker-walk'#

env_instance: Any = None#

image_size: tuple[int, int] = (64, 64)#

gym_render_mode: str = 'rgb_array'#

frame_stack: int = 1#

dmlab_action_repeat: int = 4#

dmlab_action_set: Any = None#

dmlab_observations: Any = None#

dmlab_config: Any = None#

dmlab_renderer: str = 'hardware'#

procgen_distribution_mode: str = 'easy'#

procgen_num_levels: int = 0#

procgen_start_level: Any = None#

mujoco_xml_path: Any = None#

mujoco_xml_string: Any = None#

mujoco_binary_path: Any = None#

mujoco_camera: Any = None#

mujoco_frame_skip: int = 1#

mujoco_reset_noise_scale: float = 0.0#

brax_backend: str = 'generalized'#

brax_jit: bool = True#

brax_auto_reset: bool = False#

brax_suppress_warp_warnings: bool = True#

unity_file_name: Any = None#

unity_behavior_name: Any = None#

unity_worker_id: int = 0#

unity_base_port: int = 5005#

unity_no_graphics: bool = True#

unity_time_scale: float = 20.0#

unity_quality_level: int = 1#

algo: str = 'Dreamerv1'#

exp_name: str = 'lr1e-3'#

train: bool = True#

evaluate: bool = False#

seed: int = 1#

no_gpu: bool = False#

max_episode_length: int = 1000#

buffer_size: int = 800000#

time_limit: int = 1000#

cnn_activation_function: str = 'relu'#

dense_activation_function: str = 'elu'#

obs_embed_size: int = 1024#

num_units: int = 400#

deter_size: int = 200#

stoch_size: int = 30#

action_repeat: int = 2#

action_noise: float = 0.3#

total_steps: int = 5000000#

seed_steps: int = 5000#

update_steps: int = 100#

collect_steps: int = 1000#

batch_size: int = 50#

train_seq_len: int = 50#

imagine_horizon: int = 15#

use_disc_model: bool = False#

free_nats: float = 3.0#

discount: float = 0.99#

reward_scale: float = 1.0#

td_lambda: float = 0.95#

kl_loss_coeff: float = 1.0#

kl_alpha: float = 0.8#

disc_loss_coeff: float = 10.0#

num_buckets: int = 255#

symlog_range: float = 10.0#

model_learning_rate: float = 0.0006#

actor_learning_rate: float = 8e-05#

value_learning_rate: float = 8e-05#

adam_epsilon: float = 1e-07#

grad_clip_norm: float = 100.0#

use_amp: bool = False#

perf_defaults: bool = True#

tf32: bool = True#

test: bool = False#

test_interval: int = 10000#

test_episodes: int = 10#

scalar_freq: int = 1000#

log_video_freq: int = -1#

max_videos_to_save: int = 2#

video_format: str = 'gif'#

video_fps: int = 20#

checkpoint_interval: int = 10000#

early_stopping: bool = False#

patience: int = 10#

min_delta: float = 0.0001#

checkpoint_path: str = ''#

restore: bool = False#

experience_replay: str = ''#

render: bool = False#

enable_wandb: bool = False#

wandb_project: str = 'simulusion'#

wandb_entity: str = ''#

log_dir: str = 'runs'#

logdir: Any = None#

data_dir: Any = None#

log_level: str = 'INFO'#

log_file: Any = None#

enable_tensorboard: bool = False#

enable_console_metrics: bool = True#

enable_jsonl: bool = True#

jsonl_filename: str = 'metrics.jsonl'#

log_system_stats_freq: int = 1000#

detect_anomaly: bool = False#

class simulusion.configs.jepa_config.JEPAConfig(use_bfloat16=False, model_name='vit_base', load_checkpoint=False, read_checkpoint=None, copy_data=False, pred_depth=None, pred_emb_dim=384, loss_type='l2', dataset='imagenet', val_split=None, use_gaussian_blur=False, use_horizontal_flip=False, use_color_distortion=False, color_jitter_strength=0.0, batch_size=2048, pin_mem=True, num_workers=8, root_path='/data/imagenet', image_folder='train', crop_size=224, crop_scale=(0.3, 1.0), download=False, allow_overlap=False, patch_size=16, num_enc_masks=1, min_keep=10, enc_mask_scale=(0.85, 1.0), num_pred_masks=4, pred_mask_scale=(0.15, 0.2), aspect_ratio=(0.75, 1.5), ema=(0.996, 1.0), ipe_scale=1.0, weight_decay=0.04, final_weight_decay=0.4, epochs=300, warmup=15, start_lr=0.0001, lr=0.001, final_lr=1e-06, lr_reference_batch_size=2048, early_stopping=False, patience=10, min_delta=0.0001, folder='results/jepa', write_tag='jepa_run', enable_wandb=False, wandb_project='simulusion', wandb_entity='', enable_sweep=False, sweep_config=<factory>)[source]#

Bases: SerializableConfigMixin

Minimal configuration container for JEPA training. Converts to the nested dict expected by train_jepa.main.

Parameters:

use_bfloat16 (bool)
model_name (str)
load_checkpoint (bool)
read_checkpoint (str | None)
copy_data (bool)
pred_depth (int | None)
pred_emb_dim (int)
loss_type (str)
dataset (str)
val_split (float | None)
use_gaussian_blur (bool)
use_horizontal_flip (bool)
use_color_distortion (bool)
color_jitter_strength (float)
batch_size (int)
pin_mem (bool)
num_workers (int)
root_path (str)
image_folder (str)
crop_size (int)
crop_scale (Tuple[float, float])
download (bool)
allow_overlap (bool)
patch_size (int)
num_enc_masks (int)
min_keep (int)
enc_mask_scale (Tuple[float, float])
num_pred_masks (int)
pred_mask_scale (Tuple[float, float])
aspect_ratio (Tuple[float, float])
ema (Tuple[float, float])
ipe_scale (float)
weight_decay (float)
final_weight_decay (float)
epochs (int)
warmup (int)
start_lr (float)
lr (float)
final_lr (float)
lr_reference_batch_size (int | None)
early_stopping (bool)
patience (int)
min_delta (float)
folder (str)
write_tag (str)
enable_wandb (bool)
wandb_project (str)
wandb_entity (str)
enable_sweep (bool)
sweep_config (Dict[str, Any])

use_bfloat16: bool = False#

model_name: str = 'vit_base'#

load_checkpoint: bool = False#

read_checkpoint: str | None = None#

copy_data: bool = False#

pred_depth: int | None = None#

pred_emb_dim: int = 384#

loss_type: str = 'l2'#

dataset: str = 'imagenet'#

val_split: float | None = None#

use_gaussian_blur: bool = False#

use_horizontal_flip: bool = False#

use_color_distortion: bool = False#

color_jitter_strength: float = 0.0#

batch_size: int = 2048#

pin_mem: bool = True#

num_workers: int = 8#

root_path: str = '/data/imagenet'#

image_folder: str = 'train'#

crop_size: int = 224#

crop_scale: Tuple[float, float] = (0.3, 1.0)#

download: bool = False#

allow_overlap: bool = False#

patch_size: int = 16#

num_enc_masks: int = 1#

min_keep: int = 10#

enc_mask_scale: Tuple[float, float] = (0.85, 1.0)#

num_pred_masks: int = 4#

pred_mask_scale: Tuple[float, float] = (0.15, 0.2)#

aspect_ratio: Tuple[float, float] = (0.75, 1.5)#

ema: Tuple[float, float] = (0.996, 1.0)#

ipe_scale: float = 1.0#

weight_decay: float = 0.04#

final_weight_decay: float = 0.4#

epochs: int = 300#

warmup: int = 15#

start_lr: float = 0.0001#

lr: float = 0.001#

final_lr: float = 1e-06#

lr_reference_batch_size: int | None = 2048#

early_stopping: bool = False#

patience: int = 10#

min_delta: float = 0.0001#

folder: str = 'results/jepa'#

write_tag: str = 'jepa_run'#

enable_wandb: bool = False#

wandb_project: str = 'simulusion'#

wandb_entity: str = ''#

enable_sweep: bool = False#

sweep_config: Dict[str, Any]#

to_dict()[source]#

Return type:: Dict[str, Dict[str, Any]]

classmethod from_dict(values)[source]#

Load flat field values or the nested trainer dictionary.

Parameters:: values (Dict[str, Any])
Return type:: JEPAConfig

to_train_dict()[source]#

Return the nested dictionary expected by train_jepa.

Return type:: Dict[str, Dict[str, Any]]

to_nested_dict()[source]#

Backward-compatible alias for the nested JEPA trainer dictionary.

Return type:: Dict[str, Dict[str, Any]]

class simulusion.configs.iris_config.IRISConfig(frame_height=64, frame_width=64, frame_channels=3, vocab_size=512, tokens_per_frame=16, token_embedding_dim=512, encoder_channels=64, encoder_layers=4, encoder_residual_blocks=2, decoder_depth=64, reconstruction_weight=1.0, commitment_weight=1.0, perceptual_weight=1.0, perceptual_blocks=5, perceptual_linear_weights='', quantizer='gradient', transformer_timesteps=20, transformer_embed_dim=256, transformer_layers=10, transformer_heads=4, transformer_dropout=0.1, reward_transform='sign', reward_loss='cross_entropy', imagination_horizon=20, burn_in_length=20, discount=0.995, td_lambda=0.95, entropy_coef=0.001, actor_hidden_size=512, actor_layers=1, total_epochs=600, collection_epochs=500, env_steps_per_epoch=200, training_steps_per_epoch=200, transformer_steps_per_epoch=200, actor_critic_steps_per_epoch=200, model_learning_rate=0.0001, actor_learning_rate=0.0001, value_learning_rate=0.0001, adam_beta1=0.9, adam_beta2=0.999, weight_decay=0.01, grad_clip_norm=10.0, use_amp=True, gradient_checkpointing=True, collect_epsilon=0.01, eval_temperature=0.5, collect_temperature=1.0, start_autoencoder_after=5, start_transformer_after=25, start_actor_critic_after=50, autoencoder_batch_size=256, transformer_batch_size=64, actor_critic_batch_size=64, atari_100k=True, max_env_steps=100000, env_backend='gym', env='ALE/Pong-v5', action_repeat=4, repeat_action_probability=0.0, max_episode_steps=27000, log_interval=1000, eval_episodes=100, checkpoint_interval=50, early_stopping=False, patience=10, min_delta=0.0001)[source]#

Bases: SerializableConfigMixin

Configuration for IRIS (Imagination with auto-Regression over an Inner Speech)

Based on paper: “Transformers are Sample-Efficient World Models” Implements discrete autoencoder + autoregressive Transformer for sample-efficient RL.

Parameters:

frame_height (int)
frame_width (int)
frame_channels (int)
vocab_size (int)
tokens_per_frame (int)
token_embedding_dim (int)
encoder_channels (int)
encoder_layers (int)
encoder_residual_blocks (int)
decoder_depth (int)
reconstruction_weight (float)
commitment_weight (float)
perceptual_weight (float)
perceptual_blocks (int)
perceptual_linear_weights (str)
quantizer (str)
transformer_timesteps (int)
transformer_embed_dim (int)
transformer_layers (int)
transformer_heads (int)
transformer_dropout (float)
reward_transform (str)
reward_loss (str)
imagination_horizon (int)
burn_in_length (int)
discount (float)
td_lambda (float)
entropy_coef (float)
actor_hidden_size (int)
actor_layers (int)
total_epochs (int)
collection_epochs (int)
env_steps_per_epoch (int)
training_steps_per_epoch (int)
transformer_steps_per_epoch (int)
actor_critic_steps_per_epoch (int)
model_learning_rate (float)
actor_learning_rate (float)
value_learning_rate (float)
adam_beta1 (float)
adam_beta2 (float)
weight_decay (float)
grad_clip_norm (float)
use_amp (bool)
gradient_checkpointing (bool)
collect_epsilon (float)
eval_temperature (float)
collect_temperature (float)
start_autoencoder_after (int)
start_transformer_after (int)
start_actor_critic_after (int)
autoencoder_batch_size (int)
transformer_batch_size (int)
actor_critic_batch_size (int)
atari_100k (bool)
max_env_steps (int)
env_backend (str)
env (str)
action_repeat (int)
repeat_action_probability (float)
max_episode_steps (int)
log_interval (int)
eval_episodes (int)
checkpoint_interval (int)
early_stopping (bool)
patience (int)
min_delta (float)

frame_height: int = 64#

frame_width: int = 64#

frame_channels: int = 3#

vocab_size: int = 512#

tokens_per_frame: int = 16#

token_embedding_dim: int = 512#

encoder_channels: int = 64#

encoder_layers: int = 4#

encoder_residual_blocks: int = 2#

decoder_depth: int = 64#

reconstruction_weight: float = 1.0#

commitment_weight: float = 1.0#

perceptual_weight: float = 1.0#

perceptual_blocks: int = 5#

perceptual_linear_weights: str = ''#

quantizer: str = 'gradient'#

transformer_timesteps: int = 20#

transformer_embed_dim: int = 256#

transformer_layers: int = 10#

transformer_heads: int = 4#

transformer_dropout: float = 0.1#

reward_transform: str = 'sign'#

reward_loss: str = 'cross_entropy'#

imagination_horizon: int = 20#

burn_in_length: int = 20#

discount: float = 0.995#

td_lambda: float = 0.95#

entropy_coef: float = 0.001#

actor_hidden_size: int = 512#

actor_layers: int = 1#

total_epochs: int = 600#

collection_epochs: int = 500#

env_steps_per_epoch: int = 200#

training_steps_per_epoch: int = 200#

transformer_steps_per_epoch: int = 200#

actor_critic_steps_per_epoch: int = 200#

model_learning_rate: float = 0.0001#

actor_learning_rate: float = 0.0001#

value_learning_rate: float = 0.0001#

adam_beta1: float = 0.9#

adam_beta2: float = 0.999#

weight_decay: float = 0.01#

grad_clip_norm: float = 10.0#

use_amp: bool = True#

gradient_checkpointing: bool = True#

collect_epsilon: float = 0.01#

eval_temperature: float = 0.5#

collect_temperature: float = 1.0#

start_autoencoder_after: int = 5#

start_transformer_after: int = 25#

start_actor_critic_after: int = 50#

autoencoder_batch_size: int = 256#

transformer_batch_size: int = 64#

actor_critic_batch_size: int = 64#

atari_100k: bool = True#

max_env_steps: int = 100000#

env_backend: str = 'gym'#

env: str = 'ALE/Pong-v5'#

action_repeat: int = 4#

repeat_action_probability: float = 0.0#

max_episode_steps: int = 27000#

log_interval: int = 1000#

eval_episodes: int = 100#

checkpoint_interval: int = 50#

early_stopping: bool = False#

patience: int = 10#

min_delta: float = 0.0001#

get_frame_shape()[source]#

Return type:: tuple

get_autoencoder_config()[source]#

Return type:: dict

get_transformer_config()[source]#

Return type:: dict

get_rl_config()[source]#

Return type:: dict

class simulusion.configs.genie_config.GenieConfig(num_frames=8, image_size=32, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, tokenizer_encoder_dim=256, tokenizer_decoder_dim=512, tokenizer_encoder_depth=4, tokenizer_decoder_depth=8, tokenizer_num_heads=16, action_vocab_size=8, action_embedding_dim=32, action_encoder_dim=256, action_decoder_dim=1024, action_encoder_depth=4, action_num_heads=16, action_pooling='mean', window_attention_heads=1, dynamics_dim=512, dynamics_depth=8, dynamics_num_heads=8, batch_size=4, learning_rate=3e-05, weight_decay=0.0001, warmup_steps=5000, max_steps=125000, mask_prob_min=0.5, mask_prob_max=1.0, sample_temperature=2.0, maskgit_steps=25, use_amp=False, early_stopping=False, patience=10, min_delta=0.0001, val_split=0.1)[source]#

Bases: SerializableConfigMixin

Configuration for Genie model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
tokenizer_encoder_dim (int)
tokenizer_decoder_dim (int)
tokenizer_encoder_depth (int)
tokenizer_decoder_depth (int)
tokenizer_num_heads (int)
action_vocab_size (int)
action_embedding_dim (int)
action_encoder_dim (int)
action_decoder_dim (int)
action_encoder_depth (int)
action_num_heads (int)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
batch_size (int)
learning_rate (float)
weight_decay (float)
warmup_steps (int)
max_steps (int)
mask_prob_min (float)
mask_prob_max (float)
sample_temperature (float)
maskgit_steps (int)
use_amp (bool)
early_stopping (bool)
patience (int)
min_delta (float)
val_split (float)

num_frames: int = 8#

image_size: int = 32#

in_channels: int = 3#

tokenizer_vocab_size: int = 1024#

tokenizer_embedding_dim: int = 32#

tokenizer_encoder_dim: int = 256#

tokenizer_decoder_dim: int = 512#

tokenizer_encoder_depth: int = 4#

tokenizer_decoder_depth: int = 8#

tokenizer_num_heads: int = 16#

action_vocab_size: int = 8#

action_embedding_dim: int = 32#

action_encoder_dim: int = 256#

action_decoder_dim: int = 1024#

action_encoder_depth: int = 4#

action_num_heads: int = 16#

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'#

window_attention_heads: int = 1#

dynamics_dim: int = 512#

dynamics_depth: int = 8#

dynamics_num_heads: int = 8#

batch_size: int = 4#

learning_rate: float = 3e-05#

weight_decay: float = 0.0001#

warmup_steps: int = 5000#

max_steps: int = 125000#

mask_prob_min: float = 0.5#

mask_prob_max: float = 1.0#

sample_temperature: float = 2.0#

maskgit_steps: int = 25#

use_amp: bool = False#

early_stopping: bool = False#

patience: int = 10#

min_delta: float = 0.0001#

val_split: float = 0.1#

class simulusion.configs.genie_config.GenieSmallConfig(num_frames=16, image_size=64, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, tokenizer_encoder_dim=256, tokenizer_decoder_dim=512, tokenizer_encoder_depth=4, tokenizer_decoder_depth=8, tokenizer_num_heads=16, action_vocab_size=8, action_embedding_dim=32, action_encoder_dim=512, action_decoder_dim=1024, action_encoder_depth=8, action_num_heads=16, action_pooling='mean', window_attention_heads=1, dynamics_dim=512, dynamics_depth=8, dynamics_num_heads=8, batch_size=2, learning_rate=0.0001, weight_decay=0.0001, warmup_steps=1000, max_steps=50000, mask_prob_min=0.5, mask_prob_max=1.0, sample_temperature=2.0, maskgit_steps=25, use_amp=False, early_stopping=False, patience=10, min_delta=0.0001, val_split=0.1)[source]#

Bases: SerializableConfigMixin

Small configuration for development/testing.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
tokenizer_encoder_dim (int)
tokenizer_decoder_dim (int)
tokenizer_encoder_depth (int)
tokenizer_decoder_depth (int)
tokenizer_num_heads (int)
action_vocab_size (int)
action_embedding_dim (int)
action_encoder_dim (int)
action_decoder_dim (int)
action_encoder_depth (int)
action_num_heads (int)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
batch_size (int)
learning_rate (float)
weight_decay (float)
warmup_steps (int)
max_steps (int)
mask_prob_min (float)
mask_prob_max (float)
sample_temperature (float)
maskgit_steps (int)
use_amp (bool)
early_stopping (bool)
patience (int)
min_delta (float)
val_split (float)

num_frames: int = 16#

image_size: int = 64#

in_channels: int = 3#

tokenizer_vocab_size: int = 1024#

tokenizer_embedding_dim: int = 32#

tokenizer_encoder_dim: int = 256#

tokenizer_decoder_dim: int = 512#

tokenizer_encoder_depth: int = 4#

tokenizer_decoder_depth: int = 8#

tokenizer_num_heads: int = 16#

action_vocab_size: int = 8#

action_embedding_dim: int = 32#

action_encoder_dim: int = 512#

action_decoder_dim: int = 1024#

action_encoder_depth: int = 8#

action_num_heads: int = 16#

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'#

window_attention_heads: int = 1#

dynamics_dim: int = 512#

dynamics_depth: int = 8#

dynamics_num_heads: int = 8#

batch_size: int = 2#

learning_rate: float = 0.0001#

weight_decay: float = 0.0001#

warmup_steps: int = 1000#

max_steps: int = 50000#

mask_prob_min: float = 0.5#

mask_prob_max: float = 1.0#

sample_temperature: float = 2.0#

maskgit_steps: int = 25#

use_amp: bool = False#

early_stopping: bool = False#

patience: int = 10#

min_delta: float = 0.0001#

val_split: float = 0.1#

class simulusion.configs.genie_config.STTransformerConfig(num_frames=16, num_patches_per_frame=256, dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0)[source]#

Bases: SerializableConfigMixin

Configuration for Spatiotemporal Transformer.

Parameters:

num_frames (int)
num_patches_per_frame (int)
dim (int)
depth (int)
num_heads (int)
mlp_ratio (float)
qkv_bias (bool)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)

num_frames: int = 16#

num_patches_per_frame: int = 256#

dim: int = 768#

depth: int = 12#

num_heads: int = 12#

mlp_ratio: float = 4.0#

qkv_bias: bool = True#

drop_rate: float = 0.0#

attn_drop_rate: float = 0.0#

drop_path_rate: float = 0.0#

class simulusion.configs.genie_config.VideoTokenizerConfig(num_frames=16, image_size=64, in_channels=3, encoder_dim=512, decoder_dim=1024, encoder_depth=12, decoder_depth=20, num_heads=16, patch_size=4, vocab_size=1024, embedding_dim=32, use_ema=False, ema_decay=0.99, commitment_weight=0.25)[source]#

Bases: SerializableConfigMixin

Configuration for Video Tokenizer.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
decoder_dim (int)
encoder_depth (int)
decoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
use_ema (bool)
ema_decay (float)
commitment_weight (float)

num_frames: int = 16#

image_size: int = 64#

in_channels: int = 3#

encoder_dim: int = 512#

decoder_dim: int = 1024#

encoder_depth: int = 12#

decoder_depth: int = 20#

num_heads: int = 16#

patch_size: int = 4#

vocab_size: int = 1024#

embedding_dim: int = 32#

use_ema: bool = False#

ema_decay: float = 0.99#

commitment_weight: float = 0.25#

class simulusion.configs.genie_config.LatentActionModelConfig(num_frames=16, image_size=64, in_channels=3, encoder_dim=1024, encoder_depth=20, num_heads=16, patch_size=16, vocab_size=8, embedding_dim=32, commitment_weight=1.0, action_pooling='mean', window_attention_heads=1)[source]#

Bases: SerializableConfigMixin

Configuration for Latent Action Model.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
encoder_dim (int)
encoder_depth (int)
num_heads (int)
patch_size (int)
vocab_size (int)
embedding_dim (int)
commitment_weight (float)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)

num_frames: int = 16#

image_size: int = 64#

in_channels: int = 3#

encoder_dim: int = 1024#

encoder_depth: int = 20#

num_heads: int = 16#

patch_size: int = 16#

vocab_size: int = 8#

embedding_dim: int = 32#

commitment_weight: float = 1.0#

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'#

window_attention_heads: int = 1#

class simulusion.configs.genie_config.DynamicsModelConfig(num_frames=16, image_size=64, vocab_size=1024, embedding_dim=32, action_vocab_size=8, dim=5120, depth=48, num_heads=40, patch_size=4, mlp_ratio=4.0, qkv_bias=True, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0)[source]#

Bases: SerializableConfigMixin

Configuration for Dynamics Model.

Parameters:

num_frames (int)
image_size (int)
vocab_size (int)
embedding_dim (int)
action_vocab_size (int)
dim (int)
depth (int)
num_heads (int)
patch_size (int)
mlp_ratio (float)
qkv_bias (bool)
drop_rate (float)
attn_drop_rate (float)
drop_path_rate (float)

num_frames: int = 16#

image_size: int = 64#

vocab_size: int = 1024#

embedding_dim: int = 32#

action_vocab_size: int = 8#

dim: int = 5120#

depth: int = 48#

num_heads: int = 40#

patch_size: int = 4#

mlp_ratio: float = 4.0#

qkv_bias: bool = True#

drop_rate: float = 0.0#

attn_drop_rate: float = 0.0#

drop_path_rate: float = 0.0#

class simulusion.configs.dit_config.DiTConfig(DATASET='CIFAR10', BATCH=256, EPOCHS=3, LR=0.0001, WEIGHT_DECAY=0.0, IMG_SIZE=32, CHANNELS=3, PATCH=4, WIDTH=384, DEPTH=12, HEADS=6, DROP=0.0, NUM_CLASSES=0, CLASS_DROPOUT_PROB=0.1, LEARN_SIGMA=True, BETA_START=0.0001, BETA_END=0.02, TIMESTEPS=1000, EMA=True, EMA_DECAY=0.9999, WORKDIR='./dit_demo', ROOT_PATH='./data', CHECKPOINT_EVERY=0, EARLY_STOPPING=False, PATIENCE=10, MIN_DELTA=0.0001, VAL_SPLIT=0.05, CROP_SIZE=None, NUM_WORKERS=4)[source]#

Bases: SerializableConfigMixin

Default configuration values for Diffusion Transformer (DiT) training.

The fields define dataset selection, model architecture, diffusion schedule, optimization hyperparameters, and output paths used by the built-in training entrypoints.

Field names use UPPER_CASE for backward compatibility with the original DiT codebase. Snake-case aliases are accepted via __getattr__ and get_dit_config().

Parameters:

DATASET (str)
BATCH (int)
EPOCHS (int)
LR (float)
WEIGHT_DECAY (float)
IMG_SIZE (int)
CHANNELS (int)
PATCH (int)
WIDTH (int)
DEPTH (int)
HEADS (int)
DROP (float)
NUM_CLASSES (int)
CLASS_DROPOUT_PROB (float)
LEARN_SIGMA (bool)
BETA_START (float)
BETA_END (float)
TIMESTEPS (int)
EMA (bool)
EMA_DECAY (float)
WORKDIR (str)
ROOT_PATH (str)
CHECKPOINT_EVERY (int)
EARLY_STOPPING (bool)
PATIENCE (int)
MIN_DELTA (float)
VAL_SPLIT (float)
CROP_SIZE (Any)
NUM_WORKERS (int)

DATASET: str = 'CIFAR10'#

BATCH: int = 256#

EPOCHS: int = 3#

LR: float = 0.0001#

WEIGHT_DECAY: float = 0.0#

IMG_SIZE: int = 32#

CHANNELS: int = 3#

PATCH: int = 4#

WIDTH: int = 384#

DEPTH: int = 12#

HEADS: int = 6#

DROP: float = 0.0#

NUM_CLASSES: int = 0#

CLASS_DROPOUT_PROB: float = 0.1#

LEARN_SIGMA: bool = True#

BETA_START: float = 0.0001#

BETA_END: float = 0.02#

TIMESTEPS: int = 1000#

EMA: bool = True#

EMA_DECAY: float = 0.9999#

WORKDIR: str = './dit_demo'#

ROOT_PATH: str = './data'#

CHECKPOINT_EVERY: int = 0#

EARLY_STOPPING: bool = False#

PATIENCE: int = 10#

MIN_DELTA: float = 0.0001#

VAL_SPLIT: float = 0.05#

CROP_SIZE: Any = None#

NUM_WORKERS: int = 4#

simulusion.configs.dit_config.canonical_dit_key(name)[source]#

Map a snake_case alias to its UPPER_CASE field name.

DiTConfig keeps the original DiT codebase’s UPPER_CASE field names, so dot-list overrides composed by the training entrypoint have to be translated before they reach the strict config loader. Names that are already canonical (or unknown) are returned unchanged.

Parameters:: name (str)
Return type:: str

simulusion.configs.dit_config.dit_preset_config(name, patch_size, **overrides)[source]#

Build a config for a named Table 1 model, e.g. dit_preset_config("DiT-XL", 2).

Parameters:

name (str) – One of DiT-S, DiT-B, DiT-L, DiT-XL (case-insensitive, and the DiT- prefix is optional).
patch_size (int) – Latent patch size; the paper explores 2, 4 and 8.
**overrides (Any) – Further config fields, UPPER_CASE or snake_case.

Returns:

A DiTConfig for that model. Defaults target latent diffusion of 256x256 ImageNet: a 32x32x4 latent with 1000 classes.

Return type:

DiTConfig

simulusion.configs.dit_config.list_dit_presets()[source]#

Return the Table 1 model names.

Return type:: list[str]

simulusion.configs.dit_config.get_dit_config(**overrides)[source]#

Returns a DiTConfig instance with default values overridden by the provided keyword arguments.

Both UPPER_CASE and snake_case override keys are accepted.

Example usage:: cfg = get_dit_config(BATCH=64, EPOCHS=10, LR=1e-3) cfg = get_dit_config(batch=64, epochs=10, lr=1e-3)

Parameters:: overrides (Any)
Return type:: DiTConfig

simulusion.configs.diamond_config.get_default_device()[source]#

Return type:: str

class simulusion.configs.diamond_config.ModelPreset(diffusion_channels, diffusion_res_blocks, diffusion_cond_dim, reward_channels, reward_lstm_dim, actor_channels, actor_lstm_dim)[source]#

Bases: SerializableConfigMixin

Model architecture preset for different hardware tiers.

Parameters:

diffusion_channels (List[int])
diffusion_res_blocks (int)
diffusion_cond_dim (int)
reward_channels (List[int])
reward_lstm_dim (int)
actor_channels (List[int])
actor_lstm_dim (int)

diffusion_channels: List[int]#

diffusion_res_blocks: int#

diffusion_cond_dim: int#

reward_channels: List[int]#

reward_lstm_dim: int#

actor_channels: List[int]#

actor_lstm_dim: int#

class simulusion.configs.diamond_config.DiamondConfig(preset: str | None = None, game: str = 'Breakout-v5', seed: int = 0, obs_size: int = 64, frameskip: int = 4, max_noop: int = 30, terminate_on_life_loss: bool = True, reward_clip: List[int] = <factory>, max_episode_steps: int = 27000, num_conditioning_frames: int = 4, diffusion_channels: List[int] = <factory>, diffusion_res_blocks: int = 2, diffusion_cond_dim: int = 256, sigma_data: float = 0.5, sigma_min: float = 0.002, sigma_max: float = 80.0, rho: int = 7, p_mean: float = -0.4, p_std: float = 1.2, sampling_method: str = 'euler', num_sampling_steps: int = 3, reward_channels: List[int] = <factory>, reward_res_blocks: int = 2, reward_cond_dim: int = 128, reward_lstm_dim: int = 512, burn_in_length: int = 4, actor_channels: List[int] = <factory>, actor_res_blocks: int = 1, actor_lstm_dim: int = 512, num_epochs: int = 1000, training_steps_per_epoch: int = 400, batch_size: int = 32, environment_steps_per_epoch: int = 100, epsilon_greedy: float = 0.01, data_loader_num_workers: int = 4, pin_memory: bool = True, persistent_workers: bool = True, imagination_horizon: int = 15, discount_factor: float = 0.985, entropy_weight: float = 0.001, lambda_returns: float = 0.95, learning_rate: float = 0.0001, adam_epsilon: float = 1e-08, weight_decay_diffusion: float = 0.01, weight_decay_reward: float = 0.01, weight_decay_actor: float = 0.0, use_amp: bool = True, device: str = <factory>, log_interval: int = 10, eval_interval: int = 50, save_interval: int = 100, checkpoint_dir: str = 'checkpoints/diamond', early_stopping: bool = False, patience: int = 10, min_delta: float = 0.0001)[source]#

Bases: SerializableConfigMixin

Parameters:

preset (str | None)
game (str)
seed (int)
obs_size (int)
frameskip (int)
max_noop (int)
terminate_on_life_loss (bool)
reward_clip (List[int])
max_episode_steps (int)
num_conditioning_frames (int)
diffusion_channels (List[int])
diffusion_res_blocks (int)
diffusion_cond_dim (int)
sigma_data (float)
sigma_min (float)
sigma_max (float)
rho (int)
p_mean (float)
p_std (float)
sampling_method (str)
num_sampling_steps (int)
reward_channels (List[int])
reward_res_blocks (int)
reward_cond_dim (int)
reward_lstm_dim (int)
burn_in_length (int)
actor_channels (List[int])
actor_res_blocks (int)
actor_lstm_dim (int)
num_epochs (int)
training_steps_per_epoch (int)
batch_size (int)
environment_steps_per_epoch (int)
epsilon_greedy (float)
data_loader_num_workers (int)
pin_memory (bool)
persistent_workers (bool)
imagination_horizon (int)
discount_factor (float)
entropy_weight (float)
lambda_returns (float)
learning_rate (float)
adam_epsilon (float)
weight_decay_diffusion (float)
weight_decay_reward (float)
weight_decay_actor (float)
use_amp (bool)
device (str)
log_interval (int)
eval_interval (int)
save_interval (int)
checkpoint_dir (str)
early_stopping (bool)
patience (int)
min_delta (float)

preset: str | None = None#

game: str = 'Breakout-v5'#

seed: int = 0#

obs_size: int = 64#

frameskip: int = 4#

max_noop: int = 30#

terminate_on_life_loss: bool = True#

reward_clip: List[int]#

max_episode_steps: int = 27000#

num_conditioning_frames: int = 4#

diffusion_channels: List[int]#

diffusion_res_blocks: int = 2#

diffusion_cond_dim: int = 256#

sigma_data: float = 0.5#

sigma_min: float = 0.002#

sigma_max: float = 80.0#

rho: int = 7#

p_mean: float = -0.4#

p_std: float = 1.2#

sampling_method: str = 'euler'#

num_sampling_steps: int = 3#

reward_channels: List[int]#

reward_res_blocks: int = 2#

reward_cond_dim: int = 128#

reward_lstm_dim: int = 512#

burn_in_length: int = 4#

actor_channels: List[int]#

actor_res_blocks: int = 1#

actor_lstm_dim: int = 512#

num_epochs: int = 1000#

training_steps_per_epoch: int = 400#

batch_size: int = 32#

environment_steps_per_epoch: int = 100#

epsilon_greedy: float = 0.01#

data_loader_num_workers: int = 4#

pin_memory: bool = True#

persistent_workers: bool = True#

imagination_horizon: int = 15#

discount_factor: float = 0.985#

entropy_weight: float = 0.001#

lambda_returns: float = 0.95#

learning_rate: float = 0.0001#

adam_epsilon: float = 1e-08#

weight_decay_diffusion: float = 0.01#

weight_decay_reward: float = 0.01#

weight_decay_actor: float = 0.0#

use_amp: bool = True#

device: str#

log_interval: int = 10#

eval_interval: int = 50#

save_interval: int = 100#

checkpoint_dir: str = 'checkpoints/diamond'#

early_stopping: bool = False#

patience: int = 10#

min_delta: float = 0.0001#

Training entry points#

Key classes and functions: DiamondAgent, train_diamond, train_dreamer, GenieTrainer, IRISTrainer, and related training entry points.

Training for every Simulusion family.

Trainer runs any TrainingTask (one or more per family, see simulusion.training.tasks) from a TrainerConfig; run() does it from a YAML run config. The per-family trainers below (IRISTrainer, train_diamond, …) are the older entrypoints, kept for existing scripts.

Complete World Model training pipeline for any Gym environment.

This script trains a complete World Model pipeline consisting of: 1. ConvVAE - Encodes observations into latent space 2. MDNRNN - Predicts future latent states given actions 3. Controller - Linear controller trained with CMA-ES

Usage:: python train_world_model.py –env CarRacing-v3 –data_dir ./data –logdir ./results python train_world_model.py –env BipedalWalker-v3 –action_size 4 # if env loading fails

The script will: 1. Generate rollout data (if not already present) 2. Train VAE 3. Train MDNRNN 4. Train Controller

simulusion.training.train_world_model.generate_rollouts(data_dir, env_name, num_rollouts=1000, seq_len=1000, num_workers=8)[source]#

Generate random rollouts from the specified environment.

Parameters:

data_dir (str) – Directory to save rollout files
env_name (str) – Name of the Gym environment
num_rollouts (int) – Total number of rollouts to generate
seq_len (int) – Maximum length per rollout
num_workers (int) – Number of parallel workers

Return type:

None

simulusion.training.train_world_model.env_action_size(env_name)[source]#

The action size the models see in env_name.

The size of a continuous action, or the number of choices of a discrete one (the controller emits a score per choice).

Parameters:: env_name (str)
Return type:: int

simulusion.training.train_world_model.run_training_pipeline(args, action_size)[source]#

Execute the complete World Model training pipeline.

Parameters:

args (Any)
action_size (int)

Return type:

None

simulusion.training.train_world_model.test_trained_model(logdir, env_name, action_size=None, num_episodes=5, *, time_limit=1000, env=None)[source]#

Run the trained V + M + C in the environment; the episode returns.

Reads <logdir>/{vae,mdrnn,ctrl}/best.tar. The models’ sizes come from the checkpoints, so action_size is only checked against them. Returns an empty list (after saying which) when a checkpoint is missing.

Parameters:

logdir (str) – The pipeline’s --logdir.
env_name (str) – The environment to act in.
action_size (int | None) – The action size the models must have, if given.
num_episodes (int) – Episodes to run.
time_limit (int) – Steps per episode at most.
env (Any) – An environment to use instead of making env_name (it is closed afterwards).

Return type:

list[float]

simulusion.training.train_world_model.main(argv=None)[source]#

Parameters:: argv (Sequence[str] | None)
Return type:: None

Training script for Convolutional Variational Autoencoder (ConvVAE).

This module provides functions to train a ConvVAE model on observation data for world model learning.

simulusion.training.train_convvae.train_convae(config)[source]#

Train a Convolutional VAE model.

Runs on the shared Trainer through the vae task, on the rollouts in config.data_dir (.npz files with observations, or any episode folder). As before, everything goes to <logdir>/vae:

checkpoint.tar every epoch and best.tar for the lowest held-out loss (the last 600 rollouts are held out when there are more), resumed from unless noreload is set

ReduceLROnPlateau and early stopping on that loss

samples/sample_epoch_<N>.png every sample_interval epochs unless nosamples

convae_final.pth, the final weights

The trainer adds checkpoints/, run.json, run_config.yaml and train_metrics.jsonl.

Example

>>> config = WMVAEConfig(
...     height=64, width=64, latent_size=32, num_epochs=100, logdir='results'
... )
>>> train_convae(config)

Parameters:: config (WMVAEConfig)
Return type:: None

Training script for Mixture Density Recurrent Neural Network (MDRNN).

This module provides functions to train an MDRNN model for sequence prediction in world models. The MDRNN predicts future latent states using a Gaussian Mixture Model (GMM) based on current latent states and actions.

simulusion.training.train_mdn_rnn.get_loss(mdrnn, latent_obs, action, reward, terminal, latent_next_obs, include_reward, latent_size)[source]#

Compute MDRNN loss.

Computes the combined loss for the MDRNN model: - GMM loss for next latent state prediction - BCE loss for terminal state prediction - MSE loss for reward prediction (if include_reward is True)

Parameters:

mdrnn (MDRNN) – MDRNN model.
latent_obs (Tensor) – Current latent observations.
action (Tensor) – Actions taken.
reward (Tensor) – Rewards received.
terminal (Tensor) – Terminal state flags.
latent_next_obs (Tensor) – Next latent observations (target).
include_reward (bool) – Whether to include reward prediction in loss.
latent_size (int) – Size of latent space.

Returns:

Dictionary containing gmm, bce, mse, and total loss values.

Return type:

dict[str, Tensor]

simulusion.training.train_mdn_rnn.train_mdn_rnn(vae_config, mdrnn_config, use_precomputed_latents=True, use_amp=True)[source]#

Train an MDRNN model.

Runs on the shared Trainer through the mdrnn task, on the rollouts in mdrnn_config.data_dir, encoded by the VAE in <vae_config.logdir>/vae/best.tar. As before, everything goes to <mdrnn_config.logdir>/mdrnn: checkpoint.tar every epoch, best.tar for the lowest held-out loss (resumed from unless noreload), ReduceLROnPlateau and early stopping on that loss, and mdrnn_final.pth.

Parameters:

vae_config (WMVAEConfig) – WMVAEConfig for loading pretrained VAE.
mdrnn_config (WMMDNRNNConfig) – WMMDNRNNConfig with the MDN-RNN’s settings.
use_precomputed_latents (bool) – Accepted for compatibility. Latents are sampled from the VAE as each batch is drawn, so every epoch sees fresh samples instead of one fixed draw stored on disk.
use_amp (bool) – fp16 with loss scaling, on CUDA.

Return type:

None

Example

>>> vae_config = WMVAEConfig(
...     height=64, width=64, latent_size=32, logdir='results'
... )
>>> mdrnn_config = WMMDNRNNConfig(
...     latent_size=32, action_size=3, hidden_size=256,
...     gmm_components=5, logdir='results'
... )
>>> train_mdn_rnn(vae_config, mdrnn_config)

Training a linear controller on latent + recurrent state with CMA-ES.

This module provides functions to train a linear controller using Covariance Matrix Adaptation Evolution Strategy (CMA-ES). The controller maps latent and hidden states to actions for the learned world model.

Reference:: Ha & Schmidhuber (2018). Recurrent World Models Facilitate Policy Evolution. https://arxiv.org/abs/1805.11111

simulusion.training.train_controller.make_controller_env(env_name)[source]#

env_name as the controller sees it: 64x64 image observations.

Environments that take continuous (CarRacing) get the continuous action space the World Models controller outputs. Environments are made with render_mode="rgb_array", so those whose observations are vectors (Pendulum, BipedalWalker) are seen through their rendered frames.

Parameters:: env_name (str)
Return type:: Any

simulusion.training.train_controller.train_controller(config)[source]#

Train a linear controller using CMA-ES.

Runs on the shared Trainer through the controller task, with the VAE and MDN-RNN in <logdir>/vae/best.tar and <logdir>/mdrnn/best.tar. One step is one CMA-ES generation of pop_size candidates, each scored over n_samples episodes of up to time_limit steps in env_name. The best controller so far is kept in <logdir>/ctrl/best.tar (and started from when it exists); training stops when CMA-ES converges or the best return passes target_return.

Candidates are scored one after another in this process, so max_workers and display no longer apply. The controller’s sizes come from the trained models, whatever latent_size / hidden_size / action_size say.

Parameters:: config (WMControllerConfig) – WMControllerConfig containing training hyperparameters, including env_name and action_size.
Return type:: None

simulusion.training.train_jepa.build_loss_fn(loss_type)[source]#

Return the prediction loss used to compare predictor and target tokens.

The I-JEPA paper (Sec. 3) defines the objective as the squared L2 distance between predicted and target patch representations, averaged over the M target blocks. Two reductions of that same objective are offered:

"l2" (default) averages the squared error over patches and channels as well. It has the same minimizer and the same gradient direction as the paper’s formula, but a magnitude that does not grow with block size – which is what the paper’s learning rates were tuned against, since the reference implementation also reduces by mean.
"l2_sum" is the literal formula: summed over patches within a block, averaged over blocks. Its gradients are larger than "l2"’s by roughly the number of patches times the embedding dimension, so lower the learning rate accordingly.

"smooth_l1" reproduces the reference implementation, which uses Smooth-L1 rather than the paper’s L2.

Parameters:: loss_type (str)
Return type:: Any

simulusion.training.train_jepa.make_jepa_loaders(args, *, transform, mask_collator, world_size=1, rank=0, with_val=False)[source]#

I-JEPA’s (train_loader, train_sampler, val_loader) for a nested config.

The loaders shard across world_size processes themselves (a DistributedSampler per rank). val_loader is built only with with_val, from CIFAR-10’s test split or data.val_split of an image folder, with the same mask collator, so its loss is the training objective.

Parameters:

args (Any)
transform (Any)
mask_collator (Any)
world_size (int)
rank (int)
with_val (bool)

Return type:

tuple

simulusion.training.train_jepa.main(args=None, resume_preempt=False)[source]#

Run JEPA training using a CLI argv, nested dict, or JEPAConfig instance.

This entrypoint initializes distributed context, data pipeline, masking, models, optimizers/schedulers, checkpointing, and the full epoch loop.

Parameters:

args (Any)
resume_preempt (bool)

Return type:

Any

simulusion.training.train_jepa.sweep_train()[source]#

Function for WandB sweep agent.

Return type:: None

simulusion.training.train_jepa.main_from_cli(argv=None)[source]#

Compose JEPA config from YAML/dot-list overrides and launch training.

Parameters:: argv (list[str] | None)
Return type:: Any

simulusion.training.train_iris.default_collect_temperature(game, configured)[source]#

Collection sampling temperature, applying the paper’s Freeway exception.

Only overrides a temperature left at the default of 1.0, so an explicit setting in a config file or CLI override still wins.

Parameters:

game (str)
configured (float)

Return type:

float

class simulusion.training.train_iris.IRISTrainer(game='ALE/Pong-v5', device='cuda', seed=42, config=None, env=None)[source]#

Bases: object

Training loop for IRIS on Atari 100k benchmark.

Parameters:

game (str)
device (str)
seed (int)
config (IRISConfig | None)
env (Any | None)

preprocess_frame(frame)[source]#

Preprocess frame: resize to 64x64, return CHW uint8.

Frames are kept as uint8 because that is the dtype of the replay buffer (IRISReplayBuffer.observations). Returning floats in [0, 1] here silently truncated every pixel to 0 on insertion, so the world model and policy trained exclusively on black images. Normalisation to [0, 1] happens at consumption time via to_float_tensor().

Parameters:: frame (ndarray)
Return type:: ndarray

to_float_tensor(obs)[source]#

Convert a uint8 CHW frame (or batch) to a float tensor in [0, 1].

Parameters:: obs (ndarray)
Return type:: Tensor

collect_experience(num_steps, epsilon=0.01)[source]#

Collect exactly num_steps environment steps (paper Algorithm 1).

The step count is exact because the Atari 100k budget is defined in environment steps: collection_epochs * env_steps_per_epoch must equal 100k. Episodes span epochs, so the environment, the policy’s LSTM state and the running episode return all persist between calls and are reset only on a real episode boundary.

Two details follow the paper rather than convenience:

Frames are passed through the discrete autoencoder before reaching the policy (A.1), because the policy is trained purely on reconstructions.
The LSTM state is threaded across steps (A.3); the policy is recurrent and a per-step reset would leave it unable to perceive motion.

Parameters:

num_steps (int) – Exact number of environment steps to collect
epsilon (float) – Random action probability

Returns:

Mean return of episodes that finished during this call, or the most recently completed episode’s return if none finished.

Return type:

float

train_epoch(epoch)[source]#

Train for one epoch.

Parameters:: epoch (int) – Current epoch number
Returns:: Dictionary of metrics
Return type:: dict

get_epsilon(epoch)[source]#

Exploration epsilon for collection.

Paper Table 5 uses a fixed epsilon-greedy parameter of 0.01 combined with sampling from the policy, rather than the decaying schedule most Atari 100k baselines use (Appendix H discusses the tradeoff). A decaying schedule starting near-random makes the early world model fit a data distribution the policy never revisits.

Parameters:: epoch (int)
Return type:: float

evaluate(num_episodes=100, render=False)[source]#

Evaluate agent performance.

Parameters:

num_episodes (int) – Number of evaluation episodes
render (bool) – If True, also return video frames and per-step latent vectors

Returns:

dict with evaluation metrics If render is True: tuple (episode_returns_array, videos_list, latents_array)

Return type:

If render is False (default)

train(total_epochs=None, eval_interval=50, save_dir='checkpoints/iris')[source]#

Train for total_epochs epochs (config.total_epochs by default).

Runs on the shared Trainer through the iris task, with this trainer’s environment, replay and agent. Every eval_interval epochs, once the policy trains, it evaluates (monitoring only); the best of those is written to <save_dir>/best_<game>.pt, and early stopping on it applies when config.early_stopping is set. checkpoint_<epoch>.pt is written every config.checkpoint_interval epochs. The paper’s benchmark number is the final evaluation over config.eval_episodes (self.final_eval_return, model in final_<game>.pt). The trainer adds resumable checkpoints under <save_dir>/checkpoints/, run.json, run_config.yaml and train_metrics.jsonl.

Returns:

every epoch’s metrics, by name.

Return type:

self.metrics

Parameters:

total_epochs (int | None)
eval_interval (int)
save_dir (str)

simulusion.training.train_iris.main(argv=None)[source]#

Run IRIS training with YAML config files and Hydra dot-list overrides.

Parameters:: argv (list[str] | None)
Return type:: IRISConfig

class simulusion.training.train_genie.GenieConfig(num_frames=16, image_size=64, in_channels=3, tokenizer_vocab_size=1024, tokenizer_embedding_dim=32, tokenizer_encoder_dim=512, tokenizer_decoder_dim=1024, tokenizer_encoder_depth=12, tokenizer_decoder_depth=20, action_vocab_size=8, action_embedding_dim=32, action_encoder_dim=1024, action_decoder_dim=1024, action_encoder_depth=20, action_pooling='mean', window_attention_heads=1, dynamics_dim=512, dynamics_depth=8, dynamics_num_heads=8, batch_size=4, learning_rate=3e-05, weight_decay=0.0001, warmup_steps=5000, max_steps=125000, mask_prob_min=0.5, mask_prob_max=1.0, sample_temperature=2.0, maskgit_steps=25, early_stopping=False, patience=10, min_delta=0.0001, val_split=0.1)[source]#

Bases: SerializableConfigMixin

Configuration for Genie training.

Parameters:

num_frames (int)
image_size (int)
in_channels (int)
tokenizer_vocab_size (int)
tokenizer_embedding_dim (int)
tokenizer_encoder_dim (int)
tokenizer_decoder_dim (int)
tokenizer_encoder_depth (int)
tokenizer_decoder_depth (int)
action_vocab_size (int)
action_embedding_dim (int)
action_encoder_dim (int)
action_decoder_dim (int)
action_encoder_depth (int)
action_pooling (Literal['mean', 'windowed_attention'])
window_attention_heads (int)
dynamics_dim (int)
dynamics_depth (int)
dynamics_num_heads (int)
batch_size (int)
learning_rate (float)
weight_decay (float)
warmup_steps (int)
max_steps (int)
mask_prob_min (float)
mask_prob_max (float)
sample_temperature (float)
maskgit_steps (int)
early_stopping (bool)
patience (int)
min_delta (float)
val_split (float)

num_frames: int = 16#

image_size: int = 64#

in_channels: int = 3#

tokenizer_vocab_size: int = 1024#

tokenizer_embedding_dim: int = 32#

tokenizer_encoder_dim: int = 512#

tokenizer_decoder_dim: int = 1024#

tokenizer_encoder_depth: int = 12#

tokenizer_decoder_depth: int = 20#

action_vocab_size: int = 8#

action_embedding_dim: int = 32#

action_encoder_dim: int = 1024#

action_decoder_dim: int = 1024#

action_encoder_depth: int = 20#

action_pooling: Literal['mean', 'windowed_attention'] = 'mean'#

window_attention_heads: int = 1#

dynamics_dim: int = 512#

dynamics_depth: int = 8#

dynamics_num_heads: int = 8#

batch_size: int = 4#

learning_rate: float = 3e-05#

weight_decay: float = 0.0001#

warmup_steps: int = 5000#

max_steps: int = 125000#

mask_prob_min: float = 0.5#

mask_prob_max: float = 1.0#

sample_temperature: float = 2.0#

maskgit_steps: int = 25#

early_stopping: bool = False#

patience: int = 10#

min_delta: float = 0.0001#

val_split: float = 0.1#

class simulusion.training.train_genie.VideoDataset(video_paths, num_frames=16, image_size=64)[source]#

Bases: Dataset

Video clips for Genie training, returned as (C, T, H, W) float tensors.

Each entry in video_paths may be:

a .npy / .npz array of shape (T, H, W, C) or (C, T, H, W)
a .pt / .pth tensor of the same layouts
a video file (.mp4, .avi, .mkv, .webm, .mov) loaded with OpenCV when opencv-python is installed (the viz extra)

Frames are uniformly sampled to num_frames and resized to image_size.

Parameters:

video_paths (list)
num_frames (int)
image_size (int)

class simulusion.training.train_genie.GenieTrainer(model, config, device=None, *, precision=None)[source]#

Bases: object

Trainer for Genie model.

Parameters:

model (Module)
config (GenieConfig)
device (device | None)
precision (MixedPrecision | None)

train_step(batch)[source]#

Single training step.

Parameters:: batch (Tensor) – (B, C, T, H, W) video batch
Returns:: Dictionary of losses
Return type:: Dict[str, Tensor | float | None]

validate(val_batch)[source]#

Validation step.

Parameters:: val_batch (Tensor) – (B, C, T, H, W) validation video batch
Returns:: Dictionary of validation metrics
Return type:: Dict[str, Tensor]

validate_epoch(val_dataloader)[source]#

Mean held-out reconstruction loss over the whole validation loader.

validate scores a single batch, which is far too noisy to drive a plateau test – the batch-to-batch spread swamps the epoch-to-epoch trend, so early stopping on it would fire on noise. Averaging over the loader gives a comparable number per validation.

Parameters:: val_dataloader (DataLoader)
Return type:: float

train(train_dataloader, val_dataloader=None, num_steps=None, log_interval=100, val_interval=1000, checkpoint_dir=None, checkpoint_interval=0)[source]#

Train until global_step reaches num_steps (config.max_steps).

Runs on the shared Trainer through the genie task, with this trainer’s model, optimizer and schedule.

Parameters:

train_dataloader (DataLoader) – Training data loader of (B, C, T, H, W) clips.
val_dataloader (DataLoader | None) – Held-out loader, scored over its whole length every val_interval steps (required by early_stopping).
num_steps (int | None) – Number of training steps (uses config.max_steps if None)
log_interval (int) – Logging frequency
val_interval (int) – Validation frequency
checkpoint_dir (str | None) – Where genie_step_<N>.pt checkpoints go, and the run directory for the trainer’s own files (checkpoints/, run.json, train_metrics.jsonl). Without it the run directory is runs/genie_<time> and no genie_step files are written.
checkpoint_interval (int) – Steps between checkpoints. 0 writes one only at the end (with checkpoint_dir).

Return type:

None

save_checkpoint(path)[source]#

Save model checkpoint.

Parameters:: path (str)
Return type:: None

load_checkpoint(path)[source]#

Load model checkpoint.

Parameters:: path (str)
Return type:: None

simulusion.training.train_genie.build_genie(config=None)[source]#

The Genie model a training config describes.

Parameters:: config (GenieConfig | None)
Return type:: Genie

simulusion.training.train_genie.create_genie_trainer(config=None, device=None)[source]#

Factory function to create Genie trainer and model.

Parameters:

config (GenieConfig | None)
device (device | None)

Return type:

Tuple[GenieTrainer, Module]

simulusion.training.train_genie.main(argv=None)[source]#

Console entrypoint for Genie trainer setup.

VideoDataset loads .npy / .pt clips, or video files when OpenCV is installed. For the TinyWorlds HDF5 path, use scripts/train_genie_tinyworlds.py.

Parameters:: argv (list[str] | None)
Return type:: None

simulusion.training.train_planet.train(memory, rssm, optimizer, device, N=32, H=50, beta=1.0, grads=False, update=None, precision=None)[source]#

Training implementation as indicated in: Learning Latent Dynamics for Planning from Pixels arXiv:1811.04551

(a.) The Standard Variational Bound Method: using only single step predictions.

update steps optimizer (the shared Trainer passes one at the run’s precision); without it the step is plain fp32 with the gradient norm clipped at 1000, as it always was. precision sets the autocast for the loss.

Parameters:

memory (Any)
rssm (Any)
optimizer (Any)
device (device)
N (int)
H (int)
beta (float)
grads (bool)
update (GradientUpdate | None)
precision (MixedPrecision | None)

Return type:

dict

simulusion.training.train_planet.build_parser()[source]#

CLI for the PlaNet trainer.

Every default below is the value this script previously hard-coded, so a bare python -m simulusion.training.train_planet behaves as before. They are flags because nothing else could reach them: sweeps had no way to shorten a run or redirect its output.

Return type:: ArgumentParser

simulusion.training.train_planet.main(argv=None)[source]#

PlaNet/RSSM training from the command line.

Builds a Planet on --env and runs Planet.train() (so the shared Trainer): --iters gradient steps per epoch, one exploration and one evaluation episode per epoch, videos and ckpt_<epoch>.pth checkpoints in --outdir. The optimizer is the one this script always used: Adam at 1e-3, eps 1e-4, no weight decay, no learning-rate schedule.

Parameters:: argv (list[str] | None)
Return type:: None

simulusion.training.train_rssm.train_rssm(memory, model, optimizer, record_grads=True, *, iterations=10, batch_size=32)[source]#

Train an RSSM on replayed trajectories for one optimization phase.

Samples iterations batches of batch_size sequences from memory, computes reconstruction and KL objectives across rollout steps, and returns the per-iteration losses.

Parameters:

memory (Any)
model (Any)
optimizer (Any)
record_grads (bool)
iterations (int)
batch_size (int)

Return type:

dict

simulusion.training.train_rssm.evaluate(memory, model, path, eps)[source]#

Run one RSSM reconstruction/prediction evaluation and save visual outputs.

Decodes priors/posteriors for a sampled sequence and writes frame grids for qualitative inspection.

Parameters:

memory (Any)
model (Any)
path (str)
eps (Any)

Return type:

None

simulusion.training.train_rssm.main(argv=None)[source]#

Standalone training loop for RSSM with generated replay fallback support.

Initializes environment/policy/memory, trains over episodes, logs metrics, and periodically evaluates and checkpoints the model.

Parameters:: argv (Sequence[str] | None)
Return type:: None

class simulusion.training.train_diamond.DiamondAgent(config, *, env=None)[source]#

Bases: object

DIAMOND: DIffusion As a Model Of eNvironment Dreams

RL agent trained entirely within a diffusion world model.

Parameters:

config (DiamondConfig)
env (Any | None)

classmethod from_config(config=None, **overrides)[source]#

Build a DIAMOND agent from a config object, dict, YAML file, or YAML string.

Parameters:

config (DiamondConfig | dict | str | Path | None)
overrides (Any)

Return type:

DiamondAgent

classmethod from_pretrained(pretrained_model_name_or_path, *, config=None, checkpoint_filename=None, config_filename='config.yaml', repo_type=None, revision=None, **overrides)[source]#

Load a DIAMOND checkpoint from a local path/directory or HF Hub.

Parameters:

pretrained_model_name_or_path (str | Path)
config (DiamondConfig | dict | str | Path | None)
checkpoint_filename (str | None)
config_filename (str)
repo_type (str | None)
revision (str | None)
overrides (Any)

Return type:

DiamondAgent

parameter_count(trainable_only=False)[source]#

Parameters:: trainable_only (bool)
Return type:: int

summary()[source]#

Return type:: dict

set_precision(precision)[source]#

Train at precision from now on (fresh loss scalers).

Parameters:: precision (MixedPrecision)
Return type:: None

train()[source]#

Train for config.num_epochs epochs of Algorithm 1.

Runs on the shared Trainer through the diamond task, with this agent’s environment, replay and networks. Evaluation runs every eval_interval epochs, early stopping on its reward when early_stopping is set, and checkpoint_<epoch>.pt is written to checkpoint_dir every save_interval epochs and at the end. The trainer adds resumable checkpoints under <checkpoint_dir>/checkpoints/, run.json, run_config.yaml and train_metrics.jsonl.

Return type:: None

train_epoch()[source]#

One epoch of Algorithm 1: collect, then train every component.

Collects environment_steps_per_epoch steps, then runs training_steps_per_epoch updates of the diffusion model, the reward/termination model and the actor-critic. Returns the epoch’s mean losses, or None while the replay holds too little to train.

Return type:: Dict[str, float] | None

evaluate(num_episodes=1)[source]#

Evaluate the agent.

Parameters:: num_episodes (int)
Return type:: float

save_checkpoint(path=None)[source]#

Save model checkpoint.

Parameters:: path (str | PathLike | None) – Optional path where to write the checkpoint. If path is None or a bare filename, the file is written into config.checkpoint_dir/<filename>. If path contains a directory component or is an absolute/relative path, it is used directly. When path is None the checkpoint is written to config.checkpoint_dir/checkpoint.pt.
Return type:: None

load_checkpoint(path=None)[source]#

Load model checkpoint.

Parameters:: path (str | None) – Optional path to checkpoint. If None, config.checkpoint_dir/checkpoint.pt is loaded. If a bare filename is provided, we try config.checkpoint_dir/<filename>; if a path with directory components is provided we use it directly.
Return type:: None

simulusion.training.train_diamond.train_diamond(game=None, seed=None, preset=None, device=None, config=None)[source]#

Train DIAMOND on a specific game or a composed experiment config.

Parameters:

game (str | None)
seed (int | None)
preset (str | None)
device (str | None)
config (DiamondConfig | None)

Return type:

None

simulusion.training.train_diamond.main(argv=None)[source]#

Compose DIAMOND config from YAML/dot-list overrides and launch training.

Parameters:: argv (list[str] | None)
Return type:: Any

class simulusion.training.rl_harness.ActorCritic(obs_shape, action_dim, hidden_dim=256)[source]#

Bases: Module

Simple actor-critic network for RL harness.

Parameters:

obs_shape (tuple)
action_dim (int)
hidden_dim (int)

forward(obs)[source]#

Forward pass through CNN, then actor and critic heads.

Parameters:: obs (Tensor)
Return type:: tuple[Tensor, Tensor]

get_action(obs)[source]#

Sample action from policy.

Parameters:: obs (Tensor)
Return type:: tuple[Tensor, Tensor, Tensor]

class simulusion.training.rl_harness.PPOTrainer(vec_env, device='cpu', lr=0.0003, gamma=0.99, gae_lambda=0.95, clip_ratio=0.2, num_epochs=10, batch_size=64, max_grad_norm=0.5, entropy_coeff=0.01, value_coeff=0.5)[source]#

Bases: object

Simple PPO trainer for testing vectorized environments.

Parameters:

vec_env (TorchVectorizedEnv)
device (str)
lr (float)
gamma (float)
gae_lambda (float)
clip_ratio (float)
num_epochs (int)
batch_size (int)
max_grad_norm (float)
entropy_coeff (float)
value_coeff (float)

collect_trajectories(num_steps)[source]#

Collect trajectories using the vectorized environment.

Parameters:: num_steps (int)
Return type:: Dict[str, Tensor]

compute_gae(rewards, values, dones)[source]#

Compute Generalized Advantage Estimation.

Parameters:

rewards (Tensor)
values (Tensor)
dones (Tensor)

Return type:

Tensor

train_step(trajectories)[source]#

Perform one training step using PPO.

Parameters:: trajectories (Dict[str, Tensor])
Return type:: None

train(total_timesteps, log_interval=1000)[source]#

Main training loop.

Parameters:

total_timesteps (int)
log_interval (int)

Return type:

None

simulusion.training.rl_harness.create_rl_harness_example()[source]#

Example function to create and run the RL harness. Usage: Call this with your environment factory.

Return type:: None

Memory and controllers#

class simulusion.memory.dreamer_memory.ReplayBuffer(size, obs_shape, action_size, seq_len, batch_size)[source]#

Bases: object

Fixed-size replay buffer for Dreamer with image observations and transitions.

Stores (observation, action, reward, terminal) tuples in a ring buffer and supports sampling contiguous sequences for world-model training.

Key Features

Ring buffer with fixed capacity (FIFO eviction when full)
Stores raw uint8 images to save memory
Samples sequences (not single transitions) for temporal modeling
Validates sampled sequences don’t span episode boundaries

Memory Layout

observations: (capacity, C, H, W) uint8 images
actions: (capacity, action_dim) float32
rewards: (capacity,) float32
terminals: (capacity,) float32 (1.0 = terminal, 0.0 = continue)

Sampling Process

Random start index (avoiding episode boundaries)
Collect sequence of length seq_len with wraparound
Validate no terminal in middle of sequence
Return batch of sequences

Usage with Dreamer:

buffer = ReplayBuffer(
    size=100000,           # Max transitions to store
    obs_shape=(3, 64, 64), # RGB images
    action_size=6,         # Continuous action dim
    seq_len=50,            # Sequence length for training
    batch_size=50          # Parallel sequences per batch
)

# Add transitions during interaction
buffer.add(obs, action, reward, done)

# Sample batch for world model training
obs_batch, action_batch, reward_batch, term_batch = buffer.sample()

Memory Efficiency

Uses uint8 for images (1 byte per pixel vs 4 for float32)
Sequences share observations (overlapping windows)
Configurable capacity based on available system memory

Note

The buffer stores observations as {“image”: …} dicts but returns just the image arrays for training efficiency.

Parameters:

size (int)
obs_shape (Tuple[int, ...])
action_size (int)
seq_len (int)
batch_size (int)

add(obs, ac, rew, done, terminated=None)[source]#

Add a transition to the buffer.

Parameters:

obs (dict) – Observation dict with ‘image’ key containing the observation
ac (ndarray) – Action taken, shape (action_size,)
rew (float) – Reward received, scalar
done (float) – Episode-boundary flag, 1.0 if the episode ended for any reason (termination or truncation), 0.0 otherwise
terminated (float | None) – 1.0 only if the environment genuinely terminated. Defaults to done when the caller cannot tell the two apart.

Return type:

None

sample(include_terminated=False)[source]#

Sample a batch of sequences for training.

Parameters:

include_terminated (bool) – Also return the true-termination flags as a fifth element.

Returns:

(observations, actions, rewards, terminals[, terminated])

observations: (seq_len, batch, C, H, W)
actions: (seq_len, batch, action_dim)
rewards: (seq_len, batch)
terminals: (seq_len, batch) episode boundaries
terminated: (seq_len, batch) true terminations only

Return type:

tuple

class simulusion.memory.dreamer_memory.Memory(capacity=10000)[source]#

Bases: object

Simple deque-based memory for storing transitions.

Used by PlaNet for online planning. Stores recent transitions and provides random sampling for policy updates.

Parameters:: capacity (int) – Maximum number of transitions to store

Usage:

memory = Memory(capacity=10000)
memory.append(obs, action, reward, done, info)
batch = random.sample(memory, batch_size=32)

append(*args)[source]#

Append a transition to memory.

Parameters:: *args (Any) – Variable length argument list containing transition data. Typically (observation, action, reward, done, info).
Return type:: None

sample(batch_size)[source]#

Sample random batch of transitions from memory.

Parameters:: batch_size (int) – Number of transitions to sample.
Returns:: List of sampled transitions.
Return type:: list

class simulusion.memory.dreamer_memory.Episode(observation, action=None, reward=None, terminal=None, info=None)[source]#

Bases: object

Stores a single episode for PlaNet’s imagination and planning.

An episode is a sequence of (observation, action, reward) tuples collected during environment interaction. Episodes are used for computing returns and training value functions.

Parameters:

obs – Initial observation
action (Any) – First action (optional)
reward (Any) – Initial reward (optional)
info (Any) – Additional info dict (optional)
observation (Any)
terminal (Any)

Usage:

episode = Episode(obs, info=info)
episode.append(action, obs, reward, done, info)
episodes = [episode for _ in range(num_episodes)]

# Use with Planet agent for planning
imag_state, imag_reward, imag_action = planet.imagine(episodes)

append(action, observation, reward, terminal, info=None)[source]#

Parameters:

action (Any)
observation (Any)
reward (Any)
terminal (Any)
info (Any)

Return type:

None

class simulusion.memory.planet_memory.Episode(postprocess_fn=None)[source]#

Bases: object

Records the agent’s interaction with the environment for a single episode.

Stores observations, actions, rewards, and terminal flags during a single trajectory. At termination, converts all lists to numpy arrays for efficient batch processing.

Variables:

x (list or np.ndarray) – Observations collected during the episode.
u (list or np.ndarray) – Actions taken.
r (list or np.ndarray) – Rewards received.
t (list or np.ndarray) – Terminal flags (0.0 = continue, 1.0 = terminal).
info (dict) – Additional episode metadata.

Parameters:

postprocess_fn (callable, optional) – Function to apply to observations before storing (e.g., normalization). Default: identity function.

Example:

episode = Episode()
episode.append(obs, action, reward, False)
episode.append(obs, action, reward, True)
episode.terminate(final_obs)
print(episode.x.shape)  # Now a numpy array

property size: int#

append(obs, act, reward, terminal)[source]#

Parameters:

obs (Any)
act (Any)
reward (Any)
terminal (Any)

Return type:

None

terminate(obs)[source]#

Parameters:: obs (Any)
Return type:: None

class simulusion.memory.planet_memory.Memory(size=None)[source]#

Bases: deque

Episode-based replay memory for PlaNet/RSSM training.

Stores episodes as variable-length trajectories and supports sampling sub-sequences for training. Implements a ring-buffer style eviction when capacity is reached.

Stores complete episodes as lists of transitions
Samples contiguous sub-sequences for sequence models
Supports time-major formatting (time-first) for RNN input
Memory usage estimation to prevent OOM errors

Parameters:

size (int, optional) – Maximum number of episodes to store. If None, deque grows without limit (useful for unpickling).

Variables:

episodes (deque) – Collection of Episode objects.
eps_lengths (deque) – Length of each episode.
size (property) – Total number of transitions across all episodes.

Example:

memory = Memory(size=100)
memory.append([episode1, episode2])
batch, lengths = memory.sample(batch_size=32, tracelen=50)

property size: int#

append(episodes)[source]#

Parameters:: episodes (list[Episode])
Return type:: None

sample(batch_size, tracelen=1, time_first=False)[source]#

Sample random sub-sequences from stored episodes.

Randomly selects episodes and starting positions to create batches of contiguous sequences for training sequence models.

Parameters:

batch_size (int) – Number of sequences to sample.
tracelen (int) – Length of each sequence (default: 1).
time_first (bool) – If True, returns tensors with time dimension first (T, B, …) instead of batch first (B, T, …).

Returns:

(observations, actions, rewards, terminals, lengths)

observations: (batch, tracelen+1, *obs_shape) or (tracelen+1, batch, …)
actions: (batch, tracelen, action_dim) or (tracelen, batch, …)
rewards: (batch, tracelen) or (tracelen, batch)
terminals: (batch, tracelen) or (tracelen, batch)
lengths: (batch,) original episode lengths for each sample

Return type:

tuple

Raises:

ValueError – If memory is empty or no episodes meet minimum length.
MemoryError – If estimated memory usage exceeds 200 MiB threshold.

class simulusion.memory.iris_memory.IRISReplayBuffer(size, obs_shape, action_size, seq_len=20, batch_size=64)[source]#

Bases: object

Replay buffer for IRIS (Imagination with auto-Regression over an Inner Speech) training.

Stores (observation, action, reward, terminal) tuples in a ring buffer and supports sampling contiguous sequences for world model training.

Features:

Ring buffer with fixed capacity (FIFO eviction when full)
Stores uint8 images for memory efficiency
Samples sequences with validation to avoid episode boundaries
Supports sequence sampling for temporal learning

Memory Layout:

observations: (capacity, C, H, W) uint8
actions: (capacity, action_size) float32
rewards: (capacity,) float32
terminals: (capacity,) float32

Parameters:

size (int) – Maximum number of transitions to store.
obs_shape (tuple) – Shape of observations as (C, H, W).
action_size (int) – Dimension of actions.
seq_len (int) – Length of sequences to sample (default: 20).
batch_size (int) – Number of sequences per batch (default: 64).

Variables:

size (int) – Buffer capacity.
obs_shape (tuple) – Observation shape.
action_size (int) – Action dimension.
seq_len (int) – Sequence length.
batch_size (int) – Batch size.
steps (int) – Total transitions added.
episodes (int) – Number of episode terminations observed.

add(obs, action, reward, terminal)[source]#

Add a transition to the buffer.

Parameters:

obs (ndarray) – Observation array with shape (C, H, W).
action (ndarray) – Action array with shape (action_size,).
reward (float) – Scalar reward value.
terminal (bool) – Boolean indicating if episode terminated.

Return type:

None

sample_sequence(seq_len=None)[source]#

Sample a batch of sequences for world model training.

Returns:: (batch_size, seq_len+1, C, H, W) actions: (batch_size, seq_len, action_size) rewards: (batch_size, seq_len) terminals: (batch_size, seq_len)
Return type:: observations
Parameters:: seq_len (int | None)

sample_with_burn_in(batch_size, burn_in)[source]#

Sample start frames together with the frames that precede them.

Episode boundaries are respected: any burn-in frame at or before a terminal is replaced by a repeat of the oldest valid frame, so context never bleeds across episodes.

Parameters:

batch_size (int) – Number of start frames to draw.
burn_in (int) – Number of preceding frames to return per start frame.

Returns:

(batch_size, C, H, W) uint8 frames to imagine from. burn_in_obs: (batch_size, burn_in, C, H, W) uint8 preceding frames.

Empty along axis 1 when burn_in is 0.

Return type:

start_obs

sample_single()[source]#

Sample a single transition for online updates.

Return type:: Tuple[ndarray, ndarray, float, float]

property buffer_capacity: int#: Returns the total capacity of the buffer.

class simulusion.memory.iris_memory.IRISOnPolicyBuffer(max_steps=1000)[source]#

Bases: object

On-policy buffer for collecting trajectories during environment interaction.

Useful for:

Collecting complete episode trajectories
Storing data before batch processing
Temporary storage during environment interaction

Parameters:

max_steps (int) – Maximum number of steps to store (default: 1000).

Variables:

max_steps (int) – Maximum buffer capacity.
observations (list) – List of observations.
actions (list) – List of actions.
rewards (list) – List of rewards.
terminals (list) – List of terminal flags.

add(obs, action, reward, terminal)[source]#

Parameters:

obs (ndarray)
action (ndarray)
reward (float)
terminal (bool)

Return type:

None

clear()[source]#

Return type:: None

get_arrays()[source]#

Return type:: tuple[ndarray, ndarray, ndarray, ndarray]

RSSM-based policy for model-predictive control.

This module provides the RSSMPolicy class that implements model-predictive control using the RSSM (Recurrent State Space Model) latent dynamics model. The policy uses a Cross-Entropy Method (CEM) for planning actions in latent space.

Reference:: Ha & Schmidhuber (2018). Recurrent World Models Facilitate Policy Evolution. https://arxiv.org/abs/1805.11111

class simulusion.controller.rssm_policy.RSSMPolicy(model, planning_horizon, num_candidates, num_iterations, top_candidates, device, compile_rollout=False)[source]#

Bases: object

Model-predictive controller using Cross-Entropy Method (CEM) with RSSM.

Plans actions by optimizing a sequence of future actions in the RSSM’s latent space. Uses Cross-Entropy Method to refine action sequences based on predicted returns.

Algorithm:

Initialize Gaussian distribution over action sequences
Sample N candidate action sequences
Rollout each sequence in RSSM latent space
Score by predicted cumulative rewards
Keep top K candidates, fit Gaussian to them
Repeat for T iterations
Execute first action from best sequence

Variables:

rssm – The RSSM world model.
N – Number of candidate action sequences to sample.
K – Number of top candidates to use for updating the proposal.
T – Number of CEM iterations per planning step.
H – Planning horizon (number of future steps to consider).
d – Action dimensionality.
device – Device to run computations on.
state_size – Hidden state dimensionality.
latent_size – Latent state dimensionality.

Parameters:

model (Any)
planning_horizon (int)
num_candidates (int)
num_iterations (int)
top_candidates (int)
device (device | str)
compile_rollout (bool)

Example

>>> policy = RSSMPolicy(
...     model=rssm,
...     planning_horizon=12,
...     num_candidates=1000,
...     num_iterations=5,
...     top_candidates=100,
...     device='cuda'
... )
>>> policy.reset()
>>> action = policy.poll(observation)

reset()[source]#

Reset the policy state.

Initializes the hidden state, latent state, and action to zeros. Should be called at the beginning of each episode.

Return type:: None

poll(observation, explore=False)[source]#

Get action for given observation.

Parameters:

observation (Tensor) – Current observation tensor of shape (channels, height, width).
explore (bool) – If True, add exploration noise to the selected action.

Returns:

Action tensor of shape (1, action_size).

Return type:

Tensor

class simulusion.controller.iris_policy.IRISActor(action_size, hidden_size=512, num_layers=4, frame_shape=(3, 64, 64))[source]#

Bases: ExportableAgentMixin, Module

Actor network for the IRIS (Imagination with auto-Regression over an Inner Speech) policy.

Architecture:

CNN: Extracts features from input frames (3x64x64 -> 512)
LSTM: Processes temporal sequences with configurable layers
Linear: Maps hidden states to action logits

Parameters:

action_size (int) – Number of discrete actions.
hidden_size (int) – LSTM hidden state size (default: 512).
num_layers (int) – Number of LSTM layers (default: 4).
frame_shape (tuple) – Shape of input frames as (C, H, W) (default: (3, 64, 64)).

Variables:

action_size (int) – Number of discrete actions.
hidden_size (int) – LSTM hidden state size.
num_layers (int) – Number of LSTM layers.
frame_shape (tuple) – Input frame shape.

forward(frames, hidden_state=None, burn_in_frames=None)[source]#

Forward pass through actor.

Parameters:

frames (Tensor) – Input frames (B, T, C, H, W) or (B, C, H, W)
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) tuple for LSTM state
burn_in_frames (Tensor | None) – Frames to use for initializing hidden state

Returns:

Action logits (B, T, action_size) or (B, action_size) hidden_state: Updated (h, c) tuple

Return type:

action_logits

init_hidden_state(batch_size, device)[source]#

Initialize LSTM hidden state.

Parameters:

batch_size (int)
device (device)

Return type:

Tuple[Tensor, Tensor]

get_action(frame, temperature=1.0, deterministic=False)[source]#

Get action from a single frame.

Parameters:

frame (Tensor) – Single frame (B, C, H, W)
temperature (float) – Softmax temperature (higher = more random)
deterministic (bool) – If True, return argmax; else sample

Returns:

Selected action indices (B,)

Return type:

class simulusion.controller.iris_policy.IRISCritic(hidden_size=512, num_layers=4, frame_shape=(3, 64, 64))[source]#

Bases: ExportableAgentMixin, Module

Critic network for IRIS value estimation.

Architecture:

CNN: Feature extractor with the same architecture as the actor (3x64x64 -> 512)
LSTM: Temporal processing with same architecture as actor
Linear: Maps hidden states to scalar values

Parameters:

hidden_size (int) – LSTM hidden state size (default: 512).
num_layers (int) – Number of LSTM layers (default: 4).
frame_shape (tuple) – Shape of input frames as (C, H, W) (default: (3, 64, 64)).

Variables:

hidden_size (int) – LSTM hidden state size.
num_layers (int) – Number of LSTM layers.
frame_shape (tuple) – Input frame shape.

Returns:

Value estimates with shape (B, T). hidden_state: Updated LSTM hidden state (h, c) tuple.

Return type:

values

forward(frames, hidden_state=None)[source]#

Forward pass through critic.

Parameters:

frames (Tensor) – Input frames (B, T, C, H, W)
hidden_state (Tuple[Tensor, Tensor] | None) – Optional (h, c) tuple

Returns:

Value estimates (B, T) hidden_state: Updated (h, c) tuple

Return type:

values

init_hidden_state(batch_size, device)[source]#

Initialize LSTM hidden state.

Parameters:

batch_size (int)
device (device)

Return type:

Tuple[Tensor, Tensor]

class simulusion.controller.iris_policy.CNNFeatureExtractor(frame_shape=(3, 64, 64), output_size=512, channels=(32, 64, 128, 256))[source]#

Bases: ExportableAgentMixin, Module

CNN feature extractor shared between actor and critic networks.

Reproduces the convolutional block of the IRIS actor-critic (paper A.3):

“The convolutional block consists of the same layer repeated four times: a 3x3 convolution with stride 1 and padding 1, a ReLU activation, and 2x2 max-pooling with stride 2.”

The paper does not state the channel widths; 32 -> 64 -> 128 -> 256 is kept from the previous implementation.

Parameters:

frame_shape (tuple) – Shape of input frames as (C, H, W) (default: (3, 64, 64)).
output_size (int) – Size of output feature vector (default: 512).
channels (tuple) – Per-layer output channel counts.

Variables:

frame_shape (tuple) – Input frame shape.
output_size (int) – Output feature dimension.

Returns:

Feature vectors with shape (B, output_size).

Return type:

features

forward(x)[source]#

Extract features from frames.

Parameters:: x (Tensor) – Frames (B, C, H, W)
Returns:: Feature vectors (B, output_size)
Return type:: features

class simulusion.controller.iris_policy.IRISPolicy(action_size, hidden_size=512, num_layers=4, frame_shape=(3, 64, 64))[source]#

Bases: ExportableAgentMixin, Module

Combined policy module for IRIS (Imagination with auto-Regression over an Inner Speech).

Parameters:

action_size (int) – Number of discrete actions.
hidden_size (int) – LSTM hidden state size (default: 512).
num_layers (int) – Number of LSTM layers (default: 4).
frame_shape (tuple) – Shape of input frames as (C, H, W) (default: (3, 64, 64)).

Variables:

actor (IRISActor) – The actor network for action selection.
hidden_size (int) – LSTM hidden state size.
num_layers (int) – Number of LSTM layers.
frame_shape (tuple) – Input frame shape.

Example

>>> policy = IRISPolicy(
...     action_size=18,
...     hidden_size=512,
...     num_layers=4,
...     frame_shape=(3, 64, 64)
... )
>>> action = policy.act(frame, temperature=1.0, deterministic=False)

forward(frames)[source]#

Get action logits from frames.

Parameters:: frames (Tensor)
Return type:: Tensor

act(frame, temperature=1.0, deterministic=False)[source]#

Sample action from policy.

Parameters:

frame (Tensor)
temperature (float)
deterministic (bool)

Return type:

Tensor

init_hidden(batch_size, device)[source]#

Initialize hidden state.

Parameters:

batch_size (int)
device (device)

Return type:

tuple[Tensor, Tensor]

Rollout generation utilities for World Models.

This module provides the RolloutGenerator class for collecting episode experience using trained policies in environments.

class simulusion.controller.rollout_generator.RolloutGenerator(env, device, policy=None, max_episode_steps=None, episode_gen=None, name='', enable_streaming_video=False, streaming_video_path=None, streaming_video_fps=20, streaming_video_format='mp4')[source]#

Bases: object

Generator for collecting environment rollouts.

This class handles environment interactions and rollout collection, supporting both random and policy-based action selection.

Variables:

env – The environment to interact with.
device – Device to run computations on.
policy – The policy to use for action selection (optional).
episode_gen – Factory for creating episode objects.
name – Name identifier for the generator.
max_episode_steps – Maximum steps per episode.

Parameters:

env (Any)
device (device | str)
policy (Any)
max_episode_steps (int | None)
episode_gen (Any)
name (str)
enable_streaming_video (bool)
streaming_video_path (str | None)
streaming_video_fps (int)
streaming_video_format (str)

Example

>>> generator = RolloutGenerator(
...     env=env,
...     device='cuda',
...     policy=policy,
...     max_episode_steps=1000
... )
>>> episode = generator.rollout_once()

rollout_once(random_policy=False, explore=False)[source]#

Perform a single rollout of the environment.

Parameters:

random_policy (bool) – If True, use random actions instead of policy.
explore (bool) – If True, add exploration noise to policy actions.

Returns:

Episode object containing the rollout experience.

Return type:

GymImageEnv | MuJoCoImageEnv

rollout_n(n=1, random_policy=False)[source]#

Perform multiple rollouts.

Parameters:

n (int) – Number of rollouts to perform.
random_policy (bool) – If True, use random actions.

Returns:

List of Episode objects.

Return type:

list

rollout_eval_n(n)[source]#

Perform multiple evaluation rollouts with metrics.

Parameters:: n (int) – Number of evaluation rollouts.
Returns:: Tuple of (episodes, frames, metrics).
Return type:: tuple

rollout_eval(collect_latents=False)[source]#

Parameters:: collect_latents (bool)
Return type:: tuple

Datasets, environments, and transforms#

Environment adapters#

The environment APIs below mirror the dedicated environment guide pages: DMC, DeepMind Lab, Gym/Gymnasium, Atari/ALE, Procgen, MuJoCo, Unity ML-Agents, and vectorization utilities. DIAMOND-style Atari support is intentionally not listed as an environment adapter because it is Atari preprocessing rather than a separate environment family.

simulusion.envs.make_env(env_id, **kwargs)[source]

Parameters:

env_id (str)
kwargs (Any)

Return type:

Any

class simulusion.envs.dmc.DeepMindControlEnv(name, seed, size=(64, 64), camera=None)[source]#

Bases: object

Gym-style adapter for DeepMind Control Suite tasks.

The wrapper exposes DMC observations and actions through Gym spaces and adds a rendered RGB image to each observation dict so image-based world model pipelines can train consistently across backends.

Features:

Parses domain-task names (e.g., “cheetah-run” -> domain=”cheetah”, task=”run”)
Automatically handles special cases like “cup” -> “ball_in_cup”
Renders RGB images at configurable resolution
Returns observations as dict with both state vectors and images

Parameters:

name (str) – Environment name in format “domain-task” (e.g., “cheetah-run”).
seed (int) – Random seed for environment initialization.
size (tuple) – Target image size as (height, width) (default: (64, 64)).
camera (int, optional) – Camera ID for rendering. Defaults to 0 for most domains, 2 for quadruped.

Variables:

observation_space (gym.spaces.Dict) – Dict space with state keys and “image”.
action_space (gym.spaces.Box) – Continuous action space from DMC spec.

Example

>>> env = DeepMindControlEnv("cheetah-run", seed=0, size=(64, 64))
>>> obs = env.reset()
>>> print(obs.keys())  # dict_keys(['position', 'velocity', 'image'])

property observation_space: Dict#

property action_space: Box#

step(action)[source]#

Parameters:: action (ndarray)
Return type:: tuple[dict, float, bool, dict]

reset(seed=None)[source]#

Parameters:: seed (int | None)
Return type:: dict

render(*args, **kwargs)[source]#

Parameters:

args (Any)
kwargs (Any)

Return type:

ndarray

simulusion.envs.dmc.make_dmc_env(env_id, seed=0, size=(64, 64), camera=None)[source]#

Create a DeepMind Control Suite task, e.g. make_dmc_env("walker-walk").

Parameters:

env_id (str)
seed (int)
size (tuple[int, int])
camera (int | None)

Return type:

DeepMindControlEnv

simulusion.envs.dmlab.make_dmlab_env(level, **kwargs)[source]#

Create a DeepMind Lab environment adapter for Simulusion.

Parameters:

level (str) – DeepMind Lab level name, for example "rooms_collect_good_objects_train".
**kwargs (Any) – Additional keyword arguments passed to DMLabEnv.

Returns:

A Gym-like wrapper returning {"image": (C, H, W)} uint8 observations and normalized one-hot discrete actions.

Return type:

DMLabEnv

class simulusion.envs.dmlab.DMLabEnv(level, seed=0, size=(64, 64), action_repeat=4, action_set=None, observations=None, config=None, renderer='hardware', **lab_kwargs)[source]#

Bases: object

Gym-style adapter for DeepMind Lab 3D environments.

Parameters:

level (str)
seed (int)
size (tuple[int, int])
action_repeat (int)
action_set (Sequence[Sequence[int]] | np.ndarray | None)
observations (Sequence[str] | None)
config (dict[str, Any] | None)
renderer (str)
lab_kwargs (Any)

property observation_space: Dict#

property action_space: _OneHotActionSpace#

property max_episode_steps: int#

reset(seed=None)[source]#

Parameters:: seed (int | None)
Return type:: dict[str, ndarray]

step(action)[source]#

Parameters:: action (ndarray)
Return type:: tuple[dict[str, ndarray], float, bool, dict[str, Any]]

render(*args, **kwargs)[source]#

Parameters:

args (Any)
kwargs (Any)

Return type:

ndarray

close()[source]#

Return type:: None

simulusion.envs.gym_env.make_gym_env(env, **kwargs)[source]#

Create a GymImageEnv wrapper for generic Gym/Gymnasium environments.

Parameters:

env (Any) – Either a string environment ID (e.g., “Pendulum-v1”) or a pre-built gym environment instance.
**kwargs (Any) – Additional keyword arguments passed to GymImageEnv, including: - seed (int): Random seed for environment (default: 0) - size (tuple): Target image size as (height, width) (default: (64, 64)) - render_mode (str): Render mode for environment (default: “rgb_array”)

Returns:

A wrapper that always returns image observations in the: format {“image”: (C, H, W)} suitable for pixel-based world models.

Return type:

GymImageEnv

class simulusion.envs.gym_env.GymImageEnv(env, seed=0, size=(64, 64), render_mode='rgb_array', include_state=False)[source]#

Bases: object

Gym-like environment wrapper that always returns image observations.

This wrapper normalizes diverse environment interfaces to return consistent image-based observations suitable for pixel-based world models like Dreamer.

Features:

Supports environment IDs (string) and pre-built environment objects.
For vector observations, uses the environment’s own render() frame, and only synthesizes band images from the vector when rendering fails.
Exposes continuous action spaces mapped to [-1, 1] range.
Converts discrete actions to one-hot vectors.
Returns observations as dicts with required key "image" and optional key "state" when include_state=True and a vector observation is available.

Parameters:

env (Any) – Either a string environment ID (e.g., “Pendulum-v1”) or a pre-built gym environment instance.
seed (int) – Random seed for environment reset (default: 0).
size (tuple) – Target image size as (height, width) (default: (64, 64)).
render_mode (str) – Render mode for environment (default: “rgb_array”).
include_state (bool) – Include a flattened low-dimensional "state" key in observations when one can be derived from the underlying env.

Variables:

observation_space – Dict space with required "image" key and optional "state" key when enabled.
action_space – Box space with actions in [-1, 1] range.
max_episode_steps – Maximum steps per episode (default: 1000).

property observation_space: Dict#

property action_space: Box#

property max_episode_steps: int#

reset(seed=None)[source]#

Parameters:: seed (int | None)
Return type:: dict[str, Any]

step(action)[source]#

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], float, bool, dict[str, Any]]

render(*args, **kwargs)[source]#

Parameters:

args (Any)
kwargs (Any)

Return type:

ndarray

close()[source]#

Return type:: None

simulusion.envs.ale_atari_env.make_atari_env(env_id, obs_type='rgb', frameskip=4, repeat_action_probability=0.25, full_action_space=False, max_episode_steps=None, **kwargs)[source]#

Create any Atari environment from Arcadic Learning Environment (ALE).

Parameters:

env_id (str) – The id of the Atari environment to create.
obs_type (str) – The type of observation to return (“rgb” or “ram”).
frameskip (int) – The number of frames to skip between actions.
repeat_action_probability (float) – The probability of repeating the last action.
full_action_space (bool) – Whether to use the full action space.
max_episode_steps (Optional[int]) – Maximum number of steps per episode.
**kwargs – Additional keyword arguments for environment configuration.

Returns:

The created Atari environment.

Return type:

gym.Env

simulusion.envs.ale_atari_env.list_available_atari_envs()[source]#

Get a list of all available Atari environments in Arcadic Learning Environment (ALE).

Returns:: List of available Atari environment IDs.
Return type:: list[str]

simulusion.envs.ale_atari_vector_env.make_atari_vector_env(game, num_envs, obs_type='rgb', frameskip=4, repeat_action_probability=0.25, full_action_space=False, max_episode_steps=None, seed=None, **kwargs)[source]#

Create vectorized Atari environments using ALE’s native AtariVectorEnv.

Parameters:

game (str) – The name of the Atari game (e.g., “pong”, “breakout”).
num_envs (int) – Number of parallel environments.
obs_type (str) – The type of observation to return (“rgb” or “ram”).
frameskip (int) – The number of frames to skip between actions.
repeat_action_probability (float) – The probability of repeating the last action.
full_action_space (bool) – Whether to use the full action space.
max_episode_steps (Optional[int]) – Maximum number of steps per episode.
seed (Optional[int]) – Random seed for reproducibility.
**kwargs – Additional keyword arguments for environment configuration.

Returns:

The vectorized Atari environment.

Return type:

AtariVectorEnv

Procgen environment adapter for Simulusion image-based agents.

simulusion.envs.procgen_env.list_procgen_envs()[source]#

Return the Procgen game names understood by ProcgenImageEnv.

Return type:: list[str]

simulusion.envs.procgen_env.normalize_procgen_env_name(env)[source]#

Normalize Procgen Gym ids and shorthand names to Procgen game names.

Accepted forms include "coinrun", "procgen-coinrun-v0", and "procgen:procgen-coinrun-v0".

Parameters:: env (str)
Return type:: str

simulusion.envs.procgen_env.make_procgen_env(env, **kwargs)[source]#

Create a single-environment Procgen adapter.

Parameters:

env (str) – Procgen game name or Gym-style id.
**kwargs (Any) – Options forwarded to ProcgenImageEnv.

Returns:

Simulusion-compatible image wrapper exposing {"image": (3, H, W) uint8} observations and one-hot-like actions.

Return type:

ProcgenImageEnv

class simulusion.envs.procgen_env.ProcgenImageEnv(env, seed=0, size=(64, 64), distribution_mode='easy', num_levels=0, start_level=None, action_n=15, **procgen_kwargs)[source]#

Bases: object

Adapt Procgen’s vector API to Simulusion’s single-env image interface.

The upstream procgen.ProcgenEnv API is vectorized, so this wrapper builds a one-environment vector and unwraps the leading batch dimension. Actions are exposed as a continuous one-hot-like Box[-1, 1] with one element per discrete Procgen action, matching Simulusion’s other discrete image adapters.

Parameters:

env (str)
seed (int)
size (tuple[int, int])
distribution_mode (str)
num_levels (int)
start_level (int | None)
action_n (int)
procgen_kwargs (Any)

property observation_space: Dict#

property action_space: _ProcgenActionSpace#

property max_episode_steps: int#

reset(seed=None)[source]#

Parameters:: seed (int | None)
Return type:: dict[str, NDArray[uint8]]

step(action)[source]#

Parameters:: action (Any)
Return type:: tuple[dict[str, NDArray[uint8]], float, bool, dict[str, Any]]

render(*args, **kwargs)[source]#

Parameters:

args (Any)
kwargs (Any)

Return type:

NDArray[uint8]

close()[source]#

Return type:: None

simulusion.envs.mujoco_env.make_mujoco_env_from_config(args, size)[source]#

Build a MuJoCo image environment from a DreamerConfig-like object.

Parameters:

args (Any)
size (tuple[int, int])

Return type:

Any

class simulusion.envs.mujoco_env.MuJoCoImageEnv(xml_path=None, *, xml_string=None, binary_path=None, assets=None, seed=0, size=(64, 64), camera=None, reward_fn=None, terminal_fn=None, frame_skip=1, reset_noise_scale=0.0, default_control_range=(-1.0, 1.0), include_state=False)[source]#

Bases: object

Native MuJoCo environment adapter for pixel-based world-model training.

Parameters:

xml_path (str | Path | None)
xml_string (str | None)
binary_path (str | Path | None)
assets (dict[str, bytes] | None)
seed (int)
size (tuple[int, int])
camera (str | int | None)
reward_fn (RewardFn | None)
terminal_fn (TerminalFn | None)
frame_skip (int)
reset_noise_scale (float)
default_control_range (tuple[float, float])
include_state (bool)

property observation_space: Dict#

property action_space: Box#

reset(seed=None)[source]#

Parameters:: seed (int | None)
Return type:: dict[str, ndarray]

step(action)[source]#

Parameters:: action (Any)
Return type:: tuple[dict[str, ndarray], float, bool, dict[str, Any]]

render()[source]#

Return type:: Any

close()[source]#

Return type:: None

simulusion.envs.mujoco_env.make_mujoco_env(model=None, *, backend='auto', seed=0, size=(64, 64), render_mode='rgb_array', gym_kwargs=None, **kwargs)[source]#

Create one MuJoCo image environment factory for tasks and MJCF/MJB models.

Parameters:

model (str | Path | None) – Either a Gymnasium MuJoCo task id such as "Humanoid-v4", an MJCF XML path/string, or an MJB binary path.
backend (str) – "auto" infers native vs Gymnasium task mode. Use "native" for MJCF/MJB, "gymnasium" for task ids, or "robotics" for Gymnasium Robotics registrations.
seed (int) – Seed forwarded to the image wrapper.
size (tuple[int, int]) – Target (height, width) image size.
render_mode (str) – Render mode used for Gymnasium MuJoCo task ids.
gym_kwargs (dict[str, Any] | None) – Optional keyword arguments forwarded to gymnasium.make in task-id mode. Extra **kwargs are also forwarded there.
**kwargs (Any) – Native MuJoCoImageEnv options for MJCF/MJB mode, or environment-constructor options for Gymnasium task-id mode.

Returns:

A Simulusion image environment returning {"image": uint8[C, H, W]}.

Return type:

simulusion.envs.robotics_env.is_moved_mujoco_error(exc)[source]#

Return whether Gymnasium reported the v2/v3 MuJoCo move.

Parameters:: exc (BaseException)
Return type:: bool

simulusion.envs.robotics_env.register_gymnasium_robotics_envs()[source]#

Import Gymnasium Robotics so its environments are registered with Gymnasium.

Return type:: Any

simulusion.envs.robotics_env.list_gymnasium_robotics_envs()[source]#

List all Gymnasium Robotics ids registered by the installed package.

Return type:: list[str]

simulusion.envs.robotics_env.make_gymnasium_env_with_robotics_fallback(env, *, render_mode='rgb_array', gym_kwargs=None, **kwargs)[source]#

Create a Gymnasium env and retry after Robotics registration if needed.

Parameters:

env (str)
render_mode (str)
gym_kwargs (dict[str, Any] | None)
kwargs (Any)

Return type:

Any

simulusion.envs.robotics_env.make_robotics_env(env, *, seed=0, size=(64, 64), render_mode='rgb_array', gym_kwargs=None, **kwargs)[source]#

Create a Simulusion image wrapper for a Gymnasium Robotics environment.

Parameters:

env (str) – Any environment id registered by gymnasium-robotics.
seed (int) – Seed forwarded to GymImageEnv.
size (tuple[int, int]) – Target (height, width) image size.
render_mode (str) – Render mode forwarded to gymnasium.make.
gym_kwargs (dict[str, Any] | None) – Optional keyword arguments forwarded to gymnasium.make.
**kwargs (Any) – Additional keyword arguments forwarded to gymnasium.make.

Returns:

A GymImageEnv that emits {"image": uint8[C, H, W]} observations.

Return type:

GymImageEnv

simulusion.envs.unity_env.make_unity_mlagents_env(env_id=None, **kwargs)[source]#

Create a Unity ML-Agents environment wrapper.

Factory function that instantiates a UnityMLAgentsEnv with the provided keyword arguments. Suitable for integrating Unity-based environments with Dreamer-style world model pipelines.

Parameters:

**kwargs (Any) – Keyword arguments passed to UnityMLAgentsEnv, including: - file_name (str): Path to the Unity environment binary. - behavior_name (str, optional): Name of the behavior to use. - seed (int): Random seed (default: 0). - size (tuple): Image size as (height, width) (default: (64, 64)). - worker_id (int): Worker ID for multi-environment setup (default: 0). - base_port (int): Base port for communication (default: 5005). - no_graphics (bool): Disable graphics rendering (default: True). - time_scale (float): Simulation time scale (default: 20.0). - quality_level (int): Graphics quality level (default: 1). - max_episode_steps (int): Max steps per episode (default: 1000).
env_id (str | None)
**kwargs

Returns:

A Gym-compatible wrapper for Unity environments.

Return type:

UnityMLAgentsEnv

class simulusion.envs.unity_env.UnityMLAgentsEnv(file_name, behavior_name=None, seed=0, size=(64, 64), worker_id=0, base_port=5005, no_graphics=True, time_scale=20.0, quality_level=1, max_episode_steps=1000, include_state=False)[source]#

Bases: object

Gym-like wrapper for Unity ML-Agents environments.

Provides a unified interface for Unity-based environments, converting observations to image format compatible with pixel-based world models.

Features:

Supports single-agent control with continuous action spaces.
Returns observations as dicts with required key "image".
Normalizes actions to [-1, 1] range.
Exposes non-visual sensors in info["vector_observation"] for debugging.

Parameters:

file_name (str) – Path to the Unity environment binary.
behavior_name (str, optional) – Name of the behavior to use. If None, uses the first available behavior.
seed (int) – Random seed for environment (default: 0).
size (tuple) – Target image size as (height, width) (default: (64, 64)).
worker_id (int) – Worker ID for multi-environment setup (default: 0).
base_port (int) – Base port for Unity environment communication (default: 5005).
no_graphics (bool) – Disable graphics rendering for faster simulation (default: True).
time_scale (float) – Simulation time scale multiplier (default: 20.0).
quality_level (int) – Graphics quality level 0-5 (default: 1).
max_episode_steps (int) – Maximum steps per episode (default: 1000).
include_state (bool) – Include a flattened non-visual "state" key in observations when the Unity behavior exposes non-image sensors.

Variables:

observation_space – Dict space with required "image" key containing (3, H, W) uint8 frames.
action_space – Box space with actions in [-1, 1] range.
max_episode_steps – Maximum steps per episode.

Raises:

ValueError – If no behaviors found or action space is not continuous.
RuntimeError – If no agents available after reset.

property observation_space: Dict#

property action_space: Box#

property max_episode_steps: int#

reset(seed=None)[source]#

Parameters:: seed (int | None)
Return type:: dict[str, Any]

step(action)[source]#

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], float, bool, dict[str, Any]]

render(*args, **kwargs)[source]#

Parameters:

args (Any)
kwargs (Any)

Return type:

Any

close()[source]#

Return type:: None

class simulusion.envs.vector_env.SimWorker(worker_id, env_factory, num_envs, command_queue, result_queue, seed=None)[source]#

Bases: Process

Worker process that manages a batch of environment instances. Handles batched stepping for parallel rollouts.

Parameters:

worker_id (int)
env_factory (Callable)
num_envs (int)
command_queue (Queue)
result_queue (Queue)
seed (Optional[int])

run()[source]#

Main worker loop.

Return type:: None

class simulusion.envs.vector_env.VectorizedEnv(env_factory, num_workers=2, envs_per_worker=4, seed=None)[source]#

Bases: ABC

Abstract base class for vectorized environments. Manages multiple worker processes for parallel simulation.

Parameters:

env_factory (Callable)
num_workers (int)
envs_per_worker (int)
seed (Optional[int])

abstractmethod step_batch(actions)[source]#

Step all environments with batched actions.

Parameters:: actions (Tensor)
Return type:: Dict[str, Any]

abstractmethod reset_batch()[source]#

Reset all environments.

Return type:: Dict[str, Any]

render_batch()[source]#

Render all environments.

Return type:: List[ndarray]

close()[source]#

Shutdown all workers.

Return type:: None

class simulusion.envs.vector_env.TorchVectorizedEnv(*args, **kwargs)[source]#

Bases: VectorizedEnv

Simulusion-compatible vectorized environment. Returns batched tensors suitable for PyTorch training.

Parameters:

args (Any)
kwargs (Any)

step_batch(actions)[source]#

Step all environments with batched actions.

Parameters:: actions (Tensor) – Tensor of shape (total_envs, action_dim)
Returns:: Dict with ‘obs’, ‘reward’, ‘done’, ‘info’ tensors
Return type:: Dict[str, Any]

reset_batch()[source]#

Reset all environments and return initial observations.

Return type:: Dict[str, Any]

class simulusion.envs.wrappers.TimeLimit(env, duration)[source]#

Bases: object

Terminate episodes after a fixed number of wrapper steps.

If the wrapped environment does not provide a discount flag at timeout, the wrapper injects a default discount of 1.0 for downstream learners.

Parameters:

env (Any)
duration (int)

step(action)[source]#

Parameters:: action (Any)
Return type:: tuple[Any, Any, bool, dict[str, Any]]

reset(*, seed=None)[source]#

Parameters:: seed (int | None)
Return type:: Any

class simulusion.envs.wrappers.FrameStack(env, num_frames)[source]#

Bases: object

Stack the most recent image observations along the channel axis.

The wrapped environment must emit dict observations with an "image" key in (C, H, W) layout. Non-image keys such as "state" are passed through unchanged from the latest observation.

Parameters:

env (Any)
num_frames (int)

property observation_space: Dict#

property action_space: Any#

reset(*, seed=None)[source]#

Parameters:: seed (int | None)
Return type:: dict[str, Any]

step(action)[source]#

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], Any, bool, dict[str, Any]]

class simulusion.envs.wrappers.ActionRepeat(env, amount)[source]#

Bases: object

Repeat each action for a fixed number of environment steps.

Rewards are accumulated and the loop stops early if the environment terminates, mirroring common action-repeat behavior in world model papers.

Parameters:

env (Any)
amount (int)

step(action)[source]#

Parameters:: action (Any)
Return type:: tuple[Any, float, bool, dict[str, Any]]

class simulusion.envs.wrappers.NormalizeActions(env)[source]#

Bases: object

Expose a normalized [-1, 1] action space for bounded continuous controls.

Incoming normalized actions are mapped back to the wrapped environment action bounds before stepping the environment.

Parameters:: env (Any)

property action_space: Box#

step(action)[source]#

Parameters:: action (ndarray)
Return type:: tuple[Any, Any, bool, dict[str, Any]]

class simulusion.envs.wrappers.ObsDict(env, key='obs')[source]#

Bases: object

Convert scalar/array observations into a dictionary observation format.

This harmonizes outputs for code paths that expect keyed observations (for example {“image”: …} style world model inputs).

Parameters:

env (Any)
key (str)

property observation_space: Dict#

property action_space: Any#

step(action)[source]#

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], Any, bool, dict[str, Any]]

reset(*, seed=None)[source]#

Parameters:: seed (int | None)
Return type:: dict[str, Any]

class simulusion.envs.wrappers.OneHotAction(env)[source]#

Bases: object

Wrap discrete-action environments to accept one-hot action vectors.

The wrapper validates one-hot inputs and converts them to integer action indices before forwarding to the underlying environment.

Parameters:: env (Any)

property action_space: Box#

step(action)[source]#

Parameters:: action (ndarray)
Return type:: tuple[Any, Any, bool, dict[str, Any]]

reset(*, seed=None)[source]#

Parameters:: seed (int | None)
Return type:: Any

class simulusion.envs.wrappers.RewardObs(env)[source]#

Bases: object

Augment observations with the latest scalar reward under obs[“reward”].

Useful for agents that consume reward as part of the observation stream during model learning or recurrent policy inference.

Parameters:: env (Any)

property observation_space: Dict#

step(action)[source]#

Parameters:: action (Any)
Return type:: tuple[dict[str, Any], Any, bool, dict[str, Any]]

reset(*, seed=None)[source]#

Parameters:: seed (int | None)
Return type:: dict[str, Any]

class simulusion.envs.wrappers.ResizeImage(env, size=(64, 64))[source]#

Bases: object

Resize image-like observation entries to a target spatial size.

The wrapper discovers image keys from env.obs_space, applies nearest neighbor resizing, and updates the advertised observation space shapes.

Parameters:

env (Any)
size (tuple[int, int])

property obs_space: dict[str, Any]#

step(action)[source]#

Parameters:: action (Any)
Return type:: Any

reset(*, seed=None)[source]#

Parameters:: seed (int | None)
Return type:: Any

class simulusion.envs.wrappers.RenderImage(env, key='image')[source]#

Bases: object

Inject RGB renders from env.render(“rgb_array”) into observations.

This is useful when the base environment returns non-image observations but a rendered camera view is needed for world-model training.

Parameters:

env (Any)
key (str)

property obs_space: dict[str, Any]#

step(action)[source]#

Parameters:: action (Any)
Return type:: Any

reset(*, seed=None)[source]#

Parameters:: seed (int | None)
Return type:: Any

class simulusion.envs.wrappers.UUID(env)[source]#

Bases: Wrapper

Gym wrapper that tracks a unique run identifier per environment reset.

The ID combines timestamp and UUID and can be used to tag episodes or artifacts generated during data collection.

Parameters:: env (Any)

reset(**kwargs)[source]#

Parameters:: kwargs (Any)
Return type:: Any

class simulusion.envs.wrappers.SelectAction(env, key)[source]#

Bases: Wrapper

Gym wrapper for dictionary actions that forwards a selected key only.

This enables integration with policies that emit action dicts while the environment expects a single tensor/array action payload.

Parameters:

env (Any)
key (str)

step(action)[source]#

Parameters:: action (dict[str, Any])
Return type:: Any

Atari preprocessing helpers#

These helpers wrap Atari environments for specific training recipes. They are not separate environment families.

class simulusion.envs.diamond_atari.DiamondAtariWrapper(env, frameskip=4, max_noop=30, terminate_on_life_loss=True, reward_clip=True, resize=(64, 64), seed=None)[source]#

Bases: object

Atari wrapper for DIAMOND following the paper specifications: - frameskip: number of frames to skip (default 4) - max_noop: maximum number of noop actions at reset (default 30) - terminate_on_life_loss: terminate episode when life is lost (default True) - reward_clip: clip rewards to [-1, 0, 1] (default True) - resize: resize observations to specified size (default 64x64)

Parameters:

env (Any)
frameskip (int)
max_noop (int)
terminate_on_life_loss (bool)
reward_clip (bool)
resize (Optional[Tuple[int, int]])
seed (int | None)

seed(seed=None)[source]#

Parameters:: seed (int | None)
Return type:: None

close()[source]#

Release the underlying environment.

Callers reasonably expect the gym close() contract, and without this the usual env.close() teardown raises AttributeError.

Return type:: None

render(*args, **kwargs)[source]#

Forward rendering to the wrapped environment.

Parameters:

args (Any)
kwargs (Any)

Return type:

Any

step(action)[source]#

Step the environment.

For backwards compatibility with older gym APIs this wrapper returns a 4-tuple: (obs, reward, done, info). Internally it supports gymnasium’s 5-tuple and collapses (terminated, truncated) into a single done bool.

info["terminated"] and info["truncated"] stay separate: done ends the episode either way, but only terminated is a genuine episode end and therefore the correct target for the termination head and the correct place to cut the lambda-return’s bootstrap.

Parameters:: action (int)
Return type:: Any

reset(**kwargs)[source]#

Parameters:: kwargs (Any)
Return type:: Tuple[Any, Dict[str, Any]]

simulusion.envs.diamond_atari.make_diamond_atari_env(game, frameskip=4, max_noop=30, terminate_on_life_loss=True, reward_clip=True, resize=(64, 64), seed=None)[source]#

Create a DIAMOND-compatible Atari environment.

Parameters:

game (str) – Atari game name. Accepts either the bare "Breakout-v5" or the namespaced "ALE/Breakout-v5"; the ALE/ prefix is added when missing. DIAMOND checkpoints store the bare form, so requiring the namespaced one here would make them unloadable.
frameskip (int) – Number of frames to skip between actions
max_noop (int) – Maximum number of noop actions at reset
terminate_on_life_loss (bool) – Whether to terminate on life loss
reward_clip (bool) – Whether to clip rewards to [-1, 0, 1]
resize (Tuple[int, int]) – Target size for observations
seed (int | None) – Random seed

Returns:

Configured Atari environment

Return type:

DiamondAtariWrapper

Datasets and transforms#

Data generation and dataset classes for World Models.

This module provides utilities for generating rollout data from environments and PyTorch dataset classes for loading observation sequences.

class simulusion.datasets.wm_dataset.RolloutDataset(root, transform, train=True, buffer_size=1000, num_test_files=600)[source]#

Bases: Dataset

PyTorch Dataset for loading rollout data.

This dataset loads pre-collected rollout trajectories from disk, providing a buffer-based mechanism for efficient data loading. It supports train/test splits and custom transforms.

Variables:

root – Root directory containing rollout .npz files.
transform – Albumentations transform to apply to observations.
train – If True, use training split; otherwise use test split.
buffer_size – Maximum number of files to keep in memory.
num_test_files – Number of files to use for test set.

Parameters:

root (str)
transform (Compose)
train (bool)
buffer_size (int)
num_test_files (int)

Example

>>> transform = transforms.Compose([transforms.ToTensor()])
>>> dataset = RolloutDataset(
...     root='data/carracing',
...     transform=transform,
...     train=True,
...     buffer_size=100,
... )
>>> obs, action, reward, terminal = dataset[0]

load_next_buffer()[source]#

Load the next batch of rollout files into memory.

This method implements a circular buffer, loading buffer_size files at a time and advancing through the dataset sequentially.

Return type:: None

class simulusion.datasets.wm_dataset.ObservationDataset(root, transform, train=True, buffer_size=1000, num_test_files=600)[source]#

Bases: RolloutDataset

Dataset for single observation samples (not sequences).

This dataset extends RolloutDataset to provide individual observations rather than sequences, suitable for VAE training.

Example

>>> dataset = ObservationDataset(
...     root='data/carracing',
...     transform=transform,
...     train=True,
... )
>>> obs = dataset[0]

Parameters:

root (str)
transform (Compose)
train (bool)
buffer_size (int)
num_test_files (int)

class simulusion.datasets.wm_dataset.SequenceDataset(root, transform, train, buffer_size, num_test_files, seq_len)[source]#

Bases: RolloutDataset

Dataset for sequential rollout data.

This dataset provides sequences of observations, actions, rewards, and terminal flags suitable for training recurrent models like MDRNN.

Variables:

seq_len – Length of sequences to return.

Parameters:

root (str)
transform (Compose)
train (bool)
buffer_size (int)
num_test_files (int)
seq_len (int)

Example

>>> dataset = SequenceDataset(
...     root='data/carracing',
...     transform=transform,
...     train=True,
...     seq_len=32,
... )
>>> obs, action, reward, terminal, next_obs = dataset[0]

class simulusion.datasets.wm_dataset.LatentSequenceDataset(latents_arr, actions, rewards, terminals, train, buffer_size, num_test_files, seq_len)[source]#

Bases: Dataset

Dataset for pre-computed latent sequences.

This dataset uses pre-encoded latent representations instead of raw images, which significantly reduces memory usage during RNN training.

Parameters:

latents_arr (ndarray)
actions (ndarray)
rewards (ndarray)
terminals (ndarray)
train (bool)
buffer_size (int)
num_test_files (int)
seq_len (int)

class simulusion.datasets.video_datasets.DatasetConfig(num_frames=16, image_size=64, batch_size=4, num_workers=4, pin_memory=True, shuffle=True)[source]#

Bases: object

Base configuration for datasets.

Parameters:

num_frames (int)
image_size (int)
batch_size (int)
num_workers (int)
pin_memory (bool)
shuffle (bool)

num_frames: int = 16#

image_size: int = 64#

batch_size: int = 4#

num_workers: int = 4#

pin_memory: bool = True#

shuffle: bool = True#

class simulusion.datasets.video_datasets.VideoDatasetBase(data_source, num_frames=16, image_size=64, transform=None, normalize=True)[source]#

Bases: Dataset

Base class for video datasets.

All video datasets should inherit from this class and implement the _load_video method.

Parameters:

data_source (str | Path | List[str] | List[Path])
num_frames (int)
image_size (int)
transform (Callable | None)
normalize (bool)

data_source: str | Path | List[str] | List[Path]#

video_paths: Sequence[Path | int]#

class simulusion.datasets.video_datasets.VideoFolderDataset(data_source, num_frames=16, image_size=64, transform=None, normalize=True, extensions=('.mp4', '.avi', '.mkv', '.webm', '.mov'), recursive=True)[source]#

Dataset that loads videos from a folder.

Supports common video formats: .mp4, .avi, .mkv, .webm

Usage:

dataset = VideoFolderDataset(
    data_source="/path/to/videos",
    num_frames=16,
    image_size=64
)

Parameters:

data_source (str | Path | List[str] | List[Path])
num_frames (int)
image_size (int)
transform (Callable | None)
normalize (bool)
extensions (Tuple[str, ...])
recursive (bool)

class simulusion.datasets.video_datasets.ImageFolderDataset(data_source, num_frames=16, image_size=64, transform=None, normalize=True, extensions=('.jpg', '.jpeg', '.png', '.bmp'), image_sort_key=None)[source]#

Dataset that loads image sequences from folders.

Each subfolder is treated as a video sequence.

Usage:

dataset = ImageFolderDataset(
    data_source="/path/to/images",
    num_frames=16,
    image_size=64
)

Parameters:

data_source (str | Path | List[str] | List[Path])
num_frames (int)
image_size (int)
transform (Callable | None)
normalize (bool)
extensions (Tuple[str, ...])
image_sort_key (Callable | None)

class simulusion.datasets.video_datasets.NumPyDataset(data_source, num_frames=16, image_size=64, transform=None, normalize=True, key=None)[source]#

Dataset that loads videos from numpy files.

Supports .npy and .npz files.

Usage:

dataset = NumPyDataset(
    data_source="/path/to/videos.npy",
    num_frames=16,
    image_size=64
)

Parameters:

data_source (str | Path)
num_frames (int)
image_size (int)
transform (Callable | None)
normalize (bool)
key (str | None)

class simulusion.datasets.video_datasets.RLEnvironmentDataset(data_source, num_frames=16, image_size=64, transform=None, normalize=True, obs_key='observations')[source]#

Dataset for RL environment recordings.

Loads trajectories stored as: - .npz files with ‘observations’ and ‘actions’ keys - Directory with episode folders

Usage:

dataset = RLEnvironmentDataset(
    data_source="/path/to/rl_episodes",
    num_frames=16,
    image_size=64
)

Parameters:

data_source (str | Path)
num_frames (int)
image_size (int)
transform (Callable | None)
normalize (bool)
obs_key (str)

class simulusion.datasets.video_datasets.HDF5Dataset(data_source, num_frames=16, image_size=64, transform=None, normalize=True, key='videos', memmap=False)[source]#