Keisei

Deep Reinforcement Learning for Shogi, powered by a Rust game engine.

Keisei (形成, "to give form to, to mold, to shape") trains neural network agents to play Shogi (Japanese chess) using Proximal Policy Optimization (PPO). The game engine and RL environment are written in Rust for performance; the training harness is Python with PyTorch.

Status: Early rebuild. The Rust engine is functional; the Python training harness is under active development with a KataGo-inspired multi-head architecture as the primary training target.

Prior Art and Attribution

Keisei's neural network design draws heavily from two published systems:

• KataGo (David Wu, 2019) — The SE-ResNet trunk with global pooling bias, the Win/Draw/Loss value head, and the score prediction auxiliary head are adapted from KataGo's architecture. KataGo was designed for Go; Keisei adapts these ideas for Shogi's different action space, piece-drop mechanics, and promotion rules.
• AlphaZero (Silver et al., 2018) — The general approach of residual convolutional networks with separate policy and value heads, trained via self-play, originates from DeepMind's AlphaGo/AlphaZero line of work.

What Keisei adds on top of these foundations:

• A Rust game engine with vectorized environments exposed to Python via PyO3, replacing the typical C++ engine pattern.
• Shogi-specific observation encoding (50-channel, perspective-relative) with hand-piece normalization and repetition/check planes.
• A spatial action decomposition (81 squares x 139 move types) tailored to Shogi's move semantics, including drops and promotions.
• A dual-contract model system with adapter pattern, allowing the training loop to support both simple scalar-value models and KataGo-style multi-head models without branching.
• PPO-based training (KataGo uses a custom self-play + training pipeline; AlphaZero uses MCTS + supervised learning from self-play).

Architecture

Rust Engine (shogi-engine/)

Two workspace crates providing the core game logic:

• shogi-core — Full Shogi implementation: board representation, legal move generation, rule enforcement (check, checkmate, repetition, impasse), SFEN parsing, and Zobrist hashing.
• shogi-gym — RL environment layer: vectorized environment (VecEnv), observation encoding (46 or 50 channels), action mapping (11,259 spatial actions), spectator data for live visualization, and step/reset API exposed to Python via PyO3.

Python Training Harness (keisei/)

• Config — TOML-based configuration with dataclass validation.
• Models — Four neural network architectures with a registry-based dispatch system. See Neural Network Architecture below.
• KataGo-PPO — PPO with W/D/L value head + score auxiliary head, GAE, clipped surrogate objective, and mini-batch updates. Supports torch.compile for fused kernel execution.
• Value Adapters — Adapter pattern that abstracts over scalar vs. multi-head value outputs, so the training loop is model-agnostic.
• Training Loop — Orchestrates environment interaction, PPO updates, metric logging to SQLite, checkpointing, and resume-from-checkpoint support.
• Database — SQLite layer (WAL mode) storing training metrics, game snapshots, and training state for the spectator WebUI.

Neural Network Architecture

Keisei supports four architectures via a registry. The SE-ResNet (adapted from KataGo) is the primary training target; the others serve as baselines and ablations.

The SE-ResNet outputs three heads: a spatial policy (B, 9, 9, 139) over 81 squares x 139 move types, a W/D/L value classification (Win/Draw/Loss — a KataGo innovation), and a score prediction for material balance (auxiliary task). Two PPO variants match the two model contracts: standard PPO for scalar models, KataGo-PPO for multi-head.

Built With

• css
• html
• javascript
• python
• rust
• shell
• svelte