From df39400c7bad23ae6ce5ca4fc538f52fccfa6d2c Mon Sep 17 00:00:00 2001 From: Spencer Cheng Date: Tue, 7 Apr 2026 18:53:05 +0000 Subject: [PATCH 1/8] fix sweep for WSL users --- pufferlib/pufferl.py | 7 +++++++ pufferlib/sweep.py | 11 +++++++++++ 2 files changed, 18 insertions(+) diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index 8fc0c03a89..43b736f1c4 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -333,8 +333,15 @@ def train(env_name, args=None, gpus=None, **kwargs): if rank == 0 and not subprocess: _train(env_name, worker_args, verbose=True) else: + # Spawn pickling uses CUDA IPC, unsupported on WSL2. Move to CPU first. + sweep_obj = kwargs.get('sweep_obj') + device = getattr(sweep_obj, 'device', None) + if device and device.type != 'cpu': + sweep_obj.to('cpu') ctx.Process(target=_train, args=(env_name, worker_args), kwargs=kwargs).start() + if device and device.type != 'cpu': + sweep_obj.to(device) def sweep(env_name, args=None, pareto=False): '''Train entry point. Handles single-GPU, multi-GPU DDP, and sweeps.''' 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)) From 03c5573527adf014e978ccbf02bf1e81a248ac27 Mon Sep 17 00:00:00 2001 From: Spencer Cheng Date: Thu, 9 Apr 2026 14:33:04 +0000 Subject: [PATCH 2/8] boulder env first commit --- config/boulder.ini | 51 ++++ ocean/boulder/binding.c | 22 ++ ocean/boulder/boulder.c | 45 +++ ocean/boulder/boulder.h | 638 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 756 insertions(+) create mode 100644 config/boulder.ini create mode 100644 ocean/boulder/binding.c create mode 100644 ocean/boulder/boulder.c create mode 100644 ocean/boulder/boulder.h diff --git a/config/boulder.ini b/config/boulder.ini new file mode 100644 index 0000000000..2f37acee9d --- /dev/null +++ b/config/boulder.ini @@ -0,0 +1,51 @@ +[base] +env_name = boulder + +[vec] +total_agents = 1024 +num_buffers = 2.31152 +num_threads = 16 + +[env] +num_agents = 2 +width = 900 +height = 600 + +[policy] +hidden_size = 256 +num_layers = 3.63347 +expansion_factor = 1 + +[train] +gpus = 1 +seed = 42 +total_timesteps = 77332504 +learning_rate = 0.00762627 +anneal_lr = 1 +min_lr_ratio = 0 +gamma = 0.993349 +gae_lambda = 0.97086 +replay_ratio = 2.76336 +clip_coef = 0.414034 +vf_coef = 2.90403 +vf_clip_coef = 0.01 +max_grad_norm = 1.19707 +ent_coef = 0.0020498 +beta1 = 0.937322 +beta2 = 0.999958 +eps = 2.42555e-12 +minibatch_size = 4096 +horizon = 128 +vtrace_rho_clip = 2.2155 +vtrace_c_clip = 1.21309 +prio_alpha = 0.24272 +prio_beta0 = 0.366241 + +[sweep.train.total_timesteps] +distribution = log_normal +min = 5e7 +max = 5e8 +mean = 1e8 +scale = auto + + diff --git a/ocean/boulder/binding.c b/ocean/boulder/binding.c new file mode 100644 index 0000000000..6af3fde41d --- /dev/null +++ b/ocean/boulder/binding.c @@ -0,0 +1,22 @@ +#include "boulder.h" +#define NUM_ATNS 2 +#define ACT_SIZES {2, 8} +#define OBS_TENSOR_T FloatTensor + +#define Env Boulder +#include "vecenv.h" + +void my_init(Env* env, Dict* kwargs) { + env->num_agents = dict_get(kwargs, "num_agents")->value; + env->width = dict_get(kwargs, "width")->value; + env->height = dict_get(kwargs, "height")->value; + init(env); +} + +void my_log(Log* log, Dict* out) { + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); +} + diff --git a/ocean/boulder/boulder.c b/ocean/boulder/boulder.c new file mode 100644 index 0000000000..819d86cea9 --- /dev/null +++ b/ocean/boulder/boulder.c @@ -0,0 +1,45 @@ +#include +#include "boulder.h" +#include "puffernet.h" + +void demo() { + Boulder env = { + .width = 900, + .height = 600, + .num_agents = 2, + }; + allocate(&env); + + env.client = make_client(&env); + + c_reset(&env); + int frame = 0; + SetTargetFPS(60); + while (!WindowShouldClose()) { + int kl = IsKeyDown(KEY_LEFT), kr = IsKeyDown(KEY_RIGHT); + int ku = IsKeyDown(KEY_UP), kd = IsKeyDown(KEY_DOWN); + env.actions[0] = (kl || kr || ku || kd) ? 1 : 0; // throttle + int dir = 0; + if (kr && !ku && !kd) dir = 0; // E + if (kr && ku) dir = 1; // NE + if (ku && !kl && !kr) dir = 2; // N + if (kl && ku) dir = 3; // NW + if (kl && !ku && !kd) dir = 4; // W + if (kl && kd) dir = 5; // SW + if (kd && !kl && !kr) dir = 6; // S + if (kr && kd) dir = 7; // SE + env.actions[1] = dir; + + c_step(&env); + c_render(&env); + } + //free_puffernet(net); + //free(weights); + free_allocated(&env); + close_client(env.client); +} + +int main() { + demo(); +} + diff --git a/ocean/boulder/boulder.h b/ocean/boulder/boulder.h new file mode 100644 index 0000000000..c27c4f7323 --- /dev/null +++ b/ocean/boulder/boulder.h @@ -0,0 +1,638 @@ +#include +#include +#include +#include +#include +#include +#include +#include "raylib.h" + +#define NOOP 0 +#define LEFT 1 +#define RIGHT 2 +#define UP 3 +#define DOWN 4 + +// ─── Physics & game constants +#define TICK_RATE (1.0f / 60.0f) +#define GRAVITY 980.0f // px/s² slope-gravity scale + +#define AGENT_RADIUS 12.0f +#define BOULDER_RADIUS 48.0f +#define GOAL_RADIUS 40.0f + +#define AGENT_MASS 1.0f +#define BOULDER_MASS 20.0f + +#define AGENT_ACCEL 700.0f // px/s² per directional action +#define AGENT_DRAG 6.0f // per-second linear velocity decay +#define BOULDER_LINEAR_DRAG 0.4f +#define BOULDER_ROLLING_DRAG 0.8f // per-second angular velocity decay + +#define RESTITUTION_BODIES 0.05f // agent ↔ boulder (nearly inelastic) +#define RESTITUTION_WALL 0.15f +#define FRICTION_COEF 0.45f // Coulomb friction coefficient + +#define MAX_SPEED 1200.0f +#define MAX_ANGULAR_VEL 30.0f // rad/s normalisation for boulder spin +#define MAX_TICKS 600 +#define GOAL_CENTER_RADIUS 16.0f +#define GOAL_HOLD_TICKS 10 + +#define MAX_AGENTS 2 +// Per-agent obs layout: self(6) + boulder(8) + goal(5) + (MAX_AGENTS-1) other agents(8 each) +#define OBS_SIZE (6 + 8 + 5 + (MAX_AGENTS - 1) * 8) + +typedef struct Log { + float perf; + float score; + float episode_return; + float episode_length; + float n; +} Log; + +typedef struct Client { + float width; + float height; + Texture2D sprites; +} Client; + +// Circular rigid body — used for both agents and the boulder. +// All physics functions operate on this type. +typedef struct Entity { + float x, y; + float vx, vy; + float mass; + float radius; + float angle; // cumulative rotation (rad), increases CCW + float angular_vel; // angular velocity (rad/s), + = CCW + float inv_inertia; // 1/I; solid disk: I = ½·mass·radius² → inv = 2/(m·r²) +} Entity; + +typedef struct Boulder { + Client* client; + Log log; + Log* logs; + float* observations; + float* actions; + float* rewards; + float* terminals; + Entity boulder; + Entity* agents; + int num_agents; + float goal_x, goal_y; + int goal_hold; // consecutive ticks boulder has been within GOAL_RADIUS + float* heightmap; // flat [hmap_rows × hmap_cols] elevation values (px) + int hmap_cols; + int hmap_rows; + int width; + int height; + int tick; + float score; + unsigned int rng; + int* moving_boulder; +} Boulder; + +// Helper functions + +// Returns a uniform float in [0, 1) using rand_r seeded from env->rng. +static inline float randf(Boulder* env) { + return (float)rand_r(&env->rng) / ((float)RAND_MAX + 1.0f); +} + +static inline float clampf(float x, float lo, float hi) { + return x < lo ? lo : (x > hi ? hi : x); +} + +// Initialise a circular rigid body as a solid disk at rest. +// Precomputes inv_inertia = 2 / (mass · radius²). +static inline void entity_init(Entity* e,float x, float y,float mass, float radius) { + *e = (Entity){ + .x = x, .y = y, + .mass = mass, + .radius = radius, + .inv_inertia = 2.0f / (mass * radius * radius), + }; +} + +// Semi-implicit Euler integration with exponential velocity decay (drag). +// ax, ay : net acceleration (px/s²) this tick. +// linear_drag : linear damping rate (s⁻¹); applied as exp(−k·dt) factor. +// rolling_drag : angular damping rate (s⁻¹). +static inline void integrate_entity(Entity* e,float ax, float ay,float linear_drag,float rolling_drag) { + const float dt = TICK_RATE; + e->vx += ax * dt; + e->vy += ay * dt; + float ld = expf(-linear_drag * dt); + float rd = expf(-rolling_drag * dt); + e->vx *= ld; + e->vy *= ld; + e->x += e->vx * dt; + e->y += e->vy * dt; + e->angular_vel *= rd; + e->angle += e->angular_vel * dt; +} + +// ─── Height-map + +// Bilinear sample of the elevation map at world position (x, y). +// Returns 0 if heightmap is NULL or position is out of range. +static inline float heightmap_sample(const Boulder* env, float x, float y) { + if (!env->heightmap) return 0.0f; + float cx = (x / env->width) * (float)(env->hmap_cols - 1); + float cy = (y / env->height) * (float)(env->hmap_rows - 1); + int ix = (int)cx, iy = (int)cy; + if (ix < 0 || iy < 0 || ix >= env->hmap_cols - 1 || iy >= env->hmap_rows - 1) + return 0.0f; + float tx = cx - ix, ty = cy - iy; + int s = env->hmap_cols; + float h00 = env->heightmap[ iy * s + ix ]; + float h10 = env->heightmap[ iy * s + ix + 1]; + float h01 = env->heightmap[(iy+1) * s + ix ]; + float h11 = env->heightmap[(iy+1) * s + ix + 1]; + return h00*(1-tx)*(1-ty) + h10*tx*(1-ty) + h01*(1-tx)*ty + h11*tx*ty; +} + +// Computes slope-induced gravitational acceleration (px/s²) at world position +// (x, y) via central-difference gradient of the height field. +// Fills *out_ax and *out_ay with the downhill acceleration components. +static inline void heightmap_gravity(const Boulder* env, + float x, float y, + float* out_ax, float* out_ay) { + const float h = 4.0f; + float dh_dx = (heightmap_sample(env, x + h, y) - + heightmap_sample(env, x - h, y)) / (2.0f * h); + float dh_dy = (heightmap_sample(env, x, y + h) - + heightmap_sample(env, x, y - h)) / (2.0f * h); + *out_ax = -GRAVITY * dh_dx; + *out_ay = -GRAVITY * dh_dy; +} + +// ─── Sphere–sphere impulse resolution +// +// Convention: n is the unit normal FROM a TO b. +// vrel_n > 0 ⟹ contact points approaching ⟹ apply impulse. +// +// Normal restitution + Coulomb tangential friction. +// Both bodies are modified in place. +// +// Derivation notes (2D rigid-body impulse): +// • Contact arm: r_a = +a->radius·n, r_b = −b->radius·n +// • Velocity at contact: v_c = v_cm + ω × r → (vx − ω·ry , vy + ω·rx) +// • r_a × n = 0 ⟹ angular term drops out of the normal-impulse denominator. +// • cross2d(r_a, t) = +a->radius, cross2d(r_b, t) = −b->radius +// where t = (−ny, nx) is the CCW tangent. +static inline int resolve_sphere_sphere(Entity* a, Entity* b, + float restitution, float friction) { + float dx = b->x - a->x; + float dy = b->y - a->y; + float dist2 = dx*dx + dy*dy; + float rsum = a->radius + b->radius; + if (dist2 >= rsum * rsum || dist2 < 1e-8f) return 0; + + float dist = sqrtf(dist2); + float nx = dx / dist; // unit normal a → b + float ny = dy / dist; + float tx = -ny; // unit tangent (CCW 90° from n) + float ty = nx; + + // Positional correction: push apart proportional to inverse mass + float overlap = rsum - dist; + float inv_msum = 1.0f / (a->mass + b->mass); + a->x -= nx * overlap * (b->mass * inv_msum); + a->y -= ny * overlap * (b->mass * inv_msum); + b->x += nx * overlap * (a->mass * inv_msum); + b->y += ny * overlap * (a->mass * inv_msum); + + // Contact-point arm vectors (centre → contact point) + float ra_x = a->radius * nx, ra_y = a->radius * ny; + float rb_x = -b->radius * nx, rb_y = -b->radius * ny; + + // Velocity at each contact point: v_c = v_cm + ω × r (2D: vx−ω·ry, vy+ω·rx) + float vca_x = a->vx - a->angular_vel * ra_y; + float vca_y = a->vy + a->angular_vel * ra_x; + float vcb_x = b->vx - b->angular_vel * rb_y; + float vcb_y = b->vy + b->angular_vel * rb_x; + + float vrx = vca_x - vcb_x; + float vry = vca_y - vcb_y; + float vrel_n = vrx * nx + vry * ny; + if (vrel_n <= 0.0f) return 0; // separating; no impulse needed + + float vrel_t = vrx * tx + vry * ty; + + float inv_meff = 1.0f/a->mass + 1.0f/b->mass; + + // Normal impulse (angular inertia term is 0 for sphere–sphere) + float j_n = (1.0f + restitution) * vrel_n / inv_meff; + + // Tangential (friction) impulse including angular inertia in denominator + // cross2d(ra, t) = +a->radius; cross2d(rb, t) = −b->radius + float denom_t = inv_meff + + a->radius * a->radius * a->inv_inertia + + b->radius * b->radius * b->inv_inertia; + float j_t = vrel_t / denom_t; + float j_t_max = friction * j_n; + if (j_t > j_t_max) j_t = j_t_max; + if (j_t < -j_t_max) j_t = -j_t_max; + + // Linear impulse: impulse on a = −J, impulse on b = +J + float Jx = j_n * nx + j_t * tx; + float Jy = j_n * ny + j_t * ty; + a->vx -= Jx / a->mass; + a->vy -= Jy / a->mass; + b->vx += Jx / b->mass; + b->vy += Jy / b->mass; + + // Angular impulse: Δω = cross2d(r, F) · inv_I + // F on a = −J, F on b = +J + a->angular_vel += (ra_x * (-Jy) - ra_y * (-Jx)) * a->inv_inertia; + b->angular_vel += (rb_x * Jy - rb_y * Jx ) * b->inv_inertia; + + return 1; +} + +// ─── Sphere–horizontal-wall impulse resolution +// +// wall_y : y-coordinate of the wall line. +// normal_y : outward wall normal y-component (screen coords: y increases DOWN). +// +1 → top / ceiling wall (normal points downward toward entity). +// −1 → bottom / floor wall (normal points upward toward entity). +// +// Derivation: +// Contact arm rc = (0, −normal_y · radius) → toward wall face. +// x-velocity at contact: vx_c = vx + ω·(−rc_y) = vx + ω·normal_y·radius. +// cross2d(rc, t_x) = −rc_y = normal_y·radius. +static inline void resolve_sphere_hwall(Entity* e, + float wall_y, float normal_y, + float restitution, float friction) { + // Signed distance from wall to entity centre along outward normal + float sd = (e->y - wall_y) * normal_y; + float pen = e->radius - sd; + if (pen <= 0.0f) return; + + e->y += pen * normal_y; // push entity off wall + + float vn = e->vy * normal_y; // velocity along outward normal + if (vn >= 0.0f) return; // already leaving + + float j_n = -(1.0f + restitution) * vn * e->mass; // positive + + float rc_y = -normal_y * e->radius; // arm y-component + float vx_c = e->vx + e->angular_vel * (-rc_y); // x-vel at contact + float rc_ct = -rc_y; // cross2d(rc, (1,0)) + float denom_t = 1.0f/e->mass + rc_ct * rc_ct * e->inv_inertia; + float j_t = -vx_c / denom_t; + float j_t_max = friction * j_n; + if (j_t > j_t_max) j_t = j_t_max; + if (j_t < -j_t_max) j_t = -j_t_max; + + e->vy += j_n * normal_y / e->mass; + e->vx += j_t / e->mass; + // cross2d((0, rc_y), (j_t, j_n·normal_y)) = −rc_y·j_t + e->angular_vel += (-rc_y * j_t) * e->inv_inertia; +} + +// ─── Sphere–vertical-wall impulse resolution ───────────────────────────────── +// +// wall_x : x-coordinate of the wall line. +// normal_x : outward wall normal x-component. +// +1 → left wall (normal points right toward entity). +// −1 → right wall (normal points left toward entity). +// +// Derivation: +// Contact arm rc = (−normal_x · radius, 0). +// y-velocity at contact: vy_c = vy + ω·rc_x. +// cross2d(rc, t_y) = rc_x = −normal_x·radius. +static inline void resolve_sphere_vwall(Entity* e, + float wall_x, float normal_x, + float restitution, float friction) { + float sd = (e->x - wall_x) * normal_x; + float pen = e->radius - sd; + if (pen <= 0.0f) return; + + e->x += pen * normal_x; + + float vn = e->vx * normal_x; + if (vn >= 0.0f) return; + + float j_n = -(1.0f + restitution) * vn * e->mass; // positive + + float rc_x = -normal_x * e->radius; // arm x-component + float vy_c = e->vy + e->angular_vel * rc_x; // y-vel at contact + float rc_ct = rc_x; // cross2d(rc, (0,1)) + float denom_t = 1.0f/e->mass + rc_ct * rc_ct * e->inv_inertia; + float j_t = -vy_c / denom_t; + float j_t_max = friction * j_n; + if (j_t > j_t_max) j_t = j_t_max; + if (j_t < -j_t_max) j_t = -j_t_max; + + e->vx += j_n * normal_x / e->mass; + e->vy += j_t / e->mass; + // cross2d((rc_x, 0), (j_n·normal_x, j_t)) = rc_x·j_t + e->angular_vel += (rc_x * j_t) * e->inv_inertia; +} + + +void init(Boulder* env) { + env->tick = 0; + env->score = 0.0f; + env->goal_hold = 0; + env->agents = (Entity*)calloc(env->num_agents, sizeof(Entity)); + env->logs = (Log*)calloc(env->num_agents, sizeof(Log)); + env->moving_boulder = (int*)calloc(env->num_agents, sizeof(int)); +} + +void allocate(Boulder* env) { + init(env); + int n = env->num_agents; + env->observations = (float*)calloc(n * OBS_SIZE, sizeof(float)); + env->actions = (float*)calloc(n * 2, sizeof(float)); + env->rewards = (float*)calloc(n, sizeof(float)); + env->terminals = (float*)calloc(n, sizeof(float)); + env->agents = (Entity*)calloc(n, sizeof(Entity)); +} + +void c_close(Boulder* env) {} + +void free_allocated(Boulder* env) { + free(env->observations); + free(env->actions); + free(env->rewards); + free(env->terminals); + free(env->agents); + if (env->heightmap) free(env->heightmap); + c_close(env); +} + +void add_log(Boulder* env) { + for(int i= 0;inum_agents; i++){ + env->log.episode_length += env->logs[i].episode_length; + env->log.episode_return += env->logs[i].episode_return; + env->log.score += env->score; + env->log.perf += env->score; + env->log.n += 1; + } +} + +void compute_observations(Boulder* env) { + float diag = hypotf((float)env->width, (float)env->height); + float dist_scale = env->width > env->height ? env->width: env->height; + dist_scale *= 0.5; + for (int i = 0; i < env->num_agents; i++) { + float* obs = env->observations + i * OBS_SIZE; + Entity* a = &env->agents[i]; + int idx = 0; + + // Self (4): position relative to walls and velocity + obs[idx++] = a->x / env->width; + obs[idx++] = (env->width - a->x) / env->width; + obs[idx++] = a->y / env->height; + obs[idx++] = (env->height - a->y) / env->height; + obs[idx++] = a->vx / MAX_SPEED; + obs[idx++] = a->vy / MAX_SPEED; + + // Boulder (8): ego-relative pos, dist, angle, velocity, spin + float bdx = env->boulder.x - a->x; + float bdy = env->boulder.y - a->y; + float bdist = hypotf(bdx, bdy); + float bang = atan2f(bdy, bdx); + obs[idx++] = clampf(bdx / dist_scale, -1.0f, 1.0f); + obs[idx++] = clampf(bdy / dist_scale, -1.0f, 1.0f); + obs[idx++] = clampf(bdist / diag, 0.0f, 1.0f); + obs[idx++] = sinf(bang); + obs[idx++] = cosf(bang); + obs[idx++] = env->boulder.vx / MAX_SPEED; + obs[idx++] = env->boulder.vy / MAX_SPEED; + obs[idx++] = env->boulder.angular_vel / MAX_ANGULAR_VEL; + + // Goal (5): ego-relative pos, dist, angle + float gdx = env->goal_x - env->boulder.x; + float gdy = env->goal_y - env->boulder.y; + float gdist = hypotf(gdx, gdy); + float gang = atan2f(gdy, gdx); + obs[idx++] = clampf(gdx / dist_scale, -1.0f, 1.0f); + obs[idx++] = clampf(gdy / dist_scale, -1.0f, 1.0f); + obs[idx++] = clampf(gdist / diag, 0.0f, 1.0f); + obs[idx++] = sinf(gang); + obs[idx++] = cosf(gang); + + // Other agents (8 each, up to MAX_AGENTS-1 slots, zero-padded if fewer) + int filled = 0; + for (int j = 0; j < env->num_agents && filled < MAX_AGENTS - 1; j++) { + if (j == i) continue; + Entity* o = &env->agents[j]; + float odx = o->x - a->x; + float ody = o->y - a->y; + float odist = hypotf(odx, ody); + float oang = atan2f(ody, odx); + obs[idx++] = clampf(odx / dist_scale, -1.0f, 1.0f); + obs[idx++] = clampf(ody / dist_scale, -1.0f, 1.0f); + obs[idx++] = clampf(odist / diag, 0.0f, 1.0f); + obs[idx++] = sinf(oang); + obs[idx++] = cosf(oang); + obs[idx++] = (o->vx - a->vx) / MAX_SPEED; + obs[idx++] = (o->vy - a->vy) / MAX_SPEED; + obs[idx++] = hypotf(o->vx, o->vy) / MAX_SPEED; + filled++; + } + while (filled++ < MAX_AGENTS - 1) { + for (int k = 0; k < 8; k++) obs[idx++] = 0.0f; + } + } +} + +void c_reset(Boulder* env) { + env->score = 0.0f; + env->tick = 0; + env->goal_hold = 0; + + entity_init(&env->boulder,env->width * 0.5f,env->height * 0.5f,BOULDER_MASS, BOULDER_RADIUS); + + float pi2 = 2.0f * 3.14159265f; + for (int i = 0; i < env->num_agents; i++) { + env->logs[i] = (Log){0}; + float angle = (float)i * (pi2 / env->num_agents); + float dist = BOULDER_RADIUS + AGENT_RADIUS + 24.0f; + float start_x = env->boulder.x + cosf(angle) * dist; + float start_y = env->boulder.y + sinf(angle) * dist; + entity_init(&env->agents[i], start_x, start_y, AGENT_MASS, AGENT_RADIUS); + } + + // Random goal far enough from the boulder start + float gx, gy; + do { + gx = GOAL_RADIUS + randf(env) * (env->width - 2.0f * GOAL_RADIUS); + gy = GOAL_RADIUS + randf(env) * (env->height - 2.0f * GOAL_RADIUS); + } while (hypotf(gx - env->boulder.x, gy - env->boulder.y) < BOULDER_RADIUS * 3.0f); + env->goal_x = gx; + env->goal_y = gy; + + memset(env->moving_boulder, 0, 2*sizeof(int)); + compute_observations(env); +} + +void c_step(Boulder* env) { + float dist_to_goal_before = hypotf(env->boulder.x - env->goal_x, env->boulder.y - env->goal_y); + // 8-direction unit vectors (screen coords: y increases down) + // index: 0=E, 1=NE, 2=N, 3=NW, 4=W, 5=SW, 6=S, 7=SE + static const float DIR_AX[8] = { 1.0f, 0.707f, 0.0f, -0.707f, -1.0f, -0.707f, 0.0f, 0.707f}; + static const float DIR_AY[8] = { 0.0f, -0.707f, -1.0f, -0.707f, 0.0f, 0.707f, 1.0f, 0.707f}; + + // Agent acceleration from discrete actions + height-map slope + for (int i = 0; i < env->num_agents; i++) { + int throttle = (int)env->actions[i * 2 + 0]; // 0=off, 1=on + int dir = (int)env->actions[i * 2 + 1]; // 0..7 + float ax = 0.0f, ay = 0.0f; + if (throttle && dir >= 0 && dir < 8) { + ax = DIR_AX[dir] * AGENT_ACCEL; + ay = DIR_AY[dir] * AGENT_ACCEL; + } + + float hgx = 0.0f, hgy = 0.0f; + heightmap_gravity(env, env->agents[i].x, env->agents[i].y, &hgx, &hgy); + integrate_entity(&env->agents[i], ax + hgx, ay + hgy, + AGENT_DRAG, 0.0f); + } + + // Boulder integration (height-map slope only; agents push via collision) + { + float hgx = 0.0f, hgy = 0.0f; + heightmap_gravity(env, env->boulder.x, env->boulder.y, &hgx, &hgy); + integrate_entity(&env->boulder, hgx, hgy,BOULDER_LINEAR_DRAG, BOULDER_ROLLING_DRAG); + } + + // Collision: agents ↔ boulder + for (int i = 0; i < env->num_agents; i++) { + int collided = resolve_sphere_sphere(&env->agents[i], &env->boulder,RESTITUTION_BODIES, FRICTION_COEF); + env->moving_boulder[i] = 1.0f; + } + + // Collision: agents ↔ agents + for (int i = 0; i < env->num_agents; i++) { + for (int j = i + 1; j < env->num_agents; j++) { + resolve_sphere_sphere(&env->agents[i], &env->agents[j],RESTITUTION_BODIES, FRICTION_COEF); + } + } + + // All entities ↔ map boundary walls (sphere-on-line) + float wx0 = 0.0f; + float wy0 = 0.0f; + float wx1 = env->width; + float wy1 = env->height; + + for (int i = 0; i < env->num_agents; i++) { + resolve_sphere_vwall(&env->agents[i], wx0, +1.0f, RESTITUTION_WALL, FRICTION_COEF); + resolve_sphere_vwall(&env->agents[i], wx1, -1.0f, RESTITUTION_WALL, FRICTION_COEF); + resolve_sphere_hwall(&env->agents[i], wy0, +1.0f, RESTITUTION_WALL, FRICTION_COEF); + resolve_sphere_hwall(&env->agents[i], wy1, -1.0f, RESTITUTION_WALL, FRICTION_COEF); + } + resolve_sphere_vwall(&env->boulder, wx0, +1.0f, RESTITUTION_WALL, FRICTION_COEF); + resolve_sphere_vwall(&env->boulder, wx1, -1.0f, RESTITUTION_WALL, FRICTION_COEF); + resolve_sphere_hwall(&env->boulder, wy0, +1.0f, RESTITUTION_WALL, FRICTION_COEF); + resolve_sphere_hwall(&env->boulder, wy1, -1.0f, RESTITUTION_WALL, FRICTION_COEF); + + // Goal: reward every tick boulder overlaps goal or approaches goal + float dist_to_goal = hypotf(env->boulder.x - env->goal_x, env->boulder.y - env->goal_y); + int on_goal = (dist_to_goal < GOAL_CENTER_RADIUS); + float boulder_move_reward = 0.01f*(dist_to_goal_before - dist_to_goal); + + for (int i = 0; i < env->num_agents; i++){ + float r = on_goal ? 1.0f : env->moving_boulder[i] ? boulder_move_reward : 0.0f; + env->rewards[i] += r; + env->logs[i].episode_return += r; + + } + if (on_goal) { + env->goal_hold++; + } else { + env->goal_hold = 0; + } + + env->tick++; + int goal_held_for_full = (env->goal_hold >= GOAL_HOLD_TICKS); + if(goal_held_for_full){ + env->score += 1.0f; + } + int done = goal_held_for_full || (env->tick >= MAX_TICKS); + if (done) { + for (int i = 0; i < env->num_agents; i++){ + env->terminals[i] = 1.0f; + env->logs[i].episode_length = env->tick; + } + add_log(env); + c_reset(env); + } + + compute_observations(env); +} + +// Rendering + +Client* make_client(Boulder* env) { + Client* client = (Client*)calloc(1, sizeof(Client)); + client->width = env->width; + client->height = env->height; + InitWindow(env->width, env->height, "PufferLib Boulder"); + SetTargetFPS(60); + client->sprites = LoadTexture("resources/shared/puffers.png"); + return client; +} + +void close_client(Client* client) { + UnloadTexture(client->sprites); + CloseWindow(); + free(client); +} + +void c_render(Boulder* env) { + if (env->client == NULL) env->client = make_client(env); + if (IsKeyDown(KEY_ESCAPE)) exit(0); + if (IsKeyPressed(KEY_TAB)) ToggleFullscreen(); + + BeginDrawing(); + ClearBackground((Color){6, 24, 24, 255}); + + // Goal ring + DrawCircle((int)env->goal_x, (int)env->goal_y, (int)GOAL_RADIUS, + (Color){200, 200, 0, 60}); + DrawCircleLines((int)env->goal_x, (int)env->goal_y, (int)GOAL_RADIUS, YELLOW); + + // Boulder: large dark sphere with rotation indicator line + DrawCircle((int)env->boulder.x, (int)env->boulder.y, + (int)env->boulder.radius, (Color){120, 75, 30, 255}); + DrawCircleLines((int)env->boulder.x, (int)env->boulder.y, + (int)env->boulder.radius, (Color){180, 120, 60, 255}); + { + float lx = env->boulder.x + cosf(env->boulder.angle) * env->boulder.radius * 0.75f; + float ly = env->boulder.y + sinf(env->boulder.angle) * env->boulder.radius * 0.75f; + DrawLine((int)env->boulder.x, (int)env->boulder.y, + (int)lx, (int)ly, (Color){220, 180, 100, 255}); + } + + // Agents: sprites from puffers.png, tinted per agent index, with velocity arrow + // Sprite sheet row: y=576 facing left, y=608 facing right (32px rows) + // Use sprite column 0 for all agents; differentiate via tint color + for (int i = 0; i < env->num_agents; i++) { + Entity* a = &env->agents[i]; + int src_x = 32 * i; // different column = different color + int src_y = (a->vx >= 0.0f) ? 576 : 608; // facing right or left + DrawTexturePro( + env->client->sprites, + (Rectangle){ src_x, src_y, 32, 32 }, + (Rectangle){ a->x - 16, a->y - 16, 32, 32 }, + (Vector2){0, 0}, + 0, + WHITE + ); + float speed = hypotf(a->vx, a->vy); + if (speed > 20.0f) { + float dx = a->vx / speed * (a->radius + 8.0f); + float dy = a->vy / speed * (a->radius + 8.0f); + DrawLine((int)a->x, (int)a->y, + (int)(a->x + dx), (int)(a->y + dy), WHITE); + } + } + + EndDrawing(); +} From a0f5ed8e2da3338a8867f8306a79232d2b25a5b2 Mon Sep 17 00:00:00 2001 From: Spencer Cheng Date: Thu, 9 Apr 2026 20:22:42 +0000 Subject: [PATCH 3/8] rand boulder spawn --- config/boulder.ini | 8 +++++++- ocean/boulder/binding.c | 1 + ocean/boulder/boulder.h | 41 +++++++++++++++++++++++------------------ 3 files changed, 31 insertions(+), 19 deletions(-) diff --git a/config/boulder.ini b/config/boulder.ini index 2f37acee9d..6b5b87873a 100644 --- a/config/boulder.ini +++ b/config/boulder.ini @@ -10,6 +10,7 @@ num_threads = 16 num_agents = 2 width = 900 height = 600 +dist_scale = 0.5 [policy] hidden_size = 256 @@ -48,4 +49,9 @@ max = 5e8 mean = 1e8 scale = auto - +[sweep.env.dist_scale] +distribution = uniform +max = 1.0 +min = 0.0 +mean = 0.5 +scale = auto diff --git a/ocean/boulder/binding.c b/ocean/boulder/binding.c index 6af3fde41d..028b1e313a 100644 --- a/ocean/boulder/binding.c +++ b/ocean/boulder/binding.c @@ -10,6 +10,7 @@ void my_init(Env* env, Dict* kwargs) { env->num_agents = dict_get(kwargs, "num_agents")->value; env->width = dict_get(kwargs, "width")->value; env->height = dict_get(kwargs, "height")->value; + env->dist_scale = dict_get(kwargs, "dist_scale")->value; init(env); } diff --git a/ocean/boulder/boulder.h b/ocean/boulder/boulder.h index c27c4f7323..72411d6a66 100644 --- a/ocean/boulder/boulder.h +++ b/ocean/boulder/boulder.h @@ -24,7 +24,7 @@ #define AGENT_MASS 1.0f #define BOULDER_MASS 20.0f -#define AGENT_ACCEL 700.0f // px/s² per directional action +#define AGENT_ACCEL 1400.0f // px/s² per directional action #define AGENT_DRAG 6.0f // per-second linear velocity decay #define BOULDER_LINEAR_DRAG 0.4f #define BOULDER_ROLLING_DRAG 0.8f // per-second angular velocity decay @@ -91,6 +91,7 @@ typedef struct Boulder { float score; unsigned int rng; int* moving_boulder; + float dist_scale; } Boulder; // Helper functions @@ -377,8 +378,8 @@ void add_log(Boulder* env) { void compute_observations(Boulder* env) { float diag = hypotf((float)env->width, (float)env->height); - float dist_scale = env->width > env->height ? env->width: env->height; - dist_scale *= 0.5; + float dist_scale_init = env->width > env->height ? env->width: env->height; + dist_scale_init *= env->dist_scale; for (int i = 0; i < env->num_agents; i++) { float* obs = env->observations + i * OBS_SIZE; Entity* a = &env->agents[i]; @@ -397,8 +398,8 @@ void compute_observations(Boulder* env) { float bdy = env->boulder.y - a->y; float bdist = hypotf(bdx, bdy); float bang = atan2f(bdy, bdx); - obs[idx++] = clampf(bdx / dist_scale, -1.0f, 1.0f); - obs[idx++] = clampf(bdy / dist_scale, -1.0f, 1.0f); + obs[idx++] = clampf(bdx / dist_scale_init, -1.0f, 1.0f); + obs[idx++] = clampf(bdy / dist_scale_init, -1.0f, 1.0f); obs[idx++] = clampf(bdist / diag, 0.0f, 1.0f); obs[idx++] = sinf(bang); obs[idx++] = cosf(bang); @@ -411,8 +412,8 @@ void compute_observations(Boulder* env) { float gdy = env->goal_y - env->boulder.y; float gdist = hypotf(gdx, gdy); float gang = atan2f(gdy, gdx); - obs[idx++] = clampf(gdx / dist_scale, -1.0f, 1.0f); - obs[idx++] = clampf(gdy / dist_scale, -1.0f, 1.0f); + obs[idx++] = clampf(gdx / dist_scale_init, -1.0f, 1.0f); + obs[idx++] = clampf(gdy / dist_scale_init, -1.0f, 1.0f); obs[idx++] = clampf(gdist / diag, 0.0f, 1.0f); obs[idx++] = sinf(gang); obs[idx++] = cosf(gang); @@ -426,8 +427,8 @@ void compute_observations(Boulder* env) { float ody = o->y - a->y; float odist = hypotf(odx, ody); float oang = atan2f(ody, odx); - obs[idx++] = clampf(odx / dist_scale, -1.0f, 1.0f); - obs[idx++] = clampf(ody / dist_scale, -1.0f, 1.0f); + obs[idx++] = clampf(odx / dist_scale_init, -1.0f, 1.0f); + obs[idx++] = clampf(ody / dist_scale_init, -1.0f, 1.0f); obs[idx++] = clampf(odist / diag, 0.0f, 1.0f); obs[idx++] = sinf(oang); obs[idx++] = cosf(oang); @@ -447,24 +448,28 @@ void c_reset(Boulder* env) { env->tick = 0; env->goal_hold = 0; - entity_init(&env->boulder,env->width * 0.5f,env->height * 0.5f,BOULDER_MASS, BOULDER_RADIUS); + // Boulder: random position with padding for its radius + float bx = BOULDER_RADIUS + randf(env) * (env->width - 2.0f * BOULDER_RADIUS); + float by = BOULDER_RADIUS + randf(env) * (env->height - 2.0f * BOULDER_RADIUS); + entity_init(&env->boulder, bx, by, BOULDER_MASS, BOULDER_RADIUS); - float pi2 = 2.0f * 3.14159265f; + // Agents: random positions outside boulder radius for (int i = 0; i < env->num_agents; i++) { env->logs[i] = (Log){0}; - float angle = (float)i * (pi2 / env->num_agents); - float dist = BOULDER_RADIUS + AGENT_RADIUS + 24.0f; - float start_x = env->boulder.x + cosf(angle) * dist; - float start_y = env->boulder.y + sinf(angle) * dist; - entity_init(&env->agents[i], start_x, start_y, AGENT_MASS, AGENT_RADIUS); + float ax, ay; + do { + ax = AGENT_RADIUS + randf(env) * (env->width - 2.0f * AGENT_RADIUS); + ay = AGENT_RADIUS + randf(env) * (env->height - 2.0f * AGENT_RADIUS); + } while (hypotf(ax - bx, ay - by) < BOULDER_RADIUS + AGENT_RADIUS); + entity_init(&env->agents[i], ax, ay, AGENT_MASS, AGENT_RADIUS); } - // Random goal far enough from the boulder start + // Goal: random position outside boulder radius * 3 float gx, gy; do { gx = GOAL_RADIUS + randf(env) * (env->width - 2.0f * GOAL_RADIUS); gy = GOAL_RADIUS + randf(env) * (env->height - 2.0f * GOAL_RADIUS); - } while (hypotf(gx - env->boulder.x, gy - env->boulder.y) < BOULDER_RADIUS * 3.0f); + } while (hypotf(gx - bx, gy - by) < BOULDER_RADIUS * 3.0f); env->goal_x = gx; env->goal_y = gy; From f314e33afa794a0090e6a3279ce7cd0c36500b27 Mon Sep 17 00:00:00 2001 From: Spencer Cheng Date: Fri, 10 Apr 2026 14:26:00 +0000 Subject: [PATCH 4/8] new boulder config --- config/boulder.ini | 5 ----- 1 file changed, 5 deletions(-) diff --git a/config/boulder.ini b/config/boulder.ini index 6b5b87873a..9954f3c1a3 100644 --- a/config/boulder.ini +++ b/config/boulder.ini @@ -1,11 +1,6 @@ [base] env_name = boulder -[vec] -total_agents = 1024 -num_buffers = 2.31152 -num_threads = 16 - [env] num_agents = 2 width = 900 From 29b42b9ccc5a0406eb6c2e3f687dec7746b49d7b Mon Sep 17 00:00:00 2001 From: Spencer Cheng Date: Fri, 10 Apr 2026 14:27:31 +0000 Subject: [PATCH 5/8] oops --- config/boulder.ini | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/config/boulder.ini b/config/boulder.ini index 9954f3c1a3..6b5b87873a 100644 --- a/config/boulder.ini +++ b/config/boulder.ini @@ -1,6 +1,11 @@ [base] env_name = boulder +[vec] +total_agents = 1024 +num_buffers = 2.31152 +num_threads = 16 + [env] num_agents = 2 width = 900 From 69afdef2aa3b228a22c0ef0d30bacc471580cd1a Mon Sep 17 00:00:00 2001 From: l1onh3art88 Date: Sat, 16 May 2026 17:10:45 -0500 Subject: [PATCH 6/8] selfplay + basic chess --- config/chess.ini | 101 ++ config/default.ini | 55 +- ocean/chess/binding.c | 218 +++ ocean/chess/chess.h | 3167 ++++++++++++++++++++++++++++++++++++++ pufferlib/pufferl.py | 183 ++- pufferlib/selfplay.py | 317 ++++ src/bindings.cu | 66 +- src/pufferlib.cu | 602 +++++++- src/vecenv.h | 160 +- tests/profile_kernels.cu | 2 +- 10 files changed, 4742 insertions(+), 129 deletions(-) create mode 100644 config/chess.ini create mode 100644 ocean/chess/binding.c create mode 100644 ocean/chess/chess.h create mode 100644 pufferlib/selfplay.py 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_ / hist_n_bank_ 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 +#include +#include +#include +#include +#include +#include +#include +#include +#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 \ngo nodes \n" +// child → parent: ... "info ..." ... "bestmove \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 +#include +#include +#include +#include + +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 43b736f1c4..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,15 +397,11 @@ def train(env_name, args=None, gpus=None, **kwargs): if rank == 0 and not subprocess: _train(env_name, worker_args, verbose=True) else: - # Spawn pickling uses CUDA IPC, unsupported on WSL2. Move to CPU first. - sweep_obj = kwargs.get('sweep_obj') - device = getattr(sweep_obj, 'device', None) - if device and device.type != 'cpu': - sweep_obj.to('cpu') + # 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() - if device and device.type != 'cpu': - sweep_obj.to(device) def sweep(env_name, args=None, pareto=False): '''Train entry point. Handles single-GPU, multi-GPU DDP, and sweeps.''' @@ -434,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', @@ -503,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) @@ -517,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/src/bindings.cu b/src/bindings.cu index 4469cb512c..64be61194d 100644 --- a/src/bindings.cu +++ b/src/bindings.cu @@ -2,6 +2,7 @@ #include #include +#include #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(); + 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_load_frozen_bank(&pufferl, bank_idx, path.c_str()); +} + +void py_set_agent_perm(py::object pufferl_obj, py::array_t perm) { + PuffeRL& pufferl = pufferl_obj.cast(); + 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 tags) { + PuffeRL& pufferl = pufferl_obj.cast(); + 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(); + return pufferl_count_aligned(&pufferl, tag_value, reset_flags); +} + +int py_num_envs(py::object pufferl_obj) { + PuffeRL& pufferl = pufferl_obj.cast(); + 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 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_(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 #include #include +#include #include #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_ / hist_n_bank_ 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<<>>( 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<<>>( + 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<<>>( - 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<<>>( - 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<<>>( + 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<<>>( + 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(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<<>>( + 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<<>>( 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<<>>( + 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<<>>( @@ -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<<>>( + int channels = (graph.mb_action_mask.data != nullptr) ? 6 : 5; + select_copy<<>>( 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 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<<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 create_pufferl_impl(HypersT& hypers, const std::string& env_name, Dict* vec_kwargs, Dict* env_kwargs) { auto pufferl = std::make_unique(); 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 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 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 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<<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) { From 03673370698bffc1cb352e2f0338dda54a17a4fd Mon Sep 17 00:00:00 2001 From: l1onh3art88 Date: Sat, 16 May 2026 17:15:15 -0500 Subject: [PATCH 7/8] delete stupid env --- ocean/boulder/binding.c | 23 -- ocean/boulder/boulder.c | 45 --- ocean/boulder/boulder.h | 643 ---------------------------------------- 3 files changed, 711 deletions(-) delete mode 100644 ocean/boulder/binding.c delete mode 100644 ocean/boulder/boulder.c delete mode 100644 ocean/boulder/boulder.h diff --git a/ocean/boulder/binding.c b/ocean/boulder/binding.c deleted file mode 100644 index 028b1e313a..0000000000 --- a/ocean/boulder/binding.c +++ /dev/null @@ -1,23 +0,0 @@ -#include "boulder.h" -#define NUM_ATNS 2 -#define ACT_SIZES {2, 8} -#define OBS_TENSOR_T FloatTensor - -#define Env Boulder -#include "vecenv.h" - -void my_init(Env* env, Dict* kwargs) { - env->num_agents = dict_get(kwargs, "num_agents")->value; - env->width = dict_get(kwargs, "width")->value; - env->height = dict_get(kwargs, "height")->value; - env->dist_scale = dict_get(kwargs, "dist_scale")->value; - init(env); -} - -void my_log(Log* log, Dict* out) { - dict_set(out, "perf", log->perf); - dict_set(out, "score", log->score); - dict_set(out, "episode_return", log->episode_return); - dict_set(out, "episode_length", log->episode_length); -} - diff --git a/ocean/boulder/boulder.c b/ocean/boulder/boulder.c deleted file mode 100644 index 819d86cea9..0000000000 --- a/ocean/boulder/boulder.c +++ /dev/null @@ -1,45 +0,0 @@ -#include -#include "boulder.h" -#include "puffernet.h" - -void demo() { - Boulder env = { - .width = 900, - .height = 600, - .num_agents = 2, - }; - allocate(&env); - - env.client = make_client(&env); - - c_reset(&env); - int frame = 0; - SetTargetFPS(60); - while (!WindowShouldClose()) { - int kl = IsKeyDown(KEY_LEFT), kr = IsKeyDown(KEY_RIGHT); - int ku = IsKeyDown(KEY_UP), kd = IsKeyDown(KEY_DOWN); - env.actions[0] = (kl || kr || ku || kd) ? 1 : 0; // throttle - int dir = 0; - if (kr && !ku && !kd) dir = 0; // E - if (kr && ku) dir = 1; // NE - if (ku && !kl && !kr) dir = 2; // N - if (kl && ku) dir = 3; // NW - if (kl && !ku && !kd) dir = 4; // W - if (kl && kd) dir = 5; // SW - if (kd && !kl && !kr) dir = 6; // S - if (kr && kd) dir = 7; // SE - env.actions[1] = dir; - - c_step(&env); - c_render(&env); - } - //free_puffernet(net); - //free(weights); - free_allocated(&env); - close_client(env.client); -} - -int main() { - demo(); -} - diff --git a/ocean/boulder/boulder.h b/ocean/boulder/boulder.h deleted file mode 100644 index 72411d6a66..0000000000 --- a/ocean/boulder/boulder.h +++ /dev/null @@ -1,643 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include "raylib.h" - -#define NOOP 0 -#define LEFT 1 -#define RIGHT 2 -#define UP 3 -#define DOWN 4 - -// ─── Physics & game constants -#define TICK_RATE (1.0f / 60.0f) -#define GRAVITY 980.0f // px/s² slope-gravity scale - -#define AGENT_RADIUS 12.0f -#define BOULDER_RADIUS 48.0f -#define GOAL_RADIUS 40.0f - -#define AGENT_MASS 1.0f -#define BOULDER_MASS 20.0f - -#define AGENT_ACCEL 1400.0f // px/s² per directional action -#define AGENT_DRAG 6.0f // per-second linear velocity decay -#define BOULDER_LINEAR_DRAG 0.4f -#define BOULDER_ROLLING_DRAG 0.8f // per-second angular velocity decay - -#define RESTITUTION_BODIES 0.05f // agent ↔ boulder (nearly inelastic) -#define RESTITUTION_WALL 0.15f -#define FRICTION_COEF 0.45f // Coulomb friction coefficient - -#define MAX_SPEED 1200.0f -#define MAX_ANGULAR_VEL 30.0f // rad/s normalisation for boulder spin -#define MAX_TICKS 600 -#define GOAL_CENTER_RADIUS 16.0f -#define GOAL_HOLD_TICKS 10 - -#define MAX_AGENTS 2 -// Per-agent obs layout: self(6) + boulder(8) + goal(5) + (MAX_AGENTS-1) other agents(8 each) -#define OBS_SIZE (6 + 8 + 5 + (MAX_AGENTS - 1) * 8) - -typedef struct Log { - float perf; - float score; - float episode_return; - float episode_length; - float n; -} Log; - -typedef struct Client { - float width; - float height; - Texture2D sprites; -} Client; - -// Circular rigid body — used for both agents and the boulder. -// All physics functions operate on this type. -typedef struct Entity { - float x, y; - float vx, vy; - float mass; - float radius; - float angle; // cumulative rotation (rad), increases CCW - float angular_vel; // angular velocity (rad/s), + = CCW - float inv_inertia; // 1/I; solid disk: I = ½·mass·radius² → inv = 2/(m·r²) -} Entity; - -typedef struct Boulder { - Client* client; - Log log; - Log* logs; - float* observations; - float* actions; - float* rewards; - float* terminals; - Entity boulder; - Entity* agents; - int num_agents; - float goal_x, goal_y; - int goal_hold; // consecutive ticks boulder has been within GOAL_RADIUS - float* heightmap; // flat [hmap_rows × hmap_cols] elevation values (px) - int hmap_cols; - int hmap_rows; - int width; - int height; - int tick; - float score; - unsigned int rng; - int* moving_boulder; - float dist_scale; -} Boulder; - -// Helper functions - -// Returns a uniform float in [0, 1) using rand_r seeded from env->rng. -static inline float randf(Boulder* env) { - return (float)rand_r(&env->rng) / ((float)RAND_MAX + 1.0f); -} - -static inline float clampf(float x, float lo, float hi) { - return x < lo ? lo : (x > hi ? hi : x); -} - -// Initialise a circular rigid body as a solid disk at rest. -// Precomputes inv_inertia = 2 / (mass · radius²). -static inline void entity_init(Entity* e,float x, float y,float mass, float radius) { - *e = (Entity){ - .x = x, .y = y, - .mass = mass, - .radius = radius, - .inv_inertia = 2.0f / (mass * radius * radius), - }; -} - -// Semi-implicit Euler integration with exponential velocity decay (drag). -// ax, ay : net acceleration (px/s²) this tick. -// linear_drag : linear damping rate (s⁻¹); applied as exp(−k·dt) factor. -// rolling_drag : angular damping rate (s⁻¹). -static inline void integrate_entity(Entity* e,float ax, float ay,float linear_drag,float rolling_drag) { - const float dt = TICK_RATE; - e->vx += ax * dt; - e->vy += ay * dt; - float ld = expf(-linear_drag * dt); - float rd = expf(-rolling_drag * dt); - e->vx *= ld; - e->vy *= ld; - e->x += e->vx * dt; - e->y += e->vy * dt; - e->angular_vel *= rd; - e->angle += e->angular_vel * dt; -} - -// ─── Height-map - -// Bilinear sample of the elevation map at world position (x, y). -// Returns 0 if heightmap is NULL or position is out of range. -static inline float heightmap_sample(const Boulder* env, float x, float y) { - if (!env->heightmap) return 0.0f; - float cx = (x / env->width) * (float)(env->hmap_cols - 1); - float cy = (y / env->height) * (float)(env->hmap_rows - 1); - int ix = (int)cx, iy = (int)cy; - if (ix < 0 || iy < 0 || ix >= env->hmap_cols - 1 || iy >= env->hmap_rows - 1) - return 0.0f; - float tx = cx - ix, ty = cy - iy; - int s = env->hmap_cols; - float h00 = env->heightmap[ iy * s + ix ]; - float h10 = env->heightmap[ iy * s + ix + 1]; - float h01 = env->heightmap[(iy+1) * s + ix ]; - float h11 = env->heightmap[(iy+1) * s + ix + 1]; - return h00*(1-tx)*(1-ty) + h10*tx*(1-ty) + h01*(1-tx)*ty + h11*tx*ty; -} - -// Computes slope-induced gravitational acceleration (px/s²) at world position -// (x, y) via central-difference gradient of the height field. -// Fills *out_ax and *out_ay with the downhill acceleration components. -static inline void heightmap_gravity(const Boulder* env, - float x, float y, - float* out_ax, float* out_ay) { - const float h = 4.0f; - float dh_dx = (heightmap_sample(env, x + h, y) - - heightmap_sample(env, x - h, y)) / (2.0f * h); - float dh_dy = (heightmap_sample(env, x, y + h) - - heightmap_sample(env, x, y - h)) / (2.0f * h); - *out_ax = -GRAVITY * dh_dx; - *out_ay = -GRAVITY * dh_dy; -} - -// ─── Sphere–sphere impulse resolution -// -// Convention: n is the unit normal FROM a TO b. -// vrel_n > 0 ⟹ contact points approaching ⟹ apply impulse. -// -// Normal restitution + Coulomb tangential friction. -// Both bodies are modified in place. -// -// Derivation notes (2D rigid-body impulse): -// • Contact arm: r_a = +a->radius·n, r_b = −b->radius·n -// • Velocity at contact: v_c = v_cm + ω × r → (vx − ω·ry , vy + ω·rx) -// • r_a × n = 0 ⟹ angular term drops out of the normal-impulse denominator. -// • cross2d(r_a, t) = +a->radius, cross2d(r_b, t) = −b->radius -// where t = (−ny, nx) is the CCW tangent. -static inline int resolve_sphere_sphere(Entity* a, Entity* b, - float restitution, float friction) { - float dx = b->x - a->x; - float dy = b->y - a->y; - float dist2 = dx*dx + dy*dy; - float rsum = a->radius + b->radius; - if (dist2 >= rsum * rsum || dist2 < 1e-8f) return 0; - - float dist = sqrtf(dist2); - float nx = dx / dist; // unit normal a → b - float ny = dy / dist; - float tx = -ny; // unit tangent (CCW 90° from n) - float ty = nx; - - // Positional correction: push apart proportional to inverse mass - float overlap = rsum - dist; - float inv_msum = 1.0f / (a->mass + b->mass); - a->x -= nx * overlap * (b->mass * inv_msum); - a->y -= ny * overlap * (b->mass * inv_msum); - b->x += nx * overlap * (a->mass * inv_msum); - b->y += ny * overlap * (a->mass * inv_msum); - - // Contact-point arm vectors (centre → contact point) - float ra_x = a->radius * nx, ra_y = a->radius * ny; - float rb_x = -b->radius * nx, rb_y = -b->radius * ny; - - // Velocity at each contact point: v_c = v_cm + ω × r (2D: vx−ω·ry, vy+ω·rx) - float vca_x = a->vx - a->angular_vel * ra_y; - float vca_y = a->vy + a->angular_vel * ra_x; - float vcb_x = b->vx - b->angular_vel * rb_y; - float vcb_y = b->vy + b->angular_vel * rb_x; - - float vrx = vca_x - vcb_x; - float vry = vca_y - vcb_y; - float vrel_n = vrx * nx + vry * ny; - if (vrel_n <= 0.0f) return 0; // separating; no impulse needed - - float vrel_t = vrx * tx + vry * ty; - - float inv_meff = 1.0f/a->mass + 1.0f/b->mass; - - // Normal impulse (angular inertia term is 0 for sphere–sphere) - float j_n = (1.0f + restitution) * vrel_n / inv_meff; - - // Tangential (friction) impulse including angular inertia in denominator - // cross2d(ra, t) = +a->radius; cross2d(rb, t) = −b->radius - float denom_t = inv_meff - + a->radius * a->radius * a->inv_inertia - + b->radius * b->radius * b->inv_inertia; - float j_t = vrel_t / denom_t; - float j_t_max = friction * j_n; - if (j_t > j_t_max) j_t = j_t_max; - if (j_t < -j_t_max) j_t = -j_t_max; - - // Linear impulse: impulse on a = −J, impulse on b = +J - float Jx = j_n * nx + j_t * tx; - float Jy = j_n * ny + j_t * ty; - a->vx -= Jx / a->mass; - a->vy -= Jy / a->mass; - b->vx += Jx / b->mass; - b->vy += Jy / b->mass; - - // Angular impulse: Δω = cross2d(r, F) · inv_I - // F on a = −J, F on b = +J - a->angular_vel += (ra_x * (-Jy) - ra_y * (-Jx)) * a->inv_inertia; - b->angular_vel += (rb_x * Jy - rb_y * Jx ) * b->inv_inertia; - - return 1; -} - -// ─── Sphere–horizontal-wall impulse resolution -// -// wall_y : y-coordinate of the wall line. -// normal_y : outward wall normal y-component (screen coords: y increases DOWN). -// +1 → top / ceiling wall (normal points downward toward entity). -// −1 → bottom / floor wall (normal points upward toward entity). -// -// Derivation: -// Contact arm rc = (0, −normal_y · radius) → toward wall face. -// x-velocity at contact: vx_c = vx + ω·(−rc_y) = vx + ω·normal_y·radius. -// cross2d(rc, t_x) = −rc_y = normal_y·radius. -static inline void resolve_sphere_hwall(Entity* e, - float wall_y, float normal_y, - float restitution, float friction) { - // Signed distance from wall to entity centre along outward normal - float sd = (e->y - wall_y) * normal_y; - float pen = e->radius - sd; - if (pen <= 0.0f) return; - - e->y += pen * normal_y; // push entity off wall - - float vn = e->vy * normal_y; // velocity along outward normal - if (vn >= 0.0f) return; // already leaving - - float j_n = -(1.0f + restitution) * vn * e->mass; // positive - - float rc_y = -normal_y * e->radius; // arm y-component - float vx_c = e->vx + e->angular_vel * (-rc_y); // x-vel at contact - float rc_ct = -rc_y; // cross2d(rc, (1,0)) - float denom_t = 1.0f/e->mass + rc_ct * rc_ct * e->inv_inertia; - float j_t = -vx_c / denom_t; - float j_t_max = friction * j_n; - if (j_t > j_t_max) j_t = j_t_max; - if (j_t < -j_t_max) j_t = -j_t_max; - - e->vy += j_n * normal_y / e->mass; - e->vx += j_t / e->mass; - // cross2d((0, rc_y), (j_t, j_n·normal_y)) = −rc_y·j_t - e->angular_vel += (-rc_y * j_t) * e->inv_inertia; -} - -// ─── Sphere–vertical-wall impulse resolution ───────────────────────────────── -// -// wall_x : x-coordinate of the wall line. -// normal_x : outward wall normal x-component. -// +1 → left wall (normal points right toward entity). -// −1 → right wall (normal points left toward entity). -// -// Derivation: -// Contact arm rc = (−normal_x · radius, 0). -// y-velocity at contact: vy_c = vy + ω·rc_x. -// cross2d(rc, t_y) = rc_x = −normal_x·radius. -static inline void resolve_sphere_vwall(Entity* e, - float wall_x, float normal_x, - float restitution, float friction) { - float sd = (e->x - wall_x) * normal_x; - float pen = e->radius - sd; - if (pen <= 0.0f) return; - - e->x += pen * normal_x; - - float vn = e->vx * normal_x; - if (vn >= 0.0f) return; - - float j_n = -(1.0f + restitution) * vn * e->mass; // positive - - float rc_x = -normal_x * e->radius; // arm x-component - float vy_c = e->vy + e->angular_vel * rc_x; // y-vel at contact - float rc_ct = rc_x; // cross2d(rc, (0,1)) - float denom_t = 1.0f/e->mass + rc_ct * rc_ct * e->inv_inertia; - float j_t = -vy_c / denom_t; - float j_t_max = friction * j_n; - if (j_t > j_t_max) j_t = j_t_max; - if (j_t < -j_t_max) j_t = -j_t_max; - - e->vx += j_n * normal_x / e->mass; - e->vy += j_t / e->mass; - // cross2d((rc_x, 0), (j_n·normal_x, j_t)) = rc_x·j_t - e->angular_vel += (rc_x * j_t) * e->inv_inertia; -} - - -void init(Boulder* env) { - env->tick = 0; - env->score = 0.0f; - env->goal_hold = 0; - env->agents = (Entity*)calloc(env->num_agents, sizeof(Entity)); - env->logs = (Log*)calloc(env->num_agents, sizeof(Log)); - env->moving_boulder = (int*)calloc(env->num_agents, sizeof(int)); -} - -void allocate(Boulder* env) { - init(env); - int n = env->num_agents; - env->observations = (float*)calloc(n * OBS_SIZE, sizeof(float)); - env->actions = (float*)calloc(n * 2, sizeof(float)); - env->rewards = (float*)calloc(n, sizeof(float)); - env->terminals = (float*)calloc(n, sizeof(float)); - env->agents = (Entity*)calloc(n, sizeof(Entity)); -} - -void c_close(Boulder* env) {} - -void free_allocated(Boulder* env) { - free(env->observations); - free(env->actions); - free(env->rewards); - free(env->terminals); - free(env->agents); - if (env->heightmap) free(env->heightmap); - c_close(env); -} - -void add_log(Boulder* env) { - for(int i= 0;inum_agents; i++){ - env->log.episode_length += env->logs[i].episode_length; - env->log.episode_return += env->logs[i].episode_return; - env->log.score += env->score; - env->log.perf += env->score; - env->log.n += 1; - } -} - -void compute_observations(Boulder* env) { - float diag = hypotf((float)env->width, (float)env->height); - float dist_scale_init = env->width > env->height ? env->width: env->height; - dist_scale_init *= env->dist_scale; - for (int i = 0; i < env->num_agents; i++) { - float* obs = env->observations + i * OBS_SIZE; - Entity* a = &env->agents[i]; - int idx = 0; - - // Self (4): position relative to walls and velocity - obs[idx++] = a->x / env->width; - obs[idx++] = (env->width - a->x) / env->width; - obs[idx++] = a->y / env->height; - obs[idx++] = (env->height - a->y) / env->height; - obs[idx++] = a->vx / MAX_SPEED; - obs[idx++] = a->vy / MAX_SPEED; - - // Boulder (8): ego-relative pos, dist, angle, velocity, spin - float bdx = env->boulder.x - a->x; - float bdy = env->boulder.y - a->y; - float bdist = hypotf(bdx, bdy); - float bang = atan2f(bdy, bdx); - obs[idx++] = clampf(bdx / dist_scale_init, -1.0f, 1.0f); - obs[idx++] = clampf(bdy / dist_scale_init, -1.0f, 1.0f); - obs[idx++] = clampf(bdist / diag, 0.0f, 1.0f); - obs[idx++] = sinf(bang); - obs[idx++] = cosf(bang); - obs[idx++] = env->boulder.vx / MAX_SPEED; - obs[idx++] = env->boulder.vy / MAX_SPEED; - obs[idx++] = env->boulder.angular_vel / MAX_ANGULAR_VEL; - - // Goal (5): ego-relative pos, dist, angle - float gdx = env->goal_x - env->boulder.x; - float gdy = env->goal_y - env->boulder.y; - float gdist = hypotf(gdx, gdy); - float gang = atan2f(gdy, gdx); - obs[idx++] = clampf(gdx / dist_scale_init, -1.0f, 1.0f); - obs[idx++] = clampf(gdy / dist_scale_init, -1.0f, 1.0f); - obs[idx++] = clampf(gdist / diag, 0.0f, 1.0f); - obs[idx++] = sinf(gang); - obs[idx++] = cosf(gang); - - // Other agents (8 each, up to MAX_AGENTS-1 slots, zero-padded if fewer) - int filled = 0; - for (int j = 0; j < env->num_agents && filled < MAX_AGENTS - 1; j++) { - if (j == i) continue; - Entity* o = &env->agents[j]; - float odx = o->x - a->x; - float ody = o->y - a->y; - float odist = hypotf(odx, ody); - float oang = atan2f(ody, odx); - obs[idx++] = clampf(odx / dist_scale_init, -1.0f, 1.0f); - obs[idx++] = clampf(ody / dist_scale_init, -1.0f, 1.0f); - obs[idx++] = clampf(odist / diag, 0.0f, 1.0f); - obs[idx++] = sinf(oang); - obs[idx++] = cosf(oang); - obs[idx++] = (o->vx - a->vx) / MAX_SPEED; - obs[idx++] = (o->vy - a->vy) / MAX_SPEED; - obs[idx++] = hypotf(o->vx, o->vy) / MAX_SPEED; - filled++; - } - while (filled++ < MAX_AGENTS - 1) { - for (int k = 0; k < 8; k++) obs[idx++] = 0.0f; - } - } -} - -void c_reset(Boulder* env) { - env->score = 0.0f; - env->tick = 0; - env->goal_hold = 0; - - // Boulder: random position with padding for its radius - float bx = BOULDER_RADIUS + randf(env) * (env->width - 2.0f * BOULDER_RADIUS); - float by = BOULDER_RADIUS + randf(env) * (env->height - 2.0f * BOULDER_RADIUS); - entity_init(&env->boulder, bx, by, BOULDER_MASS, BOULDER_RADIUS); - - // Agents: random positions outside boulder radius - for (int i = 0; i < env->num_agents; i++) { - env->logs[i] = (Log){0}; - float ax, ay; - do { - ax = AGENT_RADIUS + randf(env) * (env->width - 2.0f * AGENT_RADIUS); - ay = AGENT_RADIUS + randf(env) * (env->height - 2.0f * AGENT_RADIUS); - } while (hypotf(ax - bx, ay - by) < BOULDER_RADIUS + AGENT_RADIUS); - entity_init(&env->agents[i], ax, ay, AGENT_MASS, AGENT_RADIUS); - } - - // Goal: random position outside boulder radius * 3 - float gx, gy; - do { - gx = GOAL_RADIUS + randf(env) * (env->width - 2.0f * GOAL_RADIUS); - gy = GOAL_RADIUS + randf(env) * (env->height - 2.0f * GOAL_RADIUS); - } while (hypotf(gx - bx, gy - by) < BOULDER_RADIUS * 3.0f); - env->goal_x = gx; - env->goal_y = gy; - - memset(env->moving_boulder, 0, 2*sizeof(int)); - compute_observations(env); -} - -void c_step(Boulder* env) { - float dist_to_goal_before = hypotf(env->boulder.x - env->goal_x, env->boulder.y - env->goal_y); - // 8-direction unit vectors (screen coords: y increases down) - // index: 0=E, 1=NE, 2=N, 3=NW, 4=W, 5=SW, 6=S, 7=SE - static const float DIR_AX[8] = { 1.0f, 0.707f, 0.0f, -0.707f, -1.0f, -0.707f, 0.0f, 0.707f}; - static const float DIR_AY[8] = { 0.0f, -0.707f, -1.0f, -0.707f, 0.0f, 0.707f, 1.0f, 0.707f}; - - // Agent acceleration from discrete actions + height-map slope - for (int i = 0; i < env->num_agents; i++) { - int throttle = (int)env->actions[i * 2 + 0]; // 0=off, 1=on - int dir = (int)env->actions[i * 2 + 1]; // 0..7 - float ax = 0.0f, ay = 0.0f; - if (throttle && dir >= 0 && dir < 8) { - ax = DIR_AX[dir] * AGENT_ACCEL; - ay = DIR_AY[dir] * AGENT_ACCEL; - } - - float hgx = 0.0f, hgy = 0.0f; - heightmap_gravity(env, env->agents[i].x, env->agents[i].y, &hgx, &hgy); - integrate_entity(&env->agents[i], ax + hgx, ay + hgy, - AGENT_DRAG, 0.0f); - } - - // Boulder integration (height-map slope only; agents push via collision) - { - float hgx = 0.0f, hgy = 0.0f; - heightmap_gravity(env, env->boulder.x, env->boulder.y, &hgx, &hgy); - integrate_entity(&env->boulder, hgx, hgy,BOULDER_LINEAR_DRAG, BOULDER_ROLLING_DRAG); - } - - // Collision: agents ↔ boulder - for (int i = 0; i < env->num_agents; i++) { - int collided = resolve_sphere_sphere(&env->agents[i], &env->boulder,RESTITUTION_BODIES, FRICTION_COEF); - env->moving_boulder[i] = 1.0f; - } - - // Collision: agents ↔ agents - for (int i = 0; i < env->num_agents; i++) { - for (int j = i + 1; j < env->num_agents; j++) { - resolve_sphere_sphere(&env->agents[i], &env->agents[j],RESTITUTION_BODIES, FRICTION_COEF); - } - } - - // All entities ↔ map boundary walls (sphere-on-line) - float wx0 = 0.0f; - float wy0 = 0.0f; - float wx1 = env->width; - float wy1 = env->height; - - for (int i = 0; i < env->num_agents; i++) { - resolve_sphere_vwall(&env->agents[i], wx0, +1.0f, RESTITUTION_WALL, FRICTION_COEF); - resolve_sphere_vwall(&env->agents[i], wx1, -1.0f, RESTITUTION_WALL, FRICTION_COEF); - resolve_sphere_hwall(&env->agents[i], wy0, +1.0f, RESTITUTION_WALL, FRICTION_COEF); - resolve_sphere_hwall(&env->agents[i], wy1, -1.0f, RESTITUTION_WALL, FRICTION_COEF); - } - resolve_sphere_vwall(&env->boulder, wx0, +1.0f, RESTITUTION_WALL, FRICTION_COEF); - resolve_sphere_vwall(&env->boulder, wx1, -1.0f, RESTITUTION_WALL, FRICTION_COEF); - resolve_sphere_hwall(&env->boulder, wy0, +1.0f, RESTITUTION_WALL, FRICTION_COEF); - resolve_sphere_hwall(&env->boulder, wy1, -1.0f, RESTITUTION_WALL, FRICTION_COEF); - - // Goal: reward every tick boulder overlaps goal or approaches goal - float dist_to_goal = hypotf(env->boulder.x - env->goal_x, env->boulder.y - env->goal_y); - int on_goal = (dist_to_goal < GOAL_CENTER_RADIUS); - float boulder_move_reward = 0.01f*(dist_to_goal_before - dist_to_goal); - - for (int i = 0; i < env->num_agents; i++){ - float r = on_goal ? 1.0f : env->moving_boulder[i] ? boulder_move_reward : 0.0f; - env->rewards[i] += r; - env->logs[i].episode_return += r; - - } - if (on_goal) { - env->goal_hold++; - } else { - env->goal_hold = 0; - } - - env->tick++; - int goal_held_for_full = (env->goal_hold >= GOAL_HOLD_TICKS); - if(goal_held_for_full){ - env->score += 1.0f; - } - int done = goal_held_for_full || (env->tick >= MAX_TICKS); - if (done) { - for (int i = 0; i < env->num_agents; i++){ - env->terminals[i] = 1.0f; - env->logs[i].episode_length = env->tick; - } - add_log(env); - c_reset(env); - } - - compute_observations(env); -} - -// Rendering - -Client* make_client(Boulder* env) { - Client* client = (Client*)calloc(1, sizeof(Client)); - client->width = env->width; - client->height = env->height; - InitWindow(env->width, env->height, "PufferLib Boulder"); - SetTargetFPS(60); - client->sprites = LoadTexture("resources/shared/puffers.png"); - return client; -} - -void close_client(Client* client) { - UnloadTexture(client->sprites); - CloseWindow(); - free(client); -} - -void c_render(Boulder* env) { - if (env->client == NULL) env->client = make_client(env); - if (IsKeyDown(KEY_ESCAPE)) exit(0); - if (IsKeyPressed(KEY_TAB)) ToggleFullscreen(); - - BeginDrawing(); - ClearBackground((Color){6, 24, 24, 255}); - - // Goal ring - DrawCircle((int)env->goal_x, (int)env->goal_y, (int)GOAL_RADIUS, - (Color){200, 200, 0, 60}); - DrawCircleLines((int)env->goal_x, (int)env->goal_y, (int)GOAL_RADIUS, YELLOW); - - // Boulder: large dark sphere with rotation indicator line - DrawCircle((int)env->boulder.x, (int)env->boulder.y, - (int)env->boulder.radius, (Color){120, 75, 30, 255}); - DrawCircleLines((int)env->boulder.x, (int)env->boulder.y, - (int)env->boulder.radius, (Color){180, 120, 60, 255}); - { - float lx = env->boulder.x + cosf(env->boulder.angle) * env->boulder.radius * 0.75f; - float ly = env->boulder.y + sinf(env->boulder.angle) * env->boulder.radius * 0.75f; - DrawLine((int)env->boulder.x, (int)env->boulder.y, - (int)lx, (int)ly, (Color){220, 180, 100, 255}); - } - - // Agents: sprites from puffers.png, tinted per agent index, with velocity arrow - // Sprite sheet row: y=576 facing left, y=608 facing right (32px rows) - // Use sprite column 0 for all agents; differentiate via tint color - for (int i = 0; i < env->num_agents; i++) { - Entity* a = &env->agents[i]; - int src_x = 32 * i; // different column = different color - int src_y = (a->vx >= 0.0f) ? 576 : 608; // facing right or left - DrawTexturePro( - env->client->sprites, - (Rectangle){ src_x, src_y, 32, 32 }, - (Rectangle){ a->x - 16, a->y - 16, 32, 32 }, - (Vector2){0, 0}, - 0, - WHITE - ); - float speed = hypotf(a->vx, a->vy); - if (speed > 20.0f) { - float dx = a->vx / speed * (a->radius + 8.0f); - float dy = a->vy / speed * (a->radius + 8.0f); - DrawLine((int)a->x, (int)a->y, - (int)(a->x + dx), (int)(a->y + dy), WHITE); - } - } - - EndDrawing(); -} From f79d811832a1e2898560357f4db893e03f8af5db Mon Sep 17 00:00:00 2001 From: l1onh3art88 Date: Sat, 16 May 2026 17:16:26 -0500 Subject: [PATCH 8/8] one more rm --- config/boulder.ini | 57 ---------------------------------------------- 1 file changed, 57 deletions(-) delete mode 100644 config/boulder.ini diff --git a/config/boulder.ini b/config/boulder.ini deleted file mode 100644 index 6b5b87873a..0000000000 --- a/config/boulder.ini +++ /dev/null @@ -1,57 +0,0 @@ -[base] -env_name = boulder - -[vec] -total_agents = 1024 -num_buffers = 2.31152 -num_threads = 16 - -[env] -num_agents = 2 -width = 900 -height = 600 -dist_scale = 0.5 - -[policy] -hidden_size = 256 -num_layers = 3.63347 -expansion_factor = 1 - -[train] -gpus = 1 -seed = 42 -total_timesteps = 77332504 -learning_rate = 0.00762627 -anneal_lr = 1 -min_lr_ratio = 0 -gamma = 0.993349 -gae_lambda = 0.97086 -replay_ratio = 2.76336 -clip_coef = 0.414034 -vf_coef = 2.90403 -vf_clip_coef = 0.01 -max_grad_norm = 1.19707 -ent_coef = 0.0020498 -beta1 = 0.937322 -beta2 = 0.999958 -eps = 2.42555e-12 -minibatch_size = 4096 -horizon = 128 -vtrace_rho_clip = 2.2155 -vtrace_c_clip = 1.21309 -prio_alpha = 0.24272 -prio_beta0 = 0.366241 - -[sweep.train.total_timesteps] -distribution = log_normal -min = 5e7 -max = 5e8 -mean = 1e8 -scale = auto - -[sweep.env.dist_scale] -distribution = uniform -max = 1.0 -min = 0.0 -mean = 0.5 -scale = auto