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

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

Battle-state evaluation

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

Can a representation learned from battles recover a missing teammate? I tested this by training a species-prediction head on the JEPA encoder: hide one species, then ask the head to reconstruct it. The target is the species used in a recorded team, not the addition that would maximize its win rate.

For these demonstrations, the input is five species at full HP, an empty opposing team, and no weather or hazards. Movesets, items, and abilities are not supplied. The code excludes already-selected species and special tokens, but does not check format legality. “Sand” and “PU offense” are labels for the chosen examples, not validated team formats.

The percentages below are the head’s softmax scores after filtering. They are not estimated win rates or calibrated measures of team fit. The candidate notes use Showdown’s species data to point out possible fits and conflicts; any benefit still depends on the actual sets.

In the saved budget sweep, the largest run (793 training replays) recovered the hidden species as its first choice 3% of the time, and within its top five 10% of the time. Across all six budgets, those figures ranged from 0–3% and 1–11%, with no steady improvement. The evaluation considers up to 100 distinct held-out teams. The illustrated cores come from a separate full-training-set run, so these accuracies are context for the approach, not measurements of the two examples below.

Team completion

Sand / hazard core

Hippowdon and Excadrill suggest a sand partnership if their abilities are Sand Stream and Sand Rush. Three Ground types also make Water and Ice coverage worth examining. These species names alone do not specify a complete strategy.

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 2.6%

    Thick Fat could soften Ice attacks; the ability is unspecified.

  2. 2 Zoroark sprite Zoroark 2.5%

    Adds special offense, without a Water or Ice resistance.

  3. 3 Arcanine sprite Arcanine 2.4%

    Resists Ice but adds another Water weakness.

  4. 4 Wo Chien sprite Wo Chien 2.4%

    Resists Water but is weak to Ice.

  5. 5 Alolan Sandslash sprite Alolan Sandslash 2.3%

    Slush Rush needs snow, not sand; this is no obvious weather fit.

  1. The top five contain only 12.3% of the probability mass. Hariyama leads Zoroark by 0.14 percentage points before rounding: the model barely separates them.
  2. Some candidates have plausible uses, but the ranking does not consistently resolve the core’s needs. The notes above are a human interpretation, not an explanation extracted from the model.

Team completion

PU offense core

Delphox, Florges, and Rotom-Mow lean toward special offense; Mudsdale brings physical bulk, while Skuntank already supplies Dark typing. An added attacker could help, but there is no uniquely correct sixth species.

Five-slot P1 core

Delphox sprite Delphox 100%
Florges sprite Florges 100%
Rotom Mow sprite Rotom Mow 100%
Mudsdale sprite Mudsdale 100%
Skuntank sprite Skuntank 100%
?

Masked-head top five

  1. 1 Zoroark sprite Zoroark 3.0%

    Adds speed, but overlaps with existing Dark typing and special offense.

  2. 2 Wo Chien sprite Wo Chien 2.8%

    Overlaps with both Rotom-Mow’s Grass and Skuntank’s Dark typing.

  3. 3 Hariyama sprite Hariyama 2.7%

    Adds physical Fighting pressure, with low base Speed.

  4. 4 Alolan Sandslash sprite Alolan Sandslash 2.6%

    Adds Steel typing, but the core supplies no snow for Slush Rush.

  5. 5 Houndstone sprite Houndstone 2.5%

    Adds physical Ghost offense; Sand Rush would lack a sand setter.

  1. The top five contain 13.5% of the probability mass. Four candidates recur from the first core; Houndstone replaces Arcanine, and Zoroark moves to first.
  2. This shows a small change in ranking, not demonstrated strategic adaptation. A species-frequency baseline and tests with shuffled or removed cores would help separate learned team relationships from a general preference for these species.

I would describe this as an exploratory missing-species probe. The next useful test is whether it beats a frequency-based shortlist on held-out teams; a stronger teambuilder claim would also need legal sets and battles comparing the suggested additions. Saved scores and the inference code make the current result inspectable.

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