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 variantslookup/{states,latents}.bin— latent<->state lookup table (float32 LE) for decoder-free visualizationmanifest.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.