Adithya Parthasarathy

Pokemon JEPA: Learning a Battle World Model

I rebuilt a small Transformer world model over Pokemon Showdown replays to answer a stricter question: can an action-conditioned JEPA learn battle dynamics that beat simply copying the current board?

Intro

Pokemon Showdown records competitive battles as text logs. Each turn has a board state, two actions, and a next state.

The simple baseline is surprisingly strong. It is a persistence model: predict no change by copying yt+1 = yt. Most bench Pokemon do not change on a turn, so this reaches 0.020 MSE. Any dynamics model must beat it.

The model uses a true JEPA. It predicts the next latent state from the board and both actions. The full run reaches 0.016 MSE, an 18% improvement over persistence.

1

Encode each battle as 13 tokens: 12 canonical Pokemon slots plus the shared field.

2

Predict the next latent state from the current board and both players' actions.

3

Beat the naive persistence baseline, then reuse the latent for winner prediction and team completion.

Background

A winner-only model can learn useful correlations. It can also miss momentum. This experiment splits the replays first, then freezes the encoder before evaluation.

JEPA asks a second question: what will the board look like next? Here, next-turn HP is the simple probe for damage, recovery, switches, and tempo.

Schematic showing a Pokemon battle state as twelve object tokens
Figure 1. The model keeps Pokemon in canonical slots and adds field and action features.

Model

The parser creates 12 fixed Pokemon slots and one field token. Slots include species, HP, status, active state, stat boosts, and position. The predictor also receives both actions.

Parsed Replays 1,000

800 train / 200 test split before vocab extraction

Transitions 21,822

17,542 train / 4,280 test turn transitions

Board Geometry 13 tokens

12 canonical slots + 1 global field/weather token

GPU 11 GB VRAM

NVIDIA GeForce GTX 1080 Ti

Diagram of the Pokemon JEPA architecture
Figure 2. The encoder predicts the next latent from the current board and both actions.

Input

Species, HP, status, boosts, field state, and both actions.

Latent

A Transformer latent with explicit side, slot, and active embeddings.

Objectives

JEPA prediction with EMA targets and VICReg. Small frozen probes test the result.

Model Training signal Why it matters
JEPA Predicts the next latent from the board and both actions. Main model. Tests whether dynamics shape the representation.
Supervised Same model size, trained on final winner labels only. Tests what winner prediction learns without next-state prediction.
Persistence Copies current HP into the next turn. Simple dynamics floor. Any useful predictor should beat it.
Random encoder Untrained encoder with frozen readouts. Checks whether training adds signal beyond random features.

HP and winner results use small frozen linear probes. The probes read the representations; they do not train the encoders.

The online encoder sees the current board. An EMA target encoder supplies the next latent. The predictor uses both actions. VICReg keeps the representation from collapsing.

Sample Pokemon JEPA input state and next-turn HP target
Figure 3. The training pair is a current board and its next state.

Results

The study uses 1,000 public [Gen 9] PU replays from the Pokemon Showdown replays dataset, with an 80/20 replay split. The full results and code are in the public project repo.

Line chart comparing JEPA next-HP prediction error against naive persistence
Figure 4. Next-HP probe MSE versus the naive persistence baseline.

Result 1

The persistence trap is real

Most bench slots stay unchanged. Copying the current HP vector reaches 0.020 MSE. The model reaches 0.016: an 18% improvement.

It learns something about attacks and switches. It is not just copying the board.

Sample-efficiency chart comparing JEPA, supervised, random, and baseline winner accuracy
Figure 5. Winner accuracy across replay budgets.

Result 2

Winner accuracy is useful, but noisy

At full budget, the frozen JEPA probe reaches 56% accuracy versus 49% for the supervised baseline. Its Brier score is 0.26 versus 0.49. Lower is better.

Accuracy moves around with the replay budget. The dynamics result is the stronger claim.

Metric JEPA Reference Takeaway
Next-HP probe MSE 0.016 0.020 Full-budget JEPA vs naive persistence
R² over persistence +18% — Improvement beyond copying the current board
Winner probe accuracy 56% 49% Frozen JEPA probe vs supervised baseline
Winner Brier score 0.26 0.49 Lower is better; JEPA is less overconfident

The summary table shows the main story. The model improves next-state prediction and produces less overconfident winner probabilities. Detailed budget-by-budget metrics are in the repo.

Examples

The model can read a battle state and suggest a missing team slot.

State Evaluation

Dynamics Example 1

A healthy team that is already losing

The label-only model sees six healthy Pokemon and several recognizable attackers. The JEPA model sees lost tempo, hazard damage, and a stalled trajectory.

P1 board

Salazzle sprite Salazzle 100% special pressure
Bellibolt sprite Bellibolt 100% bulky special
Cryogonal sprite Cryogonal 100%
Grimmsnarl sprite Grimmsnarl 100%
Paldean Tauros Aqua sprite Paldean Tauros Aqua 100%
Rotom Mow sprite Rotom Mow 100%

P2 board

Froslass sprite Froslass 100%
Coalossal sprite Coalossal 100% Salazzle answer
Scream Tail sprite Scream Tail 100% special wall
Grimmsnarl sprite Grimmsnarl 100%
Dusknoir sprite Dusknoir 100%
Clawitzer sprite Clawitzer 100%
Supervised 59% P1
JEPA 38% P1
Actual P2 wins
  1. Surface read: P1 has three strong-looking special attackers.
  2. JEPA read: Scream Tail and Coalossal absorb the pressure and take tempo.
  3. The JEPA call is cautious: 38% P1. P2 wins.

Team Completion

The teambuilder masks slot 6 and predicts a distribution over species. It filters duplicates and illegal picks.

Teambuilder Demo 1

Sand / hazard core

The masked-species head ranks a sixth slot from the latent state of a competitive core. Probabilities stay diffuse because the vocabulary contains 331 species.

Five-slot P1 core

Hippowdon sprite Hippowdon 100% sand setter
Excadrill sprite Excadrill 100%
Corviknight sprite Corviknight 100%
Slowking sprite Slowking 100%
Gliscor sprite Gliscor 100%
?

Masked-head top five

  1. 1 Hariyama sprite Hariyama 3%
  2. 2 Zoroark sprite Zoroark 3%
  3. 3 Arcanine sprite Arcanine 2%
  4. 4 Wo Chien sprite Wo Chien 2%
  5. 5 Alolan Sandslash sprite Alolan Sandslash 2%
  1. The top predicted candidates—led by Hariyama and Zoroark (each ~3%) - demonstrate that the latent space captures role complementarity.
  2. It has selected for physical resistance to entry-hazard exploiters and high-tempo offensive pivots.

Discussion

The result is a useful state representation, not a perfect winner predictor. This is still a small study on 1,000 replays. The next step would be to connect it to action selection.

References