Add Chain Reaction env

build.sh

@@ -238,8 +238,13 @@ done
 export CCACHE_DIR="${CCACHE_DIR:-$HOME/.ccache}"
 export CCACHE_BASEDIR="$(pwd)"
 export CCACHE_COMPILERCHECK=content
-NVCC="ccache $CUDA_HOME/bin/nvcc"
-CC="${CC:-$(command -v ccache >/dev/null && echo 'ccache clang' || echo 'clang')}"
+if command -v ccache >/dev/null 2>&1; then
+    NVCC="ccache $CUDA_HOME/bin/nvcc"
+    CC="${CC:-ccache clang}"
+else
+    NVCC="$CUDA_HOME/bin/nvcc"
+    CC="${CC:-clang}"
+fi
 ARCH=${NVCC_ARCH:-native}
 
 PYTHON_INCLUDE=$(python -c "import sysconfig; print(sysconfig.get_path('include'))")

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,125 @@
+#include "chain_reaction.h"
+
+#define OBS_SIZE CHAINENV_OBS_SIZE
+#define NUM_ATNS 1
+#define ACT_SIZES {CHAINENV_MAX_ACTIONS}
+#define OBS_TENSOR_T FloatTensor
+#define MY_ACTION_MASK CHAINENV_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;
+    }
+}
+
+static int chainenv_binding_num_agents(Dict* kwargs) {
+    int num_agents = (int)dict_get(kwargs, "num_agents")->value;
+    if (num_agents < 1) num_agents = 1;
+    if (num_agents > CHAINENV_MAX_SLOTS) num_agents = CHAINENV_MAX_SLOTS;
+    return num_agents;
+}
+
+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;
+    int agents_per_env = chainenv_binding_num_agents(env_kwargs);
+    if (total_agents % agents_per_env != 0) {
+        fprintf(stderr,
+            "chain_reaction requires total_agents (%d) divisible by num_agents (%d)\n",
+            total_agents, agents_per_env);
+        exit(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++) {
+        envs[i].rng = i;
+        my_init(&envs[i], env_kwargs);
+    }
+
+    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_init(Env* env, Dict* kwargs) {
+    env->num_agents = chainenv_binding_num_agents(kwargs);
+    env->rows = (int)dict_get(kwargs, "rows")->value;
+    env->cols = (int)dict_get(kwargs, "cols")->value;
+    env->max_steps = (int)dict_get(kwargs, "max_steps")->value;
+    env->opponent_policy = (int)dict_get(kwargs, "opponent_policy")->value;
+    env->win_reward = (float)dict_get(kwargs, "win_reward")->value;
+    env->loss_reward = (float)dict_get(kwargs, "loss_reward")->value;
+    env->invalid_move_reward = (float)dict_get(kwargs, "invalid_move_reward")->value;
+
+    if (env->rows < 2) env->rows = 2;
+    if (env->cols < 2) env->cols = 2;
+    if (env->rows > CHAINENV_MAX_ROWS) env->rows = CHAINENV_MAX_ROWS;
+    if (env->cols > CHAINENV_MAX_COLS) env->cols = CHAINENV_MAX_COLS;
+    if (env->max_steps < 1) env->max_steps = 1;
+
+    init_chainenv(env);
+}
+
+void my_log(Log* log, Dict* out) {
+    float n = log->n > 0.0f ? log->n : 1.0f;
+    dict_set(out, "perf", log->perf / n);
+    dict_set(out, "score", log->score / n);
+    dict_set(out, "episode_return", log->episode_return / n);
+    dict_set(out, "episode_length", log->episode_length / n);
+    dict_set(out, "chain_bursts", log->chain_bursts / n);
+    dict_set(out, "invalid_rate", log->invalid_rate / n);
+    dict_set(out, "hist_score", log->hist_score / n);
+    dict_set(out, "hist_n", log->hist_n / n);
+    dict_set(out, "hist_score_bank_0", log->hist_score_bank[0] / n);
+    dict_set(out, "hist_score_bank_1", log->hist_score_bank[1] / n);
+    dict_set(out, "hist_score_bank_2", log->hist_score_bank[2] / n);
+    dict_set(out, "hist_score_bank_3", log->hist_score_bank[3] / n);
+    dict_set(out, "hist_score_bank_4", log->hist_score_bank[4] / n);
+    dict_set(out, "hist_score_bank_5", log->hist_score_bank[5] / n);
+    dict_set(out, "hist_score_bank_6", log->hist_score_bank[6] / n);
+    dict_set(out, "hist_score_bank_7", log->hist_score_bank[7] / n);
+    dict_set(out, "hist_n_bank_0", log->hist_n_bank[0] / n);
+    dict_set(out, "hist_n_bank_1", log->hist_n_bank[1] / n);
+    dict_set(out, "hist_n_bank_2", log->hist_n_bank[2] / n);
+    dict_set(out, "hist_n_bank_3", log->hist_n_bank[3] / n);
+    dict_set(out, "hist_n_bank_4", log->hist_n_bank[4] / n);
+    dict_set(out, "hist_n_bank_5", log->hist_n_bank[5] / n);
+    dict_set(out, "hist_n_bank_6", log->hist_n_bank[6] / n);
+    dict_set(out, "hist_n_bank_7", log->hist_n_bank[7] / n);
+    dict_set(out, "slot_0_score", log->slot_0_score / n);
+    dict_set(out, "slot_1_score", log->slot_1_score / n);
+    dict_set(out, "draw_rate", log->draw_rate / n);
+    dict_set(out, "n", log->n);
+}

ocean/chain_reaction/chain_reaction.c

@@ -0,0 +1,65 @@
+#include <time.h>
+
+#include "chain_reaction.h"
+
+static void demo(void) {
+    ChainEnv env = {0};
+    env.num_agents = 1;
+    env.rows = 9;
+    env.cols = 6;
+    env.max_steps = 132;
+    env.opponent_policy = CHAINENV_OPPONENT_HEURISTIC;
+    env.win_reward = 1.0f;
+    env.loss_reward = -1.0f;
+    env.invalid_move_reward = -1.0f;
+    env.rng = (unsigned int)time(NULL);
+    init_chainenv(&env);
+
+    float observation_buf[CHAINENV_OBS_SIZE] = {0};
+    float action_buf[1] = {0};
+    float reward_buf[1] = {0};
+    float terminal_buf[1] = {0};
+    unsigned char action_mask_buf[CHAINENV_MAX_ACTIONS] = {0};
+    env.observations = observation_buf;
+    env.actions = action_buf;
+    env.rewards = reward_buf;
+    env.terminals = terminal_buf;
+    env.action_mask = action_mask_buf;
+    env.client = make_client();
+    c_reset(&env);
+
+    while (true) {
+        if (env.client != NULL) {
+            chainenv_layout_board(env.client, &env);
+        }
+
+        bool allow_input = !chainenv_animation_active(&env);
+        if (IsMouseButtonPressed(MOUSE_LEFT_BUTTON)) {
+            if (env.end_game) {
+                c_reset(&env);
+            } else if (allow_input && env.client != NULL) {
+                int action = chainenv_pick_action(&env, env.client, GetMousePosition());
+                if (action >= 0) {
+                    if (!chainenv_is_legal_move(&env, action, CHAINENV_PLAYER_RED)) {
+                        chainenv_show_invalid_hint(&env);
+                    } else {
+                        chainenv_clear_invalid_hint(&env);
+                        env.actions[0] = (float)action;
+                        c_step(&env);
+                    }
+                }
+            }
+        }
+
+        if (IsKeyPressed(KEY_R)) {
+            c_reset(&env);
+        }
+
+        c_render(&env);
+    }
+}
+
+int main(void) {
+    demo();
+    return 0;
+}