RookWorld-LM-124M
A unified 124M parameter model combining chess policy (ROOK) and environment simulation (Arbiter) in a single transformer, enabling closed-loop self-play without external engines.
Model Details
Model Description
RookWorld-LM is a breakthrough in unified modeling - a single transformer that can both play chess (policy) and simulate the chess environment (world model) through different prompt prefixes.
Developed by: Jonathan Rahn, Jenia Jitsev (LAION/JSC), Qi Sun (Tokyo Tech/Sakana AI)Model type: GPT-2 (multi-task autoregressive)Language(s): Chess notation, game states, rewardsLicense: MITRepository: GitHub Paper: LAION Research Note Demo: HuggingFace Space
Model Architecture
Parameters: 124MArchitecture: GPT-2Context Length: up to 2048 tokensMulti-task: Policy + Environment in one modelTraining Framework: llm.c
Uses
Direct Use
Self-play Chess: Complete games without external enginePosition Analysis: Chess move prediction with reasoningEnvironment Simulation: Next state prediction from movesResearch: Unified agent-environment modeling
Unique Capabilities
Closed-loop Self-improvement: Can generate training data for itselfWorld Modeling: Predicts game outcomes and state transitionsMulti-task Performance: Excels at both policy and simulation tasks
Training Details
Training Data
Datasets:
rookworld_7m : interleaved ROOK policy and Arbiter environment samplesrookworld_46m: scaled-up interleaved dataset (per LAION note)
Training interleaves both tasks in mixed batches.
Training Procedure
Hardware: 2x NVIDIA RTX 4090Framework: llm.c (karpathy/llm.c)Trained for multiple epochs with llm.c; rookworld_7m results reported at 3 epochs; rookworld_46m at 5 epochs (LAION note)
Evaluation
Policy Performance (ROOK Mode)
Action accuracy (rookworld_46m, 5 epochs): 26.2% (LAION note)BIG-bench Checkmate-in-One: 32.1% (LAION note)
Environment Performance (Arbiter Mode)
Per repository evaluation scripts (RookWorld/README):
Next State Accuracy: 99.61% (baseline 92.3%)State NLS: 99.99% (baseline 99.76%)Reward Accuracy: 99.11% (baseline 98.93%)Terminated Accuracy: 99.13% (baseline 99.04%)
Technical Details
Prompt-based Task Selection
RookWorld-LM uses task prefixes to switch between policy and environment modes:
Policy Mode (ROOK) - Playing Chess Format: P: <FEN position>
Example:
# Input (prompt)
P: r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3
# Output (model generation)
M: d2d4 b1c3 f1c4 f1b5 d2d3 E: 0.6 0.5 0.4 0.3 0.2 B: d2d4The model generates:
M: Candidate moves in UCI notationE: Evaluation scores for each candidateB: Best move selection
Environment Mode (Arbiter) - Simulating Chess Format: A: <current_state>+<action>+<move_history>+
Example:
# Input (prompt with chess state and action)
A: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1+e2e4+e2e4+
# Output (model generation - next state, reward, terminated, truncated)
rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq e3 0 1+0+False+FalseThe model generates:
Next state: Updated FEN position after the moveReward: -1 (loss), 0 (ongoing), or 1 (win)Terminated: Whether the game endedTruncated: Whether max moves reached
Unified Tokenization
Shared vocabulary across tasks:
FEN notation tokens
UCI move notation
State transition tokens
Reward tokens (-1, 0, 1)
Special task prefixes
Self-play Implementation
1 def self_play_game ( model , tokenizer ) :
2 """Complete chess game using RookWorld-LM for both policy and environment"""
3 state = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
4 moves = [ ]
5 history = [ ]
6
7 while True :
8 # Step 1: Get move from policy mode
9 policy_prompt = f"P: { state } "
10 policy_output = model . generate ( tokenizer ( policy_prompt ) [ "input_ids" ] )
11 policy_text = tokenizer . decode ( policy_output [ 0 ] )
12
13 # Extract best move from "B: <move>" in output
14 if "B: " in policy_text :
15 move = policy_text . split ( "B: " ) [ 1 ] . split ( ) [ 0 ]
16 moves . append ( move )
17 history . append ( move )
18 else :
19 break # Invalid generation
20
21 # Step 2: Update state using environment mode
22 env_prompt = f"A: { state } + { move } + { ' ' . join ( history [ - 10 : ] ) } +"
23 env_output = model . generate ( tokenizer ( env_prompt ) [ "input_ids" ] )
24 env_text = tokenizer . decode ( env_output [ 0 ] )
25
26 # Parse environment response
27 parts = env_text . split ( "+" )
28 if len ( parts ) >= 3 :
29 next_state = parts [ 0 ] . replace ( "A: " , "" ) . strip ( )
30 reward = int ( parts [ 1 ] )
31 terminated = parts [ 2 ] == "True"
32
33 state = next_state
34 if terminated :
35 return moves , reward
36 else :
37 break # Invalid environment response
38
39 return moves , 0 # Game incomplete
Evolution Capability
RookWorld-LM supports self-improvement through:
Self-play game generation
Filtering winning trajectories
Continued training on successful games
Iterative performance improvement
See
RookWorld Evol for details.
Limitations
Compute Constraints: No deep search or Monte Carlo methodsContext Window: Limited to 2048 tokensTraining Data: Performance bounded by initial training distributionEvaluation Depth: Single-step lookahead only
Significance
RookWorld-LM demonstrates:
Unified Architectures: Single model for multiple strategic tasksEmergent Capabilities: World modeling from language model trainingSelf-sufficiency: Complete game playing without external toolsScalability: Performance improves with model size
Related Models
ROOK-CLF-9M :ROOK-LM-124M :Arbiter-2M :
Citation
1 @article{rookworld2024,
2 title={RookWorld: Unified Agent and Environment Modeling for Chess},
3 author={Rahn, Jonathan and Jitsev, Jenia and Sun, Qi},
4 journal={LAION Research Notes},
5 year={2024},
6 url={https://laion.ai/notes/rook/}
7 }
Model Card Contact
Jonathan Rahn -
GitHub |
Research Page Metrics Source
