Dreamer4

Install the dreamer4 package using pip.

$ pip install dreamer4

The model uses signal_levels and step_sizes_log2 to control the diffusion/flow process.

• Inference: You must provide signal_levels (to define the noise level) or step_sizes_log2.
• Training: If neither is provided, the model generates them randomly using prob_shortcut_train.
• Constraints: step_sizes must be powers of 2. If you provide step_sizes, the model internally derives step_sizes_log2.
• Relationship: You cannot provide both signal_levels and step_sizes_log2 simultaneously; one must be absent to determine if the mode is inference or training.

The model supports several policy optimization objectives via the objective parameter (used internally during training/optimization steps):

• ppo (Proximal Policy Optimization): Uses the standard clipped surrogate loss to prevent large policy updates.
• spo (Scaled Policy Optimization): A clipped surrogate loss variant.
• pmpo (Parallelized/Policy-specific variant): A specialized objective that weights positive and negative advantages equally by taking the sign of the advantage rather than its magnitude. It also supports an optional pmpo_reverse_kl for KL divergence calculation and uses a pmpo_pos_to_neg_weight ($\alpha$) to balance the loss components.

When providing rewards or terminals to the model, they must align with the sequence length of the latents. If the provided sequence is one step shorter than the latent sequence (length time - 1), the model automatically left-pads them with 0. for rewards or False for terminals to ensure alignment with the states.

During training, if return_pred_only is False, rewards must perfectly match the video length time unless they are being padded.

The BaseRecordEnvWrapper and its subclasses use a standardized process to extract and format images from observations to ensure they are ready for video encoding or model input:

1. Extraction: It uses obs_image_dotpath to traverse the observation dictionary. If no path is provided, it defaults to looking for 'image' or 'state' keys in a dictionary.
2. Batching: If a 4D tensor (B, C, H, W) is provided, it selects the first item in the batch.
3. Channel Ordering: It converts channel-first tensors (C, H, W) to channel-last (H, W, C).
4. Scaling: If the image is floating point and in the range [0, 1], it scales it to [0, 255] and converts to uint8 (byte).
5. RGB Conversion: If the image is grayscale (H, W), it expands it to (H, W, 1) and then repeats the channel to create a 3-channel RGB image (H, W, 3).

The Transformer class supports efficient inference/generation through a KV-cache mechanism. When calling forward, you can pass a time_cache to avoid recomputing keys and values for past time steps.

• time_cache: A tuple (past_len, tuple(next_kv_cache)).
• return_time_cache: If True, the forward method returns the updated cache for the next step.

This is particularly useful for autoregressive generation or long-sequence processing where causal=True is used.

# Using Transformer with caching
x = torch.randn(seq_len, dim)
time_cache = None

# First pass
x, (next_len, next_cache) = transformer(x, return_time_cache=True)

# Subsequent pass with cache
x_new = torch.randn(new_seq_len, dim)
x_out, (new_next_len, new_cache) = transformer(
    x_new, 
    time_cache=(next_len, next_cache), 
    return_time_cache=True
)

The implementation uses specific shorthand for tensor dimensions in its internal logic and documentation:

• b: Batch size
• n: Sequence length / Number of tokens
• h: Attention heads
• d: Feature dimension
• f: Frequencies (for rotary embeddings)
• l: Logit / Predicted bins
• y: Layers of transformer
• p: Positions (e.g., 3 for spacetime)
• t: Time
• na: Action dimension (number of discrete and continuous actions)
• g: Groups of query heads to key heads (GQA)

The DynamicsWorldModelWrapper.step() method accepts actions that can be formatted in several ways depending on the DynamicsWorldModel's configuration (number of discrete vs continuous action types):

• Both Discrete and Continuous: Pass a tuple (discrete_action, continuous_action).
• Discrete Only: Pass the discrete action tensor/array.
• Continuous Only: Pass the continuous action tensor/array.

Actions are automatically cast to tensors and moved to the correct device. The wrapper handles both unbatched and batched inputs, ensuring they are reshaped to (batch, 1, dim) for the model's internal processing.

The VideoTokenizer is used to compress video data into a latent space, typically learned through MAE (Masked Autoencoders) and LPIPS.

Key parameters:

• dim: Embedding dimension.
• dim_latent: Latent dimension.
• patch_size: Size of patches for the tokenizer.
• image_height / image_width: Input image dimensions.

You can instantiate it manually or load a saved checkpoint using init_and_load().

import torch
from dreamer4 import VideoTokenizer

tokenizer = VideoTokenizer(
    dim = 512,
    dim_latent = 32,
    patch_size = 32,
    image_height = 256,
    image_width = 256
)

video = torch.randn(2, 3, 10, 256, 256)
loss = tokenizer(video)
loss.backward()

# To load a saved tokenizer:
tokenizer = VideoTokenizer.init_and_load('./checkpoints/experiment_name/tokenizer.pt')

The action_embedder within DynamicsWorldModel handles the conversion of discrete and continuous actions into model dimensions. It supports:

• Discrete actions: Via embedding.
• Continuous actions: Via normalization and distribution types ('gaussian', 'squashed_gaussian', or 'beta').
• Multi-token prediction: If multi_token_pred_len > 1, the embedder can support unembedding for predicting future actions.