💬 NLP

UniMaia: Steering Chess Policies with Language for Human-like Play

The paper introduces UniMaia, a parameter-efficient framework that steers a frozen Lc0 chess policy network using language prompts to achieve human-like controllability over gameplay style and strength while preserving domain-specific performance without requiring end-to-end multimodal training.

Original authors: Sherman Siu (University of Waterloo), Lesley Istead (University of Waterloo)

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sherman Siu (University of Waterloo), Lesley Istead (University of Waterloo)

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a world-class chess coach who is incredibly smart and knows every possible move, but they are also a bit of a robot. They always play the "perfect" game, which is great for winning, but boring if you want to see how a human might play, or if you want them to play like a specific person (like a beginner or a grandmaster).

On the other hand, you have a creative storyteller who loves to chat and follow instructions like "Play like a nervous beginner" or "Use the Sicilian Defense," but this storyteller isn't very good at chess and often makes illegal moves or misses obvious threats.

UniMaia is a new framework that tries to get the best of both worlds. It takes the "robot coach" (a frozen, pre-trained chess AI called Lc0) and gives it a "translator" (a text encoder) that listens to natural language instructions.

Here is how it works, using some simple analogies:

1. The Frozen Coach and the New Translator

Think of the chess engine (Lc0) as a frozen statue of a perfect player. You can't melt it down or retrain it from scratch because that would take too much time and money, and you'd lose its incredible chess skills.

UniMaia doesn't melt the statue. Instead, it builds a ControlNet-style "control panel" around it. This is like attaching a set of dials and levers to the statue.

The Levers: These are controlled by a text encoder (a small language model).
The Input: You type a prompt like, "Play as a 1200-rated player who loves the King's Gambit opening."
The Action: The text encoder reads your words and turns the dials on the control panel. These dials send tiny, precise adjustments (residual updates) to the statue's brain, nudging it to play more like a human or follow a specific strategy, without changing the statue's core ability to understand the board.

2. The "Steering Wheel" Analogy

Imagine the chess engine is a high-performance race car. It knows how to drive perfectly on a track.

Old Way: To make the car drive like a taxi or a delivery truck, you had to rebuild the entire engine from scratch.
UniMaia Way: You keep the race car engine exactly as it is. You just install a smart steering wheel that can be programmed via voice commands. If you say, "Drive like a cautious taxi driver," the steering wheel adjusts the car's path slightly to be more conservative. If you say, "Drive like a reckless racer," it adjusts the path to be more aggressive. The engine is still the same powerful engine, but the behavior changes based on your voice.

3. The "Human-like" Goal

The researchers wanted to see if they could make the AI play like a human. Humans make mistakes, have favorite openings, and play differently depending on their skill level (their "Elo" rating).

The Result: UniMaia successfully learned to "steer" the chess engine. When asked to play like a beginner, it made more mistakes. When asked to play like a grandmaster, it played more accurately. It could even switch between different chess openings based on a text description.
The Trade-off: There is a slight balance to strike. The more you ask the AI to "act human" or follow complex instructions, the slightly less perfect its top move prediction becomes compared to a pure, unguided chess engine. However, it remains very competitive with models specifically built just to predict human moves.

4. The "Time Machine" Feature (UniMaia-Aux)

The researchers added a second version called UniMaia-Aux. This is like giving the coach a time machine and a stopwatch.

In addition to the text instructions, this version also looks at when the move was made, how much time the player had left on their clock, and how long they took to think about the previous move.
The Result: This helped the AI understand human behavior even better. For example, it got much better at predicting when a human would "resign" (give up) based on the game situation and time pressure, something the standard version was less good at.

5. The Data Diet

To teach this system, the researchers didn't just feed it random games. They built a massive library of 5.2 billion chess games from Lichess.

They created a "prompt generator" that automatically wrote millions of different instructions for these games, such as "A player rated 1500 playing the French Defense."
This allowed them to train the system to understand a huge variety of human playing styles without needing humans to write every single instruction by hand.

Summary

UniMaia proves that you don't need to rebuild a powerful AI from scratch to make it follow instructions. By attaching a lightweight "language translator" to a frozen expert, you can steer its behavior with natural language. It creates a chess AI that is not just a calculator, but a flexible partner that can mimic different human players, follow specific strategies, and even understand the pressure of the clock, all while keeping its core chess genius intact.

Technical Summary: UniMaia

Problem Statement

In structured decision-making domains like chess, a fundamental tension exists between controllability and domain grounding. Specialized policy networks (e.g., Leela Chess Zero, Lc0) achieve strong performance and robust inductive biases through structured board encodings but lack semantic controllability; they cannot be easily steered by natural language instructions. Conversely, prompt-conditioned large language models (LLMs) offer flexible semantic control but often exhibit weaker domain grounding, lower move-prediction accuracy, and a tendency to generate illegal moves. Furthermore, end-to-end multimodal training to unify these capabilities typically requires substantial computational resources and large-scale curated datasets.

The authors investigate whether natural language prompts can steer a pretrained, domain-specific chess policy while preserving its inductive biases, without resorting to full end-to-end multimodal training.

Methodology: UniMaia

The paper proposes UniMaia, a framework for prompt-conditioned policy modulation. The architecture combines a frozen domain-specific backbone with a lightweight, parameter-efficient conditioning mechanism.

Architecture

Frozen Backbone: UniMaia utilizes a pretrained Lc0-ChessFormer (Lc0-CF) policy network (initialized from BT4 weights) as a frozen backbone. This network encodes board states using 112 spatial feature planes and predicts move probabilities over a fixed legal move space.
Text Encoder: A ChessGPT-base text encoder (a 3B parameter model) processes natural language prompts. To maintain parameter efficiency, the encoder is fine-tuned using LoRA (Low-Rank Adaptation) with rsLoRA and LoRA+ adapters.
ControlNet-Style Conditioning: Instead of early fusion or deep fusion, UniMaia employs a layer-coupled ControlNet branch.
- A trainable copy of the encoder layers runs in parallel to the frozen backbone.
- At each layer $n$ , the ControlNet branch receives the residual-updated representation from the previous layer ( $h_{n-1} + r_{n-1}$ ) and the prompt embeddings ( $e$ ).
- The branch produces a residual update $r_n$ via a zero-initialized linear projection $Z_n$ , ensuring the model initially behaves identically to the pretrained backbone.
- Controllable Policy Head: A cloned policy embedder and head are added at the output. These layers apply cross-attention over the prompt embeddings to the intermediate policy representation, allowing direct modulation of the final move distribution.

Training and Data

Dataset: The authors constructed LICHESSGAMES, a large-scale dataset derived from 5.2 billion Lichess games (2013–2023).
Prompt Construction: To enable large-scale training without human annotation, the authors developed a semi-automated pipeline using GPT-4 Turbo to generate synthetic natural language templates. These templates are instantiated with game metadata (player ratings, time controls, openings, etc.). Two families of templates were created: LICHESSTEMPLATES-PRETRAIN (160 metadata-rich templates) and LICHESSTEMPLATES-INSTRUCT (2,048 instruction-oriented templates).
Optimization: The Lc0 backbone remains frozen throughout training. Only the text encoder adapters and the ControlNet conditioning branch are updated. The authors utilized the NorMuon optimizer for 2D non-embedding parameters and AdamW for others.
UniMaia-Aux: An extended version introduces auxiliary temporal conditioning and behavioral prediction objectives (game outcome, termination type, remaining plies, move delay) during a fine-tuning stage.

Key Contributions

UniMaia Framework: A parameter-efficient method for modulating pretrained chess policies using natural language prompts, preserving the inductive biases of specialized models while enabling semantic control.
Dataset and Pipeline: The creation of a large-scale, metadata-augmented Lichess dataset and a semi-automated prompt-generation pipeline for language-conditioned policy learning.
Unified Benchmarking: The introduction of a comprehensive evaluation framework spanning both prompt-conditioned (instruction following, opening control) and metadata-conditioned (human move prediction) settings, allowing for systematic comparison between language-based and structured-conditioning approaches.
Empirical Validation: Demonstration that lightweight prompt-conditioned modulation can achieve controllable, human-like chess behavior while remaining competitive with structured-conditioning approaches on human move prediction benchmarks.

Results

The authors evaluated UniMaia against metadata-conditioned models (Maia-1, Maia-2, Allie) and prompt-conditioned baselines (ChessGPT, GPT-3.5 Turbo).

Instruction Following: UniMaia achieves state-of-the-art Acc@1 (top-move accuracy) on prompt-conditioned instruction-following benchmarks (LIF and LGB), outperforming prior prompt-conditioned baselines.
Opening Control: On opening-specific benchmarks (LOB-P and LOB-C), UniMaia ranks second, trailing only ChessGPT-base, which specializes in opening reconstruction but generalizes poorly to broader gameplay.
Human Move Prediction: UniMaia remains competitive with dedicated metadata-conditioned models (like Allie) on human move prediction benchmarks (e.g., ABB, M1-S), despite being optimized primarily for language control.
Expected Accuracy: UniMaia-Aux achieves the strongest Expected Accuracy (E[Acc]) across most settings, suggesting that auxiliary temporal conditioning improves policy calibration even when temporal metadata is unavailable at inference time.
Continuity: Analysis shows that the induced policy varies continuously with prompted Elo ratings, suggesting the model learns a continuous notion of player strength rather than isolated rating-specific behaviors.
Efficiency: The approach achieves strong performance with approximately 67M plies of training data, significantly less than competing approaches (e.g., Maia-2 uses 9.15B plies), though the LLM backbone imposes higher computational costs per step.

Significance and Claims

The paper claims that prompt-conditioned control of domain-specific policy networks is feasible without end-to-end multimodal training. By combining a frozen expert policy with lightweight language conditioning, UniMaia successfully bridges the gap between the flexibility of natural language interfaces and the robustness of specialized decision-making models.

The authors highlight a trade-off: while language-based control offers superior semantic steerability (e.g., selecting openings or adjusting player strength via text), it comes with modest trade-offs in raw predictive accuracy compared to models trained exclusively on structured metadata. However, the results suggest that conditioning design (data composition, integration with pretrained policies) is more critical for controllability in structured domains than raw model scale.

The work positions lightweight language-conditioned adaptation as a practical alternative to full multimodal training for structured decision-making systems, offering a general framework for building controllable policies on top of pretrained expert models.

Limitations

The authors acknowledge several limitations:

Prompt Sensitivity: The model is sensitive to prompt phrasing and tokenization; small lexical changes can alter policy outputs.
Synthetic Data: Prompts are synthetically generated, which may introduce artifacts and fail to capture the full variability of real human instructions.
Evaluation Scope: Experiments rely on offline automated benchmarks rather than human studies or live gameplay, leaving questions about perceived human likeness in real-time interaction.
Compute Overhead: The frozen 3B-parameter LLM backbone introduces significant computational overhead compared to purely structured models, limiting reproducibility and accessibility for smaller-scale experiments.