Jaxcalibur

Jaxcalibur is a very strong Pokémon bot. It plays random singles battles, which is the most popular format on Pokémon Showdown. In this format, each player is given a randomly generated team to play with.

Although there had been growing interest in using Pokémon battles as a test-bed for AI, such as with the Neurips PokeAgent Challenge, Jaxcalibur is the first bot to match top human players. It was on-and-off top of the Pokémon Showdown random battles ladder from August 6-27 before it stopped actively seeking games.

This write-up describes how Jaxcalibur works. I may also release the code and model weights, but first I want to understand how the community feels about making a bot this strong widely available.

Jaxcalibur’s high-level design resembles AlphaGo Zero. Given a game state, a transformer-based neural network predicts which action (a move or switch) is best. The network is trained through self-play reinforcement learning: Jaxcalibur plays millions of battles against itself and learns which strategies lead to wins. To make this fast, I built a JAX-based Pokémon battle engine that can run many games in parallel on a GPU.

At test time, Jaxcalibur combines the network with search, allowing it to look ahead several turns. Its search algorithm is a variant of pUCT designed to handle the stochastic environment, imperfect information, and simultaneous decision-making in Pokémon battles. Jaxcalibur uses poke-env to connect to Pokémon Showdown and play online.

Contents: Architecture · Training · Game Engine · Search · Results · Parting Thoughts

Architecture

Jaxcalibur uses a non-causal transformer with various output heads. Most of the model follows a standard pre-RMSNorm architecture. The main exception is a position-based mixture-of-experts setup where (for example) positions representing Pokémon use different MLP weights from positions representing moves, events, etc.

Its inputs are

• 1 token for the battlefield state (weather, hazards, etc.)
• 2 for the active Pokémon (stats, moves, damage calcs, etc.)
• 5 for its candidate moves (Struggle gets its own extra token)
• 6 for the Pokémon on its team
• 6 for the opposing team (masked if they aren’t known yet)
• 16 tokens containing a history of the last 16 game events (moves/switch-ins)
• 24 matchup tokens, each containing information such as damage calculations for a pair of Pokémon. It matches each of its Pokémon to 4 of its opponent’s (an efficiency trick to save 12 tokens, as usually a couple opponent Pokémon are unrevealed or fainted). These are concatenated to the current residual stream half-way through the network, partially for efficiency and partially so they benefit from fresh move/item/ability predictions.

For action selection, each team position emits a logit for switching to that slot. Each move position emits two logits: one for using the move and one for Terastallizing and using the move.

Jaxcalibur uses a tiny model by modern standards: just 8.5M parameters! That is <20% of the size of AlphaGo Zero and <0.005% of the size of GPT-3. I did not perform a full scaling-law analysis, but in my experiments, smaller networks trained for longer outperformed larger networks trained for fewer steps. I’m sure a scaled-up model with more compute would perform better, but I trained Jaxcalibur on just one GPU.

Training

Jaxcalibur is trained purely with self-play reinforcement learning. The RL algorithm is PPO with generalized advantage estimation. Rewards are undiscounted: 1 for a win and 0 for a loss.

It uses two regularization terms to encourage exploration: entropy regularization and a zero-avoiding term, which is essentially KL divergence toward a uniform policy. The latter prevents the model from assigning effectively zero probability to a very situational move such as Encore and never reconsidering it.

The model also has auxiliary losses for predicting:

• Its probability of winning, which serves as the value function
• The item, ability, moves, and Tera type of each opposing Pokémon. It maintains the item and ability predictions across turns so it can remember revealed information.
• The opponent’s next action. This one made the bot substantially stronger, likely because it can react to these predictions.

Unlike AlphaZero, Jaxcalibur does not use search during self-play. Training only the base policy is much faster, and search appears to be less valuable for Pokémon than for Go or chess. The final training run took about a week on one H100, during which Jaxcalibur played almost 100 million games against itself.

Game Engine

The bottleneck in reinforcement learning is often the environment (i.e., running the battles) rather than the neural network itself. I solved this in a kind of crazy way: I reimplemented Pokémon battling in pure JAX so games can be vmapped and jitted to run quickly on a GPU. This is where the name “Jaxcalibur” came from.

The engine always executes every possible game step (moves, switches, end-of-turn effects, and so on) and keeps only the relevant results. This wastes FLOPs, but gives every action the same computational structure, making large batches efficient on a GPU.

Building the engine would have been a massive undertaking without Claude Code. Pokémon battles are conceptually simple, but there is a long tail of moves, abilities, and special-case interactions. I probably would not have attempted this a year ago.

Search

Search allows Jaxcalibur to simulate future actions for both players and evaluate the resulting states. It uses AlphaZero-style pUCT, but with several modifications to account for important properties of Pokémon:

• Simultaneous actions: In Pokémon, both players often choose actions at the same time. As a result, the Nash equilibrium is a mixed policy: there is no single best action because the strongest choice depends on what the opponent does. To handle this, Jaxcalibur samples edges during search according to a mixture of the positive regret and network prior. This formulation makes it shift towards a regret-matching policy as the number of node visits increases. For those familiar with the pUCT equations, Jaxcalibur samples according to:

  \[\pi(s,a) \propto \max(Q(s,a)-V(s), 0) + \frac{c_\text{puct}}{\sqrt{N(s)}}P(s,a)\]
• Stochastic environment: There is a lot of randomness in Pokémon battles. Rather than branching on each possible random outcome, Jaxcalibur rolls out every trajectory from scratch with a different random seed. Different simulations of the same action sequence can therefore produce different outcomes, so instead of maintaining an explicit search tree Jaxcalibur stores statistics in a dictionary mapping {game_state_hash -> node}. To reduce the number of nodes, Pokémon HP is binned into 10 buckets. Nodes where Jaxcalibur acts are hashed using only revealed information about the opponent, while opponent nodes are additionally hashed on the world ID described below.
• Jaxcalibur’s missing information: In random battles, the opponent’s team and sets aren’t known in advance. As in Foul Play, Jaxcalibur handles this by sampling possible “worlds” of the opponent’s unknown Pokémon/moves/items/etc. using its network’s predictions. Search essentially averages over these possibilities.
• The opponent’s missing information: Search for the opponent is truncated whenever Jaxcalibur does something “surprising,” such as switching to an unrevealed Pokémon or using a move predicted with <65% likelihood from the opponent’s perspective. From that point onward, the opponent only uses the network prior. This prevents search from leaking hidden information such as Tera types to the simulated opponent.

Search is used only at test time because running it during training would have required much more compute. Jaxcalibur currently searches to depth 4 (usually 2 full turns ahead), using 20,480 rollouts over 32 possible worlds. This takes just a few seconds on a GPU.

Overall, I think most of Jaxcalibur’s strength comes from having a really good neural net rather than from its search; my current estimate is that search adds 100-150 Elo over its base policy.

Results

I ran Jaxcalibur on the Pokémon Showdown random battles ladder for three weeks from August 6-27. During this time it was top of the ladder about 20% of the time and maintained a 68% win rate against other top players. It reached a peak Elo of 2557 and a peak GXE of 95.3, meaning its estimated chance of winning against an average ladder player is 95.3%. For comparison, the strongest previous bot I’m aware of, Foul Play, peaked at around 2341 Elo and 88 GXE, which is roughly top-50 rather than top-1.

The ladder on August 10, when Jaxcalibur reached its peak Elo (but not GXE).

Despite these great results, I can’t say Jaxcalibur is the world’s best random battler because one human outperformed Jaxcalibur over this time period! MichaelderBeste2 was on top of the ladder 50% of the time and has a current GXE of 95.5. I am really impressed that anyone can play better than a bot tuned over millions of games and simulating thousands of possibilities each turn; that is an incredible feat!

Parting Thoughts

Jaxcalibur began as a small side project and became much more sophisticated than I initially planned. Its strength illustrates how far AI algorithms and hardware have progressed: AlphaGo required a large research team and thousands of accelerators, while Jaxcalibur was built as an independent project and trained on one GPU.

Strong chess engines have transformed chess in both good ways and bad. I think the same could be true for Pokémon bots, so I want to make sure the community has a say in what happens with Jaxcalibur. For now, I am holding off on releasing the code and weights. For me, the biggest benefit of open-sourcing Jaxcalibur is that it would let others learn from and build on top of it. Strong bots could also make single-player Pokémon more challenging, serve as practice partners, and be powerful tools for team building or balancing formats. Or would having easily accessible information on what teams, leads, and initial plays work best make preparing less fun? I’m also worried about cheating, which has been a problem in chess: people could use Jaxcalibur to climb the ladder, obtain suspect reqs, or potentially even win tournaments in a way that’s hard to detect. If you have thoughts on whether or how Jaxcalibur should be released, please comment on the Smogon Forum post.