#su-JEPA
This began as an experiment to get comfortable training a world model end to end — data collection, dynamics model, decoder, evaluation — on ARC-AGI-3 gameplay. It is not a solver: there is no agent here and no planning. This page is the short version of what the exercise found, the recipe, and where it does not hold up.
ARC-AGI-3 hands you five never-seen interactive games and scores how well an agent solves them. The bet I wanted to test was modest: can a JEPA-style latent dynamics model, trained only on exploration recordings, learn gameplay dynamics well enough to be a substrate for a later agent — one that explores, plans in latent space, and fine-tunes on a new game as it plays. Code is at github.com/DenisSud/arc-agi-3.
What it is
A per-frame CNN encoder compresses each 64×64 frame into one 512-d latent; a causal transformer predictor takes (latent, action) and predicts the next latent as a residual delta. There is no EMA teacher and no stop-gradient — LayerNorm plus SIGReg (λ=1) are what keep the latents from collapsing. A separate frozen-encoder readout decodes latents back to pixels for visualization and measurement.
The recipe
- Data: ~573k exploration episodes / ~50.5M transitions over 20 games, collected with a Just-Explore policy. Collection saturated — fresh-state yield fell 59% → 36% → 25% — so the data axis is effectively closed. Five test games are held out entirely.
- Model: 96M parameters — 12-layer / 12-head causal transformer over 512-d latents, 768 hidden, residual-delta predictor, SIGReg λ=1.
- Packing: real contiguous windows at
seq_len 32,B=128— short windows with big batches give ~2.5× the throughput per transition of the L=512 regime — a single XLA compile for the whole run. - Optimization: AdamW, lr 1e-4 → 1e-5 cosine, all-bf16, ~30 epochs (~450k steps) on a single RTX 5070.
- Decoder: frozen encoder, class-weighted cross-entropy over the 16 colors, conditioned on the previous frame (decode the new latent given what the last frame looked like), trained on clean latents.
What worked
On held-out episodes of the trained games the one-step dynamics is solid: prediction loss fell from ~1.0 to ~0.17, predicted and true next latents agree at cosine ~0.89, and the decoder renders ~90% of changed pixels correctly — and it was still improving when I stopped it. On games the model has seen, it both predicts the next state and can draw it.
Honest footnote: an early loss-averaging bug inflated every reported prediction loss by the latent dimension. Gradients were unaffected, but the λ·SIGReg term was diluted, so old numbers need rescaling by the latent dim to be comparable.
Where it breaks
- Unseen games — the wall. On the five held-out test games, pixel-level reconstruction and prediction accuracy is ≈ 0; the decoder simply does not transfer. Only the latent cosine improves with context (e.g. 0.63 → 0.74 over a rollout) — faint in-context adaptation with no pixel-level payoff. Never-seen games are the whole point of ARC-AGI-3, so this is the result that matters.
- No diversity. All of it trained on 20 games, exhausted by exploration; more rollouts of the same games add ~0% unique content.
- Repetition, not capacity. The corpus dedups to ~37M unique transitions, which the run saw ~50 times — well past the point where repeated tokens stop paying (most value is spent by ~16 epochs). The bottleneck was never model size or steps.
- No end-to-end check. With no agent, there is no measurement of whether any of this helps solving — only proxies.
Takeaway
The pipeline is sound and measured, and the model learns the dynamics of the games it saw — it can predict them and imagine them. But generalization to new games, the only thing the benchmark rewards, did not come out of this recipe: 20 games of pure exploration is too narrow a distribution, and repetition cannot buy the missing diversity. If this line continues, the unlock is more diverse environments — new games, new mechanics, a synthetic-world generator — before more compute.
The DAG edge-packing scheme I designed along the way — packing whole prefix trees of gameplay so shared prefixes are paid for once — is gone from the training path; with short windows and a flat per-position loss profile its shared-prefix savings stopped mattering. The write-up lives here.