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]

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 synora.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 synora.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 synora.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)

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]

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

evaluate()[source]

Return type:: tuple

class synora.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 synora.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]

High-level training loop. Delegates single-step training to the existing train function.

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 synora.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]

Return type:: None

class synora.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)

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]

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:

actions

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 synora.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)[source]

Allocate an empty KV cache sized to this model’s context capacity.

Parameters:

batch_size (int)
device (device)
dtype (dtype | 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)[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)

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 synora.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 synora.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

synora.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

synora.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 synora.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 synora.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 synora.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]

synora.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 synora.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 synora.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 synora.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

synora.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:

synora.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:

synora.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:

synora.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

synora.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 synora.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 synora.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

synora.DreamerRSSM: alias of RSSM

class synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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

synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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

synora.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 synora.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 synora.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 synora.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)[source]

Parameters:

model (Module)
x_t (Tensor)
t (Tensor)

Return type:

Tensor

sample(model, n, img_size, channels)[source]

Parameters:

model (Module)
n (int)
img_size (int)
channels (int)

Return type:

Tensor

class synora.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:

action

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:

actions

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 synora.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]

synora.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 synora.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:

GymImageEnv | MuJoCoImageEnv

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 synora.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

synora.MultiHeadAttention: alias of MultiHeadSelfAttention

class synora.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 synora.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 synora.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 synora.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:

Episode

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 synora.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:

action

class synora.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 synora.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 synora.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 synora.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='synora', 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 = 'synora'

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 synora.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='synora', 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 = 'synora'

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 synora.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

synora.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

synora.list_dit_presets()[source]

Return the Table 1 model names.

Return type:: list[str]

synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 Synora’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

synora.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

synora.list_available_bsuite_ids()[source]

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

Return type:: list[str]

synora.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 synora.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 Synora’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

synora.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 Synora image environment returning {"image": uint8[C, H, W]}.

Return type:

synora.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

synora.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]

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

Create a Synora 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

synora.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 synora.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

synora.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 synora.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

synora.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 synora.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 Synora 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 Synora 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

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

Create a Synora 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 synora.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 synora.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 Synora-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

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

Create a DeepMind Lab environment adapter for Synora.

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 synora.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

synora.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

synora.MujocoEnv: alias of MuJoCoImageEnv

class synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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

synora.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]

synora.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 synora.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 synora.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

synora.DreamerRewardModel: alias of RewardModel

synora.DreamerValueModel: alias of ValueModel

synora.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

synora.deregister_world_model(name)[source]

Remove a previously registered model from the plugin registry.

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

synora.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

synora.list_registered_models()[source]

Return names of all externally registered world models.

Return type:: list[str]

synora.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

synora.deregister_env_backend(name)[source]

Remove a previously registered environment backend.

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

synora.list_registered_env_backends()[source]

Return names of all externally registered env backends.

Return type:: list[str]

synora.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]

synora.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]

synora.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 synora.Logger(log_dir, enable_wandb=False, wandb_project='synora', 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 synora.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])

synora.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

synora.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 synora.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

synora.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]

synora.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

synora.enable_performance_defaults(*, tf32=True, cudnn_benchmark=True)[source]

Turn on the CUDA throughput settings Synora 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

synora.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

synora.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 synora.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

synora.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

synora.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]

synora.make_stepper(model, **kwargs)[source]

Build the stepper matching model’s family.

Parameters:

model (Any)
kwargs (Any)

Return type:

WorldModelStepper

class synora.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 synora.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 synora.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

synora.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

synora.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

synora.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]

synora.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 synora.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

synora.load_bundle(directory)[source]

Read a bundle written by save_bundle().

Parameters:: directory (str | Path)
Return type:: Bundle

class synora.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

User-facing convenience APIs for Synora.

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 synora.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 synora.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

synora.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

synora.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

synora.api.get_env_backend_spec(name)[source]#

Return metadata for an environment backend name or alias.

Parameters:: name (str)
Return type:: EnvBackendSpec

synora.api.get_model_spec(name)[source]#

Return metadata for a model name or alias.

Parameters:: name (str)
Return type:: ModelSpec

synora.api.list_env_backends()[source]#

Return canonical backend names accepted by make_env().

Return type:: list[str]

synora.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]]

synora.api.list_models()[source]#

Return canonical model names accepted by create_model().

Return type:: list[str]

synora.api.make_env(env_id, backend='auto', **kwargs)[source]#

Create an environment with a consistent Synora 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 Synora’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.

Synora’s public models and agents get .export() from ExportableAgentMixin. Any other nn.Module goes through export_model(); Synora does not modify torch.nn.Module.

class synora.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 synora.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 synora.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]

synora.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 Synora 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

synora.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

synora.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.

synora.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 Synora 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 synora.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

synora.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 Synora 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.

synora.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]

synora.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

synora.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 Synora 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 synora.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 synora.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 synora.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 synora.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 synora.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

synora.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 synora.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 synora.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

synora.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

synora.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 synora.inference.rollout_drift().

class synora.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

synora.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]

synora.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 synora.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]

synora.inference.bundle.load_bundle(directory)[source]#

Read a bundle written by save_bundle().

Parameters:: directory (str | Path)
Return type:: Bundle

synora.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 synora.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

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 synora.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.

synora.models.dreamer.get_available_memory()[source]#

Get available physical memory in bytes.

Return type:: int

synora.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

synora.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 synora.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)

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]

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 synora.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)

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]#

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

evaluate()[source]#

Return type:: tuple

class synora.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 synora.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 synora.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 synora.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

High-level training loop. Delegates single-step training to the existing train function.

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

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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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]

synora.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 synora.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]#

Return type:: None

synora.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

synora.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

synora.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

synora.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

synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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

synora.models.vit.vit_predictor(**kwargs)[source]#

Factory for a JEPA predictor transformer with sensible defaults.

Parameters:: kwargs (Any)
Return type:: VisionTransformerPredictor

synora.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

synora.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

synora.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

synora.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

synora.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

synora.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

synora.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 synora.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)

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]

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:

actions

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 synora.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

class synora.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 synora.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 synora.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)[source]#

Allocate an empty KV cache sized to this model’s context capacity.

Parameters:

batch_size (int)
device (device)
dtype (dtype | 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)[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)

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 synora.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 synora.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

synora.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]

synora.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:

synora.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:

synora.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 synora.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]

synora.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 synora.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 synora.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

synora.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 synora.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 synora.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)[source]#

Parameters:

model (Module)
x_t (Tensor)
t (Tensor)

Return type:

Tensor

sample(model, n, img_size, channels)[source]#

Parameters:

model (Module)
n (int)
img_size (int)
channels (int)

Return type:

Tensor

synora.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

synora.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 synora.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 synora.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

synora.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 synora.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 synora.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 synora.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 synora.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 synora.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

synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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:

action

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:

actions

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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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

synora.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 synora.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 synora.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 synora.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

synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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:

GymImageEnv | MuJoCoImageEnv

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

synora.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 synora.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='synora', 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 = 'synora'

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 synora.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='synora', 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 = 'synora'

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 synora.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

synora.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

synora.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

synora.configs.list_dit_presets()[source]

Return the Table 1 model names.

Return type:: list[str]

class synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.configs.wm_config.WMControllerConfig(latent_size=32, hidden_size=200, action_size=3, env_name='CarRacing-v2', 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-v2'#

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 synora.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='synora', 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 = 'synora'#

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 synora.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='synora', 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 = 'synora'#

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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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#

synora.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

synora.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

synora.configs.dit_config.list_dit_presets()[source]#

Return the Table 1 model names.

Return type:: list[str]

synora.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

synora.configs.diamond_config.get_default_device()[source]#

Return type:: str

class synora.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 synora.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 modules for World Models.

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-v2 –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

synora.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

synora.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

synora.training.train_world_model.test_trained_model(logdir, env_name, action_size, num_episodes=5)[source]#

Test the trained world model with controller in the environment.

Parameters:

logdir (str)
env_name (str)
action_size (int)
num_episodes (int)

Return type:

None

synora.training.train_world_model.main()[source]#

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.

synora.training.train_convvae.save_checkpoint(state, is_best, filename, best_filename)[source]#

Save model checkpoint.

Parameters:

state (dict) – Dictionary containing model state to save.
is_best (bool) – If True, also save as best checkpoint.
filename (str) – Path to save checkpoint.
best_filename (str) – Path to save best checkpoint.

Return type:

None

synora.training.train_convvae.test_epoch(model, test_loader, device, loss_fn)[source]#

Run one epoch of validation.

Parameters:

model (ConvVAE) – The VAE model to evaluate.
test_loader (DataLoader) – DataLoader for test/validation data.
device (device) – Device to run evaluation on.
loss_fn (Any) – Loss function to use.

Returns:

Average test loss for the epoch.

Return type:

float

synora.training.train_convvae.train_epoch(epoch, model, optimizer, train_loader, device, train_dataset, loss_fn, use_amp=False, scaler=None)[source]#

Run one epoch of training.

Parameters:

epoch (int) – Current epoch number.
model (Any) – The VAE model to train.
optimizer (Any) – Optimizer for training.
train_loader (Any) – DataLoader for training data.
device (Any) – Device to run training on.
train_dataset (Any) – Training dataset (used to load next buffer if applicable).
loss_fn (Any) – Loss function to use.
use_amp (bool) – Whether to use automatic mixed precision.
scaler (GradScaler | None) – GradScaler for mixed precision training.

Return type:

float

synora.training.train_convvae.train_convae(config)[source]#

Train a Convolutional VAE model.

This function trains a ConvVAE on observation data using the provided configuration. It handles data loading, model initialization, training loop, checkpointing, and sample generation.

Parameters:: config (WMVAEConfig) – WMVAEConfig object containing all training hyperparameters.
Return type:: None

The training process includes:

Loading pretrained VAE if available (unless noreload is True)
Training for specified number of epochs
Validating after each epoch
Learning rate scheduling with ReduceLROnPlateau
Early stopping based on validation loss
Checkpointing best and current models
Generating sample images at specified intervals

Example

>>> config = WMVAEConfig(
...     height=64, width=64, latent_size=32, num_epochs=100, logdir='results'
... )
>>> train_convae(config)

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.

synora.training.train_mdn_rnn.precompute_latents(vae_config, mdrnn_config)[source]#

Pre-compute and save VAE latents to disk for memory-efficient RNN training.

This function encodes all observations using the VAE and saves the latent representations to disk. This allows RNN training without keeping the VAE in GPU memory.

Parameters:

vae_config (WMVAEConfig) – WMVAEConfig for loading pretrained VAE.
mdrnn_config (WMMDNRNNConfig) – WMMDNRNNConfig containing latent_size and device settings.

Return type:

None

synora.training.train_mdn_rnn.save_checkpoint(state, is_best, filename, best_filename)[source]#

Save model checkpoint.

Parameters:

state (Any) – Dictionary containing model state to save.
is_best (bool) – If True, also save as best checkpoint.
filename (str) – Path to save checkpoint.
best_filename (str) – Path to save best checkpoint.

Return type:

None

synora.training.train_mdn_rnn.to_latent(vae, obs, next_obs, device, red_size=64)[source]#

Transform observations to latent space using VAE encoder.

This function encodes observations into the latent space using the VAE’s encoder network. It applies the reparameterization trick to sample from the learned latent distribution.

Parameters:

vae (ConvVAE) – Trained VAE model with encoder.
obs (Tensor) – Batch of current observations.
next_obs (Tensor) – Batch of next observations.
device (device) – Device to run encoding on.
red_size (int) – Target size for resizing images (default: 64).

Returns:

Tuple of (latent_obs, latent_next_obs) tensors in latent space.

Return type:

tuple[Tensor, Tensor]

synora.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]

synora.training.train_mdn_rnn.data_pass(epoch, mdrnn, vae, train_loader, test_loader, optimizer, device, include_reward, test_every=10, epochs=1, use_amp=False, scaler=None, latent_size=32, max_seq_len=50, log_wandb=False, prev_val_loss=1000000.0, early_stop=None, lr_scheduler=None, batch_size=50, train=True)[source]#

Run one epoch of training or validation.

Parameters:

epoch (int) – Current epoch number.
mdrnn (Any) – MDRNN model.
vae (Any) – VAE model for encoding observations (None if using precomputed latents).
train_loader (Any) – Training data loader.
test_loader (Any) – Test/validation data loader.
optimizer (Any) – Optimizer (used only for training).
device (Any) – Device to run on.
include_reward (bool) – Whether to include reward in loss.
latent_size (int) – Size of latent space.
batch_size (int) – Batch size.
train (bool) – If True, run training pass; otherwise run validation.
use_amp (bool) – If True, use automatic mixed precision.
scaler (Any) – GradScaler for mixed precision training.
test_every (int)
epochs (int)
max_seq_len (int)
log_wandb (bool)
prev_val_loss (float)
early_stop (Any)
lr_scheduler (Any)

Returns:

Average loss for the epoch.

Return type:

float

synora.training.train_mdn_rnn.train_mdn_rnn(vae_config, mdrnn_config, use_precomputed_latents=True, use_amp=True)[source]#

Train an MDRNN model.

This function trains an MDRNN on sequence data using the provided configurations. It loads a pretrained VAE for encoding observations into latent space, then trains the MDRNN to predict future latent states given current latent states and actions.

Parameters:

vae_config (WMVAEConfig) – WMVAEConfig for loading pretrained VAE.
mdrnn_config (WMMDNRNNConfig) – WMMDNRNNConfig containing MDRNN training hyperparameters.
use_precomputed_latents (bool) – If True, use pre-encoded latents from disk.
use_amp (bool) – If True, use automatic mixed precision for memory efficiency.

Return type:

None

The training process includes:

Loading pretrained VAE from vae_config.logdir
Training for specified number of epochs
Validating after each epoch
Learning rate scheduling with ReduceLROnPlateau
Early stopping based on validation loss
Checkpointing best and current models

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

synora.training.train_controller.flatten_parameters(parameters)[source]#

Parameters:: parameters (Any)
Return type:: ndarray

synora.training.train_controller.load_parameters(params, controller)[source]#

Parameters:

params (Any)
controller (Any)

Return type:

None

synora.training.train_controller.slave_routine(p_queue, r_queue, e_queue, p_index, config, time_limit)[source]#

Worker process routine for parallel rollout evaluation.

Parameters:

p_queue (Any) – Queue containing (s_id, parameters) to evaluate.
r_queue (Any) – Queue where to place results (s_id, reward).
e_queue (Any) – End queue - when non-empty, process terminates.
p_index (int) – Process index for GPU assignment.
config (Any) – Controller configuration (must include env_name and action_size).
time_limit (int) – Maximum steps per episode.

Return type:

None

synora.training.train_controller.evaluate(solutions, results, rollouts, p_queue, r_queue)[source]#

Evaluate current controller.

Parameters:

solutions (Any)
results (Any)
rollouts (int)
p_queue (Any)
r_queue (Any)

Return type:

Any

synora.training.train_controller.train_controller(config)[source]#

Train a linear controller using CMA-ES.

Parameters:: config (WMControllerConfig) – WMControllerConfig containing training hyperparameters, including env_name and action_size.
Return type:: None

The training process includes:

Setting up parallel evaluation workers (each loads VAE + MDRNN)
Running CMA-ES optimization with parallel rollout evaluation
Evaluating and saving best controller checkpoint

synora.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

synora.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

synora.training.train_jepa.sweep_train()[source]#

Function for WandB sweep agent.

Return type:: None

synora.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

synora.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 synora.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]#

Full training loop.

Parameters:

total_epochs (int | None) – Total training epochs
eval_interval (int) – Evaluate every N epochs
save_dir (str) – Directory to save checkpoints

Return type:

None

synora.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 synora.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 synora.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 synora.training.train_genie.GenieTrainer(model, config, device=None)[source]#

Bases: object

Trainer for Genie model.

Parameters:

model (Module)
config (GenieConfig)
device (device | 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]#

Full training loop.

Parameters:

train_dataloader (DataLoader) – Training data loader
val_dataloader (DataLoader | None) – Validation data loader (optional)
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 periodic checkpoints go. Without it nothing is written until the caller saves, so a run killed by a timeout, an OOM or Ctrl+C leaves nothing behind.
checkpoint_interval (int) – Steps between checkpoints. 0 disables them, which is the previous behaviour.

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

synora.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]

synora.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

synora.training.train_planet.train(memory, rssm, optimizer, device, N=32, H=50, beta=1.0, grads=False)[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.

Parameters:

memory (Any)
rssm (Any)
optimizer (Any)
device (device)
N (int)
H (int)
beta (float)
grads (bool)

Return type:

dict

synora.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 synora.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

synora.training.train_planet.main(argv=None)[source]#

Example PlaNet/RSSM training script with rollout collection and evaluation.

Builds environment/model/policy objects, iteratively trains on replayed episodes, and periodically saves videos and checkpoints.

Parameters:: argv (list[str] | None)
Return type:: None

synora.training.train_rssm.train_rssm(memory, model, optimizer, record_grads=True)[source]#

Train an RSSM on replayed trajectories for one optimization phase.

Samples batches from memory, computes reconstruction and KL objectives across rollout steps, and returns aggregated loss metrics.

Parameters:

memory (Any)
model (Any)
optimizer (Any)
record_grads (bool)

Return type:

dict

synora.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

synora.training.train_rssm.main()[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.

Return type:: None

class synora.training.train_diamond.DiamondAgent(config)[source]#

Bases: object

DIAMOND: DIffusion As a Model Of eNvironment Dreams

RL agent trained entirely within a diffusion world model.

Parameters:: config (DiamondConfig)

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

train()[source]#

Main training loop following Algorithm 1.

Return type:: 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

synora.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

synora.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 synora.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 synora.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

synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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:

action

class synora.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 synora.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 synora.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 synora.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:

Episode

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.

synora.envs.make_env(env_id, **kwargs)[source]

Parameters:

env_id (str)
kwargs (Any)

Return type:

Any

class synora.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

synora.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

synora.envs.dmlab.make_dmlab_env(level, **kwargs)[source]#

Create a DeepMind Lab environment adapter for Synora.

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 synora.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

synora.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 synora.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

synora.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

synora.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]

synora.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 Synora image-based agents.

synora.envs.procgen_env.list_procgen_envs()[source]#

Return the Procgen game names understood by ProcgenImageEnv.

Return type:: list[str]

synora.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

synora.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:

Synora-compatible image wrapper exposing {"image": (3, H, W) uint8} observations and one-hot-like actions.

Return type:

ProcgenImageEnv

class synora.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 Synora’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 Synora’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

synora.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 synora.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

synora.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 Synora image environment returning {"image": uint8[C, H, W]}.

Return type:

synora.envs.robotics_env.is_moved_mujoco_error(exc)[source]#

Return whether Gymnasium reported the v2/v3 MuJoCo move.

Parameters:: exc (BaseException)
Return type:: bool

synora.envs.robotics_env.register_gymnasium_robotics_envs()[source]#

Import Gymnasium Robotics so its environments are registered with Gymnasium.

Return type:: Any

synora.envs.robotics_env.list_gymnasium_robotics_envs()[source]#

List all Gymnasium Robotics ids registered by the installed package.

Return type:: list[str]

synora.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

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

Create a Synora 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

synora.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 synora.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 synora.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 synora.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 synora.envs.vector_env.TorchVectorizedEnv(*args, **kwargs)[source]#

Bases: VectorizedEnv

Synora-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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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]]

synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.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 synora.datasets.video_datasets.HDF5Dataset(data_source, num_frames=16, image_size=64, transform=None, normalize=True, key='videos', memmap=False)[source]#