#su-JEPA
Work in progress — a snapshot, not a finished write-up.
The plan
ARC-AGI-3 gives you five never-seen-before interactive games and scores you on how well you solve them. My bet: a JEPA world model trained on exploration recordings can serve as a general game-understanding engine, and everything else — acting, planning, adaptation — is downstream of how good that latent dynamics model is. Roughly: pre-train the world model on gameplay recordings, with a context window long enough to hold an entire game's action history — so a never-seen game is handled purely by in-context learning, no test-time training. A single exploration-based solving algorithm (currently under development) plays the new game: it maximally explores while solving, filling the context window as it goes; data from finished games feeds back into continued training between games.
What works right now
The core model and training pipeline are up and stable. Architecture: a
per-frame CNN encoder into a single latent vector per frame, a causal transformer predictor
with RoPE where the action embedding is appended to the latent sequence
(projected by the transformer's embedding layer), residual-delta parameterization
(pred = lat + delta — copying is the zero baseline), and no EMA /
stop-gradient — LayerNorm regularization plus SIGReg (λ=1) fight latent collapse. Training: LLM-style
episode packing with block-diagonal causal masks, and trie-based DAG packing
over deduplicated frames — every unique state transition is seen exactly once
per epoch, one XLA compile for the whole run. The full pipeline runs
end-to-end on a single RTX 5070 at ~30 steps/s after warmup (B=8, seq 512) over 20 games /
50M transitions, with a frozen latent decoder for visualization
(changed-pixel accuracy 99.9%).
Currently working on
- A proper large pre-training run (~9M params, ~50 epochs) — the 20 real games are saturated, more rollouts buy nothing. DAG packing is still in testing — it needs more work on position encodings before it's actually worth it over plain document packing.
- The solver itself: an exploration-based solving algorithm that uses the world model to act in a never-seen game — maximal exploration while filling the context window (in-context learning), latent rollouts for planning. Data from finished games feeds continued training between games. This is the whole point of ARC-AGI-3 and the hardest, least-defined part.
Longer write-up once the solver exists and I have something honest to say about it.