Add Chain Reaction env

config/chain_reaction.ini

@@ -0,0 +1,61 @@
+[base]
+env_name = chain_reaction
+
+[selfplay]
+enabled = 1
+max_size = 16
+swap_winrate = 0.55
+min_games = 512
+elo_init = 0.0
+elo_k = 16.0
+seed = 42
+snapshot_interval = 10_000_000
+opp_timeout_steps = 20_000_000
+
+[vec]
+total_agents = 1024
+num_buffers = 4
+num_threads = 2
+num_frozen_banks = 1
+frozen_bank_pct = 0.1
+
+[env]
+num_agents = 2
+rows = 9
+cols = 6
+max_steps = 132
+opponent_policy = 1
+win_reward = 1.0
+loss_reward = -1.0
+invalid_move_reward = -1.0
+
+[policy]
+hidden_size = 256
+num_layers = 1
+expansion_factor = 1
+
+[train]
+gpus = 1
+seed = 42
+total_timesteps = 10000000
+learning_rate = 0.001
+anneal_lr = 1
+min_lr_ratio = 0
+gamma = 0.99
+gae_lambda = 0.95
+replay_ratio = 2
+clip_coef = 0.2
+vf_coef = 1.0
+vf_clip_coef = 1.0
+max_grad_norm = 0.5
+ent_coef = 0.01
+beta1 = 0.9
+beta2 = 0.999
+eps = 1e-08
+minibatch_size = 4096
+horizon = 32
+vtrace_rho_clip = 1.0
+vtrace_c_clip = 1.0
+prio_alpha = 0.0
+prio_beta0 = 0.0
+use_rnn = 1

ocean/chain_reaction/binding.c

@@ -0,0 +1,121 @@
+#include <unistd.h>
+
+#include "chain_reaction.h"
+
+#define NUM_ATNS 1
+#define ACT_SIZES {MAX_ACTIONS}
+#define OBS_TENSOR_T FloatTensor
+#define MY_ACTION_MASK MAX_ACTIONS
+#define MY_VEC_INIT
+#define MY_USES_PERM
+#define MY_USES_TAGS
+
+#define Env ChainEnv
+#include "vecenv.h"
+
+void my_setup_perm(StaticVec* vec, Env* env, int slot_base) {
+    size_t obs_elem_size = obs_element_size();
+    for (int slot = 0; slot < env->num_agents; slot++) {
+        int phys = vec->agent_perm != NULL
+            ? vec->agent_perm[slot_base + slot]
+            : slot_base + slot;
+        env->obs_ptr[slot] = (float*)((char*)vec->observations
+            + (size_t)phys * OBS_SIZE * obs_elem_size);
+        env->action_ptr[slot] = vec->actions + (size_t)phys * NUM_ATNS;
+        env->reward_ptr[slot] = vec->rewards + phys;
+        env->terminal_ptr[slot] = vec->terminals + phys;
+        env->action_mask_ptr[slot] = vec->action_mask + (size_t)phys * MY_ACTION_MASK;
+    }
+}
+
+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 = dict_get(vec_kwargs, "total_agents")->value;
+    int num_buffers = dict_get(vec_kwargs, "num_buffers")->value;
+    int agents_per_env = dict_get(env_kwargs, "num_agents")->value;
+    if (agents_per_env < 1 || agents_per_env > MAX_SLOTS) {
+        fprintf(stderr, "chain_reaction num_agents must be 1 or 2, got %d\n", agents_per_env);
+        exit(1);
+    }
+    if (total_agents < 1 || num_buffers < 1 || total_agents % num_buffers != 0) {
+        fprintf(stderr,
+            "chain_reaction requires positive total_agents (%d) divisible by num_buffers (%d)\n",
+            total_agents, num_buffers);
+        exit(1);
+    }
+
+    int agents_per_buffer = total_agents / num_buffers;
+    if (agents_per_buffer % agents_per_env != 0) {
+        fprintf(stderr,
+            "chain_reaction requires %d agents per buffer divisible by num_agents (%d)\n",
+            agents_per_buffer, agents_per_env);
+        exit(1);
+    }
+
+    int num_envs = total_agents / agents_per_env;
+    int envs_per_buffer = agents_per_buffer / agents_per_env;
+    Env* envs = calloc(num_envs, sizeof(Env));
+    unsigned int process_rng = (unsigned int)getpid();
+    for (int i = 0; i < num_envs; i++) {
+        envs[i].rng = rand_r(&process_rng);
+        my_init(&envs[i], env_kwargs);
+    }
+    for (int buffer = 0; buffer < num_buffers; buffer++) {
+        buffer_env_starts[buffer] = buffer * envs_per_buffer;
+        buffer_env_counts[buffer] = envs_per_buffer;
+    }
+
+    *num_envs_out = num_envs;
+    return envs;
+}
+
+void my_init(Env* env, Dict* kwargs) {
+    env->num_agents = dict_get(kwargs, "num_agents")->value;
+    env->rows = dict_get(kwargs, "rows")->value;
+    env->cols = dict_get(kwargs, "cols")->value;
+    env->max_steps = dict_get(kwargs, "max_steps")->value;
+    env->opponent_policy = dict_get(kwargs, "opponent_policy")->value;
+    env->win_reward = dict_get(kwargs, "win_reward")->value;
+    env->loss_reward = dict_get(kwargs, "loss_reward")->value;
+    env->invalid_move_reward = dict_get(kwargs, "invalid_move_reward")->value;
+
+    if (env->rows < 2 || env->rows > MAX_ROWS ||
+            env->cols < 2 || env->cols > MAX_COLS || env->max_steps < 1) {
+        fprintf(stderr,
+            "chain_reaction requires rows 2-%d, cols 2-%d, and positive max_steps; got %d, %d, %d\n",
+            MAX_ROWS, MAX_COLS, env->rows, env->cols, env->max_steps);
+        exit(1);
+    }
+
+    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);
+    dict_set(out, "chain_bursts", log->chain_bursts);
+    dict_set(out, "invalid_rate", log->invalid_rate);
+    dict_set(out, "hist_score", log->hist_score);
+    dict_set(out, "hist_n", log->hist_n);
+    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, "slot_0_score", log->slot_0_score);
+    dict_set(out, "slot_1_score", log->slot_1_score);
+    dict_set(out, "draw_rate", log->draw_rate);
+}

ocean/chain_reaction/chain_reaction.c

@@ -0,0 +1,67 @@
+#include <time.h>
+#include <unistd.h>
+
+#include "chain_reaction.h"
+
+int main(void) {
+    ChainEnv env = {0};
+    env.num_agents = 1;
+    env.rows = 9;
+    env.cols = 6;
+    env.max_steps = 132;
+    env.opponent_policy = HEURISTIC_OPPONENT;
+    env.win_reward = 1.0f;
+    env.loss_reward = -1.0f;
+    env.invalid_move_reward = -1.0f;
+    env.rng = (unsigned int)time(NULL) ^ (unsigned int)getpid();
+    init(&env);
+
+    float observation_buf[OBS_SIZE] = {0};
+    float action_buf[1] = {0};
+    float reward_buf[1] = {0};
+    float terminal_buf[1] = {0};
+    unsigned char action_mask_buf[MAX_ACTIONS] = {0};
+    env.obs_ptr[0] = observation_buf;
+    env.action_ptr[0] = action_buf;
+    env.reward_ptr[0] = reward_buf;
+    env.terminal_ptr[0] = terminal_buf;
+    env.action_mask_ptr[0] = action_mask_buf;
+    c_render(&env);
+    c_reset(&env);
+
+    while (true) {
+        bool allow_input = env.client->animation_clock
+            >= env.client->animation_total - 1.0e-4f;
+        if (IsMouseButtonPressed(MOUSE_LEFT_BUTTON)) {
+            if (env.end_game) {
+                c_reset(&env);
+            } else if (allow_input) {
+                Vector2 mouse = GetMousePosition();
+                float board_w = env.client->cell_size * env.cols;
+                float board_h = env.client->cell_size * env.rows;
+                bool on_board = mouse.x >= env.client->board_x
+                    && mouse.x < env.client->board_x + board_w
+                    && mouse.y >= env.client->board_y
+                    && mouse.y < env.client->board_y + board_h;
+                if (on_board) {
+                    int col = (int)((mouse.x - env.client->board_x) / env.client->cell_size);
+                    int row = (int)((mouse.y - env.client->board_y) / env.client->cell_size);
+                    int action = row * env.cols + col;
+                    if (!is_legal_move(&env, action, RED_PLAYER)) {
+                        env.client->invalid_hint_time = 1.9f;
+                    } else {
+                        env.client->invalid_hint_time = 0.0f;
+                        *env.action_ptr[0] = (float)action;
+                        c_step(&env);
+                    }
+                }
+            }
+        }
+
+        if (IsKeyPressed(KEY_R)) {
+            c_reset(&env);
+        }
+
+        c_render(&env);
+    }
+}