JEPA-Guidance → Cosmos context predictor (jepa-guided-diffusion)

JEPA-Guidance is a lightweight predictor bridge that maps a WTS traffic video (plus a text query) into the frozen Cosmos-Reason1 (512, 1024) context space, so the predicted embedding can condition Cosmos-Predict2.5 video generation without running the 7B Reason1 text encoder at inference time.

• Input: video / frame directory + a WTS-format query.
• Output: a (512, 1024) context tensor (the Cosmos cross-attention conditioning set).
• Objective: permutation-invariant set-match Hungarian MSE — Cosmos consumes the context as an unordered set, so the predicted 512 tokens are matched to the target 512 tokens (Hungarian assignment) before MSE.

Architecture

Part	Role	Source
X-Encoder (frozen)	VJEPA 2.1 ViT-L/384 (vjepa2_1_vit_large_384), video mode → visual tokens	torch.hub (facebookresearch/vjepa2)
Compressor (trained)	BLIP-2-style QFormer: 512 query tokens, 4 layers, cross-attn every layer	in model.pt
Backbone (trained)	last 4 layers of Llama-3.2-1B, made bidirectional	arch from hub, weights in model.pt
Head (trained)	output_projection → (512, 1024)	in model.pt
Y-Encoder (frozen, target only)	Cosmos-Reason1-7B full_concat + crossattn_proj	not needed at inference

Only the compressor, the last 4 Llama layers, the vision projection, and the output projection are trained; VJEPA2 and Reason1 are frozen.

Training metrics

Full logs for the run that produced these weights are in training_metrics/. Training ran 28 epochs; the published best_model/ is epoch 24, selected on minimum validation set-match MSE.

Metric (epoch 24)	Value
Train set-match MSE	5.8287
Val set-match MSE	6.7279
Retrieval top-1	0.3606
Retrieval MRR	0.4954
Set top-1	0.3613
Set MRR	0.5049
Alignment diag	0.1603
Alignment off-diag	0.1561
LR at best epoch	2.40e-05

Validation set-match MSE fell from 21.50 (epoch 1) to 6.73 (epoch 24).

training_metrics/epoch_metrics.csv    # 28 rows — per-epoch loss, MSE, LR, grad-norm, retrieval/alignment
training_metrics/batch_metrics.csv    # 4452 rows — per-step trace
training_metrics/events.out.tfevents.*  # TensorBoard event file

tensorboard --logdir training_metrics

Note that align_diag (0.160) and align_offdiag (0.156) sit close together: the cached Reason1 targets are strongly anisotropic, so a large shared mean dominates every vector. Training subtracts a precomputed all-token mean before normalization to recover the per-caption residual.

Repo layout

config.yaml                         # full model/eval config the loader reads
best_model/model.pt                 # trained weights (predictor only)
best_model/model_config.yaml        # model config sidecar
best_model/query_tokenizer/         # Llama query tokenizer
fixed/negative_<view>.pt            # the two fixed negative embeddings (vehicle / overhead)
training_metrics/                   # per-epoch + per-batch CSVs and the TensorBoard events

Usage

Install the JEPA-Guidance package (uv-managed), then pass this repo id as --checkpoint — it is downloaded and cached automatically:

# whole test set: writes <scenario>/embedding.pt for the Cosmos handoff
uv run jepa-guidance seq-infer --checkpoint AlterraLaniakea/jepa-guided-diffusion --test-root wts/test

# single scenario -> print the (512,1024) embedding
uv run jepa-guidance infer \
  --checkpoint AlterraLaniakea/jepa-guided-diffusion \
  --visual-path wts/test/<scenario>/input \
  --output embedding \
  --query "The current preset is from overhead view with 2 target subjects. Pedestrian: A man in his 30s stands on the road facing the oncoming vehicle. Vehicle: The vehicle goes straight at a constant speed."

Query format. For in-distribution embeddings, use the same query format as training (seq-infer builds this automatically from caption.json):

The current preset is from <overhead|vehicle> view with <N> target subject(s). Pedestrian: <ped caption> Vehicle: <veh caption>

• <overhead|vehicle> — overhead for CCTV, vehicle for dashcam (video*) clips.
• <N> — subject types present: 1 (pedestrian or vehicle) or 2 (both).

Requirements

• Gated Llama-3.2-1B: run hf auth login (or export HF_TOKEN) before first use — the loader rebuilds the Llama backbone architecture from the hub and overlays the trained weights. (Set hf_token: true in config.yaml, or provide the token via env.)
• VJEPA2 is fetched from facebookresearch/vjepa2 (torch.hub / GitHub source).
• A GPU is recommended; VJEPA2 runs at crop size 384.

Downstream

Feed the predicted embedding.pt into Cosmos-Predict2.5 via scripts/generate_cosmos.py (see the repo README) to generate the WTS video for each scenario.