RoboJEPA-Decoder-0.6B

A 0.6 B-parameter DiT diffusion decoder that turns frozen V-JEPA 2.1 features back into video. It is the pixel decoder of the robo_jepa action-conditioned robot world model: the world model predicts future V-JEPA features, and this head renders them as a 33-frame 240x320 clip so you can watch what the latent world model imagines.

It decodes encoder features (train_heads_on_predictor: false), so it works just as well as a standalone "V-JEPA features -> video" inverter, independent of any predictor.

Checkpoint: 219 400 optimization steps (epoch 1097).

Architecture

kind	DiT (transformer diffusion), x0-prediction
conditioning	V-JEPA 2.1 ViT-G/384 features, 1664-dim, grid 15x20, 9 frames (see Usage)
latent space	NVIDIA Cosmos-0.1-Tokenizer-CV4x8x8, 16 channels, 9 x 30x40
output	33 frames at 240x320 (8x spatial, 4x temporal detokenization); any 16h x 16w at test time, see Resolution
depth / width	32 blocks, 1280-dim, 20 heads, 10 KV heads (GQA), MLP ratio 4
position encoding	RoPE; local temporal attention window of 4 frames
diffusion	cosine_simple_diffusion, zero-terminal-SNR, 1000 train steps, 128 at sampling
condition dropout	0.1 (classifier-free guidance capable)
precision on disk	fp32

Parameter count

component	params
blocks (32 DiT blocks)	577.45 M
timestep_embedding	11.81 M
null_y (CFG null condition)	4.49 M
cond_mixing (1664 -> 1280)	3.77 M
out_proj	0.02 M
patch_embed	0.02 M
total (this repo)	597 568 016

The Cosmos tokenizer it decodes through is frozen and not included here (30.1 M encoder + 43.1 M decoder); download it separately, see Usage. The V-JEPA encoder that produces the conditioning features is also separate and frozen.

Files

file	what it is
model.safetensors	the 494 decoder tensors, fp32, 2.39 GB. Keys match WorldModelDiTCosmosHead.ema_model (the DDP module. prefix has been stripped)
config.json	head config in the exact shape WorldModel.load_pretrained expects for decoder/config.json, plus epoch / global_step / num_parameters

Optimizer state is not included, so this is an inference/fine-tuning-from-weights artifact, not a resumable training checkpoint.

sha256(model.safetensors) = 4954d5fc2150cd759a3e961d829a8a07f5b82b6146573b9602be6c8268bb76f6

Usage

You need the robo_jepa code (app.robo_jepa.world_model_heads) and the Cosmos tokenizer:

# the frozen Cosmos CV4x8x8 tokenizer (NVIDIA, separate license)
huggingface-cli download nvidia/Cosmos-0.1-Tokenizer-CV4x8x8 \
    encoder.jit decoder.jit --local-dir cosmos/

import json, torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from app.robo_jepa.world_model_heads import WorldModelDiTCosmosHead

REPO = "artemZholus/RoboJEPA-Decoder-0.6B"
cfg = json.load(open(hf_hub_download(REPO, "config.json")))
for k in ("kind", "epoch", "global_step", "num_parameters"):
    cfg.pop(k, None)
diffusion_kwargs = cfg.pop("diffusion_kwargs")
cfg["encoder_path"] = "cosmos/encoder.jit"      # where you downloaded the tokenizer
cfg["decoder_path"] = "cosmos/decoder.jit"

head = WorldModelDiTCosmosHead(head_config=cfg, device="cuda", diffusion_kwargs=diffusion_kwargs)
sd = load_file(hf_hub_download(REPO, "model.safetensors"))
missing, unexpected = head.ema_model.load_state_dict(sd, strict=False)
assert not unexpected, unexpected          # fail loudly instead of decoding random weights
head.eval()

# features: [B, T, V, H, W, C] V-JEPA 2.1 features; T = 9, V = 1, H, W = 15, 20, C = 1664
video = head.decode(features, denoising_steps=128)   # -> uint8 [B, 33, V, 240, 320, 3]

Decode as the features arrive, one frame at a time, by carrying the KV cache:

cache = None
for t in range(features.shape[1]):
    frame, cache = head.decode(features[:, t : t + 1], denoising_steps=4,
                               use_cache=True, cache=cache)   # uint8 [B, 1, V, 240, 320, 3]

Three things that will silently or loudly bite you:

• Pass exactly 9 feature frames. The sampler always builds num_frames + 1 = 9 Cosmos latent frames (the tokenizer's extra first frame), while the attention mask is keyed off the conditioning depth, so 8 frames in gives RuntimeError: The expanded size of the tensor (10800) must match the existing size (9600). Upstream does the same thing: decode(video_features[:, -eval_rollout_steps - 1:]) with eval_rollout_steps = 8.

• The weights go into head.ema_model, not head.model. For kind: dit that is where load_pretrained puts them and where decode() samples from. During training the EMA is refreshed with decay=0.0 (a straight copy), so EMA and online weights are identical.

• load_state_dict(..., strict=False) is how the head is loaded upstream, which means a key mismatch loads nothing and raises no error. Assert on unexpected (and check missing) as above. The keys here were exported to match, and verified by loading.

• Keep attn_impl: "sdpa". local_attn_window makes the model pass a custom attention mask, and the fmha path raises ValueError: fmha is supported with custom attention mask!. Note the SDPA kernel this needs is not available on every GPU — the same call that works on an A100 fails with RuntimeError: No available kernel on older cards.

The snippet above was run end to end on an A100-80GB before publishing: 494/494 tensors loaded with 0 missing and 0 unexpected keys, bitwise identical to the training checkpoint, and decode on random features returned a uint8 tensor of shape (1, 33, 1, 240, 320, 3). Sampling 1 clip at 8 denoising steps peaked at 6.0 GB of GPU memory.

Resolution at test time

Nothing in the decoder is tied to the 240x320 it was trained at. The conditioning features are nearest-upsampled onto the latent grid and positions come from RoPE, so a h x w feature grid renders 16h x 16w video — feed it V-JEPA features from a larger frame and you get a larger clip out, from the same weights.

The one catch is bookkeeping: the grid is read from the module at construction time, so retarget it before decoding (and use a fresh process per resolution — the attention and RoPE caches keep the previous grid):

m = head.ema_model
m.condition_h, m.condition_w = h, w          # your V-JEPA feature grid
m.grid_height, m.grid_width = 2 * h, 2 * w   # Cosmos latent grid
m.img_size = [2 * h, 2 * w]
m.attn_mask_cache.clear()
for mod in m.modules():                      # each attention module keeps its own copy
    if mod is not m and hasattr(mod, "grid_height"):
        mod.grid_height, mod.grid_width = 2 * h, 2 * w

Measured on an A100-80GB, 33-frame clips, 4 denoising steps:

feature grid	video	full-sequence peak GPU	per-frame (use_cache=True) peak GPU
15x20 (trained)	240x320	6.0 GB	6.1 GB
23x30	368x480	9.0 GB	7.3 GB
30x40	480x640	14.2 GB	10.7 GB
38x50	608x800	24.5 GB	10.7 GB

Memory for the full-sequence path grows with the square of the token count, so decode frame by frame with use_cache=True at the larger grids. Without the retargeting above, a non-native grid either raises non-autoregressive inference without attention mask detected!!! or silently mis-infers the grid.

Training

Only this head was trained. The V-JEPA 2.1 ViT-G/384 encoder is frozen, and the world model's predictor was not trained in this run, so the decoder never saw predicted features.

steps	219 400 (epoch 1097)
hardware	1 node x 4 A100-80GB, ~3.02 s/iter, 56 GB peak per GPU
global batch	8 clips/step (2 per GPU, no grad accumulation), 8 feature frames per clip
optimizer	AdamW, constant lr 1e-4, betas (0.9, 0.95), weight decay 0, grad clip 1.0
precision	bf16 mixed precision, torch.compiled
objective	x0-prediction diffusion loss in Cosmos latent space
encoder_image_loss	0.0204 at step 219 400 (0.038 at step 140 600)

Data

A weighted mixture of three real-robot datasets, sampled in proportion to video hours:

dataset	weight	episodes	raw fps	sampled fps	intermediate fps
LeRobot (single-arm SO-100/101)	0.0818	22 998	30	7.5	30
DROID	0.4257	73 338	15	1.875	7.5
AgiBot World	0.4925	34 512	30	7.5	15

intermediate_frames_fps / fps is 4 for every dataset by construction — the ratio sets the clip length the DiT head expects (33 frames), so it cannot differ across datasets in one run.

Augmentation: random resized crop to 240x320 (aspect 1.333-1.777, scale 0.5-1.0), no flip.

Caveats inherited from the data pipeline, worth knowing before you trust a sample:

• The LeRobot mixture is read with a single raw_fps: 30, but ~4.5 % of its episodes are not 30 fps (15, 10, 25, 50, 20, 5), so those clips play at the wrong speed.
• AgiBot proprioception is adapted from end/+effector/ into the loader's flat layout; the quaternion is read as scalar-last (x, y, z, w) and the gripper opening is scaled by 120 without polarity inversion. Neither is confirmed against AgiBot documentation. This affects action/state conditioning, not the pixel decoding path used by this head.
• ~3 % of AgiBot episodes are dexterous-hand (12-DoF) and have their gripper columns zeroed.

Intended use and limits

Made for research: visualizing and debugging latent robot world models, and as a general V-JEPA-feature-to-pixel decoder. It is trained only on third-person and wrist robot-manipulation video at 240x320, and it reconstructs what the features encode — it is not a text-to-video model, has no language conditioning, and will not hold identity or text detail that V-JEPA features drop. Samples of unseen scene types, other resolutions, or other feature extractors are out of distribution. Not validated for any safety-critical or on-robot control use.

Provenance

Built with robo_jepa, config configs/decoder/video_dit_240_3ds_mila.yaml, run video_dit_240_3ds_mila_26_09_23_02_35_30, trained on the Mila cluster.

Third-party components, each under its own license, none redistributed here: V-JEPA 2.1 ViT-G/384 (Meta) for the conditioning features, Cosmos-0.1-Tokenizer-CV4x8x8 (NVIDIA) for the latent space, and the DROID / AgiBot World / LeRobot community datasets for training data. Check those terms before using these weights beyond research — in particular the training datasets' terms, which can be more restrictive than the code's.