```
├── .gitignore
├── LICENSE (omitted)
├── README.md (3.2k tokens)
├── __init__.py
├── assets/
├── actiontokenizer.png
├── datasets_stylized.png
├── dynamicsmodel.png
├── finitescalarquantizer.png
├── spacetimetransformer.png
├── tinyworlds.png
├── tinyworldsarchv3.png
├── tinyworldsdark.png
├── tinyworldslight.png
├── videotokenizer.png
├── configs/
├── dev/
├── dev_training.yaml (200 tokens)
├── dynamics.yaml (100 tokens)
├── inference.yaml (100 tokens)
├── latent_actions.yaml (100 tokens)
├── training.yaml (200 tokens)
├── video_tokenizer.yaml (100 tokens)
├── datasets/
├── __init__.py
├── data_utils.py (1500 tokens)
├── datasets.py (1800 tokens)
├── models/
├── __init__.py
├── dynamics.py (2.2k tokens)
├── fsq.py (600 tokens)
├── latent_actions.py (1300 tokens)
├── muon.py (600 tokens)
├── norms.py (600 tokens)
├── patch_embed.py (400 tokens)
├── positional_encoding.py (500 tokens)
├── st_transformer.py (2.3k tokens)
├── utils.py
├── video_tokenizer.py (1000 tokens)
├── requirements.txt
├── scripts/
├── download_assets.py (1100 tokens)
├── full_train.py (600 tokens)
├── run_inference.py (1500 tokens)
├── train_dynamics.py (2.5k tokens)
├── train_latent_actions.py (1500 tokens)
├── train_video_tokenizer.py (1500 tokens)
├── visualize_batch.py (900 tokens)
├── utils/
├── config.py (2.3k tokens)
├── distributed.py (700 tokens)
├── inference_utils.py (1700 tokens)
├── optimizer_utils.py (500 tokens)
├── scheduler_utils.py (200 tokens)
├── utils.py (2.8k tokens)
├── wandb_utils.py (700 tokens)
```
## /.gitignore
```gitignore path="/.gitignore"
*results*/
*pycache*/
data/
batch_visualizations/*
wandb/*
setup.sh
```
## /README.md
<div align="center">
<picture>
<source media="(prefers-color-scheme: light)" srcset="/assets/tinyworldslight.png">
<img alt="tiny corp logo" src="assets/tinyworldsdark.png" width="80%" height="80%">
</picture>
</div>
TinyWorlds is a minimal autoregressive world model built on Google Deepmind's [Genie Architecture](https://arxiv.org/pdf/2402.15391).
World models can't use action-less internet video to scale like [VEO3](https://deepmind.google/models/veo/). Deepmind's [Genie](https://deepmind.google/discover/blog/genie-3-a-new-frontier-for-world-models/) solves this by inferring the actions between frames using **no prior action data**.
TinyWorlds is meant to help people understand the clever autoregressive, unsupervised method Deepmind likely used to achieve **scalable world models**.
## Table of Contents
- [Getting Started](#getting-started)
- [Overview](#architecture-overview)
- [Building Blocks](#architecture-building-blocks)
- [Space-Time Transformer](#space-time-transformer-stt)
- [Variational Autoencoder](#vaes)
- [Finite Scalar Quantization](#finite-scalar-quantization)
- [Architecture](#architecture)
- [Video Tokenizer](#video-tokenizer)
- [Action Tokenizer](#action-tokenizer)
- [Dynamics Model](#dynamics-model)
- [TinyWorlds Inference](#full-tinyworlds-inference)
- [Data](#data)
- [Training/Inference Acceleration](#traininginference-acceleration)
- [Shape Annotation Key](#shape-annotation-key)
- [Contributing](#contributing)
# Getting Started
```bash
# Installation
git clone https://github.com/AlmondGod/tinyworlds.git
cd tinyworlds
pip install -r requirements.txt
export WANDB_API_KEY=<YOUR_WANDB_API_KEY>
export PYTHONPATH="/workspace/tinyworlds:$PYTHONPATH"
# Training
# 1. download data from huggingface
python scripts/download_assets.py datasets --pattern "zelda_frames.h5"
# 2. run training
python scripts/full_train.py --config configs/training.yaml -- --dataset=ZELDA
# Inference
# 1. pull pretrained sonic checkpoints from huggingface
python scripts/download_assets.py models --suite-name sonic
# 2. run inference
python scripts/run_inference.py --config configs/inference.yaml -- use_latest_checkpoints=true dataset=SONIC
```
# Overview
### Why World Models?
*To shape the world, generate it*
A [world model](https://arxiv.org/pdf/1803.10122) is a function mapping the current state of an environment to the next state of an environment.
To predict the next environment state accurately, this function must compress all information in the world into a set of laws.
So the world model **captures all the inherent structure and emergent phenomena of the world.**
All of deep learning, and all of intelligence, is [trying to compress the universe into a model](https://arxiv.org/pdf/0812.4360). A model that can predict important aspects of the next state of the universe, by learning heuristics about how it operates.
Our universe can also be thought of as a world model. It is a map from state to state executing every moment by following a set of laws. Humans experience the many layers of emergent behavior over these foundational laws.
As of 2025, video-based world models have been practically applied as:
1. cortexes to give physical world understanding to robots
2. simulators for models to interact with physics fully-online
3. experiences with new structures of reality for humans to interact with
But humans are only at the very beginning of modeling our own worlds.
TinyWorlds is built to help you to understand world modeling better, and to [learn by contributing](#contributing).
### Architecture Overview
TinyWorlds is an autoregressive transformer over discrete tokens, so we can also use SOTA LLM techniques to improve our world model.
Why discrete tokens? Discretization makes our dynamics prediction problem much easier, because instead of predicting an image a near-infinite continuous space, it need only select one of the ~1000 tokens in our vocabulary (aka codebook).
TinyWorlds consists of three modules:
**Video Tokenizer:** This tokenizer reconstructs a sequence of video with a small discrete bottleneck (our video tokens) in the middle. **This layer compresses the important information from video to tokens.**
**Action Tokenizer:** This tokenizer **infers the discrete action token between two frames**. It trains by reconstructing the next frame using the previous frame and a discrete action token that sees the next frame.
**Dynamics Model:** Given past action and frame tokens, this predicts our next frame tokens. **This should capture the physics of our tiny video game worlds.**
# Building Blocks
### Space-Time Transformer
[Space-Time Transformer](https://arxiv.org/pdf/2001.02908) (STT) is a transformer for video. Each STT block contains a spatial attention layer, a temporal attention layer, and a FeedForward Network (FFN). For a brush up on self-attention, see Karpathy's [GPT From Scratch Video](https://youtu.be/kCc8FmEb1nY?si=tvfcBnGHBbEiS70v&t=3748)
In the spatial layer, each token attends to all other tokens in the same frame. In the temporal layer, each token attends to tokens in the same position but previous timesteps.
The FFN is a multi-layer perceptron on each embedding vector. Inspired by divine benevolence, I used [SwiGLU](https://arxiv.org/pdf/2002.05202) for the FFN. SwiGLU adds a Gated Linear Unit (GLU) to [Swish](https://en.wikipedia.org/wiki/Swish_function), and is computed as
$x_t = W_3[\sigma(W_1x + b1) * W_2x + b2] + b3$ (see SwiGLU diagram for clarity)
For regular STT, I used [Root Mean Squared Normalization (RMSNorm)](https://docs.pytorch.org/docs/stable/generated/torch.nn.modules.normalization.RMSNorm.html) as the normalizer, which is less sensitive to extreme outliers than 0-variance norm. In RMS, we divide our input by
$\sqrt(\epsilon + x / \sum x^2)$.
For STT conditioned on actions, I used [Feature-wise Linear Modulation (FiLM)](https://arxiv.org/pdf/1709.07871). FiLM passes actions for each timestep through an FFN to transform each action latent into Gamma ($\gamma$) and Beta ($\beta$) vectors. Our norm is then
$(x - \mu) / \sigma * (1 + \gamma) + \beta$
### Variational Autoencoder
*\*VAEs are complex, but below is an overview with many details omitted*
[Variational Autoencoders]((https://arxiv.org/pdf/1906.02691)) (VAEs) are defined by:
1. An encoder network to parameterize the approximate posterior distribution $q(z | x)$ of latent variables $z$ given data $x$
3. A decoder network to parameterize the likelihood $p(x | z)$ over input data x given latent z
VAEs maximize $log(p(x | z))$, the likelihood the decoder exactly reconstructs the input x given latent z from the encoder.
The important takeaway is that $z$ is low dimensional, so for reconstruction, it will compress all the important information from $x$.
### Finite Scalar Quantization
Since we want a set of discrete tokens, we quantize continous $z$ to one of a finite set of possible $z$.
If vectors are points in high dimensional space, [Finite Scalar Quantization](https://arxiv.org/pdf/2309.15505) (FSQ) is a quantization method that divides space into hypercubes, and the hypercube a vector falls into becomes its quantized representation.
Concretely, we quantize a vector in FSQ by:
1. tanh(x) which bounds to [-1,1]
2. scale/shift to [0, L]
3. round to the nearest integer (quantization step)
4. scale/shift back to [-1,1]
The token vocabulary has size ${L^{D}}$ where $L$ is bins per dimension and $D$ is the dimensionality of the hypercube. With 3 dimensions and 2 levels per dimension, we'd have 8 regions in the cube and size 8 token vocabulary.
FSQ VAEs let us learn structured hypercubes to use as our token vocabularies that encode information about the input. In our context, maybe one of these hypercubes represents moving left, another jumping, another crouching, et cetera.
To allow gradients to flow to the encoder (since quantization is non-differentiable), we pass the post-quantization gradients directly to the pre-quantization layer.
Precisely, the decoder takes as input $z + stopgrad(z_q - z)$ where stopgrad is, in pytorch, `.detach()`. The decoder only uses $z_q$ (since $z - z = 0$), but the gradient is taken only on $z$.
# Architecture
### Video Tokenizer
The video tokenizer is an FSQ VAE that compresses videos into discrete tokens. It reduces the dimensionality of dynamics while enabling high quality video generation.
It converts patches to embeddings with pixel-mixing 2D Convolutions.
It then uses an STTransformer over the embeddings to produce quantized tokens.
Each video token contains information about both its own patch and other patches in the same location or timestep.
Finally, it decodes the video tokens into a reconstructed image.
### Action Tokenizer
The Action Tokenizer is also an FSQ VAE, and is the key to scalability. It allows us to train without action labels by learning to infer actions between two frames. We then condition the dynamics on these actions.
The encoder takes in a sequence of frames and outputs action tokens between the frames.
The decoder takes in all previous frames $(x_1...x_t-1)$ and quantized action latent vectors $(a_1...a_t-1)$ as input and predicts the next frames $(x_2...x_t)$.
Action tokens should learn to encode the most meaningful change between the past and current frame, which should correspond to some high-level action.
In practice, the action decoder tries to ignore actions and infer purely from images. To counteract this,
1. we mask most frames except the first, so the decoder must learn to use the string of actions as signal for reconstruction
2. we encourage batch-wise variance in the encoder through an auxiliary loss
At inference time, we map each key to one of the action tokens that conditions the dynamics for the user to influence video generation.
### Dynamics Model
At timestep $t$, the dynamics model should take in tokenized video tokens $z_{1..t - 1}$ and action tokens $a_{_1..t - 1}$ and predict next frame tokens $z_{t}$.
In practice, we train dynamics like [MaskGIT](https://arxiv.org/pdf/2202.04200) and [BERT](https://arxiv.org/pdf/1810.04805).
We mask a subset of tokens and train our model to predict the masked tokens, conditioned on all current and previous frame and action tokens.
To infer dynamics at a given step, we first append a fully masked frame to our context sequence. Then, for T steps we:
1. Predict logits at each masked position
2. Compute token probabilities with softmax
3. Sample the k most likely tokens out of the still unmasked positions
4. Place them into the context tensor, removing corresponding mask tokens
5. Repeat
I chose an exponential schedule for k (first step samples ~1 token, then ~2, then ~5, then ~20, then ~50, etc)
### TinyWorlds Inference
Given initial context frames from the training distribution, we first tokenize them.
We then run the following loop:
1. The player specifies one of the n_actions action tokens to use by choosing integer in $[0, |A|]$
2. Condition the dynamics model with context window c on the video tokens t-c...t and action tokens t-c..t and run dynamics inference
3. Detokenize the predicted video tokens into a new video frame for the user
We repeat this process autoregressively over the time dimension as actions are passed to the model, tokens are predicted by the dynamics model, we detokenize them into frames to display to the user.
This process also lets us predict multiple future frames at once (bounded by memory and the training distribution), which can improve inference quality.
### Data
The data is processed and downsampled from gameplay `.mp4s` into `.h5` files.
You can download existing datasets from [Huggingface TinyWorlds Datasets](https://huggingface.co/datasets/AlmondGod/tinyworlds/tree/main) with the datasets command in [getting started](#getting-started).
Available are:
1. **PicoDoom** (`picodoom_frames.h5`), a minimal version of Doom
2. **Pong** (`pong_frames.h5`), the classic
3. **Zelda Ocarina of Time** (`zelda_frames.h5`), one of the originl 2D Zelda games
4. **Pole Position** (`pole_position_frames.h5`), a pixel racing game
5. **Sonic** (`sonic_frames.h5`), the original game
To create a new dataset, create a new dataclass in [datasets.py](datasets/datasets.py) and specify mp4 path. PR or dm me to upload your dataset to the HF repo so others can use it :)
### Training/Inference Acceleration
TinyWorlds supports the following torch features to accelerate training and/or inference:
1. **Torch compile**, which allows us to use faster CUDA kernels for certain pre-optimized operations like attention and matmuls
2. **Distributed data parallel (DDP)**, which allows us to train using multiple gpus by using same model different data per-gpu
3. **Automatic mixed precision (AMP)**, which scales certain ops from FP32 to BF16 based on the current nodes used floating point range
4. **TF32 training**, which lets us use NVIDIA TensorFloat32 for tensor-core-optimized matmuls and convolutions
### Shape Annotation Key
All tensors are shape-annotated and use einops tensor manipulation operations with the following abbreviations:
**B:** batch size \
**T:** time/sequence dimension (number of frames) \
**P:** number of patch tokens per frame \
**E:** embedding dim (d_model) \
**L:** Video Tokenizer latent dim \
**A:** Action Tokenizer latent dim (action dim) \
**D:** number of bins for each video tokenizer dim \
**L^D:** Size of the video tokenizer vocabulary \
**C:** image channels \
**H:** pixel-grid height \
**W:** pixel-grid width \
**Hp:** patch-grid height \
**Wp:** patch-grid width \
**S:** patch size
# Contributing
When you make a PR, please:
1. Retain backwards compatibility
2. Visualize trained model inference and ensure coherence, **put inference visualizations in the PR**
3. Ensure code is easy to read for someone with no context, including shape annotations and reasoning
4. Keep code as lean as possible
There are still many TODOs which may offer significant performance gains...
- [ ] Try `RoPE`/`AliBi` Position Embeddings
- [ ] Add more datasets (Terraria, Street Fighter, \<your favorite retro videogame\>)
- [ ] Try [AdaLN-Zero](https://arxiv.org/pdf/2212.09748) instead of `FiLM` (adds a pre-scale parameter)
- [ ] Add new schedulers for MaskGIT like cosine and [Halton](https://github.com/valeoai/Halton-MaskGIT)
- [ ] Replace `mean pool + concat` in the action tokenizer with `length-2 windowed attention + mean`
- [ ] Spend more compute on a much larger training run, scale to multi-billions of parameters
- [ ] Accelerate dynamics training by producing, saving, and loading pre-processed image patch embeddings instead of full frames
- [x] Implement Mixture of Experts in the Feedforward Network - added by [eren23](https://github.com/eren23) in [#20](https://github.com/AlmondGod/tinyworlds/pull/20)
- [x] Try different optimizers (`Muon`, `SOAP`) - added by [eren23](https://github.com/eren23) in [#20](https://github.com/AlmondGod/tinyworlds/pull/20)
- [x] Train on more GPUs by adding `FSDP` Support — added by [alekseymalakhov11](https://github.com/alekseymalakhov11) in [#11](https://github.com/AlmondGod/tinyworlds/pull/11)
### *Miscellanea*
TinyWorlds (excluding datasets and external assets) is licensed under the MIT [LICENSE](LICENSE). TinyWorlds is an independent research project and is not affiliated with, endorsed by, or sponsored by DeepMind or Google.
*aesthetic inspired by [Tinygrad](https://tinygrad.org/) and [Tinygpu](https://github.com/adam-maj/tiny-gpu)*
## /__init__.py
```py path="/__init__.py"
```
## /assets/actiontokenizer.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/actiontokenizer.png
## /assets/datasets_stylized.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/datasets_stylized.png
## /assets/dynamicsmodel.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/dynamicsmodel.png
## /assets/finitescalarquantizer.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/finitescalarquantizer.png
## /assets/spacetimetransformer.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/spacetimetransformer.png
## /assets/tinyworlds.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/tinyworlds.png
## /assets/tinyworldsarchv3.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/tinyworldsarchv3.png
## /assets/tinyworldsdark.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/tinyworldsdark.png
## /assets/tinyworldslight.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/tinyworldslight.png
## /assets/videotokenizer.png
Binary file available at https://raw.githubusercontent.com/AlmondGod/tinyworlds/refs/heads/main/assets/videotokenizer.png
## /configs/dev/dev_training.yaml
```yaml path="/configs/dev/dev_training.yaml"
# wandb
use_wandb: true
wandb_project: nano-genie-pipeline
# dataset
dataset: PICODOOM
fps: 30
preload_ratio: 0.05
# smaller shared model params
patch_size: 4
context_length: 4
frame_size: 64
latent_dim: 6 # for video tokenizer
num_bins: 2 # for video tokenizer
n_actions: 8
embed_dim: 32
num_heads: 4
hidden_dim: 64
num_blocks: 1
# simplest train settings
compile: false
tf32: true
amp: true
log_interval: 1
n_updates: 2
learning_rate: 0.0001
batch_size_per_gpu: 8
gradient_accumulation_steps: 1
# Stage configs
video_tokenizer_config: configs/video_tokenizer.yaml
latent_actions_config: configs/latent_actions.yaml
dynamics_config: configs/dynamics.yaml
# Which stages to run
run_video_tokenizer: true
run_latent_actions: true
run_dynamics: true
```
## /configs/dynamics.yaml
```yaml path="/configs/dynamics.yaml"
# Training
batch_size_per_gpu: 500
gradient_accumulation_steps: 1
n_updates: 300000
learning_rate: 0.01
log_interval: 2000
use_actions: true
# these can vary but for convenience, here
embed_dim: 32
num_heads: 8
hidden_dim: 128
num_blocks: 8
# Paths
video_tokenizer_path:
latent_actions_path:
# resume from checkpoint
checkpoint:
```
## /configs/inference.yaml
```yaml path="/configs/inference.yaml"
# paths (need if not get latest)
video_tokenizer_path:
latent_actions_path:
dynamics_path:
use_latest_checkpoints: true
# inference params
dataset: PONG
preload_ratio:
device: mps # make sure to use correct device
generation_steps: 10
context_window: 2
fps: 2
temperature: 0.5 # 0 is argmax
teacher_forced: false # for testing/dev
prediction_horizon: 1 # how many masked frames to append and decode at a time
# action selection modes (lower takes priority)
use_actions: false # use random actions
use_gt_actions: false # use lam-inferred actions
use_interactive_mode: true # use user-inputted actions
# inference acceleration
amp: false
tf32: false
compile: false
```
## /configs/latent_actions.yaml
```yaml path="/configs/latent_actions.yaml"
# Training
batch_size_per_gpu: 350
gradient_accumulation_steps: 1
n_updates: 10000
learning_rate: 0.0001
log_interval: 500
# these can vary but for convenience, here
embed_dim: 32
num_heads: 8
hidden_dim: 128
num_blocks: 2
# resume from checkpoint
checkpoint: null
```
## /configs/training.yaml
```yaml path="/configs/training.yaml"
# wandb
use_wandb: true
wandb_project: tinyworlds
# dataset
dataset: PICODOOM
preload_ratio: 0.005
# shared model params
patch_size: 4
context_length: 4
frame_size: 64
latent_dim: 5 # for video tokenizer
num_bins: 4 # for video tokenizer
n_actions: 4
# performance
amp: true
tf32: true
compile: true
# distributed launch (torchrun)
distributed:
use_ddp: False
use_fsdp: False
reshard_after_forward: False
nproc_per_node: 1 # how many GPUs per node
standalone: true # if we're using one node, this eliminates extraneous processes
# stage configs
video_tokenizer_config: configs/video_tokenizer.yaml
latent_actions_config: configs/latent_actions.yaml
dynamics_config: configs/dynamics.yaml
# which stages to run
run_video_tokenizer: true
run_latent_actions: true
run_dynamics: true
# optimizer: "adamw" (default) or "muon" (Newton-Schulz orthogonalized)
optimizer: "adamw"
muon_momentum: 0.95
muon_backend_steps: 5
# MoE (dynamics model only): replaces SwiGLU FFN with top-k routed experts
use_moe: false
num_experts: 4
top_k_experts: 2
moe_aux_loss_coeff: 0.01
```
## /configs/video_tokenizer.yaml
```yaml path="/configs/video_tokenizer.yaml"
# Training
batch_size_per_gpu: 350
gradient_accumulation_steps: 1
n_updates: 40000
learning_rate: 0.001
log_interval: 2500
# these can vary but for convenience, here
embed_dim: 32
num_heads: 8
hidden_dim: 128
num_blocks: 4
# resume from checkpoint
checkpoint:
```
## /datasets/__init__.py
```py path="/datasets/__init__.py"
```
## /datasets/data_utils.py
```py path="/datasets/data_utils.py"
import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
import time
import os
import numpy as np
import matplotlib.pyplot as plt
from torchvision.utils import make_grid
from datasets.datasets import PongDataset, SonicDataset, PolePositionDataset, PicoDoomDataset, ZeldaDataset
DEFAULT_NUM_WORKERS = 2
DEFAULT_PREFETCH_FACTOR = 2
DEFAULT_PIN_MEMORY = False
DEFAULT_PERSISTENT_WORKERS = True
def _default_video_transform():
return transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
def _load_video_dataset_pair(dataset_cls, video_rel_path, h5_rel_path, num_frames, transform=None, fps=30, preload_ratio=1, **kwargs):
current_folder_path = os.getcwd()
video_path = current_folder_path + video_rel_path
preprocessed_path = current_folder_path + h5_rel_path
transform = _default_video_transform() if transform is None else transform
train = dataset_cls(
video_path,
transform=transform,
save_path=preprocessed_path,
train=True,
num_frames=num_frames,
fps=fps,
preload_ratio=preload_ratio,
**kwargs
)
val = dataset_cls(
video_path,
transform=transform,
save_path=preprocessed_path,
train=False,
num_frames=num_frames,
fps=fps,
preload_ratio=preload_ratio,
**kwargs
)
return train, val
def load_pong(num_frames=1, fps=15, preload_ratio=1):
return _load_video_dataset_pair(
PongDataset,
'/data/pong.mp4',
'/data/pong_frames.h5',
num_frames=num_frames,
fps=fps,
preload_ratio=preload_ratio
)
def load_sonic(num_frames=4, fps=15, preload_ratio=1):
return _load_video_dataset_pair(
SonicDataset,
'/data/sonic_frames.mp4',
'/data/sonic_frames.h5',
num_frames=num_frames,
fps=fps,
preload_ratio=preload_ratio
)
def load_pole_position(num_frames=4, fps=15, preload_ratio=1):
return _load_video_dataset_pair(
PolePositionDataset,
'/data/pole_position.mp4',
'/data/pole_position_frames.h5',
num_frames=num_frames,
fps=fps,
preload_ratio=preload_ratio
)
def load_picodoom(num_frames=4, fps=30, preload_ratio=1):
return _load_video_dataset_pair(
PicoDoomDataset,
'/data/picodoom cleaned.mp4',
'/data/picodoom_frames.h5',
num_frames=num_frames,
fps=30,
preload_ratio=preload_ratio
)
def load_zelda(num_frames=4, fps=15, preload_ratio=1):
return _load_video_dataset_pair(
ZeldaDataset,
'/data/Zelda oot2d 1 Cut.mp4',
'/data/zelda_frames.h5',
num_frames=num_frames,
fps=fps,
preload_ratio=preload_ratio
)
def data_loaders(train_data, val_data, batch_size, distributed=False, rank=0, world_size=1):
train_sampler = None
val_sampler = None
if distributed:
train_sampler = DistributedSampler(train_data, num_replicas=world_size, rank=rank, shuffle=True, drop_last=True)
val_sampler = DistributedSampler(val_data, num_replicas=world_size, rank=rank, shuffle=False, drop_last=True)
train_loader = DataLoader(
train_data,
batch_size=batch_size,
shuffle=False if train_sampler is not None else True,
sampler=train_sampler,
num_workers=DEFAULT_NUM_WORKERS,
pin_memory=DEFAULT_PIN_MEMORY,
persistent_workers=DEFAULT_PERSISTENT_WORKERS,
prefetch_factor=DEFAULT_PREFETCH_FACTOR,
drop_last=True
)
val_loader = DataLoader(
val_data,
batch_size=batch_size,
shuffle=False if val_sampler is not None else True,
sampler=val_sampler,
num_workers=DEFAULT_NUM_WORKERS,
pin_memory=DEFAULT_PIN_MEMORY,
persistent_workers=DEFAULT_PERSISTENT_WORKERS,
prefetch_factor=DEFAULT_PREFETCH_FACTOR,
drop_last=True
)
return train_loader, val_loader
def load_data_and_data_loaders(dataset, batch_size, num_frames=1, distributed=False, rank=0, world_size=1, fps=15, preload_ratio=1):
if dataset == 'PONG':
training_data, validation_data = load_pong(num_frames=num_frames, fps=fps, preload_ratio=preload_ratio)
elif dataset == 'SONIC':
training_data, validation_data = load_sonic(num_frames=num_frames, fps=fps, preload_ratio=preload_ratio)
elif dataset == 'POLE_POSITION':
training_data, validation_data = load_pole_position(num_frames=num_frames, fps=fps, preload_ratio=preload_ratio)
elif dataset == 'PICODOOM':
training_data, validation_data = load_picodoom(num_frames=num_frames, fps=fps, preload_ratio=preload_ratio)
elif dataset == 'ZELDA':
training_data, validation_data = load_zelda(num_frames=num_frames, fps=fps, preload_ratio=preload_ratio)
else:
raise ValueError('Invalid dataset')
training_loader, validation_loader = data_loaders(
training_data, validation_data, batch_size,
distributed=distributed, rank=rank, world_size=world_size
)
x_train_var = np.var(training_data.data)
return training_data, validation_data, training_loader, validation_loader, x_train_var
def readable_timestamp():
return time.ctime().replace(' ', ' ').replace(
' ', '_').replace(':', '_').lower()
def visualize_reconstruction(original, reconstruction, save_path=None):
# original: (B, C, H, W) or (B, T, C, H, W)
# reconstruction: (B, C, H, W) or (B, T, C, H, W)
# move tensors to CPU and convert to float32 for matplotlib compatibility
original = original.detach().to('cpu', dtype=torch.float32)
reconstruction = reconstruction.detach().to('cpu', dtype=torch.float32)
# handle single frames by expanding to sequences
if original.dim() == 4: # (B, C, H, W)
original = original.unsqueeze(1) # Add sequence dimension
if reconstruction.dim() == 4: # (B, C, H, W)
reconstruction = reconstruction.unsqueeze(1) # Add sequence dimension
# take first 4 sequences, each of length 4 (or available length)
num_sequences = min(4, original.shape[0])
seq_length = min(4, original.shape[1])
original = original[:num_sequences, :seq_length] # (B, T, C, H, W)
reconstruction = reconstruction[:num_sequences, :seq_length] # (B, T, C, H, W)
# create a figure with two subplots side by side
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))
# for original sequences
# reshape to (B * T, C, H, W) for make_grid
orig_flat = original.reshape(-1, *original.shape[2:]) # (B*T, C, H, W)
grid_orig = make_grid(orig_flat, nrow=seq_length, normalize=True, padding=2).clamp(0, 1)
ax1.imshow(grid_orig.permute(1, 2, 0).contiguous().numpy())
ax1.axis('off')
ax1.set_title(f'Original Sequences (4 sequences × {seq_length} frames)')
# for reconstructed sequences
recon_flat = reconstruction.reshape(-1, *reconstruction.shape[2:]) # (B*T, C, H, W)
grid_recon = make_grid(recon_flat, nrow=seq_length, normalize=True, padding=2).clamp(0, 1)
ax2.imshow(grid_recon.permute(1, 2, 0).contiguous().numpy())
ax2.axis('off')
ax2.set_title(f'Reconstructed Sequences (4 sequences × {seq_length} frames)')
plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.close()
else:
plt.show()
plt.close()
```
## /datasets/datasets.py
```py path="/datasets/datasets.py"
from torch.utils.data import Dataset
import cv2
import h5py
import os
from tqdm import tqdm
import numpy as np
import torch
from typing import Optional, Tuple, Union
# TODO: Try pre-caching video tokens and have dataloader load video tokens instead of frames
class VideoHDF5Dataset(Dataset):
def __init__(
self,
video_path: str,
transform=None, # postnormalization
save_path: Optional[str] = None,
train: bool = True,
disable_test_split: bool = True,
num_frames: int = 4, # context length
resize_to: Tuple[int, int] = (64, 64),
fps: int = 30,
sequence_stride: Optional[int] = None, # default 60//fps
fraction_of_dataset: float = 1.0, # fraction of valid starting indices to expose
load_chunk_size: int = 1000, # chunk size when reading from HDF5
load_start_index: int = 0, # skip initial frames when reading cached HDF5
preload_ratio: Optional[float] = None, # if set, only load this ratio of cached frames
preprocess_read_step: int = 1, # step to subsample raw video during preprocessing
preprocess_slice: Optional[Tuple[Union[int, float], Union[int, float]]] = None, # optional slice applied to preprocessed frames; can be (start_idx, end_idx) ints or (start_ratio, end_ratio) floats in [0,1]
) -> None:
self.transform = transform
self.train = train
self.num_frames = num_frames
self.fps = fps
self.frame_skip = max(1, (sequence_stride if sequence_stride is not None else max(1, 60 // fps)))
self.fraction_of_dataset = float(fraction_of_dataset)
self.resize_to = resize_to
if save_path and os.path.exists(save_path):
with h5py.File(save_path, 'r') as h5_file:
frames_dset = h5_file['frames']
total = len(frames_dset)
n_frames = int(total if preload_ratio is None else max(0, min(total, int(total * preload_ratio))))
self.data = []
for i in tqdm(range(load_start_index, n_frames, load_chunk_size), desc=f"Loading {n_frames} frames"):
chunk = frames_dset[i:min(i + load_chunk_size, n_frames)][:]
self.data.extend(chunk)
self.data = np.array(self.data)
else:
frames = self._preprocess_video(
video_path=video_path,
resize_to=resize_to,
read_step=preprocess_read_step,
slice_spec=preprocess_slice,
)
if save_path:
print(f"Saving preprocessed frames to {save_path}")
with h5py.File(save_path, 'w') as f:
f.create_dataset('frames', data=frames, compression='lzf')
# Reload into memory to ensure consistent path
with h5py.File(save_path, 'r') as h5_file:
frames = h5_file['frames'][:]
self.data = frames
if not disable_test_split:
split_idx = int(0.9 * len(self.data))
self.data = self.data[:split_idx] if train else self.data[split_idx:]
def _preprocess_video(
self,
video_path: str,
resize_to: Tuple[int, int],
read_step: int = 1,
slice_spec: Optional[Tuple[Union[int, float], Union[int, float]]] = None,
) -> np.ndarray:
print(f"Preprocessing video {video_path}")
video = cv2.VideoCapture(video_path)
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
frames = []
step = max(1, int(read_step))
for i in tqdm(range(0, total_frames, step), desc="Processing video frames"):
video.set(cv2.CAP_PROP_POS_FRAMES, i)
ret, frame = video.read()
if not ret:
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame = cv2.resize(frame, resize_to, interpolation=cv2.INTER_AREA)
frames.append(frame)
video.release()
frames = np.array(frames)
if slice_spec is not None and len(frames) > 0:
start, end = slice_spec
n = len(frames)
if isinstance(start, float) or isinstance(end, float):
# interpret as ratios
s = 0 if start is None else int(n * max(0.0, min(1.0, float(start))))
e = n if end is None else int(n * max(0.0, min(1.0, float(end))))
else:
s = 0 if start is None else int(start)
e = n if end is None else int(end)
s = max(0, min(n, s))
e = max(s, min(n, e))
frames = frames[s:e]
return frames
def __len__(self) -> int:
max_valid_index = int((len(self.data) - (self.num_frames * self.frame_skip)) * self.fraction_of_dataset)
return max(0, max_valid_index)
def __getitem__(self, index: int):
if index >= len(self):
raise IndexError(f"Index {index} out of bounds for dataset of length {len(self)}")
frame_sequence = self.data[index:index + (self.num_frames * self.frame_skip):self.frame_skip]
if len(frame_sequence) != self.num_frames:
raise ValueError(f"Expected {self.num_frames} frames, got {len(frame_sequence)} frames")
frame_sequence = frame_sequence.astype(np.float32) / 255.0
if self.transform:
transformed_frames = []
for frame in frame_sequence:
transformed_frame = self.transform(frame)
transformed_frames.append(transformed_frame)
frame_sequence = torch.stack(transformed_frames, dim=0)
else:
frame_sequence = torch.from_numpy(frame_sequence).permute(0, 3, 1, 2)
return frame_sequence, 0
def __del__(self):
if hasattr(self, 'h5_file'):
self.h5_file.close()
# TODO: add more datasets
class PongDataset(VideoHDF5Dataset):
def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=1, resolution=(64, 64), fps=30, preload_ratio=1):
super().__init__(
video_path=video_path,
transform=transform,
save_path=save_path,
train=train,
num_frames=num_frames,
resize_to=resolution,
fps=fps,
preload_ratio=preload_ratio,
load_chunk_size=1000,
load_start_index=0,
preprocess_read_step=10, # keep every 10th frame from raw
preprocess_slice=None,
)
class PolePositionDataset(VideoHDF5Dataset):
def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=4, resolution=(64, 64), fps=30, preload_ratio=1):
super().__init__(
video_path=video_path,
transform=transform,
save_path=save_path,
train=train,
num_frames=num_frames,
resize_to=resolution,
fps=fps,
preload_ratio=preload_ratio,
sequence_stride=None,
load_chunk_size=1000,
load_start_index=0,
preprocess_read_step=1,
preprocess_slice=(1/50, 1/4),
)
class SonicDataset(VideoHDF5Dataset):
def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=4, resolution=(128, 128), fps=15, preload_ratio=1):
super().__init__(
video_path=video_path,
transform=transform,
save_path=save_path,
train=train,
num_frames=num_frames,
resize_to=resolution,
fps=fps,
preload_ratio=preload_ratio,
sequence_stride=None,
load_chunk_size=1000,
load_start_index=100,
preprocess_read_step=1,
preprocess_slice=None,
)
class PicoDoomDataset(VideoHDF5Dataset):
def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=4, resolution=(128, 128), fps=30, preload_ratio=0.3):
super().__init__(
video_path=video_path,
transform=transform,
save_path=save_path,
train=train,
num_frames=num_frames,
resize_to=resolution,
fps=fps,
preload_ratio=preload_ratio,
sequence_stride=None,
load_chunk_size=1000,
load_start_index=300,
preprocess_read_step=1,
preprocess_slice=None,
)
class ZeldaDataset(VideoHDF5Dataset):
def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=4, resolution=(128, 128), fps=15, preload_ratio=0.2):
super().__init__(
video_path=video_path,
transform=transform,
save_path=save_path,
train=train,
num_frames=num_frames,
resize_to=resolution,
fps=fps,
preload_ratio=preload_ratio,
sequence_stride=None,
load_chunk_size=1000,
load_start_index=1000,
preprocess_read_step=1,
preprocess_slice=None,
)
```
## /models/__init__.py
```py path="/models/__init__.py"
```
## /models/dynamics.py
```py path="/models/dynamics.py"
from models.utils import ModelType
import torch
import torch.nn as nn
import math
from models.positional_encoding import build_spatial_only_pe
from models.st_transformer import STTransformer
from einops import repeat
class DynamicsModel(nn.Module):
def __init__(self, frame_size=(128, 128), patch_size=4, embed_dim=128, num_heads=8,
hidden_dim=128, num_blocks=4, num_bins=4, n_actions=8, conditioning_dim=3, latent_dim=5,
use_moe=False, num_experts=4, top_k_experts=2, moe_aux_loss_coeff=0.01):
super().__init__()
H, W = frame_size
codebook_size = num_bins**latent_dim
self.latent_embed = nn.Linear(latent_dim, embed_dim)
self.transformer = STTransformer(
embed_dim, num_heads, hidden_dim, num_blocks, causal=True,
conditioning_dim=conditioning_dim,
use_moe=use_moe, num_experts=num_experts,
top_k_experts=top_k_experts, moe_aux_loss_coeff=moe_aux_loss_coeff,
)
self.output_mlp = nn.Linear(embed_dim, codebook_size)
# shared spatial-only PE (zeros in temporal tail)
pe_spatial = build_spatial_only_pe((H, W), patch_size, embed_dim, device='cpu', dtype=torch.float32) # [1,P,E]
self.register_buffer("pos_spatial_dec", pe_spatial, persistent=False)
# learnable mask token latent
# TODO; try leanable mask embedding in embed space instead of latent space
self.mask_token = nn.Parameter(torch.randn(1, 1, 1, latent_dim) * 0.02) # [1, 1, 1, L]
def forward(self, discrete_latents, training=True, conditioning=None, targets=None):
# discrete_latents: [B, T, P, L]
# targets: [B, T, P] indices
# conditioning: [B, T, A]
B, T, P, L = discrete_latents.shape
# convert latents to float for embedding
discrete_latents = discrete_latents.to(dtype=torch.float32)
# apply MaskGIT random masking during training
if training and self.training:
# per-batch mask ratio in [0.5, 1.0)
mask_ratio = 0.5 + torch.rand((), device=discrete_latents.device) * 0.5
mask_positions = (torch.rand(B, T, P, device=discrete_latents.device) < mask_ratio) # [B, T, P]
# guarantee at least one unmasked temporal anchor per (B, P)
# pick a random timestep for each (B,P) and force it to unmask
anchor_idx = torch.randint(0, T, (B, P), device=discrete_latents.device) # [B, P]
mask_positions[torch.arange(B)[:, None], anchor_idx, torch.arange(P)[None, :]] = False # [B, T, P]
# replace selected latents with mask tokens
mask_token = repeat(self.mask_token.to(discrete_latents.device, discrete_latents.dtype), '1 1 1 L -> B T P L', B=B, T=T, P=P) # [B, T, P, L]
discrete_latents = torch.where(mask_positions.unsqueeze(-1), mask_token, discrete_latents) # [B, T, P, L]
else:
mask_positions = None
embeddings = self.latent_embed(discrete_latents) # [B, T, P, E]
# add spatial PE (affects only first 2/3 of dimensions)
# STTransformer adds temporal PE to last 1/3 of dimensions
embeddings = embeddings + self.pos_spatial_dec.to(embeddings.device, embeddings.dtype)
transformed = self.transformer(embeddings, conditioning=conditioning) # [B, T, P, E]
# transform to logits for each token in codebook
predicted_logits = self.output_mlp(transformed) # [B, T, P, L^D]
# compute masked cross-entropy loss
loss = None
if training and self.training:
assert targets is not None, "target indices are needed for training"
Ld = predicted_logits.shape[-1] # L^D
logits_flat = predicted_logits.reshape(-1, Ld) # [(B*T*P), L^D]
targets_flat = targets.reshape(-1) # [(B*T*P)]
mask_flat = mask_positions.reshape(-1).to(torch.float32) # [(B*T*P)]
loss_per = nn.functional.cross_entropy(logits_flat, targets_flat, reduction='none') # [(B*T*P)]
denom = mask_flat.sum().clamp_min(1.0)
loss = (loss_per * mask_flat).sum() / denom
return predicted_logits, mask_positions, loss # logits, mask, optional loss
def exp_schedule_torch(self, t, T, P_total, k, device):
# t: current step, T: total steps, P_total: total masked positions across the horizon window
# exp schedule is P_total * (1 - exp(k * t / T)) / (1 - exp(k))
x = t / max(T, 1)
k_tensor = torch.tensor(k, device=device)
result = P_total * torch.expm1(k_tensor * x) / torch.expm1(k_tensor)
if t == T - 1:
return torch.tensor(P_total, dtype=result.dtype, device=device)
return result
@torch.no_grad()
def forward_inference(self, context_latents, prediction_horizon, num_steps, index_to_latents_fn, conditioning=None, schedule_k=5.0, temperature: float = 0.0):
# MaskGIT-style iterative decoding across all prediction horizon steps
# context_latents: [B, T_ctx, P, L]
# T_ctx=context timesteps, H=prediction horizon, K=codebook size
device = context_latents.device
dtype = context_latents.dtype
B, T_ctx, P, L = context_latents.shape # B, T_ctx, P, L
H = int(prediction_horizon) # number of horizon steps to decode
# append prediction_horizon masked frame latents to predict dynamics on
mask_latents = self.mask_token.to(device, dtype).expand(B, H, P, -1) # [B, H, P, L]
input_latents = torch.cat([context_latents, mask_latents], dim=1) # [B, T_ctx+H, P, L]
mask = torch.ones(B, H, P, 1, dtype=torch.bool, device=device) # [B, H, P, 1]
P_total = H * P # total masked positions across the horizon window
for m in range(num_steps):
n_tokens_raw = self.exp_schedule_torch(m, num_steps, P_total, schedule_k, device)
# predict logits for current input
logits, _, _ = self.forward(input_latents, training=False, conditioning=conditioning, targets=None) # [B, T_ctx+H, P, L^D]
# temperature scaling
if temperature and temperature > 0:
scaled_logits = logits / float(temperature)
else:
scaled_logits = logits
probs = torch.softmax(scaled_logits, dim=-1) # [B, T_ctx+H, P, L^D]
# confidence for unmask selection always from max probability
max_probs, _ = torch.max(probs, dim=-1) # [B, T_ctx+H, P]
# choose indices either via argmax (temperature==0) or sampling
if temperature and temperature > 0:
Bc, Tc, Pc, Ld = probs.shape # Bc=B, Tc=T_ctx+H, Pc=P, Ld=L^D
sampled = torch.distributions.Categorical(probs=probs.reshape(-1, Ld)).sample()
predicted_indices = sampled.view(Bc, Tc, Pc) # [B, T_ctx+H, P]
else:
_, predicted_indices = torch.max(probs, dim=-1) # [B, T_ctx+H, P]
horizon_probs = max_probs[:, -H:, :] # [B, H, P]
# for each batch element, select tokens to unmask from all masked positions
for b in range(B):
masked_mask_all = mask[b, :, :, 0] # [H, P]
masked_flat = masked_mask_all.view(-1) # [H*P]
masked_flat_idx = torch.where(masked_flat)[0] # [num_masked]
if masked_flat_idx.numel() == 0:
continue
# TODO: try to clean this
num_masked_b = int(masked_flat_idx.numel())
prev_b = P_total - num_masked_b
target_unmasked = int(torch.ceil(n_tokens_raw).item())
k_floor = max(P_total // 16, 1)
k_b = max(k_floor, min(max(target_unmasked - prev_b, 0), num_masked_b))
pos_probs_flat = horizon_probs[b].contiguous().view(-1)[masked_flat_idx] # [num_masked]
if pos_probs_flat.numel() > k_b:
top_idx = torch.topk(pos_probs_flat, k_b, largest=True).indices
sel_flat = masked_flat_idx[top_idx]
else:
sel_flat = masked_flat_idx
# map back to (h, p)
h_sel = torch.div(sel_flat, P, rounding_mode='floor') # [k_b]
p_sel = sel_flat % P # [k_b]
# group by unique h and write sampled tokens to input tensor
if h_sel.numel() > 0:
unique_h = torch.unique(h_sel, sorted=True)
for uh in unique_h:
mask_h = (h_sel == uh)
p_list = p_sel[mask_h]
if p_list.numel() == 0:
continue
t_abs = T_ctx + int(uh.item()) # absolute time index in [T_ctx, T_ctx+H-1] (for current horizon step)
idx_sel = predicted_indices[b:b+1, t_abs:t_abs+1, p_list] # [1,1,P_sel]
pred_latents_sel = index_to_latents_fn(idx_sel) # [1,1,P_sel,L]
input_latents[b:b+1, t_abs:t_abs+1, p_list] = pred_latents_sel
mask[b, int(uh.item()), p_list, 0] = False
# early exit if all horizon tokens are unmasked
if not mask[:, :, :, 0].any(): # mask: [B,H,P,1]
break
# final completion: fill any remaining masked tokens across all horizon steps via argmax
# TODO: try removing
if mask[:, :, :, 0].any():
logits, _, _ = self.forward(input_latents, training=False, conditioning=conditioning, targets=None) # [B, T_ctx+H, P, L^D]
if temperature and temperature > 0:
scaled_logits = logits / float(temperature)
else:
scaled_logits = logits
probs = torch.softmax(scaled_logits, dim=-1) # [B, T_ctx+H, P, L^D]
_, predicted_indices = torch.max(probs, dim=-1) # [B, T_ctx+H, P]
for b in range(B):
h_idx, p_idx = torch.where(mask[b, :, :, 0]) # both [N_remaining]
if h_idx.numel() == 0:
continue
unique_h = torch.unique(h_idx, sorted=True)
for uh in unique_h:
mask_h = (h_idx == uh)
p_list = p_idx[mask_h]
if p_list.numel() == 0:
continue
t_abs = T_ctx + int(uh.item()) # absolute time index
idx_sel = predicted_indices[b:b+1, t_abs:t_abs+1, p_list] # [1,1,P_sel]
pred_latents_sel = index_to_latents_fn(idx_sel) # [1,1,P_sel,L]
input_latents[b:b+1, t_abs:t_abs+1, p_list] = pred_latents_sel
mask[b, int(uh.item()), p_list, 0] = False
return input_latents # [B, T_ctx + H, P, L]
@property
def model_type(self) -> str:
return ModelType.DynamicsModel
```
## /models/fsq.py
```py path="/models/fsq.py"
# Finite Scalar Quantization from https://arxiv.org/pdf/2309.15505
# quantizes each dimension independently by bounding to 0, num_bins then rounding to nearest integer
# prevents token collapse and no auxiliary losses necessary
import torch
import torch.nn as nn
from einops import rearrange
class FiniteScalarQuantizer(nn.Module):
def __init__(self, latent_dim=5, num_bins=4):
super().__init__()
self.num_bins = num_bins # D
self.levels_np = torch.tensor(latent_dim * [num_bins])
self.codebook_size = num_bins**latent_dim # L^D
# fsq basis [L^0, L^1, ..., L^(L-1)] for converting between indices and latents
self.register_buffer('basis', (num_bins**torch.arange(latent_dim, dtype=torch.long)))
def scale_and_shift(self, z):
# scale and shift z from [-1, 1] to [0, num_bins - 1]
return 0.5 * (z + 1) * (self.num_bins - 1)
def unscale_and_unshift(self, z):
# unscale and unshift z from [0, num_bins - 1] to [-1, 1]
return 2 * z / (self.num_bins - 1) - 1
def forward(self, z):
# z: [B, T, P, L]
# apply 0.5 * (tanh(z) + 1) to go from z range to [0, num_bins - 1]
tanh_z = torch.tanh(z)
bounded_z = self.scale_and_shift(tanh_z)
# round to nearest integer
rounded_z = torch.round(bounded_z)
# stopgrad for straight-through gradient bypassing quantization
quantized_z = bounded_z + (rounded_z - bounded_z).detach()
# normalize back to [-1, 1]
quantized_z = self.unscale_and_unshift(quantized_z)
return quantized_z
def get_codebook_usage(self, quantized_z):
unique_bins = torch.unique(quantized_z).shape[0]
return unique_bins / self.num_bins
def get_indices_from_latents(self, latents, dim=-1):
# to get fsq indices, for each dimension, we get the index and add it (so multiply by L then sum)
# for each dimension of each latent, get sum of (value * L^current_dim) along latent dim which is the index of that latent in the codebook
# codebook size = L^latent_dim
# latents: [*, L]
# go from [-1, 1] to [0, num_bins - 1] in each dimension
digits = torch.round(self.scale_and_shift(latents)).clamp(0, self.num_bins-1)
# get indices for each latent by summing (value * L^current_dim_idx) along latent dim
indices = torch.sum(digits * self.basis.to(latents.device), dim=dim).long() # [*]
return indices
def get_latents_from_indices(self, indices, dim=-1):
# indices: [*]
# recover each entry of latent in range [0, num_bins - 1] by repeatedly dividing by L^current_dim and taking mod
digits = (indices.unsqueeze(-1) // self.basis) % self.num_bins # [*, L]
# go from [0, num_bins - 1] to [-1, 1] in each dimension
latents = self.unscale_and_unshift(digits) # [*, L]
return latents
```
## /models/latent_actions.py
```py path="/models/latent_actions.py"
from models.utils import ModelType
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
import math
from einops import rearrange, repeat, reduce
from models.st_transformer import STTransformer, PatchEmbedding
from models.fsq import FiniteScalarQuantizer
NUM_LATENT_ACTIONS_BINS = 2
class LatentActionsEncoder(nn.Module):
def __init__(self, frame_size=(128, 128), patch_size=8, embed_dim=128, num_heads=8,
hidden_dim=256, num_blocks=4, action_dim=3):
super().__init__()
self.patch_embed = PatchEmbedding(frame_size, patch_size, embed_dim)
self.transformer = STTransformer(embed_dim, num_heads, hidden_dim, num_blocks, causal=True)
# embeddings to discrete latent bottleneck actions
self.action_head = nn.Sequential(
nn.LayerNorm(embed_dim * 2),
nn.Linear(embed_dim * 2, 4 * action_dim),
nn.GELU(),
nn.Linear(4 * action_dim, action_dim)
)
def forward(self, frames):
# frames: [B, T, C, H, W]
batch_size, seq_len, C, H, W = frames.shape
embeddings = self.patch_embed(frames) # [B, T, P, E]
transformed = self.transformer(embeddings)
# TODO: try attention pooling + mean instead of mean + concat
# mean pool over patches (since one action per frame)
pooled = transformed.mean(dim=2) # [B, T, E]
# combine features from current and next frame
actions = []
for t in range(seq_len - 1):
# concat current and next frame features
combined = torch.cat([pooled[:, t], pooled[:, t+1]], dim=1) # [B, E*2]
action = self.action_head(combined) # [B, A]
actions.append(action)
actions = torch.stack(actions, dim=1) # [B, T-1, A]
return actions
class LatentActionsDecoder(nn.Module):
def __init__(self, frame_size=(128, 128), patch_size=8, embed_dim=128, num_heads=8,
hidden_dim=256, num_blocks=4, conditioning_dim=3):
super().__init__()
self.patch_embed = PatchEmbedding(frame_size, patch_size, embed_dim)
self.transformer = STTransformer(embed_dim, num_heads, hidden_dim, num_blocks, causal=True, conditioning_dim=conditioning_dim)
# embeddings to mixed frame output patches
self.frame_head = nn.Sequential(
nn.LayerNorm(embed_dim),
nn.Linear(embed_dim, 3 * patch_size * patch_size),
nn.Tanh()
)
self.frame_size = frame_size
self.patch_size = patch_size
self.num_patches = (frame_size[0] // patch_size) * (frame_size[1] // patch_size)
self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, embed_dim))
def forward(self, frames, actions, training=True):
# frames: [B, T, C, H, W]
# actions: [B, T - 1, A]
B, T, C, H, W = frames.shape
frames = frames[:, :-1] # [B, T-1, C, H, W]
video_embeddings = self.patch_embed(frames) # [B, T-1, P, E]
_, _, P, E = video_embeddings.shape
# mask certain tokens from all frames except first frame
# this strongly forces actions to contain most useful info (I recommend to keep based on experiments)
if training and self.training:
keep_rate = 0.0
keep = (torch.rand(B, T-1, P, 1, device=frames.device) < keep_rate)
keep[:, 0] = 1 # never mask first frame tokens (anchor) TODO: try rid of ablation
video_embeddings = torch.where(
keep, video_embeddings,
self.mask_token.to(video_embeddings.dtype).expand_as(video_embeddings)
)
transformed = self.transformer(video_embeddings, conditioning=actions) # [B, T-1, P, E]
patches = self.frame_head(transformed) # [B, T-1, P, 3 * S * S]
patches = rearrange(
patches, 'b t p (c p1 p2) -> b t c p p1 p2', c=3, p1=self.patch_size, p2=self.patch_size
) # [B, T-1, C, P, S, S]
pred_frames = rearrange(
patches, 'b t c (h w) p1 p2 -> b t c (h p1) (w p2)', h=H//self.patch_size, w=W//self.patch_size
) # [B, T-1, C, H, W]
return pred_frames # [B, T-1, C, H, W]
class LatentActionModel(nn.Module):
def __init__(self, frame_size=(128, 128), n_actions=8, patch_size=8, embed_dim=128,
num_heads=8, hidden_dim=256, num_blocks=4):
super().__init__()
assert math.log(n_actions, NUM_LATENT_ACTIONS_BINS).is_integer(), f"n_actions must be a power of {NUM_LATENT_ACTIONS_BINS}"
self.action_dim=int(math.log(n_actions, NUM_LATENT_ACTIONS_BINS))
self.encoder = LatentActionsEncoder(frame_size, patch_size, embed_dim, num_heads, hidden_dim, num_blocks, action_dim=self.action_dim)
self.quantizer = FiniteScalarQuantizer(latent_dim=self.action_dim, num_bins=NUM_LATENT_ACTIONS_BINS)
self.decoder = LatentActionsDecoder(frame_size, patch_size, embed_dim, num_heads, hidden_dim, num_blocks, conditioning_dim=self.action_dim)
self.var_target = 0.01
self.var_lambda = 100.0
def forward(self, frames):
# frames: [B, T, C, H, W]
# get quantized action latents
action_latents = self.encoder(frames) # [B, T - 1, A]
action_latents_quantized = self.quantizer(action_latents) # [B, T - 1, A]
# decode to get predicted frames
pred_frames = self.decoder(frames, action_latents_quantized, training=True) # [B, T - 1, C, H, W]
# reconstruction loss
target_frames = frames[:, 1:] # All frames except first [B, T - 1, C, H, W]
recon_loss = F.smooth_l1_loss(pred_frames, target_frames)
# variance loss across batch dim for pre-quant encoder outputs (helps prevent action collapse)
z_var = action_latents.var(dim=0, unbiased=False).mean()
var_penalty = F.relu(self.var_target - z_var)
total_loss = recon_loss + self.var_lambda * var_penalty
return total_loss, pred_frames
def encode(self, frames):
action_latents = self.encoder(frames) # [B, T, A]
action_latents_quantized = self.quantizer(action_latents) # [B, T, A]
return action_latents_quantized
@property
def model_type(self) -> str:
return ModelType.LatentActionModel
```
## /models/muon.py
```py path="/models/muon.py"
from __future__ import annotations
import torch
import torch.distributed as dist
from torch import Tensor
def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor:
a, b, c = (3.4445, -4.7750, 2.0315)
X = G.clone().bfloat16() # clone to avoid corrupting p.grad when G is already bf16
X /= X.norm() + eps
transposed = G.size(0) > G.size(1)
if transposed:
X = X.T
for _ in range(steps):
A = X @ X.T
B = b * A + c * A @ A
X = a * X + B @ X
return X.T if transposed else X
class Muon(torch.optim.Optimizer):
def __init__(self, params, lr: float, momentum: float = 0.95,
backend_steps: int = 5, nesterov: bool = True,
weight_decay: float = 0.0):
super().__init__(
params,
dict(lr=lr, momentum=momentum, backend_steps=backend_steps,
nesterov=nesterov, weight_decay=weight_decay),
)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
distributed = dist.is_available() and dist.is_initialized()
world_size = dist.get_world_size() if distributed else 1
rank = dist.get_rank() if distributed else 0
for group in self.param_groups:
params = group["params"]
if not params:
continue
lr = group["lr"]
momentum = group["momentum"]
backend_steps = group["backend_steps"]
nesterov = group["nesterov"]
wd = group["weight_decay"]
# decoupled weight decay
if wd > 0:
for p in params:
p.data.mul_(1 - lr * wd)
total_params = sum(int(p.numel()) for p in params)
updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16)
curr = 0
for i, p in enumerate(params):
if i % world_size == rank and p.grad is not None:
g = p.grad
state = self.state[p]
if "momentum_buffer" not in state:
state["momentum_buffer"] = torch.zeros_like(g)
buf = state["momentum_buffer"]
buf.mul_(momentum).add_(g)
if nesterov:
g = g.add(buf, alpha=momentum)
g = zeropower_via_newtonschulz5(g, steps=backend_steps)
g *= max(1, g.size(0) / g.size(1)) ** 0.5
updates_flat[curr : curr + p.numel()] = g.reshape(-1)
curr += p.numel()
if distributed:
dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM)
curr = 0
for p in params:
g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype)
p.add_(g, alpha=-lr)
curr += p.numel()
return loss
```
## /models/norms.py
```py path="/models/norms.py"
import torch
import torch.nn as nn
from einops import repeat
class RMSNorm(nn.Module):
def __init__(self, embed_dim, eps=1e-5):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(embed_dim)) # learned scale for rmsnorm (gamma)
def forward(self, x):
# root mean squared norm of x = x / sqrt(mean(x^2) + eps)
mean_squared = torch.mean(x**2, dim=-1, keepdim=True)
# torch.rsqrt is 1 / sqrt (faster and numerically stable vs manual 1 / sqrt)
rms_normed = x * torch.rsqrt(mean_squared + self.eps)
return rms_normed * self.weight
class SimpleLayerNorm(nn.Module):
def __init__(self, embed_dim, eps=1e-5):
super().__init__()
self.eps = eps
def forward(self, x):
# only center at 0 and make stddev 1 (for film)
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, unbiased=False, keepdim=True)
return (x - mean) * torch.rsqrt(var + self.eps)
class AdaptiveNormalizer(nn.Module):
# either conditioned FiLM or unconditioned RMSNorm
def __init__(self, embed_dim, conditioning_dim=None):
super().__init__()
self.ln = None
self.rms = None
self.to_gamma_beta = None
if conditioning_dim is None:
# RMSNorm when unconditioned (can do ln but this is better)
self.rms = RMSNorm(embed_dim)
else:
self.ln = SimpleLayerNorm(embed_dim)
self.to_gamma_beta = nn.Sequential(
nn.SiLU(),
nn.Linear(conditioning_dim, 2 * embed_dim)
)
# for lam and dynamics we initialize with small non-zero weights
# so conditioning is active from step 1 (helps prevent ignoring conditioning)
nn.init.normal_(self.to_gamma_beta[-1].weight, mean=0.0, std=1e-3)
nn.init.zeros_(self.to_gamma_beta[-1].bias)
def forward(self, x, conditioning=None):
# x: [B, T, P, E]
# conditioning: [B, T, C] or [B, T - 1, C]
if self.to_gamma_beta is None or conditioning is None:
normed = self.rms(x) if self.rms is not None else self.ln(x)
return normed
x = self.ln(x)
B, T, P, E = x.shape
out = self.to_gamma_beta(conditioning) # [B, T, 2 * E]
out = repeat(out, 'b t twoe -> b t p twoe', p=P) # [B, T, P, 2 * E]
gamma, beta = out.chunk(2, dim=-1) # each [B, T, P, E]
# preppend action tensor with zeros in T since a_t-1 should impact z_t (and a_0 is used for z_1 etc)
if gamma.shape[1] == x.shape[1] - 1 and beta.shape[1] == x.shape[1] - 1:
gamma = torch.cat([torch.zeros_like(gamma[:, :1]), gamma], dim=1)
beta = torch.cat([torch.zeros_like(beta[:, :1]), beta], dim=1)
assert gamma.shape[1] == x.shape[1], f"gamma shape: {gamma.shape} != x shape: {x.shape}"
assert beta.shape[1] == x.shape[1], f"beta shape: {beta.shape} != x shape: {x.shape}"
x = x * (1 + gamma) + beta # [B, T, P, E]
return x
```
## /models/patch_embed.py
```py path="/models/patch_embed.py"
import torch
import torch.nn as nn
from einops import rearrange
from models.positional_encoding import build_spatial_only_pe
class PatchEmbedding(nn.Module):
def __init__(self, frame_size=(128, 128), patch_size=8, embed_dim=128):
super().__init__()
H, W = frame_size
self.frame_size = frame_size
self.patch_size = patch_size
self.embed_dim = embed_dim
self.Hp, self.Wp = H // patch_size, W // patch_size
self.num_patches = self.Hp * self.Wp
# split embed dim into thirds for spatial x, spatial y, and temporal
base_split = (embed_dim // 3) & ~1
remaining_dim = embed_dim - base_split
self.spatial_x_dim = (remaining_dim // 2) & ~1
self.spatial_y_dim = remaining_dim - self.spatial_x_dim
self.temporal_dim = base_split
# ensure the embed dim is split wholy into thirds and each third is even
assert (self.spatial_x_dim + self.spatial_y_dim + self.temporal_dim) == embed_dim, \
f"Dimension mismatch: {self.spatial_x_dim} + {self.spatial_y_dim} + {self.temporal_dim} != {embed_dim}"
assert self.spatial_x_dim % 2 == 0 and self.spatial_y_dim % 2 == 0 and self.temporal_dim % 2 == 0, \
f"Embed dim x={self.spatial_x_dim}, y={self.spatial_y_dim}, t={self.temporal_dim}"
pe_spatial = build_spatial_only_pe(self.frame_size, self.patch_size, self.embed_dim, device='cpu', dtype=torch.float32) # [1,P,E]
self.register_buffer("pos_spatial", pe_spatial, persistent=False)
# pixel patches to embeddings
self.proj = nn.Conv2d(3 * self.patch_size * self.patch_size, self.embed_dim, 1)
def forward(self, frames):
B, T, C, H, W = frames.shape
# go from frames to patches
x = rearrange(frames, 'b t c (hp p1) (wp p2) -> (b t) (c p1 p2) hp wp', p1=self.patch_size, p2=self.patch_size) # [(B*T), 3*p*p, Hp, Wp]
x = self.proj(x) # [(B*T), E, Hp, Wp]
x = rearrange(x, '(b t) e hp wp -> b t (hp wp) e', b=B, t=T) # [B, T, P, E]
# add 2d spatial pos encoding (first 2/3 of embed dim)
x = x + self.pos_spatial.to(dtype=x.dtype, device=x.device) # [B, T, P, E]
return x
```
## /models/positional_encoding.py
```py path="/models/positional_encoding.py"
import torch
from einops import rearrange, repeat
# TODO: Try RoPE / AliBi
def sincos_1d(L, D, device, dtype):
# 1d sinusoidal position encoding where element j of ith patch embedding is encoded as:
# PE[i, 2j] = sin(i / 10000^(2j/D)) # even indices
# PE[i, 2j+1] = cos(i / 10000^(2j/D)) # odd indices
assert D % 2 == 0, "Encoding dimension must be even"
# position indices [L, 1] and dimension indices [1, D/2]
pos = rearrange(torch.arange(L, device=device, dtype=dtype), 'l -> l 1') # [L,1]
i = rearrange(torch.arange(D // 2, device=device, dtype=dtype), 'd -> 1 d')# [1,D/2]
# angular frequencies: 1/10000^(2i/D) for each dimension
div = torch.pow(torch.tensor(10000.0, device=device, dtype=dtype), (2*i)/D)
# angles: pos * freq for each position-dimension pair
angles = pos / div # [L, D/2] (broadcasted together)
pe = torch.zeros(L, D, device=device, dtype=dtype)
pe[:, 0::2] = torch.sin(angles) # even indices
pe[:, 1::2] = torch.cos(angles) # odd indices
return pe # [L, D]
def sincos_time(T, D, device, dtype):
# temporal PE (1d sinusoidal PE across time)
return sincos_1d(T, D, device, dtype) # reuse the same 1D builder
def build_spatial_only_pe(frame_size, patch_size, embed_dim, device='cpu', dtype=torch.float32):
# spatial positional encodings for a grid of patches in first 2/3 of embed dim (evenly into x and y axes)
# last 1/3 for temporal PE padded with 0s
H, W = frame_size
Hp, Wp = H // patch_size, W // patch_size
# split dimensions (ensure temporal even)
temporal_dim = (embed_dim // 3) & ~1
spatial_dims = embed_dim - temporal_dim
# split spatial dims between x and y (ensure both even)
spatial_x_dim = (spatial_dims // 2) & ~1
spatial_y_dim = spatial_dims - spatial_x_dim
assert spatial_x_dim % 2 == 0 and spatial_y_dim % 2 == 0 and temporal_dim % 2 == 0
# 2d PE for x and y axes
pe_x = sincos_1d(Wp, spatial_x_dim, device, dtype) # [Wp, Dx]
pe_y = sincos_1d(Hp, spatial_y_dim, device, dtype) # [Hp, Dy]
pe_x = repeat(pe_x, 'wp dx -> hp wp dx', hp=Hp) # [Hp, Wp, Dx]
pe_y = repeat(pe_y, 'hp dy -> hp wp dy', wp=Wp) # [Hp, Wp, Dy]
pe_spatial = torch.cat([
pe_x,
pe_y,
torch.zeros(Hp, Wp, temporal_dim, device=device, dtype=dtype) # zero temporal tail
], dim=-1) # [Hp, Wp, E]
pe_spatial = rearrange(pe_spatial, 'hp wp e -> 1 (hp wp) e') # [1, P, E]
return pe_spatial # [1, P, E]
```
## /models/st_transformer.py
```py path="/models/st_transformer.py"
import torch
import torch.nn as nn
from einops import rearrange
from models.positional_encoding import build_spatial_only_pe, sincos_time
from models.norms import AdaptiveNormalizer
from models.patch_embed import PatchEmbedding
import math
import torch.nn.functional as F
class SpatialAttention(nn.Module):
def __init__(self, embed_dim, num_heads, conditioning_dim=None):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
assert self.head_dim * num_heads == embed_dim, f"embed dim must be divisible by num heads"
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
self.norm = AdaptiveNormalizer(embed_dim, conditioning_dim)
def forward(self, x, conditioning=None):
B, T, P, E = x.shape
# project to Q, K, V and split into heads: [B, T, P, E] -> [(B*T), H, P, E/H]
# (4 dims to work with torch compile attention)
q = rearrange(self.q_proj(x), 'B T P (H D) -> (B T) H P D', H=self.num_heads)
k = rearrange(self.k_proj(x), 'B T P (H D) -> (B T) H P D', H=self.num_heads)
v = rearrange(self.v_proj(x), 'B T P (H D) -> (B T) H P D', H=self.num_heads)
k_t = k.transpose(-2, -1) # [(B*T), H, P, D, P]
# attention(q, k, v) = softmax(qk^T / sqrt(d)) v
scores = torch.matmul(q, k_t) / math.sqrt(self.head_dim) # [(B*T), H, P, P]
attn_weights = F.softmax(scores, dim=-1) # [(B*T), H, P, P]
attn_output = torch.matmul(attn_weights, v) # [(B*T), H, P, D]
attn_output = rearrange(attn_output, '(B T) H P D -> B T P (H D)', B=B, T=T) # [B, T, P, E]
# out proj to mix head information
attn_out = self.out_proj(attn_output) # [B, T, P, E]
# residual and optionally conditioned norm
out = self.norm(x + attn_out, conditioning) # [B, T, P, E]
return out # [B, T, P, E]
class TemporalAttention(nn.Module):
def __init__(self, embed_dim, num_heads, causal=True, conditioning_dim=None):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
assert self.head_dim * num_heads == embed_dim
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
self.norm = AdaptiveNormalizer(embed_dim, conditioning_dim)
self.causal = causal
def forward(self, x, conditioning=None):
B, T, P, E = x.shape
# project to Q, K, V and split into heads: [B, T, P, E] -> [(B*P), H, T, D]
# (4 dims to work with torch compile attention)
q = rearrange(self.q_proj(x), 'b t p (h d) -> (b p) h t d', h=self.num_heads)
k = rearrange(self.k_proj(x), 'b t p (h d) -> (b p) h t d', h=self.num_heads)
v = rearrange(self.v_proj(x), 'b t p (h d) -> (b p) h t d', h=self.num_heads) # [B, P, H, T, D]
k_t = k.transpose(-2, -1) # [(B*P), H, T, D, T]
# attention(q, k, v) = softmax(qk^T / sqrt(d)) v
scores = torch.matmul(q, k_t) / math.sqrt(self.head_dim) # [(B*P), H, T, T]
# causal mask for each token t in seq, mask out all tokens to the right of t (after t)
if self.causal:
mask = torch.triu(torch.ones(T, T), diagonal=1).bool().to(x.device)
scores = scores.masked_fill(mask, -torch.inf) # [(B*P), H, T, T]
attn_weights = F.softmax(scores, dim=-1) # [(B*P), H, T, T]
attn_output = torch.matmul(attn_weights, v) # [(B*P), H, T, D]
attn_output = rearrange(attn_output, '(b p) h t d -> b t p (h d)', b=B, p=P) # [B, T, P, E]
# out proj to mix head information
attn_out = self.out_proj(attn_output) # [B, T, P, E]
# residual and optionally conditioned norm
out = self.norm(x + attn_out, conditioning) # [B, T, P, E]
return out # [B, T, P, E]
class SwiGLUFFN(nn.Module):
# swiglu(x) = W3(sigmoid(W1(x) + b1) * (W2(x) + b2)) + b3
def __init__(self, embed_dim, hidden_dim, conditioning_dim=None):
super().__init__()
h = math.floor(2 * hidden_dim / 3)
self.w_v = nn.Linear(embed_dim, h)
self.w_g = nn.Linear(embed_dim, h)
self.w_o = nn.Linear(h, embed_dim)
self.norm = AdaptiveNormalizer(embed_dim, conditioning_dim)
def forward(self, x, conditioning=None):
v = F.silu(self.w_v(x)) # [B, T, P, h]
g = self.w_g(x) # [B, T, P, h]
out = self.w_o(v * g) # [B, T, P, E]
return self.norm(x + out, conditioning) # [B, T, P, E]
class SwiGLUExpert(nn.Module):
def __init__(self, embed_dim, hidden_dim):
super().__init__()
h = math.floor(2 * hidden_dim / 3)
self.w_v = nn.Linear(embed_dim, h)
self.w_g = nn.Linear(embed_dim, h)
self.w_o = nn.Linear(h, embed_dim)
def forward(self, x):
return self.w_o(F.silu(self.w_v(x)) * self.w_g(x))
class MoESwiGLUFFN(nn.Module):
def __init__(self, embed_dim, hidden_dim, num_experts=4, top_k=2,
aux_loss_coeff=0.01, conditioning_dim=None):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.aux_loss_coeff = aux_loss_coeff
self.router = nn.Linear(embed_dim, num_experts, bias=False)
self.experts = nn.ModuleList([
SwiGLUExpert(embed_dim, hidden_dim) for _ in range(num_experts)
])
self.norm = AdaptiveNormalizer(embed_dim, conditioning_dim)
self._aux_loss = None
self._expert_counts = None # per-expert token fractions from last forward
@property
def aux_loss(self):
if self._aux_loss is None:
device = next(self.parameters()).device
return torch.zeros((), device=device)
return self._aux_loss
@property
def expert_utilization(self):
return self._expert_counts
def forward(self, x, conditioning=None):
# x: [B, T, P, E]
B, T, P, E = x.shape
residual = x
# flatten spatial dims for routing: [B*T*P, E]
flat = x.reshape(-1, E)
N = flat.shape[0]
# route tokens to top-k experts
logits = self.router(flat) # [N, num_experts]
top_k_logits, top_k_indices = logits.topk(self.top_k, dim=-1) # [N, top_k]
top_k_weights = F.softmax(top_k_logits, dim=-1) # [N, top_k]
# load-balancing auxiliary loss
if self.training:
router_probs = F.softmax(logits, dim=-1) # [N, num_experts]
tokens_per_expert = torch.zeros(self.num_experts, device=x.device)
for k in range(self.top_k):
tokens_per_expert.scatter_add_(
0, top_k_indices[:, k],
torch.ones(N, device=x.device),
)
fraction_dispatched = tokens_per_expert / (N * self.top_k)
fraction_probs = router_probs.mean(dim=0)
self._aux_loss = self.aux_loss_coeff * self.num_experts * (
fraction_dispatched * fraction_probs
).sum()
self._expert_counts = fraction_dispatched.detach()
# compute expert outputs
output = torch.zeros_like(flat)
for k in range(self.top_k):
expert_idx = top_k_indices[:, k] # [N]
weight = top_k_weights[:, k].unsqueeze(-1) # [N, 1]
for e in range(self.num_experts):
mask = (expert_idx == e)
if not mask.any():
continue
expert_input = flat[mask]
expert_output = self.experts[e](expert_input)
output[mask] += weight[mask] * expert_output
# reshape back and apply residual + norm
out = output.reshape(B, T, P, E)
return self.norm(residual + out, conditioning) # [B, T, P, E]
class STTransformerBlock(nn.Module):
def __init__(self, embed_dim, num_heads, hidden_dim, causal=True, conditioning_dim=None,
use_moe=False, num_experts=4, top_k_experts=2, moe_aux_loss_coeff=0.01):
super().__init__()
self.spatial_attn = SpatialAttention(embed_dim, num_heads, conditioning_dim)
self.temporal_attn = TemporalAttention(embed_dim, num_heads, causal, conditioning_dim)
if use_moe:
self.ffn = MoESwiGLUFFN(
embed_dim, hidden_dim,
num_experts=num_experts, top_k=top_k_experts,
aux_loss_coeff=moe_aux_loss_coeff,
conditioning_dim=conditioning_dim,
)
else:
self.ffn = SwiGLUFFN(embed_dim, hidden_dim, conditioning_dim)
def forward(self, x, conditioning=None):
# x: [B, T, P, E]
# out: [B, T, P, E]
x = self.spatial_attn(x, conditioning)
x = self.temporal_attn(x, conditioning)
x = self.ffn(x, conditioning)
return x
class STTransformer(nn.Module):
def __init__(self, embed_dim, num_heads, hidden_dim, num_blocks, causal=True, conditioning_dim=None,
use_moe=False, num_experts=4, top_k_experts=2, moe_aux_loss_coeff=0.01):
super().__init__()
# calculate temporal PE dim
self.temporal_dim = (embed_dim // 3) & ~1 # round down to even number
self.spatial_dims = embed_dim - self.temporal_dim # rest goes to spatial
self.blocks = nn.ModuleList([
STTransformerBlock(
embed_dim, num_heads, hidden_dim, causal, conditioning_dim,
use_moe=use_moe, num_experts=num_experts,
top_k_experts=top_k_experts, moe_aux_loss_coeff=moe_aux_loss_coeff,
)
for _ in range(num_blocks)
])
def forward(self, x, conditioning=None):
# x: [B, T, P, E]
# conditioning: [B, T, E]
B, T, P, E = x.shape
tpe = sincos_time(T, self.temporal_dim, x.device, x.dtype) # [T, E/3]
# temporal PE (pad with 0s for first 2/3s spatial PE, last 1/3 temporal PE)
tpe_padded = torch.cat([
torch.zeros(T, self.spatial_dims, device=x.device, dtype=x.dtype),
tpe
], dim=-1) # [T, E]
x = x + tpe_padded[None, :, None, :] # [B,T,P,E]
# apply transformer blocks
for block in self.blocks:
x = block(x, conditioning)
return x
def moe_aux_loss(self):
device = next(self.parameters()).device
total = torch.zeros((), device=device)
for block in self.blocks:
if isinstance(block.ffn, MoESwiGLUFFN):
total = total + block.ffn.aux_loss
return total
def moe_expert_utilization(self):
"""Per-block expert token fractions. Returns dict of block_idx -> [num_experts] tensor."""
util = {}
for idx, block in enumerate(self.blocks):
if isinstance(block.ffn, MoESwiGLUFFN) and block.ffn.expert_utilization is not None:
util[idx] = block.ffn.expert_utilization
return util
```
## /models/utils.py
```py path="/models/utils.py"
from enum import Enum
class ModelType(str, Enum):
VideoTokenizer: str = 'VideoTokenizer'
LatentActionModel: str = 'LatentAction'
DynamicsModel: str = 'Dynamic'
```
## /models/video_tokenizer.py
```py path="/models/video_tokenizer.py"
from models.utils import ModelType
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
from models.st_transformer import STTransformer
from models.fsq import FiniteScalarQuantizer
from models.patch_embed import PatchEmbedding
from models.positional_encoding import build_spatial_only_pe
class VideoTokenizerEncoder(nn.Module):
def __init__(self, frame_size=(128, 128), patch_size=8, embed_dim=128, num_heads=8,
hidden_dim=256, num_blocks=4, latent_dim=5):
super().__init__()
self.patch_embed = PatchEmbedding(frame_size, patch_size, embed_dim)
self.transformer = STTransformer(embed_dim, num_heads, hidden_dim, num_blocks, causal=True)
self.latent_head = nn.Sequential(
nn.LayerNorm(embed_dim),
nn.Linear(embed_dim, latent_dim)
)
def forward(self, frames):
# frames: [B, T, C, H, W]
# frames to patch embeddings, pass through transformer, project to latent dim
embeddings = self.patch_embed(frames) # [B, T, P, E]
transformed = self.transformer(embeddings) # [B, T, P, E]
predicted_latents = self.latent_head(transformed) # [B, T, P, L]
return predicted_latents
class PixelShuffleFrameHead(nn.Module):
# conv2D embeddings to pixels head
def __init__(self, embed_dim, patch_size=8, channels=3, H=128, W=128):
super().__init__()
self.patch_size = patch_size
self.Hp, self.Wp = H // patch_size, W // patch_size
self.to_pixels = nn.Conv2d(embed_dim, channels * (patch_size ** 2), kernel_size=1)
def forward(self, tokens): # [B, T, P, E]
B, T, P, E = tokens.shape
x = rearrange(tokens, 'b t (hp wp) e -> (b t) e hp wp', hp=self.Hp, wp=self.Wp) # [(B*T), E, Hp, Wp]
x = self.to_pixels(x) # [(B*T), C*p^2, Hp, Wp]
x = rearrange(x, '(b t) (c p1 p2) hp wp -> b t c (hp p1) (wp p2)', p1=self.patch_size, p2=self.patch_size, b=B, t=T) # [B, T, C, H, W]
return x
class VideoTokenizerDecoder(nn.Module):
def __init__(self, frame_size=(128, 128), patch_size=8, embed_dim=128, num_heads=8,
hidden_dim=256, num_blocks=4, latent_dim=5):
super().__init__()
H, W = frame_size
self.patch_size = patch_size
self.Hp, self.Wp = H // patch_size, W // patch_size
self.num_patches = self.Hp * self.Wp
self.latent_embed = nn.Linear(latent_dim, embed_dim)
self.transformer = STTransformer(embed_dim, num_heads, hidden_dim, num_blocks, causal=True)
self.frame_head = PixelShuffleFrameHead(embed_dim, patch_size=patch_size, channels=3, H=H, W=W)
# first 2/3 spatial PE (temporal is last 1/3)
pe_spatial_dec = build_spatial_only_pe((H, W), self.patch_size, embed_dim, device='cpu', dtype=torch.float32) # [1,P,E]
self.register_buffer("pos_spatial_dec", pe_spatial_dec, persistent=False)
def forward(self, latents):
# latents: [B, T, P, L]
# embed latents and add spatial PE
embedding = self.latent_embed(latents) # [B, T, P, E]
embedding = embedding + self.pos_spatial_dec.to(dtype=embedding.dtype, device=embedding.device)
# apply transformer (temporal PE added inside)
embedding = self.transformer(embedding) # [B, T, P, E]
# reconstruct frames using patch-wise head
frames_out = self.frame_head(embedding) # [B, T, C, H, W]
return frames_out
class VideoTokenizer(nn.Module):
def __init__(self, frame_size=(128, 128), patch_size=8, embed_dim=128, num_heads=8,
hidden_dim=256, num_blocks=4, latent_dim=3, num_bins=4):
super().__init__()
self.encoder = VideoTokenizerEncoder(frame_size, patch_size, embed_dim, num_heads, hidden_dim, num_blocks, latent_dim)
self.decoder = VideoTokenizerDecoder(frame_size, patch_size, embed_dim, num_heads, hidden_dim, num_blocks, latent_dim)
self.quantizer = FiniteScalarQuantizer(latent_dim, num_bins)
self.codebook_size = num_bins**latent_dim
def forward(self, frames):
# encode frames to latent representations, quantize, and decode back to frames
embeddings = self.encoder(frames) # [B, T, P, L]
quantized_z = self.quantizer(embeddings)
x_hat = self.decoder(quantized_z) # [B, T, C, H, W]
recon_loss = F.smooth_l1_loss(x_hat, frames)
return recon_loss, x_hat
def tokenize(self, frames):
# encode frames to latent representations, quantize, and return indices
embeddings = self.encoder(frames) # [B, T, P, L]
quantized_z = self.quantizer(embeddings)
indices = self.quantizer.get_indices_from_latents(quantized_z, dim=-1)
return indices
def detokenize(self, quantized_z):
# decode quantized latents back to frames
x_hat = self.decoder(quantized_z) # [B, T, C, H, W]
return x_hat
@property
def model_type(self) -> str:
return ModelType.VideoTokenizer
```
## /requirements.txt
# Weights & Biases integration requirements
wandb>=0.15.0
torch>=2.8.0
torchvision>=0.10.0
numpy>=1.21.0
matplotlib>=3.3.0
tqdm>=4.62.0
einops>=0.3.0
h5py>=3.1.0
opencv-python>=4.5.0
omegaconf
huggingface_hub>=0.23.0
## /scripts/download_assets.py
```py path="/scripts/download_assets.py"
import argparse
import sys
from datetime import datetime
from pathlib import Path
from typing import List, Tuple
from huggingface_hub import hf_hub_download, list_repo_files
DATASET_REPO_ID_DEFAULT = "AlmondGod/tinyworlds"
MODELS_REPO_ID_DEFAULT = "AlmondGod/tinyworlds-models"
VALID_TYPES = {"video_tokenizer", "actions", "dynamics", "all"}
def repo_root() -> Path:
return Path(__file__).resolve().parents[1]
def expand_patterns(repo_id: str, patterns: List[str], repo_type: str) -> List[str]:
if not patterns:
return []
files = list_repo_files(repo_id=repo_id, repo_type=repo_type)
matched: List[str] = []
from fnmatch import fnmatch
for f in files:
if any(fnmatch(f, pat) for pat in patterns):
matched.append(f)
return matched
def download_pairs(pairs: List[Tuple[str, str]], output_dir: Path, resume: bool = True, repo_type: str = "model") -> List[Path]:
output_dir.mkdir(parents=True, exist_ok=True)
paths: List[Path] = []
for repo_id, filename in pairs:
p = hf_hub_download(
repo_id=repo_id,
filename=filename,
repo_type=repo_type,
local_dir=output_dir,
local_dir_use_symlinks=False,
resume_download=resume,
)
paths.append(Path(p))
return paths
def cmd_datasets(args) -> int:
# Defaults to datasets repo and repo_root/data
repo_id = args.repo or DATASET_REPO_ID_DEFAULT
data_dir = args.out or (repo_root() / "data")
patterns = args.pattern or ["*.h5", "*.json", "*.md", "assets/*"]
matched = expand_patterns(repo_id, patterns, repo_type="dataset")
if not matched:
print(f"No files matched in dataset repo {repo_id} for patterns {patterns}")
return 1
pairs = [(repo_id, m) for m in matched]
paths = download_pairs(pairs, data_dir, resume=(not args.no_resume), repo_type="dataset")
for p in paths:
print(p)
return 0
def cmd_models(args) -> int:
# Defaults to model repo and results/<timestamp>_<suite>/
repo_id = args.repo or MODELS_REPO_ID_DEFAULT
base_results = args.out or (repo_root() / "results")
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
suite = args.suite_name or "models"
folder_name = f"{timestamp}_{suite}"
model_root = base_results / folder_name
# Validate and normalize types to fetch
type_arg = args.type
if type_arg == "all":
types_to_fetch = list(VALID_TYPES)
elif type_arg in VALID_TYPES:
types_to_fetch = [type_arg]
else:
print("--type must be one of: video_tokenizer, actions, dynamics, all")
return 2
# Pre-fetch file listing once
all_files = list_repo_files(repo_id=repo_id, repo_type="model")
any_found = False
for model_type in types_to_fetch:
# Determine local save subdirectory (pattern now controls save path, not search)
pattern_arg = args.pattern[0] if isinstance(args.pattern, list) else args.pattern
save_subdir = pattern_arg or f"{model_type}/checkpoints"
# Search strategy: look under <suite>/ and match any file path containing the model_type string
matched = [
f for f in all_files
if f.startswith(f"{suite}/") and (model_type in f) and f.endswith(".pth")
]
if not matched:
print(f"No checkpoint files found in model repo {repo_id} under '{suite}/' containing '{model_type}'")
continue
any_found = True
# Target directory: results/<timestamp>_<suite>/<save_subdir>
target_root = model_root / save_subdir
pairs = [(repo_id, m) for m in matched]
paths = download_pairs(pairs, target_root, resume=(not args.no_resume), repo_type="model")
for p in paths:
print(p)
return 0 if any_found else 1
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Download datasets/checkpoints from Hugging Face Hub")
subparsers = parser.add_subparsers(dest="cmd")
# datasets subcommand
p_data = subparsers.add_parser("datasets", help="Download dataset files into repo_root/data")
p_data.add_argument("--repo", type=str, help="HF dataset repo_id (default: AlmondGod/tinyworlds)")
p_data.add_argument("--pattern", action="append", help="Glob(s) to match within dataset repo (default: *.h5, assets/*)")
p_data.add_argument("--out", type=Path, help="Target directory (default: repo_root/data)")
p_data.add_argument("--no-resume", action="store_true", help="Disable resume for interrupted downloads")
p_data.set_defaults(func=cmd_datasets)
# models subcommand
p_models = subparsers.add_parser("models", help="Download model checkpoints into results/<timestamp>_<suite>/<type>/checkpoints")
p_models.add_argument("--repo", type=str, help="HF model repo_id (default: AlmondGod/tinyworlds-models)")
p_models.add_argument("--type", default="all", help=f"Model type: {', '.join(VALID_TYPES)}, all")
p_models.add_argument("--suite-name", dest="suite_name", type=str, help="Folder name suffix (e.g., 'sonic')")
p_models.add_argument("--pattern", action="append", help="Save subdirectory under results folder (default: <type>/checkpoints)")
p_models.add_argument("--out", type=Path, help="Base results directory (default: repo_root/results)")
p_models.add_argument("--no-resume", action="store_true", help="Disable resume for interrupted downloads")
p_models.set_defaults(func=cmd_models)
return parser
def main(argv=None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
if not getattr(args, "cmd", None):
parser.print_help()
return 2
return args.func(args)
if __name__ == "__main__":
sys.exit(main())
```
## /scripts/full_train.py
```py path="/scripts/full_train.py"
import sys
import os
from utils.utils import run_command, find_latest_checkpoint, prepare_pipeline_run_root
from utils.config import TrainingConfig, load_config
DEFAULT_TRAINING_CONFIG = os.path.join(os.getcwd(), 'configs', 'training.yaml')
def main():
# load training config
argv = sys.argv
try:
idx = argv.index('--config')
training_cfg_path = argv[idx + 1] if idx + 1 < len(argv) else DEFAULT_TRAINING_CONFIG
except ValueError:
print(f"No training config path provided, using default: {DEFAULT_TRAINING_CONFIG}")
training_cfg_path = DEFAULT_TRAINING_CONFIG
train_config: TrainingConfig = load_config(TrainingConfig, default_config_path=training_cfg_path)
# torchrun if distributed, else python
if train_config.nproc_per_node > 1:
launcher = [
"torchrun",
"--nproc_per_node", str(train_config.nproc_per_node),
]
if train_config.standalone:
launcher += ["--standalone"]
else:
launcher = [sys.executable]
# top-level run root and export child processes can use
run_root, run_name = prepare_pipeline_run_root(base_cwd=os.getcwd())
os.environ['NG_RUN_ROOT_DIR'] = run_root
if train_config.run_video_tokenizer:
v_cmd = launcher + [
"scripts/train_video_tokenizer.py",
"--config", train_config.video_tokenizer_config,
"--training_config", training_cfg_path,
]
if not run_command(v_cmd, "Video Tokenizer Training"):
return
if train_config.run_latent_actions:
latent_actions_cmd = launcher + [
"scripts/train_latent_actions.py",
"--config", train_config.latent_actions_config,
"--training_config", training_cfg_path,
]
if not run_command(latent_actions_cmd, "Latent Actions Training"):
return
# need to get above checkpoints and pass in to dynamics
video_tokenizer_checkpoint = find_latest_checkpoint(".", "video_tokenizer")
latent_actions_checkpoint = find_latest_checkpoint(".", "latent_actions")
if train_config.run_dynamics:
dyn_cmd = launcher + [
"scripts/train_dynamics.py",
"--config", train_config.dynamics_config,
"--training_config", training_cfg_path,
f"video_tokenizer_path={video_tokenizer_checkpoint}",
f"latent_actions_path={latent_actions_checkpoint}",
]
if not run_command(dyn_cmd, "Dynamics Model Training"):
return
dynamics_checkpoint = find_latest_checkpoint(".", "dynamics")
print("\n📁 Results Summary:")
print(f"Video Tokenizer: {video_tokenizer_checkpoint}")
print(f"Latent Actions: {latent_actions_checkpoint}")
print(f"Dynamics Model: {dynamics_checkpoint}")
if __name__ == "__main__":
main()
```
## /scripts/run_inference.py
```py path="/scripts/run_inference.py"
import torch
from datasets.data_utils import load_data_and_data_loaders
import matplotlib.pyplot as plt
import time
import os
import random
import glob
import re
from utils.utils import load_videotokenizer_from_checkpoint, load_latent_actions_from_checkpoint, load_dynamics_from_checkpoint, find_latest_checkpoint
from utils.config import InferenceConfig, load_config
from utils.inference_utils import load_models, visualize_inference, sample_random_action, get_action_latent
from einops import repeat
from typing import Optional
def main():
# load inference config
args: InferenceConfig = load_config(InferenceConfig, default_config_path=os.path.join(os.getcwd(), 'configs', 'inference.yaml'))
# enable tf32 if requested
if args.tf32:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# whether any setting requires using action tokens
use_latent_actions = (args.use_actions or args.use_gt_actions or args.use_interactive_mode)
# check if any path is missing
def missing(path: Optional[str]) -> bool:
return (path is None) or (not os.path.exists(path))
# resolve latest checkpoints if requested or any path missing
base_dir = os.getcwd()
if args.use_latest_checkpoints or missing(args.video_tokenizer_path):
vt_ckpt = find_latest_checkpoint(base_dir, "video_tokenizer")
args.video_tokenizer_path = vt_ckpt
if (args.use_latest_checkpoints or missing(args.latent_actions_path)) and use_latent_actions:
lam_ckpt = find_latest_checkpoint(base_dir, "latent_actions")
args.latent_actions_path = lam_ckpt
if args.use_latest_checkpoints or missing(args.dynamics_path):
dyn_ckpt = find_latest_checkpoint(base_dir, "dynamics")
args.dynamics_path = dyn_ckpt
# confirm which ckpts are being used
print(f"Using video_tokenizer checkpoint: {args.video_tokenizer_path}")
if use_latent_actions:
print(f"Using latent_actions checkpoint: {args.latent_actions_path}")
print(f"Using dynamics checkpoint: {args.dynamics_path}")
# validate required paths
if missing(args.video_tokenizer_path):
raise FileNotFoundError("video_tokenizer_path is not set or not a file. Set it in configs/inference.yaml or enable use_latest_checkpoints with available runs.")
if use_latent_actions and missing(args.latent_actions_path):
raise FileNotFoundError("latent_actions_path is not set or not a file while actions are requested. Set it in configs/inference.yaml or enable use_latest_checkpoints.")
if missing(args.dynamics_path):
raise FileNotFoundError("dynamics_path is not set or not a file. Set it in configs/inference.yaml or enable use_latest_checkpoints.")
# load models, optionally compile
video_tokenizer, latent_action_model, dynamics_model = load_models(args.video_tokenizer_path, args.latent_actions_path, args.dynamics_path, args.device, use_actions=use_latent_actions)
if args.compile:
video_tokenizer = torch.compile(video_tokenizer, mode="reduce-overhead", fullgraph=False, dynamic=True)
if use_latent_actions:
latent_action_model = torch.compile(latent_action_model, mode="reduce-overhead", fullgraph=False, dynamic=True)
dynamics_model = torch.compile(dynamics_model, mode="reduce-overhead", fullgraph=False, dynamic=True)
print("Compiled all models for inference.")
# determine how many ground-truth frames we need in each batch: context + generation steps + prediction horizon
frames_to_load = args.context_window + args.generation_steps * args.prediction_horizon
# dataloader
if hasattr(args, 'preload_ratio') and args.preload_ratio is not None:
data_overrides = {'preload_ratio': args.preload_ratio}
else:
data_overrides = {}
_, _, data_loader, _, _ = load_data_and_data_loaders(
dataset=args.dataset, batch_size=1, num_frames=frames_to_load, **data_overrides)
# sample random batch
random_idx = random.randint(0, len(data_loader.dataset) - 1)
og_ground_truth_frames = data_loader.dataset[random_idx][0] # full sequence
og_ground_truth_frames = og_ground_truth_frames.unsqueeze(0).to(args.device) # [1, frames_to_load, C, H, W]
ground_truth_frames = og_ground_truth_frames[:, :frames_to_load, :, :, :] # [1, frames_to_load, C, H, W]
# start with initial context (first context_window GT frames)
context_frames = ground_truth_frames[:, :args.context_window, :, :, :]
generated_frames = context_frames.clone()
# initialize actions
n_actions = None
inferred_actions = []
if use_latent_actions:
n_actions = latent_action_model.quantizer.codebook_size
# ensure we don’t exceed available GT frames if in teacher-forced mode
max_possible_steps = ground_truth_frames.shape[1] - args.context_window
if args.teacher_forced and args.generation_steps > max_possible_steps:
print(f"[WARN] Requested {args.generation_steps} generation steps but only {max_possible_steps} are possible with teacher-forced context. Clamping.")
effective_steps = args.generation_steps if not args.teacher_forced else min(args.generation_steps, max_possible_steps)
for i in range(effective_steps):
print(f"Inferring frame {i+1}/{effective_steps}")
# select context depending on teacher-forced flag
if args.teacher_forced:
context_start = i # shift window along ground truth
context_frames = ground_truth_frames[:, context_start:context_start+args.context_window, :, :, :] # [1, context_window, C, H, W]
else:
# autoregressive: last context_window frames from generated sequence
context_frames = generated_frames[:, -args.context_window:, :, :, :] # [1, context_window, C, H, W]
# encode context frames each iteration
video_indices = video_tokenizer.tokenize(context_frames)
video_latents = video_tokenizer.quantizer.get_latents_from_indices(video_indices)
sampled_action_index, action_latent = get_action_latent(args, inferred_actions, n_actions, context_frames, latent_action_model, i)
# dynamics forward inference needs idx -> latents fun
def idx_to_latents(idx):
return video_tokenizer.quantizer.get_latents_from_indices(idx, dim=-1)
# autocast for inference if amp enabled (bfloat16 on CUDA by default)
autocast_dtype = torch.bfloat16 if args.amp else None
with torch.amp.autocast('cuda', enabled=args.amp, dtype=autocast_dtype):
next_video_latents = dynamics_model.forward_inference(
context_latents=video_latents,
prediction_horizon=args.prediction_horizon,
num_steps=10,
index_to_latents_fn=idx_to_latents,
conditioning=action_latent,
temperature=args.temperature,
)
# decode next video tokens to frames
next_frames = video_tokenizer.detokenize(next_video_latents) # [1, T, C, H, W]
generated_frames = torch.cat([generated_frames, next_frames[:, -args.prediction_horizon:, :, :]], dim=1)
# TODO: if using interactive mode, visualize next_frames[:, -1] (recently inferred frame) every time, probably with matplotlib is easiest
# point is for user to be able to interact with it in real time
# visualize inference
visualize_inference(generated_frames, ground_truth_frames, inferred_actions, args.fps, use_actions=use_latent_actions)
if __name__ == "__main__":
main()
```
## /scripts/train_dynamics.py
```py path="/scripts/train_dynamics.py"
from contextlib import nullcontext
import torch
import os
from tqdm import tqdm
from einops import rearrange
from models.dynamics import DynamicsModel
from datasets.data_utils import visualize_reconstruction, load_data_and_data_loaders
from tqdm import tqdm
from einops import rearrange
from utils.wandb_utils import (
init_wandb, log_training_metrics, log_learning_rate, log_system_metrics, finish_wandb, log_action_distribution
)
from utils.scheduler_utils import create_cosine_scheduler
from utils.utils import (
readable_timestamp,
save_training_state,
load_videotokenizer_from_checkpoint,
load_latent_actions_from_checkpoint,
load_dynamics_from_checkpoint,
prepare_pipeline_run_root,
prepare_stage_dirs,
)
from utils.config import DynamicsConfig, load_stage_config_merged
import wandb
from dataclasses import asdict
from utils.distributed import init_distributed_from_env, prepare_model_for_distributed, unwrap_model, print_param_count_if_main, cleanup_distributed
from torch.distributed.fsdp import FSDPModule
def main():
# dynamics config merged with training_config.yaml (training takes priority), plus CLI overrides
args: DynamicsConfig = load_stage_config_merged(DynamicsConfig, default_config_path=os.path.join(os.getcwd(), 'configs', 'dynamics.yaml'))
# DDP setup
dist_setup = init_distributed_from_env()
# run save dir if it doesn't exist (running not from full train)
is_main = dist_setup['is_main']
run_root = os.environ.get('NG_RUN_ROOT_DIR')
if not run_root:
run_root, _ = prepare_pipeline_run_root(base_cwd=os.getcwd())
stage_dir, checkpoints_dir, visualizations_dir = prepare_stage_dirs(run_root, 'dynamics')
if is_main:
print(f"Dynamics Training")
print(f"Results will be saved in {stage_dir}")
# load video tokenizer and latent action model
if os.path.isdir(args.video_tokenizer_path):
video_tokenizer, vq_ckpt = load_videotokenizer_from_checkpoint(
checkpoint_path=args.video_tokenizer_path,
device=args.device,
is_distributed=dist_setup['is_distributed'],
)
video_tokenizer.eval()
for p in video_tokenizer.parameters():
p.requires_grad = False
else:
raise FileNotFoundError(f"Video tokenizer checkpoint not found at {args.video_tokenizer_path}")
if os.path.isdir(args.latent_actions_path):
latent_action_model, latent_action_ckpt = load_latent_actions_from_checkpoint(
checkpoint_path=args.latent_actions_path,
device=args.device,
is_distributed=dist_setup['is_distributed'],
)
unwrap_model(latent_action_model).eval()
for p in unwrap_model(latent_action_model).parameters():
p.requires_grad = False
else:
raise FileNotFoundError(f"Latent Action Model checkpoint not found at {args.latent_actions_path}")
# init dynamics model and optional ckpt load
dynamics_model = DynamicsModel(
frame_size=(args.frame_size, args.frame_size),
patch_size=args.patch_size,
embed_dim=args.embed_dim,
num_heads=args.num_heads,
hidden_dim=args.hidden_dim,
num_blocks=args.num_blocks,
conditioning_dim=unwrap_model(latent_action_model).action_dim,
latent_dim=args.latent_dim,
num_bins=args.num_bins,
use_moe=getattr(args, 'use_moe', False),
num_experts=getattr(args, 'num_experts', 4),
top_k_experts=getattr(args, 'top_k_experts', 2),
moe_aux_loss_coeff=getattr(args, 'moe_aux_loss_coeff', 0.01),
).to(args.device)
if args.checkpoint:
dynamics_model, _ = load_dynamics_from_checkpoint(
checkpoint_path=args.checkpoint,
device=args.device,
model=dynamics_model,
is_distributed=dist_setup['is_distributed'],
)
# optional DDP, compile, param count, tf32
print_param_count_if_main(dynamics_model, "DynamicsModel", is_main)
if args.compile:
video_tokenizer = torch.compile(video_tokenizer, mode="reduce-overhead", fullgraph=False, dynamic=True)
latent_action_model = torch.compile(latent_action_model, mode="reduce-overhead", fullgraph=False, dynamic=True)
dynamics_model = torch.compile(dynamics_model, mode="reduce-overhead", fullgraph=False, dynamic=True)
print("Compiled all models for training.")
dynamics_model = prepare_model_for_distributed(
dynamics_model,
args.distributed,
model_type=dynamics_model.model_type,
device_mesh=dist_setup['device_mesh'],
)
if args.tf32:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# create optimizer(s) — AdamW or Muon+AdamW split
from utils.optimizer_utils import create_optimizer
optimizers = create_optimizer(dynamics_model, args)
# cosine scheduler for lr warmup and AMP grad scaler
schedulers = [create_cosine_scheduler(opt, args.n_updates) for opt in optimizers]
train_ctx = torch.amp.autocast(args.device, enabled=True, dtype=torch.bfloat16) if args.amp and not args.distributed.use_fsdp else nullcontext()
results = {
'n_updates': 0,
'dynamics_losses': [],
'loss_vals': [],
}
# init wandb
if args.use_wandb and is_main:
run_name = f"dynamics_{readable_timestamp()}"
init_wandb(args.wandb_project, asdict(args), run_name)
unwrap_model(dynamics_model).train()
# dataloader
data_overrides = {}
if hasattr(args, 'fps') and args.fps is not None:
data_overrides['fps'] = args.fps
if hasattr(args, 'preload_ratio') and args.preload_ratio is not None:
data_overrides['preload_ratio'] = args.preload_ratio
_, _, training_loader, _, _ = load_data_and_data_loaders(
dataset=args.dataset,
batch_size=args.batch_size_per_gpu,
num_frames=args.context_length,
distributed=dist_setup['is_distributed'],
rank=dist_setup['device_mesh'].get_rank() if dist_setup['device_mesh'] is not None else 0,
world_size=dist_setup['world_size'],
**data_overrides,
)
train_iter = iter(training_loader)
use_moe = getattr(args, 'use_moe', False)
for i in tqdm(range(0, args.n_updates), disable=not is_main):
for opt in optimizers:
opt.zero_grad(set_to_none=True)
if isinstance(dynamics_model, FSDPModule):
dynamics_model.set_requires_gradient_sync(False)
if args.compile:
torch.compiler.cudagraph_mark_step_begin()
for micro_batch in range(args.gradient_accumulation_steps):
try:
x, _ = next(train_iter)
except StopIteration:
train_iter = iter(training_loader) # reset iterator when epoch ends
x, _ = next(train_iter)
x = x.to(args.device, non_blocking=True) # [batch_size, seq_len, channels, height, width]
# get video tokens for batch
video_tokens = video_tokenizer.tokenize(x) # [B, T, P]
video_latents = video_tokenizer.quantizer.get_latents_from_indices(video_tokens, dim=-1) # [B, T, P, L]
if args.use_actions:
quantized_actions = latent_action_model.encode(x) # [B, T - 1, A]
else:
quantized_actions = None
# predict masked frame latents with dynamics model (masking in dynamics model)
with train_ctx:
predicted_next_logits, mask_positions, loss = dynamics_model(
video_latents, training=True, conditioning=quantized_actions, targets=video_tokens
)
# add MoE load-balancing auxiliary loss
moe_aux_scalar = 0.0
if use_moe:
aux_loss = unwrap_model(dynamics_model).transformer.moe_aux_loss()
moe_aux_scalar = aux_loss.item()
loss = loss + aux_loss
if isinstance(dynamics_model, FSDPModule):
if (micro_batch + 1) % args.gradient_accumulation_steps == 0:
dynamics_model.set_requires_gradient_sync(True)
loss /= args.gradient_accumulation_steps
loss.backward()
results['n_updates'] = i
results['dynamics_losses'].append(loss.detach().cpu())
results['loss_vals'].append(loss.detach().cpu())
torch.nn.utils.clip_grad_norm_(unwrap_model(dynamics_model).parameters(), max_norm=1.0)
for opt in optimizers:
opt.step()
for sched in schedulers:
sched.step()
# wandb logging
if args.use_wandb and is_main:
log_dict = {'train/loss': loss.item()}
if use_moe:
log_dict['train/moe_aux_loss'] = moe_aux_scalar
# per-expert utilization across blocks
expert_util = unwrap_model(dynamics_model).transformer.moe_expert_utilization()
for block_idx, fracs in expert_util.items():
for expert_i, frac in enumerate(fracs.tolist()):
log_dict[f'moe/block{block_idx}_expert{expert_i}_frac'] = frac
wandb.log(log_dict, step=i)
log_system_metrics(i)
log_learning_rate(optimizers[0], i)
if args.use_actions:
action_indices = latent_action_model.quantizer.get_indices_from_latents(quantized_actions)
log_action_distribution(action_indices, i, args.n_actions)
# save model and visualize results
if i % args.log_interval == 0:
if args.use_wandb:
predicted_next_indices = torch.argmax(predicted_next_logits, dim=-1)
predicted_next_latents = video_tokenizer.quantizer.get_latents_from_indices(predicted_next_indices, dim=-1)
with torch.no_grad():
predicted_frames = video_tokenizer.decoder(predicted_next_latents[:16]) # [B, T, C, H, W]
# convert mask_positions to patch-level mask for visualization
B, T, P = mask_positions.shape
patch_size = args.patch_size
H, W = args.frame_size, args.frame_size
pixel_mask = torch.zeros(B, T, H, W, device=mask_positions.device)
# for each pixel patch, mask if equivalent token is masked
for b in range(B):
for t in range(T):
for p in range(P):
if mask_positions[b, t, p]:
patch_row = (p // (W // patch_size)) * patch_size
patch_col = (p % (W // patch_size)) * patch_size
pixel_mask[b, t, patch_row:patch_row+patch_size, patch_col:patch_col+patch_size] = 1 # assigning 1 to the patch in the mask of dim [1, 1, Hp, Wp]
pixel_mask_expanded = rearrange(pixel_mask, 'b t h w -> b t 1 h w')
masked_frames = x * (1 - pixel_mask_expanded)
hyperparameters = args.__dict__
ckpt_path = save_training_state(dynamics_model, optimizers[0], schedulers[0], hyperparameters, checkpoints_dir, prefix='dynamics', step=i)
# save secondary optimizer/scheduler state when using split optimizers (Muon+AdamW)
if len(optimizers) > 1:
import pathlib
torch.save(
{f'optimizer_{i}': o.state_dict() for i, o in enumerate(optimizers)},
pathlib.Path(ckpt_path) / 'all_optimizers.pt',
)
torch.save(
{f'scheduler_{i}': s.state_dict() for i, s in enumerate(schedulers)},
pathlib.Path(ckpt_path) / 'all_schedulers.pt',
)
if is_main:
save_path = os.path.join(visualizations_dir, f'dynamics_prediction_step_{i}.png')
visualize_reconstruction(masked_frames[:16].cpu(), predicted_frames[:16].cpu(), save_path)
print('\n Step', i, 'Loss:', torch.mean(torch.stack(results["loss_vals"][-args.log_interval:])).item())
# finish wandb
if args.use_wandb and is_main:
finish_wandb()
cleanup_distributed(dist_setup['is_distributed'])
if __name__ == "__main__":
main()
```
## /scripts/train_latent_actions.py
```py path="/scripts/train_latent_actions.py"
from contextlib import nullcontext
import torch
import os
from models.latent_actions import LatentActionModel
from datasets.data_utils import load_data_and_data_loaders, visualize_reconstruction
from utils.scheduler_utils import create_cosine_scheduler
from tqdm import tqdm
import wandb
from utils.utils import readable_timestamp, save_training_state, prepare_stage_dirs, prepare_pipeline_run_root
from utils.config import LatentActionsConfig, load_stage_config_merged
from utils.utils import save_training_state, load_latent_actions_from_checkpoint
from utils.wandb_utils import init_wandb, log_system_metrics, finish_wandb, log_action_distribution, log_learning_rate
from dataclasses import asdict
from utils.distributed import init_distributed_from_env, prepare_model_for_distributed, unwrap_model, print_param_count_if_main, cleanup_distributed
from torch.distributed.fsdp import FSDPModule
def main():
# latent actions config merged with training_config.yaml (training takes priority), plus CLI overrides
args: LatentActionsConfig = load_stage_config_merged(LatentActionsConfig, default_config_path=os.path.join(os.getcwd(), 'configs', 'latent_actions.yaml'))
# DDP setup
dist_setup = init_distributed_from_env()
# run save dir if it doesn't exist (running not from full train)
timestamp = readable_timestamp()
run_root = os.environ.get('NG_RUN_ROOT_DIR')
if not run_root:
run_root, _ = prepare_pipeline_run_root(base_cwd=os.getcwd())
is_main = dist_setup['is_main']
stage_dir, checkpoints_dir, visualizations_dir = prepare_stage_dirs(run_root, 'latent_actions')
if is_main:
print(f"Latent Actions Training")
print(f'Results will be saved in {stage_dir}')
# dataloader
data_overrides = {}
if hasattr(args, 'fps') and args.fps is not None:
data_overrides['fps'] = args.fps
if hasattr(args, 'preload_ratio') and args.preload_ratio is not None:
data_overrides['preload_ratio'] = args.preload_ratio
training_data, validation_data, training_loader, validation_loader, x_train_var = load_data_and_data_loaders(
dataset=args.dataset,
batch_size=args.batch_size_per_gpu,
num_frames=args.context_length,
distributed=dist_setup['is_distributed'],
rank=dist_setup['device_mesh'].get_rank() if dist_setup['device_mesh'] is not None else 0,
world_size=dist_setup['world_size'],
**data_overrides,
)
# init model and optional ckpt load
model = LatentActionModel(
frame_size=(args.frame_size, args.frame_size),
patch_size=args.patch_size,
embed_dim=args.embed_dim,
num_heads=args.num_heads,
hidden_dim=args.hidden_dim,
num_blocks=args.num_blocks,
n_actions=args.n_actions,
).to(args.device)
if args.checkpoint:
model, _ = load_latent_actions_from_checkpoint(
args.checkpoint,
args.device,
model,
dist_setup['is_distributed'],
)
# optional DDP, compile, param count, tf32
print_param_count_if_main(model, "LatentActionModel", is_main)
if args.compile:
model = torch.compile(model, mode="reduce-overhead", fullgraph=False, dynamic=True)
model = prepare_model_for_distributed(
model,
args.distributed,
model_type=model.model_type,
device_mesh=dist_setup['device_mesh'],
)
if args.tf32:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# create optimizer(s) — AdamW or Muon+AdamW split
from utils.optimizer_utils import create_optimizer
optimizers = create_optimizer(model, args)
# cosine scheduler for lr warmup and AMP
schedulers = [create_cosine_scheduler(opt, args.n_updates) for opt in optimizers]
train_ctx = torch.amp.autocast(args.device, enabled=True, dtype=torch.bfloat16) if args.amp and not args.distributed.use_fsdp else nullcontext()
results = {
'n_updates': 0,
'loss_vals': [],
}
# init wandb
if args.use_wandb and is_main:
cfg = asdict(args)
cfg.update({'timestamp': timestamp})
run_name = f"latent_actions_{timestamp}"
init_wandb(args.wandb_project, cfg, run_name)
unwrap_model(model).train()
train_iter = iter(training_loader)
for i in tqdm(range(args.n_updates), disable=not is_main):
for opt in optimizers:
opt.zero_grad(set_to_none=True)
if isinstance(model, FSDPModule):
model.set_requires_gradient_sync(False)
if args.compile:
torch.compiler.cudagraph_mark_step_begin()
for micro_batch in range(args.gradient_accumulation_steps):
try:
(x, _) = next(train_iter)
except StopIteration:
train_iter = iter(training_loader)
(x, _) = next(train_iter)
x = x.to(args.device, non_blocking=True)
with train_ctx:
loss, pred_frames = model(x)
loss /= args.gradient_accumulation_steps
if isinstance(model, FSDPModule):
if (micro_batch + 1) % args.gradient_accumulation_steps == 0:
model.set_requires_gradient_sync(True)
loss.backward()
torch.nn.utils.clip_grad_norm_(unwrap_model(model).parameters(), max_norm=1.0)
for opt in optimizers:
opt.step()
for sched in schedulers:
sched.step()
results['n_updates'] = i
results['loss_vals'].append(loss.detach().cpu())
if args.use_wandb and is_main:
wandb.log({
'train/loss': loss.item(),
}, step=i)
log_system_metrics(i)
log_learning_rate(optimizers[0], i)
# save model and visualize results
if i % args.log_interval == 0:
if args.use_wandb:
with torch.no_grad():
actions = unwrap_model(model).encoder(x)
actions_quantized = unwrap_model(model).quantizer(actions)
idx = unwrap_model(model).quantizer.get_indices_from_latents(actions_quantized)
codebook_usage = idx.unique().numel() / unwrap_model(model).quantizer.codebook_size
z_e_var = actions.var(dim=0, unbiased=False).mean().item()
pred_frames_var = pred_frames.var(dim=0, unbiased=False).mean().item()
if args.use_wandb and is_main:
wandb.log({
"latent_actions/codebook_usage": codebook_usage,
"latent_actions/encoder_variance": z_e_var,
"latent_actions/decoder_variance": pred_frames_var,
}, step=i)
log_action_distribution(idx, i, args.n_actions)
hyperparameters = vars(args)
save_training_state(model, optimizers[0], None, hyperparameters, checkpoints_dir, prefix='latent_actions', step=i)
if is_main:
save_path = os.path.join(visualizations_dir, f'reconstructions_latent_actions_step_{i}.png')
visualize_reconstruction(x, pred_frames, save_path)
print('\n Step', i, 'Loss:', loss.item(), 'Codebook Usage:', codebook_usage, 'Encoder Variance:', z_e_var, 'Decoder Variance:', pred_frames_var)
# finish wandb
if args.use_wandb and is_main:
finish_wandb()
cleanup_distributed(dist_setup['is_distributed'])
if __name__ == "__main__":
main()
```
## /scripts/train_video_tokenizer.py
```py path="/scripts/train_video_tokenizer.py"
from contextlib import nullcontext
import torch
import os
from models.video_tokenizer import VideoTokenizer
from datasets.data_utils import visualize_reconstruction, load_data_and_data_loaders
from utils.scheduler_utils import create_cosine_scheduler
from tqdm import tqdm
import wandb
from utils.utils import readable_timestamp, save_training_state, prepare_stage_dirs, prepare_pipeline_run_root
from utils.config import VideoTokenizerConfig, load_stage_config_merged
from utils.utils import save_training_state, load_videotokenizer_from_checkpoint
from utils.wandb_utils import init_wandb, log_training_metrics, log_system_metrics, log_learning_rate, finish_wandb
from dataclasses import asdict
from utils.distributed import init_distributed_from_env, prepare_model_for_distributed, unwrap_model, print_param_count_if_main, cleanup_distributed
from torch.distributed.fsdp import FSDPModule
def main():
# vidtokenizer config merged with training_config.yaml (training takes priority), plus CLI overrides
args: VideoTokenizerConfig = load_stage_config_merged(VideoTokenizerConfig, default_config_path=os.path.join(os.getcwd(), 'configs', 'video_tokenizer.yaml'))
# DDP setup
dist_setup = init_distributed_from_env()
# run save dir if it doesn't exist (running not from full train)
timestamp = readable_timestamp()
run_root = os.environ.get('NG_RUN_ROOT_DIR')
if not run_root:
run_root, _ = prepare_pipeline_run_root(base_cwd=os.getcwd())
is_main = dist_setup['is_main']
stage_dir, checkpoints_dir, visualizations_dir = prepare_stage_dirs(run_root, 'video_tokenizer')
if is_main:
print(f"Video Tokenizer Training")
print(f'Results will be saved in {stage_dir}')
# dataloader
data_overrides = {}
if hasattr(args, 'fps') and args.fps is not None:
data_overrides['fps'] = args.fps
if hasattr(args, 'preload_ratio') and args.preload_ratio is not None:
data_overrides['preload_ratio'] = args.preload_ratio
training_data, validation_data, training_loader, validation_loader, x_train_var = load_data_and_data_loaders(
dataset=args.dataset,
batch_size=args.batch_size_per_gpu,
num_frames=args.context_length,
distributed=dist_setup['is_distributed'],
rank=dist_setup['device_mesh'].get_rank() if dist_setup['device_mesh'] is not None else 0,
world_size=dist_setup['world_size'],
**data_overrides,
)
# print("Length of training data:", len(training_data))
# print("Length of validation data:", len(validation_data))
# init model and optional ckpt load
model = VideoTokenizer(
frame_size=(args.frame_size, args.frame_size),
patch_size=args.patch_size,
embed_dim=args.embed_dim,
num_heads=args.num_heads,
hidden_dim=args.hidden_dim,
num_blocks=args.num_blocks,
latent_dim=args.latent_dim,
num_bins=args.num_bins,
).to(args.device)
if args.checkpoint:
model, _ = load_videotokenizer_from_checkpoint(
args.checkpoint,
args.device,
model,
dist_setup['is_distributed'],
)
# optional DDP, compile, param count, tf32
print_param_count_if_main(model, "VideoTokenizer", is_main)
if args.compile:
model = torch.compile(model, mode="reduce-overhead", fullgraph=False, dynamic=True)
model = prepare_model_for_distributed(
model,
args.distributed,
model_type=model.model_type,
device_mesh=dist_setup['device_mesh'],
)
if args.tf32:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# create optimizer(s) — AdamW or Muon+AdamW split
from utils.optimizer_utils import create_optimizer
optimizers = create_optimizer(model, args)
# cosine scheduler for lr warmup
schedulers = [create_cosine_scheduler(opt, args.n_updates) for opt in optimizers]
train_ctx = torch.amp.autocast(args.device, enabled=True, dtype=torch.bfloat16) if args.amp and not args.distributed.use_fsdp else nullcontext()
results = {
'n_updates': 0,
'loss_vals': [],
}
# init wandb
if args.use_wandb and is_main:
cfg = asdict(args)
cfg.update({'timestamp': timestamp})
run_name = f"video_tokenizer_{timestamp}"
init_wandb(args.wandb_project, cfg, run_name)
unwrap_model(model).train()
train_iter = iter(training_loader)
for i in tqdm(range(args.n_updates), disable=not is_main):
for opt in optimizers:
opt.zero_grad(set_to_none=True)
if isinstance(model, FSDPModule):
model.set_requires_gradient_sync(False)
if args.compile:
torch.compiler.cudagraph_mark_step_begin()
for micro_batch in range(args.gradient_accumulation_steps):
try:
(x, _) = next(train_iter)
except StopIteration:
train_iter = iter(training_loader) # reset iterator when epoch ends
(x, _) = next(train_iter)
x = x.to(args.device, non_blocking=True)
with train_ctx:
loss, x_hat = model(x)
loss /= args.gradient_accumulation_steps
if isinstance(model, FSDPModule):
if (micro_batch + 1) % args.gradient_accumulation_steps == 0:
model.set_requires_gradient_sync(True)
loss.backward()
torch.nn.utils.clip_grad_norm_(unwrap_model(model).parameters(), max_norm=1.0)
for opt in optimizers:
opt.step()
for sched in schedulers:
sched.step()
results["loss_vals"].append(loss.cpu().detach())
results["n_updates"] = i
# wandb logging
if args.use_wandb and is_main:
metrics = {
'loss': loss.item(),
'learning_rate': schedulers[0].get_last_lr()[0],
}
log_training_metrics(i, metrics, prefix='train')
log_system_metrics(i)
log_learning_rate(optimizers[0], i)
# save model and visualize results
if i % args.log_interval == 0:
if args.use_wandb:
with torch.no_grad():
indices = unwrap_model(model).tokenize(x)
unique_codes = torch.unique(indices).numel()
codebook_usage = unique_codes / unwrap_model(model).codebook_size
if is_main:
wandb.log({'train/codebook_usage': codebook_usage}, step=i)
hyperparameters = args.__dict__
save_training_state(model, optimizers[0], schedulers[0], hyperparameters, checkpoints_dir, prefix='video_tokenizer', step=i)
if is_main:
x_hat_vis = x_hat.detach().cpu()
x_vis = x.detach().cpu()
save_path = os.path.join(visualizations_dir, f'video_tokenizer_recon_step_{i}.png')
visualize_reconstruction(x_vis[:16], x_hat_vis[:16], save_path)
print('\n Step', i, 'Loss:', torch.mean(torch.stack(results["loss_vals"][-args.log_interval:])).item())
# finish wandb
if args.use_wandb and is_main:
finish_wandb()
cleanup_distributed(dist_setup['is_distributed'])
if __name__ == "__main__":
main()
```
## /scripts/visualize_batch.py
```py path="/scripts/visualize_batch.py"
import torch
import matplotlib.pyplot as plt
import numpy as np
import os
import sys
import argparse
from datasets.data_utils import load_data_and_data_loaders
def visualize_batch(frames, save_path=None, title="Video Sequences Batch", max_batch_size=8, max_seq_length=8):
# move to CPU and get dimensions
frames = frames.detach().cpu()
batch_size, seq_len, C, H, W = frames.shape
batch_size = min(batch_size, max_batch_size)
seq_len = min(seq_len, max_seq_length)
frames = frames[:batch_size, :seq_len]
# denormalize from [-1, 1] to [0, 1]
frames = (frames + 1) / 2
frames = torch.clamp(frames, 0, 1)
fig, axes = plt.subplots(batch_size, seq_len, figsize=(2 * seq_len, 2 * batch_size))
# single row/column case
if batch_size == 1:
axes = axes.reshape(1, -1)
if seq_len == 1:
axes = axes.reshape(-1, 1)
# plot frames
for i in range(batch_size):
for j in range(seq_len):
frame = frames[i, j].permute(1, 2, 0).numpy() # [H, W, C]
axes[i, j].imshow(frame)
axes[i, j].set_title(f'B{i}, T{j}', fontsize=10)
axes[i, j].axis('off')
# row and column labels
for i in range(batch_size):
axes[i, 0].set_ylabel(f'Batch {i}', fontsize=12, fontweight='bold')
for j in range(seq_len):
axes[0, j].set_title(f'Timestep {j}', fontsize=12, fontweight='bold')
plt.suptitle(title, fontsize=16, fontweight='bold')
plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150, bbox_inches='tight')
print(f"Visualization saved to: {save_path}")
plt.show()
def visualize_batch_with_stats(frames, save_path=None, title="Video Sequences Batch with Statistics"):
# Move to CPU and get dimensions
frames = frames.detach().cpu()
batch_size, seq_len, C, H, W = frames.shape
frame_mean = frames.mean().item()
frame_std = frames.std().item()
frame_min = frames.min().item()
frame_max = frames.max().item()
visualize_batch(frames, save_path, title)
return {
'mean': frame_mean,
'std': frame_std,
'min': frame_min,
'max': frame_max,
'shape': frames.shape
}
def visualize_multiple_batches(dataloader, num_batches=3, save_dir="batch_visualizations"):
os.makedirs(save_dir, exist_ok=True)
for batch_idx, (frames, _) in enumerate(dataloader):
if batch_idx >= num_batches:
break
stats = visualize_batch_with_stats(
frames,
save_path=os.path.join(save_dir, f"batch_{batch_idx + 1}.png"),
title=f"Batch {batch_idx + 1} - Video Sequences"
)
def main():
parser = argparse.ArgumentParser(description="Visualize video sequence batches")
parser.add_argument("--dataset", type=str, default="SONIC", help="Dataset to use")
parser.add_argument("--batch_size", type=int, default=4, help="Batch size")
parser.add_argument("--context_length", type=int, default=4, help="Context length")
parser.add_argument("--num_batches", type=int, default=4, help="Number of batches to visualize")
parser.add_argument("--save_dir", type=str, default="batch_visualizations", help="Directory to save visualizations")
parser.add_argument("--max_batch_display", type=int, default=8, help="Maximum batch elements to display")
parser.add_argument("--max_seq_display", type=int, default=8, help="Maximum timesteps to display")
args = parser.parse_args()
print(f"Loading {args.dataset} dataset...")
# load data
_, _, validation_loader, _, _ = load_data_and_data_loaders(
dataset=args.dataset,
batch_size=args.batch_size,
num_frames=args.context_length
)
# visualize batches
for batch_idx, (frames, _) in enumerate(validation_loader):
if batch_idx >= args.num_batches:
break
# calculate statistics
frames_cpu = frames.detach().cpu()
batch_size, seq_len, C, H, W = frames_cpu.shape
# visualize batch
os.makedirs(args.save_dir, exist_ok=True)
save_path = os.path.join(args.save_dir, f"{args.dataset}_batch_{batch_idx + 1}.png")
visualize_batch(
frames,
save_path=save_path,
title=f"Batch {batch_idx + 1} - {args.dataset} Sequences",
max_batch_size=args.max_batch_display,
max_seq_length=args.max_seq_display
)
if __name__ == "__main__":
main()
```
## /utils/config.py
```py path="/utils/config.py"
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
import argparse
import os
from omegaconf import OmegaConf
from torch.distributed.fsdp import MixedPrecisionPolicy, CPUOffloadPolicy
import torch
class DeviceType(str, Enum):
CUDA: str = 'cuda'
CPU: str = 'cpu'
@dataclass
class FSDPMixedPrecisionConfig:
param_dtype: str = "bfloat16"
reduce_dtype: str = "float32"
output_dtype: str = "float32"
cast_forward_inputs: bool = True
def _resolve_dtype(self, value: str | torch.dtype) -> torch.dtype:
if isinstance(value, torch.dtype):
return value
try:
return getattr(torch, value)
except AttributeError as exc:
raise ValueError(f"Unknown torch dtype '{value}'") from exc
def to_policy(self) -> MixedPrecisionPolicy:
return MixedPrecisionPolicy(
param_dtype=self._resolve_dtype(self.param_dtype),
reduce_dtype=self._resolve_dtype(self.reduce_dtype),
output_dtype=self._resolve_dtype(self.output_dtype),
cast_forward_inputs=self.cast_forward_inputs,
)
@dataclass
class DistributedConfig:
use_ddp: bool = False
use_fsdp: bool = False
reshard_after_forward: bool = False
fsdp_mixed_precision: FSDPMixedPrecisionConfig | None = field(default_factory=FSDPMixedPrecisionConfig)
offload_policy: CPUOffloadPolicy | None = None
def __post_init__(self) -> None:
if self.use_ddp and self.use_fsdp:
raise ValueError("DistributedConfig cannot enable both DDP and FSDP; choose only one.")
def get_mixed_precision_policy(self) -> MixedPrecisionPolicy | None:
if self.fsdp_mixed_precision is None:
return None
if isinstance(self.fsdp_mixed_precision, MixedPrecisionPolicy):
return self.fsdp_mixed_precision
return self.fsdp_mixed_precision.to_policy()
def _validate_amp_fsdp(amp: bool, distributed: DistributedConfig) -> None:
if amp and distributed.use_fsdp:
raise ValueError(
"Disable AMP when using FSDP; configure mixed precision via distributed.fsdp_mixed_precision instead."
)
def _validate_distibuted_training(nproc_per_node: int, distributed: DistributedConfig) -> None:
if nproc_per_node > 1 and not (distributed.use_ddp or distributed.use_fsdp):
raise ValueError(
"nproc_per_node > 1 requires enabling distributed.use_ddp or distributed.use_fsdp."
)
def _validate_distributed_device(device: DeviceType, distributed: DistributedConfig) -> None:
current_device = device if isinstance(device, DeviceType) else DeviceType(device)
if (distributed.use_ddp or distributed.use_fsdp) and current_device is not DeviceType.CUDA:
raise ValueError("Distributed training (DDP/FSDP) requires device=cuda.")
@dataclass
class VideoTokenizerConfig:
# Training
batch_size_per_gpu: int
gradient_accumulation_steps: int
n_updates: int # number of optimizer.step(), excluding grad_accum_step
learning_rate: float
log_interval: int
dataset: str
context_length: int
frame_size: int
# Model
patch_size: int
embed_dim: int
num_heads: int
hidden_dim: int
num_blocks: int
latent_dim: int
num_bins: int
amp: bool
tf32: bool
compile: bool
# distributed
distributed: DistributedConfig
nproc_per_node: int
standalone: bool
# W&B
use_wandb: bool
wandb_project: str
# resume from checkpoint
checkpoint: Optional[str]
# Optimizer
optimizer: str = "adamw"
muon_momentum: float = 0.95
muon_backend_steps: int = 5
# device
device: DeviceType = DeviceType.CUDA
# other params
fps: Optional[int] = None
preload_ratio: Optional[float] = None
def __post_init__(self) -> None:
_validate_amp_fsdp(self.amp, self.distributed)
_validate_distibuted_training(self.nproc_per_node, self.distributed)
_validate_distributed_device(self.device, self.distributed)
@dataclass
class LatentActionsConfig:
# Training
batch_size_per_gpu: int
gradient_accumulation_steps: int
n_updates: int # number of optimizer.step(), excluding grad_accum_step
learning_rate: float
log_interval: int
dataset: str
context_length: int
frame_size: int
# Model
n_actions: int
patch_size: int
embed_dim: int
num_heads: int
hidden_dim: int
num_blocks: int
amp: bool
tf32: bool
compile: bool
# distributed
distributed: DistributedConfig
nproc_per_node: int
standalone: bool
# W&B
use_wandb: bool
wandb_project: str
# resume from checkpoint
checkpoint: Optional[str]
# Optimizer
optimizer: str = "adamw"
muon_momentum: float = 0.95
muon_backend_steps: int = 5
# device
device: DeviceType = DeviceType.CUDA
# other params
fps: Optional[int] = None
preload_ratio: Optional[float] = None
def __post_init__(self) -> None:
_validate_amp_fsdp(self.amp, self.distributed)
_validate_distibuted_training(self.nproc_per_node, self.distributed)
_validate_distributed_device(self.device, self.distributed)
@dataclass
class DynamicsConfig:
# Training
batch_size_per_gpu: int
gradient_accumulation_steps: int
n_updates: int # number of optimizer.step(), excluding grad_accum_step
learning_rate: float
log_interval: int
dataset: str
context_length: int
frame_size: int
# Model (must match tokenizer)
patch_size: int
embed_dim: int
num_heads: int
hidden_dim: int
num_blocks: int
latent_dim: int
num_bins: int
n_actions: int
use_actions: bool
# Paths
video_tokenizer_path: Optional[str]
latent_actions_path: Optional[str]
# Perf
amp: bool
tf32: bool
compile: bool
# distributed
distributed: DistributedConfig
nproc_per_node: int
standalone: bool
# W&B
use_wandb: bool
wandb_project: str
# resume from checkpoint
checkpoint: Optional[str]
# MoE
use_moe: bool = False
num_experts: int = 4
top_k_experts: int = 2
moe_aux_loss_coeff: float = 0.01
# Optimizer
optimizer: str = "adamw"
muon_momentum: float = 0.95
muon_backend_steps: int = 5
# device
device: DeviceType = DeviceType.CUDA
# other params
fps: Optional[int] = None
preload_ratio: Optional[float] = None
def __post_init__(self) -> None:
_validate_amp_fsdp(self.amp, self.distributed)
_validate_distibuted_training(self.nproc_per_node, self.distributed)
_validate_distributed_device(self.device, self.distributed)
@dataclass
class TrainingConfig:
# WandB
use_wandb: bool
wandb_project: str
# Dataset
dataset: str
# Config paths for stages
video_tokenizer_config: str
latent_actions_config: str
dynamics_config: str
# Which stages to run
run_video_tokenizer: bool
run_latent_actions: bool
run_dynamics: bool
# Shared model params
patch_size: int
context_length: int
frame_size: int
latent_dim: int
num_bins: int
n_actions: int
# Performance
amp: bool
tf32: bool
compile: bool
# distributed
distributed: DistributedConfig
nproc_per_node: int
standalone: bool
# device
device: DeviceType = DeviceType.CUDA
# These can vary per model
embed_dim: Optional[int] = None
num_heads: Optional[int] = None
hidden_dim: Optional[int] = None
num_blocks: Optional[int] = None
learning_rate: Optional[float] = None
batch_size_per_gpu: Optional[int] = None
gradient_accumulation_steps: Optional[int] = None
log_interval: Optional[int] = None
n_updates: Optional[int] = None # number of optimizer.step(), excluding grad_accum_step
fps: Optional[int] = None
preload_ratio: Optional[float] = None
# MoE (dynamics only)
use_moe: bool = False
num_experts: int = 4
top_k_experts: int = 2
moe_aux_loss_coeff: float = 0.01
# Optimizer
optimizer: str = "adamw"
muon_momentum: float = 0.95
muon_backend_steps: int = 5
def __post_init__(self) -> None:
_validate_amp_fsdp(self.amp, self.distributed)
_validate_distibuted_training(self.nproc_per_node, self.distributed)
_validate_distributed_device(self.device, self.distributed)
@dataclass
class InferenceConfig:
video_tokenizer_path: Optional[str]
latent_actions_path: Optional[str]
dynamics_path: Optional[str]
device: str
generation_steps: int
context_window: int
fps: int
temperature: float
use_actions: bool
teacher_forced: bool
use_latest_checkpoints: bool
prediction_horizon: int
dataset: str
use_gt_actions: bool
# Inference performance options
amp: bool
tf32: bool
compile: bool
# Interactive mode (user enters action ids)
use_interactive_mode: bool
preload_ratio: Optional[float] = None
def load_config(config_cls, default_config_path: Optional[str] = None):
parser = argparse.ArgumentParser(add_help=True)
parser.add_argument("--config", type=str, default=default_config_path)
# Accept dotlist overrides like key=value
parser.add_argument("overrides", nargs=argparse.REMAINDER)
args = parser.parse_args()
# Build a structured schema from the dataclass TYPE, not an instance.
# This allows Python-side defaults to be omitted and provided solely via YAML.
base = OmegaConf.structured(config_cls)
cfg = base
if args.config is not None:
if not os.path.isfile(args.config):
raise FileNotFoundError(f"Config file not found: {args.config} (cwd: {os.getcwd()})")
file_cfg = OmegaConf.load(args.config)
cfg = OmegaConf.merge(cfg, file_cfg)
# Merge dotlist overrides if any (ignore leading '--')
dot_overrides = [s.lstrip('-') for s in (args.overrides or []) if '=' in s]
if dot_overrides:
cli_cfg = OmegaConf.from_dotlist(dot_overrides)
cfg = OmegaConf.merge(cfg, cli_cfg)
# Return a typed dataclass instance
return OmegaConf.to_object(cfg)
def load_stage_config_merged(config_cls, default_config_path: Optional[str] = None, training_config_path: Optional[str] = None):
"""Load a stage config YAML, then overlay values from training_config.yaml (priority),
restricted to keys that exist in the stage schema. CLI dotlist overrides still have highest priority.
"""
parser = argparse.ArgumentParser(add_help=True)
parser.add_argument("--config", type=str, default=default_config_path)
parser.add_argument("--training_config", type=str, default=training_config_path or os.path.join(os.getcwd(), 'configs', 'training.yaml'))
parser.add_argument("overrides", nargs=argparse.REMAINDER)
args = parser.parse_args()
# Build a structured schema from the dataclass TYPE, not an instance.
base = OmegaConf.structured(config_cls)
cfg = base
# Load stage file
if args.config is not None:
if not os.path.isfile(args.config):
raise FileNotFoundError(f"Config file not found: {args.config} (cwd: {os.getcwd()})")
stage_file_cfg = OmegaConf.load(args.config)
cfg = OmegaConf.merge(cfg, stage_file_cfg)
# Load training config and filter to known keys
if args.training_config and os.path.isfile(args.training_config):
training_file_cfg = OmegaConf.load(args.training_config)
# Convert to plain dict to filter
training_container = OmegaConf.to_container(training_file_cfg, resolve=True) or {}
allowed_keys = set(cfg.keys())
filtered_training = {k: v for k, v in training_container.items() if k in allowed_keys and v is not None}
if filtered_training:
training_cfg_filtered = OmegaConf.create(filtered_training)
# Training config takes priority over stage
cfg = OmegaConf.merge(cfg, training_cfg_filtered)
# Merge CLI overrides if any (ignore leading '--')
dot_overrides = [s.lstrip('-') for s in (args.overrides or []) if '=' in s]
if dot_overrides:
cli_cfg = OmegaConf.from_dotlist(dot_overrides)
cfg = OmegaConf.merge(cfg, cli_cfg)
# Return a typed dataclass instance
return OmegaConf.to_object(cfg)
```
## /utils/distributed.py
```py path="/utils/distributed.py"
import os
from models.utils import ModelType
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh, DeviceMesh
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import fully_shard
from typing import Dict, Iterable
from utils.config import DistributedConfig
def init_distributed_from_env() -> Dict[str, object]:
"""Initialize DeviceMesh from torchrun env vars.
Returns a context dict with is_distributed, world_size, is_main, device_mesh.
"""
world_size = int(os.environ.get('WORLD_SIZE', '1'))
is_distributed = world_size > 1 and torch.cuda.is_available()
if is_distributed:
device_mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=('fsdp',))
else:
device_mesh = None
rank = device_mesh.get_rank() if is_distributed else 0
is_main = (rank == 0)
return {
'is_distributed': is_distributed,
'world_size': world_size,
'is_main': is_main,
'device_mesh': device_mesh,
}
def prepare_model_for_distributed(model: torch.nn.Module, config: DistributedConfig, model_type: ModelType, device_mesh: DeviceMesh) -> torch.nn.Module:
if config.use_ddp:
return DDP(model, device_ids=[device_mesh.get_local_rank()], output_device=device_mesh.get_local_rank(), find_unused_parameters=False)
if config.use_fsdp:
mp_policy = config.get_mixed_precision_policy()
fsdp_kwargs = {
"reshard_after_forward": config.reshard_after_forward,
"mp_policy": mp_policy,
"offload_policy": config.offload_policy,
}
def shard_layers(layers: Iterable[torch.nn.Module]) -> None:
for layer in layers:
fully_shard(layer, mesh=device_mesh, **fsdp_kwargs)
if model_type in {ModelType.VideoTokenizer, ModelType.LatentActionModel}:
shard_layers(model.encoder.transformer.blocks)
if model_type == ModelType.VideoTokenizer:
shard_layers(model.encoder.latent_head)
if model_type == ModelType.LatentActionModel:
shard_layers(model.encoder.action_head)
shard_layers(model.decoder.transformer.blocks)
shard_layers(model.decoder.frame_head)
fully_shard(model.encoder, mesh=device_mesh, **fsdp_kwargs)
fully_shard(model.decoder, mesh=device_mesh, **fsdp_kwargs)
fully_shard(model.quantizer, mesh=device_mesh, **fsdp_kwargs)
elif model_type == ModelType.DynamicsModel:
shard_layers(model.transformer.blocks)
fully_shard(model.latent_embed, mesh=device_mesh, **fsdp_kwargs)
fully_shard(model.output_mlp, mesh=device_mesh, **fsdp_kwargs)
else:
raise ValueError('Unknown model type')
fully_shard(
model,
mesh=device_mesh,
**fsdp_kwargs,
)
return model
def unwrap_model(model: torch.nn.Module) -> torch.nn.Module:
return model.module if isinstance(model, DDP) else model
def print_param_count_if_main(model: torch.nn.Module, model_name: str, is_main: bool) -> None:
if not is_main:
return
try:
params = sum(p.numel() for p in model.parameters())
print(f"{model_name} parameters: {params/1e6:.2f}M ({params})")
except Exception:
pass
def cleanup_distributed(is_distributed: bool) -> None:
if is_distributed and dist.is_initialized():
dist.destroy_process_group()
```
## /utils/inference_utils.py
```py path="/utils/inference_utils.py"
import torch
import time
import os
import matplotlib.pyplot as plt
import cv2
import numpy as np
from utils.utils import load_videotokenizer_from_checkpoint, load_latent_actions_from_checkpoint, load_dynamics_from_checkpoint
from einops import repeat
def load_models(video_tokenizer_path, latent_actions_path, dynamics_path, device, use_actions=True):
# Load tokenizer and dynamics, and Latent Actions if using actions
video_tokenizer, _vt_ckpt = load_videotokenizer_from_checkpoint(video_tokenizer_path, device)
video_tokenizer.eval()
latent_action_model = None
if use_actions:
latent_action_model, _latent_action_ckpt = load_latent_actions_from_checkpoint(latent_actions_path, device)
latent_action_model.eval()
dynamics_model, _dyn_ckpt = load_dynamics_from_checkpoint(dynamics_path, device)
dynamics_model.eval()
return video_tokenizer, latent_action_model, dynamics_model
def visualize_inference(predicted_frames, ground_truth_frames, inferred_actions, fps, use_actions=True):
# Move to CPU and convert to numpy
predicted_frames = predicted_frames.detach().cpu()
ground_truth_frames = ground_truth_frames.detach().cpu()
# Denormalize frames from [-1, 1] to [0, 1]
predicted_frames = (predicted_frames + 1) / 2
predicted_frames = torch.clamp(predicted_frames, 0, 1)
ground_truth_frames = (ground_truth_frames + 1) / 2
ground_truth_frames = torch.clamp(ground_truth_frames, 0, 1)
# Get dimensions
B, T, C, H, W = predicted_frames.shape
_, num_gt_frames, _, _, _ = ground_truth_frames.shape
# Create figure with ground truth and predictions side by side
fig, axes = plt.subplots(2, T, figsize=(4 * T, 8))
# Handle single subplot case
if T == 1:
axes = axes.reshape(2, 1)
# Plot ground truth frames (top row)
for i in range(num_gt_frames):
frame = ground_truth_frames[0, i].permute(1, 2, 0).numpy() # [H, W, C]
axes[0, i].imshow(frame)
axes[0, i].set_title(f'Ground Truth {i+1}', fontsize=12, color='green')
axes[0, i].axis('off')
# Plot predicted frames (bottom row)
for i in range(T):
frame = predicted_frames[0, i].permute(1, 2, 0).numpy() # [H, W, C]
axes[1, i].imshow(frame)
title = f'Predicted {i+1}'
if use_actions and i < len(inferred_actions):
title += f'\nAction {inferred_actions[i].item()}' if i < len(inferred_actions) else ''
axes[1, i].set_title(title, fontsize=12, color='red')
axes[1, i].axis('off')
plt.suptitle('Ground Truth vs Predicted Frames', fontsize=16, fontweight='bold')
# Save the visualization
timestamp = time.strftime("%Y%m%d_%H%M%S")
save_dir = "inference_results"
os.makedirs(save_dir, exist_ok=True)
if use_actions:
save_path = os.path.join(save_dir, f"inference_results_gt_vs_pred_{timestamp}.png")
mp4_path = os.path.join(save_dir, f"inference_video_{timestamp}.mp4")
else:
save_path = os.path.join(save_dir, f"inference_results_gt_vs_pred_no_actions_{timestamp}.png")
mp4_path = os.path.join(save_dir, f"inference_video_no_actions_{timestamp}.mp4")
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.close()
print(f"Visualization saved to: {save_path}")
all_frames = torch.cat([ground_truth_frames, predicted_frames], dim=1)
save_frames_as_mp4(all_frames, mp4_path, fps)
# Calculate and display reconstruction error
mse_error = torch.mean((predicted_frames - ground_truth_frames) ** 2).item()
print(f"\nInference stats:")
print(f"Total frames generated: {T}")
print(f"Mean Squared Error (GT vs Pred): {mse_error:.6f}")
if use_actions:
print(f"Actions used: {[action.item() for action in inferred_actions]}")
else:
print(f"No actions used.")
# TODO: get working mp4
def save_frames_as_mp4(frames, output_path, fps=2):
B, T, C, H, W = frames.shape
# OpenCV expects (W, H)
fourcc = cv2.VideoWriter_fourcc(*'avc1')
out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
for i in range(T):
frame = frames[0, i].detach().cpu().permute(1, 2, 0).numpy() # [H, W, C]
# Ensure float32
frame = frame.astype(np.float32)
# Clamp and scale
frame = np.clip(frame, 0, 1)
frame = (frame * 255).astype(np.uint8)
# If grayscale, convert to 3 channels
if frame.shape[2] == 1:
frame = np.repeat(frame, 3, axis=2)
# Convert RGB to BGR for OpenCV
frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
out.write(frame_bgr)
out.release()
print(f"MP4 video saved to: {output_path}")
def sample_random_action(n_actions):
random_action = torch.randint(0, n_actions, (1,))
return random_action
def get_action_latent(args, inferred_actions, n_actions, context_frames, latent_action_model, step):
if args.use_interactive_mode: # let user input actions
print("using interactive mode")
user_input = input(f"Enter action id [0..{n_actions-1}] for step {step+1}: ").strip()
assert user_input.isdigit() and 0 <= int(user_input) < n_actions, f"Invalid input. Please enter an integer in [0,{n_actions-1}]"
val = int(user_input)
sampled_action_index = torch.tensor([val], device=args.device)
inferred_actions.append(sampled_action_index)
recent = inferred_actions[-args.context_window:] if len(inferred_actions) > args.context_window else inferred_actions
recent_tensor = repeat(torch.tensor(recent, device=args.device), 'i -> 1 i') # [1, i or T_ctx]
action_latent = latent_action_model.quantizer.get_latents_from_indices(recent_tensor)
if args.prediction_horizon > 1:
action_latent = repeat(action_latent, 'b 1 a -> b ph a', ph=args.prediction_horizon)
if len(recent) < args.context_window:
gt_pad_actions = latent_action_model.encode(context_frames[:, :args.context_window - len(recent) + 1])
quantized_gt_pad_actions = latent_action_model.quantizer(gt_pad_actions)
action_latent = torch.cat([quantized_gt_pad_actions, action_latent], dim=1)
elif args.use_gt_actions: # use action tokenizer actions
print("using gt actions")
gt_action_latents = latent_action_model.encode(context_frames) # [1, T - 1, A]
sampled_action_index = sample_random_action(n_actions) # [1]
inferred_actions.append(sampled_action_index) # [i]
sampled_action_index_tensor = repeat(torch.tensor(sampled_action_index, device=args.device), 'i -> 1 i') # [1, i]
sampled_action_latent = latent_action_model.quantizer.get_latents_from_indices(sampled_action_index_tensor) # [1, i, A]
action_latent = torch.cat([gt_action_latents, sampled_action_latent], dim=1) # [1, T, A]
elif args.use_actions: # use random actions
print(f"using random actions")
sampled_action_index = sample_random_action(n_actions) # [1]
inferred_actions.append(sampled_action_index) # [i]
recent_inferred_actions = inferred_actions[-args.context_window:] if len(inferred_actions) > args.context_window else inferred_actions # [i or T_ctx]
recent_inferred_actions_tensor = repeat(torch.tensor(recent_inferred_actions, device=args.device), 'i -> 1 i') # [1, i or T_ctx]
action_latent = latent_action_model.quantizer.get_latents_from_indices(recent_inferred_actions_tensor) # [1, i or T_ctx, A]
if args.prediction_horizon > 1:
action_latent = repeat(action_latent, 'b 1 a -> b ph a', ph=args.prediction_horizon)
if len(recent_inferred_actions) < args.context_window:
# if we dont have enough inferred actions (in the beginning) add enough gt to fill the sequence
gt_pad_actions = latent_action_model.encode(context_frames[:, :args.context_window - len(recent_inferred_actions) + 1]) # [1, context_window - len(inferred_actions), A]
quantized_gt_pad_actions = latent_action_model.quantizer(gt_pad_actions) # [1, context_window - len(inferred_actions), A]
action_latent = torch.cat([quantized_gt_pad_actions, action_latent], dim=1) # [1, S, A]
else:
sampled_action_index = None
action_latent = None
return sampled_action_index, action_latent
```
## /utils/optimizer_utils.py
```py path="/utils/optimizer_utils.py"
from __future__ import annotations
import torch
import torch.optim as optim
def create_optimizer(model, args):
from torch.nn.parallel import DistributedDataParallel as DDP
raw_model = model.module if isinstance(model, DDP) else model
optimizer_name = getattr(args, "optimizer", "adamw")
if optimizer_name == "muon":
return _create_muon_split(raw_model, args)
else:
return _create_adamw(raw_model, args)
def _create_adamw(model, args):
decay, no_decay = _split_decay_params(model)
optimizer = optim.AdamW([
{"params": decay, "weight_decay": 0.01},
{"params": no_decay, "weight_decay": 0},
], lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-8, fused=True)
return [optimizer]
def _create_muon_split(model, args):
from models.muon import Muon
muon_params = []
adamw_decay = []
adamw_no_decay = []
for name, param in model.named_parameters():
if not param.requires_grad:
continue
# Muon only makes sense for 2D weight matrices (not embeddings, not biases)
if param.ndim == 2 and "embed" not in name:
muon_params.append(param)
elif param.ndim == 1 or name.endswith(".bias") or "norm" in name:
adamw_no_decay.append(param)
else:
adamw_decay.append(param)
lr = args.learning_rate
momentum = getattr(args, "muon_momentum", 0.95)
backend_steps = getattr(args, "muon_backend_steps", 5)
optimizers = []
if muon_params:
muon_opt = Muon(
muon_params, lr=lr, momentum=momentum,
backend_steps=backend_steps, weight_decay=0.01,
)
optimizers.append(muon_opt)
# AdamW for the rest
adamw_groups = []
if adamw_decay:
adamw_groups.append({"params": adamw_decay, "weight_decay": 0.01})
if adamw_no_decay:
adamw_groups.append({"params": adamw_no_decay, "weight_decay": 0})
if adamw_groups:
adamw_opt = optim.AdamW(adamw_groups, lr=lr, betas=(0.9, 0.999), eps=1e-8, fused=True)
optimizers.append(adamw_opt)
return optimizers
def _split_decay_params(model):
decay = []
no_decay = []
for name, param in model.named_parameters():
if param.requires_grad:
if len(param.shape) == 1 or name.endswith(".bias") or "norm" in name:
no_decay.append(param)
else:
decay.append(param)
return decay, no_decay
```
## /utils/scheduler_utils.py
```py path="/utils/scheduler_utils.py"
import math
import torch.optim as optim
def cosine_with_warmup(step, *, warmup_steps, total_steps, min_lr=0.0):
if step < warmup_steps:
return step / max(1, warmup_steps)
# calculate step within warmup
progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
# calculate where on cosine curve we are
cosine = 0.5 * (1 + math.cos(math.pi * progress))
return min_lr + (1 - min_lr) * cosine
def create_cosine_scheduler(optimizer, total_steps, warmup_fraction=0.05, min_lr_fraction=0.01):
warmup_steps = int(total_steps * warmup_fraction)
return optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda s: cosine_with_warmup(
s,
warmup_steps=warmup_steps,
total_steps=total_steps,
min_lr=min_lr_fraction
)
)
```
## /utils/utils.py
```py path="/utils/utils.py"
from pathlib import Path
import time
import glob
import subprocess
import os
import re
from typing import Optional
from torch.distributed.checkpoint.state_dict import (
get_model_state_dict,
get_optimizer_state_dict,
set_model_state_dict,
StateDictOptions,
)
from torch.distributed.fsdp import FSDPModule
from torch.nn.parallel import DistributedDataParallel as DDP
MODEL_CHECKPOINT = "model_state_dict.pt"
OPTIMIZER_CHECKPOINT = "optim_state_dict.pt"
STATE = "state.pt"
def readable_timestamp():
"""Generate a sortable timestamp for filenames (no weekday)."""
return time.strftime("%Y_%m_%d_%H_%M_%S")
def find_latest_checkpoint(base_dir, model_name, run_root_dir: Optional[str] = None, stage_name: Optional[str] = None):
"""Find latest checkpoint.
If run_root_dir (and optional stage_name) are provided, search only under
<run_root_dir>/<stage_name>/checkpoints (or <run_root_dir>/**/checkpoints if stage_name None).
Otherwise, fall back to project-wide model type-based search.
Newest run dir first, then highest step within it.
If the newest run root has none for this model, keep searching older runs.
"""
def collect_checkpoint_paths(roots, model_name):
alias_map = {
'video_tokenizer': ['video_tokenizer'],
'latent_actions': ['latent_actions', 'lam', 'actions', 'action_tokenizer'],
'dynamics': ['dynamics'],
}
aliases = alias_map.get(model_name, [model_name])
candidates = []
seen = set()
def add_candidate(path: str) -> None:
norm = os.path.normpath(path)
if norm in seen:
return
if os.path.isdir(norm):
# require at least state or model file
state_file = os.path.join(norm, STATE)
model_file = os.path.join(norm, MODEL_CHECKPOINT)
if os.path.isfile(state_file) or os.path.isfile(model_file):
candidates.append(norm)
seen.add(norm)
else:
_, ext = os.path.splitext(norm)
if ext in ('.pt', '.pth'):
candidates.append(norm)
seen.add(norm)
patterns = []
for alias in aliases:
patterns.append(f"*{alias}_step_*")
patterns.append(f"*{alias}_checkpoint_*")
for root in roots:
for pat in patterns:
search_pattern = os.path.join(root, "**", pat)
for match in glob.glob(search_pattern, recursive=True):
add_candidate(match)
return candidates
def run_dir_of(path: str) -> str:
# Walk up until we reach the 'checkpoints' directory, then return its parent (the stage dir)
d = path if os.path.isdir(path) else os.path.dirname(path)
while d and os.path.basename(d) != 'checkpoints':
parent = os.path.dirname(d)
if parent == d:
break
d = parent
# If we found 'checkpoints', return its parent; otherwise, fallback to two-level up
if d and os.path.basename(d) == 'checkpoints':
return os.path.dirname(d)
# Fallback: use the directory containing the checkpoint path (or its parent)
candidate_dir = path if os.path.isdir(path) else os.path.dirname(path)
return candidate_dir
def project_wide_search():
# Fallback to model_type/results search in repository
# Allow multiple possible directory names per model
model_type_dirs = {
'video_tokenizer': ['video_tokenizer'],
'latent_actions': ['latent_actions', 'lam', 'actions', 'action_tokenizer'],
'dynamics': ['dynamics',],
}
dirs = model_type_dirs.get(model_name)
if not dirs:
roots = [os.path.join(base_dir, 'results', '**', 'checkpoints')]
else:
roots = [os.path.join(base_dir, 'results', '**', d, 'checkpoints') for d in dirs]
files = collect_checkpoint_paths(roots, model_name)
if not files:
# Generic fallback: search all checkpoints regardless of stage dir name
generic_roots = [os.path.join(base_dir, 'results', '**', 'checkpoints')]
files = collect_checkpoint_paths(generic_roots, model_name)
if not files:
raise Exception(f"No checkpoint files found for {model_name}")
run_dir_to_files = {}
for p in files:
rd = run_dir_of(p)
run_dir_to_files.setdefault(rd, []).append(p)
newest_run_dir = max(run_dir_to_files.keys(), key=lambda d: os.path.getctime(d))
candidate_files = run_dir_to_files[newest_run_dir]
return candidate_files
if run_root_dir is not None:
if stage_name:
roots = [os.path.join(run_root_dir, stage_name, 'checkpoints')]
else:
roots = [os.path.join(run_root_dir, '**', 'checkpoints')]
files = collect_checkpoint_paths(roots, model_name)
if not files:
# Fallback: search project-wide older runs until found
candidate_files = project_wide_search()
else:
# Group by run dir within the provided root and choose newest run dir
run_dir_to_files = {}
for p in files:
rd = run_dir_of(p)
run_dir_to_files.setdefault(rd, []).append(p)
newest_run_dir = max(run_dir_to_files.keys(), key=lambda d: os.path.getctime(d))
candidate_files = run_dir_to_files[newest_run_dir]
else:
candidate_files = project_wide_search()
def extract_step(path: str) -> int:
fname = os.path.basename(path)
m = re.search(r"_step_(\d+)", fname)
return int(m.group(1)) if m else -1
candidate_files.sort(key=lambda p: (extract_step(p), os.path.getctime(p)))
return candidate_files[-1]
def run_command(cmd, description):
# empirical max dataLoader throughput settings (I used 1-6 H100s)
env = os.environ.copy()
env.setdefault("NG_NUM_WORKERS", str(max(2, (os.cpu_count() or 4) - 2)))
env.setdefault("NG_PREFETCH_FACTOR", "4")
env.setdefault("NG_PIN_MEMORY", "1")
env["NG_PERSISTENT_WORKERS"] = "0"
env.setdefault("TORCH_CUDNN_V8_API_ENABLED", "1")
try:
result = subprocess.run(cmd, check=True, capture_output=False, env=env)
return True
except subprocess.CalledProcessError as e:
print(f"Error: {e.stderr}")
return False
except KeyboardInterrupt:
return False
def save_training_state(model, optimizer, scheduler, config, checkpoints_dir, prefix, step):
"""Save a checkpoint with model/optimizer/scheduler and the exact config.
The filename includes the global step and a timestamp for uniqueness.
"""
import torch
ts = readable_timestamp()
if isinstance(model, (FSDPModule, DDP)):
state_dict = get_model_state_dict(
model=model,
options=StateDictOptions(
full_state_dict=True,
cpu_offload=True,
),
)
optimizer_state_dict = get_optimizer_state_dict(
model=model,
optimizers=optimizer,
options=StateDictOptions(
full_state_dict=True,
cpu_offload=True,
),
)
else:
# Avoid saving the model with _orig_mod prefix if it's compiled
state_dict = getattr(model, '_orig_mod', model).state_dict()
optimizer_state_dict = optimizer.state_dict()
state = {
'scheduler_state_dict': scheduler.state_dict() if scheduler is not None else None,
'config': config,
'step': int(step) if step is not None else None,
'timestamp': ts,
}
os.makedirs(checkpoints_dir, exist_ok=True)
ckpt_path = os.path.join(checkpoints_dir, f"{prefix}_step_{int(step) if step is not None else 0}")
os.makedirs(ckpt_path, exist_ok=True)
torch.save(state_dict, Path(ckpt_path) / MODEL_CHECKPOINT)
torch.save(optimizer_state_dict, Path(ckpt_path) / OPTIMIZER_CHECKPOINT)
torch.save(state, Path(ckpt_path) / STATE)
return ckpt_path
def load_videotokenizer_from_checkpoint(checkpoint_path, device, model = None, is_distributed = False):
"""Instantiate VideoTokenizer from a checkpoint's saved config and load weights."""
import torch
from models.video_tokenizer import VideoTokenizer
model_sd = torch.load(Path(checkpoint_path) / MODEL_CHECKPOINT, map_location='cpu', weights_only=True)
state_cfg = torch.load(Path(checkpoint_path) / STATE, map_location='cpu', weights_only=False)
cfg = state_cfg.get('config', {}) or {}
frame_size = cfg.get('frame_size', 128)
kwargs = {
'frame_size': (frame_size, frame_size),
'patch_size': cfg.get('patch_size', 8),
'embed_dim': cfg.get('embed_dim', 128),
'num_heads': cfg.get('num_heads', 8),
'hidden_dim': cfg.get('hidden_dim', 256),
'num_blocks': cfg.get('num_blocks', 4),
'latent_dim': cfg.get('latent_dim', 6),
'num_bins': cfg.get('num_bins', 4),
}
if model is None:
model = VideoTokenizer(**kwargs)
set_model_state_dict(
model=model,
model_state_dict=model_sd,
options=StateDictOptions(
full_state_dict=True,
broadcast_from_rank0=is_distributed,
),
)
model = model.to(device)
return model, state_cfg
def load_latent_actions_from_checkpoint(checkpoint_path, device, model = None, is_distributed = False):
"""Instantiate LatentActionModel from a checkpoint's saved config and load weights."""
import torch
from models.latent_actions import LatentActionModel
model_sd = torch.load(Path(checkpoint_path) / MODEL_CHECKPOINT, map_location='cpu', weights_only=True)
state_cfg = torch.load(Path(checkpoint_path) / STATE, map_location='cpu', weights_only=False)
cfg = state_cfg.get('config', {}) or {}
frame_size = cfg.get('frame_size', 128)
kwargs = {
'frame_size': (frame_size, frame_size),
'n_actions': cfg.get('n_actions', 8),
'patch_size': cfg.get('patch_size', 8),
'embed_dim': cfg.get('embed_dim', 128),
'num_heads': cfg.get('num_heads', 8),
'hidden_dim': cfg.get('hidden_dim', 256),
'num_blocks': cfg.get('num_blocks', 4),
}
if model is None:
model = LatentActionModel(**kwargs)
set_model_state_dict(
model=model,
model_state_dict=model_sd,
options=StateDictOptions(
full_state_dict=True,
broadcast_from_rank0=is_distributed,
),
)
model = model.to(device)
return model, state_cfg
def load_dynamics_from_checkpoint(checkpoint_path, device, model = None, is_distributed = False):
"""Instantiate DynamicsModel from a checkpoint's saved config and load weights."""
import torch
from models.dynamics import DynamicsModel
model_sd = torch.load(Path(checkpoint_path) / MODEL_CHECKPOINT, map_location='cpu', weights_only=True)
state_cfg = torch.load(Path(checkpoint_path) / STATE, map_location='cpu', weights_only=False)
cfg = state_cfg.get('config', {}) or {}
frame_size = cfg.get('frame_size', 128)
# Infer conditioning_dim from checkpoint if missing
conditioning_dim = cfg.get('conditioning_dim', None)
if conditioning_dim is None:
cond_inferred = None
for k, v in model_sd.items():
# Linear weight shape: [out_features, in_features]; in_features is conditioning dim
if k.endswith('to_gamma_beta.1.weight'):
cond_inferred = int(v.shape[1])
break
conditioning_dim = cond_inferred if cond_inferred is not None else 3
kwargs = {
'frame_size': (frame_size, frame_size),
'patch_size': cfg.get('patch_size', 8),
'embed_dim': cfg.get('embed_dim', 128),
'num_heads': cfg.get('num_heads', 8),
'hidden_dim': cfg.get('hidden_dim', 256),
'num_blocks': cfg.get('num_blocks', 4),
'conditioning_dim': conditioning_dim,
'latent_dim': cfg.get('latent_dim', 6),
'num_bins': cfg.get('num_bins', 4),
'use_moe': cfg.get('use_moe', False),
'num_experts': cfg.get('num_experts', 4),
'top_k_experts': cfg.get('top_k_experts', 2),
'moe_aux_loss_coeff': cfg.get('moe_aux_loss_coeff', 0.01),
}
if model is None:
model = DynamicsModel(**kwargs)
set_model_state_dict(
model=model,
model_state_dict=model_sd,
options=StateDictOptions(
full_state_dict=True,
broadcast_from_rank0=is_distributed,
)
)
model = model.to(device)
return model, state_cfg
def prepare_pipeline_run_root(run_name: Optional[str] = None, base_cwd: Optional[str] = None):
"""Create a top-level run root directory results/<timestamp_or_name>"""
cwd = base_cwd or os.getcwd()
ts = readable_timestamp()
name = run_name or ts
run_root = os.path.join(cwd, 'results', name)
os.makedirs(run_root, exist_ok=True)
return run_root, name
def prepare_stage_dirs(run_root_dir: str, stage_name: str):
"""Create stage subdirectories under the given run root.
Structure:
<run_root_dir>/<stage_name>/checkpoints
<run_root_dir>/<stage_name>/visualizations
Returns (stage_dir, checkpoints_dir, visualizations_dir).
"""
stage_dir = os.path.join(run_root_dir, stage_name)
checkpoints_dir = os.path.join(stage_dir, 'checkpoints')
visualizations_dir = os.path.join(stage_dir, 'visualizations')
os.makedirs(checkpoints_dir, exist_ok=True)
os.makedirs(visualizations_dir, exist_ok=True)
return stage_dir, checkpoints_dir, visualizations_dir
```
## /utils/wandb_utils.py
```py path="/utils/wandb_utils.py"
import wandb
import torch
import os
import time
from typing import Dict, Any, Optional
def init_wandb(project_name: str, config: Dict[str, Any], run_name: Optional[str] = None) -> wandb.run:
# Generate run name if not provided
if run_name is None:
run_name = f"{project_name}_{time.strftime('%Y%m%d_%H%M%S')}"
# Initialize wandb
run = wandb.init(
project=project_name,
config=config,
name=run_name,
tags=[project_name, "training"]
)
print(f"🚀 W&B run initialized: {run.name}")
print(f"📊 Project: {project_name}")
print(f"🔗 View at: {run.url}")
return run
def log_training_metrics(step: int, metrics: Dict[str, float], prefix: str = "train"):
# Add prefix to metric names
prefixed_metrics = {f"{prefix}/{k}": (v.item() if hasattr(v, "item") else float(v)) for k, v in metrics.items()}
wandb.log(prefixed_metrics, step=step)
def log_learning_rate(optimizer: torch.optim.Optimizer, step: int):
for i, param_group in enumerate(optimizer.param_groups):
wandb.log({
f"learning_rate/group_{i}": float(param_group['lr']),
}, step=step)
def log_codebook_usage(codebook_usage: float, step: int, model_name: str = "model"):
wandb.log({
f"{model_name}/codebook_usage": float(codebook_usage),
}, step=step)
def log_action_distribution(action_indices: torch.Tensor, step: int, n_actions: int):
# Convert to CPU and numpy
flat = action_indices.detach().cpu().reshape(-1)
# Compute distribution
counts = torch.bincount(flat.long(), minlength=n_actions).float()
probs = counts / counts.sum().clamp_min(1)
wandb.log({
"action_distribution": wandb.Histogram(flat.numpy()),
"action_entropy": float(-(probs * (probs + 1e-8).log()).sum()),
"unique_actions": int((counts > 0).sum().item()),
}, step=step)
def log_system_metrics(step: int):
if torch.cuda.is_available():
wandb.log({
"system/gpu_memory_allocated": float(torch.cuda.memory_allocated() / 1024**3), # GB
"system/gpu_memory_reserved": float(torch.cuda.memory_reserved() / 1024**3), # GB
}, step=step)
def finish_wandb():
"""Finish the W&B run"""
if wandb.run is not None:
wandb.finish()
def create_wandb_config(args, model_config: Dict[str, Any]) -> Dict[str, Any]:
config = {
# Training parameters
"batch_size": args.batch_size,
"n_updates": args.n_updates,
"learning_rate": args.learning_rate,
"log_interval": getattr(args, 'log_interval', 100),
# Dataset parameters
"dataset": getattr(args, 'dataset', 'SONIC'),
"context_length": getattr(args, 'context_length', 4),
# Model architecture
"model_architecture": model_config,
# System parameters
"device": str(torch.device("cuda" if torch.cuda.is_available() else "cpu")),
"timestamp": time.strftime("%Y%m%d_%H%M%S"),
}
# Add any additional args that might exist
for attr in dir(args):
if not attr.startswith('_') and not callable(getattr(args, attr)):
if attr not in config:
config[attr] = getattr(args, attr)
return config
```
The better and more specific the context, the better the LLM can follow instructions. If the context seems verbose, the user can refine the filter using uithub. Thank you for using https://uithub.com - Perfect LLM context for any GitHub repo.