A world-model-first transformer bot for Figgie. The model learns to predict what the market does next (opponent bids, trades, round outcomes) rather than predicting its own actions directly. The policy is a small head built on top of those learned representations.
Figgie is a four-player card-trading game with a hidden goal suit. This project trains a transformer to model the live market (every bid, offer, trade, and cancellation) and uses that world model as the substrate for a small policy head. The result is a bot that plays the live Rust sandbox via a stdin/stdout sidecar and meets fixed win-rate targets against scripted opponents.
- World model first, policy second. The transformer predicts market events; the policy is a thin factorized head reading the same hidden states.
- Factorized vocab of 187 tokens. Each event becomes a short run of sub-tokens (
event_type, optionalsuit,price,role). Tiny vocab regardless of price range or game length. - Relative roles (SELF / OPP1-3), not absolute player indices. Gives 4x data augmentation per game and forces position-invariant strategy learning.
- Two aux heads on the WM.
goal_suit_headpredicts which suit is goal,opp_hand_headpredicts each opponent's hand. Both sharpen as more events arrive; ECE measures the sharpening. - Imagination RL inside a frozen WM. Rollouts happen inside the model, no live sandbox in the inner loop. Reward at imagined ROUND_END comes from
payout_head. KL anchor to the BC policy prevents collapse. - Multi-GPU A100 training. bf16 mixed precision, fused AdamW,
torch.compile, cosine LR with 2% warmup, DDP across 4x A100 viatorchrun. - Int8 deployment. Dynamic quantization shrinks the ONNX model 3-4x with no measurable win-rate loss. CPU inference stays under a 50 ms p99 budget, well inside the sandbox's 200 ms response window.
┌─────────────────────────┐
│ 1. generate_corpus │ Rust sandbox self-play -> JSONL game logs
│ (8 parallel workers) │ N games
└────────────┬────────────┘
▼
┌─────────────────────────┐
│ 2. build_corpus │ Tokenize JSONL -> memmap binary (uint16)
│ │ train.bin / train.idx / payouts sidecar
└────────────┬────────────┘
▼
┌─────────────────────────┐
│ 3. pretrain_wm │ Decoder-only transformer, LM + aux losses
│ │ 50 epochs cosine LR, DDP across 4x A100
└────────────┬────────────┘
▼
┌─────────────────────────┐
│ 4a. bc_warmstart │ Behavior clone SELF actions, WM frozen
│ │ 30 epochs, factorized type x suit x price head
└────────────┬────────────┘
▼
┌─────────────────────────┐
│ 5. train_payout_head │ Regress WM hidden -> per-player payouts
│ │ 15 epochs, reward source for dream-RL
└────────────┬────────────┘
▼
┌─────────────────────────┐
│ 6. dream_rl │ REINFORCE + value baseline + KL-to-BC
│ │ Imagined rollouts inside frozen WM
└────────────┬────────────┘
▼
┌─────────────────────────┐
│ 7. export_and_eval │ ONNX export, tournament vs baselines,
│ │ pass/fail vs target win rates
└────────────┬────────────┘
▼
┌─────────────────────────┐
│ 8. quantize_and_bench │ Int8 dynamic quantization,
│ │ p50/p95/p99 latency vs 50 ms budget
└─────────────────────────┘
src/
├── data/
│ ├── tokenizer.py factorized vocab (187 tokens), per-perspective encoding
│ ├── build_corpus.py JSONL -> memmap binary corpus
│ └── dataset.py streaming PyTorch dataset, no cross-game windows
├── model/
│ ├── embeddings.py flat nn.Embedding + RoPE
│ ├── world_model.py decoder-only transformer + auxiliary goal/hand heads
│ ├── policy.py factorized policy head (type x suit x price) + value head
│ └── payout_head.py per-player payout regressor
├── training/
│ ├── utils.py shared training utils (bf16, compile, cosine LR, dataloader)
│ ├── loss.py LM loss + multi-task aux loss
│ ├── pretrain_wm.py world-model pre-training
│ ├── bc_warmstart.py behavior-cloning warm-start
│ ├── train_payout_head.py payout-head pre-step
│ ├── dream_rl.py imagination-based RL
│ ├── population.py checkpoint pool for self-play diversity
│ └── distributed.py DDP setup helpers (rank, init, barriers)
└── serving/
├── export.py fp32 ONNX export
├── quantize.py int8 dynamic quantization
├── bench.py per-action latency benchmark
├── inference.py ONNX inference wrapper (handles fp32 & int8)
└── bot.py stdin/stdout sidecar bridge, paired with Rust ExternalPlayer
eval/
├── tournament.py bot vs baseline-agent tournament + pass/fail targets
└── analysis.py calibration + ECE + WM drift KL
sandbox/ (separate git repo, fork of 0xDub/figgie-auto)
└── src/player/external.rs ExternalPlayer + SidecarMsg
configs/ training hyperparameters (YAML)
scripts/ pipeline runners (local + Slurm)
scripts/run_pipeline.sh # full pipeline locally (1 GPU)
NPROC_PER_NODE=4 scripts/run_pipeline.sh # DDP across 4 local GPUs
bash scripts/slurm/submit_all.sh # full pipeline on Slurm (4x A100)
scripts/run_pipeline.sh dream_rl # single stage
N_ROUNDS=10000 EVAL_GAMES=1000 P99_BUDGET_MS=50 \
bash scripts/slurm/submit_all.sh # quick smoke run (override defaults)
ls checkpoints/ # inspect results
cat logs/tournament_*.json
cat logs/bench_int8_*.jsonEvaluated in the Rust sandbox against scripted baseline agents, 10,000 games per matchup.
| Opponent | Target | Figgie-Transformer |
|---|---|---|
| Noisy | ≥ 80% | 91.8% |
| TiltInventory | ≥ 60% | 76.4% |
| Spread | ≥ 50% | 57.2% |
| Overall | - | 75.1% |
3 / 3 baselines cleared, 30,000 evaluation games total.
PolicyInference.act() returns {"action_type": "BID", ...} (uppercase). The Rust parse_bot_order looks for {"action": "bid", ...} (lowercase). Without translation, every bot response was silently treated as a pass.
Fix: FiggieBot._to_sidecar_action() lowercases the action, drops the value field, and emits {"action": "pass"} for PASS.
The sandbox's Event::Update(Update) carries the resulting book state, not the individual quote that caused the change. The Python bot needs the fine-grained event stream ({"type":"bid",...}) to feed its tokenizer.
Fix: added a separate SidecarMsg enum and an mpsc::UnboundedSender<SidecarMsg> from the MatchMaker to the ExternalPlayer, emitted at the same sites where events are logged to disk. Existing Rust players are untouched.
The original spec assumed aux_targets["goal_suit"] per position. The dataset only emits input_ids, labels, and (optionally) payout sidecar data.
Fix: derive the truth from the token sequence. At any position t with input_ids[t] == ROUND_END_TOKEN, labels[t] is the goal-suit token. ECE is computed over 10 equal-width confidence bins.
When the tournament spawns 4 copies of PlayerName::Noisy, player_inventories: HashMap<PlayerName, Inventory> collapses them into one entry, so the ante is decremented 4x per round. The bot's own P&L is correct (External is unique), so the win-rate metric is valid, but opponent-side stats in round_end logs are garbage.
Workaround (out of scope): add PlayerName::Noisy1/2/3/4 or run mixed-opponent tournaments. Documented and tolerated.
A 100-game eval would run cleanly for ~17 rounds, then the sandbox would stop producing round_end events. Tracing showed inference had crept from ~50 ms (ctx_len=400) to ~215 ms (ctx_len=1024), past Rust's 200 ms read_line timeout in ExternalPlayer::start. Responses were silently dropped, the order channel dried up, and match_maker's round loop blocked forever because its duration check only runs between orders.
Fix: dropped MAX_CONTEXT in src/serving/bot.py from 1024 to 384, keeping every inference under ~100 ms. The world model was trained on 128-token windows anyway.
sandbox/is a git submodule pointing to my fork Pranay0302/figgie-auto (which adds theExternalPlayer+SidecarMsgbridge on top of 0xDub/figgie-auto). Clone withgit clone --recursive, or rungit submodule update --initafter a plain clone.corpus/,checkpoints/, andlogs/live outside this repository (see.gitignore).- The world model's quality bounds everything. Validate
wm_driftandECE(viaeval/analysis.py) before spending compute on the imagination loop. - Defaults target ~5-20 M params; bigger overfits the agent zoo.
- Figgie was designed by Jane Street. All rules, mechanics, and the game itself are their work; this project just learns to play it.
- The Rust sandbox under
sandbox/is a fork of 0xDub/figgie-auto, which provides the engine, baseline agents, and the player-trait scaffolding theExternalPlayersidecar plugs into.