latent-lab — Two Rooms JEPA world models (ONNX)

Action-conditioned JEPA world models trained on a "Two Rooms" navigation environment, exported to ONNX for in-browser inference (onnxruntime-web, WebGPU/WASM). Part of latent-lab, an interactive playground for understanding JEPA world models and latent planning.

▶ Live demo (these exact weights, running in your browser): https://adimunot21.github.io/latent-lab/

What's here

• models/<id>/encoder.onnx — CNN encoder: 64x64 grayscale frame -> 128-d latent (0.91M params)
• models/<id>/predictor.onnx — residual MLP: (latent, action) -> next latent (0.13M params)
• models/<id>/*.int8.onnx — weight-only dynamic-int8 variants
• lookup/{states,latents}.bin — latent<->state lookup table (float32 LE) for decoder-free visualization
• manifest.json — normalization stats, env config, PCA projection, per-file sha256

Checkpoints: healthy (MSE + SIGReg, 97% planning success), healthy_early (epoch 1), collapsed (lambda_reg = 0 — deliberate representation collapse, a demo feature), collapsed_early.

How they were trained

Joint-embedding predictive architecture with NO EMA target and NO stop-gradient; collapse is prevented solely by SIGReg (Epps-Pulley characteristic-function statistic on random 1-D projections of the latent batch, pushing toward an isotropic Gaussian). Next-embedding MSE + SIGReg, AdamW, AMP, 15 epochs on 60k transitions from a scripted mixed random/goal-directed policy. Trained on a single GTX 1650 (peak VRAM 0.41 GB).

Recorded metrics (held-out): healthy linear position probe R^2 = 0.9997; CEM planning success 97% (N=100). Collapsed: latent std 0.001 (vs 1.16 healthy), planning 44%. fp32 ONNX parity vs PyTorch < 2e-6 max abs diff; int8 errors recorded in manifest.json.

Limitations

• Toy environment: a 2-DoF point agent; these weights model nothing else.
• The int8 predictor's quantization error (~0.24 max abs) compounds over multi-step rollouts; prefer the fp32 predictor.
• Deterministic env: the world model has never seen stochastic dynamics.

Use

Fetch from a pinned revision (see the latent-lab site config for the current pin), verify sha256 against manifest.json, run with onnxruntime. Input normalization: (uint8_frame / 255 - frame_mean) / frame_std with the stats in manifest.json.