diff --git a/config/chess.ini b/config/chess.ini
new file mode 100644
index 0000000000..92f9630b13
--- /dev/null
+++ b/config/chess.ini
@@ -0,0 +1,101 @@
+[base]
+env_name = chess
+
+[selfplay]
+enabled = 1
+max_size = 500
+swap_winrate = 0.6
+min_games = 4096
+snapshot_interval = 1_000_000_000
+opp_timeout_steps = 4_000_000_000
+
+[vec]
+total_agents = 8192
+num_buffers = 1
+num_threads = 2
+num_frozen_banks = 1
+frozen_bank_pct = 0.1
+
+[env]
+max_moves = 5000
+reward_draw = 0
+reward_invalid_piece = 0
+reward_invalid_move = 0
+reward_repetition = 0
+enable_50_move_rule = 1
+enable_threefold_repetition = 1
+mode = 1
+random_fen = 0
+render_fps = 30
+fen_curric_pct = 0.9
+
+[policy]
+hidden_size = 512
+num_layers = 3
+expansion_factor = 1
+
+[train]
+gpus = 1
+seed = 42
+total_timesteps = 1_000_000_000_000
+learning_rate = 0.000572786
+anneal_lr = 1
+min_lr_ratio = 0.086914
+gamma = 0.994795
+gae_lambda = 0.754641
+replay_ratio = 0.25
+clip_coef = 0.557019
+vf_coef = 4.37465
+vf_clip_coef = 1.69524
+max_grad_norm = 3.54293
+ent_coef = 0.0984801
+anneal_ent_coef = 1
+min_ent_coef_ratio = 0.1
+beta1 = 0.972205
+beta2 = 0.9
+eps = 8.51435e-11
+minibatch_size = 32768
+horizon = 64
+vtrace_rho_clip = 1.99404
+vtrace_c_clip = 2.19484
+prio_alpha = 1
+prio_beta0 = 0.474524
+torch_deterministic = 1
+adam_beta1 = 0.963206
+adam_beta2 = 0.99999
+adam_eps = 5.09326e-08
+update_epochs = 1
+[sweep]
+# Score each trial by winrate in a 2-policy match against a fixed enemy rather
+# than the training-time self-play env/score. 'latest' resolves to the newest
+# .bin in checkpoint_dir/chess/**. Enemy arch must match the checkpoint's arch.
+match_enemy_model_path = 'resources/chess/10b_weights.bin'
+match_num_games = 4096
+match_enemy_hidden_size = 512
+match_enemy_num_layers = 3
+
+[sweep.train.total_timesteps]
+distribution = log_normal
+min = 7e9
+max = 15e9
+mean = 10e9
+scale = time
+
+[sweep.selfplay.swap_winrate]
+distribution = uniform
+min = 0.55
+max = 0.90
+scale = auto
+
+[sweep.vec.num_buffers]
+distribution = uniform_pow2
+min = 1
+max = 8
+scale = auto
+
+[sweep.vec.frozen_bank_pct]
+distribution = uniform
+min = 0.05
+max = 0.5
+scale = auto
+
diff --git a/config/default.ini b/config/default.ini
index fd53bf7d95..29bc1808b7 100644
--- a/config/default.ini
+++ b/config/default.ini
@@ -9,7 +9,7 @@ nccl_id = 'None'
 profile = False
 checkpoint_dir = checkpoints
 log_dir = logs
-checkpoint_interval = 200
+checkpoint_interval = 500
 eval_episodes = 10000
 
 # Epoch at which to capture CUDA graphs. -1 to disable.
@@ -25,6 +25,25 @@ total_agents = 4096
 num_buffers = 2
 num_threads = 16
 
+# Selfplay-pool training. When enabled, frozen_bank_pct of agent slots (see
+# [vec]) play against a snapshot of an older policy from a disk-backed pool.
+# Primary advances to the next pool entry once it beats the current opponent
+# at >= swap_winrate over at least min_games (counted only on historical envs).
+[selfplay]
+enabled = 0
+max_size = 16
+swap_winrate = 0.8
+min_games = 2048
+elo_init = 0.0
+elo_k = 16.0
+seed = 42
+# Add a snapshot to the pool every snapshot_interval global steps, independent
+# of swap. 0 disables interval snapshotting (pool stays at bootstrap).
+snapshot_interval = 1_000_000_000
+# Force an opponent swap if the current opponent has been active for this many
+# global steps without a winrate-driven swap. 0 disables the timeout.
+opp_timeout_steps = 500_000_000
+
 # Args used by your env's binding.c go here
 [env]
 
@@ -53,6 +72,12 @@ vf_coef = 2.0
 vf_clip_coef = 0.2
 max_grad_norm = 1.5
 ent_coef = 0.001
+# Cosine-anneal ent_coef like lr does. Off by default (matches old behavior).
+# When on: ent_coef decays from its base value to min_ent_coef_ratio * ent_coef
+# over total_timesteps. Useful late in training to let a learned policy commit
+# harder on its preferences once the entropy bonus is no longer load-bearing.
+anneal_ent_coef = 0
+min_ent_coef_ratio = 0.1
 beta1 = 0.95
 beta2 = 0.999
 eps = 1e-12
@@ -75,6 +100,13 @@ downsample = 5
 use_gpu = True
 prune_pareto = True
 early_stop_quantile = 0.3
+# When set, each sweep trial is scored by winrate in a match against a fixed
+# enemy checkpoint rather than by the training-time env/score. Score key emitted
+# as env/match_score; set match_enemy_model_path to '' to disable.
+match_enemy_model_path = ''
+match_num_games = 1024
+match_enemy_hidden_size = 0
+match_enemy_num_layers = 0
 
 [sweep.train.total_timesteps]
 distribution = log_normal
@@ -94,11 +126,11 @@ min = 1
 max = 8
 scale = auto
 
-[sweep.vec.total_agents]
-distribution = uniform_pow2
-min = 256
-max = 16384
-scale = auto
+#[sweep.vec.total_agents]
+#distribution = uniform_pow2
+#min = 256
+#max = 16384
+#scale = auto
 
 [sweep.vec.num_buffers]
 distribution = uniform
@@ -112,11 +144,11 @@ min = 8
 max = 1024
 scale = auto
 
-[sweep.train.minibatch_size]
-distribution = uniform_pow2
-min = 4096
-max = 65536
-scale = auto
+#[sweep.train.minibatch_size]
+#distribution = uniform_pow2
+#min = 4096
+#max = 65536
+#scale = auto
 
 [sweep.train.learning_rate]
 distribution = log_normal
@@ -139,7 +171,6 @@ scale = auto
 [sweep.train.gae_lambda]
 distribution = logit_normal
 min = 0.2
-#min = 0.6
 max = 0.995
 scale = auto
 
diff --git a/ocean/chess/binding.c b/ocean/chess/binding.c
new file mode 100644
index 0000000000..1a3cd74d63
--- /dev/null
+++ b/ocean/chess/binding.c
@@ -0,0 +1,218 @@
+#include "chess.h"
+// Before embedding approach
+// #define OBS_SIZE 1082
+#define OBS_SIZE 167
+#define NUM_ATNS 1
+#define ACT_SIZES {97}
+#define OBS_TENSOR_T ByteTensor
+#define MY_ACTION_MASK 97
+
+#define MY_VEC_INIT
+#define MY_VEC_CLOSE
+#define MY_USES_PERM
+#define MY_USES_TAGS
+#define Env Chess
+#include "vecenv.h"
+
+void my_setup_perm(StaticVec* vec, Env* env, int slot_base) {
+    size_t obs_elem_size = obs_element_size();
+    for (int s = 0; s < env->num_agents; s++) {
+        int phys = vec->agent_perm ? vec->agent_perm[slot_base + s] : (slot_base + s);
+        env->obs_ptr[s]         = (uint8_t*)vec->observations + (size_t)phys * OBS_SIZE * obs_elem_size;
+        env->action_mask_ptr[s] = vec->action_mask + (size_t)phys * MY_ACTION_MASK;
+        env->action_ptr[s]      = vec->actions + (size_t)phys * NUM_ATNS;
+        env->reward_ptr[s]      = vec->rewards + phys;
+        env->terminal_ptr[s]    = vec->terminals + phys;
+    }
+}
+
+#define DEFAULT_STARTING_FEN "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
+#define FEN_CURRICULUM_PATH "resources/chess/fens.txt"
+
+static char** SHARED_FEN_CURRICULUM = NULL;
+static int SHARED_NUM_FENS = 0;
+
+static char** load_fen_file(const char* path, int* num_fens_out) {
+    FILE* f = fopen(path, "r");
+    if (f == NULL) {
+        *num_fens_out = 0;
+        return NULL;
+    }
+
+    int num_fens = 0;
+    char line[256];
+    while (fgets(line, sizeof(line), f)) {
+        if (line[0] != '#' && line[0] != '\n' && line[0] != '\r') {
+            num_fens++;
+        }
+    }
+    if (num_fens == 0) {
+        fclose(f);
+        *num_fens_out = 0;
+        return NULL;
+    }
+
+    char** fens = (char**)malloc(num_fens * sizeof(char*));
+    rewind(f);
+    int idx = 0;
+    while (fgets(line, sizeof(line), f) && idx < num_fens) {
+        if (line[0] != '#' && line[0] != '\n' && line[0] != '\r') {
+            size_t len = strlen(line);
+            while (len > 0 && (line[len-1] == '\n' || line[len-1] == '\r')) {
+                line[--len] = '\0';
+            }
+            fens[idx++] = strdup(line);
+        }
+    }
+    fclose(f);
+    *num_fens_out = num_fens;
+    return fens;
+}
+
+static void apply_kwargs(Env* env, Dict* kwargs) {
+    env->max_moves = (int)dict_get(kwargs, "max_moves")->value;
+    env->reward_draw = (float)dict_get(kwargs, "reward_draw")->value;
+    env->reward_invalid_piece = (float)dict_get(kwargs, "reward_invalid_piece")->value;
+    env->reward_invalid_move = (float)dict_get(kwargs, "reward_invalid_move")->value;
+    env->reward_repetition = (float)dict_get(kwargs, "reward_repetition")->value;
+    env->render_fps = (int)dict_get(kwargs, "render_fps")->value;
+    env->mode = (int)dict_get(kwargs, "mode")->value;
+    env->enable_50_move_rule = (int)dict_get(kwargs, "enable_50_move_rule")->value;
+    env->enable_threefold_repetition = (int)dict_get(kwargs, "enable_threefold_repetition")->value;
+    env->random_fen = (int)dict_get(kwargs, "random_fen")->value;
+    env->fen_curric_pct = (float)dict_get(kwargs, "fen_curric_pct")->value;
+
+    env->client = NULL;
+    env->legal_dirty = 1;
+    env->human_color = -1;
+    env->log_pgn = 0;
+    env->log_pgn_choice_made = 1;
+    env->pgn_filename[0] = '\0';
+    env->pgn_game_number = 0;
+    env->maia_pid = 0;
+    env->maia_stdin_fd = -1;
+    env->maia_stdout_fd = -1;
+    env->maia_phase = 0;
+    strcpy(env->starting_fen, DEFAULT_STARTING_FEN);
+    strcpy(env->last_result, "Game starting...");
+}
+
+Env* my_vec_init(int* num_envs_out, int* buffer_env_starts, int* buffer_env_counts,
+                 Dict* vec_kwargs, Dict* env_kwargs) {
+    int total_agents = (int)dict_get(vec_kwargs, "total_agents")->value;
+    int num_buffers = (int)dict_get(vec_kwargs, "num_buffers")->value;
+    int agents_per_buffer = total_agents / num_buffers;
+
+    float curric_pct = (float)dict_get(env_kwargs, "fen_curric_pct")->value;
+    if (curric_pct > 0.0f && SHARED_FEN_CURRICULUM == NULL) {
+        SHARED_FEN_CURRICULUM = load_fen_file(FEN_CURRICULUM_PATH, &SHARED_NUM_FENS);
+        if (SHARED_FEN_CURRICULUM != NULL) {
+            printf("Loaded %d FENs from %s\n", SHARED_NUM_FENS, FEN_CURRICULUM_PATH);
+        }
+    }
+
+    int mode = (int)dict_get(env_kwargs, "mode")->value;
+    int agents_per_env = (mode == CHESS_MODE_SELFPLAY) ? 2 : 1;
+    int num_envs = total_agents / agents_per_env;
+    Env* envs = (Env*)calloc(num_envs, sizeof(Env));
+    for (int i = 0; i < num_envs; i++) {
+        Env* env = &envs[i];
+        apply_kwargs(env, env_kwargs);
+        env->num_agents = agents_per_env;
+        env->rng = i;
+        // In selfplay, learner_color is unused; the slot↔color mapping is per-env
+        // randomized so policies in fixed slots see both colors equally.
+        env->learner_color = (agents_per_env == 1) ? (i % 2) : CHESS_WHITE;
+        if (agents_per_env == 2 && (i & 1)) {
+            env->slot_for_color[CHESS_WHITE] = 1;
+            env->slot_for_color[CHESS_BLACK] = 0;
+        } else {
+            env->slot_for_color[CHESS_WHITE] = 0;
+            env->slot_for_color[CHESS_BLACK] = 1;
+        }
+        env->fen_curriculum = SHARED_FEN_CURRICULUM;
+        env->num_fens = SHARED_NUM_FENS;
+        init_bitboards();
+    }
+
+    int buf = 0;
+    int buf_agents = 0;
+    buffer_env_starts[0] = 0;
+    buffer_env_counts[0] = 0;
+    for (int i = 0; i < num_envs; i++) {
+        buf_agents += agents_per_env;
+        buffer_env_counts[buf]++;
+        if (buf_agents >= agents_per_buffer && buf < num_buffers - 1) {
+            buf++;
+            buffer_env_starts[buf] = i + 1;
+            buffer_env_counts[buf] = 0;
+            buf_agents = 0;
+        }
+    }
+
+    *num_envs_out = num_envs;
+    return envs;
+}
+
+void my_vec_close(Env* envs) {
+    if (SHARED_FEN_CURRICULUM != NULL) {
+        for (int i = 0; i < SHARED_NUM_FENS; i++) {
+            free(SHARED_FEN_CURRICULUM[i]);
+        }
+        free(SHARED_FEN_CURRICULUM);
+        SHARED_FEN_CURRICULUM = NULL;
+        SHARED_NUM_FENS = 0;
+    }
+}
+
+void my_init(Env* env, Dict* kwargs) {
+    apply_kwargs(env, kwargs);
+    env->num_agents = (env->mode == CHESS_MODE_SELFPLAY) ? 2 : 1;
+    env->learner_color = (env->num_agents == 1) ? CHESS_WHITE : CHESS_WHITE;
+    env->slot_for_color[CHESS_WHITE] = 0;
+    env->slot_for_color[CHESS_BLACK] = 1;
+    env->fen_curriculum = NULL;
+    env->num_fens = 0;
+    init_bitboards();
+}
+
+void my_log(Log* log, Dict* out) {
+    dict_set(out, "perf", log->perf);
+    dict_set(out, "score", log->score);
+    dict_set(out, "draw_rate", log->draw_rate);
+    dict_set(out, "timeout_rate", log->timeout_rate);
+    dict_set(out, "chess_moves", log->chess_moves);
+    dict_set(out, "episode_length", log->episode_length);
+    dict_set(out, "episode_return", log->episode_return);
+    dict_set(out, "invalid_action_rate", log->invalid_action_rate);
+    dict_set(out, "slot_0_score", log->slot_0_score);
+    dict_set(out, "slot_1_score", log->slot_1_score);
+    dict_set(out, "hist_score", log->hist_score);
+    dict_set(out, "hist_n", log->hist_n);
+    // Per-bank historical stats for multi-bank selfplay. selfplay.py reads
+    // hist_score_bank_<b> / hist_n_bank_<b> to drive each bank's swap decision.
+    // dict_set stores the key pointer (vecenv.h:61), not a copy, so we MUST
+    // use string literals here — a stack buffer in a loop aliases and collapses
+    // all 16 entries into one. Sized to CHESS_MAX_BANKS = 8.
+    dict_set(out, "hist_score_bank_0", log->hist_score_bank[0]);
+    dict_set(out, "hist_score_bank_1", log->hist_score_bank[1]);
+    dict_set(out, "hist_score_bank_2", log->hist_score_bank[2]);
+    dict_set(out, "hist_score_bank_3", log->hist_score_bank[3]);
+    dict_set(out, "hist_score_bank_4", log->hist_score_bank[4]);
+    dict_set(out, "hist_score_bank_5", log->hist_score_bank[5]);
+    dict_set(out, "hist_score_bank_6", log->hist_score_bank[6]);
+    dict_set(out, "hist_score_bank_7", log->hist_score_bank[7]);
+    dict_set(out, "hist_n_bank_0", log->hist_n_bank[0]);
+    dict_set(out, "hist_n_bank_1", log->hist_n_bank[1]);
+    dict_set(out, "hist_n_bank_2", log->hist_n_bank[2]);
+    dict_set(out, "hist_n_bank_3", log->hist_n_bank[3]);
+    dict_set(out, "hist_n_bank_4", log->hist_n_bank[4]);
+    dict_set(out, "hist_n_bank_5", log->hist_n_bank[5]);
+    dict_set(out, "hist_n_bank_6", log->hist_n_bank[6]);
+    dict_set(out, "hist_n_bank_7", log->hist_n_bank[7]);
+    dict_set(out, "wins_as_white", log->wins_as_white);
+    dict_set(out, "wins_as_black", log->wins_as_black);
+    dict_set(out, "games_as_white", log->games_as_white);
+    dict_set(out, "games_as_black", log->games_as_black);
+    dict_set(out, "maia_failures", log->maia_failures);
+}
diff --git a/ocean/chess/chess.h b/ocean/chess/chess.h
new file mode 100644
index 0000000000..2853232f1e
--- /dev/null
+++ b/ocean/chess/chess.h
@@ -0,0 +1,3167 @@
+#include <stdlib.h>
+#include <stdio.h>
+#include <string.h>
+#include <stdint.h>
+#include <stdbool.h>
+#include <assert.h>
+#include <math.h>
+#include <time.h>
+#include <unistd.h> 
+#include "raylib.h"
+
+typedef uint64_t Bitboard;
+typedef uint64_t Key;
+typedef uint32_t Square;
+typedef uint32_t Move;
+typedef uint32_t Piece;
+typedef uint8_t ChessColor;
+
+enum {
+    SQ_A1, SQ_B1, SQ_C1, SQ_D1, SQ_E1, SQ_F1, SQ_G1, SQ_H1,
+    SQ_A2, SQ_B2, SQ_C2, SQ_D2, SQ_E2, SQ_F2, SQ_G2, SQ_H2,
+    SQ_A3, SQ_B3, SQ_C3, SQ_D3, SQ_E3, SQ_F3, SQ_G3, SQ_H3,
+    SQ_A4, SQ_B4, SQ_C4, SQ_D4, SQ_E4, SQ_F4, SQ_G4, SQ_H4,
+    SQ_A5, SQ_B5, SQ_C5, SQ_D5, SQ_E5, SQ_F5, SQ_G5, SQ_H5,
+    SQ_A6, SQ_B6, SQ_C6, SQ_D6, SQ_E6, SQ_F6, SQ_G6, SQ_H6,
+    SQ_A7, SQ_B7, SQ_C7, SQ_D7, SQ_E7, SQ_F7, SQ_G7, SQ_H7,
+    SQ_A8, SQ_B8, SQ_C8, SQ_D8, SQ_E8, SQ_F8, SQ_G8, SQ_H8,
+    SQ_NONE = 64
+};
+
+enum { PAWN = 1, KNIGHT, BISHOP, ROOK, QUEEN, KING };
+
+enum {
+    NO_PIECE = 0,
+    W_PAWN = 1, W_KNIGHT, W_BISHOP, W_ROOK, W_QUEEN, W_KING,
+    B_PAWN = 9, B_KNIGHT, B_BISHOP, B_ROOK, B_QUEEN, B_KING
+};
+
+enum { CHESS_WHITE = 0, CHESS_BLACK = 1 };
+
+enum {
+    NO_CASTLING = 0,
+    WHITE_OO = 1, WHITE_OOO = 2,
+    BLACK_OO = 4, BLACK_OOO = 8,
+    WHITE_CASTLING = 3, BLACK_CASTLING = 12
+};
+
+
+enum { NORMAL, PROMOTION, ENPASSANT, CASTLING };
+
+enum {
+    NORTH = 8, EAST = 1, SOUTH = -8, WEST = -1,
+    NORTH_EAST = 9, SOUTH_EAST = -7,
+    NORTH_WEST = 7, SOUTH_WEST = -9
+};
+
+#define MOVE_NONE 0
+#define MOVE_NULL 65
+
+static inline Move make_move(Square from, Square to) {
+    return (Move)(to | (from << 6));
+}
+
+static inline Move make_promotion(Square from, Square to, int pt) {
+    return (Move)(to | (from << 6) | (PROMOTION << 14) | ((pt - KNIGHT) << 12));
+}
+
+static inline Move make_enpassant(Square from, Square to) {
+    return (Move)(to | (from << 6) | (ENPASSANT << 14));
+}
+
+static inline Move make_castling(Square from, Square to) {
+    return (Move)(to | (from << 6) | (CASTLING << 14));
+}
+
+static inline Square from_sq(Move m) {
+    return (Square)((m >> 6) & 0x3f);
+}
+
+static inline Square to_sq(Move m) {
+    return (Square)(m & 0x3f);
+}
+
+static inline int type_of_m(Move m) {
+    return (int)(m >> 14);
+}
+
+static inline int promotion_type(Move m) {
+    return (int)(((m >> 12) & 3) + KNIGHT);
+}
+
+static inline Square make_square(int f, int r) {
+    return (Square)((r << 3) + f);
+}
+
+static inline int file_of(Square s) {
+    return (int)(s & 7);
+}
+
+static inline int rank_of(Square s) {
+    return (int)(s >> 3);
+}
+
+static inline Piece make_piece(int c, int pt) {
+    return (Piece)((c << 3) + pt);
+}
+
+static inline int type_of_p(Piece p) {
+    return (int)(p & 7);
+}
+
+static inline int color_of(Piece p) {
+    return (int)(p >> 3);
+}
+#define MAX_GAME_PLIES 2048
+
+#define FileABB 0x0101010101010101ULL
+#define FileBBB (FileABB << 1)
+#define FileCBB (FileABB << 2)
+#define FileDBB (FileABB << 3)
+#define FileEBB (FileABB << 4)
+#define FileFBB (FileABB << 5)
+#define FileGBB (FileABB << 6)
+#define FileHBB (FileABB << 7)
+
+#define Rank1BB 0xFFULL
+#define Rank2BB (Rank1BB << 8)
+#define Rank3BB (Rank1BB << 16)
+#define Rank4BB (Rank1BB << 24)
+#define Rank5BB (Rank1BB << 32)
+#define Rank6BB (Rank1BB << 40)
+#define Rank7BB (Rank1BB << 48)
+#define Rank8BB (Rank1BB << 56)
+
+#define SQ_FEATURES 15
+
+static const char* PIECE_CHARS[] = {
+    "",
+    "P", "N", "B", "R", "Q", "K",
+    "", "",
+    "p", "n", "b", "r", "q", "k"
+};
+
+static const char* PIECE_FILLED[] = {
+    "",
+    "♟", "♞", "♝", "♜", "♛", "♚",
+    "", "",
+    "♟", "♞", "♝", "♜", "♛", "♚"
+};
+
+
+static uint64_t prng_state = 1070372;
+static inline uint64_t prng_rand(void) {
+    prng_state ^= prng_state >> 12;
+    prng_state ^= prng_state << 25;
+    prng_state ^= prng_state >> 27;
+    return prng_state * 2685821657736338717ULL;
+}
+
+extern Bitboard SquareBB[65];
+extern Bitboard PawnAttacks[2][64];
+extern Bitboard KnightAttacks[64];
+extern Bitboard KingAttacks[64];
+extern Bitboard BetweenBB[64][64];
+extern Bitboard LineBB[64][64];
+
+static Bitboard BishopMasks[64];
+static uint64_t BishopMagics[64];
+static int BishopShifts[64];
+static Bitboard BishopTable[64 * 512];
+static Bitboard* BishopAttacks[64];
+static const uint64_t BISHOP_MAGICS[64] = {
+    9368648609924554880ULL, 9009475591934976ULL,     4504776450605056ULL,
+    1130334595844096ULL,    1725202480235520ULL,     288516396277699584ULL,
+    613618303369805920ULL,  10168455467108368ULL,    9046920051966080ULL,
+    36031066926022914ULL,   1152925941509587232ULL,  9301886096196101ULL,
+    290536121828773904ULL,  5260205533369993472ULL,  7512287909098426400ULL,
+    153141218749450240ULL,  9241386469758076456ULL,  5352528174448640064ULL,
+    2310346668982272096ULL, 1154049638051909890ULL,  282645627930625ULL,
+    2306405976892514304ULL, 11534281888680707074ULL, 72339630111982113ULL,
+    8149474640617539202ULL, 2459884588819024896ULL,  11675583734899409218ULL,
+    1196543596102144ULL,    5774635144585216ULL,     145242600416216065ULL,
+    2522607328671633440ULL, 145278609400071184ULL,   5101802674455216ULL,
+    650979603259904ULL,     9511646410653040801ULL,  1153493285013424640ULL,
+    18016048314974752ULL,   4688397299729694976ULL,  9226754220791842050ULL,
+    4611969694574863363ULL, 145532532652773378ULL,   5265289125480634376ULL,
+    288239448330604544ULL,  2395019802642432ULL,     14555704381721968898ULL,
+    2324459974457168384ULL, 23652833739932677ULL,    282583111844497ULL,
+    4629880776036450560ULL, 5188716322066279440ULL,  146367151686549765ULL,
+    1153170821083299856ULL, 2315697107408912522ULL,  2342448293961403408ULL,
+    2309255902098161920ULL, 469501395595331584ULL,   4615626809856761874ULL,
+    576601773662552642ULL,  621501155230386208ULL,   13835058055890469376ULL,
+    3748138521932726784ULL, 9223517207018883457ULL,  9237736128969216257ULL,
+    1127068154855556ULL,
+};
+
+static Bitboard RookMasks[64];
+static uint64_t RookMagics[64];
+static int RookShifts[64];
+static Bitboard RookTable[64 * 4096];  
+static Bitboard* RookAttacks[64];
+static const uint64_t ROOK_MAGICS[64] = {
+    612498416294952992ULL,  2377936612260610304ULL,  36037730568766080ULL,
+    72075188908654856ULL,   144119655536003584ULL,   5836666216720237568ULL,
+    9403535813175676288ULL, 1765412295174865024ULL,  3476919663777054752ULL,
+    288300746238222339ULL,  9288811671472386ULL,     146648600474026240ULL,
+    3799946587537536ULL,    704237264700928ULL,      10133167915730964ULL,
+    2305983769267405952ULL, 9223634270415749248ULL,  10344480540467205ULL,
+    9376496898355021824ULL, 2323998695235782656ULL,  9241527722809755650ULL,
+    189159985010188292ULL,  2310421375767019786ULL,  4647717014536733827ULL,
+    5585659813035147264ULL, 1442911135872321664ULL,  140814801969667ULL,
+    1188959108457300100ULL, 288815318485696640ULL,   758869733499076736ULL,
+    234750139167147013ULL,  2305924931420225604ULL,  9403727128727390345ULL,
+    9223970239903959360ULL, 309094713112139074ULL,   38290492990967808ULL,
+    3461016597114651648ULL, 181289678366835712ULL,   4927518981226496513ULL,
+    1155212901905072225ULL, 36099167912755202ULL,    9024792514543648ULL,
+    4611826894462124048ULL, 291045264466247688ULL,   83880127713378308ULL,
+    1688867174481936ULL,    563516973121544ULL,      9227888831703941123ULL,
+    703691741225216ULL,     45203259517829248ULL,    693563138976596032ULL,
+    4038638777286134272ULL, 865817582546978176ULL,   13835621555058516608ULL,
+    11541041685463296ULL,   288511853443695360ULL,   283749161902275ULL,
+    176489098445378ULL,     2306124759338845321ULL,  720584805193941061ULL,
+    4977040710267061250ULL, 10097633331715778562ULL, 325666550235288577ULL,
+    1100057149646ULL,
+};
+
+typedef struct {
+    Key psq[16][64];
+    Key enpassant[8];
+    Key castling[16];
+    Key side;
+} Zobrist;
+
+extern Zobrist zob;
+
+typedef struct {
+    Bitboard byTypeBB[7];    // [0]=all, [1-6]=PAWN,KNIGHT,BISHOP,ROOK,QUEEN,KING
+    Bitboard byColorBB[2];
+    uint8_t board[64];
+    uint8_t pieceCount[16];
+    ChessColor sideToMove;
+    uint8_t castlingRights;
+    uint8_t epSquare;
+    uint8_t rule50;
+    Key key;
+} Position;
+
+static inline Bitboard pieces(const Position* pos) {
+    return pos->byTypeBB[0];
+}
+
+static inline Bitboard pieces_p(const Position* pos, int p) {
+    return pos->byTypeBB[p];
+}
+
+static inline Bitboard pieces_c(const Position* pos, int c) {
+    return pos->byColorBB[c];
+}
+
+static inline Bitboard pieces_cp(const Position* pos, int c, int p) {
+    return pieces_p(pos, p) & pieces_c(pos, c);
+}
+
+static inline Piece piece_on(const Position* pos, Square s) {
+    return (Piece)pos->board[s];
+}
+
+typedef struct {
+    Move move;
+} ExtMove;
+
+typedef struct {
+    ExtMove moves[256];
+    int count;
+} MoveList;
+/*
+enum {
+    // Relational NNUE tokens
+    O_TOKEN_COUNT = 0,                    
+    O_TOKEN_DATA  = 2,                  
+    // Meta data
+    O_SIDE = 130,                      
+    O_CASTLE = 132,                   
+    O_EP = 148,                      
+    O_PICK_PHASE = 213,             
+    O_SELECTED_PIECE = 215,        
+    O_VALID_PIECES = 279,         
+    O_VALID_DESTS = 343,         
+    O_VALID_PROMOS = 407,       
+
+    O_SELF_CHECK = 439,
+    O_OPP_CHECK = 440,
+    O_RULE50 = 441,
+    O_REPETITION = 442,
+    O_PASS_VALID = 443,
+
+    OBS_SIZE = 444
+};
+*/
+/*enum {
+    O_SQUARES = 0,
+    O_VALID_PROMOS = 960,
+    O_CASTLE = 992,
+    O_EP = 993,
+    O_PICK_PHASE = 994,
+    O_RULE50 = 995,
+    O_REPETITION = 996,
+    O_PASS_VALID = 997,
+    OBS_SIZE = 998
+};
+*/
+
+/*
+Selfplay branch obs layout before embedding approach
+enum {
+    O_BOARD = 0,
+    O_SIDE = 768,
+    O_CASTLE = 770,
+    O_EP = 786,
+    O_PICK_PHASE = 851,
+    O_SELECTED_PIECE = 853,
+    O_VALID_PIECES = 917,
+    O_VALID_DESTS = 981,
+    O_VALID_PROMOS = 1045,
+    O_SELF_CHECK = 1077,
+    O_OPP_CHECK = 1078,
+    O_RULE50 = 1079,
+    O_REPETITION = 1080,
+    O_PASS_VALID = 1081,
+    OBS_SIZE = 1082
+};
+*/
+enum {
+    O_BOARD = 0,
+    O_SIDE = 64,
+    O_CASTLE = 65,
+    O_EP = 69,
+    O_RULE50 = 78,
+    O_REPETITION = 79,
+    O_SELF_CHECK = 80,
+    O_OPP_CHECK = 81,
+    O_PICK_PHASE = 82,
+    O_SELECTED_PIECE = 83,
+    O_VALID_FROM_COUNT = 84,
+    O_VALID_FROM = 85,
+    O_VALID_TO_COUNT = 101,
+    O_VALID_TO = 102,
+    O_VALID_PROMOS = 134,
+    O_PASS_VALID = 166,
+    OBS_SIZE = 167
+};
+
+#define CHESS_MAX_VALID_FROM 16
+#define CHESS_MAX_VALID_TO   32
+#define CHESS_NULL_SQ        64
+
+#define PASS_ACTION 96
+#define NUM_ACTIONS 97
+enum {
+    CHESS_MODE_RANDOM = 0,
+    CHESS_MODE_SELFPLAY = 1,
+    CHESS_MODE_HUMAN = 2,
+    CHESS_MODE_HUMAN_RANDOM = 3,
+    CHESS_MODE_MAIA = 4
+};
+
+#define CHESS_TAG_SELFPLAY 0
+#define CHESS_TAG_HISTORICAL 1
+// Multi-bank selfplay: env tags are 1..CHESS_MAX_BANKS, one tag value per
+// frozen bank. tag = 0 means pure selfplay env (no historical opponent).
+// Backward compat: with num_frozen_banks=1, tag=1 still means "play bank 0".
+#define CHESS_MAX_BANKS 8
+
+typedef struct {
+    float perf;
+    float score;
+    float draw_rate;
+    float timeout_rate;
+    float chess_moves;
+    float episode_length;
+    float episode_return;
+    float invalid_action_rate;
+    // Per-slot scores (selfplay only). In match: slot 0 = primary policy A,
+    // slot 1 = frozen policy B. In selfplay training both should average ~0.5.
+    float slot_0_score;
+    float slot_1_score;
+    // Per-bank historical tracking. hist_score_bank[b] sums primary's score
+    // (1.0 win / 0.5 draw / 0 loss) on historical envs tagged b+1; hist_n_bank
+    // counts those games. Python recovers per-bank winrate as score/n. The
+    // legacy aggregates hist_score / hist_n sum across all banks for backward
+    // compat with single-bank dashboards.
+    float hist_score;
+    float hist_n;
+    float hist_score_bank[CHESS_MAX_BANKS];
+    float hist_n_bank[CHESS_MAX_BANKS];
+    float n;
+    // Eval diagnostics (non-selfplay modes). Per-color score/games let Python
+    // compute per-color win rate and surface obs-flip / perspective bugs as
+    // lopsided splits. maia_failures counts how often maia_get_move returned
+    // MOVE_NONE and we fell back to a random legal move — non-zero means part
+    // of the eval is degraded.
+    float wins_as_white;
+    float wins_as_black;
+    float games_as_white;
+    float games_as_black;
+    float maia_failures;
+} Log;
+
+typedef struct {
+    int cell_size;
+    Font piece_font;
+    int use_unicode_pieces;
+} Client;
+
+typedef struct {
+    Piece captured;
+    uint8_t castlingRights;
+    uint8_t epSquare;
+    uint8_t rule50;
+    Key key;
+    uint8_t pliesFromNull;
+} UndoInfo;
+
+typedef struct {
+    Log log;
+    Client* client;
+    uint8_t* observations;
+    float* actions;
+    float* rewards;
+    float* terminals;
+    unsigned char* action_mask;   // (97,) — NULL unless MY_ACTION_MASK is defined
+
+    // Per-slot pointers used by the env body. Identity perm = same addresses as
+    // base+stride; non-identity perm = where this slot actually lives in vec
+    // global buffers. Populated by my_setup_perm.
+    uint8_t* obs_ptr[2];
+    unsigned char* action_mask_ptr[2];
+    float* action_ptr[2];
+    float* reward_ptr[2];
+    float* terminal_ptr[2];
+
+    unsigned int rng;
+    int num_agents;
+
+    Position pos;
+    MoveList legal_moves;
+    int legal_dirty;
+    int game_result;
+    int tick;
+    int chess_moves;
+    int max_moves;
+    float reward_draw;
+    float episode_reward;
+    int render_fps;
+    int mode;
+    
+    char starting_fen[128];
+    char** fen_curriculum;
+    float fen_curric_pct;
+    int num_fens;
+    int random_fen;
+    
+    UndoInfo undo_stack[MAX_GAME_PLIES];
+    int undo_stack_ptr;
+    uint8_t repetition_matches;
+    
+    int invalid_actions_this_episode;
+    
+    int pick_phase[2];
+    Square selected_square[2];
+    MoveList valid_destinations[2];
+    Bitboard valid_from_mask[2];
+    Bitboard valid_to_mask[2];
+    Bitboard obs_selected_view_mask[2];
+    Bitboard obs_valid_from_view_mask[2];
+    Bitboard obs_valid_to_view_mask[2];
+    uint32_t obs_valid_promo_mask[2];
+    float reward_invalid_piece;
+    float reward_invalid_move;
+    float reward_repetition;
+    
+    int enable_50_move_rule;
+    int enable_threefold_repetition;
+    
+    int learner_color;
+    // Selfplay-pool tagging. tag = 0 selfplay, tag = 1 historical (slot 0 =
+    // primary, slot 1 = frozen). boundary_reached set on game-end so Python can
+    // detect when historical envs have all completed at least one game since
+    // the last swap arm.
+    int tag;
+    int boundary_reached;
+    // Selfplay only: slot_for_color[c] = which slot (0 or 1) plays color c.
+    // Default (slot 0 = WHITE, slot 1 = BLACK); randomized per env to remove
+    // white-bias when running matched policies in different slots.
+    int slot_for_color[2];
+    int human_color;
+    float white_score;
+    float black_score;
+    float learner_wins;
+    float learner_losses; 
+    float learner_draws;
+    char last_result[32];
+    
+    Move pgn_moves[MAX_GAME_PLIES];
+    int pgn_move_count;
+    int show_game_end_popup;
+    
+    int log_pgn;
+    int log_pgn_choice_made;
+    char pgn_filename[128];
+    int pgn_game_number;
+    
+    int white_captured[6];
+    int black_captured[6];
+    int render_paused;
+    int has_last_move_highlight;
+    Square last_move_from;
+    Square last_move_to;
+
+    // CHESS_MODE_MAIA: per-env lc0 subprocess pipes. Initialized lazily on the
+    // first opponent move. -1 / 0 means "not yet spawned".
+    int maia_pid;
+    int maia_stdin_fd;
+    int maia_stdout_fd;
+    // Maia commits its move in one UCI round-trip, but selfplay training shows
+    // the learner a 2-step opponent wait (pick + place phases). Splitting
+    // Maia's move into 2 c_steps (no-op + commit) keeps the learner's LSTM in
+    // the training distribution. 0 = next c_step is the no-op phase, 1 = the
+    // commit phase.
+    int maia_phase;
+} Chess;
+
+static inline Bitboard sq_bb(Square s) {
+    return SquareBB[s];
+}
+
+static inline int popcount(Bitboard b) {
+    return __builtin_popcountll(b);
+}
+
+static inline Square lsb(Bitboard b) {
+    assert(b);
+    return __builtin_ctzll(b);
+}
+
+static inline Square pop_lsb(Bitboard* b) {
+    Square s = lsb(*b);
+    *b &= *b - 1;
+    return s;
+}
+
+static inline Bitboard shift_bb(int Direction, Bitboard b) {
+    return Direction == NORTH ? b << 8
+         : Direction == SOUTH ? b >> 8
+         : Direction == EAST ? (b & ~FileHBB) << 1
+         : Direction == WEST ? (b & ~FileABB) >> 1
+         : Direction == NORTH_EAST ? (b & ~FileHBB) << 9
+         : Direction == SOUTH_EAST ? (b & ~FileHBB) >> 7
+         : Direction == NORTH_WEST ? (b & ~FileABB) << 7
+         : Direction == SOUTH_WEST ? (b & ~FileABB) >> 9
+         : 0;
+}
+
+static inline Bitboard pawn_attacks_bb(ChessColor c, Square s) {
+    return PawnAttacks[c][s];
+}
+
+static inline Bitboard knight_attacks_bb(Square s) {
+    return KnightAttacks[s];
+}
+
+static inline Bitboard king_attacks_bb(Square s) {
+    return KingAttacks[s];
+}
+
+
+static inline Bitboard rook_attacks_bb(Square s, Bitboard occupied) {
+    occupied &= RookMasks[s];
+    return RookAttacks[s][(occupied * RookMagics[s]) >> RookShifts[s]];
+}
+
+static inline Bitboard bishop_attacks_bb(Square s, Bitboard occupied) {
+    occupied &= BishopMasks[s];
+    return BishopAttacks[s][(occupied * BishopMagics[s]) >> BishopShifts[s]];
+}
+
+static inline Bitboard queen_attacks_bb(Square s, Bitboard occupied) {
+    return rook_attacks_bb(s, occupied) | bishop_attacks_bb(s, occupied);
+}
+
+
+Bitboard SquareBB[65];
+Bitboard PawnAttacks[2][64];
+Bitboard KnightAttacks[64];
+Bitboard KingAttacks[64];
+Bitboard BetweenBB[64][64];
+Bitboard LineBB[64][64];
+Zobrist zob;
+
+static bool bitboards_initialized = false;
+
+static Bitboard index_to_occupancy(int index, Bitboard mask) {
+    Bitboard occ = 0;
+    int bits = popcount(mask);
+    for (int i = 0; i < bits; i++) {
+        Square sq = lsb(mask);
+        mask &= mask - 1;
+        if (index & (1 << i)) occ |= sq_bb(sq);
+    }
+    return occ;
+}
+static Bitboard compute_bishop_mask(Square s) {
+    Bitboard mask = 0;
+    int r = rank_of(s), f = file_of(s);
+    for (int rr = r + 1, ff = f + 1; rr < 7 && ff < 7; rr++, ff++) mask |= sq_bb(make_square(ff, rr));
+    for (int rr = r - 1, ff = f + 1; rr > 0 && ff < 7; rr--, ff++) mask |= sq_bb(make_square(ff, rr));
+    for (int rr = r - 1, ff = f - 1; rr > 0 && ff > 0; rr--, ff--) mask |= sq_bb(make_square(ff, rr));
+    for (int rr = r + 1, ff = f - 1; rr < 7 && ff > 0; rr++, ff--) mask |= sq_bb(make_square(ff, rr));
+    return mask;
+}
+
+static void init_bishop_magics(void) {
+    Bitboard* table_ptr = BishopTable;
+    
+    for (Square sq = 0; sq < 64; sq++) {
+        BishopMasks[sq] = compute_bishop_mask(sq);
+        BishopMagics[sq] = BISHOP_MAGICS[sq];
+        
+        int bits = popcount(BishopMasks[sq]);
+        BishopShifts[sq] = 64 - bits;
+        BishopAttacks[sq] = table_ptr;
+        
+        int num_entries = 1 << bits;
+        memset(table_ptr, 0, num_entries * sizeof(Bitboard));
+        
+        for (int i = 0; i < num_entries; i++) {
+            Bitboard occ = index_to_occupancy(i, BishopMasks[sq]);
+            
+            Bitboard attacks = 0;
+            int r = rank_of(sq), f = file_of(sq);
+            for (int rr = r + 1, ff = f + 1; rr < 8 && ff < 8; rr++, ff++) {
+                Square tsq = make_square(ff, rr);
+                attacks |= sq_bb(tsq);
+                if (occ & sq_bb(tsq)) break;
+            }
+            for (int rr = r - 1, ff = f + 1; rr >= 0 && ff < 8; rr--, ff++) {
+                Square tsq = make_square(ff, rr);
+                attacks |= sq_bb(tsq);
+                if (occ & sq_bb(tsq)) break;
+            }
+            for (int rr = r - 1, ff = f - 1; rr >= 0 && ff >= 0; rr--, ff--) {
+                Square tsq = make_square(ff, rr);
+                attacks |= sq_bb(tsq);
+                if (occ & sq_bb(tsq)) break;
+            }
+            for (int rr = r + 1, ff = f - 1; rr < 8 && ff >= 0; rr++, ff--) {
+                Square tsq = make_square(ff, rr);
+                attacks |= sq_bb(tsq);
+                if (occ & sq_bb(tsq)) break;
+            }
+            
+            uint64_t idx = (occ * BishopMagics[sq]) >> BishopShifts[sq];
+            table_ptr[idx] = attacks;
+        }
+        table_ptr += num_entries;
+    }
+}
+
+static Bitboard compute_rook_mask(Square s) {
+    Bitboard mask = 0;
+    int r = rank_of(s), f = file_of(s);
+    for (int rr = r + 1; rr < 7; rr++) mask |= sq_bb(make_square(f, rr));
+    for (int rr = r - 1; rr > 0; rr--) mask |= sq_bb(make_square(f, rr));
+    for (int ff = f + 1; ff < 7; ff++) mask |= sq_bb(make_square(ff, r));
+    for (int ff = f - 1; ff > 0; ff--) mask |= sq_bb(make_square(ff, r));
+    return mask;
+}
+
+static void init_rook_magics(void) {
+    Bitboard* table_ptr = RookTable;
+    
+    for (Square sq = 0; sq < 64; sq++) {
+        RookMasks[sq] = compute_rook_mask(sq);
+        RookMagics[sq] = ROOK_MAGICS[sq];
+        
+        int bits = popcount(RookMasks[sq]);
+        RookShifts[sq] = 64 - bits;
+        RookAttacks[sq] = table_ptr;
+        
+        int num_entries = 1 << bits;
+        memset(table_ptr, 0, num_entries * sizeof(Bitboard));
+        
+        for (int i = 0; i < num_entries; i++) {
+            Bitboard occ = index_to_occupancy(i, RookMasks[sq]);
+            
+            Bitboard attacks = 0;
+            int r = rank_of(sq), f = file_of(sq);
+            for (int rr = r + 1; rr < 8; rr++) {
+                Square tsq = make_square(f, rr);
+                attacks |= sq_bb(tsq);
+                if (occ & sq_bb(tsq)) break;
+            }
+            for (int rr = r - 1; rr >= 0; rr--) {
+                Square tsq = make_square(f, rr);
+                attacks |= sq_bb(tsq);
+                if (occ & sq_bb(tsq)) break;
+            }
+            for (int ff = f + 1; ff < 8; ff++) {
+                Square tsq = make_square(ff, r);
+                attacks |= sq_bb(tsq);
+                if (occ & sq_bb(tsq)) break;
+            }
+            for (int ff = f - 1; ff >= 0; ff--) {
+                Square tsq = make_square(ff, r);
+                attacks |= sq_bb(tsq);
+                if (occ & sq_bb(tsq)) break;
+            }
+            
+            uint64_t idx = (occ * RookMagics[sq]) >> RookShifts[sq];
+            table_ptr[idx] = attacks;
+        }
+        table_ptr += num_entries;
+    }
+}
+
+void init_bitboards(void) {
+    if (bitboards_initialized) return;
+    
+    for (int c = 0; c < 2; c++) {
+        for (int pt = PAWN; pt <= KING; pt++) {
+            for (int s = 0; s < 64; s++) {
+                zob.psq[make_piece(c, pt)][s] = prng_rand();
+            }
+        }
+    }
+    for (int f = 0; f < 8; f++) {
+        zob.enpassant[f] = prng_rand();
+    }
+    for (int cr = 0; cr < 16; cr++) {
+        zob.castling[cr] = prng_rand();
+    }
+    zob.side = prng_rand();
+    
+    for (int i = 0; i < 64; i++) {
+        SquareBB[i] = 1ULL << i;
+    }
+    SquareBB[64] = 0;
+    
+    for (int s = 0; s < 64; s++) {
+        Bitboard bb = sq_bb(s);
+        PawnAttacks[CHESS_WHITE][s] = shift_bb(NORTH_WEST, bb) | shift_bb(NORTH_EAST, bb);
+        PawnAttacks[CHESS_BLACK][s] = shift_bb(SOUTH_WEST, bb) | shift_bb(SOUTH_EAST, bb);
+    }
+    
+    int knight_dirs[] = {-17, -15, -10, -6, 6, 10, 15, 17};
+    for (int s = 0; s < 64; s++) {
+        Bitboard attack = 0;
+        int file = file_of(s);
+        int rank = rank_of(s);
+        
+        for (int i = 0; i < 8; i++) {
+            int to = s + knight_dirs[i];
+            if (to >= 0 && to < 64) {
+                int to_file = file_of(to);
+                int to_rank = rank_of(to);
+                if (abs(to_file - file) <= 2 && abs(to_rank - rank) <= 2) {
+                    attack |= sq_bb(to);
+                }
+            }
+        }
+        KnightAttacks[s] = attack;
+    }
+    
+    int king_dirs[] = {-9, -8, -7, -1, 1, 7, 8, 9};
+    for (int s = 0; s < 64; s++) {
+        Bitboard attack = 0;
+        int file = file_of(s);
+        
+        for (int i = 0; i < 8; i++) {
+            int to = s + king_dirs[i];
+            if (to >= 0 && to < 64) {
+                int to_file = file_of(to);
+                if (abs(to_file - file) <= 1) {
+                    attack |= sq_bb(to);
+                }
+            }
+        }
+        KingAttacks[s] = attack;
+    }
+    
+    for (int s1 = 0; s1 < 64; s1++) {
+        for (int s2 = 0; s2 < 64; s2++) {
+            BetweenBB[s1][s2] = 0;
+            LineBB[s1][s2] = 0;
+            
+            if (s1 == s2) continue;
+            
+            int f1 = file_of(s1), r1 = rank_of(s1);
+            int f2 = file_of(s2), r2 = rank_of(s2);
+            int df = f2 - f1, dr = r2 - r1;
+            
+            if (df == 0 || dr == 0 || abs(df) == abs(dr)) {
+                int step_f = df == 0 ? 0 : (df > 0 ? 1 : -1);
+                int step_r = dr == 0 ? 0 : (dr > 0 ? 1 : -1);
+                
+                // BetweenBB: squares strictly between s1 and s2
+                int f = f1 + step_f;
+                int r = r1 + step_r;
+                while (f != f2 || r != r2) {
+                    Square sq = make_square(f, r);
+                    BetweenBB[s1][s2] |= sq_bb(sq);
+                    f += step_f;
+                    r += step_r;
+                }
+
+                f = f1;
+                r = r1;
+                while (f - step_f >= 0 && f - step_f < 8 && r - step_r >= 0 && r - step_r < 8) {
+                    f -= step_f;
+                    r -= step_r;
+                }
+                while (f >= 0 && f < 8 && r >= 0 && r < 8) {
+                    LineBB[s1][s2] |= sq_bb(make_square(f, r));
+                    f += step_f;
+                    r += step_r;
+                }
+            }
+        }
+    }
+    init_bishop_magics(); 
+    init_rook_magics();
+    bitboards_initialized = true;
+}
+
+static void pos_set(Position* pos, const char* fen) {
+    memset(pos, 0, sizeof(Position));
+    
+    int rank = 7, file = 0;
+    const char* ptr = fen;
+    
+    while (*ptr && *ptr != ' ') {
+        char c = *ptr++;
+        
+        if (c == '/') {
+            rank--;
+            file = 0;
+        } else if (c >= '1' && c <= '8') {
+            file += c - '0';
+        } else {
+            Square sq = make_square(file, rank);
+            Piece pc = NO_PIECE;
+            int pt = 0, color = 0;
+            
+            switch (c) {
+                case 'P': pc = W_PAWN; pt = PAWN; color = CHESS_WHITE; break;
+                case 'N': pc = W_KNIGHT; pt = KNIGHT; color = CHESS_WHITE; break;
+                case 'B': pc = W_BISHOP; pt = BISHOP; color = CHESS_WHITE; break;
+                case 'R': pc = W_ROOK; pt = ROOK; color = CHESS_WHITE; break;
+                case 'Q': pc = W_QUEEN; pt = QUEEN; color = CHESS_WHITE; break;
+                case 'K': pc = W_KING; pt = KING; color = CHESS_WHITE; break;
+                case 'p': pc = B_PAWN; pt = PAWN; color = CHESS_BLACK; break;
+                case 'n': pc = B_KNIGHT; pt = KNIGHT; color = CHESS_BLACK; break;
+                case 'b': pc = B_BISHOP; pt = BISHOP; color = CHESS_BLACK; break;
+                case 'r': pc = B_ROOK; pt = ROOK; color = CHESS_BLACK; break;
+                case 'q': pc = B_QUEEN; pt = QUEEN; color = CHESS_BLACK; break;
+                case 'k': pc = B_KING; pt = KING; color = CHESS_BLACK; break;
+            }
+            
+            if (pc != NO_PIECE) {
+                pos->board[sq] = pc;
+                pos->byTypeBB[pt] |= sq_bb(sq);
+                pos->byColorBB[color] |= sq_bb(sq);
+                pos->byTypeBB[0] |= sq_bb(sq);
+                pos->pieceCount[pc]++;
+            }
+            file++;
+        }
+    }
+    
+    if (*ptr == ' ') ptr++;
+    
+    pos->sideToMove = (*ptr == 'w') ? CHESS_WHITE : CHESS_BLACK;
+    ptr += 2;
+    
+    pos->castlingRights = NO_CASTLING;
+    while (*ptr && *ptr != ' ') {
+        if (*ptr == 'K') pos->castlingRights |= WHITE_OO;
+        else if (*ptr == 'Q') pos->castlingRights |= WHITE_OOO;
+        else if (*ptr == 'k') pos->castlingRights |= BLACK_OO;
+        else if (*ptr == 'q') pos->castlingRights |= BLACK_OOO;
+        ptr++;
+    }
+    
+
+    if (*ptr == ' ') ptr++;
+    
+    pos->epSquare = SQ_NONE;
+    if (*ptr != '-') {
+        int ep_file = ptr[0] - 'a';
+        int ep_rank = ptr[1] - '1';
+        pos->epSquare = make_square(ep_file, ep_rank);
+    }
+    
+    pos->key = 0;
+    for (Square sq = SQ_A1; sq <= SQ_H8; sq++) {
+        Piece pc = pos->board[sq];
+        if (pc != NO_PIECE) {
+            pos->key ^= zob.psq[pc][sq];
+        }
+    }
+    if (pos->sideToMove == CHESS_BLACK) {
+        pos->key ^= zob.side;
+    }
+    if (pos->castlingRights) {
+        pos->key ^= zob.castling[pos->castlingRights];
+    }
+    if (pos->epSquare != SQ_NONE) {
+        pos->key ^= zob.enpassant[file_of(pos->epSquare)];
+    }
+}
+
+static void do_move(Position* pos, Move m, UndoInfo* undo_stack, int* undo_stack_ptr) {
+    if (m == MOVE_NULL) {
+        undo_stack[*undo_stack_ptr].captured = NO_PIECE;
+        undo_stack[*undo_stack_ptr].castlingRights = pos->castlingRights;
+        undo_stack[*undo_stack_ptr].epSquare = pos->epSquare;
+        undo_stack[*undo_stack_ptr].rule50 = pos->rule50;
+        undo_stack[*undo_stack_ptr].key = pos->key;
+        undo_stack[*undo_stack_ptr].pliesFromNull = 0;
+        (*undo_stack_ptr)++;
+        
+        if (pos->epSquare != SQ_NONE) {
+            pos->key ^= zob.enpassant[file_of(pos->epSquare)];
+            pos->epSquare = SQ_NONE;
+        }
+        pos->sideToMove = !pos->sideToMove;
+        pos->key ^= zob.side;
+        return;
+    }
+    
+    Square from = from_sq(m);
+    Square to = to_sq(m);
+    int move_type = type_of_m(m);
+    Piece pc = piece_on(pos, from);
+    Piece captured = piece_on(pos, to);
+    int pt = type_of_p(pc);
+    ChessColor us = pos->sideToMove;
+    ChessColor them = !us;
+    
+    undo_stack[*undo_stack_ptr].captured = captured;
+    undo_stack[*undo_stack_ptr].castlingRights = pos->castlingRights;
+    undo_stack[*undo_stack_ptr].epSquare = pos->epSquare;
+    undo_stack[*undo_stack_ptr].rule50 = pos->rule50;
+    undo_stack[*undo_stack_ptr].key = pos->key;
+    undo_stack[*undo_stack_ptr].pliesFromNull = (*undo_stack_ptr > 0) ? undo_stack[*undo_stack_ptr - 1].pliesFromNull + 1 : 0;
+    (*undo_stack_ptr)++;
+    
+    if (pt == PAWN || captured != NO_PIECE) {
+        pos->rule50 = 0;
+        undo_stack[*undo_stack_ptr - 1].pliesFromNull = 0;
+    }
+    else {
+        pos->rule50++;
+    }
+    
+    if (pos->epSquare != SQ_NONE) {
+        pos->key ^= zob.enpassant[file_of(pos->epSquare)];
+    }
+    pos->epSquare = SQ_NONE;
+    
+    switch (move_type) {
+        case CASTLING: {
+            pos->key ^= zob.psq[pc][from];
+            
+            pos->board[from] = NO_PIECE;
+            pos->board[to] = pc;
+            pos->byTypeBB[pt] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byColorBB[us] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byTypeBB[0] ^= sq_bb(from) ^ sq_bb(to);
+            pos->key ^= zob.psq[pc][to];
+            
+            Square rook_from, rook_to;
+            if (to > from) { 
+                rook_from = from + 3;
+                rook_to = from + 1;
+            } else {
+                rook_from = from - 4;
+                rook_to = from - 1;
+            }
+            
+            Piece rook = piece_on(pos, rook_from);
+            pos->key ^= zob.psq[rook][rook_from];
+            pos->board[rook_from] = NO_PIECE;
+            pos->board[rook_to] = rook;
+            pos->byTypeBB[ROOK] ^= sq_bb(rook_from) ^ sq_bb(rook_to);
+            pos->byColorBB[us] ^= sq_bb(rook_from) ^ sq_bb(rook_to);
+            pos->byTypeBB[0] ^= sq_bb(rook_from) ^ sq_bb(rook_to);
+            pos->key ^= zob.psq[rook][rook_to];
+            break;
+        }
+        case ENPASSANT: {
+            pos->key ^= zob.psq[pc][from];
+            
+            pos->board[from] = NO_PIECE;
+            pos->board[to] = pc;
+            pos->byTypeBB[pt] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byColorBB[us] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byTypeBB[0] ^= sq_bb(from) ^ sq_bb(to);
+            pos->key ^= zob.psq[pc][to];
+            
+            Square cap_sq = to + (us == CHESS_WHITE ? SOUTH : NORTH);
+            Piece cap_pawn = piece_on(pos, cap_sq);
+            pos->key ^= zob.psq[cap_pawn][cap_sq];
+            pos->board[cap_sq] = NO_PIECE;
+            pos->byTypeBB[PAWN] ^= sq_bb(cap_sq);
+            pos->byColorBB[them] ^= sq_bb(cap_sq);
+            pos->byTypeBB[0] ^= sq_bb(cap_sq);
+            pos->pieceCount[cap_pawn]--;
+            break;
+        }
+        case NORMAL:
+        case PROMOTION: {
+            pos->key ^= zob.psq[pc][from];
+            
+            if (captured != NO_PIECE) {
+                pos->key ^= zob.psq[captured][to];
+                int cap_pt = type_of_p(captured);
+                pos->byTypeBB[cap_pt] ^= sq_bb(to);
+                pos->byColorBB[them] ^= sq_bb(to);
+                pos->byTypeBB[0] ^= sq_bb(to);
+                pos->pieceCount[captured]--;
+            }
+            
+            pos->board[from] = NO_PIECE;
+            pos->board[to] = pc;
+            pos->byTypeBB[pt] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byColorBB[us] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byTypeBB[0] ^= sq_bb(from) ^ sq_bb(to);
+            pos->key ^= zob.psq[pc][to];
+            
+            if (move_type == PROMOTION) {
+                int promo_pt = promotion_type(m);
+                Piece promo_pc = make_piece(us, promo_pt);
+                pos->key ^= zob.psq[pc][to];
+                pos->board[to] = promo_pc;
+                pos->byTypeBB[pt] ^= sq_bb(to);
+                pos->byTypeBB[promo_pt] ^= sq_bb(to);
+                pos->pieceCount[pc]--;
+                pos->pieceCount[promo_pc]++;
+                pos->key ^= zob.psq[promo_pc][to];
+            }
+            
+            if (pt == PAWN) {
+                int diff = to - from;
+                if (diff == 16 || diff == -16) {
+                    Square ep_sq = (from + to) / 2;
+                    if (pawn_attacks_bb(us, ep_sq) & pieces_cp(pos, them, PAWN)) {
+                        pos->epSquare = ep_sq;
+                        pos->key ^= zob.enpassant[file_of(ep_sq)];
+                    }
+                }
+            }
+            break;
+        }
+        default:
+            break;
+    }
+    
+    uint8_t old_castling = pos->castlingRights;
+    if (pt == KING) {
+        pos->castlingRights &= us == CHESS_WHITE ? ~WHITE_CASTLING : ~BLACK_CASTLING;
+    }
+    if (from == SQ_A1 || to == SQ_A1) pos->castlingRights &= ~WHITE_OOO;
+    if (from == SQ_H1 || to == SQ_H1) pos->castlingRights &= ~WHITE_OO;
+    if (from == SQ_A8 || to == SQ_A8) pos->castlingRights &= ~BLACK_OOO;
+    if (from == SQ_H8 || to == SQ_H8) pos->castlingRights &= ~BLACK_OO;
+    
+    if (old_castling != pos->castlingRights) {
+        pos->key ^= zob.castling[old_castling];
+        pos->key ^= zob.castling[pos->castlingRights];
+    }
+    
+    pos->sideToMove = them;
+    pos->key ^= zob.side;
+}
+
+static void undo_move(Position* pos, Move m, UndoInfo* undo_stack, int* undo_stack_ptr) {
+    (*undo_stack_ptr)--;
+    UndoInfo* undo = &undo_stack[*undo_stack_ptr];
+    
+    if (m == MOVE_NULL) {
+        pos->castlingRights = undo->castlingRights;
+        pos->epSquare = undo->epSquare;
+        pos->rule50 = undo->rule50;
+        pos->key = undo->key;
+        pos->sideToMove = !pos->sideToMove;
+        return;
+    }
+    
+    Square from = from_sq(m);
+    Square to = to_sq(m);
+    int move_type = type_of_m(m);
+    ChessColor us = !pos->sideToMove;
+    ChessColor them = pos->sideToMove;
+    
+    Piece pc = piece_on(pos, to);
+    int pt = type_of_p(pc);
+    
+    pos->castlingRights = undo->castlingRights;
+    pos->epSquare = undo->epSquare;
+    pos->rule50 = undo->rule50;
+    pos->key = undo->key;
+    pos->sideToMove = us;
+    
+    switch (move_type) {
+        case CASTLING: {
+            pos->board[to] = NO_PIECE;
+            pos->board[from] = pc;
+            pos->byTypeBB[pt] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byColorBB[us] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byTypeBB[0] ^= sq_bb(from) ^ sq_bb(to);
+            
+            Square rook_from, rook_to;
+            if (to > from) {
+                rook_from = from + 3;
+                rook_to = from + 1;
+            } else {
+                rook_from = from - 4;
+                rook_to = from - 1;
+            }
+            
+            Piece rook = piece_on(pos, rook_to);
+            pos->board[rook_to] = NO_PIECE;
+            pos->board[rook_from] = rook;
+            pos->byTypeBB[ROOK] ^= sq_bb(rook_from) ^ sq_bb(rook_to);
+            pos->byColorBB[us] ^= sq_bb(rook_from) ^ sq_bb(rook_to);
+            pos->byTypeBB[0] ^= sq_bb(rook_from) ^ sq_bb(rook_to);
+            break;
+        }
+        case ENPASSANT: {
+            pos->board[to] = NO_PIECE;
+            pos->board[from] = pc;
+            pos->byTypeBB[pt] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byColorBB[us] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byTypeBB[0] ^= sq_bb(from) ^ sq_bb(to);
+
+            Square cap_sq = to + (us == CHESS_WHITE ? SOUTH : NORTH);
+            Piece cap_pawn = make_piece(them, PAWN);
+            pos->board[cap_sq] = cap_pawn;
+            pos->byTypeBB[PAWN] ^= sq_bb(cap_sq);
+            pos->byColorBB[them] ^= sq_bb(cap_sq);
+            pos->byTypeBB[0] ^= sq_bb(cap_sq);
+            pos->pieceCount[cap_pawn]++;
+            break;
+        }
+        case NORMAL:
+        case PROMOTION: {
+            if (move_type == PROMOTION) {
+                int promo_pt = promotion_type(m);
+                Piece promo_pc = make_piece(us, promo_pt);
+                pc = make_piece(us, PAWN);
+                pt = PAWN;
+                pos->board[to] = NO_PIECE;
+                pos->byTypeBB[promo_pt] ^= sq_bb(to);
+                pos->byTypeBB[pt] ^= sq_bb(to);
+                pos->pieceCount[promo_pc]--;
+                pos->pieceCount[pc]++;
+            }
+            
+            pos->board[to] = undo->captured;
+            pos->board[from] = pc;
+            pos->byTypeBB[pt] ^= sq_bb(from) ^ sq_bb(to);
+            pos->byColorBB[us] ^= sq_bb(from) ^ sq_bb(to);
+            
+            if (undo->captured != NO_PIECE) {
+                int cap_pt = type_of_p(undo->captured);
+                pos->byTypeBB[cap_pt] ^= sq_bb(to);
+                pos->byColorBB[them] ^= sq_bb(to);
+                pos->byTypeBB[0] ^= sq_bb(from);
+                pos->pieceCount[undo->captured]++;
+            } else {
+                pos->byTypeBB[0] ^= sq_bb(from) ^ sq_bb(to);
+            }
+            break;
+        }
+        default:
+            break;
+    }
+}
+
+static inline void add_move(MoveList* ml, Move m) {
+    ml->moves[ml->count].move = m;
+    ml->count++;
+}
+
+static void generate_pawn_moves(Position* pos, MoveList* ml, ChessColor us) {
+    ChessColor them = !us;
+    int up = (us == CHESS_WHITE) ? NORTH : SOUTH;
+    Bitboard rank7 = (us == CHESS_WHITE) ? Rank7BB : Rank2BB;
+    Bitboard rank3 = (us == CHESS_WHITE) ? Rank3BB : Rank6BB;
+    
+    Bitboard pawns = pieces_cp(pos, us, PAWN);
+    Bitboard pawnsOn7 = pawns & rank7;
+    Bitboard pawnsNotOn7 = pawns & ~rank7;
+    
+    Bitboard enemies = pieces_c(pos, them);
+    Bitboard empty = ~pieces(pos);
+    
+    Bitboard b1 = shift_bb(up, pawnsNotOn7) & empty;
+    Bitboard b2 = shift_bb(up, b1 & rank3) & empty;
+    
+    while (b1) {
+        Square to = pop_lsb(&b1);
+        add_move(ml, make_move(to - up, to));
+    }
+    
+    while (b2) {
+        Square to = pop_lsb(&b2);
+        add_move(ml, make_move(to - up - up, to));
+    }
+    
+    if (pawnsOn7) {
+        Bitboard b3 = shift_bb(up, pawnsOn7) & empty;
+        while (b3) {
+            Square to = pop_lsb(&b3);
+            Square from = to - up;
+            add_move(ml, make_promotion(from, to, QUEEN));
+            add_move(ml, make_promotion(from, to, ROOK));
+            add_move(ml, make_promotion(from, to, BISHOP));
+            add_move(ml, make_promotion(from, to, KNIGHT));
+        }
+    }
+    
+    Bitboard b4 = shift_bb(up + WEST, pawnsNotOn7) & enemies;
+    Bitboard b5 = shift_bb(up + EAST, pawnsNotOn7) & enemies;
+    
+    while (b4) {
+        Square to = pop_lsb(&b4);
+        add_move(ml, make_move(to - up - WEST, to));
+    }
+    
+    while (b5) {
+        Square to = pop_lsb(&b5);
+        add_move(ml, make_move(to - up - EAST, to));
+    }
+    
+    if (pawnsOn7) {
+        Bitboard b6 = shift_bb(up + WEST, pawnsOn7) & enemies;
+        Bitboard b7 = shift_bb(up + EAST, pawnsOn7) & enemies;
+        
+        while (b6) {
+            Square to = pop_lsb(&b6);
+            Square from = to - up - WEST;
+            add_move(ml, make_promotion(from, to, QUEEN));
+            add_move(ml, make_promotion(from, to, ROOK));
+            add_move(ml, make_promotion(from, to, BISHOP));
+            add_move(ml, make_promotion(from, to, KNIGHT));
+        }
+        
+        while (b7) {
+            Square to = pop_lsb(&b7);
+            Square from = to - up - EAST;
+            add_move(ml, make_promotion(from, to, QUEEN));
+            add_move(ml, make_promotion(from, to, ROOK));
+            add_move(ml, make_promotion(from, to, BISHOP));
+            add_move(ml, make_promotion(from, to, KNIGHT));
+        }
+    }
+    
+    if (pos->epSquare != SQ_NONE) {
+        Bitboard ep_pawns = pawnsNotOn7 & pawn_attacks_bb(them, pos->epSquare);
+        while (ep_pawns) {
+            Square from = pop_lsb(&ep_pawns);
+            add_move(ml, make_enpassant(from, pos->epSquare));
+        }
+    }
+}
+
+static void generate_castling(Position* pos, MoveList* ml, ChessColor us) {
+    Bitboard occupied = pieces(pos);
+    
+    if (us == CHESS_WHITE) {
+        if (pos->castlingRights & WHITE_OO) {
+            if (!(occupied & (sq_bb(SQ_F1) | sq_bb(SQ_G1)))) {
+                add_move(ml, make_castling(SQ_E1, SQ_G1));
+            }
+        }
+        if (pos->castlingRights & WHITE_OOO) {
+            if (!(occupied & (sq_bb(SQ_D1) | sq_bb(SQ_C1) | sq_bb(SQ_B1)))) {
+                add_move(ml, make_castling(SQ_E1, SQ_C1));
+            }
+        }
+    } else {
+        if (pos->castlingRights & BLACK_OO) {
+            if (!(occupied & (sq_bb(SQ_F8) | sq_bb(SQ_G8)))) {
+                add_move(ml, make_castling(SQ_E8, SQ_G8));
+            }
+        }
+        if (pos->castlingRights & BLACK_OOO) {
+            if (!(occupied & (sq_bb(SQ_D8) | sq_bb(SQ_C8) | sq_bb(SQ_B8)))) {
+                add_move(ml, make_castling(SQ_E8, SQ_C8));
+            }
+        }
+    }
+}
+
+static Bitboard attackers_to_sq(Position* pos, Square sq, Bitboard occupied) {
+    return (pawn_attacks_bb(CHESS_WHITE, sq) & pieces_cp(pos, CHESS_BLACK, PAWN) & occupied)
+         | (pawn_attacks_bb(CHESS_BLACK, sq) & pieces_cp(pos, CHESS_WHITE, PAWN) & occupied)
+         | (knight_attacks_bb(sq) & pieces_p(pos, KNIGHT) & occupied)
+         | (king_attacks_bb(sq) & pieces_p(pos, KING) & occupied)
+         | (bishop_attacks_bb(sq, occupied) & (pieces_p(pos, BISHOP) | pieces_p(pos, QUEEN)))
+         | (rook_attacks_bb(sq, occupied) & (pieces_p(pos, ROOK) | pieces_p(pos, QUEEN)));
+}
+
+static bool is_check(Position* pos, ChessColor c) {
+    Bitboard king_bb = pieces_cp(pos, c, KING);
+    if (!king_bb) return false;
+    Square king_sq = lsb(king_bb);
+    return (attackers_to_sq(pos, king_sq, pieces(pos)) & pieces_c(pos, !c)) != 0;
+}
+
+static Bitboard compute_pinned(Position* pos, ChessColor c) {
+    Bitboard pinned = 0;
+    Bitboard our_pieces = pieces_c(pos, c);
+    Bitboard king_bb = pieces_cp(pos, c, KING);
+    if (!king_bb) return 0;
+    
+    Square ksq = lsb(king_bb);
+    ChessColor them = !c;
+    Bitboard occupied = pieces(pos);
+    
+    Bitboard diag_pinners = (pieces_cp(pos, them, BISHOP) | pieces_cp(pos, them, QUEEN)) 
+                          & bishop_attacks_bb(ksq, 0);
+    
+    while (diag_pinners) {
+        Square pinner_sq = pop_lsb(&diag_pinners);
+        Bitboard between = BetweenBB[ksq][pinner_sq] & occupied;
+        if (popcount(between) == 1) {
+            pinned |= between & our_pieces;
+        }
+    }
+    
+    Bitboard rook_pinners = (pieces_cp(pos, them, ROOK) | pieces_cp(pos, them, QUEEN)) 
+                          & rook_attacks_bb(ksq, 0);
+    
+    while (rook_pinners) {
+        Square pinner_sq = pop_lsb(&rook_pinners);
+        Bitboard between = BetweenBB[ksq][pinner_sq] & occupied;
+        if (popcount(between) == 1) {
+            pinned |= between & our_pieces;
+        }
+    }
+    
+    return pinned;
+}
+
+static inline bool is_legal_move_fast(Position* pos, Move m, Bitboard pinned, Square ksq, ChessColor us) {
+    Square from = from_sq(m);
+    Square to = to_sq(m);
+    int mt = type_of_m(m);
+    
+    if (from == ksq) {
+        if (mt == CASTLING) {
+            ChessColor them = !us;
+            if (is_check(pos, us)) return false;
+            Square mid = (from + to) / 2;
+            Bitboard occ = pieces(pos) ^ sq_bb(from);
+            if (attackers_to_sq(pos, mid, occ) & pieces_c(pos, them)) return false;
+            if (attackers_to_sq(pos, to, occ) & pieces_c(pos, them)) return false;
+            return true;
+        }
+        Bitboard occ = pieces(pos) ^ sq_bb(from);
+        return !(attackers_to_sq(pos, to, occ) & pieces_c(pos, !us));
+    }
+    
+    if (mt == ENPASSANT) {
+        Bitboard occ = pieces(pos) ^ sq_bb(from) ^ sq_bb(to);
+        Square capsq = to + (us == CHESS_WHITE ? -8 : 8);
+        occ ^= sq_bb(capsq);
+        return !(attackers_to_sq(pos, ksq, occ) & pieces_c(pos, !us));
+    }
+    
+    if (!(pinned & sq_bb(from))) {
+        return true;
+    }
+    
+    return LineBB[ksq][from] & sq_bb(to);
+}
+
+static inline bool is_legal_move(Position* pos, Move m) {
+    ChessColor us = pos->sideToMove;
+    ChessColor them = (ChessColor)!us;
+    int mt = type_of_m(m);
+    if (mt == CASTLING) {
+        if (is_check(pos, us)) return false;
+        Square from = from_sq(m), to = to_sq(m);
+        Square mid = (from + to) / 2;
+        Bitboard occ = pieces(pos);
+        if ((attackers_to_sq(pos, mid, occ) & pieces_c(pos, them))
+         || (attackers_to_sq(pos, to, occ) & pieces_c(pos, them))) return false;
+        return true;
+    }
+    if (mt == ENPASSANT) {
+        Bitboard king_bb = pieces_cp(pos, us, KING);
+        if (!king_bb) return false;
+        Square ksq = lsb(king_bb);
+        Square from = from_sq(m), to = to_sq(m);
+        Square capsq = (us == CHESS_WHITE) ? (to - 8) : (to + 8);
+        Bitboard occ = pieces(pos) ^ sq_bb(from) ^ sq_bb(capsq) ^ sq_bb(to);
+        return (attackers_to_sq(pos, ksq, occ) & pieces_c(pos, them)) == 0;
+    }
+    UndoInfo u[1]; int p = 0;
+    do_move(pos, m, u, &p);
+    bool ok = !is_check(pos, us);
+    undo_move(pos, m, u, &p);
+    return ok;
+}
+
+static inline void generate_pseudo_legal(Position* pos, MoveList* ml, ChessColor us) {
+    ml->count = 0;
+    generate_pawn_moves(pos, ml, us);
+
+    Bitboard occupied = pieces(pos);
+    Bitboard target = ~pieces_c(pos, us);
+    Bitboard bb = pieces_cp(pos, us, KNIGHT);
+    while (bb) {
+        Square from = pop_lsb(&bb);
+        Bitboard attacks = knight_attacks_bb(from) & target;
+        while (attacks) {
+            Square to = pop_lsb(&attacks);
+            add_move(ml, make_move(from, to));
+        }
+    }
+
+    bb = pieces_cp(pos, us, BISHOP);
+    while (bb) {
+        Square from = pop_lsb(&bb);
+        Bitboard attacks = bishop_attacks_bb(from, occupied) & target;
+        while (attacks) {
+            Square to = pop_lsb(&attacks);
+            add_move(ml, make_move(from, to));
+        }
+    }
+
+    bb = pieces_cp(pos, us, ROOK);
+    while (bb) {
+        Square from = pop_lsb(&bb);
+        Bitboard attacks = rook_attacks_bb(from, occupied) & target;
+        while (attacks) {
+            Square to = pop_lsb(&attacks);
+            add_move(ml, make_move(from, to));
+        }
+    }
+
+    bb = pieces_cp(pos, us, QUEEN);
+    while (bb) {
+        Square from = pop_lsb(&bb);
+        Bitboard attacks = queen_attacks_bb(from, occupied) & target;
+        while (attacks) {
+            Square to = pop_lsb(&attacks);
+            add_move(ml, make_move(from, to));
+        }
+    }
+
+    bb = pieces_cp(pos, us, KING);
+    while (bb) {
+        Square from = pop_lsb(&bb);
+        Bitboard attacks = king_attacks_bb(from) & target;
+        while (attacks) {
+            Square to = pop_lsb(&attacks);
+            add_move(ml, make_move(from, to));
+        }
+    }
+
+    generate_castling(pos, ml, us);
+}
+
+static void generate_legal(Position* pos, MoveList* ml, UndoInfo* undo_stack, int* undo_stack_ptr) {
+    generate_pseudo_legal(pos, ml, pos->sideToMove);
+    ChessColor us = pos->sideToMove;
+    ChessColor them = (ChessColor)!us;
+    Bitboard king_bb = pieces_cp(pos, us, KING);
+    Square ksq = king_bb ? lsb(king_bb) : SQ_NONE;
+    Bitboard pinned = compute_pinned(pos, us);
+    bool in_check = is_check(pos, us);
+    int check_count = 0;
+    Bitboard evasion_mask = 0;
+    if (in_check) {
+        Bitboard checkers = attackers_to_sq(pos, ksq, pieces(pos)) & pieces_c(pos, them);
+        check_count = popcount(checkers);
+        if (check_count == 1) {
+            Square checker_sq = lsb(checkers);
+            evasion_mask = sq_bb(checker_sq);
+            Piece checker = piece_on(pos, checker_sq);
+            if (checker != NO_PIECE) {
+                int checker_type = type_of_p(checker);
+                if (checker_type == BISHOP || checker_type == ROOK || checker_type == QUEEN) {
+                    evasion_mask |= BetweenBB[ksq][checker_sq];
+                }
+            }
+        }
+    }
+    
+    int write = 0;
+    for (int i = 0; i < ml->count; i++) {
+        Move m = ml->moves[i].move;
+        Square from = from_sq(m);
+        bool legal;
+        if (!in_check) {
+            legal = is_legal_move_fast(pos, m, pinned, ksq, us);
+        } else if (from == ksq) {
+            legal = is_legal_move_fast(pos, m, pinned, ksq, us);
+        } else if (check_count > 1) {
+            legal = false;
+        } else if (type_of_m(m) == ENPASSANT) {
+            legal = is_legal_move(pos, m);
+        } else if ((sq_bb(to_sq(m)) & evasion_mask) == 0) {
+            legal = false;
+        } else {
+            legal = is_legal_move_fast(pos, m, pinned, ksq, us);
+        }
+        if (legal) {
+            ml->moves[write++] = ml->moves[i];
+        }
+    }
+    ml->count = write;
+}
+
+static inline bool is_insufficient_material(const Position* pos) {
+    if (pieces_p(pos, PAWN) | pieces_p(pos, ROOK) | pieces_p(pos, QUEEN))
+        return false;
+
+    int wN = popcount(pieces_cp(pos, CHESS_WHITE, KNIGHT));
+    int bN = popcount(pieces_cp(pos, CHESS_BLACK, KNIGHT));
+    int wB = popcount(pieces_cp(pos, CHESS_WHITE, BISHOP));
+    int bB = popcount(pieces_cp(pos, CHESS_BLACK, BISHOP));
+    int totalMinors = wN + bN + wB + bB;
+
+    if (totalMinors == 0)
+        return true;
+
+    if (totalMinors == 1)
+        return true;
+
+    if (totalMinors == 2) {
+        if ((wN == 2 && wB == 0 && bN == 0 && bB == 0) || (bN == 2 && bB == 0 && wN == 0 && wB == 0))
+            return true;
+        if ((wN + wB) == 1 && (bN + bB) == 1)
+            return true;
+    }
+
+    return false;
+}
+
+static void clear_player_selection(Chess* env, int side) {
+    env->pick_phase[side] = 0;
+    env->selected_square[side] = SQ_NONE;
+    env->valid_destinations[side].count = 0;
+    env->valid_to_mask[side] = 0;
+}
+
+static void rebuild_legal_state(Chess* env) {
+    generate_legal(&env->pos, &env->legal_moves, env->undo_stack, &env->undo_stack_ptr);
+    env->legal_dirty = 0;
+    env->valid_from_mask[0] = 0;
+    env->valid_from_mask[1] = 0;
+    env->valid_to_mask[0] = 0;
+    env->valid_to_mask[1] = 0;
+    int side = (int)env->pos.sideToMove;
+    Bitboard from_mask = 0;
+    for (int i = 0; i < env->legal_moves.count; i++) {
+        from_mask |= sq_bb(from_sq(env->legal_moves.moves[i].move));
+    }
+    env->valid_from_mask[side] = from_mask;
+}
+void populate_observations(Chess* env) {
+    Position* pos = &env->pos;
+
+    int num_players = env->mode == CHESS_MODE_SELFPLAY ? 2 : 1;
+    for (int player_iter = 0; player_iter < num_players; player_iter++) {
+        int player = env->mode == CHESS_MODE_SELFPLAY ? player_iter : env->learner_color;
+        // Selfplay: slot ↔ color mapping is randomized per env (slot_for_color).
+        // Single-agent modes: only the learner has a slot (idx 0).
+        int buffer_idx = (env->mode == CHESS_MODE_SELFPLAY) ? env->slot_for_color[player] : 0;
+        uint8_t* player_obs = env->obs_ptr[buffer_idx];
+        memset(player_obs, 0, OBS_SIZE);
+        uint8_t* board_planes = player_obs + O_BOARD;
+
+        // Selfplay: each iteration writes into its own per-slot mask.
+        // Single-agent modes: only the learner iter writes into the (single) mask.
+        unsigned char* my_mask = NULL;
+        bool fill_mask = false;
+        if (env->action_mask != NULL) {
+            if (env->mode == CHESS_MODE_SELFPLAY) {
+                my_mask = env->action_mask_ptr[buffer_idx];
+                fill_mask = true;
+            } else if (player == env->learner_color) {
+                my_mask = env->action_mask_ptr[0];
+                fill_mask = true;
+            }
+        }
+        if (fill_mask) {
+            memset(my_mask, 0, NUM_ACTIONS * sizeof(unsigned char));
+        }
+
+        ChessColor us = (ChessColor)player;  // 0=White, 1=Black
+        ChessColor them = (ChessColor)!us;
+
+        int flip = player * 56;
+
+
+        // Compact ego-centric board: one byte per square. 0 = empty,
+        // own P..K = 1..6, enemy P..K = 7..12.
+        for (int pt = PAWN; pt <= KING; pt++) {
+            Bitboard bb = pieces_cp(pos, player, pt);
+            while (bb) {
+                Square sq = pop_lsb(&bb);
+                board_planes[sq ^ flip] = (uint8_t)pt;
+            }
+        }
+        
+        for (int pt = PAWN; pt <= KING; pt++) {
+            Bitboard bb = pieces_cp(pos, them, pt);
+            while (bb) {
+                Square sq = pop_lsb(&bb);
+                board_planes[sq ^ flip] = (uint8_t)(6 + pt);
+            }
+        }
+               
+        ChessColor side_to_move = pos->sideToMove;
+        
+        player_obs[O_SIDE] = (pos->sideToMove == us) ? 1 : 0;
+        
+        uint8_t castle_rights = pos->castlingRights;
+        if (player == 1) {
+            uint8_t flipped = 0;
+            if (castle_rights & BLACK_OO) flipped |= WHITE_OO;
+            if (castle_rights & BLACK_OOO) flipped |= WHITE_OOO;
+            if (castle_rights & WHITE_OO) flipped |= BLACK_OO;
+            if (castle_rights & WHITE_OOO) flipped |= BLACK_OOO;
+            castle_rights = flipped;
+        }
+        player_obs[O_CASTLE + 0] = (castle_rights & WHITE_OO) ? 1 : 0;
+        player_obs[O_CASTLE + 1] = (castle_rights & WHITE_OOO) ? 1 : 0;
+        player_obs[O_CASTLE + 2] = (castle_rights & BLACK_OO) ? 1 : 0;
+        player_obs[O_CASTLE + 3] = (castle_rights & BLACK_OOO) ? 1 : 0;
+
+        if (pos->epSquare < 64) {
+            int ep_sq = (player == 1) ? (pos->epSquare ^ 56) : pos->epSquare;
+            player_obs[O_EP + file_of((Square)ep_sq)] = 1;
+        } else {
+            player_obs[O_EP + 8] = 1;
+        }
+    
+        uint8_t* valid_from_indices = player_obs + O_VALID_FROM;
+        uint8_t* valid_to_indices = player_obs + O_VALID_TO;
+        for (int k = 0; k < CHESS_MAX_VALID_FROM; k++) valid_from_indices[k] = CHESS_NULL_SQ;
+        for (int k = 0; k < CHESS_MAX_VALID_TO; k++)   valid_to_indices[k]   = CHESS_NULL_SQ;
+        int valid_from_count = 0;
+        int valid_to_count = 0;
+
+        int player_idx = (int)us;
+
+        if (side_to_move == us) {
+            if (env->pick_phase[player_idx] == 0) {
+                Bitboard added = 0;
+                for (int i = 0; i < env->legal_moves.count; i++) {
+                    Square from = from_sq(env->legal_moves.moves[i].move);
+                    int view_from = (player == 1) ? (from ^ 56) : from;
+                    Bitboard bit = sq_bb((Square)view_from);
+                    if (!(added & bit)) {
+                        added |= bit;
+                        if (valid_from_count < CHESS_MAX_VALID_FROM) {
+                            valid_from_indices[valid_from_count++] = (uint8_t)view_from;
+                        }
+                        if (fill_mask) my_mask[view_from] = 1;
+                    }
+                }
+            } else {
+                Bitboard added = 0;
+                for (int i = 0; i < env->valid_destinations[player_idx].count; i++) {
+                    Square to = to_sq(env->valid_destinations[player_idx].moves[i].move);
+                    int view_to = (player == 1) ? (to ^ 56) : to;
+                    Bitboard bit = sq_bb((Square)view_to);
+                    if (!(added & bit)) {
+                        added |= bit;
+                        if (valid_to_count < CHESS_MAX_VALID_TO) {
+                            valid_to_indices[valid_to_count++] = (uint8_t)view_to;
+                        }
+                        if (fill_mask) my_mask[view_to] = 1;
+                    }
+                }
+            }
+        }
+        player_obs[O_VALID_FROM_COUNT] = (uint8_t)valid_from_count;
+        player_obs[O_VALID_TO_COUNT]   = (uint8_t)valid_to_count;
+        player_obs[O_PASS_VALID] = (side_to_move != us) ? 255 : 0;
+        if (fill_mask && side_to_move != us) {
+            my_mask[PASS_ACTION] = 1;
+        }
+
+        player_obs[O_PICK_PHASE] = env->pick_phase[player_idx] ? 1 : 0;
+
+        uint8_t selected_byte = (uint8_t)CHESS_NULL_SQ;
+        if (env->pick_phase[player_idx] == 1 && env->selected_square[player_idx] != SQ_NONE) {
+            int view_selected = (player == 1)
+                ? (env->selected_square[player_idx] ^ 56)
+                : env->selected_square[player_idx];
+            selected_byte = (uint8_t)view_selected;
+        }
+        player_obs[O_SELECTED_PIECE] = selected_byte;
+        
+        uint8_t* valid_promos = player_obs + O_VALID_PROMOS;
+        
+        if (env->pick_phase[player_idx] == 1 && env->valid_destinations[player_idx].count > 0) {
+            for (int i = 0; i < env->valid_destinations[player_idx].count; i++) {
+                Move m = env->valid_destinations[player_idx].moves[i].move;
+                if (type_of_m(m) == PROMOTION) {
+                    int type_idx = QUEEN - promotion_type(m);
+                    int file_idx = file_of(to_sq(m));
+                    valid_promos[type_idx * 8 + file_idx] = 1;
+                    if (fill_mask) my_mask[64 + type_idx * 8 + file_idx] = 1;
+                }
+            }
+        }
+        
+        player_obs[O_SELF_CHECK] = is_check(pos, us) ? 255 : 0;
+        player_obs[O_OPP_CHECK] = is_check(pos, them) ? 255 : 0;
+        
+        int rule50 = pos->rule50;
+        if (rule50 > 100) rule50 = 100;
+        player_obs[O_RULE50] = (uint8_t)((rule50 * 255) / 100);
+        
+        uint8_t rep_val = 0;
+        if (env->undo_stack_ptr >= 4) {
+            uint8_t plies = env->undo_stack[env->undo_stack_ptr - 1].pliesFromNull;
+            if (plies >= 4) {
+                int repetitions = 0;
+                for (int i = 4; i <= plies; i += 2) {
+                    int idx = env->undo_stack_ptr - i;
+                    if (idx >= 0 && env->undo_stack[idx].key == pos->key) {
+                        repetitions++;
+                    }
+                }
+                if (repetitions >= 2) {
+                    rep_val = 255;
+                } else if (repetitions == 1) {
+                    rep_val = 128;
+                }
+            }
+        }
+        player_obs[O_REPETITION] = rep_val;
+    }
+}
+
+static int move_to_san(Position* pos, Move m, char* buf, UndoInfo* undo_stack, int* undo_stack_ptr) {
+    const char files[] = "abcdefgh";
+    const char ranks[] = "12345678";
+    const char piece_chars[] = ".PNBRQK";
+    char* ptr = buf;
+    
+    Square from = from_sq(m);
+    Square to = to_sq(m);
+    int move_type = type_of_m(m);
+    Piece pc = piece_on(pos, from);
+    int pt = type_of_p(pc);
+    ChessColor us = pos->sideToMove;
+    
+    if (move_type == CASTLING) {
+        if (to > from) {
+            strcpy(ptr, "O-O");
+            ptr += 3;
+        } else {
+            strcpy(ptr, "O-O-O");
+            ptr += 5;
+        }
+    } else {
+        if (pt != PAWN) {
+            *ptr++ = piece_chars[pt];
+            
+            Bitboard same_pieces = pieces_cp(pos, us, pt) & ~sq_bb(from);
+            Bitboard attackers = 0;
+            
+            if (pt == KNIGHT) {
+                attackers = knight_attacks_bb(to) & same_pieces;
+            } else if (pt == BISHOP) {
+                attackers = bishop_attacks_bb(to, pieces(pos)) & same_pieces;
+            } else if (pt == ROOK) {
+                attackers = rook_attacks_bb(to, pieces(pos)) & same_pieces;
+            } else if (pt == QUEEN) {
+                attackers = (bishop_attacks_bb(to, pieces(pos)) | rook_attacks_bb(to, pieces(pos))) & same_pieces;
+            } else if (pt == KING) {
+                attackers = king_attacks_bb(to) & same_pieces;
+            }
+            
+            Bitboard legal_attackers = 0;
+            while (attackers) {
+                Square attacker_sq = pop_lsb(&attackers);
+                Move test_move = make_move(attacker_sq, to);
+                if (is_legal_move(pos, test_move)) {
+                    legal_attackers |= sq_bb(attacker_sq);
+                }
+            }
+            
+            if (legal_attackers) {
+                int same_file = 0, same_rank = 0;
+                Bitboard temp = legal_attackers;
+                while (temp) {
+                    Square s = pop_lsb(&temp);
+                    if (file_of(s) == file_of(from)) same_file++;
+                    if (rank_of(s) == rank_of(from)) same_rank++;
+                }
+                
+                if (same_file == 0) {
+                    *ptr++ = files[file_of(from)];
+                } else if (same_rank == 0) {
+                    *ptr++ = ranks[rank_of(from)];
+                } else {
+                    *ptr++ = files[file_of(from)];
+                    *ptr++ = ranks[rank_of(from)];
+                }
+            }
+        }
+        
+        Piece captured = piece_on(pos, to);
+        bool is_capture = (captured != NO_PIECE) || (move_type == ENPASSANT);
+        
+        if (is_capture) {
+            if (pt == PAWN) {
+                *ptr++ = files[file_of(from)];
+            }
+            *ptr++ = 'x';
+        }
+        
+        *ptr++ = files[file_of(to)];
+        *ptr++ = ranks[rank_of(to)];
+        
+        if (move_type == PROMOTION) {
+            *ptr++ = '=';
+            const char promo_pieces[] = "..NBRQ";
+            *ptr++ = promo_pieces[promotion_type(m)];
+        }
+    }
+    
+    do_move(pos, m, undo_stack, undo_stack_ptr);
+    
+    ChessColor them = pos->sideToMove;
+    if (is_check(pos, them)) {
+        MoveList ml;
+        generate_legal(pos, &ml, undo_stack, undo_stack_ptr);
+        if (ml.count == 0) {
+            *ptr++ = '#';
+        } else {
+            *ptr++ = '+';
+        }
+    }
+    
+    undo_move(pos, m, undo_stack, undo_stack_ptr);
+    
+    *ptr = '\0';
+    return ptr - buf;
+}
+
+static void export_pgn_append(Chess* env, const char* filename, int append) {
+    FILE* f = fopen(filename, append ? "a" : "w");
+    if (!f) return;
+    
+    if (env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM) {
+        const char* opponent_name = env->mode == CHESS_MODE_HUMAN ? "AI" : "Random";
+        const char* event_name = env->mode == CHESS_MODE_HUMAN ? "Human vs AI" : "Human vs Random";
+        fprintf(f, "[Event \"%s\"]\n", event_name);
+        fprintf(f, "[White \"%s\"]\n", env->human_color == CHESS_WHITE ? "Human" : opponent_name);
+        fprintf(f, "[Black \"%s\"]\n", env->human_color == CHESS_BLACK ? "Human" : opponent_name);
+    } else {
+        fprintf(f, "[Event \"Selfplay Eval Game %d\"]\n", env->pgn_game_number);
+        fprintf(f, "[White \"%s\"]\n", env->learner_color == CHESS_BLACK ? "Learner" : "Opponent");
+        fprintf(f, "[Black \"%s\"]\n", env->learner_color == CHESS_BLACK ? "Opponent" : "Learner");
+    }
+    fprintf(f, "[Site \"PufferLib\"]\n");
+    fprintf(f, "[Result \"%s\"]\n\n", env->last_result);
+    
+    Position replay_pos;
+    pos_set(&replay_pos, env->starting_fen);
+    
+    UndoInfo replay_undo[MAX_GAME_PLIES];
+    int replay_undo_ptr = 0;
+    
+    char san_buf[16];
+    
+    for (int i = 0; i < env->pgn_move_count; i++) {
+        if (i % 2 == 0) {
+            fprintf(f, "%d. ", i/2 + 1);
+        }
+        
+        Move m = env->pgn_moves[i];
+        move_to_san(&replay_pos, m, san_buf, replay_undo, &replay_undo_ptr);
+        fprintf(f, "%s ", san_buf);
+        
+        do_move(&replay_pos, m, replay_undo, &replay_undo_ptr);
+        
+        if ((i + 1) % 8 == 0) fprintf(f, "\n");
+    }
+    
+    if (strcmp(env->last_result, "White Wins") == 0) {
+        fprintf(f, "1-0");
+    } else if (strcmp(env->last_result, "Black Wins") == 0) {
+        fprintf(f, "0-1");
+    } else {
+        fprintf(f, "1/2-1/2");
+    }
+    
+    fprintf(f, "\n\n");
+    fclose(f);
+}
+
+static void generate_random_fen(Chess* env, char* fen_out) {
+    char board[64];
+    memset(board, '.', 64);
+
+    int wk_sq, bk_sq;
+    do {
+        wk_sq = rand_r(&env->rng) % 64;
+        bk_sq = rand_r(&env->rng) % 64;
+        int wk_rank = wk_sq / 8, wk_file = wk_sq % 8;
+        int bk_rank = bk_sq / 8, bk_file = bk_sq % 8;
+        int rank_diff = abs(wk_rank - bk_rank);
+        int file_diff = abs(wk_file - bk_file);
+        if (wk_sq != bk_sq && (rank_diff > 1 || file_diff > 1)) break;
+    } while (1);
+
+    board[wk_sq] = 'K';
+    board[bk_sq] = 'k';
+
+    const char* white_pieces = "QRRNNBBPP";
+    const char* black_pieces = "qrrnnbbpp";
+    int num_white = rand_r(&env->rng) % 16;
+    int num_black = rand_r(&env->rng) % 16;
+
+    for (int i = 0; i < num_white; i++) {
+        int sq, rank;
+        char piece;
+        do {
+            sq = rand_r(&env->rng) % 64;
+            rank = sq / 8;
+            piece = white_pieces[rand_r(&env->rng) % 9];
+        } while (board[sq] != '.' || (piece == 'P' && (rank == 0 || rank == 7)));
+        board[sq] = piece;
+    }
+
+    for (int i = 0; i < num_black; i++) {
+        int sq, rank;
+        char piece;
+        do {
+            sq = rand_r(&env->rng) % 64;
+            rank = sq / 8;
+            piece = black_pieces[rand_r(&env->rng) % 9];
+        } while (board[sq] != '.' || (piece == 'p' && (rank == 0 || rank == 7)));
+        board[sq] = piece;
+    }
+    
+    char* ptr = fen_out;
+    for (int rank = 7; rank >= 0; rank--) {
+        int empty = 0;
+        for (int file = 0; file < 8; file++) {
+            char piece = board[rank * 8 + file];
+            if (piece == '.') {
+                empty++;
+            } else {
+                if (empty > 0) {
+                    *ptr++ = '0' + empty;
+                    empty = 0;
+                }
+                *ptr++ = piece;
+            }
+        }
+        if (empty > 0) *ptr++ = '0' + empty;
+        if (rank > 0) *ptr++ = '/';
+    }
+    strcpy(ptr, " w - - 0 1");
+}
+
+static inline int apply_move_to_env(Chess* env, Move chosen, int* is_timeout) {
+    env->chess_moves++;
+    env->last_move_from = from_sq(chosen);
+    env->last_move_to = to_sq(chosen);
+    env->has_last_move_highlight = 1;
+
+    if ((env->mode == CHESS_MODE_HUMAN
+            || env->mode == CHESS_MODE_HUMAN_RANDOM
+            || env->log_pgn) && env->pgn_move_count < MAX_GAME_PLIES) {
+        env->pgn_moves[env->pgn_move_count++] = chosen;
+    }
+    
+    ChessColor side_before = env->pos.sideToMove;
+    do_move(&env->pos, chosen, env->undo_stack, &env->undo_stack_ptr);
+    env->legal_dirty = 1;
+    clear_player_selection(env, (int)env->pos.sideToMove);
+    
+    if (env->undo_stack_ptr > 0) {
+        Piece cap = env->undo_stack[env->undo_stack_ptr - 1].captured;
+        if (cap != NO_PIECE) {
+            int pt = type_of_p(cap) - 1;
+            if (pt >= 0 && pt < 6) {
+                if (color_of(cap) == CHESS_WHITE) env->white_captured[pt]++;
+                else env->black_captured[pt]++;
+            }
+        } else if ((int)type_of_m(chosen) == ENPASSANT) {
+            Piece cap_pawn = (side_before == CHESS_WHITE) ? B_PAWN : W_PAWN;
+            int pt = type_of_p(cap_pawn) - 1;
+            if (pt >= 0 && pt < 6) {
+                if (color_of(cap_pawn) == CHESS_WHITE) env->white_captured[pt]++;
+                else env->black_captured[pt]++;
+            }
+        }
+        if (env->undo_stack[env->undo_stack_ptr - 1].pliesFromNull > 99) {
+            env->undo_stack[env->undo_stack_ptr - 1].pliesFromNull = 99;
+        }
+    }
+
+    env->repetition_matches = 0;
+    if (env->undo_stack_ptr >= 4) {
+        int max_back = env->undo_stack[env->undo_stack_ptr - 1].pliesFromNull;
+        if (max_back > env->undo_stack_ptr) max_back = env->undo_stack_ptr;
+        for (int i = 4; i <= max_back; i += 2) {
+            if (env->undo_stack[env->undo_stack_ptr - i].key != env->pos.key) continue;
+            env->repetition_matches++;
+            if (env->repetition_matches == 2) break;
+        }
+    }
+    
+    rebuild_legal_state(env);
+
+    int game_result = 0;
+    *is_timeout = 0;
+    if (env->chess_moves >= env->max_moves || env->undo_stack_ptr >= MAX_GAME_PLIES - 2) {
+        *is_timeout = 1;
+        game_result = 3;
+    } else if (env->legal_moves.count == 0) {
+        if (is_check(&env->pos, env->pos.sideToMove)) {
+            game_result = env->pos.sideToMove == CHESS_WHITE ? 1 : 2;
+        } else {
+            game_result = 3;
+        }
+    } else if (is_insufficient_material(&env->pos)) {
+        game_result = 3;
+    } else if (env->enable_50_move_rule && env->pos.rule50 >= 100) {
+        game_result = 3;
+    } else if (env->enable_threefold_repetition && env->repetition_matches >= 2) {
+        game_result = 3;
+    }
+
+    return game_result;
+}
+
+// ---- CHESS_MODE_MAIA: external lc0/Maia UCI engine ----
+// Each env owns one lc0 child process. Communication is line-based UCI:
+//   parent → child: "position fen <FEN>\ngo nodes <N>\n"
+//   child  → parent: ... "info ..." ... "bestmove <uci>\n"
+// Configured via env vars: MAIA_LC0_PATH, MAIA_WEIGHTS_PATH, MAIA_NODES,
+// MAIA_BACKEND. MAIA_NODES=1 ≈ weakest Maia setting; raise for stronger play.
+
+#include <unistd.h>
+#include <sys/wait.h>
+#include <fcntl.h>
+#include <errno.h>
+#include <signal.h>
+
+static void position_to_fen(const Position* pos, char* out) {
+    char* p = out;
+    // Indexed by Piece enum: W_* are 1..6, B_* are 9..14. Gap at 7..8.
+    static const char pchars[16] = ".PNBRQK??pnbrqk?";
+    for (int rank = 7; rank >= 0; rank--) {
+        int empty = 0;
+        for (int file = 0; file < 8; file++) {
+            Piece pc = piece_on(pos, make_square(file, rank));
+            if (pc == NO_PIECE) {
+                empty++;
+            } else {
+                if (empty > 0) { *p++ = '0' + empty; empty = 0; }
+                *p++ = pchars[pc];
+            }
+        }
+        if (empty > 0) *p++ = '0' + empty;
+        if (rank > 0) *p++ = '/';
+    }
+    *p++ = ' ';
+    *p++ = (pos->sideToMove == CHESS_WHITE) ? 'w' : 'b';
+    *p++ = ' ';
+    int wrote_castle = 0;
+    if (pos->castlingRights & 1) { *p++ = 'K'; wrote_castle = 1; }
+    if (pos->castlingRights & 2) { *p++ = 'Q'; wrote_castle = 1; }
+    if (pos->castlingRights & 4) { *p++ = 'k'; wrote_castle = 1; }
+    if (pos->castlingRights & 8) { *p++ = 'q'; wrote_castle = 1; }
+    if (!wrote_castle) *p++ = '-';
+    *p++ = ' ';
+    if (pos->epSquare != SQ_NONE && (int)pos->epSquare < 64) {
+        *p++ = 'a' + (int)file_of(pos->epSquare);
+        *p++ = '1' + (int)rank_of(pos->epSquare);
+    } else {
+        *p++ = '-';
+    }
+    *p++ = ' ';
+    // Fullmove number isn't tracked on Position; hardcode 1. Halfmove (rule50)
+    // matters for the 50-move rule and is passed through.
+    p += sprintf(p, "%d 1", pos->rule50);
+    *p = '\0';
+}
+
+// Parse a UCI move like "e2e4" or "g7g8q" against the env's current legal list.
+// Returns MOVE_NONE if not found.
+static Move uci_to_move(const char* uci, const MoveList* legal) {
+    if (uci[0] < 'a' || uci[0] > 'h' || uci[2] < 'a' || uci[2] > 'h') return MOVE_NONE;
+    int from_file = uci[0] - 'a';
+    int from_rank = uci[1] - '1';
+    int to_file   = uci[2] - 'a';
+    int to_rank   = uci[3] - '1';
+    if (from_rank < 0 || from_rank > 7 || to_rank < 0 || to_rank > 7) return MOVE_NONE;
+    Square from = make_square(from_file, from_rank);
+    Square to   = make_square(to_file,   to_rank);
+    int promo_pt = -1;
+    if (uci[4] == 'q') promo_pt = QUEEN;
+    else if (uci[4] == 'r') promo_pt = ROOK;
+    else if (uci[4] == 'b') promo_pt = BISHOP;
+    else if (uci[4] == 'n') promo_pt = KNIGHT;
+    for (int i = 0; i < legal->count; i++) {
+        Move m = legal->moves[i].move;
+        if (from_sq(m) != from || to_sq(m) != to) continue;
+        if (promo_pt >= 0) {
+            if (type_of_m(m) != PROMOTION) continue;
+            if ((int)promotion_type(m) != promo_pt) continue;
+        }
+        return m;
+    }
+    return MOVE_NONE;
+}
+
+static int maia_write_all(int fd, const char* buf, int len) {
+    int n = 0;
+    while (n < len) {
+        int w = (int)write(fd, buf + n, (size_t)(len - n));
+        if (w < 0) {
+            if (errno == EINTR) continue;
+            return -1;
+        }
+        n += w;
+    }
+    return 0;
+}
+
+// Read one line (up to '\n' or buf-1 chars). Returns # bytes (excluding NUL) or
+// -1 on error / EOF. Blocks until a line is available.
+static int maia_read_line(int fd, char* buf, int bufsz) {
+    int n = 0;
+    while (n < bufsz - 1) {
+        char c;
+        int r = (int)read(fd, &c, 1);
+        if (r == 0) return -1;  // EOF
+        if (r < 0) {
+            if (errno == EINTR) continue;
+            return -1;
+        }
+        if (c == '\n') break;
+        buf[n++] = c;
+    }
+    buf[n] = '\0';
+    return n;
+}
+
+static void maia_close(Chess* env);
+
+static void maia_init(Chess* env) {
+    if (env->maia_pid > 0) return;  // already spawned
+
+    const char* lc0_path     = getenv("MAIA_LC0_PATH");
+    const char* weights_path = getenv("MAIA_WEIGHTS_PATH");
+    const char* backend_arg  = getenv("MAIA_BACKEND");
+    if (lc0_path == NULL)     lc0_path     = "./lc0";
+    if (weights_path == NULL) weights_path = "lc0/maia-1100.pb.gz";
+
+    int in_pipe[2], out_pipe[2];
+    if (pipe(in_pipe) < 0 || pipe(out_pipe) < 0) {
+        fprintf(stderr, "maia_init: pipe() failed\n");
+        env->maia_pid = -1;
+        return;
+    }
+    pid_t pid = fork();
+    if (pid < 0) {
+        close(in_pipe[0]); close(in_pipe[1]);
+        close(out_pipe[0]); close(out_pipe[1]);
+        fprintf(stderr, "maia_init: fork() failed\n");
+        env->maia_pid = -1;
+        return;
+    }
+    if (pid == 0) {
+        // Child: stdin←in_pipe[0], stdout→out_pipe[1].
+        dup2(in_pipe[0], STDIN_FILENO);
+        dup2(out_pipe[1], STDOUT_FILENO);
+        // Redirect stderr to /dev/null so info logs don't spam the parent.
+        int devnull = open("/dev/null", O_WRONLY);
+        if (devnull >= 0) { dup2(devnull, STDERR_FILENO); close(devnull); }
+        close(in_pipe[0]);  close(in_pipe[1]);
+        close(out_pipe[0]); close(out_pipe[1]);
+        char weights_arg[512];
+        snprintf(weights_arg, sizeof(weights_arg), "--weights=%s", weights_path);
+        if (backend_arg) {
+            char backend_buf[128];
+            snprintf(backend_buf, sizeof(backend_buf), "--backend=%s", backend_arg);
+            execlp(lc0_path, "lc0", weights_arg, backend_buf, (char*)NULL);
+        } else {
+            execlp(lc0_path, "lc0", weights_arg, (char*)NULL);
+        }
+        _exit(127);
+    }
+    // Parent.
+    close(in_pipe[0]);  // we write to in_pipe[1]
+    close(out_pipe[1]); // we read from out_pipe[0]
+    env->maia_pid       = (int)pid;
+    env->maia_stdin_fd  = in_pipe[1];
+    env->maia_stdout_fd = out_pipe[0];
+
+    // UCI handshake: send "uci", drain until "uciok"; then "isready", drain
+    // until "readyok". Engine init also loads weights so this can take a few
+    // seconds the first time.
+    char line[1024];
+    if (maia_write_all(env->maia_stdin_fd, "uci\n", 4) < 0) { maia_close(env); return; }
+    while (maia_read_line(env->maia_stdout_fd, line, sizeof(line)) >= 0) {
+        if (strncmp(line, "uciok", 5) == 0) break;
+    }
+    if (maia_write_all(env->maia_stdin_fd, "isready\n", 8) < 0) { maia_close(env); return; }
+    while (maia_read_line(env->maia_stdout_fd, line, sizeof(line)) >= 0) {
+        if (strncmp(line, "readyok", 7) == 0) break;
+    }
+}
+
+static void maia_close(Chess* env) {
+    if (env->maia_pid > 0) {
+        if (env->maia_stdin_fd >= 0) {
+            (void)maia_write_all(env->maia_stdin_fd, "quit\n", 5);
+            close(env->maia_stdin_fd);
+            env->maia_stdin_fd = -1;
+        }
+        if (env->maia_stdout_fd >= 0) {
+            close(env->maia_stdout_fd);
+            env->maia_stdout_fd = -1;
+        }
+        int status;
+        if (waitpid((pid_t)env->maia_pid, &status, WNOHANG) == 0) {
+            kill((pid_t)env->maia_pid, SIGTERM);
+            waitpid((pid_t)env->maia_pid, &status, 0);
+        }
+    }
+    env->maia_pid = 0;
+}
+
+// Ask Maia for the best move at the current env position. Returns a Move that
+// is guaranteed to be in env->legal_moves (or MOVE_NONE on engine failure).
+static Move maia_get_move(Chess* env) {
+    if (env->maia_pid <= 0) {
+        maia_init(env);
+        if (env->maia_pid <= 0) return MOVE_NONE;
+    }
+
+    int nodes = 1;
+    const char* nodes_str = getenv("MAIA_NODES");
+    if (nodes_str) nodes = atoi(nodes_str);
+    if (nodes < 1) nodes = 1;
+
+    char fen[128];
+    position_to_fen(&env->pos, fen);
+    char cmd[256];
+    int n = snprintf(cmd, sizeof(cmd), "position fen %s\ngo nodes %d\n", fen, nodes);
+    if (maia_write_all(env->maia_stdin_fd, cmd, n) < 0) {
+        maia_close(env);
+        return MOVE_NONE;
+    }
+
+    char line[1024];
+    while (1) {
+        int len = maia_read_line(env->maia_stdout_fd, line, sizeof(line));
+        if (len < 0) { maia_close(env); return MOVE_NONE; }
+        if (strncmp(line, "bestmove ", 9) != 0) continue;
+        char uci[8] = {0};
+        int j = 9, k = 0;
+        while (line[j] && line[j] != ' ' && line[j] != '\n' && line[j] != '\r' && k < 7) {
+            uci[k++] = line[j++];
+        }
+        return uci_to_move(uci, &env->legal_moves);
+    }
+}
+
+void c_reset(Chess* env) {
+    env->tick = 0;
+    env->chess_moves = 0;
+    env->game_result = 0;
+    env->undo_stack_ptr = 0;
+    env->repetition_matches = 0;
+    env->invalid_actions_this_episode = 0;
+    env->episode_reward = 0.0f;
+    env->pgn_move_count = 0;
+    env->show_game_end_popup = 0;
+    env->has_last_move_highlight = 0;
+    clear_player_selection(env, 0);
+    clear_player_selection(env, 1);
+    env->valid_from_mask[0] = 0;
+    env->valid_from_mask[1] = 0;
+    
+    memset(env->white_captured, 0, sizeof(env->white_captured));
+    memset(env->black_captured, 0, sizeof(env->black_captured));
+    
+    if (env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM) {
+        env->human_color = -1;
+    } else if (env->mode != CHESS_MODE_SELFPLAY) {
+        env->learner_color = 1 - env->learner_color;
+    }
+    env->maia_phase = 0;
+    
+    if (env->fen_curriculum != NULL && env->num_fens > 0) {
+        float randvalue = (float)rand_r(&env->rng) / (float)(RAND_MAX);
+        if(env->fen_curric_pct >= randvalue){
+            int idx = rand_r(&env->rng) % env->num_fens;
+            pos_set(&env->pos, env->fen_curriculum[idx]);
+        }
+        else {
+            pos_set(&env->pos, env->starting_fen);
+        }
+
+    } else if (env->random_fen) {
+        char fen_buf[128];
+        generate_random_fen(env, fen_buf);
+        pos_set(&env->pos, fen_buf);
+    } else {
+        pos_set(&env->pos, env->starting_fen);
+    }
+    
+    rebuild_legal_state(env);
+    populate_observations(env);
+
+}
+
+void c_step(Chess* env) {
+    if (env->render_paused && env->client != NULL) {
+        return;
+    }
+    if ((env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+            && env->human_color == -1) {
+        return;
+    }
+    
+    if (env->mode == CHESS_MODE_SELFPLAY && !env->log_pgn_choice_made) {
+        if (env->client != NULL) {
+            return;
+        }
+        env->log_pgn = 0;
+        env->log_pgn_choice_made = 1;
+    }
+
+    if ((env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+            && env->show_game_end_popup) {
+        *env->reward_ptr[0] = 0.0f;
+        *env->terminal_ptr[0] = 0;
+        return;
+    }
+
+    if (env->legal_dirty) {
+        rebuild_legal_state(env);
+    }
+    
+    env->tick++;
+    int move_completed = 0;
+    ChessColor mover = env->pos.sideToMove;
+    int mover_idx = (int)mover;
+    int game_result = 0;
+    int is_timeout = 0;
+
+    if ((env->mode == CHESS_MODE_RANDOM && env->pos.sideToMove != env->learner_color)
+            || (env->mode == CHESS_MODE_HUMAN_RANDOM && env->pos.sideToMove != env->human_color)) {
+        if (env->legal_moves.count > 0) {
+            int idx = rand_r(&env->rng) % env->legal_moves.count;
+            clear_player_selection(env, mover_idx);
+            game_result = apply_move_to_env(env, env->legal_moves.moves[idx].move, &is_timeout);
+            move_completed = 1;
+        }
+    } else if (env->mode == CHESS_MODE_MAIA && env->pos.sideToMove != env->learner_color) {
+        if (env->legal_moves.count > 0) {
+            if (env->maia_phase == 0) {
+                // First c_step of Maia's "move": no-op so the learner's LSTM
+                // sees a 2-step opponent wait (matches selfplay's pick+place
+                // cadence the policy was trained on).
+                env->maia_phase = 1;
+                move_completed = 1;
+            } else {
+                Move maia_mv = maia_get_move(env);
+                if (maia_mv == MOVE_NONE) {
+                    // Engine failure / unparseable bestmove: fall back to random
+                    // so the trial completes. Logged via log.maia_failures so
+                    // Python can flag a degraded eval.
+                    env->log.maia_failures += 1.0f;
+                    int idx = rand_r(&env->rng) % env->legal_moves.count;
+                    maia_mv = env->legal_moves.moves[idx].move;
+                }
+                clear_player_selection(env, mover_idx);
+                game_result = apply_move_to_env(env, maia_mv, &is_timeout);
+                env->maia_phase = 0;
+                move_completed = 1;
+            }
+        }
+    } else {
+        // Selfplay: side-to-move's action lives in whichever slot plays that color.
+        int action = (env->mode == CHESS_MODE_SELFPLAY)
+            ? *env->action_ptr[env->slot_for_color[mover_idx]]
+            : *env->action_ptr[0];
+        if ((env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+                && env->pos.sideToMove == env->human_color) {
+            action = -1;
+            *env->action_ptr[0] = -1;
+        }
+
+        mover = env->pos.sideToMove;
+        mover_idx = (int)mover;
+
+        // In selfplay both players train, so charge whichever slot moved.
+        // In single-agent modes, only the learner's slot exists.
+        int penalty_slot = (env->mode == CHESS_MODE_SELFPLAY) ? env->slot_for_color[mover_idx]
+                         : (mover == env->learner_color) ? 0 : -1;
+
+        if (env->legal_moves.count == 0) {
+            clear_player_selection(env, mover_idx);
+        } else if (action < 0 || action >= PASS_ACTION) {
+            if (penalty_slot >= 0) {
+                *env->reward_ptr[penalty_slot] += (env->pick_phase[mover_idx] == 0)
+                    ? env->reward_invalid_piece : env->reward_invalid_move;
+                env->invalid_actions_this_episode++;
+            }
+            if (env->pick_phase[mover_idx] == 1) {
+                clear_player_selection(env, mover_idx);
+            }
+        } else {
+            bool is_promo = (action >= 64 && action < 96);
+
+            if (env->pick_phase[mover_idx] == 0) {
+                clear_player_selection(env, mover_idx);
+
+                bool valid_pick = !is_promo;
+                Square picked_sq = SQ_NONE;
+                if (valid_pick) {
+                    picked_sq = (mover == CHESS_BLACK) ? (Square)(action ^ 56) : (Square)action;
+                    Piece pc = piece_on(&env->pos, picked_sq);
+                    valid_pick = (pc != NO_PIECE && color_of(pc) == mover);
+                }
+
+                if (valid_pick) {
+                    MoveList* dests = &env->valid_destinations[mover_idx];
+                    dests->count = 0;
+                    Bitboard to_mask = 0;
+                    for (int i = 0; i < env->legal_moves.count; i++) {
+                        if (from_sq(env->legal_moves.moves[i].move) == picked_sq) {
+                            dests->moves[dests->count++] = env->legal_moves.moves[i];
+                            to_mask |= sq_bb(to_sq(env->legal_moves.moves[i].move));
+                        }
+                    }
+
+                    if (dests->count > 0) {
+                        env->selected_square[mover_idx] = picked_sq;
+                        env->pick_phase[mover_idx] = 1;
+                        env->valid_to_mask[mover_idx] = to_mask;
+                    } else {
+                        valid_pick = false;
+                        clear_player_selection(env, mover_idx);
+                    }
+                }
+
+                if (!valid_pick && penalty_slot >= 0) {
+                    *env->reward_ptr[penalty_slot] += env->reward_invalid_piece;
+                    env->invalid_actions_this_episode++;
+                }
+            } else {
+                if (env->selected_square[mover_idx] == SQ_NONE || env->valid_destinations[mover_idx].count == 0) {
+                    fprintf(stderr, "c_step: pick_phase=1 but selected_square=%u, valid_destinations.count=%d (mover=%d)\n",
+                            env->selected_square[mover_idx], env->valid_destinations[mover_idx].count, mover_idx);
+                    exit(1);
+                }
+
+                Square target_sq = SQ_NONE;
+                Move chosen_move = MOVE_NONE;
+                int desired_promo = -1;
+                int desired_file = -1;
+
+                if (is_promo) {
+                    int promo_row = (action - 64) / 8;
+                    desired_file = (action - 64) % 8;
+                    desired_promo = QUEEN - promo_row;
+                } else {
+                    target_sq = (mover == CHESS_BLACK) ? (Square)(action ^ 56) : (Square)action;
+                }
+
+                for (int i = 0; i < env->valid_destinations[mover_idx].count; i++) {
+                    Move m = env->valid_destinations[mover_idx].moves[i].move;
+                    if (!is_promo) {
+                        if ((int)to_sq(m) == (int)target_sq) {
+                            chosen_move = m;
+                            break;
+                        }
+                    } else {
+                        if ((int)type_of_m(m) == PROMOTION
+                                && (int)promotion_type(m) == desired_promo
+                                && (int)file_of(to_sq(m)) == desired_file) {
+                            chosen_move = m;
+                            break;
+                        }
+                    }
+                }
+
+                if (chosen_move == MOVE_NONE) {
+                    if (penalty_slot >= 0) {
+                        *env->reward_ptr[penalty_slot] += env->reward_invalid_move;
+                        env->invalid_actions_this_episode++;
+                    }
+                    clear_player_selection(env, mover_idx);
+                } else {
+                    game_result = apply_move_to_env(env, chosen_move, &is_timeout);
+                    if (env->reward_repetition != 0.0f
+                            && penalty_slot >= 0
+                            && env->repetition_matches >= 1) {
+                        *env->reward_ptr[penalty_slot] += env->reward_repetition;
+                    }
+                    move_completed = 1;
+                }
+            }
+        }
+    }
+
+    if (!move_completed) {
+        if (env->chess_moves >= env->max_moves || env->undo_stack_ptr >= MAX_GAME_PLIES - 2) {
+            game_result = 3;
+            is_timeout = 1;
+        } else {
+            if (env->legal_moves.count == 0) {
+                if (is_check(&env->pos, env->pos.sideToMove)) {
+                    game_result = env->pos.sideToMove == CHESS_WHITE ? 1 : 2;
+                } else {
+                    game_result = 3;
+                }
+            } else if (is_insufficient_material(&env->pos)) {
+                game_result = 3;
+            } else if (env->enable_50_move_rule && env->pos.rule50 >= 100) {
+                game_result = 3;
+            } else if (env->enable_threefold_repetition && env->repetition_matches >= 2) {
+                game_result = 3;
+            }
+        }
+    }
+
+    if (game_result != 0) {
+        *env->terminal_ptr[0] = 1;
+        if (env->mode == CHESS_MODE_SELFPLAY) {
+            *env->terminal_ptr[1] = 1;
+        }
+        env->game_result = game_result;
+        float win_value = 0.0f;
+
+        switch (game_result) {
+            case 3:
+                *env->reward_ptr[0] = env->reward_draw;
+                if (env->mode == CHESS_MODE_SELFPLAY) {
+                    *env->reward_ptr[1] = env->reward_draw;
+                    env->log.slot_0_score += 0.5f;
+                    env->log.slot_1_score += 0.5f;
+                }
+                win_value = 0.5f;
+                env->log.draw_rate += 1.0f;
+                if (is_timeout) {
+                    env->log.timeout_rate += 1.0f;
+                }
+                env->white_score += 0.5f;
+                env->black_score += 0.5f;
+                env->learner_draws += 1.0f;
+                strcpy(env->last_result, "Draw");
+                break;
+            case 1:
+                env->black_score += 1.0f;
+                if (env->mode == CHESS_MODE_SELFPLAY) {
+                    *env->reward_ptr[env->slot_for_color[CHESS_WHITE]] = -1.0f;
+                    *env->reward_ptr[env->slot_for_color[CHESS_BLACK]] = 1.0f;
+                    if (env->slot_for_color[CHESS_BLACK] == 0) env->log.slot_0_score += 1.0f;
+                    else env->log.slot_1_score += 1.0f;
+                    win_value = 0.5f;             // zero-sum: averaged across both slots
+                } else if (env->learner_color == CHESS_WHITE) {
+                    *env->reward_ptr[0] = -1.0f;
+                    env->learner_losses += 1.0f;
+                } else {
+                    *env->reward_ptr[0] = 1.0f;
+                    win_value = 1.0f;
+                    env->learner_wins += 1.0f;
+                }
+                strcpy(env->last_result, "Black Wins");
+                break;
+            case 2:
+                env->white_score += 1.0f;
+                if (env->mode == CHESS_MODE_SELFPLAY) {
+                    *env->reward_ptr[env->slot_for_color[CHESS_WHITE]] = 1.0f;
+                    *env->reward_ptr[env->slot_for_color[CHESS_BLACK]] = -1.0f;
+                    if (env->slot_for_color[CHESS_WHITE] == 0) env->log.slot_0_score += 1.0f;
+                    else env->log.slot_1_score += 1.0f;
+                    win_value = 0.5f;
+                } else if (env->learner_color == CHESS_WHITE) {
+                    *env->reward_ptr[0] = 1.0f;
+                    win_value = 1.0f;
+                    env->learner_wins += 1.0f;
+                } else {
+                    *env->reward_ptr[0] = -1.0f;
+                    env->learner_losses += 1.0f;
+                }
+                strcpy(env->last_result, "White Wins");
+                break;
+            default:
+                break;
+        }
+
+        if (env->mode == CHESS_MODE_SELFPLAY) {
+            env->episode_reward += *env->reward_ptr[0] + *env->reward_ptr[1];
+        } else {
+            env->episode_reward += *env->reward_ptr[0];
+        }
+        env->log.episode_return += env->episode_reward;
+        env->log.perf += win_value;
+        env->log.score += win_value;
+        env->log.chess_moves += env->chess_moves;
+        env->log.episode_length += env->tick;
+        env->log.invalid_action_rate += (env->tick > 0)
+            ? ((float)env->invalid_actions_this_episode / (float)env->tick) : 0.0f;
+
+        env->log.n += 1.0f;
+
+        // Per-color split for non-selfplay eval. learner_color is the color the
+        // learner just played; win_value already accounts for whether it won.
+        // A lopsided split (e.g. high white win rate, near-zero black) is the
+        // signature of a broken obs flip / wrong perspective for one color.
+        if (env->mode != CHESS_MODE_SELFPLAY) {
+            if (env->learner_color == CHESS_WHITE) {
+                env->log.wins_as_white += win_value;
+                env->log.games_as_white += 1.0f;
+            } else {
+                env->log.wins_as_black += win_value;
+                env->log.games_as_black += 1.0f;
+            }
+        }
+
+        // Per-bank historical tracking. Tag = 1..CHESS_MAX_BANKS picks the bank
+        // index (tag-1) this env was assigned to play. Tag = 0 is pure selfplay
+        // (skip historical accounting). Backward compat: tag=1 with single bank
+        // still routes into hist_score_bank[0] / hist_n_bank[0].
+        if (env->tag > 0 && env->tag <= CHESS_MAX_BANKS) {
+            int bank_idx = env->tag - 1;
+            float primary_score;
+            if (game_result == 3) {
+                primary_score = 0.5f;
+            } else if (game_result == 2) {  // White wins
+                primary_score = (env->slot_for_color[CHESS_WHITE] == 0) ? 1.0f : 0.0f;
+            } else {  // Black wins
+                primary_score = (env->slot_for_color[CHESS_BLACK] == 0) ? 1.0f : 0.0f;
+            }
+            env->log.hist_score_bank[bank_idx] += primary_score;
+            env->log.hist_n_bank[bank_idx] += 1.0f;
+            // Legacy aggregate fields — sum across all banks.
+            env->log.hist_score += primary_score;
+            env->log.hist_n += 1.0f;
+            env->boundary_reached = 1;
+        }
+
+        if (env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM) {
+            env->show_game_end_popup = 1;
+        } else {
+            if (env->log_pgn && env->pgn_filename[0] != '\0') {
+                env->pgn_game_number++;
+                export_pgn_append(env, env->pgn_filename, 1);
+            }
+            c_reset(env);
+        }
+    } else {
+        if (env->mode == CHESS_MODE_SELFPLAY) {
+            env->episode_reward += *env->reward_ptr[0] + *env->reward_ptr[1];
+        } else {
+            env->episode_reward += *env->reward_ptr[0];
+        }
+    }
+
+    populate_observations(env);
+}
+static Font load_piece_font(int cell_size, int* loaded) {
+    const char* candidates[] = {
+        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
+        "/usr/share/fonts/truetype/noto/NotoSansSymbols2-Regular.ttf",
+        "/System/Library/Fonts/Supplemental/Apple Symbols.ttf",
+        "C:\\Windows\\Fonts\\seguisym.ttf"
+    };
+
+    int codepoints[] = {0x2654, 0x2655, 0x2656, 0x2657, 0x2658, 0x2659, 0x265A, 0x265B, 0x265C, 0x265D, 0x265E, 0x265F};
+    Font font = (Font){0};
+    size_t candidate_count = sizeof(candidates) / sizeof(candidates[0]);
+    size_t codepoint_count = sizeof(codepoints) / sizeof(codepoints[0]);
+
+    for (size_t i = 0; i < candidate_count; i++) {
+        if (!FileExists(candidates[i])) {
+            continue;
+        }
+        font = LoadFontEx(candidates[i], cell_size, codepoints, (int)codepoint_count);
+        if (font.texture.id != 0) {
+            if (loaded) {
+                *loaded = 1;
+            }
+            SetTextureFilter(font.texture, TEXTURE_FILTER_BILINEAR);
+            return font;
+        }
+    }
+
+    if (loaded) {
+        *loaded = 0;
+    }
+    return GetFontDefault();
+}
+
+static void draw_piece(Chess* env, Piece pc, int file, int rank, int cell_size) {
+    if (pc == NO_PIECE) {
+        return;
+    }
+    
+    Color pc_color = color_of(pc) == CHESS_WHITE 
+        ? (Color){255, 255, 255, 255}
+        : (Color){0, 0, 0, 255};
+    
+    Color outline = (color_of(pc) == CHESS_WHITE) 
+        ? (Color){0, 0, 0, 220} 
+        : (Color){255, 255, 255, 180};
+
+    int draw_x = file * cell_size;
+    int draw_y = (7 - rank) * cell_size;
+
+    if (env->client && env->client->use_unicode_pieces) {
+        float icon_size = cell_size * 0.85f;
+        Vector2 pos = (Vector2){
+            draw_x + (cell_size - icon_size) / 2.0f,
+            draw_y + (cell_size - icon_size) / 2.0f - cell_size * 0.05f
+        };
+        const char* str = PIECE_FILLED[pc];
+        for (int dx = -1; dx <= 1; dx++) {
+            for (int dy = -1; dy <= 1; dy++) {
+                if (dx != 0 || dy != 0) {
+                    Vector2 opos = (Vector2){pos.x + dx, pos.y + dy};
+                    DrawTextEx(env->client->piece_font, str, opos, icon_size, 0, outline);
+                }
+            }
+        }
+        DrawTextEx(env->client->piece_font, str, pos, icon_size, 0, pc_color);
+    } else {
+        int x = draw_x + cell_size / 4;
+        int y = draw_y + cell_size / 8;
+        for (int dx = -1; dx <= 1; dx++) {
+            for (int dy = -1; dy <= 1; dy++) {
+                if (dx != 0 || dy != 0) {
+                    DrawText(PIECE_CHARS[pc], x + dx, y + dy, cell_size / 2, outline);
+                }
+            }
+        }
+        DrawText(PIECE_CHARS[pc], x, y, cell_size / 2, pc_color);
+    }
+}
+
+static void init_chess_client(Chess* env, int cell_size) {
+    SetConfigFlags(FLAG_MSAA_4X_HINT);
+    int board_size = 8 * cell_size;
+    InitWindow(board_size, board_size + 140, "PufferLib Chess - AI vs Opponent");
+    SetTargetFPS(env->render_fps > 0 ? env->render_fps : 30);
+    env->client = (Client*)calloc(1, sizeof(Client));
+    env->client->cell_size = cell_size;
+    int font_loaded = 0;
+    env->client->piece_font = load_piece_font(cell_size, &font_loaded);
+    env->client->use_unicode_pieces = font_loaded;
+    if (env->mode == CHESS_MODE_SELFPLAY) env->log_pgn_choice_made = 0;
+}
+
+void c_render(Chess* env) {
+    const int cell_size = 64;
+    const int board_size = 8 * cell_size;
+    const int scoreboard_y = board_size + 10;
+    static int speed_idx = 3;
+    static const int SPEED_FPS[] = {2, 5, 10, 30, 60, 120, 0};
+    static const int NUM_SPEEDS = 7;
+    static int selected_sq = -1;
+    
+    if (env->client == NULL) {
+        init_chess_client(env, cell_size);
+    }
+
+human_wait_retry:
+    if (IsKeyDown(KEY_ESCAPE) || WindowShouldClose()) { CloseWindow(); exit(0); }
+    
+    int flip_board = ((env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+        && env->human_color == CHESS_BLACK) ? 1 : 0;
+    Vector2 mouse = GetMousePosition();
+    int clicked = IsMouseButtonPressed(MOUSE_LEFT_BUTTON);
+    
+    if (IsKeyPressed(KEY_SPACE)) env->render_paused = !env->render_paused;
+    if (IsKeyPressed(KEY_EQUAL) || IsKeyPressed(KEY_KP_ADD)) {
+        if (speed_idx < NUM_SPEEDS - 1) { speed_idx++; SetTargetFPS(SPEED_FPS[speed_idx]); }
+    }
+    if (IsKeyPressed(KEY_MINUS) || IsKeyPressed(KEY_KP_SUBTRACT)) {
+        if (speed_idx > 0) { speed_idx--; SetTargetFPS(SPEED_FPS[speed_idx]); }
+    }
+    
+    if (!env->render_paused
+            && (env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+            && env->human_color != -1
+            && !env->show_game_end_popup
+            && clicked) {
+        int file = (int)(mouse.x) / cell_size;
+        int rank = 7 - ((int)(mouse.y) / cell_size);
+        if (flip_board) { file = 7 - file; rank = 7 - rank; }
+        if (file >= 0 && file < 8 && rank >= 0 && rank < 8) {
+            int clicked_sq = (int)make_square(file, rank);
+            if (selected_sq == -1) {
+                if (env->pos.sideToMove == env->human_color) {
+                    Piece pc = piece_on(&env->pos, (Square)clicked_sq);
+                    if (pc != NO_PIECE && color_of(pc) == env->human_color) {
+                        bool has_from = false;
+                        for (int i = 0; i < env->legal_moves.count; i++) {
+                            if ((int)from_sq(env->legal_moves.moves[i].move) == clicked_sq) { has_from = true; break; }
+                        }
+                        if (has_from) selected_sq = clicked_sq;
+                    }
+                }
+            } else {
+                Move chosen = MOVE_NONE;
+                for (int i = 0; i < env->legal_moves.count; i++) {
+                    Move m = env->legal_moves.moves[i].move;
+                    if ((int)from_sq(m) == selected_sq && (int)to_sq(m) == clicked_sq) { chosen = m; break; }
+                }
+                if (chosen != MOVE_NONE) {
+                    int is_timeout = 0;
+                    apply_move_to_env(env, chosen, &is_timeout);
+                    *env->action_ptr[0] = -1;
+                }
+                selected_sq = -1;
+            }
+        }
+    }
+
+    BeginDrawing();
+    ClearBackground((Color){40, 40, 40, 255});
+    
+    if ((env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+            && env->show_game_end_popup) {
+        int pw = 300, ph = 200;
+        int px = (board_size - pw) / 2, py = (board_size - ph) / 2;
+        DrawRectangle(px, py, pw, ph, (Color){60, 60, 60, 255});
+        DrawRectangleLines(px, py, pw, ph, WHITE);
+        DrawText("Game Over!", px + 70, py + 20, 24, WHITE);
+        DrawText(env->last_result, px + 80, py + 55, 18, YELLOW);
+        
+        Rectangle save_btn = {px + 20, py + 110, 120, 35};
+        Rectangle new_btn  = {px + 160, py + 110, 120, 35};
+        DrawRectangleRec(save_btn, DARKGREEN);
+        DrawRectangleLinesEx(save_btn, 2, WHITE);
+        DrawText("Save PGN", px + 35, py + 120, 16, WHITE);
+        DrawRectangleRec(new_btn, DARKBLUE);
+        DrawRectangleLinesEx(new_btn, 2, WHITE);
+        DrawText("New Game", px + 175, py + 120, 16, WHITE);
+        
+        if (clicked) {
+            if (CheckCollisionPointRec(mouse, save_btn)) {
+                char filename[64];
+                snprintf(filename, sizeof(filename), "game_%d.pgn", (int)time(NULL));
+                export_pgn_append(env, filename, 0);
+                printf("Saved PGN to %s\n", filename);
+            } else if (CheckCollisionPointRec(mouse, new_btn)) {
+                c_reset(env);
+            }
+        }
+    } else if (env->mode == CHESS_MODE_SELFPLAY && !env->log_pgn_choice_made) {
+        int cx = board_size / 2;
+        DrawText("Log PGN Files?", cx - 80, 180, 24, WHITE);
+        DrawText("Games will be appended to a timestamped file", cx - 160, 220, 14, LIGHTGRAY);
+        
+        Rectangle yes_btn = {cx - 70, 270, 140, 40};
+        Rectangle no_btn  = {cx - 70, 330, 140, 40};
+        DrawRectangleRec(yes_btn, DARKGREEN);
+        DrawRectangleLinesEx(yes_btn, 2, WHITE);
+        DrawText("Yes, Log PGN", cx - 55, 282, 16, WHITE);
+        DrawRectangleRec(no_btn, MAROON);
+        DrawRectangleLinesEx(no_btn, 2, WHITE);
+        DrawText("No Logging", cx - 45, 342, 16, WHITE);
+        
+        if (clicked) {
+            if (CheckCollisionPointRec(mouse, yes_btn)) {
+                env->log_pgn = 1;
+                env->log_pgn_choice_made = 1;
+                env->pgn_game_number = 0;
+                snprintf(env->pgn_filename, sizeof(env->pgn_filename), "run_%d_pgns.pgn", (int)time(NULL));
+                printf("PGN logging enabled: %s\n", env->pgn_filename);
+            } else if (CheckCollisionPointRec(mouse, no_btn)) {
+                env->log_pgn = 0;
+                env->log_pgn_choice_made = 1;
+                printf("PGN logging disabled\n");
+            }
+        }
+    } else if ((env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+            && env->human_color == -1) {
+        int cx = board_size / 2;
+        DrawText("Choose Your Color", cx - 100, 200, 24, WHITE);
+        
+        Rectangle white_btn = {cx - 60, 280, 120, 40};
+        Rectangle black_btn = {cx - 60, 340, 120, 40};
+        DrawRectangleRec(white_btn, LIGHTGRAY);
+        DrawRectangleLinesEx(white_btn, 2, BLACK);
+        DrawText("Play White", cx - 45, 292, 18, BLACK);
+        DrawRectangleRec(black_btn, GRAY);
+        DrawRectangleLinesEx(black_btn, 2, BLACK);
+        DrawText("Play Black", cx - 45, 352, 18, WHITE);
+        
+        if (clicked) {
+            if (CheckCollisionPointRec(mouse, white_btn)) {
+                env->human_color = CHESS_WHITE;
+                env->learner_color = CHESS_BLACK;
+            } else if (CheckCollisionPointRec(mouse, black_btn)) {
+                env->human_color = CHESS_BLACK;
+                env->learner_color = CHESS_WHITE;
+            }
+        }
+    } else {
+        Bitboard selected_destinations = 0;
+        if (selected_sq != -1) {
+            for (int i = 0; i < env->legal_moves.count; i++) {
+                Move m = env->legal_moves.moves[i].move;
+                if ((int)from_sq(m) == selected_sq) {
+                    selected_destinations |= sq_bb(to_sq(m));
+                }
+            }
+        }
+        int selected_file = -1;
+        int selected_rank = -1;
+        if (selected_sq != -1) {
+            selected_file = file_of((Square)selected_sq);
+            selected_rank = rank_of((Square)selected_sq);
+        }
+        for (int rank = 0; rank < 8; rank++) {
+            for (int file = 0; file < 8; file++) {
+                Color sq_color = ((rank + file) % 2 == 1) ? (Color){240, 217, 181, 255} : (Color){181, 136, 99, 255};
+                int draw_file = flip_board ? (7 - file) : file;
+                int draw_rank = flip_board ? (7 - rank) : rank;
+                int draw_x = draw_file * cell_size;
+                int draw_y = (7 - draw_rank) * cell_size;
+                DrawRectangle(draw_x, draw_y, cell_size, cell_size, sq_color);
+
+                if (env->has_last_move_highlight) {
+                    Square lf = env->last_move_from;
+                    Square lt = env->last_move_to;
+                    if ((file == (int)file_of(lf) && rank == (int)rank_of(lf))
+                            || (file == (int)file_of(lt) && rank == (int)rank_of(lt))) {
+                        Color last_mv = (Color){247, 247, 105, 255};
+                        DrawRectangle(draw_x, draw_y, cell_size, cell_size, Fade(last_mv, 0.52f));
+                    }
+                }
+
+                if (selected_sq != -1 && selected_file == file && selected_rank == rank) {
+                    DrawRectangleLines(draw_x, draw_y, cell_size, cell_size, (Color){255, 215, 0, 255});
+                }
+                if (selected_sq != -1 && (selected_destinations & sq_bb(make_square(file, rank)))) {
+                    DrawRectangleLines(draw_x + 2, draw_y + 2, cell_size - 4, cell_size - 4, (Color){0, 200, 0, 255});
+                }
+            }
+        }
+        for (int pt = PAWN; pt <= KING; pt++) {
+            Bitboard bb = pieces_p(&env->pos, pt);
+            while (bb) {
+                Square sq = pop_lsb(&bb);
+                Piece pc = piece_on(&env->pos, sq);
+                int f = file_of(sq), r = rank_of(sq);
+                int draw_f = flip_board ? (7 - f) : f;
+                int draw_r = flip_board ? (7 - r) : r;
+                draw_piece(env, pc, draw_f, draw_r, cell_size);
+            }
+        }
+        
+        char buf[128];
+        snprintf(buf, sizeof(buf), "White: %.1f  Black: %.1f", env->white_score, env->black_score);
+        DrawText(buf, 10, scoreboard_y, 20, WHITE);
+        
+        snprintf(buf, sizeof(buf), "Learner: %.0f-%.0f-%.0f (W-L-D)", env->learner_wins, env->learner_losses, env->learner_draws);
+        DrawText(buf, 10, scoreboard_y + 22, 16, GREEN);
+        
+        snprintf(buf, sizeof(buf), "Move: %d", env->chess_moves);
+        DrawText(buf, board_size - 100, scoreboard_y, 18, LIGHTGRAY);
+        
+        if (env->mode != CHESS_MODE_HUMAN && env->mode != CHESS_MODE_HUMAN_RANDOM) {
+            DrawText(env->learner_color == CHESS_WHITE ? "Learner: White" : "Learner: Black",
+                     board_size - 120, scoreboard_y + 22, 16, LIGHTGRAY);
+        }
+        
+        int cap_y = scoreboard_y + 42;
+        int cap_x_start = 10;
+        Color white_cap_color = (Color){240, 217, 181, 255};
+        Color black_cap_color = (Color){100, 100, 100, 255};
+        int white_x = cap_x_start;
+        int black_x = cap_x_start;
+
+        for (int pt = 0; pt < 6; pt++) {
+            int wc = env->white_captured[pt];
+            if (wc > 0) {
+                Piece wpc = (Piece)(W_PAWN + pt);
+                if (env->client && env->client->use_unicode_pieces) {
+                    DrawTextEx(env->client->piece_font, PIECE_FILLED[wpc],
+                        (Vector2){(float)white_x, (float)(cap_y - 1)}, 16.0f, 0.0f, white_cap_color);
+                    white_x += 16;
+                } else {
+                    DrawText(PIECE_CHARS[wpc], white_x, cap_y, 14, white_cap_color);
+                    white_x += 12;
+                }
+                if (wc > 1) {
+                    char mult[8];
+                    snprintf(mult, sizeof(mult), "x%d", wc);
+                    DrawText(mult, white_x, cap_y + 2, 10, white_cap_color);
+                    white_x += MeasureText(mult, 10) + 4;
+                } else {
+                    white_x += 4;
+                }
+            }
+
+            int bc = env->black_captured[pt];
+            if (bc > 0) {
+                Piece bpc = (Piece)(B_PAWN + pt);
+                Color outline = (Color){255, 255, 255, 180};
+                if (env->client && env->client->use_unicode_pieces) {
+                    Vector2 pos = {(float)black_x, (float)(cap_y + 17)};
+                    for (int dx = -1; dx <= 1; dx++) {
+                        for (int dy = -1; dy <= 1; dy++) {
+                            if (dx != 0 || dy != 0) {
+                                DrawTextEx(env->client->piece_font, PIECE_FILLED[bpc],
+                                    (Vector2){pos.x + dx, pos.y + dy}, 16.0f, 0.0f, outline);
+                            }
+                        }
+                    }
+                    DrawTextEx(env->client->piece_font, PIECE_FILLED[bpc], pos, 16.0f, 0.0f, black_cap_color);
+                    black_x += 16;
+                } else {
+                    for (int dx = -1; dx <= 1; dx++) {
+                        for (int dy = -1; dy <= 1; dy++) {
+                            if (dx != 0 || dy != 0)
+                                DrawText(PIECE_CHARS[bpc], black_x + dx, cap_y + 18 + dy, 14, outline);
+                        }
+                    }
+                    DrawText(PIECE_CHARS[bpc], black_x, cap_y + 18, 14, black_cap_color);
+                    black_x += 12;
+                }
+                if (bc > 1) {
+                    char mult[8];
+                    snprintf(mult, sizeof(mult), "x%d", bc);
+                    for (int dx = -1; dx <= 1; dx++) {
+                        for (int dy = -1; dy <= 1; dy++) {
+                            if (dx != 0 || dy != 0)
+                                DrawText(mult, black_x + dx, cap_y + 20 + dy, 10, outline);
+                        }
+                    }
+                    DrawText(mult, black_x, cap_y + 20, 10, black_cap_color);
+                    black_x += MeasureText(mult, 10) + 4;
+                } else {
+                    black_x += 4;
+                }
+            }
+        }
+
+        if (env->last_result[0] != '\0') {
+            Color rc = YELLOW;
+            if (strstr(env->last_result, "White")) rc = (Color){240, 217, 181, 255};
+            else if (strstr(env->last_result, "Black")) rc = (Color){100, 100, 100, 255};
+            DrawText(env->last_result, 10, cap_y + 40, 18, rc);
+        }
+        
+        int btn_w = 36;
+        int btn_h = 24;
+        int btn_y = scoreboard_y + 100;
+        int btn_x = (env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+            ? board_size / 2 - 100 : board_size / 2 - 70;
+        Rectangle minus_btn = {btn_x, btn_y, btn_w, btn_h};
+        Rectangle pause_btn = {btn_x + btn_w + 5, btn_y, btn_w + 10, btn_h};
+        Rectangle plus_btn = {btn_x + 2 * btn_w + 20, btn_y, btn_w, btn_h};
+        DrawRectangleRec(minus_btn, DARKGRAY);
+        DrawRectangleLinesEx(minus_btn, 2, LIGHTGRAY);
+        DrawText("-", btn_x + 14, btn_y + 4, 20, WHITE);
+        DrawRectangleRec(pause_btn, env->render_paused ? MAROON : DARKGREEN);
+        DrawRectangleLinesEx(pause_btn, 2, LIGHTGRAY);
+        DrawText(env->render_paused ? ">" : "||", btn_x + btn_w + 14, btn_y + 4, 18, WHITE);
+        DrawRectangleRec(plus_btn, DARKGRAY);
+        DrawRectangleLinesEx(plus_btn, 2, LIGHTGRAY);
+        DrawText("+", btn_x + 2 * btn_w + 32, btn_y + 4, 20, WHITE);
+        char speed_buf[32];
+        if (SPEED_FPS[speed_idx] == 0) {
+            snprintf(speed_buf, sizeof(speed_buf), "max");
+        } else {
+            snprintf(speed_buf, sizeof(speed_buf), "%dfps", SPEED_FPS[speed_idx]);
+        }
+        DrawText(speed_buf, btn_x + 3 * btn_w + 30, btn_y + 4, 14, env->render_paused ? RED : LIGHTGRAY);
+        
+        Rectangle restart_btn = {0, 0, 0, 0};
+        if (env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM) {
+            restart_btn = (Rectangle){board_size - 60, minus_btn.y, 55, minus_btn.height};
+            DrawRectangleRec(restart_btn, MAROON);
+            DrawRectangleLinesEx(restart_btn, 2, LIGHTGRAY);
+            DrawText("Exit", board_size - 53, minus_btn.y + 4, 16, WHITE);
+        }
+        
+        if (env->render_paused) {
+            DrawRectangle(0, 0, board_size, board_size, (Color){0, 0, 0, 120});
+            DrawText("PAUSED", board_size / 2 - 60, board_size / 2 - 15, 30, RED);
+        }
+        
+        if (clicked) {
+            if (CheckCollisionPointRec(mouse, minus_btn)) {
+                if (speed_idx > 0) { speed_idx--; SetTargetFPS(SPEED_FPS[speed_idx]); }
+            }
+            if (CheckCollisionPointRec(mouse, pause_btn)) env->render_paused = !env->render_paused;
+            if (CheckCollisionPointRec(mouse, plus_btn)) {
+                if (speed_idx < NUM_SPEEDS - 1) { speed_idx++; SetTargetFPS(SPEED_FPS[speed_idx]); }
+            }
+            if ((env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+                    && CheckCollisionPointRec(mouse, restart_btn)) c_reset(env);
+        }
+    }
+    
+    EndDrawing();
+
+    // Human-mode only: stay in c_render (on the window-owning thread) until
+    // the human commits a move via mouse clicks. Re-poll input + redraw each
+    // iteration. Non-human modes fall through to a single c_render call as
+    // before. Refresh obs after the commit so the next rollout inference sees
+    // the post-human-move state instead of the stale "human's turn" obs.
+    if ((env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+            && env->human_color != -1
+            && env->pos.sideToMove == env->human_color
+            && !env->show_game_end_popup
+            && !env->render_paused
+            && !WindowShouldClose()) {
+        goto human_wait_retry;
+    }
+    if ((env->mode == CHESS_MODE_HUMAN || env->mode == CHESS_MODE_HUMAN_RANDOM)
+            && env->human_color != -1
+            && env->pos.sideToMove != env->human_color
+            && !env->show_game_end_popup) {
+        if (env->legal_dirty) rebuild_legal_state(env);
+        populate_observations(env);
+    }
+}
+
+void c_close(Chess* env) {
+    if (env->client != NULL) {
+        if (env->client->use_unicode_pieces && env->client->piece_font.texture.id != 0) {
+            UnloadFont(env->client->piece_font);
+        }
+        if (IsWindowReady()) {
+            CloseWindow();
+        }
+        free(env->client);
+        env->client = NULL;
+    }
+    maia_close(env);
+    env->fen_curriculum = NULL;
+    env->num_fens = 0;
+}
diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py
index 8fc0c03a89..929c300b48 100644
--- a/pufferlib/pufferl.py
+++ b/pufferlib/pufferl.py
@@ -25,6 +25,8 @@
 except ImportError:
     raise ImportError('Failed to import PufferLib C++ backend. If you have non-default PyTorch, try installing with --no-build-isolation')
 
+from pufferlib import selfplay
+
 import rich
 import rich.traceback
 from rich.table import Table
@@ -203,6 +205,11 @@ def _train(env_name, args, sweep_obj=None, result_queue=None, verbose=False):
     total_timesteps = args['train']['total_timesteps']
     all_logs = []
 
+    # When sweeping, optionally score each trial by winrate vs a fixed enemy
+    # checkpoint (match mode) instead of the training-time self-play metric.
+    match_mode = (sweep_obj is not None
+        and bool(args.get('sweep', {}).get('match_enemy_model_path')))
+
     checkpoint_dir = os.path.join(args['checkpoint_dir'], args['env_name'], run_id)
     os.makedirs(checkpoint_dir, exist_ok=True)
 
@@ -223,6 +230,18 @@ def _train(env_name, args, sweep_obj=None, result_queue=None, verbose=False):
         flat_logs = dict(unroll_nested_dict(backend.log(pufferl)))
         print_dashboard(args, model_size, flat_logs, clear=True)
 
+    # Selfplay-pool curriculum (no-op unless selfplay.enabled). Disabled
+    # under match-mode sweeps since match() owns its own perm/frozen bank.
+    pool_state = None
+    try:
+        pool_state = selfplay.setup(pufferl, backend, args, run_id)
+    except RuntimeError as e:
+        print(f'WARNING: {e}, skipping')
+        backend.close(pufferl)
+        if result_queue is not None:
+            result_queue.put((args['gpu_id'], [], [], []))
+        return
+
     model_path = ''
     flat_logs = {}
     train_epochs = int(total_timesteps // (args['vec']['total_agents'] * args['train']['horizon']))
@@ -233,7 +252,12 @@ def _train(env_name, args, sweep_obj=None, result_queue=None, verbose=False):
         if epoch < train_epochs:
             backend.train(pufferl)
 
-        if (epoch % args['checkpoint_interval'] == 0 or epoch == train_epochs - 1) and sweep_obj is None:
+        # In match-sweep mode we need the final checkpoint to feed into match().
+        is_final = epoch == train_epochs - 1
+        should_save = (sweep_obj is None
+            and (epoch % args['checkpoint_interval'] == 0 or is_final)
+        ) or (match_mode and is_final)
+        if should_save:
             model_path = os.path.join(checkpoint_dir, f'{pufferl.global_step:016d}.bin')
             backend.save_weights(pufferl, model_path)
 
@@ -244,6 +268,9 @@ def _train(env_name, args, sweep_obj=None, result_queue=None, verbose=False):
         logs = backend.eval_log(pufferl) if epoch >= train_epochs else backend.log(pufferl)
         flat_logs = {**flat_logs, **dict(unroll_nested_dict(logs))}
 
+        if epoch < train_epochs:
+            selfplay.step(pufferl, backend, pool_state, flat_logs, epoch)
+
         if verbose:
             print_dashboard(args, model_size, flat_logs)
 
@@ -265,6 +292,11 @@ def _train(env_name, args, sweep_obj=None, result_queue=None, verbose=False):
 
 
     print_dashboard(args, model_size, flat_logs)
+    # Match-mode trials may have early-stopped before the in-loop save fired;
+    # ensure we always have a checkpoint to feed match().
+    if match_mode and not model_path:
+        model_path = os.path.join(checkpoint_dir, f'{pufferl.global_step:016d}.bin')
+        backend.save_weights(pufferl, model_path)
     backend.close(pufferl)
 
     if target_key not in flat_logs:
@@ -272,6 +304,25 @@ def _train(env_name, args, sweep_obj=None, result_queue=None, verbose=False):
             result_queue.put((args['gpu_id'], None, None, None))
         return
 
+    # Match-mode scoring: primary = trained policy (model_path); frozen bank =
+    # fixed enemy. Score is slot 0's average winrate. Creates its own pufferl
+    # so must run after the training instance is closed. Single observation per
+    # trial (mid-training curve doesn't predict final match score).
+    match_score = None
+    if match_mode:
+        sweep_cfg = args['sweep']
+        match_args = deepcopy(args)
+        match_args['enemy_hidden_size'] = int(sweep_cfg['match_enemy_hidden_size'])
+        match_args['enemy_num_layers'] = int(sweep_cfg['match_enemy_num_layers'])
+        match_logs = match(env_name,
+            policy_a_path=model_path,
+            policy_b_path=sweep_cfg['match_enemy_model_path'],
+            num_games=int(sweep_cfg['match_num_games']),
+            args=match_args, verbose=verbose)
+        match_score = float(match_logs['env/slot_0_score'])
+        if args['wandb']:
+            wandb.log({'env/match_score': match_score}, step=flat_logs['agent_steps'])
+
     # This version has the training perf logs and eval env logs
     all_logs.append(flat_logs)
 
@@ -295,6 +346,14 @@ def _train(env_name, args, sweep_obj=None, result_queue=None, verbose=False):
     for k in metrics:
         metrics[k][-1] = all_logs[-1][k]
 
+    # Match-mode: single observation at final-training cost. Protein's curve
+    # fit collapses to one point — we only trust the match winrate, not any
+    # training-time proxy. Replicate the scalar across all downsample bins so
+    # the JSON log shape matches every other metric (cache_data.py rejects
+    # length-mismatched metrics as "bad data").
+    if match_mode:
+        metrics['env/match_score'] = [match_score] * len(metrics['agent_steps'])
+
     # Save own log: config + downsampled results
     log_dir = os.path.join(args['log_dir'], args['env_name'])
     os.makedirs(log_dir, exist_ok=True)
@@ -310,7 +369,12 @@ def _train(env_name, args, sweep_obj=None, result_queue=None, verbose=False):
         wandb.run.finish()
 
     if result_queue is not None:
-        result_queue.put((args['gpu_id'], metrics['env/score'], metrics['uptime'], metrics['agent_steps']))
+        if match_mode:
+            # One observation: final hypers -> match winrate, at total training cost.
+            result_queue.put((args['gpu_id'], [match_score],
+                [metrics['uptime'][-1]], [metrics['agent_steps'][-1]]))
+        else:
+            result_queue.put((args['gpu_id'], metrics['env/score'], metrics['uptime'], metrics['agent_steps']))
 
 def train(env_name, args=None, gpus=None, **kwargs):
     args = args or load_config(env_name)
@@ -333,6 +397,9 @@ def train(env_name, args=None, gpus=None, **kwargs):
         if rank == 0 and not subprocess:
             _train(env_name, worker_args, verbose=True)
         else:
+            # Protein's GP models live on cuda:0 on non-WSL setups; spawn-pickling
+            # them works fine via CUDA IPC. On WSL, sweep.py forces device='cpu'
+            # at construction so there's nothing to move.
             ctx.Process(target=_train, args=(env_name, worker_args),
                 kwargs=kwargs).start()
 
@@ -427,10 +494,106 @@ def eval(env_name, args=None, load_path=None):
 
     backend.close(pufferl)
 
+def match(env_name, policy_a_path, policy_b_path, num_games=4096, args=None, verbose=True):
+    '''Head-to-head match between two trained policies in a 2-agent selfplay env.
+    Policy A plays slot 0 (e.g. white in chess), policy B plays slot 1 (black).
+    Both checkpoints must come from the same env / arch.
+    '''
+    args = args or load_config(env_name)
+    args['reset_state'] = False
+    args['train']['horizon'] = 1
+    args.setdefault('nccl_id', b'')  # match is always single-GPU
+    # Sweep suggestions can give odd agents_per_buffer (e.g. num_buffers=5,
+    # total_agents=4096 -> 819). Pin to a stable eval config that guarantees
+    # clean slot-0/slot-1 split; ignores trial's vec tuning (eval, not train).
+    args['vec']['num_buffers'] = 2
+    args['vec']['total_agents'] = 8192
+    backend = _resolve_backend(args)
+    if backend is not _C:
+        raise RuntimeError('match() requires the native CUDA backend')
+
+    def _resolve_latest(path):
+        if path != 'latest':
+            return path
+        pattern = os.path.join(args['checkpoint_dir'], args['env_name'], '**', '*.bin')
+        candidates = glob.glob(pattern, recursive=True)
+        if not candidates:
+            raise FileNotFoundError(f'No .bin checkpoints found in {args["checkpoint_dir"]}/{args["env_name"]}/')
+        return max(candidates, key=os.path.getctime)
+    policy_a_path = _resolve_latest(policy_a_path)
+    policy_b_path = _resolve_latest(policy_b_path)
+
+    total_agents = int(args['vec']['total_agents'])
+    num_buffers = int(args['vec']['num_buffers'])
+    agents_per_buffer = total_agents // num_buffers
+    half = agents_per_buffer // 2
+    if 2 * half != agents_per_buffer:
+        raise RuntimeError(f'agents_per_buffer ({agents_per_buffer}) must be even for 2-agent selfplay')
+
+    # Primary holds policy A (owns first half of each buffer); one frozen bank
+    # holds policy B (owns second half). Bank is created inside create_pufferl
+    # before cudagraph capture so the graph bakes in its pointers; weight loads
+    # later only update data.
+    args['vec']['num_frozen_banks'] = 1
+    args['vec']['frozen_bank_pct'] = 0.5
+    # CLI flags take precedence; fall back to [sweep].match_enemy_* so the same
+    # config drives sweep-time and CLI-time matches. 0 / None means "use primary".
+    sweep_cfg = args.get('sweep', {})
+    enemy_hidden = args.get('enemy_hidden_size') or sweep_cfg.get('match_enemy_hidden_size')
+    enemy_layers = args.get('enemy_num_layers')  or sweep_cfg.get('match_enemy_num_layers')
+    if enemy_hidden:
+        args['vec']['frozen_bank_hidden_size'] = int(enemy_hidden)
+    if enemy_layers:
+        args['vec']['frozen_bank_num_layers'] = int(enemy_layers)
+
+    pufferl = backend.create_pufferl(args)
+
+    # Per-buffer perm: each env's slot 0 lands in primary's slice [0, half),
+    # slot 1 lands in frozen bank's slice [half, agents_per_buffer). The env
+    # side randomizes slot<->color per env, so A and B each play both colors.
+    perm = np.empty(total_agents, dtype=np.int32)
+    envs_per_buffer = half
+    for b in range(num_buffers):
+        off = b * agents_per_buffer
+        for i in range(envs_per_buffer):
+            perm[off + 2*i]     = off + i
+            perm[off + 2*i + 1] = off + half + i
+    backend.set_agent_perm(pufferl, perm)
+
+    backend.load_weights(pufferl, policy_a_path)
+    backend.load_frozen_bank(pufferl, 0, policy_b_path)
+
+    logs = {}
+    while True:
+        backend.rollouts(pufferl)
+        logs = dict(unroll_nested_dict(backend.eval_log(pufferl)))
+        n = int(logs.get('env/n', 0))
+        if verbose:
+            a = logs.get('env/slot_0_score', 0.0)
+            b = logs.get('env/slot_1_score', 0.0)
+            draws = logs.get('env/draw_rate', 0.0)
+            print(f'\rgames={n}/{num_games}  A={a:.3f}  B={b:.3f}  draw={draws:.3f}', end='')
+        if n >= num_games:
+            break
+
+    if verbose:
+        print()
+
+    backend.close(pufferl)
+    return logs
+
 def load_config(env_name):
     parser = argparse.ArgumentParser(formatter_class=RichHelpFormatter, add_help=False)
     parser.add_argument('--load-model-path', type=str, default=None,
         help='Path to a pretrained checkpoint')
+    parser.add_argument('--load-enemy-model-path', type=str, default=None,
+        help='Path to opponent checkpoint for `puffer match` (slot 1 / black in chess)')
+    parser.add_argument('--num-games', type=int, default=4096,
+        help='Number of games to play in `puffer match`')
+    parser.add_argument('--enemy-hidden-size', type=int, default=None,
+        help='hidden_size of the enemy checkpoint (defaults to primary)')
+    parser.add_argument('--enemy-num-layers', type=int, default=None,
+        help='num_layers of the enemy checkpoint (defaults to primary)')
     parser.add_argument('--load-id', type=str,
         default=None, help='Kickstart/eval from from a finished Wandbrun')
     parser.add_argument('--render-mode', type=str, default='auto',
@@ -496,7 +659,7 @@ def load_config(env_name):
     return dict(args)
 
 def main():
-    err = 'Usage: puffer [train, eval, sweep, paretosweep] [env_name] [optional args]. --help for more info'
+    err = 'Usage: puffer [train, eval, sweep, paretosweep, match] [env_name] [optional args]. --help for more info'
     if len(sys.argv) < 3:
         raise ValueError(err)
 
@@ -510,6 +673,13 @@ def main():
         eval(env_name=env_name, args=args)
     elif 'sweep' in mode:
         sweep(env_name=env_name, args=args, pareto='pareto' in mode)
+    elif 'match' in mode:
+        a_path = args.get('load_model_path')
+        b_path = args.get('load_enemy_model_path')
+        if not a_path or not b_path:
+            raise ValueError('puffer match requires --load-model-path and --load-enemy-model-path')
+        match(env_name=env_name, policy_a_path=a_path, policy_b_path=b_path,
+            num_games=args.get('num_games', 4096), args=args)
     else:
         raise ValueError(err)
 
diff --git a/pufferlib/selfplay.py b/pufferlib/selfplay.py
new file mode 100644
index 0000000000..19cc4791b6
--- /dev/null
+++ b/pufferlib/selfplay.py
@@ -0,0 +1,317 @@
+"""Selfplay-pool training: a fraction of envs play primary vs a frozen historical
+snapshot, the rest are pure selfplay. Used by `_train` in pufferl.py — gated on
+`selfplay.enabled` (config section).
+
+Pool grows on two triggers:
+  - snapshot_interval: every N global steps, save primary weights as a new
+    pool entry regardless of winrate. Provides a steady cadence.
+  - winrate-driven swap: when primary beats the current opponent at >=
+    swap_winrate over >= min_games, also save primary as a pool entry, then
+    swap to a new opponent. Marks progress checkpoints in the curriculum.
+
+Swap (without a snapshot) also fires when opp_timeout_steps have elapsed
+since the current opponent was finalized. Timeout prevents stalemates from
+pinning the curriculum to a single opponent indefinitely.
+
+Pool storage is disk-only (paths held in memory; weights only on GPU when
+loaded as the frozen bank). Stride-eviction preserves temporal coverage when
+the pool exceeds its cap.
+"""
+import os
+
+import numpy as np
+
+from pufferlib import _C
+
+
+def sample_opponent(pool, rng):
+    candidates = pool if len(pool) < 6 else pool[:-5]
+    weights = np.array([(i + 1) ** 0.5 for i in range(len(candidates))], dtype=np.float64)
+    weights /= weights.sum()
+    idx = int(rng.choice(len(candidates), p=weights))
+    return candidates[idx]
+
+
+def update_elo(primary_elo, opp_elo, score_rate, k):
+    expected = 1.0 / (1.0 + 10.0 ** ((opp_elo - primary_elo) / 400.0))
+    delta = k * (score_rate - expected)
+    return primary_elo + delta, opp_elo - delta
+
+
+def evict(pool, max_size):
+    '''Drop every other entry from the oldest half once the pool exceeds max_size.
+    Newest half is preserved intact.'''
+    if len(pool) <= max_size:
+        return pool
+    half = len(pool) // 2
+    return pool[:half:2] + pool[half:]
+
+
+def build_perm_tags(num_buffers, agents_per_buffer, agents_per_env, frozen_sizes, num_envs):
+    '''Build env-slot -> rollout-row routing and per-env bank tag.
+
+    Multi-bank generalization. `frozen_sizes` is a list of per-bank agent counts
+    (per buffer). With one bank this reduces to the legacy single-bank layout.
+
+    Per-buffer physical-row layout (apb = agents_per_buffer, F = sum(frozen_sizes)):
+        [0,           apb - 2F)                       primary — selfplay envs (all slots)
+        [apb - 2F,    apb - F)                        primary — historical envs' team A
+        [apb - F,     apb - F + frozen_sizes[0])      bank 0  — historical envs' team B
+        [apb - F + frozen_sizes[0], ... + ...[1])     bank 1  — ... etc.
+
+    Env order within a buffer: selfplay envs first (tag=0), then historical
+    envs assigned to banks in block order — the first `frozen_sizes[0]/team_size`
+    historical envs play bank 0 (tag=1), next block plays bank 1 (tag=2), etc.
+
+    The C-side bank_layout (pufferlib.cu:1798-1806) lays banks out sequentially
+    after primary, so our routing matches: bank b's slice is
+    [apb - F + sum(frozen_sizes[:b]),  apb - F + sum(frozen_sizes[:b+1])).
+
+    Returns (perm, tags, num_hist_envs_per_bank) — last is a list of per-bank
+    historical-env counts across all buffers, used by selfplay.step to know how
+    many env alignments to wait for per bank during swaps.'''
+    team_size = agents_per_env // 2
+    envs_per_buffer = agents_per_buffer // agents_per_env
+    num_banks = len(frozen_sizes)
+    total_frozen = sum(frozen_sizes)
+    hist_envs_per_bank_per_buffer = [fs // team_size for fs in frozen_sizes]
+    total_hist_envs_per_buffer = sum(hist_envs_per_bank_per_buffer)
+    selfplay_envs = envs_per_buffer - total_hist_envs_per_buffer
+    perm = np.empty(num_buffers * agents_per_buffer, dtype=np.int32)
+    tags = np.zeros(num_envs, dtype=np.int32)
+    env_idx = 0
+    for b_buf in range(num_buffers):
+        buf_start          = b_buf * agents_per_buffer
+        hist_primary_start = buf_start + agents_per_buffer - 2 * total_frozen
+        bank_starts = []
+        offset = buf_start + agents_per_buffer - total_frozen
+        for bank in range(num_banks):
+            bank_starts.append(offset)
+            offset += frozen_sizes[bank]
+        h_within_buffer = 0
+        for e in range(envs_per_buffer):
+            slot_base = buf_start + e * agents_per_env
+            if e < selfplay_envs:
+                for s in range(agents_per_env):
+                    perm[slot_base + s] = slot_base + s
+                tags[env_idx] = 0
+            else:
+                # Block assignment: walk cumulative bank capacity to find which
+                # bank this historical env belongs to.
+                bank_idx = 0
+                cum = hist_envs_per_bank_per_buffer[0]
+                while h_within_buffer >= cum and bank_idx < num_banks - 1:
+                    bank_idx += 1
+                    cum += hist_envs_per_bank_per_buffer[bank_idx]
+                h_in_bank = h_within_buffer - (cum - hist_envs_per_bank_per_buffer[bank_idx])
+                team_a_offset = hist_primary_start + h_within_buffer * team_size
+                team_b_offset = bank_starts[bank_idx] + h_in_bank * team_size
+                for s in range(team_size):
+                    perm[slot_base + s] = team_a_offset + s
+                    perm[slot_base + team_size + s] = team_b_offset + s
+                tags[env_idx] = bank_idx + 1
+                h_within_buffer += 1
+            env_idx += 1
+    num_hist_envs_per_bank = [n * num_buffers for n in hist_envs_per_bank_per_buffer]
+    return perm, tags, num_hist_envs_per_bank
+
+
+def setup(pufferl, backend, args, run_id):
+    '''Wire up agent_perm/tags and bootstrap the frozen bank with the current
+    weights so historical envs have an opponent from rollout 1. Returns a
+    pool_state dict (or None if disabled).'''
+    sp = args.get('selfplay', {})
+    if not sp.get('enabled', 0):
+        return None
+    if backend is not _C:
+        raise RuntimeError('selfplay_pool requires the native CUDA backend')
+
+    total_agents = int(args['vec']['total_agents'])
+    num_buffers = int(args['vec']['num_buffers'])
+    if total_agents % num_buffers != 0:
+        raise RuntimeError(f'total_agents ({total_agents}) must be divisible by '
+                           f'num_buffers ({num_buffers})')
+    agents_per_buffer = total_agents // num_buffers
+
+    num_envs = backend.num_envs(pufferl)
+    agents_per_env = total_agents // num_envs
+    if agents_per_env % 2 != 0:
+        raise RuntimeError(f'agents_per_env ({agents_per_env}) must be even (two equal teams)')
+    if agents_per_buffer % agents_per_env != 0:
+        raise RuntimeError(f'agents_per_buffer ({agents_per_buffer}) must be divisible by '
+                           f'agents_per_env ({agents_per_env})')
+    team_size = agents_per_env // 2
+
+    num_banks = int(args['vec'].get('num_frozen_banks', 1))
+    if num_banks <= 0:
+        raise RuntimeError('selfplay.enabled requires num_frozen_banks >= 1')
+    if num_banks > 8:
+        raise RuntimeError(f'num_frozen_banks {num_banks} exceeds chess.h CHESS_MAX_BANKS=8')
+
+    # frozen_bank_pct is per-bank (matches C-side: pufferlib.cu:2069). Each bank
+    # gets floor(apb * pct) agents, total historical = num_banks * frozen_size.
+    frozen_size = int(agents_per_buffer * float(args['vec']['frozen_bank_pct']))
+    frozen_size -= frozen_size % team_size
+    if frozen_size <= 0:
+        raise RuntimeError('selfplay.enabled but frozen_bank_pct rounds to 0 slots '
+                           f'after team-size ({team_size}) alignment')
+    total_frozen = frozen_size * num_banks
+    if total_frozen >= agents_per_buffer // 2:
+        raise RuntimeError(f'total_frozen {total_frozen} (= num_banks {num_banks} '
+                           f'* per_bank {frozen_size}) >= apb/2 {agents_per_buffer//2}')
+
+    frozen_sizes = [frozen_size] * num_banks
+    perm, tags, num_hist_envs_per_bank = build_perm_tags(
+        num_buffers, agents_per_buffer, agents_per_env, frozen_sizes, num_envs)
+    backend.set_agent_perm(pufferl, perm)
+    backend.set_env_tags(pufferl, tags)
+
+    pool_dir = os.path.join(args['checkpoint_dir'], args['env_name'], run_id, 'pool')
+    os.makedirs(pool_dir, exist_ok=True)
+    bootstrap_path = os.path.join(pool_dir, f'{pufferl.global_step:016d}.bin')
+    backend.save_weights(pufferl, bootstrap_path)
+    # Load bootstrap into every bank — they'll diverge as each bank's swap fires.
+    for b in range(num_banks):
+        backend.load_frozen_bank(pufferl, b, bootstrap_path)
+
+    elo_init = float(sp.get('elo_init', 0.0))
+    elo_k    = float(sp.get('elo_k',    16.0))
+    rng = np.random.default_rng(int(sp.get('seed', 0)))
+
+    banks_state = []
+    for b in range(num_banks):
+        banks_state.append({
+            'cur_opp_path': bootstrap_path,
+            'cur_opp_elo': elo_init,
+            'hist_score': 0.0,
+            'hist_n': 0.0,
+            'pending_opp_path': None,
+            'pending_opp_elo': None,
+            'epoch_armed': 0,
+            'opp_started_step': int(pufferl.global_step),
+            'num_hist_envs': num_hist_envs_per_bank[b],
+            'last_winrate_at_swap': 0.0,
+            'last_epochs_to_align': 0,
+        })
+
+    return {
+        'pool_dir': pool_dir,
+        'pool': [{'path': bootstrap_path, 'elo': elo_init}],
+        'rng': rng,
+        'max_size': int(sp['max_size']),
+        'min_games': int(sp['min_games']),
+        'swap_winrate': float(sp['swap_winrate']),
+        'snapshot_interval': int(sp.get('snapshot_interval', 1_000_000_000)),
+        'opp_timeout_steps': int(sp.get('opp_timeout_steps', 500_000_000)),
+        'num_banks': num_banks,
+        'banks': banks_state,
+        'primary_elo': elo_init,
+        'elo_k': elo_k,
+        'last_snapshot_step': int(pufferl.global_step),
+    }
+
+
+def step(pufferl, backend, pool_state, flat_logs, epoch):
+    if pool_state is None:
+        return
+
+    n_window = float(flat_logs.get('env/n', 0.0))
+    num_banks = pool_state['num_banks']
+
+    # 1. Per-bank Elo update from the most recent rollout window.
+    for b in range(num_banks):
+        bank = pool_state['banks'][b]
+        hist_score_w = float(flat_logs.get(f'env/hist_score_bank_{b}', 0.0)) * n_window
+        hist_n_w     = float(flat_logs.get(f'env/hist_n_bank_{b}',     0.0)) * n_window
+        if hist_n_w > 0.0:
+            bank['hist_score'] += hist_score_w
+            bank['hist_n']     += hist_n_w
+            score_rate = hist_score_w / hist_n_w
+            new_p, new_o = update_elo(pool_state['primary_elo'],
+                bank['cur_opp_elo'], score_rate, pool_state['elo_k'])
+            # All banks update the shared primary Elo. Multiple banks updating
+            # primary in one step is fine — Elo is symmetric, just a few more
+            # tiny adjustments per rollout window.
+            pool_state['primary_elo'] = new_p
+            bank['cur_opp_elo'] = new_o
+            for entry in pool_state['pool']:
+                if entry['path'] == bank['cur_opp_path']:
+                    entry['elo'] = new_o
+                    break
+
+    # 2. Global snapshot cadence (shared across banks).
+    if (pool_state['snapshot_interval'] > 0
+            and pufferl.global_step - pool_state['last_snapshot_step']
+                >= pool_state['snapshot_interval']):
+        snap_path = os.path.join(pool_state['pool_dir'],
+            f'{pufferl.global_step:016d}.bin')
+        backend.save_weights(pufferl, snap_path)
+        pool_state['pool'].append({'path': snap_path, 'elo': pool_state['primary_elo']})
+        pool_state['pool'] = evict(pool_state['pool'], pool_state['max_size'])
+        pool_state['last_snapshot_step'] = int(pufferl.global_step)
+
+    # 3. Per-bank swap logic. Each bank decides independently based on its own
+    # winrate. Tags 1..num_banks correspond to bank 0..num_banks-1.
+    for b in range(num_banks):
+        bank = pool_state['banks'][b]
+        winrate = (bank['hist_score'] / bank['hist_n']
+                       if bank['hist_n'] > 0 else None)
+        winrate_met = (winrate is not None
+            and bank['hist_n'] >= pool_state['min_games']
+            and winrate >= pool_state['swap_winrate'])
+        timed_out = (pool_state['opp_timeout_steps'] > 0
+            and pufferl.global_step - bank['opp_started_step']
+                >= pool_state['opp_timeout_steps'])
+        tag_value = b + 1
+
+        if bank['pending_opp_path'] is not None:
+            if backend.count_aligned(pufferl, tag_value, 0) >= bank['num_hist_envs']:
+                backend.load_frozen_bank(pufferl, b, bank['pending_opp_path'])
+                backend.count_aligned(pufferl, tag_value, 1)
+                bank['cur_opp_path'] = bank['pending_opp_path']
+                bank['cur_opp_elo'] = bank['pending_opp_elo']
+                bank['pending_opp_path'] = None
+                bank['pending_opp_elo'] = None
+                bank['hist_score'] = 0.0
+                bank['hist_n'] = 0.0
+                bank['opp_started_step'] = int(pufferl.global_step)
+                bank['last_epochs_to_align'] = epoch - bank['epoch_armed']
+        elif winrate_met or timed_out:
+            # Winrate-driven snapshot kept while pool is small (< 10). After
+            # that, only the global interval cadence grows the pool. Timeout
+            # swaps never snapshot (stalemate, not progress).
+            if winrate_met and len(pool_state['pool']) < 10:
+                snap_path = os.path.join(pool_state['pool_dir'],
+                    f'{pufferl.global_step:016d}.bin')
+                backend.save_weights(pufferl, snap_path)
+                pool_state['pool'].append({'path': snap_path, 'elo': pool_state['primary_elo']})
+                pool_state['pool'] = evict(pool_state['pool'], pool_state['max_size'])
+                pool_state['last_snapshot_step'] = int(pufferl.global_step)
+            opp_entry = sample_opponent(pool_state['pool'], pool_state['rng'])
+            bank['pending_opp_path'] = opp_entry['path']
+            bank['pending_opp_elo'] = opp_entry['elo']
+            bank['epoch_armed'] = epoch
+            bank['last_winrate_at_swap'] = winrate if winrate is not None else 0.0
+
+    # 4. Emit logs — per-bank and aggregate.
+    flat_logs['pool/size']     = len(pool_state['pool'])
+    flat_logs['env/elo']       = pool_state['primary_elo']
+    flat_logs['pool/num_banks'] = num_banks
+    total_score = 0.0
+    total_n     = 0.0
+    for b in range(num_banks):
+        bank = pool_state['banks'][b]
+        wr = (bank['hist_score'] / bank['hist_n']
+              if bank['hist_n'] > 0 else None)
+        flat_logs[f'pool/winrate_at_swap_bank_{b}'] = bank['last_winrate_at_swap']
+        flat_logs[f'pool/epochs_to_align_bank_{b}'] = bank['last_epochs_to_align']
+        if wr is not None:
+            flat_logs[f'pool/winrate_bank_{b}']           = wr
+            flat_logs[f'env/historical_winrate_bank_{b}'] = wr
+        total_score += bank['hist_score']
+        total_n     += bank['hist_n']
+    # Aggregate winrate across all banks (legacy compat with old dashboards).
+    if total_n > 0:
+        agg = total_score / total_n
+        flat_logs['pool/winrate']           = agg
+        flat_logs['env/historical_winrate'] = agg
diff --git a/pufferlib/sweep.py b/pufferlib/sweep.py
index 1212d8e1ec..36e27bf42a 100644
--- a/pufferlib/sweep.py
+++ b/pufferlib/sweep.py
@@ -618,6 +618,17 @@ def __init__(self,
             self.gp_cost_buffer = torch.empty(self.gp_max_obs, device=self.device)
             self.infer_batch_buffer = torch.empty(self.infer_batch_size, self.hyperparameters.num, device=self.device)
 
+    def to(self, device):
+        self.device = torch.device(device)
+        for attr in ('gp_score', 'gp_cost', 'likelihood_score', 'likelihood_cost',
+                     'mll_score', 'mll_cost', 'gp_params_buffer', 'gp_score_buffer',
+                     'gp_cost_buffer', 'infer_batch_buffer'):
+            setattr(self, attr, getattr(self, attr).to(self.device))
+        for opt in (self.score_opt, self.cost_opt):
+            for state in opt.state.values():
+                state.update({k: v.to(self.device) for k, v in state.items() if isinstance(v, torch.Tensor)})
+        return self
+
     def _filter_near_duplicates(self, inputs, duplicate_threshold=EPSILON):
         if len(inputs) < 2:
             return np.arange(len(inputs))
diff --git a/src/bindings.cu b/src/bindings.cu
index 4469cb512c..64be61194d 100644
--- a/src/bindings.cu
+++ b/src/bindings.cu
@@ -2,6 +2,7 @@
 
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <pybind11/numpy.h>
 #include "pufferlib.cu"
 
 #define _PUFFER_STRINGIFY(x) #x
@@ -106,7 +107,9 @@ pybind11::dict puf_eval_log(pybind11::object pufferl_obj) {
     pufferl.last_log_step = pufferl.global_step;
  
     pybind11::dict env_dict;
-    Dict* env_out = create_dict(32);
+    // Capacity 64 to fit chess's per-bank hist_score_bank/hist_n_bank entries
+    // (16 keys across 8 banks) on top of base env-log fields.
+    Dict* env_out = create_dict(64);
     static_vec_eval_log(pufferl.vec, env_out);
     for (int i = 0; i < env_out->size; i++) {
         env_dict[env_out->items[i].key] = env_out->items[i].value;
@@ -134,11 +137,18 @@ void rollouts(pybind11::object pufferl_obj) {
     pybind11::gil_scoped_release no_gil;
     double t0 = wall_clock();
 
-    // Zero state buffers
+    // Zero state buffers (primary + every frozen bank, so all banks see fresh
+    // state symmetrically — otherwise frozen banks accumulate indefinitely while
+    // primary resets, giving primary an unfair in-distribution advantage).
     if (pufferl.hypers.reset_state) {
         for (int i = 0; i < pufferl.hypers.num_buffers; i++) {
             puf_zero(&pufferl.buffer_states[i], pufferl.default_stream);
         }
+        for (int b = 0; b < pufferl.num_frozen_banks; b++) {
+            for (int i = 0; i < pufferl.hypers.num_buffers; i++) {
+                puf_zero(&pufferl.frozen_banks[b].buffer_states[i], pufferl.default_stream);
+            }
+        }
     }
 
     static_vec_omp_step(pufferl.vec);
@@ -207,6 +217,48 @@ void load_weights(pybind11::object pufferl_obj, const std::string& path) {
     }
 }
 
+int py_add_frozen_bank(py::object pufferl_obj, int slice_size,
+                       int hidden_size, int num_layers) {
+    PuffeRL& pufferl = pufferl_obj.cast<PuffeRL&>();
+    return pufferl_add_frozen_bank(&pufferl, slice_size, hidden_size, num_layers);
+}
+
+void py_load_frozen_bank(py::object pufferl_obj, int bank_idx, const std::string& path) {
+    PuffeRL& pufferl = pufferl_obj.cast<PuffeRL&>();
+    pufferl_load_frozen_bank(&pufferl, bank_idx, path.c_str());
+}
+
+void py_set_agent_perm(py::object pufferl_obj, py::array_t<int> perm) {
+    PuffeRL& pufferl = pufferl_obj.cast<PuffeRL&>();
+    auto buf = perm.request();
+    if (buf.ndim != 1) throw std::runtime_error("agent_perm must be 1-D");
+    if ((int)buf.shape[0] != pufferl.vec->total_agents) {
+        throw std::runtime_error("agent_perm length must equal total_agents");
+    }
+    pufferl_set_agent_perm(&pufferl, (const int*)buf.ptr);
+}
+
+void py_set_env_tags(py::object pufferl_obj, py::array_t<int> tags) {
+    PuffeRL& pufferl = pufferl_obj.cast<PuffeRL&>();
+    auto buf = tags.request();
+    if (buf.ndim != 1) throw std::runtime_error("env_tags must be 1-D");
+    int num_envs = pufferl_num_envs(&pufferl);
+    if ((int)buf.shape[0] != num_envs) {
+        throw std::runtime_error("env_tags length must equal num_envs");
+    }
+    pufferl_set_env_tags(&pufferl, (const int*)buf.ptr);
+}
+
+int py_count_aligned(py::object pufferl_obj, int tag_value, int reset_flags) {
+    PuffeRL& pufferl = pufferl_obj.cast<PuffeRL&>();
+    return pufferl_count_aligned(&pufferl, tag_value, reset_flags);
+}
+
+int py_num_envs(py::object pufferl_obj) {
+    PuffeRL& pufferl = pufferl_obj.cast<PuffeRL&>();
+    return pufferl_num_envs(&pufferl);
+}
+
 void py_puff_advantage(
         long long values_ptr, long long rewards_ptr,
         long long dones_ptr,  long long importance_ptr,
@@ -367,6 +419,8 @@ std::unique_ptr<PuffeRL> create_pufferl(py::dict args) {
     hypers.vf_clip_coef = get_config(train_kwargs, "vf_clip_coef");
     hypers.vf_coef = get_config(train_kwargs, "vf_coef");
     hypers.ent_coef = get_config(train_kwargs, "ent_coef");
+    hypers.min_ent_coef_ratio = get_config(train_kwargs, "min_ent_coef_ratio");
+    hypers.anneal_ent_coef = get_config(train_kwargs, "anneal_ent_coef");
     // GAE
     hypers.gamma = get_config(train_kwargs, "gamma");
     hypers.gae_lambda = get_config(train_kwargs, "gae_lambda");
@@ -465,6 +519,12 @@ PYBIND11_MODULE(_C, m) {
     m.def("close", &puf_close);
     m.def("save_weights", &save_weights);
     m.def("load_weights", &load_weights);
+    m.def("add_frozen_bank", &py_add_frozen_bank);
+    m.def("load_frozen_bank", &py_load_frozen_bank);
+    m.def("set_agent_perm", &py_set_agent_perm);
+    m.def("set_env_tags", &py_set_env_tags);
+    m.def("count_aligned", &py_count_aligned);
+    m.def("num_envs", &py_num_envs);
     m.def("python_vec_recv", &python_vec_recv);
     m.def("python_vec_send", &python_vec_send);
     py::class_<Policy>(m, "Policy");
@@ -493,6 +553,8 @@ PYBIND11_MODULE(_C, m) {
         .def_readwrite("vf_clip_coef", &HypersT::vf_clip_coef)
         .def_readwrite("vf_coef", &HypersT::vf_coef)
         .def_readwrite("ent_coef", &HypersT::ent_coef)
+        .def_readwrite("min_ent_coef_ratio", &HypersT::min_ent_coef_ratio)
+        .def_readwrite("anneal_ent_coef", &HypersT::anneal_ent_coef)
         .def_readwrite("gamma", &HypersT::gamma)
         .def_readwrite("gae_lambda", &HypersT::gae_lambda)
         .def_readwrite("vtrace_rho_clip", &HypersT::vtrace_rho_clip)
diff --git a/src/pufferlib.cu b/src/pufferlib.cu
index 3a3e6ee00e..0014999e64 100644
--- a/src/pufferlib.cu
+++ b/src/pufferlib.cu
@@ -3,6 +3,7 @@
 #include <nvtx3/nvToolsExt.h>
 #include <nvml.h>
 #include <nccl.h>
+#include <vector>
 
 #include <time.h>
 #include "models.cu"
@@ -56,11 +57,13 @@ struct RolloutBuf {
     PrecisionTensor terminals;
     PrecisionTensor ratio;
     PrecisionTensor importance;
+    PrecisionTensor action_mask;   // (horizon, agents, mask_size); .data=nullptr when env opts out
 };
 
 // Buffers are initialized as raw structs with only shape information. alloc_register
 // stores the shape and data pointer. Memory is only allocated after all buffers are registered.
-void register_rollout_buffers(RolloutBuf& bufs, Allocator* alloc, int T, int B, int input_size, int num_atns) {
+void register_rollout_buffers(RolloutBuf& bufs, Allocator* alloc, int T, int B, int input_size,
+        int num_atns, int mask_size) {
     bufs = (RolloutBuf){
         .observations = {.shape = {T, B, input_size}},
         .actions      = {.shape = {T, B, num_atns}},
@@ -70,6 +73,7 @@ void register_rollout_buffers(RolloutBuf& bufs, Allocator* alloc, int T, int B,
         .terminals    = {.shape = {T, B}},
         .ratio        = {.shape = {T, B}},
         .importance   = {.shape = {T, B}},
+        .action_mask  = {},
     };
     alloc_register(alloc, &bufs.observations);
     alloc_register(alloc, &bufs.actions);
@@ -79,6 +83,10 @@ void register_rollout_buffers(RolloutBuf& bufs, Allocator* alloc, int T, int B,
     alloc_register(alloc, &bufs.terminals);
     alloc_register(alloc, &bufs.ratio);
     alloc_register(alloc, &bufs.importance);
+    if (mask_size > 0) {
+        bufs.action_mask = {.shape = {T, B, mask_size}};
+        alloc_register(alloc, &bufs.action_mask);
+    }
 }
 
 // Train data layout is transposed to (B, T) from rollouts layout (T, B)
@@ -95,10 +103,11 @@ struct TrainGraph {
     PrecisionTensor mb_ratio;
     PrecisionTensor mb_newvalue;
     PrecisionTensor mb_prio;        // (B,)
+    PrecisionTensor mb_action_mask; // (B, T, mask_size); .data=nullptr when disabled
 };
 
 void register_train_buffers(TrainGraph& bufs, Allocator* alloc, int B, int T, int input_size,
-        int hidden_size, int num_atns, int num_layers) {
+        int hidden_size, int num_atns, int num_layers, int mask_size) {
     bufs = (TrainGraph){
         .mb_state =         {.shape = {num_layers, B, hidden_size}},
         .mb_obs =           {.shape = {B, T, input_size}},
@@ -110,6 +119,7 @@ void register_train_buffers(TrainGraph& bufs, Allocator* alloc, int B, int T, in
         .mb_ratio =         {.shape = {B, T}},
         .mb_newvalue =      {.shape = {B, T}},
         .mb_prio =          {.shape = {B}},
+        .mb_action_mask =   {},
     };
     alloc_register(alloc, &bufs.mb_obs);
     alloc_register(alloc, &bufs.mb_state);
@@ -121,6 +131,10 @@ void register_train_buffers(TrainGraph& bufs, Allocator* alloc, int B, int T, in
     alloc_register(alloc, &bufs.mb_returns);
     alloc_register(alloc, &bufs.mb_ratio);
     alloc_register(alloc, &bufs.mb_newvalue);
+    if (mask_size > 0) {
+        bufs.mb_action_mask = {.shape = {B, T, mask_size}};
+        alloc_register(alloc, &bufs.mb_action_mask);
+    }
 }
 
 // PPO buffers + args are quite complex. We do the entire
@@ -146,6 +160,8 @@ struct PPOKernelArgs {
     const float* adv_mean;
     const float* adv_var;
     const int* act_sizes;
+    const precision_t* action_mask; // (N, T, A_total) or nullptr
+    int mask_stride_n, mask_stride_t;
     int num_atns;
     float clip_coef, vf_clip_coef, vf_coef, ent_coef;
     int T_seq, A_total, N;
@@ -219,6 +235,7 @@ struct EnvBuf {
     FloatTensor actions;   // (total_agents, num_atns)
     FloatTensor rewards;   // (total_agents,)
     FloatTensor terminals; // (total_agents,)
+    ByteTensor action_mask; // (total_agents, mask_size); .data=nullptr when env opts out
 };
 
 StaticVec* create_environments(int num_buffers, int total_agents,
@@ -231,6 +248,12 @@ StaticVec* create_environments(int num_buffers, int total_agents,
     env.actions = { .data = (float*)vec->gpu_actions, .shape = {total_agents, get_num_atns()} };
     env.rewards = { .data = (float*)vec->gpu_rewards, .shape = {total_agents} };
     env.terminals = { .data = (float*)vec->gpu_terminals, .shape = {total_agents} };
+    if (vec->action_mask_size > 0) {
+        env.action_mask = { .data = vec->gpu_action_mask,
+                            .shape = {total_agents, vec->action_mask_size} };
+    } else {
+        env.action_mask = { .data = nullptr, .shape = {0} };
+    }
     return vec;
 }
 
@@ -261,6 +284,11 @@ typedef struct {
     float vf_clip_coef;
     float vf_coef;
     float ent_coef;
+    // Entropy coefficient anneal — mirrors lr annealing. When anneal_ent_coef
+    // is set, ent_coef cosine-decays from its base value to
+    // min_ent_coef_ratio * ent_coef over total_timesteps.
+    float min_ent_coef_ratio;
+    bool anneal_ent_coef;
     // GAE
     float gamma;
     float gae_lambda;
@@ -284,6 +312,23 @@ typedef struct {
     int seed;
 } HypersT;
 
+// A frozen weight bank: same shape as the primary, but its own params buffer
+// (and per-buffer rollout states/activations). Used for match (eval) and league
+// (frozen historical opponents). Not trained; updated only via load.
+typedef struct {
+    Policy policy;  // Bank-owned Policy; lets banks have different arch than primary.
+    PolicyWeights weights;
+    Allocator params_alloc;
+    Allocator acts_alloc;
+    PrecisionTensor param_puf;
+    FloatTensor master_weights;
+    PrecisionTensor* buffer_states;         // [num_buffers]
+    PolicyActivations* buffer_activations;  // [num_buffers]
+    int slice_size;  // # agents per buffer this bank owns; sets activation/state batch dim
+    int hidden_size;
+    int num_layers;
+} WeightBank;
+
 typedef struct {
     Policy policy;
     PolicyWeights weights;       // current precision_t weights (structured)
@@ -327,10 +372,22 @@ typedef struct {
     bool train_captured;
     ulong seed;
     curandStatePhilox4_32_10_t** rng_states;  // per-buffer persistent RNG states [num_buffers]
+    // Optional frozen weight banks for match / league.
+    WeightBank* frozen_banks;  // [num_frozen_banks]
+    int num_frozen_banks;
+    std::string env_name;  // Kept for post-init bank adds (needs create_custom_encoder).
+    // Per-buffer-relative bank layout: bank_layout[b] = first agent within each
+    // buffer chunk owned by bank b. Length num_banks+1; ends at agents_per_buffer.
+    // Same shape applied to every buffer (each buffer hosts every bank), so each
+    // worker thread only writes inside its own physical chunk.
+    // Bank 0 = primary (learner). NULL = no layout set (primary owns full chunk).
+    int* bank_layout;
 } PuffeRL;
 
 Dict* log_environments_impl(PuffeRL& pufferl) {
-    Dict* out = create_dict(32);
+    // Capacity raised from 32 to 64 to accommodate chess's per-bank
+    // hist_score_bank_<b> / hist_n_bank_<b> entries (16 keys for 8 banks).
+    Dict* out = create_dict(64);
     static_vec_log(pufferl.vec, out);
     return out;
 }
@@ -371,6 +428,17 @@ __device__ __forceinline__ float safe_logit(const precision_t* logits,
     return l;
 }
 
+__device__ __forceinline__ float masked_logit(const precision_t* logits,
+        int logits_base, int logits_offset, int offset,
+        const precision_t* mask, int mask_base) {
+    float l = safe_logit(logits, logits_base, logits_offset, offset);
+    if (mask != nullptr) {
+        float m = to_float(mask[mask_base + logits_offset + offset]);
+        if (m == 0.0f) l = -1e4f;
+    }
+    return l;
+}
+
 // Expects action logits and values to be in the same contiguous buffer. See default decoder
 __global__ void sample_logits(
         PrecisionTensor dec_out,              // (B, logits_dim + 1 for values)
@@ -379,7 +447,9 @@ __global__ void sample_logits(
         precision_t* __restrict__ actions,    // (B, num_atns)
         precision_t* __restrict__ logprobs,   // (B,)
         precision_t* __restrict__ value_out,  // (B,)
-        curandStatePhilox4_32_10_t* __restrict__ rng_states) {
+        curandStatePhilox4_32_10_t* __restrict__ rng_states,
+        const precision_t* __restrict__ action_mask, // (B, A_total) or nullptr
+        int mask_stride) {                    // 0 when action_mask is nullptr
     int B = dec_out.shape[0];
     int fused_cols = dec_out.shape[1];
     int num_atns = numel(act_sizes_puf.shape);
@@ -427,6 +497,7 @@ __global__ void sample_logits(
     } else {
         // Discrete action sampling (original multinomial logic)
         int logits_offset = 0;  // offset within row for current action head
+        int mask_base = (action_mask != nullptr) ? idx * mask_stride : 0;
 
         for (int h = 0; h < num_atns; ++h) {
             int A = act_sizes[h];  // size of this action head
@@ -435,7 +506,7 @@ __global__ void sample_logits(
             float max_val = -INFINITY;
             float sum_exp = 0.0f;
             for (int a = 0; a < A; ++a) {
-                float l = safe_logit(logits, logits_base, logits_offset, a);
+                float l = masked_logit(logits, logits_base, logits_offset, a, action_mask, mask_base);
                 if (l > max_val) {
                     sum_exp *= expf(max_val - l);
                     max_val = l;
@@ -449,10 +520,10 @@ __global__ void sample_logits(
 
             // Step 4: Multinomial sampling using inverse CDF
             float cumsum = 0.0f;
-            int sampled_action = A - 1;  // default to last action
+            int sampled_action = -1;  // sentinel: no action chosen yet
 
             for (int a = 0; a < A; ++a) {
-                float l = safe_logit(logits, logits_base, logits_offset, a);
+                float l = masked_logit(logits, logits_base, logits_offset, a, action_mask, mask_base);
                 float prob = expf(l - logsumexp);
                 cumsum += prob;
                 if (rand_val < cumsum) {
@@ -461,8 +532,21 @@ __global__ void sample_logits(
                 }
             }
 
+            // Float rounding can leave cumsum < 1.0; fall back to the last legal action.
+            if (sampled_action < 0) {
+                sampled_action = A - 1;
+                if (action_mask != nullptr) {
+                    for (int a = A - 1; a >= 0; --a) {
+                        if (to_float(action_mask[mask_base + logits_offset + a]) != 0.0f) {
+                            sampled_action = a;
+                            break;
+                        }
+                    }
+                }
+            }
+
             // Step 5: Gather log probability of sampled action
-            float sampled_logit = safe_logit(logits, logits_base, logits_offset, sampled_action);
+            float sampled_logit = masked_logit(logits, logits_base, logits_offset, sampled_action, action_mask, mask_base);
             float log_prob = sampled_logit - logsumexp;
 
             // Write action for this head
@@ -528,29 +612,80 @@ extern "C" void net_callback_wrapper(void* ctx, int buf, int t) {
     cast<<<grid_size(n), BLOCK_SIZE, 0, stream>>>(
         term_dst.data, env.terminals.data + start, n);
 
-    // Policy forward pass for rollouts
-    PrecisionTensor state_puf = pufferl->buffer_states[buf];
-    PrecisionTensor dec_puf = policy_forward(&pufferl->policy, pufferl->weights, pufferl->buffer_activations[buf], obs_dst, state_puf, stream);
+    // Copy action mask from env into rollout buffer (if env opted in)
+    PrecisionTensor mask_slice = {};
+    int mask_stride = 0;
+    if (rollouts.action_mask.data != nullptr) {
+        int mask_size = rollouts.action_mask.shape[2];
+        mask_stride = mask_size;
+        mask_slice = puf_slice(rollouts.action_mask, t, start, block_size);
+        int mask_n = block_size * mask_size;
+        cast<<<grid_size(mask_n), BLOCK_SIZE, 0, stream>>>(
+            mask_slice.data,
+            env.action_mask.data + (long)start * mask_size,
+            mask_n);
+    }
+
+    // Per-bank policy forward + sampling. Each bank owns a contiguous sub-range
+    // [bank_layout[b], bank_layout[b+1]) within every buffer's chunk; layout is
+    // per-buffer-relative so each worker writes only inside its own chunk.
+    // Cudagraph capture absorbs the extra kernel launches.
+    int num_banks = 1 + pufferl->num_frozen_banks;
+    long act_cols = env.actions.shape[1];
+    for (int b = 0; b < num_banks; b++) {
+        int bank_off = pufferl->bank_layout ? pufferl->bank_layout[b] : 0;
+        int bank_end = pufferl->bank_layout ? pufferl->bank_layout[b + 1] : block_size;
+        int bank_size = bank_end - bank_off;
+        if (bank_size == 0) continue;
+
+        Policy* p_bank;
+        PolicyWeights* w_bank;
+        PolicyActivations* a_bank;
+        PrecisionTensor* s_bank;
+        if (b == 0) {
+            p_bank = &pufferl->policy;
+            w_bank = &pufferl->weights;
+            a_bank = &pufferl->buffer_activations[buf];
+            s_bank = &pufferl->buffer_states[buf];
+        } else {
+            WeightBank* fb = &pufferl->frozen_banks[b - 1];
+            p_bank = &fb->policy;
+            w_bank = &fb->weights;
+            a_bank = &fb->buffer_activations[buf];
+            s_bank = &fb->buffer_states[buf];
+        }
 
-    // Sample actions, logprobs, values into rollout buffer
-    PrecisionTensor act_slice = puf_slice(rollouts.actions, t, start, block_size);
-    PrecisionTensor lp_slice = puf_slice(rollouts.logprobs, t, start, block_size);
-    PrecisionTensor val_slice = puf_slice(rollouts.values, t, start, block_size);
-    PrecisionTensor p_logstd = {};
-    DecoderWeights* dw = (DecoderWeights*)pufferl->weights.decoder;
-    if (dw->continuous) {
-        p_logstd = dw->logstd;
-    }
+        int sub_start = start + bank_off;
+        PrecisionTensor obs_b   = puf_slice(rollouts.observations, t, sub_start, bank_size);
+        PrecisionTensor act_b   = puf_slice(rollouts.actions,      t, sub_start, bank_size);
+        PrecisionTensor lp_b    = puf_slice(rollouts.logprobs,     t, sub_start, bank_size);
+        PrecisionTensor val_b   = puf_slice(rollouts.values,       t, sub_start, bank_size);
+        PrecisionTensor mask_b  = {};
+        int mask_stride_b = 0;
+        if (rollouts.action_mask.data != nullptr) {
+            mask_b = puf_slice(rollouts.action_mask, t, sub_start, bank_size);
+            mask_stride_b = mask_stride;
+        }
 
-    sample_logits<<<grid_size(block_size), BLOCK_SIZE, 0, stream>>>(
-        dec_puf, p_logstd, pufferl->act_sizes_puf,
-        act_slice.data, lp_slice.data, val_slice.data,
-        pufferl->rng_states[buf]);
+        PrecisionTensor dec_puf = policy_forward(p_bank, *w_bank, *a_bank, obs_b, *s_bank, stream);
 
-    // Copy actions to env
-    long act_cols = env.actions.shape[1];
-    cast<<<grid_size(numel(act_slice.shape)), BLOCK_SIZE, 0, stream>>>(
-            env.actions.data + start * act_cols, act_slice.data, numel(act_slice.shape));
+        PrecisionTensor p_logstd = {};
+        DecoderWeights* dw = (DecoderWeights*)w_bank->decoder;
+        if (dw->continuous) {
+            p_logstd = dw->logstd;
+        }
+
+        // Offset RNG by bank_off so banks don't collide on per-buffer rng slots.
+        sample_logits<<<grid_size(bank_size), BLOCK_SIZE, 0, stream>>>(
+            dec_puf, p_logstd, pufferl->act_sizes_puf,
+            act_b.data, lp_b.data, val_b.data,
+            pufferl->rng_states[buf] + bank_off,
+            mask_b.data, mask_stride_b);
+
+        cast<<<grid_size(numel(act_b.shape)), BLOCK_SIZE, 0, stream>>>(
+                env.actions.data + (long)sub_start * act_cols,
+                act_b.data, numel(act_b.shape));
+    }
 
     if (capturing) {
         cudaGraph_t _graph;
@@ -565,16 +700,32 @@ extern "C" void net_callback_wrapper(void* ctx, int buf, int t) {
 }
 
 
+__device__ __forceinline__ float load_logit_masked(
+        const precision_t* __restrict__ logits, int logits_base,
+        int logits_stride_a, int logits_offset, int a,
+        const precision_t* __restrict__ mask, int mask_base) {
+    float l = to_float(logits[logits_base + (logits_offset + a) * logits_stride_a]);
+    if (mask != nullptr) {
+        float m = to_float(mask[mask_base + logits_offset + a]);
+        if (m == 0.0f) {
+            l = -1e4f;
+            return l;
+        }
+    }
+    return l;
+}
+
 __device__ __forceinline__ void ppo_discrete_head(
         const precision_t* __restrict__ logits, int logits_base,
         int logits_stride_a, int logits_offset, int A, int act,
+        const precision_t* __restrict__ mask, int mask_base,
         float* out_logsumexp, float* out_entropy, float* out_logp) {
     float max_logit = -INFINITY;
     float sum = 0.0f;
     float act_logit = 0.0f;
 
     for (int a = 0; a < A; ++a) {
-        float l = to_float(logits[logits_base + (logits_offset + a) * logits_stride_a]);
+        float l = load_logit_masked(logits, logits_base, logits_stride_a, logits_offset, a, mask, mask_base);
         if (a == act) {
             act_logit = l;
         }
@@ -588,7 +739,7 @@ __device__ __forceinline__ void ppo_discrete_head(
 
     float ent = 0.0f;
     for (int a = 0; a < A; ++a) {
-        float l = to_float(logits[logits_base + (logits_offset + a) * logits_stride_a]);
+        float l = load_logit_masked(logits, logits_base, logits_stride_a, logits_offset, a, mask, mask_base);
         float logp = l - logsumexp;
         float p = __expf(logp);
         ent -= p * logp;
@@ -685,6 +836,9 @@ __global__ void ppo_loss_compute(
     float head_entropy[MAX_ATN_HEADS];
     int head_act[MAX_ATN_HEADS];
 
+    int mask_base = (a.action_mask != nullptr)
+        ? n * a.mask_stride_n + t * a.mask_stride_t : 0;
+
     if (!a.is_continuous) {
         int logits_offset = 0;
         for (int h = 0; h < a.num_atns; ++h) {
@@ -692,7 +846,8 @@ __global__ void ppo_loss_compute(
             int act = static_cast<int>(g.actions[nt * a.num_atns + h]);
             head_act[h] = act;
             float lse, ent, lp;
-            ppo_discrete_head(a.logits, logits_base, a.logits_stride_a, logits_offset, A, act, &lse, &ent, &lp);
+            ppo_discrete_head(a.logits, logits_base, a.logits_stride_a, logits_offset, A, act,
+                              a.action_mask, mask_base, &lse, &ent, &lp);
             head_logsumexp[h] = lse;
             head_entropy[h] = ent;
             total_log_prob += lp;
@@ -738,7 +893,8 @@ __global__ void ppo_loss_compute(
             float ent = head_entropy[h];
 
             for (int j = 0; j < A; ++j) {
-                float l = to_float(a.logits[logits_base + (logits_offset + j) * a.logits_stride_a]);
+                float l = load_logit_masked(a.logits, logits_base, a.logits_stride_a,
+                                            logits_offset, j, a.action_mask, mask_base);
                 float logp = l - logsumexp;
                 float p = __expf(logp);
                 float d_logit = (j == act) ? d_new_logp : 0.0f;
@@ -909,6 +1065,7 @@ void ppo_loss_fwd_bwd(
         .returns = graph.mb_returns.data,
     };
 
+    bool has_mask = (graph.mb_action_mask.data != nullptr);
     PPOKernelArgs args = {
         .grad_logits = bufs.grad_logits.data,
         .grad_logstd = is_continuous ? bufs.grad_logstd.data : nullptr,
@@ -919,6 +1076,9 @@ void ppo_loss_fwd_bwd(
         .adv_mean = adv_mean_ptr,
         .adv_var = adv_var_ptr,
         .act_sizes = act_sizes.data,
+        .action_mask = has_mask ? graph.mb_action_mask.data : nullptr,
+        .mask_stride_n = has_mask ? T * A_total : 0,
+        .mask_stride_t = has_mask ? A_total : 0,
         .num_atns = (int)numel(act_sizes.shape),
         .clip_coef = clip_coef, .vf_clip_coef = vf_clip_coef,
         .vf_coef = vf_coef, .ent_coef = ent_coef,
@@ -1205,6 +1365,30 @@ void puff_advantage_cuda(PrecisionTensor& values, PrecisionTensor& rewards,
         advantages.data, gamma, lambda, rho_clip, c_clip, num_steps, horizon);
 }
 
+// Zero advantages on frozen-bank rows so prio_replay never samples them. Frozen
+// rollout rows hold actions/logprobs from the frozen policy — training the
+// primary's PPO on them produces garbage ratios and poisoned gradients.
+__global__ void zero_frozen_advantages_kernel(precision_t* advantages,
+        int agents_per_buffer, int primary_per_buffer, int total_rows, int horizon) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = total_rows * horizon;
+    if (idx >= total) return;
+    int row = idx / horizon;
+    int rel = row % agents_per_buffer;
+    if (rel >= primary_per_buffer) {
+        advantages[idx] = from_float(0.0f);
+    }
+}
+
+void zero_frozen_advantages_cuda(PrecisionTensor& advantages,
+        int agents_per_buffer, int primary_per_buffer, cudaStream_t stream) {
+    int total_rows = advantages.shape[0];
+    int horizon = advantages.shape[1];
+    int total = total_rows * horizon;
+    zero_frozen_advantages_kernel<<<grid_size(total), BLOCK_SIZE, 0, stream>>>(
+        advantages.data, agents_per_buffer, primary_per_buffer, total_rows, horizon);
+}
+
 // Minor copy bandwidth optimizations
 __global__ void index_copy(char* __restrict__ dst, const int* __restrict__ idx,
         const char* __restrict__ src, int num_idx, int row_bytes) {
@@ -1267,8 +1451,16 @@ __global__ void select_copy(RolloutBuf rollouts, TrainGraph graph,
     case 4:
         if (threadIdx.x == 0) {
             graph.mb_prio.data[mb] = from_float(mb_prio[mb]);
-            break;
         }
+        break;
+    case 5:
+        if (graph.mb_action_mask.data != nullptr) {
+            int mask_row_bytes = (numel(rollouts.action_mask.shape)
+                / rollouts.action_mask.shape[0]) * sizeof(precision_t);
+            copy_bytes((const char*)rollouts.action_mask.data,
+                       (char*)graph.mb_action_mask.data, src_row, mb, mask_row_bytes);
+        }
+        break;
     }
 }
 
@@ -1309,6 +1501,11 @@ void train_impl(PuffeRL& pufferl) {
         rollouts.ratio.data, src.ratio.data, T, B, 1);
     transpose_102<<<grid_size(T*B), BLOCK_SIZE, 0, train_stream>>>(
         rollouts.values.data, src.values.data, T, B, 1);
+    if (src.action_mask.data != nullptr) {
+        int mask_size = src.action_mask.shape[2];
+        transpose_102<<<grid_size(T*B*mask_size), BLOCK_SIZE, 0, train_stream>>>(
+            rollouts.action_mask.data, src.action_mask.data, T, B, mask_size);
+    }
 
     // We hard-clamp rewards to -1, 1. Our envs are mostly designed to respect this range
     clamp_precision_kernel<<<grid_size(numel(rollouts.rewards.shape)), BLOCK_SIZE, 0, train_stream>>>(
@@ -1335,6 +1532,17 @@ void train_impl(PuffeRL& pufferl) {
         cudaMemcpy(muon->lr_ptr, &lr, sizeof(float), cudaMemcpyHostToDevice);
     }
 
+    // Annealed entropy coefficient — same cosine shape as lr. With PG signal
+    // alive, the entropy bonus that kept early-training exploratory becomes
+    // load-bearing dead weight late in training; cosine-decay frees the policy
+    // to commit harder on what it has already learned.
+    float current_ent_coef = hypers.ent_coef;
+    if (hypers.anneal_ent_coef) {
+        float ent_min = hypers.min_ent_coef_ratio * hypers.ent_coef;
+        current_ent_coef = cosine_annealing(hypers.ent_coef, ent_min,
+                                            current_epoch, total_epochs);
+    }
+
     // Annealed priority exponent
     float anneal_beta = prio_beta0 + (1.0f - prio_beta0) * prio_alpha * (float)current_epoch/(float)total_epochs;
     TrainGraph& graph = pufferl.train_buf;
@@ -1349,6 +1557,11 @@ void train_impl(PuffeRL& pufferl) {
         puff_advantage_cuda(rollouts.values, rollouts.rewards, rollouts.terminals,
             rollouts.ratio, advantages_puf, hypers.gamma, hypers.gae_lambda,
             hypers.vtrace_rho_clip, hypers.vtrace_c_clip, train_stream);
+        if (pufferl.num_frozen_banks > 0 && pufferl.bank_layout != NULL) {
+            int apb = hypers.total_agents / hypers.num_buffers;
+            zero_frozen_advantages_cuda(advantages_puf, apb,
+                pufferl.bank_layout[1], train_stream);
+        }
         profile_end(hypers.profile);
 
         profile_begin("compute_prio", hypers.profile);
@@ -1365,7 +1578,8 @@ void train_impl(PuffeRL& pufferl) {
             RolloutBuf sel_src = rollouts;
             sel_src.values = rollouts.values;
             int mb_segs = pufferl.prio_bufs.idx.shape[0];
-            select_copy<<<dim3(mb_segs, 5), SELECT_COPY_THREADS, 0, train_stream>>>(
+            int channels = (graph.mb_action_mask.data != nullptr) ? 6 : 5;
+            select_copy<<<dim3(mb_segs, channels), SELECT_COPY_THREADS, 0, train_stream>>>(
                 sel_src, graph, pufferl.prio_bufs.idx.data,
                 advantages_puf.data, pufferl.prio_bufs.mb_prio.data);
         }
@@ -1394,7 +1608,7 @@ void train_impl(PuffeRL& pufferl) {
 
             ppo_loss_fwd_bwd(dec_puf, p_logstd, graph,
                 pufferl.act_sizes_puf, pufferl.losses_puf,
-                hypers.clip_coef, hypers.vf_clip_coef, hypers.vf_coef, hypers.ent_coef,
+                hypers.clip_coef, hypers.vf_clip_coef, hypers.vf_coef, current_ent_coef,
                 pufferl.ppo_bufs_puf, pufferl.is_continuous, stream);
 
             FloatTensor grad_logits_puf = pufferl.ppo_bufs_puf.grad_logits;
@@ -1460,12 +1674,243 @@ void train_impl(PuffeRL& pufferl) {
 
 }
 
+// Build a Policy value for a given env + arch. Encoder/decoder algorithms are
+// fixed by the env; hidden_size/num_layers/horizon parameterize shape. Policy
+// has no heap state so this returns by value; callers store it wherever.
+static Policy build_policy(const char* env_name, int input_size, int hidden_size,
+                           int num_layers, int decoder_output_size, int act_n,
+                           bool is_continuous, int horizon) {
+    Encoder encoder = {
+        .forward = encoder_forward,
+        .backward = encoder_backward,
+        .init_weights = encoder_init_weights,
+        .reg_params = encoder_reg_params,
+        .reg_train = encoder_reg_train,
+        .reg_rollout = encoder_reg_rollout,
+        .create_weights = encoder_create_weights,
+        .free_weights = encoder_free_weights,
+        .free_activations = encoder_free_activations,
+        .in_dim = input_size, .out_dim = hidden_size,
+        .activation_size = sizeof(EncoderActivations),
+    };
+    create_custom_encoder(env_name, &encoder);
+    Decoder decoder = {
+        .forward = decoder_forward,
+        .backward = decoder_backward,
+        .init_weights = decoder_init_weights,
+        .reg_params = decoder_reg_params,
+        .reg_train = decoder_reg_train,
+        .reg_rollout = decoder_reg_rollout,
+        .create_weights = decoder_create_weights,
+        .free_weights = decoder_free_weights,
+        .free_activations = decoder_free_activations,
+        .hidden_dim = hidden_size, .output_dim = decoder_output_size, .continuous = is_continuous,
+    };
+    Network network = {
+        .forward = mingru_forward,
+        .forward_train = mingru_forward_train,
+        .backward = mingru_backward,
+        .init_weights = mingru_init_weights,
+        .reg_params = mingru_reg_params,
+        .reg_train = mingru_reg_train,
+        .reg_rollout = mingru_reg_rollout,
+        .create_weights = mingru_create_weights,
+        .free_weights = mingru_free_weights,
+        .free_activations = mingru_free_activations,
+        .hidden = hidden_size, .num_layers = num_layers, .horizon = horizon,
+    };
+    return Policy{
+        .encoder = encoder, .decoder = decoder, .network = network,
+        .input_dim = input_size, .hidden_dim = hidden_size, .output_dim = decoder_output_size,
+        .num_atns = act_n,
+    };
+}
+
+// Allocate a fresh frozen WeightBank with its own Policy (may differ in
+// hidden_size/num_layers from primary). slice_size = how many agents per buffer
+// this bank will own. Weights are uninitialized — caller must load before use.
+static void weight_bank_create_for_pufferl(WeightBank* bank, PuffeRL* pufferl,
+        int slice_size, int hidden_size, int num_layers) {
+    int num_buffers = pufferl->hypers.num_buffers;
+
+    // Rebuild arch-varying Policy from env metadata already on pufferl.
+    int input_size = pufferl->env.obs.shape[1];
+    int num_action_heads = pufferl->env.actions.shape[1];
+    int* raw_act_sizes = get_act_sizes();
+    int act_n = 0;
+    for (int i = 0; i < num_action_heads; i++) act_n += raw_act_sizes[i];
+    int decoder_output_size = pufferl->is_continuous ? num_action_heads : act_n;
+    bank->policy = build_policy(pufferl->env_name.c_str(), input_size, hidden_size,
+        num_layers, decoder_output_size, act_n, pufferl->is_continuous, pufferl->hypers.horizon);
+    bank->hidden_size = hidden_size;
+    bank->num_layers = num_layers;
+
+    Allocator* params = &bank->params_alloc;
+    Allocator* acts = &bank->acts_alloc;
+
+    bank->slice_size = slice_size;
+    bank->weights = policy_weights_create(&bank->policy, params);
+    bank->buffer_activations = (PolicyActivations*)calloc(num_buffers, sizeof(PolicyActivations));
+    bank->buffer_states = (PrecisionTensor*)calloc(num_buffers, sizeof(PrecisionTensor));
+    for (int i = 0; i < num_buffers; i++) {
+        bank->buffer_activations[i] = policy_reg_rollout(&bank->policy, bank->weights, acts, slice_size);
+        bank->buffer_states[i] = {.shape = {num_layers, slice_size, hidden_size}};
+        alloc_register(acts, &bank->buffer_states[i]);
+    }
+
+    alloc_create(params);
+    alloc_create(acts);
+
+    bank->param_puf = {.data = (precision_t*)params->mem, .shape = {params->total_elems}};
+    if (USE_BF16) {
+        bank->master_weights = {.shape = {params->total_elems}};
+        cudaMalloc(&bank->master_weights.data, params->total_elems * sizeof(float));
+    } else {
+        bank->master_weights = {.data = (float*)bank->param_puf.data, .shape = {params->total_elems}};
+    }
+}
+
+// Mirror of weight_bank_create_for_pufferl. Frees the bank's weights, per-buffer
+// activations, allocators, and master_weights (BF16 only). Does not free the
+// WeightBank struct itself — caller owns that.
+static void weight_bank_destroy(WeightBank* bank, PuffeRL* pufferl) {
+    int num_buffers = pufferl->hypers.num_buffers;
+    policy_weights_free(&bank->policy, &bank->weights);
+    if (bank->buffer_activations != NULL) {
+        for (int i = 0; i < num_buffers; i++) {
+            policy_activations_free(&bank->policy, bank->buffer_activations[i]);
+        }
+        free(bank->buffer_activations);
+    }
+    free(bank->buffer_states);
+    alloc_free(&bank->params_alloc);
+    alloc_free(&bank->acts_alloc);
+    if (USE_BF16 && bank->master_weights.data != NULL) {
+        cudaFree(bank->master_weights.data);
+    }
+}
+
+// Append a fresh frozen bank with the given per-buffer slice size; returns its
+// index. Rebuilds bank_layout sequentially (primary first, then frozen banks in
+// add order). Must be called BEFORE cudagraph capture (pointers get baked in).
+extern "C" int pufferl_add_frozen_bank(PuffeRL* pufferl, int slice_size,
+        int hidden_size, int num_layers) {
+    int idx = pufferl->num_frozen_banks;
+    pufferl->frozen_banks = (WeightBank*)realloc(
+        pufferl->frozen_banks, (idx + 1) * sizeof(WeightBank));
+    memset(&pufferl->frozen_banks[idx], 0, sizeof(WeightBank));
+    weight_bank_create_for_pufferl(&pufferl->frozen_banks[idx], pufferl,
+        slice_size, hidden_size, num_layers);
+    pufferl->num_frozen_banks++;
+
+    // Rebuild sequential layout from declared slice_sizes.
+    int agents_per_buffer = pufferl->vec->total_agents / pufferl->hypers.num_buffers;
+    int frozen_total = 0;
+    for (int b = 0; b < pufferl->num_frozen_banks; b++) {
+        frozen_total += pufferl->frozen_banks[b].slice_size;
+    }
+    if (frozen_total > agents_per_buffer) {
+        fprintf(stderr, "pufferl_add_frozen_bank: total frozen slice (%d) exceeds "
+            "agents_per_buffer (%d)\n", frozen_total, agents_per_buffer);
+    }
+    int num_banks = 1 + pufferl->num_frozen_banks;
+    pufferl->bank_layout = (int*)realloc(pufferl->bank_layout, (num_banks + 1) * sizeof(int));
+    pufferl->bank_layout[0] = 0;
+    pufferl->bank_layout[1] = agents_per_buffer - frozen_total;  // primary
+    int cumul = pufferl->bank_layout[1];
+    for (int b = 0; b < pufferl->num_frozen_banks; b++) {
+        cumul += pufferl->frozen_banks[b].slice_size;
+        pufferl->bank_layout[2 + b] = cumul;
+    }
+    return idx;
+}
+
+// Load a frozen bank's weights from a file (same format as save_weights — flat fp32).
+// Safe to call between rollouts (in-place cudaMemcpy; cudagraphs hold the pointer,
+// not a copy of the data).
+extern "C" void pufferl_load_frozen_bank(PuffeRL* pufferl, int bank_idx, const char* path) {
+    if (bank_idx < 0 || bank_idx >= pufferl->num_frozen_banks) {
+        fprintf(stderr, "pufferl_load_frozen_bank: bank_idx %d out of range\n", bank_idx);
+        return;
+    }
+    WeightBank* bank = &pufferl->frozen_banks[bank_idx];
+    int64_t nbytes = numel(bank->master_weights.shape) * sizeof(float);
+    FILE* f = fopen(path, "rb");
+    if (!f) {
+        fprintf(stderr, "pufferl_load_frozen_bank: failed to open %s\n", path);
+        return;
+    }
+    fseek(f, 0, SEEK_END);
+    long file_size = ftell(f);
+    fseek(f, 0, SEEK_SET);
+    if (file_size != nbytes) {
+        fprintf(stderr, "pufferl_load_frozen_bank: size mismatch (expected %lld, got %ld)\n",
+            (long long)nbytes, file_size);
+        fclose(f);
+        return;
+    }
+    std::vector<char> buf(nbytes);
+    size_t nread = fread(buf.data(), 1, nbytes, f);
+    fclose(f);
+    if ((int64_t)nread != nbytes) {
+        fprintf(stderr, "pufferl_load_frozen_bank: short read on %s\n", path);
+        return;
+    }
+    cudaMemcpy(bank->master_weights.data, buf.data(), nbytes, cudaMemcpyHostToDevice);
+    if (USE_BF16) {
+        int n = numel(bank->param_puf.shape);
+        cast<<<grid_size(n), BLOCK_SIZE, 0, pufferl->default_stream>>>(
+            bank->param_puf.data, bank->master_weights.data, n);
+    }
+    cudaDeviceSynchronize();
+}
+
+// Set the agent permutation. Validates that the perm respects buffer boundaries:
+// each buffer's range [buf_start, buf_start+buf_size) must map onto itself (no
+// cross-buffer writes, since each worker only owns its physical chunk).
+extern "C" void pufferl_set_agent_perm(PuffeRL* pufferl, const int* perm) {
+    int total = pufferl->vec->total_agents;
+    int num_buffers = pufferl->hypers.num_buffers;
+    int buf_size = total / num_buffers;
+    for (int b = 0; b < num_buffers; b++) {
+        int lo = b * buf_size;
+        int hi = lo + buf_size;
+        for (int i = lo; i < hi; i++) {
+            if (perm[i] < lo || perm[i] >= hi) {
+                fprintf(stderr,
+                    "pufferl_set_agent_perm: perm[%d]=%d crosses buffer %d range [%d,%d)\n",
+                    i, perm[i], b, lo, hi);
+                return;
+            }
+        }
+    }
+    static_vec_set_perm(pufferl->vec, perm);
+}
+
+// Set per-env tags (e.g. selfplay vs historical). tags array length must equal
+// pufferl_num_envs(). Also clears each env's boundary_reached flag.
+extern "C" void pufferl_set_env_tags(PuffeRL* pufferl, const int* tags) {
+    static_vec_set_env_tags(pufferl->vec, tags);
+}
+
+// Returns count of envs with tag == tag_value AND boundary_reached. If
+// reset_flags != 0, clears boundary_reached only on envs whose tag matches
+// tag_value (so multi-bank swaps don't trample each other's alignment).
+extern "C" int pufferl_count_aligned(PuffeRL* pufferl, int tag_value, int reset_flags) {
+    return static_vec_count_aligned(pufferl->vec, tag_value, reset_flags);
+}
+
+extern "C" int pufferl_num_envs(PuffeRL* pufferl) {
+    return pufferl->vec->size;
+}
+
 std::unique_ptr<PuffeRL> create_pufferl_impl(HypersT& hypers,
         const std::string& env_name, Dict* vec_kwargs, Dict* env_kwargs) {
     auto pufferl = std::make_unique<PuffeRL>();
     pufferl->hypers = hypers;
     pufferl->nccl_comm = nullptr;
     pufferl->default_stream = 0;
+    pufferl->env_name = env_name;
 
     cudaSetDevice(hypers.gpu_id);
 
@@ -1534,50 +1979,8 @@ std::unique_ptr<PuffeRL> create_pufferl_impl(HypersT& hypers,
     int batch = total_agents / hypers.num_buffers;
     int num_buffers = hypers.num_buffers;
 
-    Encoder encoder = {
-        .forward = encoder_forward,
-        .backward = encoder_backward,
-        .init_weights = encoder_init_weights,
-        .reg_params = encoder_reg_params,
-        .reg_train = encoder_reg_train,
-        .reg_rollout = encoder_reg_rollout,
-        .create_weights = encoder_create_weights,
-        .free_weights = encoder_free_weights,
-        .free_activations = encoder_free_activations,
-        .in_dim = input_size, .out_dim = hidden_size,
-        .activation_size = sizeof(EncoderActivations),
-    };
-    create_custom_encoder(env_name, &encoder);
-    Decoder decoder = {
-        .forward = decoder_forward,
-        .backward = decoder_backward,
-        .init_weights = decoder_init_weights,
-        .reg_params = decoder_reg_params,
-        .reg_train = decoder_reg_train,
-        .reg_rollout = decoder_reg_rollout,
-        .create_weights = decoder_create_weights,
-        .free_weights = decoder_free_weights,
-        .free_activations = decoder_free_activations,
-        .hidden_dim = hidden_size, .output_dim = decoder_output_size, .continuous = is_continuous,
-    };
-    Network network = {
-        .forward = mingru_forward,
-        .forward_train = mingru_forward_train,
-        .backward = mingru_backward,
-        .init_weights = mingru_init_weights,
-        .reg_params = mingru_reg_params,
-        .reg_train = mingru_reg_train,
-        .reg_rollout = mingru_reg_rollout,
-        .create_weights = mingru_create_weights,
-        .free_weights = mingru_free_weights,
-        .free_activations = mingru_free_activations,
-        .hidden = hidden_size, .num_layers = num_layers, .horizon = hypers.horizon,
-    };
-    pufferl->policy = Policy{
-        .encoder = encoder, .decoder = decoder, .network = network,
-        .input_dim = input_size, .hidden_dim = hidden_size, .output_dim = decoder_output_size,
-        .num_atns = act_n,
-    };
+    pufferl->policy = build_policy(env_name.c_str(), input_size, hidden_size,
+        num_layers, decoder_output_size, act_n, is_continuous, hypers.horizon);
 
     // Create and allocate params
     Allocator* params = &pufferl->params_alloc;
@@ -1597,13 +2000,14 @@ std::unique_ptr<PuffeRL> create_pufferl_impl(HypersT& hypers,
         };
         alloc_register(acts, &pufferl->buffer_states[i]);
     }
+    int mask_size = pufferl->vec->action_mask_size;
     register_rollout_buffers(pufferl->rollouts,
-        acts, horizon, total_agents, input_size, num_action_heads);
+        acts, horizon, total_agents, input_size, num_action_heads, mask_size);
     register_train_buffers(pufferl->train_buf,
         acts, minibatch_segments, horizon, input_size,
-        hidden_size, num_action_heads, num_layers);
+        hidden_size, num_action_heads, num_layers, mask_size);
     register_rollout_buffers(pufferl->train_rollouts,
-        acts, total_agents, horizon, input_size, num_action_heads);
+        acts, total_agents, horizon, input_size, num_action_heads, mask_size);
     register_ppo_buffers(pufferl->ppo_bufs_puf,
         acts, minibatch_segments, hypers.horizon, decoder_output_size, is_continuous);
     register_prio_buffers(pufferl->prio_bufs,
@@ -1667,6 +2071,34 @@ std::unique_ptr<PuffeRL> create_pufferl_impl(HypersT& hypers,
     cudaMemcpy(pufferl->ppo_bufs_puf.grad_loss.data, &one, sizeof(float), cudaMemcpyHostToDevice);
     muon_post_create(&pufferl->muon);
 
+    // Set up frozen banks declared in vec_kwargs (num_frozen_banks +
+    // frozen_bank_pct: each bank gets floor(agents_per_buffer * pct) agents).
+    // Must happen BEFORE cudagraph capture so the graph bakes in their pointers
+    // and per-bank loop iterations.
+    DictItem* nb_item = dict_get_unsafe(vec_kwargs, "num_frozen_banks");
+    DictItem* fbp_item = dict_get_unsafe(vec_kwargs, "frozen_bank_pct");
+    DictItem* fbh_item = dict_get_unsafe(vec_kwargs, "frozen_bank_hidden_size");
+    DictItem* fbl_item = dict_get_unsafe(vec_kwargs, "frozen_bank_num_layers");
+    int num_frozen = nb_item ? (int)nb_item->value : 0;
+    float frozen_pct = fbp_item ? (float)fbp_item->value : 0.0f;
+    int frozen_hidden = fbh_item ? (int)fbh_item->value : hidden_size;
+    int frozen_layers = fbl_item ? (int)fbl_item->value : num_layers;
+    if (num_frozen > 0) {
+        int agents_per_buffer = total_agents / num_buffers;
+        int frozen_size = (int)((float)agents_per_buffer * frozen_pct);  // truncates = floor for positive
+        int frozen_total = num_frozen * frozen_size;
+        if (frozen_size <= 0 || frozen_total > agents_per_buffer) {
+            fprintf(stderr, "create_pufferl: invalid frozen bank config "
+                "(num=%d, pct=%.4f -> size=%d, total=%d, agents_per_buffer=%d)\n",
+                num_frozen, frozen_pct, frozen_size, frozen_total, agents_per_buffer);
+            return nullptr;
+        }
+        // add_frozen_bank auto-builds the sequential bank_layout.
+        for (int b = 0; b < num_frozen; b++) {
+            pufferl_add_frozen_bank(pufferl.get(), frozen_size, frozen_hidden, frozen_layers);
+        }
+    }
+
     // Cudagraph rolluts and entire training step
     if (hypers.cudagraphs >= 0) {
         pufferl->fused_rollout_cudagraphs = (cudaGraphExec_t*)calloc(horizon*num_buffers, sizeof(cudaGraphExec_t));
@@ -1809,6 +2241,12 @@ void close_impl(PuffeRL& pufferl) {
     free(pufferl.fused_rollout_cudagraphs);
     free(pufferl.streams);
 
+    for (int b = 0; b < pufferl.num_frozen_banks; b++) {
+        weight_bank_destroy(&pufferl.frozen_banks[b], &pufferl);
+    }
+    free(pufferl.frozen_banks);
+    free(pufferl.bank_layout);
+
     if (pufferl.nccl_comm != nullptr) {
         ncclCommDestroy(pufferl.nccl_comm);
     }
diff --git a/src/vecenv.h b/src/vecenv.h
index e8b20ba686..42958d321e 100644
--- a/src/vecenv.h
+++ b/src/vecenv.h
@@ -82,15 +82,21 @@ typedef struct StaticVec {
     float* actions;
     float* rewards;
     float* terminals;
+    unsigned char* action_mask;  // NULL unless env defines MY_ACTION_MASK
     void* gpu_observations;
     float* gpu_actions;
     float* gpu_rewards;
     float* gpu_terminals;
+    unsigned char* gpu_action_mask;  // NULL unless env defines MY_ACTION_MASK
     cudaStream_t* streams;
     StaticThreading* threading;
     int obs_size;
     int num_atns;
+    int action_mask_size;        // 0 unless env defines MY_ACTION_MASK
     int gpu;
+    // Optional permutation: agent_perm[slot] = physical agent index in global buffers.
+    // NULL = identity (current behavior). Only valid when env defines MY_USES_PERM.
+    int* agent_perm;
 } StaticVec;
 
 // Callback types
@@ -130,6 +136,17 @@ void static_vec_step(StaticVec* vec);
 void gpu_vec_step(StaticVec* vec);
 void cpu_vec_step(StaticVec* vec);
 
+// Optional permutation. Sets agent_perm and re-populates env per-slot pointers
+// via my_setup_perm. Only defined when env opted in via MY_USES_PERM; otherwise
+// emits an error and leaves the perm unset.
+void static_vec_set_perm(StaticVec* vec, const int* perm);
+
+// Optional per-env tagging + boundary tracking for selfplay-pool curricula.
+// Env must opt in via MY_USES_TAGS and provide `int tag` and
+// `int boundary_reached` fields on its Env struct.
+void static_vec_set_env_tags(StaticVec* vec, const int* tags);
+int static_vec_count_aligned(StaticVec* vec, int tag_value, int reset_flags);
+
 // Optional shared state functions
 void* my_shared(void* env, Dict* kwargs);
 void my_shared_close(void* env);
@@ -192,11 +209,13 @@ extern const char* cudaGetErrorString(cudaError_t);
 void my_init(Env* env, Dict* kwargs);
 void my_log(Log* log, Dict* out);
 
-#ifdef MY_VEC_STEP_RANGE
-void MY_VEC_STEP_RANGE(StaticVec* vec, int env_start, int env_count, int num_workers);
+#ifdef MY_USES_PERM
+// Env-provided: populate per-slot pointer arrays on env, given the global slot
+// base for slot 0. Reads vec->agent_perm (NULL = identity) to compute physical
+// indices into vec global buffers.
+void my_setup_perm(StaticVec* vec, Env* env, int slot_base);
 #endif
 
-typedef struct StaticOMPArg StaticOMPArg;
 
 struct StaticThreading {
     atomic_int* buffer_states;
@@ -205,17 +224,16 @@ struct StaticThreading {
     int num_buffers;
     pthread_t* threads;
     float* accum;  // [num_buffers * NUM_EVAL_PROF] per-buffer timing in ms
-    StaticOMPArg* thread_args;
 };
 
-struct StaticOMPArg {
+typedef struct StaticOMPArg {
     StaticVec* vec;
     int buf;
     int horizon;
     void* ctx;
     net_callback_fn net_callback;
     thread_init_fn thread_init;
-};
+} StaticOMPArg;
 
 // OMP thread manager
 static void* static_omp_threadmanager(void* arg) {
@@ -240,6 +258,8 @@ static void* static_omp_threadmanager(void* arg) {
     int num_workers = threading->num_threads / vec->buffers;
     if (num_workers < 1) num_workers = 1;
 
+    Env* envs = (Env*)vec->envs;
+
     printf("Num workers: %d\n", num_workers);
     while (true) {
         while (atomic_load(&buffer_states[buf]) != OMP_RUNNING) {
@@ -268,15 +288,10 @@ static void* static_omp_threadmanager(void* arg) {
             memset(&vec->rewards[agent_start], 0, agents_per_buffer * sizeof(float));
             memset(&vec->terminals[agent_start], 0, agents_per_buffer * sizeof(float));
             clock_gettime(CLOCK_MONOTONIC, &t0);
-            #ifdef MY_VEC_STEP_RANGE
-            MY_VEC_STEP_RANGE(vec, env_start, env_count, num_workers);
-            #else
-            Env* envs = (Env*)vec->envs;
             #pragma omp parallel for schedule(static) num_threads(num_workers)
             for (int i = env_start; i < env_start + env_count; i++) {
                 c_step(&envs[i]);
             }
-            #endif
             clock_gettime(CLOCK_MONOTONIC, &t1);
             my_accum[EVAL_ENV_STEP] += (t1.tv_sec - t0.tv_sec) * 1000.0f + (t1.tv_nsec - t0.tv_nsec) / 1e6f;
 
@@ -295,6 +310,13 @@ static void* static_omp_threadmanager(void* arg) {
                 &vec->terminals[agent_start],
                 agents_per_buffer * sizeof(float),
                 cudaMemcpyHostToDevice, stream);
+#ifdef MY_ACTION_MASK
+            cudaMemcpyAsync(
+                vec->gpu_action_mask + agent_start * MY_ACTION_MASK,
+                vec->action_mask     + agent_start * MY_ACTION_MASK,
+                agents_per_buffer * MY_ACTION_MASK * sizeof(unsigned char),
+                cudaMemcpyHostToDevice, stream);
+#endif
         }
         cudaStreamSynchronize(stream);
         atomic_store(&buffer_states[buf], OMP_WAITING);
@@ -423,6 +445,21 @@ StaticVec* create_static_vec(int total_agents, int num_buffers, int gpu, Dict* v
         vec->gpu_terminals = vec->terminals;
     }
 
+#ifdef MY_ACTION_MASK
+    vec->action_mask_size = MY_ACTION_MASK;
+    size_t mask_bytes = (size_t)total_agents * MY_ACTION_MASK * sizeof(unsigned char);
+    if (gpu) {
+        cudaHostAlloc((void**)&vec->action_mask, mask_bytes, cudaHostAllocPortable);
+        cudaMalloc((void**)&vec->gpu_action_mask, mask_bytes);
+        cudaMemset(vec->gpu_action_mask, 0, mask_bytes);
+    } else {
+        vec->action_mask = (unsigned char*)calloc(total_agents * MY_ACTION_MASK, sizeof(unsigned char));
+        vec->gpu_action_mask = vec->action_mask;
+    }
+#endif
+    // No #else: action_mask, gpu_action_mask, action_mask_size are already 0/NULL
+    // from calloc(1, sizeof(StaticVec)) above.
+
     // Streams allocated here, created in create_static_threads
     vec->streams = (cudaStream_t*)calloc(num_buffers, sizeof(cudaStream_t));
 
@@ -441,14 +478,89 @@ StaticVec* create_static_vec(int total_agents, int num_buffers, int gpu, Dict* v
             env->actions = vec->actions + slot * NUM_ATNS;
             env->rewards = vec->rewards + slot;
             env->terminals = vec->terminals + slot;
+#ifdef MY_ACTION_MASK
+            env->action_mask = vec->action_mask + slot * MY_ACTION_MASK;
+#endif
+#ifdef MY_USES_PERM
+            // Populate per-slot pointer arrays. agent_perm is NULL here (identity),
+            // so slots map to the same adjacent layout as the base+stride pointers.
+            my_setup_perm(vec, env, slot);
+#endif
             buf_agent += env->num_agents;
         }
-        assert(buf_agent == vec->agents_per_buffer && "buffer agents don't match total agents");
     }
 
     return vec;
 }
 
+void static_vec_set_perm(StaticVec* vec, const int* perm) {
+#ifndef MY_USES_PERM
+    (void)vec; (void)perm;
+    fprintf(stderr, "static_vec_set_perm: env did not opt in via MY_USES_PERM; ignoring.\n");
+    return;
+#else
+    int N = vec->total_agents;
+    if (vec->agent_perm == NULL) {
+        vec->agent_perm = (int*)malloc(N * sizeof(int));
+    }
+    memcpy(vec->agent_perm, perm, N * sizeof(int));
+
+    Env* envs = (Env*)vec->envs;
+    for (int buf = 0; buf < vec->buffers; buf++) {
+        int buf_start = buf * vec->agents_per_buffer;
+        int buf_agent = 0;
+        int env_start = vec->buffer_env_starts[buf];
+        int env_count = vec->buffer_env_counts[buf];
+        for (int e = 0; e < env_count; e++) {
+            Env* env = &envs[env_start + e];
+            my_setup_perm(vec, env, buf_start + buf_agent);
+            buf_agent += env->num_agents;
+        }
+    }
+#endif
+}
+
+#ifdef MY_USES_TAGS
+void static_vec_set_env_tags(StaticVec* vec, const int* tags) {
+    Env* envs = (Env*)vec->envs;
+    for (int i = 0; i < vec->size; i++) {
+        envs[i].tag = tags[i];
+        envs[i].boundary_reached = 0;
+    }
+}
+
+int static_vec_count_aligned(StaticVec* vec, int tag_value, int reset_flags) {
+    Env* envs = (Env*)vec->envs;
+    int count = 0;
+    for (int i = 0; i < vec->size; i++) {
+        if (envs[i].tag == tag_value && envs[i].boundary_reached) {
+            count++;
+        }
+    }
+    // Only reset matching-tag envs. Multi-bank: each bank's swap finalizes
+    // by calling count_aligned(tag=bank+1, reset=1); we must not touch
+    // boundary_reached on envs playing other banks or their alignment data
+    // would be lost.
+    if (reset_flags) {
+        for (int i = 0; i < vec->size; i++) {
+            if (envs[i].tag == tag_value) {
+                envs[i].boundary_reached = 0;
+            }
+        }
+    }
+    return count;
+}
+#else
+void static_vec_set_env_tags(StaticVec* vec, const int* tags) {
+    (void)vec; (void)tags;
+    fprintf(stderr, "static_vec_set_env_tags: env did not opt in via MY_USES_TAGS; ignoring.\n");
+}
+int static_vec_count_aligned(StaticVec* vec, int tag_value, int reset_flags) {
+    (void)vec; (void)tag_value; (void)reset_flags;
+    return 0;
+}
+#endif
+
 void static_vec_reset(StaticVec* vec) {
     Env* envs = (Env*)vec->envs;
     for (int i = 0; i < vec->size; i++) {
@@ -459,6 +571,11 @@ void static_vec_reset(StaticVec* vec) {
             vec->total_agents * OBS_SIZE * obs_element_size(), cudaMemcpyHostToDevice);
         cudaMemset(vec->gpu_rewards,   0, vec->total_agents * sizeof(float));
         cudaMemset(vec->gpu_terminals, 0, vec->total_agents * sizeof(float));
+#ifdef MY_ACTION_MASK
+        cudaMemcpy(vec->gpu_action_mask, vec->action_mask,
+            (size_t)vec->total_agents * MY_ACTION_MASK * sizeof(unsigned char),
+            cudaMemcpyHostToDevice);
+#endif
         cudaDeviceSynchronize();
     } else {
         memset(vec->rewards, 0, vec->total_agents * sizeof(float));
@@ -478,8 +595,7 @@ void create_static_threads(StaticVec* vec, int num_threads, int horizon,
     // Streams are now created by pufferlib.cu (PyTorch-managed streams)
     // Do NOT create streams here - they've already been set up
 
-    vec->threading->thread_args = (StaticOMPArg*)calloc(vec->buffers, sizeof(StaticOMPArg));
-    StaticOMPArg* args = vec->threading->thread_args;
+    StaticOMPArg* args = (StaticOMPArg*)calloc(vec->buffers, sizeof(StaticOMPArg));
     for (int i = 0; i < vec->buffers; i++) {
         args[i].vec = vec;
         args[i].buf = i;
@@ -512,7 +628,6 @@ void static_vec_close(StaticVec* vec) {
         free(vec->threading->buffer_states);
         free(vec->threading->threads);
         free(vec->threading->accum);
-        free(vec->threading->thread_args);
         free(vec->threading);
     }
     free(vec->buffer_env_starts);
@@ -528,14 +643,22 @@ void static_vec_close(StaticVec* vec) {
         cudaFreeHost(vec->actions);
         cudaFreeHost(vec->rewards);
         cudaFreeHost(vec->terminals);
+#ifdef MY_ACTION_MASK
+        cudaFree(vec->gpu_action_mask);
+        cudaFreeHost(vec->action_mask);
+#endif
     } else {
         free(vec->observations);
         free(vec->actions);
         free(vec->rewards);
         free(vec->terminals);
+#ifdef MY_ACTION_MASK
+        free(vec->action_mask);
+#endif
     }
 
     free(vec->streams);
+    if (vec->agent_perm != NULL) free(vec->agent_perm);
     free(vec);
 }
 
@@ -615,12 +738,6 @@ int get_num_act_sizes(void) { return (int)(sizeof(_act_sizes) / sizeof(_act_size
 const char* get_obs_dtype(void) { return dtype_symbol; }
 size_t get_obs_elem_size(void) { return obs_element_size(); }
 
-#ifdef MY_VEC_STEP
-void MY_VEC_STEP(StaticVec* vec);
-static inline void _static_vec_env_step(StaticVec* vec) {
-    MY_VEC_STEP(vec);
-}
-#else
 static inline void _static_vec_env_step(StaticVec* vec) {
     memset(vec->rewards, 0, vec->total_agents * sizeof(float));
     memset(vec->terminals, 0, vec->total_agents * sizeof(float));
@@ -630,7 +747,6 @@ static inline void _static_vec_env_step(StaticVec* vec) {
         c_step(&envs[i]);
     }
 }
-#endif
 
 void gpu_vec_step(StaticVec* vec) {
     assert(vec->buffers == 1);
diff --git a/tests/profile_kernels.cu b/tests/profile_kernels.cu
index 880eeab2ee..02337b3fe4 100644
--- a/tests/profile_kernels.cu
+++ b/tests/profile_kernels.cu
@@ -511,7 +511,7 @@ void run_samplelogits(SampleLogitsProfile* p) {
     sample_logits<<<grid_size(p->B), BLOCK_SIZE>>>(
         p->dec_out, p->logstd, p->act_sizes,
         p->actions_t.data, p->logprobs_t.data, p->value_out_t.data,
-        p->rng_states);
+        p->rng_states, nullptr, 0);
 }
 
 void profile_samplelogits(int B, int A) {
