JAX-native board games, batched for RL at GPU speed

What it does

Pgx is a library of GPU-accelerated game simulators built entirely in JAX. It covers discrete-state classics—Chess, Shogi, Go, Backgammon, Bridge, Poker, and a handful of arcade-style MinAtar variants—so you can batch and JIT-compile RL training loops on accelerators. The environments expose a standard state-and-step interface, and every transition is written to be vectorized and deterministic where the rules allow.

The interesting bit

The project essentially does for discrete board games what Brax does for continuous physics: it removes the Python overhead that normally bottlenecks tree-search and self-play. A recent API v2 change even forces explicit random keys for stochastic games, plugging a subtle hole where planning algorithms might accidentally peek at hidden RNG state.

Key highlights

• JIT-able step functions across all environments, ready for jax.vmap.
• Broad catalog spanning Go, Chess, Shogi, Backgammon, Bridge bidding, Kuhn Poker, Hex, and the MinAtar suite.
• Built-in SVG rendering, so you can inspect states without writing custom drawing code.
• Optional PettingZoo AEC wrapper for compatibility with existing multi-agent pipelines.
• Example training code for AlphaZero and PPO included.

Caveats

• API maturity is uneven: the catalog mixes v0, v1, and v2 environments, so stability and signatures vary by game.
• The v2 update changed the step signature for stochastic environments (2048, Backgammon, MinAtar), breaking older code that relied on internal RNG state.
• The README admits the library prioritizes speed over API polish; while the interface is general, you may still need to handle version differences across the catalog.

Verdict

Grab Pgx if you are training AlphaZero-style agents or batched policy-gradient methods in JAX and need more than a single Python game loop. If your stack is PyTorch or you only need one-off human play, the JAX dependency and vectorization model are likely overkill.