sauravvvv/jepa-arcade

JEPA Arcade

A 13M-parameter Joint Embedding Predictive Architecture that plays real Atari ROMs from pixels, in true two-player PettingZoo environments.

A frozen self-supervised world model perceives the game. A linear probe of ~1–2K parameters reads game state out of the predicted latent. A hand-written policy acts on it. No reinforcement learning, no reward signal, no fine-tuning at play time.

Code: https://github.com/saurav-34/lepong · Training data: sauravvvv/jepa-arcade (dataset)

Checkpoints

File	Game	Actions	Epochs	Contents
pong_R5.pt	Pong	6	100	Stage 1 — encoder + predictor + state head. Recommended for Pong.
pong_R3.pt	Pong	6	100	Stage 1, earlier run. Superseded by pong_R5.
boxing_R1.pt	Boxing	18	100	Stage 1. Recommended for Boxing.
lepong_tennis_v1.pt	Tennis	18	200	Stage 1. Recommended for Tennis.
lepong_atari.pt	Pong	6	100	Stage 1, early Atari run (pre-state_names schema).
lepong_atari_statehead.pt	Pong	6	20	Stage 2 — frozen-backbone state head, 1,158 trainable params.
lepong_tennis_statehead.pt	Tennis	18	20	Stage 2 — frozen-backbone state head, 1,544 trainable params.

Stage-1 checkpoints carry state_mean / state_std / state_names and are directly playable without a stage-2 head.

Results

Validation-set correlations, state head reading a fully frozen encoder + predictor.

Pong (lepong_atari_statehead.pt, 1,158 trainable params):

player_y  +0.991    ball_x  +0.970    ball_vx  +0.971
enemy_y   +0.994    ball_y  +0.981    ball_vy  +0.969

Tennis (lepong_tennis_statehead.pt, 1,544 trainable params):

player_x  +0.964    enemy_x  +0.963    ball_x  +0.899    ball_vx  +0.770
player_y  +0.969    enemy_y  +0.965    ball_y  +0.920    ball_vy  +0.782

Boxing (boxing_R1.pt), L_state 0.0533 → 0.0028 over 100 epochs:

first_x   +0.999    second_x  +0.999
first_y   +1.000    second_y  +0.998

Backbone: 13,010,276 parameters, frozen. Runs at ~20 fps on CPU.

Architecture

128×128 RGB frame ──> CNN Encoder ──> 192-d embedding
                                           │
                    action ──> ActionEncoder ──> AdaLN-zero conditioning
                                           │
                              6-layer causal Transformer
                                           │
                                 predicted embedding
                                           │
                            Linear(192, N) state head    ← the only trainable part

Component	Detail
PixelEncoder	4-layer CNN (3→32→64→128→192), stride-2 convs, BatchNorm + GELU, avg ⊕ max pool fused via Linear(384, 192)
ActionEncoder	2-layer MLP → 192-d, SiLU
ARPredictor	6-layer causal Transformer, AdaLN-zero conditioning, 16 heads, dim_head=64
ProjectorMLP	2-layer MLP with BatchNorm (192→2048→192)
SIGReg	Spectral implicit Gaussian regularization, anti-collapse
StateHead	Linear(192, N) — N = 6 (Pong), 8 (Tennis), 4 (Boxing)

Usage

git clone https://github.com/saurav-34/lepong && cd lepong
pip install -r requirements.txt
hf download sauravvvv/jepa-arcade --local-dir checkpoints/

python -m server.play_pong   --checkpoint checkpoints/pong_R5.pt          --port 8793
python -m server.play_tennis --checkpoint checkpoints/lepong_tennis_v1.pt --port 8792
python -m server.play_boxing --checkpoint checkpoints/boxing_R1.pt        --port 8794

Open the printed URL. Arrow keys move; the model plays the other seat.

Training notes

Auxiliary state loss in stage 1. An MSE over a 192-d whole-frame embedding is dominated by large, high-variance structure — background, walls, paddles. The Atari ball is one or two pixels, so the encoder can discard it entirely and the prediction loss barely moves. Ground-truth coordinates are free at collection time (RAM is already being read), so they are fed back as a supervised term inside stage 1. The head reads the predictor's output, not the raw encoder embedding, matching the inference path exactly.

Perception vs. proprioception. The model perceives the ball and opponent from pixels; its own position is read from RAM. In the collected data both players chase the ball with the same heuristic, so a probe trained to read "my own y" can score well by reading the ball's y instead — a shortcut that inverts the moment the learned policy diverges from the collector's. RAM is therefore used only for what pixels cannot express: own-body position, serve detection, Tennis court-end swaps, and the scripted Boxing opponent.

Limitations

• Trained on data from heuristic + ε-random policies, not human play. Behaviour outside that distribution is uncharacterized.
• Tennis ball_y derives from RAM byte 17, which is the ball's arc height, not its court position. Recorded faithfully but treat as suspect.
• Tennis ball_vx / ball_vy correlations (~0.77–0.78) are materially weaker than positions.
• lepong_atari.pt predates the state_names schema and lacks that metadata.
• Requires Atari ROMs via autorom; ALE cannot run in a browser, so the server owns the environment.

License

MIT

---

Built upon lepong — the original 13M-parameter JEPA world model that plays a synthetic Pong simulator from pixels.