From ea01b51de10db413d31f52cf1f37c319ac446eb2 Mon Sep 17 00:00:00 2001 From: Kinvert Date: Mon, 17 Aug 2026 00:52:33 -0400 Subject: [PATCH 1/4] Affine Lock Cleanup --- ocean/affine_lock/README.md | 16 +- ocean/affine_lock/affine_lock.c | 33 +- ocean/affine_lock/affine_lock.h | 474 ++++++++---------- .../affine_lock/affine_lock_visible_targets.h | 161 ++---- ocean/affine_lock/tests/run_all.sh | 13 +- ocean/affine_lock/tests/test_affine_lock.c | 278 +++++----- .../tests/test_affine_lock_log_export.c | 24 +- .../affine_lock/tests/test_metadata_smoke.py | 12 +- .../tests/test_visible_targets_loader.c | 15 +- src/puffercpu.c | 3 + 10 files changed, 435 insertions(+), 594 deletions(-) diff --git a/ocean/affine_lock/README.md b/ocean/affine_lock/README.md index 5fcb2b530f..3ed5b7f2de 100644 --- a/ocean/affine_lock/README.md +++ b/ocean/affine_lock/README.md @@ -30,13 +30,13 @@ manifest identify this exact transform set as `affine_lock_8action_v1`: | `7` | `reverse_each_byte` | reverse bit order within each byte | The Puffer binding exposes one discrete action slot with -`AFFINE_LOCK_NUM_ACTIONS = 8`. +`NUM_ACTIONS = 8`. ## Resets Resets always sample from the committed visible-target table. To train or test on different target distributions, generate a new table with the tool below and -point `AFFINE_LOCK_VISIBLE_TARGET_TABLE_PATH` at it when building. +point `VISIBLE_TARGET_TABLE_PATH` at it when building. ## Committed Target Table @@ -105,7 +105,7 @@ To train against a custom 8-action table, either write it to the default path or build with an explicit table path: ```bash -EXTRA_CFLAGS='-DAFFINE_LOCK_VISIBLE_TARGET_TABLE_PATH="/tmp/affine_lock_8action_visible_targets_seed42.bin"' \ +EXTRA_CFLAGS='-DVISIBLE_TARGET_TABLE_PATH="/tmp/affine_lock_8action_visible_targets_seed42.bin"' \ ./build.sh affine_lock ``` @@ -169,7 +169,7 @@ This is generator-only. The committed runtime environment does not train on this action set. It is kept as a small, explicit alternate because a four-action policy can be easier to learn, and this graph has far more unique depth-16 pairs than the committed 8-action table. To make it a runtime environment, -update the env action table, `AFFINE_LOCK_NUM_ACTIONS`, the visible-table +update the env action table, `NUM_ACTIONS`, the visible-table action-set hash/path, generated table artifact, and any policy/config expectations that assume eight actions. @@ -200,10 +200,10 @@ Example generation command: The 4-action table is not plug-compatible with the committed 8-action runtime. To make a real 4-action runtime variant: -1. Change `AFFINE_LOCK_NUM_ACTIONS` to `4`. +1. Change `NUM_ACTIONS` to `4`. 2. Change the runtime action enum/table in `affine_lock.h` to match the generator's `affine_lock_4action_v1` order. -3. Point `AFFINE_LOCK_VISIBLE_TARGET_TABLE_PATH` at a 4-action table. +3. Point `VISIBLE_TARGET_TABLE_PATH` at a 4-action table. 4. Update the expected action-set hash in `affine_lock_visible_targets.h` to the 4-action manifest's `action_set_hash`. 5. Remove runtime helpers and render labels that only exist for the old @@ -225,8 +225,8 @@ would need to: 1. Add the depth to `TARGET_DEPTHS` in `tools/generate_8action_visible_targets.c`. 2. Regenerate the `.bin` and `.json`. -3. Add the depth to `AFFINE_LOCK_CURRICULUM_DEPTHS` and update - `AFFINE_LOCK_CURRICULUM_DEPTH_COUNT`. +3. Add the depth to `CURRICULUM_DEPTHS` and update + `CURRICULUM_DEPTH_COUNT`. 4. Add matching `Log.depth_D_rate` and `Log.depth_D_solve_rate` fields plus `my_log` exports if the depth should appear in training logs. 5. Update config/docs/tests to expect the new depth and record count. diff --git a/ocean/affine_lock/affine_lock.c b/ocean/affine_lock/affine_lock.c index 9d8389b013..21229ee3e3 100644 --- a/ocean/affine_lock/affine_lock.c +++ b/ocean/affine_lock/affine_lock.c @@ -1,38 +1,15 @@ -#include -#include #include #include "affine_lock.h" -static AffineLock* g_env = NULL; - -static void demo_cleanup(void) { - if (g_env == NULL) { - return; - } - free(g_env->agents[0].observations); - free(g_env->agents[0].actions); - free(g_env->agents[0].rewards); - free(g_env->agents[0].terminals); - puf_close(g_env); - g_env = NULL; -} - int main(void) { AffineLock env; memset(&env, 0, sizeof(env)); - g_env = &env; - atexit(demo_cleanup); - env.agents[0].observations = calloc(AFFINE_LOCK_OBS_SIZE, sizeof(float)); - env.agents[0].actions = (float*)calloc(AFFINE_LOCK_NUM_ATNS, sizeof(float)); + env.agents[0].observations = (float*)calloc(OBS_SIZE, sizeof(float)); + env.agents[0].actions = (float*)calloc(NUM_ATNS, sizeof(float)); env.agents[0].rewards = (float*)calloc(1, sizeof(float)); env.agents[0].terminals = (float*)calloc(1, sizeof(float)); - if (env.agents[0].observations == NULL || env.agents[0].actions == NULL || - env.agents[0].rewards == NULL || env.agents[0].terminals == NULL) { - fprintf(stderr, "failed to allocate affine_lock demo buffers\n"); - return 1; - } Dict kwargs = {0}; dict_set(&kwargs, "start_depth", 2); @@ -50,6 +27,10 @@ int main(void) { puf_render(&env); } - demo_cleanup(); + free(env.agents[0].observations); + free(env.agents[0].actions); + free(env.agents[0].rewards); + free(env.agents[0].terminals); + puf_close(&env); return 0; } diff --git a/ocean/affine_lock/affine_lock.h b/ocean/affine_lock/affine_lock.h index f9416ce27d..e7c48a184b 100644 --- a/ocean/affine_lock/affine_lock.h +++ b/ocean/affine_lock/affine_lock.h @@ -7,39 +7,33 @@ typedef float obs_t; #include "pufferenv.h" #include "affine_lock_visible_targets.h" -#define AFFINE_LOCK_BITS 16 -#define AFFINE_LOCK_TIMER_INDEX (2 * AFFINE_LOCK_BITS) -#define AFFINE_LOCK_OBS_SIZE (AFFINE_LOCK_TIMER_INDEX + 1) -#define AFFINE_LOCK_NUM_ATNS 1 -#define AFFINE_LOCK_NUM_ACTIONS 8 -#define AFFINE_LOCK_MAX_SOLUTION_DEPTH 16 -#define AFFINE_LOCK_CURRICULUM_DEPTH_COUNT 6 -#define AFFINE_LOCK_STEP_REWARD (-0.01f) -#ifndef AFFINE_LOCK_VISIBLE_TARGET_TABLE_PATH -#define AFFINE_LOCK_VISIBLE_TARGET_TABLE_PATH \ - "ocean/affine_lock/generated/affine_lock_8action_visible_targets.bin" -#endif - -#define ACT_SIZES {AFFINE_LOCK_NUM_ACTIONS} -#define OBS_SIZE AFFINE_LOCK_OBS_SIZE -#define NUM_ATNS AFFINE_LOCK_NUM_ATNS -#define PUF_STEPS_PER_SEC 2 +#define BITS 16 +#define TIMER_INDEX (2 * BITS) +#define OBS_SIZE (TIMER_INDEX + 1) +#define NUM_ATNS 1 +#define NUM_ACTIONS 8 +#define MAX_SOLUTION_DEPTH 16 +#define CURRICULUM_DEPTH_COUNT 6 +#define STEP_REWARD (-0.01f) // TODO should this be in ini so it can be swept? +#define VISIBLE_TARGET_TABLE_PATH "ocean/affine_lock/generated/affine_lock_8action_visible_targets.bin" +#define ACT_SIZES {NUM_ACTIONS} +#define PUF_STEPS_PER_SEC 2 // TODO remove this? #define MY_VEC_INIT #define MY_VEC_CLOSE -static const int AFFINE_LOCK_CURRICULUM_DEPTHS[ - AFFINE_LOCK_CURRICULUM_DEPTH_COUNT] = {2, 4, 5, 6, 8, 16}; +// TODO should this be in ini file? So it doesn't need a build to change it. +static const int CURRICULUM_DEPTHS[CURRICULUM_DEPTH_COUNT] = {2, 4, 5, 6, 8, 16}; typedef enum AffineLockAction { - AFFINE_LOCK_ACTION_SHIFT_LEFT = 0, - AFFINE_LOCK_ACTION_SHIFT_RIGHT = 1, - AFFINE_LOCK_ACTION_INVERT_RIGHT_7 = 2, - AFFINE_LOCK_ACTION_SWAP_ADJACENT_BITS = 3, - AFFINE_LOCK_ACTION_SWAP_ADJACENT_PAIRS = 4, - AFFINE_LOCK_ACTION_SWAP_NIBBLES_EACH_BYTE = 5, - AFFINE_LOCK_ACTION_REVERSE_EACH_NIBBLE = 6, - AFFINE_LOCK_ACTION_REVERSE_EACH_BYTE = 7, + ACTION_SHIFT_LEFT = 0, + ACTION_SHIFT_RIGHT = 1, + ACTION_INVERT_RIGHT_7 = 2, + ACTION_SWAP_ADJACENT_BITS = 3, + ACTION_SWAP_ADJACENT_PAIRS = 4, + ACTION_SWAP_NIBBLES_EACH_BYTE = 5, + ACTION_REVERSE_EACH_NIBBLE = 6, + ACTION_REVERSE_EACH_BYTE = 7, } AffineLockAction; struct Log { @@ -51,16 +45,10 @@ struct Log { float episode_length; float solve_steps; float timeout_rate; - float invalid_rate; + float invalid_rate; // TODO is this even needed? float solve_efficiency; float target_distance; float solved_target_distance; - float depth_2_rate; - float depth_2_solve_rate; - float depth_4_rate; - float depth_4_solve_rate; - float depth_5_rate; - float depth_5_solve_rate; float depth_6_rate; float depth_6_solve_rate; float depth_8_rate; @@ -74,14 +62,14 @@ typedef struct AffineLockShared { int start_depth; int max_depth; int step_grace; - int num_states; + int num_states; // TODO Only used in testing uint32_t mask; uint32_t* next; AffineLockVisibleTargetTable visible_target_table; float observation_bit_patterns[256][8]; } AffineLockShared; -typedef struct Client { +typedef struct Client { // TODO see if we even need this several envs don't have it int screen_width; int screen_height; } Client; @@ -91,7 +79,7 @@ struct Env { Agent agents[1]; int tag; int boundary_reached; - int num_agents; + int num_agents; // TODO is this even needed? unsigned int rng; uint32_t state; uint32_t target; @@ -100,7 +88,7 @@ struct Env { int scramble_depth; int curriculum_depth; int solution_length; - int solution_actions[AFFINE_LOCK_MAX_SOLUTION_DEPTH]; + int solution_actions[MAX_SOLUTION_DEPTH]; int target_distance; float episode_return; int owns_shared; @@ -109,7 +97,7 @@ struct Env { }; typedef Env AffineLock; -static void affine_lock_init_shared( +static void init_shared( AffineLockShared* shared, int start_depth, int max_depth, @@ -117,8 +105,8 @@ static void affine_lock_init_shared( shared->start_depth = start_depth; shared->max_depth = max_depth; shared->step_grace = step_grace; - shared->num_states = 1 << AFFINE_LOCK_BITS; - shared->mask = (1u << AFFINE_LOCK_BITS) - 1u; + shared->num_states = 1 << BITS; + shared->mask = (1u << BITS) - 1u; for (int value = 0; value < 256; value++) { for (int bit = 0; bit < 8; bit++) { shared->observation_bit_patterns[value][bit] = @@ -126,71 +114,142 @@ static void affine_lock_init_shared( } } - shared->next = calloc( - shared->num_states * AFFINE_LOCK_NUM_ACTIONS, sizeof(uint32_t)); - assert(shared->next); + shared->next = (uint32_t*)calloc( + shared->num_states * NUM_ACTIONS, sizeof(uint32_t)); uint32_t nstates = shared->num_states; for (uint32_t state = 0; state < nstates; state++) { - for (int action = 0; action < AFFINE_LOCK_NUM_ACTIONS; action++) { + for (int action = 0; action < NUM_ACTIONS; action++) { uint32_t next = state; switch (action) { - case AFFINE_LOCK_ACTION_SHIFT_LEFT: + case ACTION_SHIFT_LEFT: next = (state >> 1) | ((state & 1u) << 15); break; - case AFFINE_LOCK_ACTION_SHIFT_RIGHT: + case ACTION_SHIFT_RIGHT: next = ((state << 1) & 0xffffu) | ((state >> 15) & 1u); break; - case AFFINE_LOCK_ACTION_INVERT_RIGHT_7: + case ACTION_INVERT_RIGHT_7: next = state ^ 0xfe00u; break; - case AFFINE_LOCK_ACTION_SWAP_ADJACENT_BITS: + case ACTION_SWAP_ADJACENT_BITS: next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); break; - case AFFINE_LOCK_ACTION_SWAP_ADJACENT_PAIRS: + case ACTION_SWAP_ADJACENT_PAIRS: next = ((state & 0x3333u) << 2) | ((state & 0xccccu) >> 2); break; - case AFFINE_LOCK_ACTION_SWAP_NIBBLES_EACH_BYTE: + case ACTION_SWAP_NIBBLES_EACH_BYTE: next = ((state & 0x0f0fu) << 4) | ((state & 0xf0f0u) >> 4); break; - case AFFINE_LOCK_ACTION_REVERSE_EACH_NIBBLE: + case ACTION_REVERSE_EACH_NIBBLE: next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); break; - case AFFINE_LOCK_ACTION_REVERSE_EACH_BYTE: + case ACTION_REVERSE_EACH_BYTE: next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); next = ((next & 0x0f0fu) << 4) | ((next & 0xf0f0u) >> 4); break; } - shared->next[state * AFFINE_LOCK_NUM_ACTIONS + action] = + shared->next[state * NUM_ACTIONS + action] = next & shared->mask; } } - char error[256]; - assert(affine_lock_visible_targets_load( - AFFINE_LOCK_VISIBLE_TARGET_TABLE_PATH, - AFFINE_LOCK_VISIBLE_TARGET_8ACTION_V1_HASH, - &shared->visible_target_table, - error, - sizeof(error)) == 0); + assert(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &shared->visible_target_table) == 0 && + "failed to load visible target table -- see 'Regenerating the Target " + "Table' in ocean/affine_lock/README.md"); } -static void affine_lock_free_shared(AffineLockShared* shared) { +static AffineLockShared* create_shared(int start_depth, int max_depth, int step_grace) { + AffineLockShared* shared = (AffineLockShared*)calloc(1, sizeof(AffineLockShared)); + init_shared(shared, start_depth, max_depth, step_grace); + return shared; +} + +static void init_env(AffineLock* env, AffineLockShared* shared, unsigned int seed) { + env->shared = shared; + env->rng = seed; + env->num_agents = 1; + env->curriculum_depth = shared->start_depth; +} + +static unsigned int env_seed( // TODO see if this is in the golden envs + unsigned int base_seed, unsigned int env_id) { + uint32_t value = 0x811c9dc5u; + value = (value ^ base_seed) * 0x01000193u; + value = (value ^ env_id) * 0x01000193u; + value ^= value >> 16; + value *= 0x7feb352du; + value ^= value >> 15; + value *= 0x846ca68bu; + value ^= value >> 16; + return value; +} + +void puf_init(Env* env, Dict* kwargs) { + int start_depth = dict_get(kwargs, "start_depth"); + int max_depth = dict_get(kwargs, "max_depth"); + int step_grace = dict_get(kwargs, "step_grace"); + int seed = dict_get(kwargs, "seed"); + AffineLockShared* shared = + create_shared(start_depth, max_depth, step_grace); + init_env(env, shared, env_seed(seed, env->rng)); + env->owns_shared = 1; +} + +static void free_shared(AffineLockShared* shared) { free(shared->next); - affine_lock_visible_targets_free(&shared->visible_target_table); + visible_targets_free(&shared->visible_target_table); } -static uint32_t affine_lock_apply_action( - const AffineLockShared* shared, uint32_t rel, int action) { - return shared->next[rel * AFFINE_LOCK_NUM_ACTIONS + action]; +void puf_close(AffineLock* env) { + if (env->client) { + if (IsWindowReady()) { + CloseWindow(); + } + free(env->client); + } + if (env->owns_shared) { + free_shared(env->shared); + free(env->shared); + } } -// Keep RNG fully local to each env so sweep runs differ only by hyperparams. -// The mixer avoids weak low bits from the LCG when sampling bounded actions or -// bit states. Do not replace this with global rand()/srand(). -static uint32_t affine_lock_random_mixed_u32(AffineLock* env) { +static void add_log(AffineLock* env, int solved, int invalid) { + AffineLockShared* shared = env->shared; + int log_depth = env->target_distance; + int at_max_depth = log_depth == shared->max_depth; + float solve_credit = solved ? log_depth / (float)shared->max_depth : 0; + env->log.perf += solve_credit; + env->log.score += solve_credit; + env->log.solve_rate += solved; + env->log.max_depth_solve += solved && at_max_depth; + env->log.episode_return += env->episode_return; + env->log.episode_length += env->step_count; + env->log.solve_steps += solved ? env->step_count : 0; + env->log.timeout_rate += !solved && !invalid; + env->log.invalid_rate += invalid; + env->log.solve_efficiency += solved ? env->step_count / (float)log_depth : 0; + env->log.target_distance += env->target_distance; + env->log.solved_target_distance += solved ? env->target_distance : 0; + env->log.depth_6_rate += log_depth == 6; + env->log.depth_6_solve_rate += solved && log_depth == 6; + env->log.depth_8_rate += log_depth == 8; + env->log.depth_8_solve_rate += solved && log_depth == 8; + env->log.depth_16_rate += log_depth == 16; + env->log.depth_16_solve_rate += solved && log_depth == 16; + env->log.n += 1; +} + +static uint32_t apply_action(const AffineLockShared* shared, uint32_t rel, int action) { + return shared->next[rel * NUM_ACTIONS + action]; +} + +// Not rand_r(): glibc's LCG has statistically weak low-order bits, and this +// env repeatedly samples individual state bits and small action ranges +// directly from those bits, where the weakness would show up as bias. +static uint32_t random_mixed_u32(AffineLock* env) { env->rng = env->rng * 1664525u + 1013904223u; uint32_t x = env->rng; x ^= x >> 16; @@ -201,30 +260,17 @@ static uint32_t affine_lock_random_mixed_u32(AffineLock* env) { return x; } -static int affine_lock_random_bounded(AffineLock* env, int bound) { +static int random_bounded(AffineLock* env, int bound) { uint32_t ubound = bound; uint32_t limit = UINT32_MAX - UINT32_MAX % ubound; - uint32_t value = affine_lock_random_mixed_u32(env); + uint32_t value = random_mixed_u32(env); while (value >= limit) { - value = affine_lock_random_mixed_u32(env); + value = random_mixed_u32(env); } return value % ubound; } -static unsigned int affine_lock_env_seed( - unsigned int base_seed, unsigned int env_id) { - uint32_t value = 0x811c9dc5u; - value = (value ^ base_seed) * 0x01000193u; - value = (value ^ env_id) * 0x01000193u; - value ^= value >> 16; - value *= 0x7feb352du; - value ^= value >> 15; - value *= 0x846ca68bu; - value ^= value >> 16; - return value; -} - -static const AffineLockVisibleTargetDepth* affine_lock_visible_target_depth( +static const AffineLockVisibleTargetDepth* visible_target_depth( const AffineLockShared* shared, uint32_t requested_depth) { const AffineLockVisibleTargetTable* table = &shared->visible_target_table; @@ -236,21 +282,21 @@ static const AffineLockVisibleTargetDepth* affine_lock_visible_target_depth( return NULL; } -static void affine_lock_reset_state(AffineLock* env) { +static void reset_state(AffineLock* env) { AffineLockShared* shared = env->shared; env->scramble_depth = env->curriculum_depth; env->step_count = 0; env->episode_return = 0; const AffineLockVisibleTargetDepth* depth = - affine_lock_visible_target_depth(shared, env->scramble_depth); - int choice = affine_lock_random_bounded(env, depth->stored_count); + visible_target_depth(shared, env->scramble_depth); + int choice = random_bounded(env, depth->stored_count); const AffineLockVisibleTargetRecord* record = &shared->visible_target_table.records[depth->first_record + choice]; env->state = record->start; env->target = record->target; env->target_distance = record->depth; env->solution_length = record->solution_length; - for (int i = 0; i < AFFINE_LOCK_MAX_SOLUTION_DEPTH; i++) { + for (int i = 0; i < MAX_SOLUTION_DEPTH; i++) { env->solution_actions[i] = -1; } for (int i = 0; i < env->solution_length; i++) { @@ -259,48 +305,7 @@ static void affine_lock_reset_state(AffineLock* env) { env->max_steps = env->target_distance + shared->step_grace; } -static void affine_lock_init_env( - AffineLock* env, AffineLockShared* shared, unsigned int seed) { - env->shared = shared; - env->rng = seed; - env->num_agents = 1; - env->curriculum_depth = shared->start_depth; -} - -static void affine_lock_add_log(AffineLock* env, int solved, int invalid) { - AffineLockShared* shared = env->shared; - int log_depth = env->target_distance; - int at_max_depth = log_depth == shared->max_depth; - float solve_credit = solved ? log_depth / (float)shared->max_depth : 0; - env->log.perf += solve_credit; - env->log.score += solve_credit; - env->log.solve_rate += solved; - env->log.max_depth_solve += solved && at_max_depth; - env->log.episode_return += env->episode_return; - env->log.episode_length += env->step_count; - env->log.solve_steps += solved ? env->step_count : 0; - env->log.timeout_rate += !solved && !invalid; - env->log.invalid_rate += invalid; - env->log.solve_efficiency += solved ? - env->step_count / (float)log_depth : 0; - env->log.target_distance += env->target_distance; - env->log.solved_target_distance += solved ? env->target_distance : 0; - env->log.depth_2_rate += log_depth == 2; - env->log.depth_2_solve_rate += solved && log_depth == 2; - env->log.depth_4_rate += log_depth == 4; - env->log.depth_4_solve_rate += solved && log_depth == 4; - env->log.depth_5_rate += log_depth == 5; - env->log.depth_5_solve_rate += solved && log_depth == 5; - env->log.depth_6_rate += log_depth == 6; - env->log.depth_6_solve_rate += solved && log_depth == 6; - env->log.depth_8_rate += log_depth == 8; - env->log.depth_8_solve_rate += solved && log_depth == 8; - env->log.depth_16_rate += log_depth == 16; - env->log.depth_16_solve_rate += solved && log_depth == 16; - env->log.n += 1; -} - -static void affine_lock_compute_observations(AffineLock* env) { +static void compute_observations(AffineLock* env) { float (*patterns)[8] = env->shared->observation_bit_patterns; uint32_t state = env->state; uint32_t target = env->target; @@ -311,21 +316,21 @@ static void affine_lock_compute_observations(AffineLock* env) { obs[16 + i] = patterns[target & 0xffu][i]; obs[24 + i] = patterns[(target >> 8) & 0xffu][i]; } - obs[AFFINE_LOCK_TIMER_INDEX] = env->step_count / (float)env->max_steps; + obs[TIMER_INDEX] = env->step_count / (float)env->max_steps; } void puf_reset(AffineLock* env) { env->agents[0].rewards[0] = 0; env->agents[0].terminals[0] = 0; - affine_lock_reset_state(env); - affine_lock_compute_observations(env); + reset_state(env); + compute_observations(env); } -static int affine_lock_next_curriculum_depth( +static int next_curriculum_depth( const AffineLockShared* shared, int current_depth) { - for (int i = 0; i < AFFINE_LOCK_CURRICULUM_DEPTH_COUNT; i++) { - int depth = AFFINE_LOCK_CURRICULUM_DEPTHS[i]; + for (int i = 0; i < CURRICULUM_DEPTH_COUNT; i++) { + int depth = CURRICULUM_DEPTHS[i]; if (depth > current_depth) { return depth < shared->max_depth ? depth : shared->max_depth; } @@ -333,16 +338,15 @@ static int affine_lock_next_curriculum_depth( return shared->max_depth; } -// Hold Left Shift + 1-8. -static int affine_lock_human_controls(AffineLock *env) { +static int human_controls(AffineLock *env) { if (!IsWindowReady() || !IsKeyDown(KEY_LEFT_SHIFT)) { return 0; } - static const int keys[AFFINE_LOCK_NUM_ACTIONS] = { + static const int keys[NUM_ACTIONS] = { KEY_ONE, KEY_TWO, KEY_THREE, KEY_FOUR, KEY_FIVE, KEY_SIX, KEY_SEVEN, KEY_EIGHT, }; - for (int i = 0; i < AFFINE_LOCK_NUM_ACTIONS; i++) { + for (int i = 0; i < NUM_ACTIONS; i++) { if (IsKeyPressed(keys[i])) { env->agents[0].actions[0] = (float)i; return 1; @@ -352,17 +356,17 @@ static int affine_lock_human_controls(AffineLock *env) { } void puf_step(AffineLock* env) { - if (affine_lock_human_controls(env) < 0) { + if (human_controls(env) < 0) { return; } AffineLockShared* shared = env->shared; float raw = env->agents[0].actions[0]; - int invalid = !isfinite(raw) || raw < 0 || raw > AFFINE_LOCK_NUM_ACTIONS - 1; + int invalid = !isfinite(raw) || raw < 0 || raw > NUM_ACTIONS - 1; int action = invalid ? -1 : raw; if (!invalid && action != raw) { invalid = 1; } - float reward = AFFINE_LOCK_STEP_REWARD; + float reward = STEP_REWARD; int terminal = 0; int solved = 0; @@ -373,7 +377,7 @@ void puf_step(AffineLock* env) { reward = -1; terminal = 1; } else { - env->state = affine_lock_apply_action(shared, env->state, action); + env->state = apply_action(shared, env->state, action); if (env->state == env->target) { reward = 1; terminal = 1; @@ -387,26 +391,77 @@ void puf_step(AffineLock* env) { env->episode_return += reward; if (terminal) { env->agents[0].terminals[0] = 1; - affine_lock_add_log(env, solved, invalid); + add_log(env, solved, invalid); env->curriculum_depth = solved ? - affine_lock_next_curriculum_depth(shared, env->scramble_depth) : + next_curriculum_depth(shared, env->scramble_depth) : shared->start_depth; - affine_lock_reset_state(env); + reset_state(env); } - affine_lock_compute_observations(env); + compute_observations(env); } -void puf_close(AffineLock* env) { - if (env->client) { - if (IsWindowReady()) { - CloseWindow(); +void puf_log(Log* log, Dict* out) { + float nsolve = log->solve_rate; + float solved_min_win_moves = nsolve ? log->solved_target_distance / nsolve : 0; + float conditional_solve_steps = nsolve ? log->solve_steps / nsolve : 0; + float conditional_solve_efficiency = nsolve ? + log->solve_efficiency / nsolve : 0; + + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "solve_rate", log->solve_rate); + dict_set(out, "max_depth_solve", log->max_depth_solve); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); + dict_set(out, "timeout_rate", log->timeout_rate); + dict_set(out, "invalid_rate", log->invalid_rate); + dict_set(out, "min_win_moves", log->target_distance); + dict_set(out, "solved_min_win_moves", solved_min_win_moves); + dict_set(out, "conditional_solve_steps", conditional_solve_steps); + dict_set(out, "conditional_solve_efficiency", conditional_solve_efficiency); + dict_set(out, "depth_6_solve_rate", log->depth_6_rate ? log->depth_6_solve_rate / log->depth_6_rate : 0); + dict_set(out, "depth_8_solve_rate", log->depth_8_rate ? log->depth_8_solve_rate / log->depth_8_rate : 0); + dict_set(out, "depth_16_solve_rate", log->depth_16_rate ? log->depth_16_solve_rate / log->depth_16_rate : 0); + dict_set(out, "n", log->n); +} + +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"); + int num_buffers = dict_get(vec_kwargs, "num_buffers"); + int agents_per_buffer = total_agents / num_buffers; + int base_seed = dict_get(env_kwargs, "seed"); + int start_depth = dict_get(env_kwargs, "start_depth"); + int max_depth = dict_get(env_kwargs, "max_depth"); + int step_grace = dict_get(env_kwargs, "step_grace"); + + AffineLockShared* shared = + create_shared(start_depth, max_depth, step_grace); + Env* envs = (Env*)calloc(total_agents, sizeof(Env)); + + int buf = 0; + int buf_agents = 0; + buffer_env_starts[0] = 0; + buffer_env_counts[0] = 0; + for (int i = 0; i < total_agents; i++) { + Env* env = &envs[i]; + init_env(env, shared, env_seed(base_seed, i)); + buf_agents += env->num_agents; + 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; } - free(env->client); - } - if (env->owns_shared) { - affine_lock_free_shared(env->shared); - free(env->shared); } + *num_envs_out = total_agents; + return envs; +} + +void my_vec_close(Env* envs) { + free_shared(envs[0].shared); + free(envs[0].shared); } void puf_render(AffineLock* env) { @@ -414,9 +469,9 @@ void puf_render(AffineLock* env) { puf_close(env); exit(0); } - affine_lock_human_controls(env); + human_controls(env); if (!env->client) { - Client* client = calloc(1, sizeof(Client)); + Client* client = (Client*)calloc(1, sizeof(Client)); client->screen_width = 780; client->screen_height = 360; InitWindow(client->screen_width, client->screen_height, @@ -452,7 +507,7 @@ void puf_render(AffineLock* env) { int row_y[2] = {138, 220}; for (int row = 0; row < 2; row++) { DrawText(row_label[row], 30, row_y[row] + 9, 20, RAYWHITE); - for (int bit = 0; bit < AFFINE_LOCK_BITS; bit++) { + for (int bit = 0; bit < BITS; bit++) { int x = 145 + bit * 34; int on = (row_value[row] >> bit) & 1u; int mismatch = ((env->state ^ env->target) >> bit) & 1u; @@ -476,96 +531,3 @@ void puf_render(AffineLock* env) { EndDrawing(); puf_web_vsync(); } - -void puf_log(Log* log, Dict* out) { - float nsolve = log->solve_rate; - float solved_min_win_moves = nsolve ? log->solved_target_distance / nsolve : 0; - float conditional_solve_steps = nsolve ? log->solve_steps / nsolve : 0; - float conditional_solve_efficiency = nsolve ? - log->solve_efficiency / nsolve : 0; - - dict_set(out, "perf", log->perf); - dict_set(out, "score", log->score); - dict_set(out, "solve_rate", log->solve_rate); - dict_set(out, "max_depth_solve", log->max_depth_solve); - dict_set(out, "episode_return", log->episode_return); - dict_set(out, "episode_length", log->episode_length); - dict_set(out, "timeout_rate", log->timeout_rate); - dict_set(out, "invalid_rate", log->invalid_rate); - dict_set(out, "min_win_moves", log->target_distance); - dict_set(out, "solved_min_win_moves", solved_min_win_moves); - dict_set(out, "conditional_solve_steps", conditional_solve_steps); - dict_set(out, "conditional_solve_efficiency", conditional_solve_efficiency); - dict_set(out, "depth_2_solve_rate", - log->depth_2_rate ? log->depth_2_solve_rate / log->depth_2_rate : 0); - dict_set(out, "depth_4_solve_rate", - log->depth_4_rate ? log->depth_4_solve_rate / log->depth_4_rate : 0); - dict_set(out, "depth_5_solve_rate", - log->depth_5_rate ? log->depth_5_solve_rate / log->depth_5_rate : 0); - dict_set(out, "depth_6_solve_rate", - log->depth_6_rate ? log->depth_6_solve_rate / log->depth_6_rate : 0); - dict_set(out, "depth_8_solve_rate", - log->depth_8_rate ? log->depth_8_solve_rate / log->depth_8_rate : 0); - dict_set(out, "depth_16_solve_rate", - log->depth_16_rate ? log->depth_16_solve_rate / log->depth_16_rate : 0); - dict_set(out, "n", log->n); -} - -static AffineLockShared* affine_lock_create_shared( - int start_depth, int max_depth, int step_grace) { - AffineLockShared* shared = calloc(1, sizeof(AffineLockShared)); - assert(shared); - affine_lock_init_shared(shared, start_depth, max_depth, step_grace); - return shared; -} - -void puf_init(Env* env, Dict* kwargs) { - int start_depth = dict_get(kwargs, "start_depth"); - int max_depth = dict_get(kwargs, "max_depth"); - int step_grace = dict_get(kwargs, "step_grace"); - int seed = dict_get(kwargs, "seed"); - AffineLockShared* shared = - affine_lock_create_shared(start_depth, max_depth, step_grace); - affine_lock_init_env(env, shared, affine_lock_env_seed(seed, env->rng)); - env->owns_shared = 1; -} - -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"); - int num_buffers = dict_get(vec_kwargs, "num_buffers"); - int agents_per_buffer = total_agents / num_buffers; - int base_seed = dict_get(env_kwargs, "seed"); - int start_depth = dict_get(env_kwargs, "start_depth"); - int max_depth = dict_get(env_kwargs, "max_depth"); - int step_grace = dict_get(env_kwargs, "step_grace"); - - AffineLockShared* shared = - affine_lock_create_shared(start_depth, max_depth, step_grace); - Env* envs = calloc(total_agents, sizeof(Env)); - assert(envs); - - int buf = 0; - int buf_agents = 0; - buffer_env_starts[0] = 0; - buffer_env_counts[0] = 0; - for (int i = 0; i < total_agents; i++) { - Env* env = &envs[i]; - affine_lock_init_env(env, shared, affine_lock_env_seed(base_seed, i)); - buf_agents += env->num_agents; - 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 = total_agents; - return envs; -} - -void my_vec_close(Env* envs) { - affine_lock_free_shared(envs[0].shared); - free(envs[0].shared); -} diff --git a/ocean/affine_lock/affine_lock_visible_targets.h b/ocean/affine_lock/affine_lock_visible_targets.h index 6388f7e3bc..63a814bc5b 100644 --- a/ocean/affine_lock/affine_lock_visible_targets.h +++ b/ocean/affine_lock/affine_lock_visible_targets.h @@ -1,15 +1,13 @@ #pragma once -#include -#include #include #include #include #include -#define AFFINE_LOCK_VISIBLE_TARGET_FORMAT_VERSION 1u -#define AFFINE_LOCK_VISIBLE_TARGET_RECORD_SIZE 16u -#define AFFINE_LOCK_VISIBLE_TARGET_8ACTION_V1_HASH 0x6e11e18fdafc0baaull +#define VISIBLE_TARGET_FORMAT_VERSION 1u +#define VISIBLE_TARGET_RECORD_SIZE 16u +#define VISIBLE_TARGET_8ACTION_V1_HASH 0x6e11e18fdafc0baaull typedef struct AffineLockVisibleTargetDepth { uint32_t depth; @@ -40,7 +38,7 @@ typedef struct AffineLockVisibleTargetTable { AffineLockVisibleTargetRecord* records; } AffineLockVisibleTargetTable; -static uint64_t affine_lock_visible_targets_mix_u64( +static uint64_t visible_targets_mix_u64( uint64_t hash, uint64_t value) { hash ^= value; @@ -48,43 +46,29 @@ static uint64_t affine_lock_visible_targets_mix_u64( return hash; } -static void affine_lock_visible_targets_set_error( - char* error, - size_t error_size, - const char* format, - ...) { - if (error == NULL || error_size == 0) { - return; - } - va_list args; - va_start(args, format); - vsnprintf(error, error_size, format, args); - va_end(args); -} - -static int affine_lock_visible_targets_read_exact( +static int visible_targets_read_exact( FILE* file, void* out, size_t size) { return fread(out, 1, size, file) == size ? 0 : -1; } -static int affine_lock_visible_targets_read_u16( +static int visible_targets_read_u16( FILE* file, uint16_t* out) { unsigned char bytes[2]; - if (affine_lock_visible_targets_read_exact(file, bytes, sizeof(bytes)) != 0) { + if (visible_targets_read_exact(file, bytes, sizeof(bytes)) != 0) { return -1; } *out = (uint16_t)bytes[0] | ((uint16_t)bytes[1] << 8); return 0; } -static int affine_lock_visible_targets_read_u32( +static int visible_targets_read_u32( FILE* file, uint32_t* out) { unsigned char bytes[4]; - if (affine_lock_visible_targets_read_exact(file, bytes, sizeof(bytes)) != 0) { + if (visible_targets_read_exact(file, bytes, sizeof(bytes)) != 0) { return -1; } *out = (uint32_t)bytes[0] | @@ -94,11 +78,11 @@ static int affine_lock_visible_targets_read_u32( return 0; } -static int affine_lock_visible_targets_read_u64( +static int visible_targets_read_u64( FILE* file, uint64_t* out) { unsigned char bytes[8]; - if (affine_lock_visible_targets_read_exact(file, bytes, sizeof(bytes)) != 0) { + if (visible_targets_read_exact(file, bytes, sizeof(bytes)) != 0) { return -1; } uint64_t value = 0; @@ -109,7 +93,7 @@ static int affine_lock_visible_targets_read_u64( return 0; } -static void affine_lock_visible_targets_free( +static void visible_targets_free( AffineLockVisibleTargetTable* table) { if (table == NULL) { return; @@ -119,38 +103,36 @@ static void affine_lock_visible_targets_free( memset(table, 0, sizeof(*table)); } -static uint64_t affine_lock_visible_targets_checksum( +static uint64_t visible_targets_checksum( const AffineLockVisibleTargetTable* table) { uint64_t hash = 1469598103934665603ull; - hash = affine_lock_visible_targets_mix_u64(hash, table->action_set_hash); + hash = visible_targets_mix_u64(hash, table->action_set_hash); for (uint32_t depth_index = 0; depth_index < table->depth_count; depth_index++) { const AffineLockVisibleTargetDepth* depth = &table->depths[depth_index]; - hash = affine_lock_visible_targets_mix_u64(hash, depth->depth); - hash = affine_lock_visible_targets_mix_u64(hash, depth->exact_pair_count); - hash = affine_lock_visible_targets_mix_u64(hash, depth->stored_count); + hash = visible_targets_mix_u64(hash, depth->depth); + hash = visible_targets_mix_u64(hash, depth->exact_pair_count); + hash = visible_targets_mix_u64(hash, depth->stored_count); for (uint32_t i = 0; i < depth->stored_count; i++) { uint32_t record_index = depth->first_record + i; const AffineLockVisibleTargetRecord* record = &table->records[record_index]; - hash = affine_lock_visible_targets_mix_u64(hash, record->start); - hash = affine_lock_visible_targets_mix_u64(hash, record->target); - hash = affine_lock_visible_targets_mix_u64( + hash = visible_targets_mix_u64(hash, record->start); + hash = visible_targets_mix_u64(hash, record->target); + hash = visible_targets_mix_u64( hash, record->packed_actions); - hash = affine_lock_visible_targets_mix_u64( + hash = visible_targets_mix_u64( hash, record->solution_length); - hash = affine_lock_visible_targets_mix_u64(hash, record->depth); + hash = visible_targets_mix_u64(hash, record->depth); } } return hash; } -static int affine_lock_visible_targets_load( +static int visible_targets_load( const char* path, uint64_t expected_action_set_hash, - AffineLockVisibleTargetTable* table, - char* error, - size_t error_size) { + AffineLockVisibleTargetTable* table) { static const unsigned char expected_magic[8] = { 'A', 'L', '7', 'T', 'G', 'T', '1', '\0' }; @@ -158,57 +140,45 @@ static int affine_lock_visible_targets_load( FILE* file = fopen(path, "rb"); if (file == NULL) { - affine_lock_visible_targets_set_error( - error, error_size, "failed to open %s: %s", path, strerror(errno)); return -1; } unsigned char magic[8]; - if (affine_lock_visible_targets_read_exact(file, magic, sizeof(magic)) != 0 || - affine_lock_visible_targets_read_u32(file, &table->version) != 0 || - affine_lock_visible_targets_read_u32(file, &table->header_size) != 0 || - affine_lock_visible_targets_read_u32(file, &table->record_size) != 0 || - affine_lock_visible_targets_read_u32(file, &table->bits) != 0 || - affine_lock_visible_targets_read_u32(file, &table->num_actions) != 0 || - affine_lock_visible_targets_read_u32(file, &table->depth_count) != 0 || - affine_lock_visible_targets_read_u32(file, &table->record_count) != 0 || - affine_lock_visible_targets_read_u64(file, &table->checksum) != 0 || - affine_lock_visible_targets_read_u64(file, &table->action_set_hash) != 0) { - affine_lock_visible_targets_set_error( - error, error_size, "truncated visible target header"); + if (visible_targets_read_exact(file, magic, sizeof(magic)) != 0 || + visible_targets_read_u32(file, &table->version) != 0 || + visible_targets_read_u32(file, &table->header_size) != 0 || + visible_targets_read_u32(file, &table->record_size) != 0 || + visible_targets_read_u32(file, &table->bits) != 0 || + visible_targets_read_u32(file, &table->num_actions) != 0 || + visible_targets_read_u32(file, &table->depth_count) != 0 || + visible_targets_read_u32(file, &table->record_count) != 0 || + visible_targets_read_u64(file, &table->checksum) != 0 || + visible_targets_read_u64(file, &table->action_set_hash) != 0) { fclose(file); return -1; } if (memcmp(magic, expected_magic, sizeof(magic)) != 0) { - affine_lock_visible_targets_set_error( - error, error_size, "invalid visible target magic"); fclose(file); return -1; } - if (table->version != AFFINE_LOCK_VISIBLE_TARGET_FORMAT_VERSION || - table->record_size != AFFINE_LOCK_VISIBLE_TARGET_RECORD_SIZE || + if (table->version != VISIBLE_TARGET_FORMAT_VERSION || + table->record_size != VISIBLE_TARGET_RECORD_SIZE || table->bits != 16 || table->num_actions == 0 || table->num_actions > 8 || table->depth_count == 0 || table->depth_count > 16) { - affine_lock_visible_targets_set_error( - error, error_size, "unsupported visible target table header"); fclose(file); return -1; } uint32_t expected_header_size = 52u + table->depth_count * 24u; if (table->header_size != expected_header_size) { - affine_lock_visible_targets_set_error( - error, error_size, "unexpected visible target header size"); fclose(file); return -1; } if (expected_action_set_hash != 0 && table->action_set_hash != expected_action_set_hash) { - affine_lock_visible_targets_set_error( - error, error_size, "visible target action set hash mismatch"); fclose(file); return -1; } @@ -218,10 +188,8 @@ static int affine_lock_visible_targets_load( table->records = (AffineLockVisibleTargetRecord*)calloc( table->record_count, sizeof(AffineLockVisibleTargetRecord)); if (table->depths == NULL || table->records == NULL) { - affine_lock_visible_targets_set_error( - error, error_size, "failed to allocate visible target table"); fclose(file); - affine_lock_visible_targets_free(table); + visible_targets_free(table); return -1; } @@ -229,18 +197,16 @@ static int affine_lock_visible_targets_load( for (uint32_t i = 0; i < table->depth_count; i++) { AffineLockVisibleTargetDepth* depth = &table->depths[i]; uint32_t reserved = 0; - if (affine_lock_visible_targets_read_u32(file, &depth->depth) != 0 || - affine_lock_visible_targets_read_u32( + if (visible_targets_read_u32(file, &depth->depth) != 0 || + visible_targets_read_u32( file, &depth->first_record) != 0 || - affine_lock_visible_targets_read_u32( + visible_targets_read_u32( file, &depth->stored_count) != 0 || - affine_lock_visible_targets_read_u32(file, &reserved) != 0 || - affine_lock_visible_targets_read_u64( + visible_targets_read_u32(file, &reserved) != 0 || + visible_targets_read_u64( file, &depth->exact_pair_count) != 0) { - affine_lock_visible_targets_set_error( - error, error_size, "truncated visible target depth table"); fclose(file); - affine_lock_visible_targets_free(table); + visible_targets_free(table); return -1; } if (reserved != 0 || @@ -248,75 +214,58 @@ static int affine_lock_visible_targets_load( depth->stored_count > table->record_count || depth->first_record + depth->stored_count > table->record_count) { - affine_lock_visible_targets_set_error( - error, error_size, "invalid visible target depth table"); fclose(file); - affine_lock_visible_targets_free(table); + visible_targets_free(table); return -1; } depth_record_total += depth->stored_count; } if (depth_record_total != table->record_count) { - affine_lock_visible_targets_set_error( - error, error_size, "visible target depth counts do not sum"); fclose(file); - affine_lock_visible_targets_free(table); + visible_targets_free(table); return -1; } for (uint32_t i = 0; i < table->record_count; i++) { AffineLockVisibleTargetRecord* record = &table->records[i]; uint16_t reserved = 0; - if (affine_lock_visible_targets_read_u16(file, &record->start) != 0 || - affine_lock_visible_targets_read_u16(file, &record->target) != 0 || - affine_lock_visible_targets_read_u64( + if (visible_targets_read_u16(file, &record->start) != 0 || + visible_targets_read_u16(file, &record->target) != 0 || + visible_targets_read_u64( file, &record->packed_actions) != 0) { - affine_lock_visible_targets_set_error( - error, error_size, "truncated visible target record"); fclose(file); - affine_lock_visible_targets_free(table); + visible_targets_free(table); return -1; } int solution_length = fgetc(file); int depth = fgetc(file); if (solution_length == EOF || depth == EOF || - affine_lock_visible_targets_read_u16(file, &reserved) != 0) { - affine_lock_visible_targets_set_error( - error, error_size, "truncated visible target record"); + visible_targets_read_u16(file, &reserved) != 0) { fclose(file); - affine_lock_visible_targets_free(table); + visible_targets_free(table); return -1; } record->solution_length = (uint8_t)solution_length; record->depth = (uint8_t)depth; if (reserved != 0 || record->solution_length != record->depth) { - affine_lock_visible_targets_set_error( - error, error_size, "invalid visible target record"); fclose(file); - affine_lock_visible_targets_free(table); + visible_targets_free(table); return -1; } } int extra = fgetc(file); if (extra != EOF) { - affine_lock_visible_targets_set_error( - error, error_size, "visible target file has trailing bytes"); fclose(file); - affine_lock_visible_targets_free(table); + visible_targets_free(table); return -1; } fclose(file); uint64_t computed_checksum = - affine_lock_visible_targets_checksum(table); + visible_targets_checksum(table); if (computed_checksum != table->checksum) { - affine_lock_visible_targets_set_error( - error, error_size, - "visible target checksum mismatch: got 0x%016llx expected 0x%016llx", - (unsigned long long)computed_checksum, - (unsigned long long)table->checksum); - affine_lock_visible_targets_free(table); + visible_targets_free(table); return -1; } return 0; diff --git a/ocean/affine_lock/tests/run_all.sh b/ocean/affine_lock/tests/run_all.sh index cb4bc34d70..b9a119cc71 100755 --- a/ocean/affine_lock/tests/run_all.sh +++ b/ocean/affine_lock/tests/run_all.sh @@ -8,8 +8,15 @@ LOG_OUT="${TMPDIR:-/tmp}/affine_lock_log_export_tests" C99_OUT="${TMPDIR:-/tmp}/affine_lock_c99_compile" CC_BIN="${CC:-clang}" RAYLIB_INC="$ROOT/raylib-5.5_linux_amd64/include" +RAYLIB_LIB="$ROOT/raylib-5.5_linux_amd64/lib/libraylib.a" if [ ! -d "$RAYLIB_INC" ]; then RAYLIB_INC="$ROOT/raylib-5.5_macos/include" + RAYLIB_LIB="$ROOT/raylib-5.5_macos/lib/libraylib.a" +fi +if [ "$(uname -s)" = "Linux" ]; then + RAYLIB_LDFLAGS=(-lGL -lpthread -ldl -lrt) +else + RAYLIB_LDFLAGS=(-framework Cocoa -framework IOKit -framework CoreVideo -framework OpenGL) fi python3 "$ROOT/ocean/affine_lock/tests/test_metadata_smoke.py" @@ -21,7 +28,7 @@ bash "$ROOT/ocean/affine_lock/tests/test_8action_visible_targets_smoke.sh" -I"$ROOT" -I"$ROOT/src" -I"$ROOT/ocean/affine_lock" -I"$ROOT/vendor" \ -I"$RAYLIB_INC" \ "$ROOT/ocean/affine_lock/tests/test_affine_lock.c" \ - -Wl,--gc-sections -lm -o "$C99_OUT" + "$RAYLIB_LIB" -Wl,--gc-sections "${RAYLIB_LDFLAGS[@]}" -lm -o "$C99_OUT" "$CC_BIN" \ -std=c11 -Wall -Wextra -Werror -Wno-unused-function \ @@ -29,7 +36,7 @@ bash "$ROOT/ocean/affine_lock/tests/test_8action_visible_targets_smoke.sh" -I"$ROOT" -I"$ROOT/src" -I"$ROOT/ocean/affine_lock" -I"$ROOT/vendor" \ -I"$RAYLIB_INC" \ "$ROOT/ocean/affine_lock/tests/test_affine_lock.c" \ - -Wl,--gc-sections -lm -o "$OUT" + "$RAYLIB_LIB" -Wl,--gc-sections "${RAYLIB_LDFLAGS[@]}" -lm -o "$OUT" "$CC_BIN" \ -std=c11 -Wall -Wextra -Werror -Wno-unused-function -Wno-unused-parameter \ @@ -38,7 +45,7 @@ bash "$ROOT/ocean/affine_lock/tests/test_8action_visible_targets_smoke.sh" -I"$ROOT" -I"$ROOT/src" -I"$ROOT/ocean/affine_lock" -I"$ROOT/vendor" \ -I"$RAYLIB_INC" \ "$ROOT/ocean/affine_lock/tests/test_affine_lock_log_export.c" \ - -Wl,--gc-sections -lm -o "$LOG_OUT" + "$RAYLIB_LIB" -Wl,--gc-sections "${RAYLIB_LDFLAGS[@]}" -lm -o "$LOG_OUT" ASAN_OPTIONS="${ASAN_OPTIONS:-detect_leaks=0}" "$OUT" ASAN_OPTIONS="${ASAN_OPTIONS:-detect_leaks=0}" "$LOG_OUT" diff --git a/ocean/affine_lock/tests/test_affine_lock.c b/ocean/affine_lock/tests/test_affine_lock.c index 0176313535..c4b0cae183 100644 --- a/ocean/affine_lock/tests/test_affine_lock.c +++ b/ocean/affine_lock/tests/test_affine_lock.c @@ -69,7 +69,7 @@ static AffineLockShared make_shared( int step_grace) { AffineLockShared shared; memset(&shared, 0, sizeof(shared)); - affine_lock_init_shared(&shared, start_depth, max_depth, step_grace); + init_shared(&shared, start_depth, max_depth, step_grace); return shared; } @@ -77,16 +77,16 @@ static void make_env( AffineLock* env, AffineLockShared* shared, unsigned int seed, - float observations[AFFINE_LOCK_OBS_SIZE], - float actions[AFFINE_LOCK_NUM_ATNS], + float observations[OBS_SIZE], + float actions[NUM_ATNS], float rewards[1], float terminals[1]) { memset(env, 0, sizeof(*env)); - memset(observations, 0, AFFINE_LOCK_OBS_SIZE * sizeof(float)); + memset(observations, 0, OBS_SIZE * sizeof(float)); actions[0] = 0.0f; rewards[0] = 0.0f; terminals[0] = 0.0f; - affine_lock_init_env(env, shared, seed); + init_env(env, shared, seed); env->agents[0].observations = observations; env->agents[0].actions = actions; env->agents[0].rewards = rewards; @@ -94,9 +94,9 @@ static void make_env( } static uint32_t bits_from_text(const char* bits) { - EXPECT_EQ_INT(strlen(bits), AFFINE_LOCK_BITS); + EXPECT_EQ_INT(strlen(bits), BITS); uint32_t value = 0u; - for (int i = 0; i < AFFINE_LOCK_BITS; i++) { + for (int i = 0; i < BITS; i++) { EXPECT_TRUE(bits[i] == '0' || bits[i] == '1'); if (bits[i] == '1') { value |= 1u << i; @@ -177,7 +177,7 @@ static void compute_test_bfs_stats( stats->farthest_distance = distance; } - for (int action = 0; action < AFFINE_LOCK_NUM_ACTIONS; action++) { + for (int action = 0; action < NUM_ACTIONS; action++) { uint32_t next = test_apply_action(state, action) & shared->mask; if (distances[next] >= 0) { continue; @@ -203,7 +203,7 @@ static void test_log_solve_credit_uses_known_target_distance(void) { env.target_distance = 8; env.step_count = 8; - affine_lock_add_log(&env, 1, 0); + add_log(&env, 1, 0); EXPECT_NEAR(env.log.perf, expected_solve_credit(&shared, 8), 0.0f); EXPECT_NEAR(env.log.score, expected_solve_credit(&shared, 8), 0.0f); @@ -211,8 +211,6 @@ static void test_log_solve_credit_uses_known_target_distance(void) { EXPECT_NEAR(env.log.solve_efficiency, 1.0f, 0.0f); EXPECT_NEAR(env.log.target_distance, 8.0f, 0.0f); EXPECT_NEAR(env.log.solved_target_distance, 8.0f, 0.0f); - EXPECT_NEAR(env.log.depth_5_rate, 0.0f, 0.0f); - EXPECT_NEAR(env.log.depth_5_solve_rate, 0.0f, 0.0f); EXPECT_NEAR(env.log.depth_6_rate, 0.0f, 0.0f); EXPECT_NEAR(env.log.depth_6_solve_rate, 0.0f, 0.0f); EXPECT_NEAR(env.log.depth_8_rate, 1.0f, 0.0f); @@ -220,33 +218,33 @@ static void test_log_solve_credit_uses_known_target_distance(void) { EXPECT_NEAR(env.log.depth_16_rate, 0.0f, 0.0f); EXPECT_NEAR(env.log.depth_16_solve_rate, 0.0f, 0.0f); - affine_lock_free_shared(&shared); + free_shared(&shared); } static void expect_observation_matches(const AffineLock* env) { float* obs = env->agents[0].observations; - for (int bit = 0; bit < AFFINE_LOCK_BITS; bit++) { + for (int bit = 0; bit < BITS; bit++) { uint32_t bit_mask = 1u << bit; float expected_current = (env->state & bit_mask) ? 1.0f : -1.0f; float expected_target = (env->target & bit_mask) ? 1.0f : -1.0f; EXPECT_NEAR(obs[bit], expected_current, 0.0f); - EXPECT_NEAR(obs[AFFINE_LOCK_BITS + bit], expected_target, 0.0f); + EXPECT_NEAR(obs[BITS + bit], expected_target, 0.0f); } - for (int i = 0; i < AFFINE_LOCK_TIMER_INDEX; i++) { + for (int i = 0; i < TIMER_INDEX; i++) { EXPECT_TRUE(obs[i] == -1.0f || obs[i] == 1.0f); } float expected_timer = env->max_steps > 0 ? (float)env->step_count / (float)env->max_steps : 0.0f; - EXPECT_TRUE(obs[AFFINE_LOCK_TIMER_INDEX] >= 0.0f); - EXPECT_TRUE(obs[AFFINE_LOCK_TIMER_INDEX] <= 1.0f); - EXPECT_NEAR(obs[AFFINE_LOCK_TIMER_INDEX], expected_timer, 0.000001f); + EXPECT_TRUE(obs[TIMER_INDEX] >= 0.0f); + EXPECT_TRUE(obs[TIMER_INDEX] <= 1.0f); + EXPECT_NEAR(obs[TIMER_INDEX], expected_timer, 0.000001f); } static int find_non_solving_action(AffineLock* env) { - for (int action = 0; action < AFFINE_LOCK_NUM_ACTIONS; action++) { - uint32_t next = affine_lock_apply_action(env->shared, env->state, action); + for (int action = 0; action < NUM_ACTIONS; action++) { + uint32_t next = apply_action(env->shared, env->state, action); if (next != env->target) { return action; } @@ -283,12 +281,6 @@ static uint64_t log_snapshot_checksum(uint64_t hash, const Log* log) { hash = mix_float(hash, log->solve_efficiency); hash = mix_float(hash, log->target_distance); hash = mix_float(hash, log->solved_target_distance); - hash = mix_float(hash, log->depth_2_rate); - hash = mix_float(hash, log->depth_2_solve_rate); - hash = mix_float(hash, log->depth_4_rate); - hash = mix_float(hash, log->depth_4_solve_rate); - hash = mix_float(hash, log->depth_5_rate); - hash = mix_float(hash, log->depth_5_solve_rate); hash = mix_float(hash, log->depth_6_rate); hash = mix_float(hash, log->depth_6_solve_rate); hash = mix_float(hash, log->depth_8_rate); @@ -312,11 +304,11 @@ static uint64_t reset_snapshot_checksum(const AffineLock* env) { hash = mix_float(hash, env->agents[0].rewards[0]); hash = mix_float(hash, env->agents[0].terminals[0]); float* obs = env->agents[0].observations; - for (int i = 0; i < AFFINE_LOCK_OBS_SIZE; i++) { + for (int i = 0; i < OBS_SIZE; i++) { hash = mix_float(hash, obs[i]); } hash = log_snapshot_checksum(hash, &env->log); - for (int i = 0; i < AFFINE_LOCK_MAX_SOLUTION_DEPTH; i++) { + for (int i = 0; i < MAX_SOLUTION_DEPTH; i++) { hash = mix_u64(hash, (uint64_t)(env->solution_actions[i] + 1)); } return hash; @@ -325,10 +317,10 @@ static uint64_t reset_snapshot_checksum(const AffineLock* env) { static void expect_env_snapshots_equal( const AffineLock* a, const AffineLock* b, - const float obs_a[AFFINE_LOCK_OBS_SIZE], - const float obs_b[AFFINE_LOCK_OBS_SIZE]) { + const float obs_a[OBS_SIZE], + const float obs_b[OBS_SIZE]) { EXPECT_EQ_U64(reset_snapshot_checksum(a), reset_snapshot_checksum(b)); - EXPECT_TRUE(memcmp(obs_a, obs_b, AFFINE_LOCK_OBS_SIZE * sizeof(float)) == 0); + EXPECT_TRUE(memcmp(obs_a, obs_b, OBS_SIZE * sizeof(float)) == 0); EXPECT_EQ_U32(a->state, b->state); EXPECT_EQ_U32(a->target, b->target); EXPECT_EQ_INT(a->scramble_depth, b->scramble_depth); @@ -353,7 +345,7 @@ static void expect_solution_reaches_target( uint32_t simulated = env->state; for (int i = 0; i < env->solution_length; i++) { int action = env->solution_actions[i]; - EXPECT_TRUE(action >= 0 && action < AFFINE_LOCK_NUM_ACTIONS); + EXPECT_TRUE(action >= 0 && action < NUM_ACTIONS); simulated = test_apply_action(simulated, action) & shared->mask; } EXPECT_EQ_U32(simulated, env->target); @@ -375,18 +367,6 @@ static void expect_depth_log_delta( const Log* after, int depth, int solved) { - EXPECT_NEAR(after->depth_2_rate, - before->depth_2_rate + (depth == 2 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_2_solve_rate, - before->depth_2_solve_rate + (solved && depth == 2 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_4_rate, - before->depth_4_rate + (depth == 4 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_4_solve_rate, - before->depth_4_solve_rate + (solved && depth == 4 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_5_rate, - before->depth_5_rate + (depth == 5 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_5_solve_rate, - before->depth_5_solve_rate + (solved && depth == 5 ? 1.0f : 0.0f), 0.0f); EXPECT_NEAR(after->depth_6_rate, before->depth_6_rate + (depth == 6 ? 1.0f : 0.0f), 0.0f); EXPECT_NEAR(after->depth_6_solve_rate, @@ -417,7 +397,7 @@ static void expect_oracle_episode_win(AffineLock* env, int depth) { env->agents[0].actions[0] = (float)env->solution_actions[step]; puf_step(env); if (step + 1 < solution_length) { - EXPECT_NEAR(env->agents[0].rewards[0], AFFINE_LOCK_STEP_REWARD, 0.0f); + EXPECT_NEAR(env->agents[0].rewards[0], STEP_REWARD, 0.0f); EXPECT_NEAR(env->agents[0].terminals[0], 0.0f, 0.0f); EXPECT_EQ_INT(env->step_count, step + 1); expect_observation_matches(env); @@ -445,7 +425,7 @@ static void expect_oracle_episode_win(AffineLock* env, int depth) { before.solved_target_distance + (float)target_distance, 0.0f); expect_depth_log_delta(&before, &env->log, depth, 1); - int next_depth = affine_lock_next_curriculum_depth(shared, depth); + int next_depth = next_curriculum_depth(shared, depth); EXPECT_EQ_INT(env->scramble_depth, next_depth); expect_observation_matches(env); } @@ -467,7 +447,7 @@ static void expect_non_solving_episode_timeout(AffineLock* env, int depth) { env->agents[0].actions[0] = (float)action; puf_step(env); if (step + 1 < max_steps) { - EXPECT_NEAR(env->agents[0].rewards[0], AFFINE_LOCK_STEP_REWARD, 0.0f); + EXPECT_NEAR(env->agents[0].rewards[0], STEP_REWARD, 0.0f); EXPECT_NEAR(env->agents[0].terminals[0], 0.0f, 0.0f); EXPECT_EQ_INT(env->step_count, step + 1); expect_observation_matches(env); @@ -505,11 +485,11 @@ static size_t read_text_file(const char* path, char* buffer, size_t capacity) { } static void test_metadata_contract(void) { - EXPECT_EQ_INT(AFFINE_LOCK_BITS, 16); - EXPECT_EQ_INT(AFFINE_LOCK_TIMER_INDEX, 32); - EXPECT_EQ_INT(AFFINE_LOCK_OBS_SIZE, 33); - EXPECT_EQ_INT(AFFINE_LOCK_NUM_ATNS, 1); - EXPECT_EQ_INT(AFFINE_LOCK_NUM_ACTIONS, 8); + EXPECT_EQ_INT(BITS, 16); + EXPECT_EQ_INT(TIMER_INDEX, 32); + EXPECT_EQ_INT(OBS_SIZE, 33); + EXPECT_EQ_INT(NUM_ATNS, 1); + EXPECT_EQ_INT(NUM_ACTIONS, 8); } static void test_config_and_binding_metadata_contract(void) { @@ -531,9 +511,9 @@ static void test_config_and_binding_metadata_contract(void) { char header[65536]; read_text_file("ocean/affine_lock/affine_lock.h", header, sizeof(header)); - EXPECT_TRUE(strstr(header, "#define OBS_SIZE AFFINE_LOCK_OBS_SIZE") != NULL); - EXPECT_TRUE(strstr(header, "#define ACT_SIZES {AFFINE_LOCK_NUM_ACTIONS}") != NULL); - EXPECT_TRUE(strstr(header, "#define NUM_ATNS AFFINE_LOCK_NUM_ATNS") != NULL); + EXPECT_TRUE(strstr(header, "#define OBS_SIZE (TIMER_INDEX + 1)") != NULL); + EXPECT_TRUE(strstr(header, "#define ACT_SIZES {NUM_ACTIONS}") != NULL); + EXPECT_TRUE(strstr(header, "#define NUM_ATNS 1") != NULL); EXPECT_TRUE(strstr(header, "typedef") != NULL && strstr(header, "obs_t") != NULL); } @@ -541,7 +521,7 @@ static void test_global_action_examples(void) { AffineLockShared shared = make_shared(2, 16, 0); uint32_t start = bits_from_text("0011011000010111"); - const char* expected[AFFINE_LOCK_NUM_ACTIONS] = { + const char* expected[NUM_ACTIONS] = { "0110110000101110", "1001101100001011", "0011011001101000", @@ -552,50 +532,50 @@ static void test_global_action_examples(void) { "0110110011101000", }; - for (int action = 0; action < AFFINE_LOCK_NUM_ACTIONS; action++) { - uint32_t next = affine_lock_apply_action(&shared, start, action); + for (int action = 0; action < NUM_ACTIONS; action++) { + uint32_t next = apply_action(&shared, start, action); EXPECT_EQ_U32(next, bits_from_text(expected[action])); } - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_actions_round_trip_for_all_states(void) { AffineLockShared shared = make_shared(2, 16, 0); - const int inverse_actions[AFFINE_LOCK_NUM_ACTIONS] = { - AFFINE_LOCK_ACTION_SHIFT_RIGHT, - AFFINE_LOCK_ACTION_SHIFT_LEFT, - AFFINE_LOCK_ACTION_INVERT_RIGHT_7, - AFFINE_LOCK_ACTION_SWAP_ADJACENT_BITS, - AFFINE_LOCK_ACTION_SWAP_ADJACENT_PAIRS, - AFFINE_LOCK_ACTION_SWAP_NIBBLES_EACH_BYTE, - AFFINE_LOCK_ACTION_REVERSE_EACH_NIBBLE, - AFFINE_LOCK_ACTION_REVERSE_EACH_BYTE, + const int inverse_actions[NUM_ACTIONS] = { + ACTION_SHIFT_RIGHT, + ACTION_SHIFT_LEFT, + ACTION_INVERT_RIGHT_7, + ACTION_SWAP_ADJACENT_BITS, + ACTION_SWAP_ADJACENT_PAIRS, + ACTION_SWAP_NIBBLES_EACH_BYTE, + ACTION_REVERSE_EACH_NIBBLE, + ACTION_REVERSE_EACH_BYTE, }; EXPECT_EQ_INT(shared.num_states, 1 << 16); EXPECT_EQ_U32(shared.mask, 0xffffu); - for (int action = 0; action < AFFINE_LOCK_NUM_ACTIONS; action++) { + for (int action = 0; action < NUM_ACTIONS; action++) { int inverse = inverse_actions[action]; - EXPECT_TRUE(inverse >= 0 && inverse < AFFINE_LOCK_NUM_ACTIONS); + EXPECT_TRUE(inverse >= 0 && inverse < NUM_ACTIONS); EXPECT_EQ_INT(inverse_actions[inverse], action); for (uint32_t state = 0; state < (uint32_t)shared.num_states; state++) { - uint32_t next = affine_lock_apply_action(&shared, state, action); + uint32_t next = apply_action(&shared, state, action); EXPECT_EQ_U32(next & ~shared.mask, 0u); - uint32_t round_trip = affine_lock_apply_action(&shared, next, inverse); + uint32_t round_trip = apply_action(&shared, next, inverse); EXPECT_EQ_U32(round_trip, state); } } - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_reset_randomizes_target_and_current(void) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, 123, observations, actions, rewards, terminals); @@ -629,15 +609,15 @@ static void test_reset_randomizes_target_and_current(void) { EXPECT_TRUE(target_changed); EXPECT_TRUE(state_changed); - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_visible_target_table_initialization_samples_reachable_target(void) { AffineLockShared shared = make_shared(8, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, 777, observations, actions, rewards, terminals); @@ -651,7 +631,7 @@ static void test_visible_target_table_initialization_samples_reachable_target(vo expect_solution_reaches_target(&shared, &env); expect_observation_matches(&env); - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_visible_target_table_depths_have_expected_distances(void) { @@ -661,8 +641,8 @@ static void test_visible_target_table_depths_have_expected_distances(void) { AffineLockShared shared = make_shared(depth, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, (unsigned int)(1900 + depth), observations, @@ -682,7 +662,7 @@ static void test_visible_target_table_depths_have_expected_distances(void) { EXPECT_NEAR(rewards[0], 1.0f, 0.0f); EXPECT_NEAR(terminals[0], 1.0f, 0.0f); - affine_lock_free_shared(&shared); + free_shared(&shared); } } @@ -694,14 +674,14 @@ static void test_visible_target_table_reset_uses_stored_records(void) { int requested_depth = requested_depths[depth_index]; AffineLockShared shared = make_shared(requested_depth, 16, 0); const AffineLockVisibleTargetDepth* table_depth = - affine_lock_visible_target_depth(&shared, requested_depth); + visible_target_depth(&shared, requested_depth); EXPECT_TRUE(table_depth != NULL); EXPECT_EQ_INT((int)table_depth->stored_count, expected_pool_sizes[depth_index]); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, (unsigned int)(2500 + requested_depth), @@ -719,7 +699,7 @@ static void test_visible_target_table_reset_uses_stored_records(void) { expect_solution_reaches_target(&shared, &env); } - affine_lock_free_shared(&shared); + free_shared(&shared); } } @@ -731,8 +711,8 @@ static void test_visible_target_table_matches_independent_bfs_over_repeated_rese AffineLockShared shared = make_shared(depth, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, (unsigned int)(1000 + depth), @@ -746,15 +726,15 @@ static void test_visible_target_table_matches_independent_bfs_over_repeated_rese expect_observation_matches(&env); } - affine_lock_free_shared(&shared); + free_shared(&shared); } } static void test_observation_encoding_is_32_signed_bit_floats_plus_timer(void) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, 7, observations, actions, rewards, terminals); @@ -763,57 +743,57 @@ static void test_observation_encoding_is_32_signed_bit_floats_plus_timer(void) { env.target = 0x0f0fu; env.step_count = 3; env.max_steps = 12; - affine_lock_compute_observations(&env); + compute_observations(&env); expect_observation_matches(&env); - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_timer_observation_progresses_and_resets_after_timeout(void) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, 19, observations, actions, rewards, terminals); puf_reset(&env); - EXPECT_NEAR(observations[AFFINE_LOCK_TIMER_INDEX], 0.0f, 0.0f); + EXPECT_NEAR(observations[TIMER_INDEX], 0.0f, 0.0f); env.target = 0u; env.state = shared.mask; env.step_count = 0; env.max_steps = 4; - affine_lock_compute_observations(&env); - EXPECT_NEAR(observations[AFFINE_LOCK_TIMER_INDEX], 0.0f, 0.0f); + compute_observations(&env); + EXPECT_NEAR(observations[TIMER_INDEX], 0.0f, 0.0f); actions[0] = 1.0f; puf_step(&env); EXPECT_NEAR(terminals[0], 0.0f, 0.0f); - EXPECT_NEAR(observations[AFFINE_LOCK_TIMER_INDEX], 0.25f, 0.000001f); + EXPECT_NEAR(observations[TIMER_INDEX], 0.25f, 0.000001f); puf_step(&env); EXPECT_NEAR(terminals[0], 0.0f, 0.0f); - EXPECT_NEAR(observations[AFFINE_LOCK_TIMER_INDEX], 0.5f, 0.000001f); + EXPECT_NEAR(observations[TIMER_INDEX], 0.5f, 0.000001f); puf_step(&env); EXPECT_NEAR(terminals[0], 0.0f, 0.0f); - EXPECT_NEAR(observations[AFFINE_LOCK_TIMER_INDEX], 0.75f, 0.000001f); + EXPECT_NEAR(observations[TIMER_INDEX], 0.75f, 0.000001f); puf_step(&env); EXPECT_NEAR(rewards[0], -1.0f, 0.0f); EXPECT_NEAR(terminals[0], 1.0f, 0.0f); EXPECT_EQ_INT(env.step_count, 0); - EXPECT_NEAR(observations[AFFINE_LOCK_TIMER_INDEX], 0.0f, 0.0f); + EXPECT_NEAR(observations[TIMER_INDEX], 0.0f, 0.0f); - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_actions_apply_to_current_state_directly(void) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, 55, observations, actions, rewards, terminals); @@ -822,7 +802,7 @@ static void test_actions_apply_to_current_state_directly(void) { uint32_t target = bits_from_text("1111000011110000"); uint32_t state = bits_from_text("0011011000010111"); int action = 1; - uint32_t expected_state = affine_lock_apply_action(&shared, state, action); + uint32_t expected_state = apply_action(&shared, state, action); EXPECT_NE_U32(expected_state, target); env.target = target; @@ -837,14 +817,14 @@ static void test_actions_apply_to_current_state_directly(void) { EXPECT_EQ_U32(env.target, target); EXPECT_EQ_U32(env.state, expected_state); - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_action_float_validation_rejects_non_discrete_values(void) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, 57, observations, actions, rewards, terminals); @@ -873,15 +853,15 @@ static void test_action_float_validation_rejects_non_discrete_values(void) { EXPECT_EQ_INT(env.step_count, 0); } - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_visible_target_table_curriculum_and_logging(void) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, 42, observations, actions, rewards, terminals); @@ -898,12 +878,6 @@ static void test_visible_target_table_curriculum_and_logging(void) { float prev_max_depth_solve = env.log.max_depth_solve; float prev_target_distance = env.log.target_distance; float prev_solved_target_distance = env.log.solved_target_distance; - float prev_depth_2 = env.log.depth_2_rate; - float prev_depth_2_solve = env.log.depth_2_solve_rate; - float prev_depth_4 = env.log.depth_4_rate; - float prev_depth_4_solve = env.log.depth_4_solve_rate; - float prev_depth_5 = env.log.depth_5_rate; - float prev_depth_5_solve = env.log.depth_5_solve_rate; float prev_depth_6 = env.log.depth_6_rate; float prev_depth_6_solve = env.log.depth_6_solve_rate; float prev_depth_8 = env.log.depth_8_rate; @@ -926,18 +900,6 @@ static void test_visible_target_table_curriculum_and_logging(void) { prev_target_distance + (float)target_distance, 0.0f); EXPECT_NEAR(env.log.solved_target_distance, prev_solved_target_distance + (float)target_distance, 0.0f); - EXPECT_NEAR(env.log.depth_2_rate, - prev_depth_2 + (metric_depth == 2 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_2_solve_rate, - prev_depth_2_solve + (metric_depth == 2 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_4_rate, - prev_depth_4 + (metric_depth == 4 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_4_solve_rate, - prev_depth_4_solve + (metric_depth == 4 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_5_rate, - prev_depth_5 + (metric_depth == 5 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_5_solve_rate, - prev_depth_5_solve + (metric_depth == 5 ? 1.0f : 0.0f), 0.0f); EXPECT_NEAR(env.log.depth_6_rate, prev_depth_6 + (metric_depth == 6 ? 1.0f : 0.0f), 0.0f); EXPECT_NEAR(env.log.depth_6_solve_rate, @@ -970,15 +932,15 @@ static void test_visible_target_table_curriculum_and_logging(void) { EXPECT_NEAR(env.log.invalid_rate, prev_invalid + 1.0f, 0.0f); EXPECT_EQ_INT(env.scramble_depth, shared.start_depth); - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_visible_target_table_oracle_wins_all_curriculum_depths_end_to_end(void) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, 4242, observations, actions, rewards, terminals); @@ -991,12 +953,6 @@ static void test_visible_target_table_oracle_wins_all_curriculum_depths_end_to_e EXPECT_EQ_INT(env.scramble_depth, shared.max_depth); EXPECT_NEAR(env.log.n, 6.0f, 0.0f); - EXPECT_NEAR(env.log.depth_2_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_2_solve_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_4_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_4_solve_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_5_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_5_solve_rate, 1.0f, 0.0f); EXPECT_NEAR(env.log.depth_6_rate, 1.0f, 0.0f); EXPECT_NEAR(env.log.depth_6_solve_rate, 1.0f, 0.0f); EXPECT_NEAR(env.log.depth_8_rate, 1.0f, 0.0f); @@ -1006,7 +962,7 @@ static void test_visible_target_table_oracle_wins_all_curriculum_depths_end_to_e EXPECT_NEAR(env.log.timeout_rate, 0.0f, 0.0f); EXPECT_NEAR(env.log.invalid_rate, 0.0f, 0.0f); - affine_lock_free_shared(&shared); + free_shared(&shared); } static void test_visible_target_table_timeouts_at_all_curriculum_depths_end_to_end(void) { @@ -1017,8 +973,8 @@ static void test_visible_target_table_timeouts_at_all_curriculum_depths_end_to_e AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, (unsigned int)(5200 + loss_depth), @@ -1035,20 +991,20 @@ static void test_visible_target_table_timeouts_at_all_curriculum_depths_end_to_e EXPECT_TRUE(env.log.solve_rate >= 0.0f); EXPECT_NEAR(env.log.invalid_rate, 0.0f, 0.0f); - affine_lock_free_shared(&shared); + free_shared(&shared); } } static int deterministic_stream_action(int episode, int step) { - return (episode * 3 + step * 7) % AFFINE_LOCK_NUM_ACTIONS; + return (episode * 3 + step * 7) % NUM_ACTIONS; } static uint64_t run_seed_sequence_checksum(unsigned int seed) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, seed, observations, actions, rewards, terminals); @@ -1072,7 +1028,7 @@ static uint64_t run_seed_sequence_checksum(unsigned int seed) { } } - affine_lock_free_shared(&shared); + free_shared(&shared); return checksum; } @@ -1081,8 +1037,8 @@ static void test_deterministic_seed_sequences(void) { AffineLock env_a; AffineLock env_b; - float obs_a[AFFINE_LOCK_OBS_SIZE], obs_b[AFFINE_LOCK_OBS_SIZE]; - float atn_a[AFFINE_LOCK_NUM_ATNS], atn_b[AFFINE_LOCK_NUM_ATNS]; + float obs_a[OBS_SIZE], obs_b[OBS_SIZE]; + float atn_a[NUM_ATNS], atn_b[NUM_ATNS]; float rew_a[1], rew_b[1]; float term_a[1], term_b[1]; make_env(&env_a, &shared, 12345, obs_a, atn_a, rew_a, term_a); @@ -1111,7 +1067,7 @@ static void test_deterministic_seed_sequences(void) { } } - affine_lock_free_shared(&shared); + free_shared(&shared); uint64_t seed_1 = run_seed_sequence_checksum(1); uint64_t seed_1_repeat = run_seed_sequence_checksum(1); @@ -1126,8 +1082,8 @@ static uint64_t run_visible_table_seed_42_golden_sequence(void) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; - float observations[AFFINE_LOCK_OBS_SIZE]; - float actions[AFFINE_LOCK_NUM_ATNS]; + float observations[OBS_SIZE]; + float actions[NUM_ATNS]; float rewards[1]; float terminals[1]; make_env(&env, &shared, 42, observations, actions, rewards, terminals); @@ -1152,13 +1108,13 @@ static uint64_t run_visible_table_seed_42_golden_sequence(void) { checksum = mix_u64(checksum, reset_snapshot_checksum(&env)); EXPECT_EQ_INT(env.scramble_depth, 2); - affine_lock_free_shared(&shared); + free_shared(&shared); return checksum; } static void test_visible_table_seed_42_golden_checksum(void) { uint64_t checksum = run_visible_table_seed_42_golden_sequence(); - EXPECT_EQ_U64(checksum, 0x1b6d67bf767fd010ull); + EXPECT_EQ_U64(checksum, 0xc5721c1259a9fd50ull); } static void test_deterministic_seed_sequences_and_distinct_env_ids(void) { @@ -1166,8 +1122,8 @@ static void test_deterministic_seed_sequences_and_distinct_env_ids(void) { AffineLock env_a; AffineLock env_b; - float obs_a[AFFINE_LOCK_OBS_SIZE], obs_b[AFFINE_LOCK_OBS_SIZE]; - float atn_a[AFFINE_LOCK_NUM_ATNS], atn_b[AFFINE_LOCK_NUM_ATNS]; + float obs_a[OBS_SIZE], obs_b[OBS_SIZE]; + float atn_a[NUM_ATNS], atn_b[NUM_ATNS]; float rew_a[1], rew_b[1]; float term_a[1], term_b[1]; make_env(&env_a, &shared, 12345, obs_a, atn_a, rew_a, term_a); @@ -1197,10 +1153,10 @@ static void test_deterministic_seed_sequences_and_distinct_env_ids(void) { AffineLock env_2; AffineLock env_1_repeat; AffineLock env_2_repeat; - float obs_1[AFFINE_LOCK_OBS_SIZE], obs_2[AFFINE_LOCK_OBS_SIZE]; - float obs_1r[AFFINE_LOCK_OBS_SIZE], obs_2r[AFFINE_LOCK_OBS_SIZE]; - float atn_1[AFFINE_LOCK_NUM_ATNS], atn_2[AFFINE_LOCK_NUM_ATNS]; - float atn_1r[AFFINE_LOCK_NUM_ATNS], atn_2r[AFFINE_LOCK_NUM_ATNS]; + float obs_1[OBS_SIZE], obs_2[OBS_SIZE]; + float obs_1r[OBS_SIZE], obs_2r[OBS_SIZE]; + float atn_1[NUM_ATNS], atn_2[NUM_ATNS]; + float atn_1r[NUM_ATNS], atn_2r[NUM_ATNS]; float rew_1[1], rew_2[1], rew_1r[1], rew_2r[1]; float term_1[1], term_2[1], term_1r[1], term_2r[1]; make_env(&env_1, &shared, 1, obs_1, atn_1, rew_1, term_1); @@ -1232,7 +1188,7 @@ static void test_deterministic_seed_sequences_and_distinct_env_ids(void) { } EXPECT_TRUE(differs); - affine_lock_free_shared(&shared); + free_shared(&shared); } int main(void) { diff --git a/ocean/affine_lock/tests/test_affine_lock_log_export.c b/ocean/affine_lock/tests/test_affine_lock_log_export.c index 85e0919c2a..6bc400d4d3 100644 --- a/ocean/affine_lock/tests/test_affine_lock_log_export.c +++ b/ocean/affine_lock/tests/test_affine_lock_log_export.c @@ -142,7 +142,7 @@ static uint64_t binding_reset_checksum(const Env* env) { hash = mix_u64_for_binding_test(hash, env->target); hash = mix_u64_for_binding_test(hash, (uint64_t)(env->target_distance + 1)); hash = mix_u64_for_binding_test(hash, (uint64_t)env->solution_length); - for (int i = 0; i < AFFINE_LOCK_MAX_SOLUTION_DEPTH; i++) { + for (int i = 0; i < MAX_SOLUTION_DEPTH; i++) { hash = mix_u64_for_binding_test( hash, (uint64_t)(env->solution_actions[i] + 1)); } @@ -152,12 +152,12 @@ static uint64_t binding_reset_checksum(const Env* env) { static void assign_binding_env_buffers( Env* envs, int total_agents, - float observations[][AFFINE_LOCK_OBS_SIZE], + float observations[][OBS_SIZE], float actions[], float rewards[], float terminals[]) { memset(observations, 0, - (size_t)total_agents * AFFINE_LOCK_OBS_SIZE * sizeof(float)); + (size_t)total_agents * OBS_SIZE * sizeof(float)); memset(actions, 0, (size_t)total_agents * sizeof(float)); memset(rewards, 0, (size_t)total_agents * sizeof(float)); memset(terminals, 0, (size_t)total_agents * sizeof(float)); @@ -174,8 +174,8 @@ static void test_vec_init_visible_targets_repeat_across_runs_and_vary_by_env_id( Env* run_a = make_binding_env_batch(42, total_agents); Env* run_b = make_binding_env_batch(42, total_agents); - float obs_a[64][AFFINE_LOCK_OBS_SIZE]; - float obs_b[64][AFFINE_LOCK_OBS_SIZE]; + float obs_a[64][OBS_SIZE]; + float obs_b[64][OBS_SIZE]; float actions_a[64], actions_b[64]; float rewards_a[64], rewards_b[64]; float terminals_a[64], terminals_b[64]; @@ -216,12 +216,6 @@ static void test_vec_init_visible_targets_repeat_across_runs_and_vary_by_env_id( static void test_depth_solve_rates_are_conditional_on_depth_attempts(void) { Log log = {0}; - log.depth_2_rate = 0.25f; - log.depth_2_solve_rate = 0.125f; - log.depth_4_rate = 0.5f; - log.depth_4_solve_rate = 0.375f; - log.depth_5_rate = 0.25f; - log.depth_5_solve_rate = 0.125f; log.depth_6_rate = 0.25f; log.depth_6_solve_rate = 0.125f; log.depth_8_rate = 0.0f; @@ -236,19 +230,13 @@ static void test_depth_solve_rates_are_conditional_on_depth_attempts(void) { Dict out = {0}; puf_log(&log, &out); - EXPECT_EQ_INT(out.size, 19); + EXPECT_EQ_INT(out.size, 16); EXPECT_NEAR(dict_value(&out, "score"), 0.75, 0.0); EXPECT_TRUE(!dict_has_key(&out, "solve_steps")); EXPECT_TRUE(!dict_has_key(&out, "solve_efficiency")); EXPECT_TRUE(!dict_has_key(&out, "scramble_unique_states")); EXPECT_NEAR(dict_value(&out, "min_win_moves"), 4.0, 0.0); EXPECT_NEAR(dict_value(&out, "solved_min_win_moves"), 4.0, 0.0); - EXPECT_TRUE(!dict_has_key(&out, "depth_2_rate")); - EXPECT_NEAR(dict_value(&out, "depth_2_solve_rate"), 0.5, 0.0); - EXPECT_TRUE(!dict_has_key(&out, "depth_4_rate")); - EXPECT_NEAR(dict_value(&out, "depth_4_solve_rate"), 0.75, 0.0); - EXPECT_TRUE(!dict_has_key(&out, "depth_5_rate")); - EXPECT_NEAR(dict_value(&out, "depth_5_solve_rate"), 0.5, 0.0); EXPECT_TRUE(!dict_has_key(&out, "depth_6_rate")); EXPECT_NEAR(dict_value(&out, "depth_6_solve_rate"), 0.5, 0.0); EXPECT_TRUE(!dict_has_key(&out, "depth_8_rate")); diff --git a/ocean/affine_lock/tests/test_metadata_smoke.py b/ocean/affine_lock/tests/test_metadata_smoke.py index 870fcddb6e..477012ecfc 100644 --- a/ocean/affine_lock/tests/test_metadata_smoke.py +++ b/ocean/affine_lock/tests/test_metadata_smoke.py @@ -22,12 +22,10 @@ "solved_min_win_moves", "conditional_solve_steps", "conditional_solve_efficiency", - "depth_2_solve_rate", - "depth_4_solve_rate", - "depth_5_solve_rate", "depth_6_solve_rate", "depth_8_solve_rate", "depth_16_solve_rate", + "n", ] @@ -138,9 +136,9 @@ def check_config(): def check_header_text(): header = (ROOT / "ocean" / "affine_lock" / "affine_lock.h").read_text() - assert "#define OBS_SIZE AFFINE_LOCK_OBS_SIZE" in header - assert "#define ACT_SIZES {AFFINE_LOCK_NUM_ACTIONS}" in header - assert "#define NUM_ATNS AFFINE_LOCK_NUM_ATNS" in header + assert "#define OBS_SIZE (TIMER_INDEX + 1)" in header + assert "#define ACT_SIZES {NUM_ACTIONS}" in header + assert "#define NUM_ATNS 1" in header assert "typedef" in header and "obs_t" in header assert "void puf_init(" in header assert "void puf_reset(" in header @@ -152,7 +150,7 @@ def check_header_text(): log_keys = re.findall(r'dict_set\(out,\s*"([^"]+)"', header) assert log_keys == EXPECTED_MY_LOG_KEYS - assert len(log_keys) + 1 <= 32 # trainer vec_log appends "n". + assert len(log_keys) <= 32 def float_buffer(ptr, count): diff --git a/ocean/affine_lock/tests/test_visible_targets_loader.c b/ocean/affine_lock/tests/test_visible_targets_loader.c index 9fed276f83..fd56accc16 100644 --- a/ocean/affine_lock/tests/test_visible_targets_loader.c +++ b/ocean/affine_lock/tests/test_visible_targets_loader.c @@ -54,15 +54,12 @@ int main(int argc, char** argv) { EXPECT_TRUE(end != argv[4] && *end == '\0'); AffineLockVisibleTargetTable table; - char error[256]; - int rc = affine_lock_visible_targets_load( + int rc = visible_targets_load( argv[1], - AFFINE_LOCK_VISIBLE_TARGET_8ACTION_V1_HASH, - &table, - error, - sizeof(error)); + VISIBLE_TARGET_8ACTION_V1_HASH, + &table); if (rc != 0) { - fprintf(stderr, "failed to load visible target table: %s\n", error); + fprintf(stderr, "failed to load visible target table: %s\n", argv[1]); return 1; } @@ -73,7 +70,7 @@ int main(int argc, char** argv) { EXPECT_EQ_U32(table.record_count, expected_record_count); EXPECT_EQ_U64( table.action_set_hash, - AFFINE_LOCK_VISIBLE_TARGET_8ACTION_V1_HASH); + VISIBLE_TARGET_8ACTION_V1_HASH); const uint32_t expected_depths[6] = {2, 4, 5, 6, 8, 16}; const uint64_t expected_exact_counts[6] = { @@ -111,6 +108,6 @@ int main(int argc, char** argv) { } } - affine_lock_visible_targets_free(&table); + visible_targets_free(&table); return 0; } diff --git a/src/puffercpu.c b/src/puffercpu.c index d343d31cb0..1a6363e842 100644 --- a/src/puffercpu.c +++ b/src/puffercpu.c @@ -895,6 +895,9 @@ int main(int argc, char** argv) { SetTargetFPS(60); } #endif + if (!headless) { + puf_render(&env); + } while (headless ? (steps < headless_max_steps && env.log.n < (float)headless_episodes) : !WindowShouldClose()) { From 7cb35b96e8a51538dfcb9626cf4e2ad47a0468ad Mon Sep 17 00:00:00 2001 From: Kinvert Date: Mon, 17 Aug 2026 03:29:13 -0400 Subject: [PATCH 2/4] More Affine Lock Cleanup --- ocean/affine_lock/affine_lock.h | 154 ++++++------------ .../affine_lock/affine_lock_visible_targets.h | 38 ++--- ocean/affine_lock/tests/run_all.sh | 2 + ocean/affine_lock/tests/test_affine_lock.c | 128 +++++++-------- .../tests/test_affine_lock_log_export.c | 26 +-- .../affine_lock/tests/test_metadata_smoke.py | 10 +- .../tests/test_visible_targets_loader.c | 4 +- 7 files changed, 152 insertions(+), 210 deletions(-) diff --git a/ocean/affine_lock/affine_lock.h b/ocean/affine_lock/affine_lock.h index e7c48a184b..c8c97beccd 100644 --- a/ocean/affine_lock/affine_lock.h +++ b/ocean/affine_lock/affine_lock.h @@ -17,7 +17,7 @@ typedef float obs_t; #define STEP_REWARD (-0.01f) // TODO should this be in ini so it can be swept? #define VISIBLE_TARGET_TABLE_PATH "ocean/affine_lock/generated/affine_lock_8action_visible_targets.bin" #define ACT_SIZES {NUM_ACTIONS} -#define PUF_STEPS_PER_SEC 2 // TODO remove this? +#define PUF_STEPS_PER_SEC 2 #define MY_VEC_INIT #define MY_VEC_CLOSE @@ -45,16 +45,15 @@ struct Log { float episode_length; float solve_steps; float timeout_rate; - float invalid_rate; // TODO is this even needed? float solve_efficiency; float target_distance; float solved_target_distance; - float depth_6_rate; - float depth_6_solve_rate; - float depth_8_rate; - float depth_8_solve_rate; - float depth_16_rate; - float depth_16_solve_rate; + float d6_rate; + float d6_solve_rate; + float d8_rate; + float d8_solve_rate; + float d16_rate; + float d16_solve_rate; float n; }; @@ -62,63 +61,50 @@ typedef struct AffineLockShared { int start_depth; int max_depth; int step_grace; - int num_states; // TODO Only used in testing uint32_t mask; uint32_t* next; - AffineLockVisibleTargetTable visible_target_table; + VisibleTargetTable visible_target_table; float observation_bit_patterns[256][8]; } AffineLockShared; -typedef struct Client { // TODO see if we even need this several envs don't have it - int screen_width; - int screen_height; -} Client; - struct Env { Log log; Agent agents[1]; int tag; int boundary_reached; - int num_agents; // TODO is this even needed? + int num_agents; unsigned int rng; uint32_t state; uint32_t target; int step_count; int max_steps; int scramble_depth; - int curriculum_depth; + int curriculum_depth; // TODO Consider only demoting by one level on a loss, perhaps both options in the ini int solution_length; int solution_actions[MAX_SOLUTION_DEPTH]; int target_distance; float episode_return; int owns_shared; AffineLockShared* shared; - Client* client; }; typedef Env AffineLock; -static void init_shared( - AffineLockShared* shared, - int start_depth, - int max_depth, - int step_grace) { +static void init_shared(AffineLockShared* shared, int start_depth, int max_depth, int step_grace) { shared->start_depth = start_depth; shared->max_depth = max_depth; shared->step_grace = step_grace; - shared->num_states = 1 << BITS; shared->mask = (1u << BITS) - 1u; for (int value = 0; value < 256; value++) { for (int bit = 0; bit < 8; bit++) { - shared->observation_bit_patterns[value][bit] = - (value & (1 << bit)) ? 1.0f : -1.0f; + shared->observation_bit_patterns[value][bit] = (value & (1 << bit)) ? 1.0f : -1.0f; } } + uint32_t num_states = 1u << BITS; shared->next = (uint32_t*)calloc( - shared->num_states * NUM_ACTIONS, sizeof(uint32_t)); + num_states * NUM_ACTIONS, sizeof(uint32_t)); - uint32_t nstates = shared->num_states; - for (uint32_t state = 0; state < nstates; state++) { + for (uint32_t state = 0; state < num_states; state++) { for (int action = 0; action < NUM_ACTIONS; action++) { uint32_t next = state; switch (action) { @@ -174,27 +160,13 @@ static void init_env(AffineLock* env, AffineLockShared* shared, unsigned int see env->curriculum_depth = shared->start_depth; } -static unsigned int env_seed( // TODO see if this is in the golden envs - unsigned int base_seed, unsigned int env_id) { - uint32_t value = 0x811c9dc5u; - value = (value ^ base_seed) * 0x01000193u; - value = (value ^ env_id) * 0x01000193u; - value ^= value >> 16; - value *= 0x7feb352du; - value ^= value >> 15; - value *= 0x846ca68bu; - value ^= value >> 16; - return value; -} - void puf_init(Env* env, Dict* kwargs) { int start_depth = dict_get(kwargs, "start_depth"); int max_depth = dict_get(kwargs, "max_depth"); int step_grace = dict_get(kwargs, "step_grace"); - int seed = dict_get(kwargs, "seed"); - AffineLockShared* shared = - create_shared(start_depth, max_depth, step_grace); - init_env(env, shared, env_seed(seed, env->rng)); + unsigned int seed = (unsigned int)dict_get(kwargs, "seed"); + AffineLockShared* shared = create_shared(start_depth, max_depth, step_grace); + init_env(env, shared, rand_r(&seed)); env->owns_shared = 1; } @@ -204,11 +176,8 @@ static void free_shared(AffineLockShared* shared) { } void puf_close(AffineLock* env) { - if (env->client) { - if (IsWindowReady()) { - CloseWindow(); - } - free(env->client); + if (IsWindowReady()) { + CloseWindow(); } if (env->owns_shared) { free_shared(env->shared); @@ -216,7 +185,7 @@ void puf_close(AffineLock* env) { } } -static void add_log(AffineLock* env, int solved, int invalid) { +static void add_log(AffineLock* env, int solved) { AffineLockShared* shared = env->shared; int log_depth = env->target_distance; int at_max_depth = log_depth == shared->max_depth; @@ -228,24 +197,19 @@ static void add_log(AffineLock* env, int solved, int invalid) { env->log.episode_return += env->episode_return; env->log.episode_length += env->step_count; env->log.solve_steps += solved ? env->step_count : 0; - env->log.timeout_rate += !solved && !invalid; - env->log.invalid_rate += invalid; + env->log.timeout_rate += !solved; env->log.solve_efficiency += solved ? env->step_count / (float)log_depth : 0; env->log.target_distance += env->target_distance; env->log.solved_target_distance += solved ? env->target_distance : 0; - env->log.depth_6_rate += log_depth == 6; - env->log.depth_6_solve_rate += solved && log_depth == 6; - env->log.depth_8_rate += log_depth == 8; - env->log.depth_8_solve_rate += solved && log_depth == 8; - env->log.depth_16_rate += log_depth == 16; - env->log.depth_16_solve_rate += solved && log_depth == 16; + env->log.d6_rate += log_depth == 6; + env->log.d6_solve_rate += solved && log_depth == 6; + env->log.d8_rate += log_depth == 8; + env->log.d8_solve_rate += solved && log_depth == 8; + env->log.d16_rate += log_depth == 16; + env->log.d16_solve_rate += solved && log_depth == 16; env->log.n += 1; } -static uint32_t apply_action(const AffineLockShared* shared, uint32_t rel, int action) { - return shared->next[rel * NUM_ACTIONS + action]; -} - // Not rand_r(): glibc's LCG has statistically weak low-order bits, and this // env repeatedly samples individual state bits and small action ranges // directly from those bits, where the weakness would show up as bias. @@ -270,10 +234,10 @@ static int random_bounded(AffineLock* env, int bound) { return value % ubound; } -static const AffineLockVisibleTargetDepth* visible_target_depth( +static const VisibleTargetDepth* visible_target_depth( const AffineLockShared* shared, uint32_t requested_depth) { - const AffineLockVisibleTargetTable* table = &shared->visible_target_table; + const VisibleTargetTable* table = &shared->visible_target_table; for (uint32_t i = 0; i < table->depth_count; i++) { if (table->depths[i].depth == requested_depth) { return &table->depths[i]; @@ -287,11 +251,9 @@ static void reset_state(AffineLock* env) { env->scramble_depth = env->curriculum_depth; env->step_count = 0; env->episode_return = 0; - const AffineLockVisibleTargetDepth* depth = - visible_target_depth(shared, env->scramble_depth); + const VisibleTargetDepth* depth = visible_target_depth(shared, env->scramble_depth); int choice = random_bounded(env, depth->stored_count); - const AffineLockVisibleTargetRecord* record = - &shared->visible_target_table.records[depth->first_record + choice]; + const VisibleTargetRecord* record = &shared->visible_target_table.records[depth->first_record + choice]; env->state = record->start; env->target = record->target; env->target_distance = record->depth; @@ -326,9 +288,7 @@ void puf_reset(AffineLock* env) { compute_observations(env); } -static int next_curriculum_depth( - const AffineLockShared* shared, - int current_depth) { +static int next_curriculum_depth( const AffineLockShared* shared, int current_depth) { for (int i = 0; i < CURRICULUM_DEPTH_COUNT; i++) { int depth = CURRICULUM_DEPTHS[i]; if (depth > current_depth) { @@ -360,24 +320,20 @@ void puf_step(AffineLock* env) { return; } AffineLockShared* shared = env->shared; - float raw = env->agents[0].actions[0]; - int invalid = !isfinite(raw) || raw < 0 || raw > NUM_ACTIONS - 1; - int action = invalid ? -1 : raw; - if (!invalid && action != raw) { - invalid = 1; - } float reward = STEP_REWARD; int terminal = 0; int solved = 0; env->agents[0].terminals[0] = 0; env->step_count += 1; - + float raw = env->agents[0].actions[0]; + int invalid = !isfinite(raw) || raw < 0 || raw > NUM_ACTIONS - 1; if (invalid) { reward = -1; terminal = 1; } else { - env->state = apply_action(shared, env->state, action); + int action = (int)raw; + env->state = shared->next[env->state * NUM_ACTIONS + action]; if (env->state == env->target) { reward = 1; terminal = 1; @@ -391,7 +347,7 @@ void puf_step(AffineLock* env) { env->episode_return += reward; if (terminal) { env->agents[0].terminals[0] = 1; - add_log(env, solved, invalid); + add_log(env, solved); env->curriculum_depth = solved ? next_curriculum_depth(shared, env->scramble_depth) : shared->start_depth; @@ -404,8 +360,7 @@ void puf_log(Log* log, Dict* out) { float nsolve = log->solve_rate; float solved_min_win_moves = nsolve ? log->solved_target_distance / nsolve : 0; float conditional_solve_steps = nsolve ? log->solve_steps / nsolve : 0; - float conditional_solve_efficiency = nsolve ? - log->solve_efficiency / nsolve : 0; + float conditional_solve_efficiency = nsolve ? log->solve_efficiency / nsolve : 0; dict_set(out, "perf", log->perf); dict_set(out, "score", log->score); @@ -414,14 +369,13 @@ void puf_log(Log* log, Dict* out) { dict_set(out, "episode_return", log->episode_return); dict_set(out, "episode_length", log->episode_length); dict_set(out, "timeout_rate", log->timeout_rate); - dict_set(out, "invalid_rate", log->invalid_rate); dict_set(out, "min_win_moves", log->target_distance); dict_set(out, "solved_min_win_moves", solved_min_win_moves); dict_set(out, "conditional_solve_steps", conditional_solve_steps); dict_set(out, "conditional_solve_efficiency", conditional_solve_efficiency); - dict_set(out, "depth_6_solve_rate", log->depth_6_rate ? log->depth_6_solve_rate / log->depth_6_rate : 0); - dict_set(out, "depth_8_solve_rate", log->depth_8_rate ? log->depth_8_solve_rate / log->depth_8_rate : 0); - dict_set(out, "depth_16_solve_rate", log->depth_16_rate ? log->depth_16_solve_rate / log->depth_16_rate : 0); + dict_set(out, "d6_solve_rate", log->d6_rate ? log->d6_solve_rate / log->d6_rate : 0); + dict_set(out, "d8_solve_rate", log->d8_rate ? log->d8_solve_rate / log->d8_rate : 0); + dict_set(out, "d16_solve_rate", log->d16_rate ? log->d16_solve_rate / log->d16_rate : 0); dict_set(out, "n", log->n); } @@ -430,13 +384,12 @@ Env* my_vec_init(int* num_envs_out, int* buffer_env_starts, int* buffer_env_coun int total_agents = dict_get(vec_kwargs, "total_agents"); int num_buffers = dict_get(vec_kwargs, "num_buffers"); int agents_per_buffer = total_agents / num_buffers; - int base_seed = dict_get(env_kwargs, "seed"); + unsigned int running_seed = (unsigned int)dict_get(env_kwargs, "seed"); int start_depth = dict_get(env_kwargs, "start_depth"); int max_depth = dict_get(env_kwargs, "max_depth"); int step_grace = dict_get(env_kwargs, "step_grace"); - AffineLockShared* shared = - create_shared(start_depth, max_depth, step_grace); + AffineLockShared* shared = create_shared(start_depth, max_depth, step_grace); Env* envs = (Env*)calloc(total_agents, sizeof(Env)); int buf = 0; @@ -445,7 +398,7 @@ Env* my_vec_init(int* num_envs_out, int* buffer_env_starts, int* buffer_env_coun buffer_env_counts[0] = 0; for (int i = 0; i < total_agents; i++) { Env* env = &envs[i]; - init_env(env, shared, env_seed(base_seed, i)); + init_env(env, shared, rand_r(&running_seed)); buf_agents += env->num_agents; buffer_env_counts[buf]++; if (buf_agents >= agents_per_buffer && buf < num_buffers - 1) { @@ -470,14 +423,11 @@ void puf_render(AffineLock* env) { exit(0); } human_controls(env); - if (!env->client) { - Client* client = (Client*)calloc(1, sizeof(Client)); - client->screen_width = 780; - client->screen_height = 360; - InitWindow(client->screen_width, client->screen_height, - "PufferLib AffineLock"); + static int window_initialized = 0; + if (!window_initialized) { + InitWindow(780, 360, "PufferLib AffineLock"); SetTargetFPS(30); - env->client = client; + window_initialized = 1; } uint32_t rel = (env->state ^ env->target) & env->shared->mask; @@ -511,10 +461,8 @@ void puf_render(AffineLock* env) { int x = 145 + bit * 34; int on = (row_value[row] >> bit) & 1u; int mismatch = ((env->state ^ env->target) >> bit) & 1u; - Color fill = on ? - (Color){80, 210, 140, 255} : (Color){38, 48, 58, 255}; - Color border = mismatch ? - (Color){238, 88, 88, 255} : (Color){182, 196, 205, 255}; + Color fill = on ? (Color){80, 210, 140, 255} : (Color){38, 48, 58, 255}; + Color border = mismatch ? (Color){238, 88, 88, 255} : (Color){182, 196, 205, 255}; DrawRectangle(x, row_y[row], 24, 34, fill); DrawRectangleLinesEx( (Rectangle){x, row_y[row], 24, 34}, diff --git a/ocean/affine_lock/affine_lock_visible_targets.h b/ocean/affine_lock/affine_lock_visible_targets.h index 63a814bc5b..a013b45e56 100644 --- a/ocean/affine_lock/affine_lock_visible_targets.h +++ b/ocean/affine_lock/affine_lock_visible_targets.h @@ -9,22 +9,22 @@ #define VISIBLE_TARGET_RECORD_SIZE 16u #define VISIBLE_TARGET_8ACTION_V1_HASH 0x6e11e18fdafc0baaull -typedef struct AffineLockVisibleTargetDepth { +typedef struct VisibleTargetDepth { uint32_t depth; uint32_t first_record; uint32_t stored_count; uint64_t exact_pair_count; -} AffineLockVisibleTargetDepth; +} VisibleTargetDepth; -typedef struct AffineLockVisibleTargetRecord { +typedef struct VisibleTargetRecord { uint16_t start; uint16_t target; uint64_t packed_actions; uint8_t solution_length; uint8_t depth; -} AffineLockVisibleTargetRecord; +} VisibleTargetRecord; -typedef struct AffineLockVisibleTargetTable { +typedef struct VisibleTargetTable { uint32_t version; uint32_t header_size; uint32_t record_size; @@ -34,9 +34,9 @@ typedef struct AffineLockVisibleTargetTable { uint32_t record_count; uint64_t checksum; uint64_t action_set_hash; - AffineLockVisibleTargetDepth* depths; - AffineLockVisibleTargetRecord* records; -} AffineLockVisibleTargetTable; + VisibleTargetDepth* depths; + VisibleTargetRecord* records; +} VisibleTargetTable; static uint64_t visible_targets_mix_u64( uint64_t hash, @@ -94,7 +94,7 @@ static int visible_targets_read_u64( } static void visible_targets_free( - AffineLockVisibleTargetTable* table) { + VisibleTargetTable* table) { if (table == NULL) { return; } @@ -104,18 +104,18 @@ static void visible_targets_free( } static uint64_t visible_targets_checksum( - const AffineLockVisibleTargetTable* table) { + const VisibleTargetTable* table) { uint64_t hash = 1469598103934665603ull; hash = visible_targets_mix_u64(hash, table->action_set_hash); for (uint32_t depth_index = 0; depth_index < table->depth_count; depth_index++) { - const AffineLockVisibleTargetDepth* depth = &table->depths[depth_index]; + const VisibleTargetDepth* depth = &table->depths[depth_index]; hash = visible_targets_mix_u64(hash, depth->depth); hash = visible_targets_mix_u64(hash, depth->exact_pair_count); hash = visible_targets_mix_u64(hash, depth->stored_count); for (uint32_t i = 0; i < depth->stored_count; i++) { uint32_t record_index = depth->first_record + i; - const AffineLockVisibleTargetRecord* record = + const VisibleTargetRecord* record = &table->records[record_index]; hash = visible_targets_mix_u64(hash, record->start); hash = visible_targets_mix_u64(hash, record->target); @@ -132,7 +132,7 @@ static uint64_t visible_targets_checksum( static int visible_targets_load( const char* path, uint64_t expected_action_set_hash, - AffineLockVisibleTargetTable* table) { + VisibleTargetTable* table) { static const unsigned char expected_magic[8] = { 'A', 'L', '7', 'T', 'G', 'T', '1', '\0' }; @@ -183,10 +183,10 @@ static int visible_targets_load( return -1; } - table->depths = (AffineLockVisibleTargetDepth*)calloc( - table->depth_count, sizeof(AffineLockVisibleTargetDepth)); - table->records = (AffineLockVisibleTargetRecord*)calloc( - table->record_count, sizeof(AffineLockVisibleTargetRecord)); + table->depths = (VisibleTargetDepth*)calloc( + table->depth_count, sizeof(VisibleTargetDepth)); + table->records = (VisibleTargetRecord*)calloc( + table->record_count, sizeof(VisibleTargetRecord)); if (table->depths == NULL || table->records == NULL) { fclose(file); visible_targets_free(table); @@ -195,7 +195,7 @@ static int visible_targets_load( uint64_t depth_record_total = 0; for (uint32_t i = 0; i < table->depth_count; i++) { - AffineLockVisibleTargetDepth* depth = &table->depths[i]; + VisibleTargetDepth* depth = &table->depths[i]; uint32_t reserved = 0; if (visible_targets_read_u32(file, &depth->depth) != 0 || visible_targets_read_u32( @@ -227,7 +227,7 @@ static int visible_targets_load( } for (uint32_t i = 0; i < table->record_count; i++) { - AffineLockVisibleTargetRecord* record = &table->records[i]; + VisibleTargetRecord* record = &table->records[i]; uint16_t reserved = 0; if (visible_targets_read_u16(file, &record->start) != 0 || visible_targets_read_u16(file, &record->target) != 0 || diff --git a/ocean/affine_lock/tests/run_all.sh b/ocean/affine_lock/tests/run_all.sh index b9a119cc71..d9521ba3bf 100755 --- a/ocean/affine_lock/tests/run_all.sh +++ b/ocean/affine_lock/tests/run_all.sh @@ -24,6 +24,7 @@ bash "$ROOT/ocean/affine_lock/tests/test_8action_visible_targets_smoke.sh" "$CC_BIN" \ -std=c99 -pedantic -Wall -Wextra -Werror -Wno-unused-function \ + -D_POSIX_C_SOURCE=200809L \ -O0 -ffunction-sections -fdata-sections \ -I"$ROOT" -I"$ROOT/src" -I"$ROOT/ocean/affine_lock" -I"$ROOT/vendor" \ -I"$RAYLIB_INC" \ @@ -32,6 +33,7 @@ bash "$ROOT/ocean/affine_lock/tests/test_8action_visible_targets_smoke.sh" "$CC_BIN" \ -std=c11 -Wall -Wextra -Werror -Wno-unused-function \ + -D_POSIX_C_SOURCE=200809L \ -O0 -g -ffunction-sections -fdata-sections -fsanitize=address,undefined \ -I"$ROOT" -I"$ROOT/src" -I"$ROOT/ocean/affine_lock" -I"$ROOT/vendor" \ -I"$RAYLIB_INC" \ diff --git a/ocean/affine_lock/tests/test_affine_lock.c b/ocean/affine_lock/tests/test_affine_lock.c index c4b0cae183..cad393dad4 100644 --- a/ocean/affine_lock/tests/test_affine_lock.c +++ b/ocean/affine_lock/tests/test_affine_lock.c @@ -148,13 +148,14 @@ static void compute_test_bfs_stats( memset(stats, 0, sizeof(*stats)); stats->shortest_distance = -1; - int* distances = (int*)malloc((size_t)shared->num_states * sizeof(int)); + int num_states = 1 << BITS; + int* distances = (int*)malloc((size_t)num_states * sizeof(int)); uint32_t* queue = - (uint32_t*)malloc((size_t)shared->num_states * sizeof(uint32_t)); + (uint32_t*)malloc((size_t)num_states * sizeof(uint32_t)); EXPECT_TRUE(distances != NULL); EXPECT_TRUE(queue != NULL); - for (int i = 0; i < shared->num_states; i++) { + for (int i = 0; i < num_states; i++) { distances[i] = -1; } @@ -203,7 +204,7 @@ static void test_log_solve_credit_uses_known_target_distance(void) { env.target_distance = 8; env.step_count = 8; - add_log(&env, 1, 0); + add_log(&env, 1); EXPECT_NEAR(env.log.perf, expected_solve_credit(&shared, 8), 0.0f); EXPECT_NEAR(env.log.score, expected_solve_credit(&shared, 8), 0.0f); @@ -211,12 +212,12 @@ static void test_log_solve_credit_uses_known_target_distance(void) { EXPECT_NEAR(env.log.solve_efficiency, 1.0f, 0.0f); EXPECT_NEAR(env.log.target_distance, 8.0f, 0.0f); EXPECT_NEAR(env.log.solved_target_distance, 8.0f, 0.0f); - EXPECT_NEAR(env.log.depth_6_rate, 0.0f, 0.0f); - EXPECT_NEAR(env.log.depth_6_solve_rate, 0.0f, 0.0f); - EXPECT_NEAR(env.log.depth_8_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_8_solve_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_16_rate, 0.0f, 0.0f); - EXPECT_NEAR(env.log.depth_16_solve_rate, 0.0f, 0.0f); + EXPECT_NEAR(env.log.d6_rate, 0.0f, 0.0f); + EXPECT_NEAR(env.log.d6_solve_rate, 0.0f, 0.0f); + EXPECT_NEAR(env.log.d8_rate, 1.0f, 0.0f); + EXPECT_NEAR(env.log.d8_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(env.log.d16_rate, 0.0f, 0.0f); + EXPECT_NEAR(env.log.d16_solve_rate, 0.0f, 0.0f); free_shared(&shared); } @@ -244,7 +245,7 @@ static void expect_observation_matches(const AffineLock* env) { static int find_non_solving_action(AffineLock* env) { for (int action = 0; action < NUM_ACTIONS; action++) { - uint32_t next = apply_action(env->shared, env->state, action); + uint32_t next = env->shared->next[env->state * NUM_ACTIONS + action]; if (next != env->target) { return action; } @@ -277,16 +278,15 @@ static uint64_t log_snapshot_checksum(uint64_t hash, const Log* log) { hash = mix_float(hash, log->episode_length); hash = mix_float(hash, log->solve_steps); hash = mix_float(hash, log->timeout_rate); - hash = mix_float(hash, log->invalid_rate); hash = mix_float(hash, log->solve_efficiency); hash = mix_float(hash, log->target_distance); hash = mix_float(hash, log->solved_target_distance); - hash = mix_float(hash, log->depth_6_rate); - hash = mix_float(hash, log->depth_6_solve_rate); - hash = mix_float(hash, log->depth_8_rate); - hash = mix_float(hash, log->depth_8_solve_rate); - hash = mix_float(hash, log->depth_16_rate); - hash = mix_float(hash, log->depth_16_solve_rate); + hash = mix_float(hash, log->d6_rate); + hash = mix_float(hash, log->d6_solve_rate); + hash = mix_float(hash, log->d8_rate); + hash = mix_float(hash, log->d8_solve_rate); + hash = mix_float(hash, log->d16_rate); + hash = mix_float(hash, log->d16_solve_rate); hash = mix_float(hash, log->n); return hash; } @@ -367,18 +367,18 @@ static void expect_depth_log_delta( const Log* after, int depth, int solved) { - EXPECT_NEAR(after->depth_6_rate, - before->depth_6_rate + (depth == 6 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_6_solve_rate, - before->depth_6_solve_rate + (solved && depth == 6 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_8_rate, - before->depth_8_rate + (depth == 8 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_8_solve_rate, - before->depth_8_solve_rate + (solved && depth == 8 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_16_rate, - before->depth_16_rate + (depth == 16 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(after->depth_16_solve_rate, - before->depth_16_solve_rate + (solved && depth == 16 ? 1.0f : 0.0f), 0.0f); + EXPECT_NEAR(after->d6_rate, + before->d6_rate + (depth == 6 ? 1.0f : 0.0f), 0.0f); + EXPECT_NEAR(after->d6_solve_rate, + before->d6_solve_rate + (solved && depth == 6 ? 1.0f : 0.0f), 0.0f); + EXPECT_NEAR(after->d8_rate, + before->d8_rate + (depth == 8 ? 1.0f : 0.0f), 0.0f); + EXPECT_NEAR(after->d8_solve_rate, + before->d8_solve_rate + (solved && depth == 8 ? 1.0f : 0.0f), 0.0f); + EXPECT_NEAR(after->d16_rate, + before->d16_rate + (depth == 16 ? 1.0f : 0.0f), 0.0f); + EXPECT_NEAR(after->d16_solve_rate, + before->d16_solve_rate + (solved && depth == 16 ? 1.0f : 0.0f), 0.0f); } static void expect_oracle_episode_win(AffineLock* env, int depth) { @@ -414,7 +414,6 @@ static void expect_oracle_episode_win(AffineLock* env, int depth) { before.score + expected_solve_credit(shared, depth), 0.0f); EXPECT_NEAR(env->log.solve_rate, before.solve_rate + 1.0f, 0.0f); EXPECT_NEAR(env->log.timeout_rate, before.timeout_rate, 0.0f); - EXPECT_NEAR(env->log.invalid_rate, before.invalid_rate, 0.0f); EXPECT_NEAR(env->log.episode_length, before.episode_length + (float)solution_length, 0.0f); EXPECT_NEAR(env->log.solve_steps, @@ -462,7 +461,6 @@ static void expect_non_solving_episode_timeout(AffineLock* env, int depth) { EXPECT_NEAR(env->log.score, before.score, 0.0f); EXPECT_NEAR(env->log.solve_rate, before.solve_rate, 0.0f); EXPECT_NEAR(env->log.timeout_rate, before.timeout_rate + 1.0f, 0.0f); - EXPECT_NEAR(env->log.invalid_rate, before.invalid_rate, 0.0f); EXPECT_NEAR(env->log.episode_length, before.episode_length + (float)max_steps, 0.0f); EXPECT_NEAR(env->log.solve_steps, before.solve_steps, 0.0f); @@ -533,7 +531,7 @@ static void test_global_action_examples(void) { }; for (int action = 0; action < NUM_ACTIONS; action++) { - uint32_t next = apply_action(&shared, start, action); + uint32_t next = shared.next[start * NUM_ACTIONS + action]; EXPECT_EQ_U32(next, bits_from_text(expected[action])); } @@ -552,7 +550,6 @@ static void test_actions_round_trip_for_all_states(void) { ACTION_REVERSE_EACH_NIBBLE, ACTION_REVERSE_EACH_BYTE, }; - EXPECT_EQ_INT(shared.num_states, 1 << 16); EXPECT_EQ_U32(shared.mask, 0xffffu); for (int action = 0; action < NUM_ACTIONS; action++) { @@ -560,10 +557,10 @@ static void test_actions_round_trip_for_all_states(void) { EXPECT_TRUE(inverse >= 0 && inverse < NUM_ACTIONS); EXPECT_EQ_INT(inverse_actions[inverse], action); - for (uint32_t state = 0; state < (uint32_t)shared.num_states; state++) { - uint32_t next = apply_action(&shared, state, action); + for (uint32_t state = 0; state < (1u << BITS); state++) { + uint32_t next = shared.next[state * NUM_ACTIONS + action]; EXPECT_EQ_U32(next & ~shared.mask, 0u); - uint32_t round_trip = apply_action(&shared, next, inverse); + uint32_t round_trip = shared.next[next * NUM_ACTIONS + inverse]; EXPECT_EQ_U32(round_trip, state); } } @@ -673,7 +670,7 @@ static void test_visible_target_table_reset_uses_stored_records(void) { for (int depth_index = 0; depth_index < 6; depth_index++) { int requested_depth = requested_depths[depth_index]; AffineLockShared shared = make_shared(requested_depth, 16, 0); - const AffineLockVisibleTargetDepth* table_depth = + const VisibleTargetDepth* table_depth = visible_target_depth(&shared, requested_depth); EXPECT_TRUE(table_depth != NULL); EXPECT_EQ_INT((int)table_depth->stored_count, @@ -802,7 +799,7 @@ static void test_actions_apply_to_current_state_directly(void) { uint32_t target = bits_from_text("1111000011110000"); uint32_t state = bits_from_text("0011011000010111"); int action = 1; - uint32_t expected_state = apply_action(&shared, state, action); + uint32_t expected_state = shared.next[state * NUM_ACTIONS + action]; EXPECT_NE_U32(expected_state, target); env.target = target; @@ -820,7 +817,7 @@ static void test_actions_apply_to_current_state_directly(void) { free_shared(&shared); } -static void test_action_float_validation_rejects_non_discrete_values(void) { +static void test_action_float_validation_rejects_out_of_range_values(void) { AffineLockShared shared = make_shared(2, 16, 0); AffineLock env; float observations[OBS_SIZE]; @@ -832,7 +829,6 @@ static void test_action_float_validation_rejects_non_discrete_values(void) { const float invalid_actions[] = { -1.0f, 8.0f, - 1.5f, NAN, INFINITY, -INFINITY, @@ -840,7 +836,7 @@ static void test_action_float_validation_rejects_non_discrete_values(void) { int count = (int)(sizeof(invalid_actions) / sizeof(invalid_actions[0])); for (int i = 0; i < count; i++) { puf_reset(&env); - float prev_invalid = env.log.invalid_rate; + float prev_timeout = env.log.timeout_rate; float prev_n = env.log.n; actions[0] = invalid_actions[i]; @@ -848,7 +844,7 @@ static void test_action_float_validation_rejects_non_discrete_values(void) { EXPECT_NEAR(rewards[0], -1.0f, 0.0f); EXPECT_NEAR(terminals[0], 1.0f, 0.0f); - EXPECT_NEAR(env.log.invalid_rate, prev_invalid + 1.0f, 0.0f); + EXPECT_NEAR(env.log.timeout_rate, prev_timeout + 1.0f, 0.0f); EXPECT_NEAR(env.log.n, prev_n + 1.0f, 0.0f); EXPECT_EQ_INT(env.step_count, 0); } @@ -878,12 +874,12 @@ static void test_visible_target_table_curriculum_and_logging(void) { float prev_max_depth_solve = env.log.max_depth_solve; float prev_target_distance = env.log.target_distance; float prev_solved_target_distance = env.log.solved_target_distance; - float prev_depth_6 = env.log.depth_6_rate; - float prev_depth_6_solve = env.log.depth_6_solve_rate; - float prev_depth_8 = env.log.depth_8_rate; - float prev_depth_8_solve = env.log.depth_8_solve_rate; - float prev_depth_16 = env.log.depth_16_rate; - float prev_depth_16_solve = env.log.depth_16_solve_rate; + float prev_depth_6 = env.log.d6_rate; + float prev_depth_6_solve = env.log.d6_solve_rate; + float prev_depth_8 = env.log.d8_rate; + float prev_depth_8_solve = env.log.d8_solve_rate; + float prev_depth_16 = env.log.d16_rate; + float prev_depth_16_solve = env.log.d16_solve_rate; int target_distance = env.target_distance; int metric_depth = target_distance > 0 ? target_distance : depth; @@ -900,17 +896,17 @@ static void test_visible_target_table_curriculum_and_logging(void) { prev_target_distance + (float)target_distance, 0.0f); EXPECT_NEAR(env.log.solved_target_distance, prev_solved_target_distance + (float)target_distance, 0.0f); - EXPECT_NEAR(env.log.depth_6_rate, + EXPECT_NEAR(env.log.d6_rate, prev_depth_6 + (metric_depth == 6 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_6_solve_rate, + EXPECT_NEAR(env.log.d6_solve_rate, prev_depth_6_solve + (metric_depth == 6 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_8_rate, + EXPECT_NEAR(env.log.d8_rate, prev_depth_8 + (metric_depth == 8 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_8_solve_rate, + EXPECT_NEAR(env.log.d8_solve_rate, prev_depth_8_solve + (metric_depth == 8 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_16_rate, + EXPECT_NEAR(env.log.d16_rate, prev_depth_16 + (metric_depth == 16 ? 1.0f : 0.0f), 0.0f); - EXPECT_NEAR(env.log.depth_16_solve_rate, + EXPECT_NEAR(env.log.d16_solve_rate, prev_depth_16_solve + (metric_depth == 16 ? 1.0f : 0.0f), 0.0f); int next_depth = episode < 5 ? expected_depths[episode + 1] : 16; @@ -920,7 +916,7 @@ static void test_visible_target_table_curriculum_and_logging(void) { float prev_n = env.log.n; float prev_perf = env.log.perf; float prev_max_depth_solve = env.log.max_depth_solve; - float prev_invalid = env.log.invalid_rate; + float prev_timeout = env.log.timeout_rate; EXPECT_EQ_INT(env.scramble_depth, shared.max_depth); actions[0] = 999.0f; puf_step(&env); @@ -929,7 +925,7 @@ static void test_visible_target_table_curriculum_and_logging(void) { EXPECT_NEAR(env.log.n, prev_n + 1.0f, 0.0f); EXPECT_NEAR(env.log.perf, prev_perf, 0.0f); EXPECT_NEAR(env.log.max_depth_solve, prev_max_depth_solve, 0.0f); - EXPECT_NEAR(env.log.invalid_rate, prev_invalid + 1.0f, 0.0f); + EXPECT_NEAR(env.log.timeout_rate, prev_timeout + 1.0f, 0.0f); EXPECT_EQ_INT(env.scramble_depth, shared.start_depth); free_shared(&shared); @@ -953,14 +949,13 @@ static void test_visible_target_table_oracle_wins_all_curriculum_depths_end_to_e EXPECT_EQ_INT(env.scramble_depth, shared.max_depth); EXPECT_NEAR(env.log.n, 6.0f, 0.0f); - EXPECT_NEAR(env.log.depth_6_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_6_solve_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_8_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_8_solve_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_16_rate, 1.0f, 0.0f); - EXPECT_NEAR(env.log.depth_16_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(env.log.d6_rate, 1.0f, 0.0f); + EXPECT_NEAR(env.log.d6_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(env.log.d8_rate, 1.0f, 0.0f); + EXPECT_NEAR(env.log.d8_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(env.log.d16_rate, 1.0f, 0.0f); + EXPECT_NEAR(env.log.d16_solve_rate, 1.0f, 0.0f); EXPECT_NEAR(env.log.timeout_rate, 0.0f, 0.0f); - EXPECT_NEAR(env.log.invalid_rate, 0.0f, 0.0f); free_shared(&shared); } @@ -989,7 +984,6 @@ static void test_visible_target_table_timeouts_at_all_curriculum_depths_end_to_e EXPECT_EQ_INT(env.scramble_depth, shared.start_depth); EXPECT_TRUE(env.log.timeout_rate >= 1.0f); EXPECT_TRUE(env.log.solve_rate >= 0.0f); - EXPECT_NEAR(env.log.invalid_rate, 0.0f, 0.0f); free_shared(&shared); } @@ -1114,7 +1108,7 @@ static uint64_t run_visible_table_seed_42_golden_sequence(void) { static void test_visible_table_seed_42_golden_checksum(void) { uint64_t checksum = run_visible_table_seed_42_golden_sequence(); - EXPECT_EQ_U64(checksum, 0xc5721c1259a9fd50ull); + EXPECT_EQ_U64(checksum, 0x733eb55fe141e600ull); } static void test_deterministic_seed_sequences_and_distinct_env_ids(void) { @@ -1205,7 +1199,7 @@ int main(void) { test_observation_encoding_is_32_signed_bit_floats_plus_timer(); test_timer_observation_progresses_and_resets_after_timeout(); test_actions_apply_to_current_state_directly(); - test_action_float_validation_rejects_non_discrete_values(); + test_action_float_validation_rejects_out_of_range_values(); test_visible_target_table_curriculum_and_logging(); test_visible_target_table_oracle_wins_all_curriculum_depths_end_to_end(); test_visible_target_table_timeouts_at_all_curriculum_depths_end_to_end(); diff --git a/ocean/affine_lock/tests/test_affine_lock_log_export.c b/ocean/affine_lock/tests/test_affine_lock_log_export.c index 6bc400d4d3..7f3bd09312 100644 --- a/ocean/affine_lock/tests/test_affine_lock_log_export.c +++ b/ocean/affine_lock/tests/test_affine_lock_log_export.c @@ -216,12 +216,12 @@ static void test_vec_init_visible_targets_repeat_across_runs_and_vary_by_env_id( static void test_depth_solve_rates_are_conditional_on_depth_attempts(void) { Log log = {0}; - log.depth_6_rate = 0.25f; - log.depth_6_solve_rate = 0.125f; - log.depth_8_rate = 0.0f; - log.depth_8_solve_rate = 0.0f; - log.depth_16_rate = 0.125f; - log.depth_16_solve_rate = 0.0f; + log.d6_rate = 0.25f; + log.d6_solve_rate = 0.125f; + log.d8_rate = 0.0f; + log.d8_solve_rate = 0.0f; + log.d16_rate = 0.125f; + log.d16_solve_rate = 0.0f; log.score = 0.75f; log.target_distance = 4.0f; log.solved_target_distance = 2.0f; @@ -230,19 +230,19 @@ static void test_depth_solve_rates_are_conditional_on_depth_attempts(void) { Dict out = {0}; puf_log(&log, &out); - EXPECT_EQ_INT(out.size, 16); + EXPECT_EQ_INT(out.size, 15); EXPECT_NEAR(dict_value(&out, "score"), 0.75, 0.0); EXPECT_TRUE(!dict_has_key(&out, "solve_steps")); EXPECT_TRUE(!dict_has_key(&out, "solve_efficiency")); EXPECT_TRUE(!dict_has_key(&out, "scramble_unique_states")); EXPECT_NEAR(dict_value(&out, "min_win_moves"), 4.0, 0.0); EXPECT_NEAR(dict_value(&out, "solved_min_win_moves"), 4.0, 0.0); - EXPECT_TRUE(!dict_has_key(&out, "depth_6_rate")); - EXPECT_NEAR(dict_value(&out, "depth_6_solve_rate"), 0.5, 0.0); - EXPECT_TRUE(!dict_has_key(&out, "depth_8_rate")); - EXPECT_NEAR(dict_value(&out, "depth_8_solve_rate"), 0.0, 0.0); - EXPECT_TRUE(!dict_has_key(&out, "depth_16_rate")); - EXPECT_NEAR(dict_value(&out, "depth_16_solve_rate"), 0.0, 0.0); + EXPECT_TRUE(!dict_has_key(&out, "d6_rate")); + EXPECT_NEAR(dict_value(&out, "d6_solve_rate"), 0.5, 0.0); + EXPECT_TRUE(!dict_has_key(&out, "d8_rate")); + EXPECT_NEAR(dict_value(&out, "d8_solve_rate"), 0.0, 0.0); + EXPECT_TRUE(!dict_has_key(&out, "d16_rate")); + EXPECT_NEAR(dict_value(&out, "d16_solve_rate"), 0.0, 0.0); dict_clear(&out); } diff --git a/ocean/affine_lock/tests/test_metadata_smoke.py b/ocean/affine_lock/tests/test_metadata_smoke.py index 477012ecfc..5d3c5690bb 100644 --- a/ocean/affine_lock/tests/test_metadata_smoke.py +++ b/ocean/affine_lock/tests/test_metadata_smoke.py @@ -17,14 +17,13 @@ "episode_return", "episode_length", "timeout_rate", - "invalid_rate", "min_win_moves", "solved_min_win_moves", "conditional_solve_steps", "conditional_solve_efficiency", - "depth_6_solve_rate", - "depth_8_solve_rate", - "depth_16_solve_rate", + "d6_solve_rate", + "d8_solve_rate", + "d16_solve_rate", "n", ] @@ -229,8 +228,7 @@ def load_affine_args(extra_argv): logs = vec.log() assert logs["n"] == 2.0 - assert logs["invalid_rate"] == 1.0 - assert logs["timeout_rate"] == 0.0 + assert logs["timeout_rate"] == 1.0 assert logs["solve_rate"] == 0.0 assert logs["episode_length"] == 1.0 assert logs["episode_return"] == -1.0 diff --git a/ocean/affine_lock/tests/test_visible_targets_loader.c b/ocean/affine_lock/tests/test_visible_targets_loader.c index fd56accc16..9e3ba91e9a 100644 --- a/ocean/affine_lock/tests/test_visible_targets_loader.c +++ b/ocean/affine_lock/tests/test_visible_targets_loader.c @@ -53,7 +53,7 @@ int main(int argc, char** argv) { unsigned long expected_d16_count = strtoul(argv[4], &end, 10); EXPECT_TRUE(end != argv[4] && *end == '\0'); - AffineLockVisibleTargetTable table; + VisibleTargetTable table; int rc = visible_targets_load( argv[1], VISIBLE_TARGET_8ACTION_V1_HASH, @@ -93,7 +93,7 @@ int main(int argc, char** argv) { } for (uint32_t i = 0; i < table.record_count; i++) { - const AffineLockVisibleTargetRecord* record = &table.records[i]; + const VisibleTargetRecord* record = &table.records[i]; EXPECT_TRUE(record->solution_length == record->depth); EXPECT_TRUE( record->depth == 2 || From 7f673ec7d5d1909b71aee6fd929c5fb636d83f6b Mon Sep 17 00:00:00 2001 From: Kinvert Date: Mon, 17 Aug 2026 04:53:13 -0400 Subject: [PATCH 3/4] Add quadratic perf_weighting option to affine_lock --- config/affine_lock.ini | 33 +++++++------------ ocean/affine_lock/affine_lock.c | 1 + ocean/affine_lock/affine_lock.h | 28 ++++++++++++---- ocean/affine_lock/tests/test_affine_lock.c | 23 ++++++++++++- .../tests/test_affine_lock_log_export.c | 1 + .../affine_lock/tests/test_metadata_smoke.py | 12 +------ 6 files changed, 59 insertions(+), 39 deletions(-) diff --git a/config/affine_lock.ini b/config/affine_lock.ini index 72b741c112..ad8af2bbe7 100644 --- a/config/affine_lock.ini +++ b/config/affine_lock.ini @@ -15,6 +15,7 @@ seed = 42 start_depth = 2 max_depth = 16 step_grace = 0 +perf_weighting = 1 [train] total_timesteps = 200_000_000 @@ -29,22 +30,19 @@ clip_coef = 0.83 vf_coef = 4.75 vf_clip_coef = 0.8 max_grad_norm = 3.0 -beta1 = 0.5 -beta2 = 0.9915 -eps = 0.0001 vtrace_rho_clip = 1.4 vtrace_c_clip = 3.75 [sweep] metric = perf goal = maximize -max_runs = 50 +max_runs = 1000 use_gpu = False [sweep.train.total_timesteps] distribution = log_normal min = 100_000_000 -max = 200_000_000 +max = 2_000_000_000 mean = 200_000_000 scale = time @@ -70,7 +68,7 @@ mean = 3 scale = auto [sweep.vec.num_buffers] -distribution = uniform +distribution = uniform_pow2 min = 1 max = 4 mean = 2 @@ -101,7 +99,11 @@ scale = auto mean = 0.012 [sweep.train.ent_coef] +distribution = log_normal +min = 0.01 +max = 0.4 mean = 0.2 +scale = auto [sweep.train.gamma] mean = 0.8 @@ -109,12 +111,6 @@ mean = 0.8 [sweep.train.gae_lambda] mean = 0.995 -[sweep.train.vtrace_rho_clip] -mean = 1.4 - -[sweep.train.vtrace_c_clip] -mean = 3.75 - [sweep.train.clip_coef] mean = 0.83 @@ -133,13 +129,8 @@ mean = 4.75 scale = auto [sweep.train.max_grad_norm] +distribution = uniform +min = 0.5 +max = 6.0 mean = 3.0 - -[sweep.train.beta1] -mean = 0.5 - -[sweep.train.beta2] -mean = 0.9915 - -[sweep.train.eps] -mean = 0.0001 +scale = auto diff --git a/ocean/affine_lock/affine_lock.c b/ocean/affine_lock/affine_lock.c index 21229ee3e3..8d2ae8cb81 100644 --- a/ocean/affine_lock/affine_lock.c +++ b/ocean/affine_lock/affine_lock.c @@ -15,6 +15,7 @@ int main(void) { dict_set(&kwargs, "start_depth", 2); dict_set(&kwargs, "max_depth", 16); dict_set(&kwargs, "step_grace", 2); + dict_set(&kwargs, "perf_weighting", PERF_WEIGHTING_QUADRATIC); dict_set(&kwargs, "seed", (double)(unsigned int)time(NULL)); puf_init(&env, &kwargs); dict_clear(&kwargs); diff --git a/ocean/affine_lock/affine_lock.h b/ocean/affine_lock/affine_lock.h index c8c97beccd..aef57f422e 100644 --- a/ocean/affine_lock/affine_lock.h +++ b/ocean/affine_lock/affine_lock.h @@ -19,6 +19,9 @@ typedef float obs_t; #define ACT_SIZES {NUM_ACTIONS} #define PUF_STEPS_PER_SEC 2 +#define PERF_WEIGHTING_LINEAR 0 +#define PERF_WEIGHTING_QUADRATIC 1 + #define MY_VEC_INIT #define MY_VEC_CLOSE @@ -61,6 +64,7 @@ typedef struct AffineLockShared { int start_depth; int max_depth; int step_grace; + int perf_weighting; uint32_t mask; uint32_t* next; VisibleTargetTable visible_target_table; @@ -89,10 +93,12 @@ struct Env { }; typedef Env AffineLock; -static void init_shared(AffineLockShared* shared, int start_depth, int max_depth, int step_grace) { +static void init_shared(AffineLockShared* shared, int start_depth, int max_depth, + int step_grace, int perf_weighting) { shared->start_depth = start_depth; shared->max_depth = max_depth; shared->step_grace = step_grace; + shared->perf_weighting = perf_weighting; shared->mask = (1u << BITS) - 1u; for (int value = 0; value < 256; value++) { for (int bit = 0; bit < 8; bit++) { @@ -147,9 +153,10 @@ static void init_shared(AffineLockShared* shared, int start_depth, int max_depth "Table' in ocean/affine_lock/README.md"); } -static AffineLockShared* create_shared(int start_depth, int max_depth, int step_grace) { +static AffineLockShared* create_shared(int start_depth, int max_depth, + int step_grace, int perf_weighting) { AffineLockShared* shared = (AffineLockShared*)calloc(1, sizeof(AffineLockShared)); - init_shared(shared, start_depth, max_depth, step_grace); + init_shared(shared, start_depth, max_depth, step_grace, perf_weighting); return shared; } @@ -164,8 +171,10 @@ void puf_init(Env* env, Dict* kwargs) { int start_depth = dict_get(kwargs, "start_depth"); int max_depth = dict_get(kwargs, "max_depth"); int step_grace = dict_get(kwargs, "step_grace"); + int perf_weighting = dict_get(kwargs, "perf_weighting"); unsigned int seed = (unsigned int)dict_get(kwargs, "seed"); - AffineLockShared* shared = create_shared(start_depth, max_depth, step_grace); + AffineLockShared* shared = + create_shared(start_depth, max_depth, step_grace, perf_weighting); init_env(env, shared, rand_r(&seed)); env->owns_shared = 1; } @@ -189,7 +198,12 @@ static void add_log(AffineLock* env, int solved) { AffineLockShared* shared = env->shared; int log_depth = env->target_distance; int at_max_depth = log_depth == shared->max_depth; - float solve_credit = solved ? log_depth / (float)shared->max_depth : 0; + float ratio = log_depth / (float)shared->max_depth; + float solve_credit = 0; + if (solved) { + solve_credit = shared->perf_weighting == PERF_WEIGHTING_QUADRATIC ? + ratio * ratio : ratio; + } env->log.perf += solve_credit; env->log.score += solve_credit; env->log.solve_rate += solved; @@ -388,8 +402,10 @@ Env* my_vec_init(int* num_envs_out, int* buffer_env_starts, int* buffer_env_coun int start_depth = dict_get(env_kwargs, "start_depth"); int max_depth = dict_get(env_kwargs, "max_depth"); int step_grace = dict_get(env_kwargs, "step_grace"); + int perf_weighting = dict_get(env_kwargs, "perf_weighting"); - AffineLockShared* shared = create_shared(start_depth, max_depth, step_grace); + AffineLockShared* shared = + create_shared(start_depth, max_depth, step_grace, perf_weighting); Env* envs = (Env*)calloc(total_agents, sizeof(Env)); int buf = 0; diff --git a/ocean/affine_lock/tests/test_affine_lock.c b/ocean/affine_lock/tests/test_affine_lock.c index cad393dad4..5e3c3aa045 100644 --- a/ocean/affine_lock/tests/test_affine_lock.c +++ b/ocean/affine_lock/tests/test_affine_lock.c @@ -69,7 +69,7 @@ static AffineLockShared make_shared( int step_grace) { AffineLockShared shared; memset(&shared, 0, sizeof(shared)); - init_shared(&shared, start_depth, max_depth, step_grace); + init_shared(&shared, start_depth, max_depth, step_grace, PERF_WEIGHTING_LINEAR); return shared; } @@ -222,6 +222,26 @@ static void test_log_solve_credit_uses_known_target_distance(void) { free_shared(&shared); } +static void test_log_solve_credit_uses_quadratic_perf_weighting(void) { + AffineLockShared shared; + memset(&shared, 0, sizeof(shared)); + init_shared(&shared, 2, 16, 0, PERF_WEIGHTING_QUADRATIC); + AffineLock env; + memset(&env, 0, sizeof(env)); + env.shared = &shared; + env.scramble_depth = 16; + env.target_distance = 8; + env.step_count = 8; + + add_log(&env, 1); + + float linear_ratio = expected_solve_credit(&shared, 8); + EXPECT_NEAR(env.log.perf, linear_ratio * linear_ratio, 0.0f); + EXPECT_NEAR(env.log.score, linear_ratio * linear_ratio, 0.0f); + + free_shared(&shared); +} + static void expect_observation_matches(const AffineLock* env) { float* obs = env->agents[0].observations; for (int bit = 0; bit < BITS; bit++) { @@ -1196,6 +1216,7 @@ int main(void) { test_visible_target_table_reset_uses_stored_records(); test_visible_target_table_matches_independent_bfs_over_repeated_resets(); test_log_solve_credit_uses_known_target_distance(); + test_log_solve_credit_uses_quadratic_perf_weighting(); test_observation_encoding_is_32_signed_bit_floats_plus_timer(); test_timer_observation_progresses_and_resets_after_timeout(); test_actions_apply_to_current_state_directly(); diff --git a/ocean/affine_lock/tests/test_affine_lock_log_export.c b/ocean/affine_lock/tests/test_affine_lock_log_export.c index 7f3bd09312..9923acb242 100644 --- a/ocean/affine_lock/tests/test_affine_lock_log_export.c +++ b/ocean/affine_lock/tests/test_affine_lock_log_export.c @@ -73,6 +73,7 @@ static void fill_env_kwargs(Dict* env_kwargs, int seed) { dict_set(env_kwargs, "start_depth", 2); dict_set(env_kwargs, "max_depth", 16); dict_set(env_kwargs, "step_grace", 0); + dict_set(env_kwargs, "perf_weighting", PERF_WEIGHTING_LINEAR); dict_set(env_kwargs, "seed", seed); } diff --git a/ocean/affine_lock/tests/test_metadata_smoke.py b/ocean/affine_lock/tests/test_metadata_smoke.py index 5d3c5690bb..436a755f60 100644 --- a/ocean/affine_lock/tests/test_metadata_smoke.py +++ b/ocean/affine_lock/tests/test_metadata_smoke.py @@ -54,6 +54,7 @@ def check_config(): assert parse_int(config["env"]["seed"]) == 42 assert parse_int(config["env"]["start_depth"]) == 2 assert parse_int(config["env"]["max_depth"]) == 16 + assert parse_int(config["env"]["perf_weighting"]) == 1 assert parse_int(config["train"]["total_timesteps"]) == 200_000_000 assert parse_int(config["train"]["horizon"]) == 64 assert parse_int(config["train"]["minibatch_size"]) == 8192 @@ -66,9 +67,6 @@ def check_config(): assert parse_float(config["train"]["vf_coef"]) == 4.75 assert parse_float(config["train"]["vf_clip_coef"]) == 0.8 assert parse_float(config["train"]["max_grad_norm"]) == 3.0 - assert parse_float(config["train"]["beta1"]) == 0.5 - assert parse_float(config["train"]["beta2"]) == 0.9915 - assert parse_float(config["train"]["eps"]) == 0.0001 assert parse_float(config["train"]["vtrace_rho_clip"]) == 1.4 assert parse_float(config["train"]["vtrace_c_clip"]) == 3.75 assert "prio_alpha" not in config["train"] @@ -89,11 +87,6 @@ def check_config(): assert_sweep_mean(config, "sweep.train.vf_coef", 4.75) assert_sweep_mean(config, "sweep.train.vf_clip_coef", 0.8) assert_sweep_mean(config, "sweep.train.max_grad_norm", 3.0) - assert_sweep_mean(config, "sweep.train.beta1", 0.5) - assert_sweep_mean(config, "sweep.train.beta2", 0.9915) - assert_sweep_mean(config, "sweep.train.eps", 0.0001) - assert_sweep_mean(config, "sweep.train.vtrace_rho_clip", 1.4) - assert_sweep_mean(config, "sweep.train.vtrace_c_clip", 3.75) assert "sweep.train.prio_alpha" not in config assert "sweep.train.prio_beta0" not in config assert config["sweep"]["metric"] == "perf" @@ -188,9 +181,6 @@ def load_affine_args(extra_argv): assert base_args["train"]["vf_coef"] == 4.75 assert base_args["train"]["vf_clip_coef"] == 0.8 assert base_args["train"]["max_grad_norm"] == 3.0 - assert base_args["train"]["beta1"] == 0.5 - assert base_args["train"]["beta2"] == 0.9915 - assert base_args["train"]["eps"] == 0.0001 assert base_args["train"]["vtrace_rho_clip"] == 1.4 assert base_args["train"]["vtrace_c_clip"] == 3.75 assert "prio_alpha" not in base_args["train"] From f16987c68ad6d6a18ca487f61af103fd7c40d398 Mon Sep 17 00:00:00 2001 From: Kinvert Date: Wed, 19 Aug 2026 17:36:55 -0700 Subject: [PATCH 4/4] Add CUDA backend and parity tests for Affine Lock --- ocean/affine_lock/affine_lock.cu | 808 ++++++++++++ ocean/affine_lock/tests/run_cuda.sh | 30 + .../tests/test_affine_lock_cuda.cu | 1109 +++++++++++++++++ 3 files changed, 1947 insertions(+) create mode 100644 ocean/affine_lock/affine_lock.cu create mode 100755 ocean/affine_lock/tests/run_cuda.sh create mode 100644 ocean/affine_lock/tests/test_affine_lock_cuda.cu diff --git a/ocean/affine_lock/affine_lock.cu b/ocean/affine_lock/affine_lock.cu new file mode 100644 index 0000000000..adb550eb5f --- /dev/null +++ b/ocean/affine_lock/affine_lock.cu @@ -0,0 +1,808 @@ +// Vibe coded by OpenAI Codex. +// GPU Affine Lock environment. This is intentionally standalone from +// affine_lock.h: --gpu builds include this file instead of the CPU source. +#ifndef PUFFER_AFFINE_LOCK_GPU_CU +#define PUFFER_AFFINE_LOCK_GPU_CU + +#define PUF_BACKEND PUF_GPU + +#include +#include + +#include +#include +#include +#include +#include + +// Environment observations are fixed bf16. All bit observations are +/-1 +// (exact in bf16); only the timer is rounded. Keeping them bf16 halves the +// rollout bandwidth compared with the CPU float representation. +typedef __nv_bfloat16 obs_t; +#include "pufferenv.h" +#include "affine_lock_visible_targets.h" + +#define BITS 16 +#define TIMER_INDEX (2 * BITS) +#define OBS_SIZE (TIMER_INDEX + 1) +#define NUM_ATNS 1 +#define NUM_ACTIONS 8 +#define MAX_SOLUTION_DEPTH 16 +#define CURRICULUM_DEPTH_COUNT 6 +#define STEP_REWARD (-0.01f) +#ifndef VISIBLE_TARGET_TABLE_PATH +#define VISIBLE_TARGET_TABLE_PATH "ocean/affine_lock/generated/affine_lock_8action_visible_targets.bin" +#endif +#define ACT_SIZES {NUM_ACTIONS} +#define PUF_STEPS_PER_SEC 2 + +#define PERF_WEIGHTING_LINEAR 0 +#define PERF_WEIGHTING_QUADRATIC 1 + +#ifndef AFFINE_LOCK_GPU_SHARED_OBS +#define AFFINE_LOCK_GPU_SHARED_OBS 1 +#endif +#ifndef AFFINE_LOCK_GPU_SHARED_BLOCK +#define AFFINE_LOCK_GPU_SHARED_BLOCK 128 +#endif +#define AFFINE_LOCK_GPU_DEPTH_LUT_SIZE (MAX_SOLUTION_DEPTH + 1) + +static_assert(AFFINE_LOCK_GPU_SHARED_BLOCK >= 32 && + AFFINE_LOCK_GPU_SHARED_BLOCK <= 256 && + AFFINE_LOCK_GPU_SHARED_BLOCK % 32 == 0, + "AFFINE_LOCK_GPU_SHARED_BLOCK must contain whole warps"); + +#if !AFFINE_LOCK_GPU_SHARED_OBS +#ifndef AFFINE_LOCK_GPU_LANES +#define AFFINE_LOCK_GPU_LANES 4 +#endif +#define AFFINE_LOCK_GPU_BLOCK 256 +static_assert(AFFINE_LOCK_GPU_LANES == 4 || AFFINE_LOCK_GPU_LANES == 8 || + AFFINE_LOCK_GPU_LANES == 16 || AFFINE_LOCK_GPU_LANES == 32, + "AFFINE_LOCK_GPU_LANES must be a power-of-two subwarp"); +static_assert(AFFINE_LOCK_GPU_BLOCK % AFFINE_LOCK_GPU_LANES == 0, + "block size must contain whole environments"); +#endif + +struct Log { + float perf; + float score; + float solve_rate; + float max_depth_solve; + float episode_return; + float episode_length; + float solve_steps; + float timeout_rate; + float solve_efficiency; + float target_distance; + float solved_target_distance; + float d6_rate; + float d6_solve_rate; + float d8_rate; + float d8_solve_rate; + float d16_rate; + float d16_solve_rate; + float n; +}; + +static_assert(sizeof(Log) == 18 * sizeof(float), + "trainer log reduction requires a packed float-only Log"); + +// The trainer only reads Env::log for a GPU backend. Runtime state is kept in +// a separate compact array so log scans do not pull state into cache and state +// updates do not stride over the relatively large log payload. +struct Env { + Log log; + Agent agents[1]; + int num_agents; + int tag; + int boundary_reached; + unsigned int rng; +}; + +// Exactly 32 bytes: four adjacent environment records fit in one 128-byte +// transaction. The default shared-observation kernel reads one record per env. +typedef struct GpuAffineLockState { + uint32_t rng; + uint16_t state; + uint16_t target; + int step_count; + int max_steps; + int scramble_depth; + int curriculum_depth; + int target_distance; + float episode_return; +} GpuAffineLockState; + +static_assert(sizeof(GpuAffineLockState) == 32, + "GpuAffineLockState layout is performance-sensitive"); + +typedef struct GpuAffineLockConfig { + int start_depth; + int max_depth; + int step_grace; + int perf_weighting; + uint32_t depth_first[AFFINE_LOCK_GPU_DEPTH_LUT_SIZE]; + uint32_t depth_counts[AFFINE_LOCK_GPU_DEPTH_LUT_SIZE]; +} GpuAffineLockConfig; + +__constant__ GpuAffineLockConfig d_affine_lock_config; + +static struct { + Env* envs; + GpuAffineLockState* states; + uint32_t* target_pairs; + int n; + obs_t* observations; + float* actions; + float* rewards; + float* terminals; + cudaStream_t stream; + GpuAffineLockConfig config; +} g_gpu; + +static void gpu_affine_lock_check(cudaError_t status, const char* operation) { + if (status != cudaSuccess) { + std::fprintf(stderr, "Affine Lock CUDA: %s failed: %s\n", + operation, cudaGetErrorString(status)); + std::exit(1); + } +} + +#if !AFFINE_LOCK_GPU_SHARED_OBS +static int gpu_affine_lock_grid(int threads) { + return (threads + AFFINE_LOCK_GPU_BLOCK - 1) / AFFINE_LOCK_GPU_BLOCK; +} +#endif + +__device__ __forceinline__ uint32_t gpu_affine_lock_random_mixed_u32( + GpuAffineLockState* env) { + env->rng = env->rng * 1664525u + 1013904223u; + uint32_t x = env->rng; + x ^= x >> 16; + x *= 0x7feb352du; + x ^= x >> 15; + x *= 0x846ca68bu; + x ^= x >> 16; + return x; +} + +__device__ __forceinline__ int gpu_affine_lock_random_bounded( + GpuAffineLockState* env, int bound) { + uint32_t ubound = (uint32_t)bound; + uint32_t limit = UINT32_MAX - UINT32_MAX % ubound; + uint32_t value = gpu_affine_lock_random_mixed_u32(env); + while (value >= limit) { + value = gpu_affine_lock_random_mixed_u32(env); + } + return (int)(value % ubound); +} + +__device__ __forceinline__ void gpu_affine_lock_reset_state( + GpuAffineLockState* env, const uint32_t* target_pairs) { + env->scramble_depth = env->curriculum_depth; + env->step_count = 0; + env->episode_return = 0.0f; + int depth = env->scramble_depth; + uint32_t count = d_affine_lock_config.depth_counts[depth]; + int choice = gpu_affine_lock_random_bounded(env, (int)count); + uint32_t record_index = d_affine_lock_config.depth_first[depth] + + (uint32_t)choice; + uint32_t pair = target_pairs[record_index]; + env->state = (uint16_t)(pair & 0xffffu); + env->target = (uint16_t)(pair >> 16); + env->target_distance = env->scramble_depth; + env->max_steps = env->target_distance + d_affine_lock_config.step_grace; +} + +__device__ __forceinline__ uint16_t gpu_affine_lock_apply_action( + uint16_t state, int action) { + uint32_t next = state; + switch (action) { + case 0: + next = (state >> 1) | ((state & 1u) << 15); + break; + case 1: + next = ((state << 1) & 0xffffu) | ((state >> 15) & 1u); + break; + case 2: + next = state ^ 0xfe00u; + break; + case 3: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + break; + case 4: + next = ((state & 0x3333u) << 2) | ((state & 0xccccu) >> 2); + break; + case 5: + next = ((state & 0x0f0fu) << 4) | ((state & 0xf0f0u) >> 4); + break; + case 6: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + break; + case 7: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + next = ((next & 0x0f0fu) << 4) | ((next & 0xf0f0u) >> 4); + break; + } + return (uint16_t)(next & 0xffffu); +} + +__device__ __forceinline__ int gpu_affine_lock_next_curriculum_depth( + int current_depth) { + constexpr int curriculum_depths[CURRICULUM_DEPTH_COUNT] = {2, 4, 5, 6, 8, 16}; +#pragma unroll + for (int i = 0; i < CURRICULUM_DEPTH_COUNT; i++) { + int depth = curriculum_depths[i]; + if (depth > current_depth) { + return depth < d_affine_lock_config.max_depth + ? depth : d_affine_lock_config.max_depth; + } + } + return d_affine_lock_config.max_depth; +} + +__device__ __forceinline__ void gpu_affine_lock_add_log( + Env* trainer_env, const GpuAffineLockState* env, int solved) { + int log_depth = env->target_distance; + int at_max_depth = log_depth == d_affine_lock_config.max_depth; + float ratio = log_depth / (float)d_affine_lock_config.max_depth; + float solve_credit = 0.0f; + if (solved) { + solve_credit = d_affine_lock_config.perf_weighting == PERF_WEIGHTING_QUADRATIC + ? ratio * ratio : ratio; + } + Log* log = &trainer_env->log; + log->perf += solve_credit; + log->score += solve_credit; + log->solve_rate += solved; + log->max_depth_solve += solved && at_max_depth; + log->episode_return += env->episode_return; + log->episode_length += env->step_count; + log->solve_steps += solved ? env->step_count : 0; + log->timeout_rate += !solved; + log->solve_efficiency += solved + ? env->step_count / (float)log_depth : 0.0f; + log->target_distance += env->target_distance; + log->solved_target_distance += solved ? env->target_distance : 0; + log->d6_rate += log_depth == 6; + log->d6_solve_rate += solved && log_depth == 6; + log->d8_rate += log_depth == 8; + log->d8_solve_rate += solved && log_depth == 8; + log->d16_rate += log_depth == 16; + log->d16_solve_rate += solved && log_depth == 16; + log->n += 1; +} + +__device__ __forceinline__ uint32_t gpu_affine_lock_step_one( + Env* trainer_env, GpuAffineLockState* env, + const uint32_t* target_pairs, float action, + float* reward_out, float* terminal_out, float* timer_out) { + float reward = STEP_REWARD; + float terminal = 0.0f; + int solved = 0; + env->step_count += 1; + int invalid = !isfinite(action) || action < 0.0f || action > NUM_ACTIONS - 1; + if (invalid) { + reward = -1.0f; + terminal = 1.0f; + } else { + env->state = gpu_affine_lock_apply_action(env->state, (int)action); + if (env->state == env->target) { + reward = 1.0f; + terminal = 1.0f; + solved = 1; + } else if (env->step_count >= env->max_steps) { + reward = -1.0f; + terminal = 1.0f; + } + } + env->episode_return += reward; + if (terminal != 0.0f) { + gpu_affine_lock_add_log(trainer_env, env, solved); + env->curriculum_depth = solved + ? gpu_affine_lock_next_curriculum_depth(env->scramble_depth) + : d_affine_lock_config.start_depth; + gpu_affine_lock_reset_state(env, target_pairs); + } + *reward_out = reward; + *terminal_out = terminal; + *timer_out = env->step_count / (float)env->max_steps; + return (uint32_t)env->state | ((uint32_t)env->target << 16); +} + +#if !AFFINE_LOCK_GPU_SHARED_OBS +__device__ __forceinline__ void gpu_affine_lock_write_observations( + obs_t* observations, uint32_t packed_bits, float timer, int lane) { +#pragma unroll + for (int bit = lane; bit < 2 * BITS; bit += AFFINE_LOCK_GPU_LANES) { + observations[bit] = __float2bfloat16( + (packed_bits & (1u << bit)) ? 1.0f : -1.0f); + } + if (lane == 0) { + observations[TIMER_INDEX] = __float2bfloat16(timer); + } +} + +__global__ __launch_bounds__(AFFINE_LOCK_GPU_BLOCK) +void gpu_affine_lock_reset_kernel(Env* envs, GpuAffineLockState* states, + const uint32_t* target_pairs, obs_t* observations, + float* rewards, float* terminals, int num_envs) { + int thread = blockIdx.x * blockDim.x + threadIdx.x; + int relative_env = thread / AFFINE_LOCK_GPU_LANES; + int lane = thread & (AFFINE_LOCK_GPU_LANES - 1); + int active = relative_env < num_envs; + uint32_t packed_bits = 0; + float timer = 0.0f; + if (active && lane == 0) { + GpuAffineLockState* env = &states[relative_env]; + gpu_affine_lock_reset_state(env, target_pairs); + rewards[relative_env] = 0.0f; + terminals[relative_env] = 0.0f; + packed_bits = (uint32_t)env->state | ((uint32_t)env->target << 16); + } + int leader = (threadIdx.x & 31) & ~(AFFINE_LOCK_GPU_LANES - 1); + packed_bits = __shfl_sync(0xffffffffu, packed_bits, leader); + timer = __shfl_sync(0xffffffffu, timer, leader); + if (active) { + gpu_affine_lock_write_observations( + observations + (size_t)relative_env * OBS_SIZE, + packed_bits, timer, lane); + } + (void)envs; +} + +__global__ __launch_bounds__(AFFINE_LOCK_GPU_BLOCK) +void gpu_affine_lock_step_kernel(Env* envs, GpuAffineLockState* states, + const uint32_t* target_pairs, const float* actions, + obs_t* observations, float* rewards, float* terminals, + int num_envs) { + int thread = blockIdx.x * blockDim.x + threadIdx.x; + int relative_env = thread / AFFINE_LOCK_GPU_LANES; + int lane = thread & (AFFINE_LOCK_GPU_LANES - 1); + int active = relative_env < num_envs; + uint32_t packed_bits = 0; + float timer = 0.0f; + if (active && lane == 0) { + packed_bits = gpu_affine_lock_step_one( + &envs[relative_env], &states[relative_env], target_pairs, + actions[(size_t)relative_env * NUM_ATNS], + &rewards[relative_env], &terminals[relative_env], &timer); + } + int leader = (threadIdx.x & 31) & ~(AFFINE_LOCK_GPU_LANES - 1); + packed_bits = __shfl_sync(0xffffffffu, packed_bits, leader); + timer = __shfl_sync(0xffffffffu, timer, leader); + if (active) { + gpu_affine_lock_write_observations( + observations + (size_t)relative_env * OBS_SIZE, + packed_bits, timer, lane); + } +} +#endif + +#if AFFINE_LOCK_GPU_SHARED_OBS +// One simulation thread per environment writes +// a conflict-free 33-float shared-memory row, then the whole block converts and +// stores a linear bf16 tile with fully coalesced global writes. +__global__ __launch_bounds__(AFFINE_LOCK_GPU_SHARED_BLOCK) +void gpu_affine_lock_shared_reset_kernel(Env* envs, + GpuAffineLockState* states, const uint32_t* target_pairs, + obs_t* observations, float* rewards, float* terminals, int num_envs) { + __shared__ float observation_tile[AFFINE_LOCK_GPU_SHARED_BLOCK * OBS_SIZE]; + int block_start = blockIdx.x * AFFINE_LOCK_GPU_SHARED_BLOCK; + int relative_env = block_start + threadIdx.x; + int active_count = num_envs - block_start; + if (active_count > AFFINE_LOCK_GPU_SHARED_BLOCK) { + active_count = AFFINE_LOCK_GPU_SHARED_BLOCK; + } + if (active_count < 0) { + active_count = 0; + } + if (threadIdx.x < active_count) { + GpuAffineLockState* env = &states[relative_env]; + gpu_affine_lock_reset_state(env, target_pairs); + rewards[relative_env] = 0.0f; + terminals[relative_env] = 0.0f; + uint32_t bits = (uint32_t)env->state | ((uint32_t)env->target << 16); + float* row = observation_tile + threadIdx.x * OBS_SIZE; +#pragma unroll + for (int bit = 0; bit < 2 * BITS; bit++) { + row[bit] = (bits & (1u << bit)) ? 1.0f : -1.0f; + } + row[TIMER_INDEX] = 0.0f; + } + __syncthreads(); + int tile_values = active_count * OBS_SIZE; + for (int value = threadIdx.x; value < tile_values; value += blockDim.x) { + observations[(size_t)block_start * OBS_SIZE + value] = + __float2bfloat16(observation_tile[value]); + } + (void)envs; +} + +__global__ __launch_bounds__(AFFINE_LOCK_GPU_SHARED_BLOCK) +void gpu_affine_lock_shared_step_kernel(Env* envs, + GpuAffineLockState* states, const uint32_t* target_pairs, + const float* actions, obs_t* observations, + float* rewards, float* terminals, int num_envs) { + __shared__ float observation_tile[AFFINE_LOCK_GPU_SHARED_BLOCK * OBS_SIZE]; + int block_start = blockIdx.x * AFFINE_LOCK_GPU_SHARED_BLOCK; + int relative_env = block_start + threadIdx.x; + int active_count = num_envs - block_start; + if (active_count > AFFINE_LOCK_GPU_SHARED_BLOCK) { + active_count = AFFINE_LOCK_GPU_SHARED_BLOCK; + } + if (active_count < 0) { + active_count = 0; + } + if (threadIdx.x < active_count) { + float timer = 0.0f; + uint32_t bits = gpu_affine_lock_step_one( + &envs[relative_env], &states[relative_env], target_pairs, + actions[(size_t)relative_env * NUM_ATNS], + &rewards[relative_env], &terminals[relative_env], &timer); + float* row = observation_tile + threadIdx.x * OBS_SIZE; +#pragma unroll + for (int bit = 0; bit < 2 * BITS; bit++) { + row[bit] = (bits & (1u << bit)) ? 1.0f : -1.0f; + } + row[TIMER_INDEX] = timer; + } + __syncthreads(); + int tile_values = active_count * OBS_SIZE; + for (int value = threadIdx.x; value < tile_values; value += blockDim.x) { + observations[(size_t)block_start * OBS_SIZE + value] = + __float2bfloat16(observation_tile[value]); + } +} +#endif + +void puf_log(Log* log, Dict* out) { + float nsolve = log->solve_rate; + float solved_min_win_moves = nsolve + ? log->solved_target_distance / nsolve : 0; + float conditional_solve_steps = nsolve ? log->solve_steps / nsolve : 0; + float conditional_solve_efficiency = nsolve + ? log->solve_efficiency / nsolve : 0; + + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "solve_rate", log->solve_rate); + dict_set(out, "max_depth_solve", log->max_depth_solve); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); + dict_set(out, "timeout_rate", log->timeout_rate); + dict_set(out, "min_win_moves", log->target_distance); + dict_set(out, "solved_min_win_moves", solved_min_win_moves); + dict_set(out, "conditional_solve_steps", conditional_solve_steps); + dict_set(out, "conditional_solve_efficiency", conditional_solve_efficiency); + dict_set(out, "d6_solve_rate", log->d6_rate + ? log->d6_solve_rate / log->d6_rate : 0); + dict_set(out, "d8_solve_rate", log->d8_rate + ? log->d8_solve_rate / log->d8_rate : 0); + dict_set(out, "d16_solve_rate", log->d16_rate + ? log->d16_solve_rate / log->d16_rate : 0); + dict_set(out, "n", log->n); +} + +static int gpu_affine_lock_host_has_depth( + const GpuAffineLockConfig* config, int depth) { + return depth >= 0 && depth < AFFINE_LOCK_GPU_DEPTH_LUT_SIZE + && config->depth_counts[depth] != 0; +} + +static int gpu_affine_lock_host_next_depth(int current_depth, int max_depth) { + static const int curriculum_depths[CURRICULUM_DEPTH_COUNT] = {2, 4, 5, 6, 8, 16}; + for (int i = 0; i < CURRICULUM_DEPTH_COUNT; i++) { + int depth = curriculum_depths[i]; + if (depth > current_depth) { + return depth < max_depth ? depth : max_depth; + } + } + return max_depth; +} + +static void gpu_affine_lock_validate_curriculum( + const GpuAffineLockConfig* config) { + if (config->start_depth <= 0 || + config->max_depth < config->start_depth || + config->max_depth > MAX_SOLUTION_DEPTH) { + std::fprintf(stderr, + "Affine Lock CUDA: invalid curriculum range start=%d max=%d\n", + config->start_depth, config->max_depth); + std::exit(1); + } + int depth = config->start_depth; + for (int i = 0; i <= CURRICULUM_DEPTH_COUNT; i++) { + if (!gpu_affine_lock_host_has_depth(config, depth)) { + std::fprintf(stderr, + "Affine Lock CUDA: target table has no depth %d section\n", depth); + std::exit(1); + } + if (depth + config->step_grace <= 0) { + std::fprintf(stderr, + "Affine Lock CUDA: depth %d with step_grace=%d has no valid steps\n", + depth, config->step_grace); + std::exit(1); + } + if (depth == config->max_depth) { + return; + } + int next = gpu_affine_lock_host_next_depth(depth, config->max_depth); + if (next == depth) { + break; + } + depth = next; + } + std::fprintf(stderr, "Affine Lock CUDA: curriculum does not reach max depth %d\n", + config->max_depth); + std::exit(1); +} + +Env* puf_vec_create(int n, Dict* env_kwargs, + obs_t* observations, float* actions, + float* rewards, float* terminals) { + if (n <= 0) { + std::fprintf(stderr, "Affine Lock CUDA: vector size must be positive\n"); + std::exit(1); + } + if (g_gpu.envs != nullptr) { + std::fprintf(stderr, "Affine Lock CUDA: vector already exists\n"); + std::exit(1); + } + + VisibleTargetTable table = {}; + if (visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table) != 0) { + std::fprintf(stderr, + "Affine Lock CUDA: failed to load visible target table %s\n", + VISIBLE_TARGET_TABLE_PATH); + std::exit(1); + } + if (table.num_actions != NUM_ACTIONS) { + std::fprintf(stderr, + "Affine Lock CUDA: target table has %u actions, expected %d\n", + table.num_actions, NUM_ACTIONS); + visible_targets_free(&table); + std::exit(1); + } + + GpuAffineLockConfig config = {}; + config.start_depth = (int)dict_get(env_kwargs, "start_depth"); + config.max_depth = (int)dict_get(env_kwargs, "max_depth"); + config.step_grace = (int)dict_get(env_kwargs, "step_grace"); + config.perf_weighting = (int)dict_get(env_kwargs, "perf_weighting"); + for (uint32_t i = 0; i < table.depth_count; i++) { + const VisibleTargetDepth* depth = &table.depths[i]; + if (depth->depth >= AFFINE_LOCK_GPU_DEPTH_LUT_SIZE || + depth->stored_count == 0 || + config.depth_counts[depth->depth] != 0) { + std::fprintf(stderr, + "Affine Lock CUDA: invalid target-table depth section %u\n", + depth->depth); + visible_targets_free(&table); + std::exit(1); + } + config.depth_first[depth->depth] = depth->first_record; + config.depth_counts[depth->depth] = depth->stored_count; + for (uint32_t record_offset = 0; + record_offset < depth->stored_count; record_offset++) { + const VisibleTargetRecord* record = + &table.records[depth->first_record + record_offset]; + if (record->depth != depth->depth) { + std::fprintf(stderr, + "Affine Lock CUDA: record depth %u does not match section %u\n", + (unsigned int)record->depth, depth->depth); + visible_targets_free(&table); + std::exit(1); + } + } + } + gpu_affine_lock_validate_curriculum(&config); + + uint32_t* host_pairs = (uint32_t*)std::malloc( + (size_t)table.record_count * sizeof(uint32_t)); + if (host_pairs == nullptr) { + std::perror("malloc"); + visible_targets_free(&table); + std::exit(1); + } + for (uint32_t i = 0; i < table.record_count; i++) { + host_pairs[i] = (uint32_t)table.records[i].start + | ((uint32_t)table.records[i].target << 16); + } + + GpuAffineLockState* host_states = (GpuAffineLockState*)std::calloc( + (size_t)n, sizeof(GpuAffineLockState)); + if (host_states == nullptr) { + std::perror("calloc"); + std::free(host_pairs); + visible_targets_free(&table); + std::exit(1); + } + unsigned int running_seed = (unsigned int)dict_get(env_kwargs, "seed"); + for (int i = 0; i < n; i++) { + host_states[i].rng = (uint32_t)rand_r(&running_seed); + host_states[i].curriculum_depth = config.start_depth; + } + + Env* device_envs = nullptr; + GpuAffineLockState* device_states = nullptr; + uint32_t* device_pairs = nullptr; + gpu_affine_lock_check(cudaMalloc((void**)&device_envs, + (size_t)n * sizeof(Env)), "cudaMalloc envs"); + gpu_affine_lock_check(cudaMalloc((void**)&device_states, + (size_t)n * sizeof(GpuAffineLockState)), "cudaMalloc states"); + gpu_affine_lock_check(cudaMalloc((void**)&device_pairs, + (size_t)table.record_count * sizeof(uint32_t)), "cudaMalloc target pairs"); + gpu_affine_lock_check(cudaMemset(device_envs, 0, + (size_t)n * sizeof(Env)), "clear env logs"); + gpu_affine_lock_check(cudaMemcpy(device_states, host_states, + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyHostToDevice), "copy states"); + gpu_affine_lock_check(cudaMemcpy(device_pairs, host_pairs, + (size_t)table.record_count * sizeof(uint32_t), cudaMemcpyHostToDevice), + "copy target pairs"); + gpu_affine_lock_check(cudaMemcpyToSymbol(d_affine_lock_config, + &config, sizeof(config)), "copy config"); + + std::free(host_pairs); + std::free(host_states); + visible_targets_free(&table); + + g_gpu.envs = device_envs; + g_gpu.states = device_states; + g_gpu.target_pairs = device_pairs; + g_gpu.n = n; + g_gpu.observations = observations; + g_gpu.actions = actions; + g_gpu.rewards = rewards; + g_gpu.terminals = terminals; + g_gpu.stream = nullptr; + g_gpu.config = config; + return device_envs; +} + +void puf_bind_stream(cudaStream_t stream) { + g_gpu.stream = stream; +} + +// GPU creation is vector-only; puf_init exists to satisfy the common API. +void puf_init(Env* env, Dict* kwargs) { + (void)env; + (void)kwargs; +} + +void puf_reset(Env* env) { + (void)env; +#if AFFINE_LOCK_GPU_SHARED_OBS + int blocks = (g_gpu.n + AFFINE_LOCK_GPU_SHARED_BLOCK - 1) + / AFFINE_LOCK_GPU_SHARED_BLOCK; + gpu_affine_lock_shared_reset_kernel<<< + blocks, AFFINE_LOCK_GPU_SHARED_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.observations, g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#else + int threads = g_gpu.n * AFFINE_LOCK_GPU_LANES; + gpu_affine_lock_reset_kernel<<< + gpu_affine_lock_grid(threads), AFFINE_LOCK_GPU_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.observations, g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#endif + gpu_affine_lock_check(cudaPeekAtLastError(), "launch reset kernel"); +} + +void puf_step(Env* env) { + (void)env; +#if AFFINE_LOCK_GPU_SHARED_OBS + int blocks = (g_gpu.n + AFFINE_LOCK_GPU_SHARED_BLOCK - 1) + / AFFINE_LOCK_GPU_SHARED_BLOCK; + gpu_affine_lock_shared_step_kernel<<< + blocks, AFFINE_LOCK_GPU_SHARED_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.actions, g_gpu.observations, + g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#else + int threads = g_gpu.n * AFFINE_LOCK_GPU_LANES; + gpu_affine_lock_step_kernel<<< + gpu_affine_lock_grid(threads), AFFINE_LOCK_GPU_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.actions, g_gpu.observations, + g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#endif + gpu_affine_lock_check(cudaPeekAtLastError(), "launch step kernel"); +} + +void puf_close(Env* env) { + (void)env; + if (IsWindowReady()) { + CloseWindow(); + } + if (g_gpu.envs != nullptr) { + gpu_affine_lock_check(cudaFree(g_gpu.envs), "cudaFree envs"); + } + if (g_gpu.states != nullptr) { + gpu_affine_lock_check(cudaFree(g_gpu.states), "cudaFree states"); + } + if (g_gpu.target_pairs != nullptr) { + gpu_affine_lock_check(cudaFree(g_gpu.target_pairs), "cudaFree target pairs"); + } + g_gpu = {}; +} + +void puf_render(Env* env) { + (void)env; + if (g_gpu.envs == nullptr || g_gpu.n < 1) { + return; + } + if (IsWindowReady() && (WindowShouldClose() || IsKeyPressed(KEY_ESCAPE))) { + puf_close(g_gpu.envs); + std::exit(0); + } + if (!IsWindowReady()) { + InitWindow(780, 360, "PufferLib AffineLock CUDA"); + SetTargetFPS(30); + } + if (g_gpu.stream != nullptr) { + gpu_affine_lock_check(cudaStreamSynchronize(g_gpu.stream), + "synchronize render stream"); + } + GpuAffineLockState state; + float reward = 0.0f; + float terminal = 0.0f; + gpu_affine_lock_check(cudaMemcpy(&state, g_gpu.states, sizeof(state), + cudaMemcpyDeviceToHost), "copy render state"); + gpu_affine_lock_check(cudaMemcpy(&reward, g_gpu.rewards, sizeof(reward), + cudaMemcpyDeviceToHost), "copy render reward"); + gpu_affine_lock_check(cudaMemcpy(&terminal, g_gpu.terminals, sizeof(terminal), + cudaMemcpyDeviceToHost), "copy render terminal"); + + uint32_t mismatches = (state.state ^ state.target) & 0xffffu; + const char* status = terminal == 0.0f + ? "running" : (reward > 0.0f ? "solved" : "failed"); + Color status_color = terminal == 0.0f + ? (Color){190, 198, 206, 255} + : (reward > 0.0f + ? (Color){80, 210, 140, 255} + : (Color){238, 88, 88, 255}); + + BeginDrawing(); + ClearBackground((Color){6, 24, 24, 255}); + DrawText("Affine Lock CUDA", 30, 24, 28, RAYWHITE); + DrawText(TextFormat("depth %d/%d step %d/%d last reward %.2f", + state.scramble_depth, g_gpu.config.max_depth, + state.step_count, state.max_steps, reward), + 30, 62, 20, (Color){180, 190, 200, 255}); + DrawText(TextFormat("status %s mismatches 0x%04x", + status, mismatches), 30, 90, 20, status_color); + + const char* row_label[2] = {"current", "target"}; + uint32_t row_value[2] = {state.state, state.target}; + int row_y[2] = {138, 220}; + for (int row = 0; row < 2; row++) { + DrawText(row_label[row], 30, row_y[row] + 9, 20, RAYWHITE); + for (int bit = 0; bit < BITS; bit++) { + int x = 145 + bit * 34; + int on = (row_value[row] >> bit) & 1u; + int mismatch = ((state.state ^ state.target) >> bit) & 1u; + Color fill = on + ? (Color){80, 210, 140, 255} + : (Color){38, 48, 58, 255}; + Color border = mismatch + ? (Color){238, 88, 88, 255} + : (Color){182, 196, 205, 255}; + DrawRectangle(x, row_y[row], 24, 34, fill); + DrawRectangleLinesEx((Rectangle){(float)x, (float)row_y[row], 24, 34}, + mismatch ? 3 : 1, border); + DrawText(TextFormat("%d", bit), x + 5, row_y[row] + 40, 10, + (Color){128, 140, 150, 255}); + } + } + DrawText("GPU-resident environment", 30, 310, 16, + (Color){160, 170, 178, 255}); + EndDrawing(); + puf_web_vsync(); +} + +#endif diff --git a/ocean/affine_lock/tests/run_cuda.sh b/ocean/affine_lock/tests/run_cuda.sh new file mode 100755 index 0000000000..f7762be9a1 --- /dev/null +++ b/ocean/affine_lock/tests/run_cuda.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "$0")/../../.." && pwd)" +cd "$ROOT" + +OUT="${TMPDIR:-/tmp}/affine_lock_cuda_tests" +CUDA_ROOT="${CUDA_HOME:-${CUDA_PATH:-/usr/local/cuda}}" +NVCC_BIN="${NVCC:-$CUDA_ROOT/bin/nvcc}" +CUDA_ARCH="${NVCC_ARCH:-native}" + +RAYLIB_ROOT="$ROOT/raylib-5.5_linux_amd64" +if [ ! -d "$RAYLIB_ROOT/include" ]; then + echo "raylib-5.5_linux_amd64 not found" >&2 + exit 1 +fi + +"$NVCC_BIN" \ + -std=c++17 -O3 -lineinfo -arch="$CUDA_ARCH" \ + -Xcompiler=-Wall,-Wextra,-Werror,-Wno-unused-function,-Wno-unused-parameter,-Wno-missing-field-initializers \ + -Xcompiler=-ffunction-sections,-fdata-sections \ + -I"$ROOT" -I"$ROOT/src" -I"$ROOT/ocean/affine_lock" \ + -I"$ROOT/vendor" -I"$RAYLIB_ROOT/include" \ + "$ROOT/ocean/affine_lock/tests/test_affine_lock_cuda.cu" \ + "$RAYLIB_ROOT/lib/libraylib.a" \ + -Xlinker=--gc-sections \ + -lGL -lpthread -ldl -lrt -lm \ + -o "$OUT" + +"$OUT" diff --git a/ocean/affine_lock/tests/test_affine_lock_cuda.cu b/ocean/affine_lock/tests/test_affine_lock_cuda.cu new file mode 100644 index 0000000000..54cb06ad34 --- /dev/null +++ b/ocean/affine_lock/tests/test_affine_lock_cuda.cu @@ -0,0 +1,1109 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "../affine_lock.cu" + +#define EXPECT_TRUE(condition) do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: expected true: %s\n", \ + __FILE__, __LINE__, #condition); \ + std::exit(1); \ + } \ +} while (0) + +#define EXPECT_EQ(actual, expected) do { \ + auto actual_value = (actual); \ + auto expected_value = (expected); \ + if (actual_value != expected_value) { \ + std::fprintf(stderr, "%s:%d: expected %s == %s, got %lld != %lld\n", \ + __FILE__, __LINE__, #actual, #expected, \ + (long long)actual_value, (long long)expected_value); \ + std::exit(1); \ + } \ +} while (0) + +#define EXPECT_NEAR(actual, expected, tolerance) do { \ + float actual_value = (float)(actual); \ + float expected_value = (float)(expected); \ + if (!std::isfinite(actual_value) || !std::isfinite(expected_value) || \ + std::fabs(actual_value - expected_value) > (tolerance)) { \ + std::fprintf(stderr, "%s:%d: expected %s ~= %.9g, got %.9g\n", \ + __FILE__, __LINE__, #actual, expected_value, actual_value); \ + std::exit(1); \ + } \ +} while (0) + +static void check_cuda(cudaError_t status, const char* operation) { + if (status != cudaSuccess) { + std::fprintf(stderr, "%s failed: %s\n", operation, cudaGetErrorString(status)); + std::exit(1); + } +} + +typedef struct OracleState { + uint32_t rng; + uint16_t state; + uint16_t target; + int step_count; + int max_steps; + int scramble_depth; + int curriculum_depth; + int target_distance; + float episode_return; + Log log; +} OracleState; + +static uint32_t oracle_random_mixed_u32(OracleState* env) { + env->rng = env->rng * 1664525u + 1013904223u; + uint32_t x = env->rng; + x ^= x >> 16; + x *= 0x7feb352du; + x ^= x >> 15; + x *= 0x846ca68bu; + x ^= x >> 16; + return x; +} + +static int oracle_random_bounded(OracleState* env, int bound) { + uint32_t ubound = (uint32_t)bound; + uint32_t limit = UINT32_MAX - UINT32_MAX % ubound; + uint32_t value = oracle_random_mixed_u32(env); + while (value >= limit) { + value = oracle_random_mixed_u32(env); + } + return (int)(value % ubound); +} + +static const VisibleTargetDepth* oracle_depth( + const VisibleTargetTable* table, int requested_depth) { + for (uint32_t i = 0; i < table->depth_count; i++) { + if ((int)table->depths[i].depth == requested_depth) { + return &table->depths[i]; + } + } + return nullptr; +} + +static void oracle_reset_state(OracleState* env, + const VisibleTargetTable* table, int step_grace) { + env->scramble_depth = env->curriculum_depth; + env->step_count = 0; + env->episode_return = 0.0f; + const VisibleTargetDepth* depth = oracle_depth(table, env->scramble_depth); + EXPECT_TRUE(depth != nullptr); + int choice = oracle_random_bounded(env, (int)depth->stored_count); + const VisibleTargetRecord* record = + &table->records[depth->first_record + (uint32_t)choice]; + env->state = record->start; + env->target = record->target; + env->target_distance = record->depth; + env->max_steps = env->target_distance + step_grace; +} + +static uint16_t oracle_apply_action(uint16_t state, int action) { + uint32_t next = state; + switch (action) { + case 0: next = (state >> 1) | ((state & 1u) << 15); break; + case 1: next = ((state << 1) & 0xffffu) | ((state >> 15) & 1u); break; + case 2: next = state ^ 0xfe00u; break; + case 3: next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); break; + case 4: next = ((state & 0x3333u) << 2) | ((state & 0xccccu) >> 2); break; + case 5: next = ((state & 0x0f0fu) << 4) | ((state & 0xf0f0u) >> 4); break; + case 6: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + break; + case 7: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + next = ((next & 0x0f0fu) << 4) | ((next & 0xf0f0u) >> 4); + break; + } + return (uint16_t)(next & 0xffffu); +} + +static int oracle_next_curriculum_depth(int current_depth, int max_depth) { + static const int curriculum_depths[] = {2, 4, 5, 6, 8, 16}; + for (int depth : curriculum_depths) { + if (depth > current_depth) { + return depth < max_depth ? depth : max_depth; + } + } + return max_depth; +} + +static void oracle_add_log(OracleState* env, int solved, + int max_depth, int perf_weighting) { + int log_depth = env->target_distance; + int at_max_depth = log_depth == max_depth; + float ratio = log_depth / (float)max_depth; + float solve_credit = 0.0f; + if (solved) { + solve_credit = perf_weighting == PERF_WEIGHTING_QUADRATIC + ? ratio * ratio : ratio; + } + env->log.perf += solve_credit; + env->log.score += solve_credit; + env->log.solve_rate += solved; + env->log.max_depth_solve += solved && at_max_depth; + env->log.episode_return += env->episode_return; + env->log.episode_length += env->step_count; + env->log.solve_steps += solved ? env->step_count : 0; + env->log.timeout_rate += !solved; + env->log.solve_efficiency += solved + ? env->step_count / (float)log_depth : 0.0f; + env->log.target_distance += env->target_distance; + env->log.solved_target_distance += solved ? env->target_distance : 0; + env->log.d6_rate += log_depth == 6; + env->log.d6_solve_rate += solved && log_depth == 6; + env->log.d8_rate += log_depth == 8; + env->log.d8_solve_rate += solved && log_depth == 8; + env->log.d16_rate += log_depth == 16; + env->log.d16_solve_rate += solved && log_depth == 16; + env->log.n += 1; +} + +static void oracle_step(OracleState* env, float action, + const VisibleTargetTable* table, int start_depth, int max_depth, + int step_grace, int perf_weighting, float* reward, float* terminal) { + *reward = STEP_REWARD; + *terminal = 0.0f; + int solved = 0; + env->step_count += 1; + int invalid = !std::isfinite(action) || action < 0.0f || action > 7.0f; + if (invalid) { + *reward = -1.0f; + *terminal = 1.0f; + } else { + env->state = oracle_apply_action(env->state, (int)action); + if (env->state == env->target) { + *reward = 1.0f; + *terminal = 1.0f; + solved = 1; + } else if (env->step_count >= env->max_steps) { + *reward = -1.0f; + *terminal = 1.0f; + } + } + env->episode_return += *reward; + if (*terminal != 0.0f) { + oracle_add_log(env, solved, max_depth, perf_weighting); + env->curriculum_depth = solved + ? oracle_next_curriculum_depth(env->scramble_depth, max_depth) + : start_depth; + oracle_reset_state(env, table, step_grace); + } +} + +static uint16_t obs_bits(obs_t value) { + uint16_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +static uint32_t float_bits(float value) { + uint32_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +static void expect_log_equal(const Log& actual, const Log& expected) { + const float* a = (const float*)&actual; + const float* e = (const float*)&expected; + for (size_t i = 0; i < sizeof(Log) / sizeof(float); i++) { + EXPECT_EQ(float_bits(a[i]), float_bits(e[i])); + } +} + +static void expect_state_equal(const GpuAffineLockState& actual, + const OracleState& expected) { + EXPECT_EQ(actual.rng, expected.rng); + EXPECT_EQ(actual.state, expected.state); + EXPECT_EQ(actual.target, expected.target); + EXPECT_EQ(actual.step_count, expected.step_count); + EXPECT_EQ(actual.max_steps, expected.max_steps); + EXPECT_EQ(actual.scramble_depth, expected.scramble_depth); + EXPECT_EQ(actual.curriculum_depth, expected.curriculum_depth); + EXPECT_EQ(actual.target_distance, expected.target_distance); + EXPECT_EQ(float_bits(actual.episode_return), + float_bits(expected.episode_return)); +} + +static void expect_observation_equal(const obs_t* actual, + const OracleState& expected) { + uint32_t bits = (uint32_t)expected.state | ((uint32_t)expected.target << 16); + for (int bit = 0; bit < 32; bit++) { + float value = (bits & (1u << bit)) ? 1.0f : -1.0f; + EXPECT_EQ(obs_bits(actual[bit]), obs_bits(__float2bfloat16(value))); + } + float timer = expected.step_count / (float)expected.max_steps; + EXPECT_EQ(obs_bits(actual[TIMER_INDEX]), + obs_bits(__float2bfloat16(timer))); +} + +static void fill_kwargs(Dict* kwargs, int seed, int step_grace, + int perf_weighting) { + std::memset(kwargs, 0, sizeof(*kwargs)); + dict_set(kwargs, "seed", seed); + dict_set(kwargs, "start_depth", 2); + dict_set(kwargs, "max_depth", 16); + dict_set(kwargs, "step_grace", step_grace); + dict_set(kwargs, "perf_weighting", perf_weighting); +} + +static void test_deterministic_reset_and_step_parity() { + constexpr int n = 257; + constexpr int seed = 42; + constexpr int step_grace = 2; + constexpr int perf_weighting = PERF_WEIGHTING_QUADRATIC; + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), "cudaMalloc observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), "cudaMalloc actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), "cudaMalloc rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), "cudaMalloc terminals"); + + Dict kwargs; + fill_kwargs(&kwargs, seed, step_grace, perf_weighting); + Env* envs = puf_vec_create(n, &kwargs, observations, actions, rewards, terminals); + EXPECT_TRUE(envs != nullptr); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "initial reset"); + + std::vector oracle(n); + unsigned int running_seed = seed; + for (int i = 0; i < n; i++) { + oracle[i].rng = (uint32_t)rand_r(&running_seed); + oracle[i].curriculum_depth = 2; + oracle_reset_state(&oracle[i], &table, step_grace); + } + + std::vector states(n); + std::vector host_envs(n); + std::vector host_obs((size_t)n * OBS_SIZE); + std::vector host_rewards(n), host_terminals(n), host_actions(n); + check_cuda(cudaMemcpy(states.data(), g_gpu.states, + n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), "copy reset states"); + check_cuda(cudaMemcpy(host_obs.data(), observations, + host_obs.size() * sizeof(obs_t), cudaMemcpyDeviceToHost), "copy reset observations"); + check_cuda(cudaMemcpy(host_rewards.data(), rewards, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy reset rewards"); + check_cuda(cudaMemcpy(host_terminals.data(), terminals, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy reset terminals"); + for (int i = 0; i < n; i++) { + expect_state_equal(states[i], oracle[i]); + expect_observation_equal(&host_obs[(size_t)i * OBS_SIZE], oracle[i]); + EXPECT_NEAR(host_rewards[i], 0.0f, 0.0f); + EXPECT_NEAR(host_terminals[i], 0.0f, 0.0f); + } + + for (int step = 0; step < 96; step++) { + for (int i = 0; i < n; i++) { + int selector = (step * 17 + i * 13) % 41; + if (selector == 0) host_actions[i] = std::numeric_limits::quiet_NaN(); + else if (selector == 1) host_actions[i] = -0.25f; + else if (selector == 2) host_actions[i] = 8.0f; + else host_actions[i] = (float)((step + 3 * i) & 7) + (selector == 3 ? 0.75f : 0.0f); + } + check_cuda(cudaMemcpy(actions, host_actions.data(), + n * sizeof(float), cudaMemcpyHostToDevice), "copy actions"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "step"); + + for (int i = 0; i < n; i++) { + oracle_step(&oracle[i], host_actions[i], &table, + 2, 16, step_grace, perf_weighting, + &host_rewards[i], &host_terminals[i]); + } + + check_cuda(cudaMemcpy(states.data(), g_gpu.states, + n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), "copy states"); + check_cuda(cudaMemcpy(host_envs.data(), envs, + n * sizeof(Env), cudaMemcpyDeviceToHost), "copy logs"); + check_cuda(cudaMemcpy(host_obs.data(), observations, + host_obs.size() * sizeof(obs_t), cudaMemcpyDeviceToHost), "copy observations"); + std::vector actual_rewards(n), actual_terminals(n); + check_cuda(cudaMemcpy(actual_rewards.data(), rewards, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy rewards"); + check_cuda(cudaMemcpy(actual_terminals.data(), terminals, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy terminals"); + for (int i = 0; i < n; i++) { + expect_state_equal(states[i], oracle[i]); + expect_log_equal(host_envs[i].log, oracle[i].log); + expect_observation_equal(&host_obs[(size_t)i * OBS_SIZE], oracle[i]); + EXPECT_NEAR(actual_rewards[i], host_rewards[i], 0.0f); + EXPECT_NEAR(actual_terminals[i], host_terminals[i], 0.0f); + } + } + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree observations"); + check_cuda(cudaFree(actions), "cudaFree actions"); + check_cuda(cudaFree(rewards), "cudaFree rewards"); + check_cuda(cudaFree(terminals), "cudaFree terminals"); +} + +static void test_reset_rejection_sampling() { + constexpr uint32_t rejection_seed = 24481u; + constexpr uint32_t expected_final_rng = 3424986747u; + static const int depths[] = {2, 16}; + static const int expected_choices[] = {44338, 15778}; + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, OBS_SIZE * sizeof(obs_t)), + "cudaMalloc rejection observations"); + check_cuda(cudaMalloc(&actions, sizeof(float)), + "cudaMalloc rejection action"); + check_cuda(cudaMalloc(&rewards, sizeof(float)), + "cudaMalloc rejection reward"); + check_cuda(cudaMalloc(&terminals, sizeof(float)), + "cudaMalloc rejection terminal"); + + Dict kwargs; + fill_kwargs(&kwargs, 1, 0, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(1, &kwargs, + observations, actions, rewards, terminals); + + for (int case_index = 0; case_index < 2; case_index++) { + int depth = depths[case_index]; + GpuAffineLockState injected = {}; + injected.rng = rejection_seed; + injected.curriculum_depth = depth; + check_cuda(cudaMemcpy(g_gpu.states, &injected, sizeof(injected), + cudaMemcpyHostToDevice), "inject rejection state"); + + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "rejection reset"); + + OracleState expected = {}; + expected.rng = rejection_seed; + expected.curriculum_depth = depth; + oracle_reset_state(&expected, &table, 0); + EXPECT_EQ(expected.rng, expected_final_rng); + const VisibleTargetDepth* table_depth = oracle_depth(&table, depth); + EXPECT_TRUE(table_depth != nullptr); + const VisibleTargetRecord* selected = &table.records[ + table_depth->first_record + (uint32_t)expected_choices[case_index]]; + EXPECT_EQ(expected.state, selected->start); + EXPECT_EQ(expected.target, selected->target); + + GpuAffineLockState actual = {}; + obs_t actual_obs[OBS_SIZE]; + float actual_reward = 123.0f; + float actual_terminal = 123.0f; + check_cuda(cudaMemcpy(&actual, g_gpu.states, sizeof(actual), + cudaMemcpyDeviceToHost), "copy rejection state"); + check_cuda(cudaMemcpy(actual_obs, observations, sizeof(actual_obs), + cudaMemcpyDeviceToHost), "copy rejection observations"); + check_cuda(cudaMemcpy(&actual_reward, rewards, sizeof(actual_reward), + cudaMemcpyDeviceToHost), "copy rejection reward"); + check_cuda(cudaMemcpy(&actual_terminal, terminals, sizeof(actual_terminal), + cudaMemcpyDeviceToHost), "copy rejection terminal"); + expect_state_equal(actual, expected); + expect_observation_equal(actual_obs, expected); + EXPECT_EQ(float_bits(actual_reward), float_bits(0.0f)); + EXPECT_EQ(float_bits(actual_terminal), float_bits(0.0f)); + } + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree rejection observations"); + check_cuda(cudaFree(actions), "cudaFree rejection action"); + check_cuda(cudaFree(rewards), "cudaFree rejection reward"); + check_cuda(cudaFree(terminals), "cudaFree rejection terminal"); +} + +static void test_puf_log_exports_cpu_contract() { + Log log = {}; + log.perf = 1.25f; + log.score = 2.5f; + log.solve_rate = 2.0f; + log.max_depth_solve = 1.0f; + log.episode_return = 3.5f; + log.episode_length = 8.0f; + log.solve_steps = 5.0f; + log.timeout_rate = 1.0f; + log.solve_efficiency = 1.75f; + log.target_distance = 20.0f; + log.solved_target_distance = 12.0f; + log.d6_rate = 2.0f; + log.d6_solve_rate = 1.0f; + log.d8_rate = 4.0f; + log.d8_solve_rate = 3.0f; + log.d16_rate = 1.0f; + log.d16_solve_rate = 1.0f; + log.n = 3.0f; + Dict out = {}; + puf_log(&log, &out); + EXPECT_EQ(out.size, 15); + EXPECT_NEAR(dict_get(&out, "perf"), 1.25f, 0.0f); + EXPECT_NEAR(dict_get(&out, "score"), 2.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "solve_rate"), 2.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "max_depth_solve"), 1.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "episode_return"), 3.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "episode_length"), 8.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "timeout_rate"), 1.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "min_win_moves"), 20.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "solved_min_win_moves"), 6.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "conditional_solve_steps"), 2.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "conditional_solve_efficiency"), 0.875f, 0.0f); + EXPECT_NEAR(dict_get(&out, "d6_solve_rate"), 0.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "d8_solve_rate"), 0.75f, 0.0f); + EXPECT_NEAR(dict_get(&out, "d16_solve_rate"), 1.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "n"), 3.0f, 0.0f); + dict_clear(&out); + + Log zero_denominators = {}; + zero_denominators.solved_target_distance = 12.0f; + zero_denominators.solve_steps = 5.0f; + zero_denominators.solve_efficiency = 1.75f; + zero_denominators.d6_solve_rate = 1.0f; + zero_denominators.d8_solve_rate = 1.0f; + zero_denominators.d16_solve_rate = 1.0f; + Dict zero_out = {}; + puf_log(&zero_denominators, &zero_out); + EXPECT_EQ(zero_out.size, 15); + EXPECT_NEAR(dict_get(&zero_out, "solved_min_win_moves"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "conditional_solve_steps"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "conditional_solve_efficiency"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "d6_solve_rate"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "d8_solve_rate"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "d16_solve_rate"), 0.0f, 0.0f); + dict_clear(&zero_out); +} + +static void test_exhaustive_action_transforms() { + constexpr int n = 1 << BITS; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMalloc exhaustive observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), + "cudaMalloc exhaustive actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), + "cudaMalloc exhaustive rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), + "cudaMalloc exhaustive terminals"); + + Dict kwargs; + fill_kwargs(&kwargs, 7, 100, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + std::vector states(n); + std::vector host_actions(n), host_rewards(n), host_terminals(n); + + for (int action = 0; action < NUM_ACTIONS; action++) { + for (int value = 0; value < n; value++) { + uint16_t expected = oracle_apply_action((uint16_t)value, action); + states[value] = {}; + states[value].rng = (uint32_t)(value + 1); + states[value].state = (uint16_t)value; + states[value].target = expected ^ 1u; + states[value].max_steps = 100; + states[value].scramble_depth = 16; + states[value].curriculum_depth = 16; + states[value].target_distance = 16; + host_actions[value] = (float)action; + } + check_cuda(cudaMemcpy(g_gpu.states, states.data(), + n * sizeof(GpuAffineLockState), cudaMemcpyHostToDevice), + "copy exhaustive states"); + check_cuda(cudaMemcpy(actions, host_actions.data(), + n * sizeof(float), cudaMemcpyHostToDevice), + "copy exhaustive actions"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "exhaustive action step"); + check_cuda(cudaMemcpy(states.data(), g_gpu.states, + n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), + "copy exhaustive results"); + check_cuda(cudaMemcpy(host_rewards.data(), rewards, + n * sizeof(float), cudaMemcpyDeviceToHost), + "copy exhaustive rewards"); + check_cuda(cudaMemcpy(host_terminals.data(), terminals, + n * sizeof(float), cudaMemcpyDeviceToHost), + "copy exhaustive terminals"); + for (int value = 0; value < n; value++) { + EXPECT_EQ(states[value].state, + oracle_apply_action((uint16_t)value, action)); + EXPECT_EQ(states[value].step_count, 1); + EXPECT_NEAR(states[value].episode_return, STEP_REWARD, 0.0f); + EXPECT_NEAR(host_rewards[value], STEP_REWARD, 0.0f); + EXPECT_NEAR(host_terminals[value], 0.0f, 0.0f); + } + } + + puf_close(envs); + dict_clear(&kwargs); + check_cuda(cudaFree(observations), "cudaFree exhaustive observations"); + check_cuda(cudaFree(actions), "cudaFree exhaustive actions"); + check_cuda(cudaFree(rewards), "cudaFree exhaustive rewards"); + check_cuda(cudaFree(terminals), "cudaFree exhaustive terminals"); +} + +static void test_action_boundaries() { + const float infinity = std::numeric_limits::infinity(); + const float actions_under_test[] = { + -infinity, + -0.25f, + -std::numeric_limits::denorm_min(), + -0.0f, + 0.0f, + 0.5f, + 0.999f, + 1.0f, + 1.5f, + 6.999f, + 7.0f, + std::nextafter(7.0f, infinity), + 8.0f, + infinity, + std::numeric_limits::quiet_NaN(), + }; + constexpr int n = sizeof(actions_under_test) / sizeof(actions_under_test[0]); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMalloc boundary observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), + "cudaMalloc boundary actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), + "cudaMalloc boundary rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), + "cudaMalloc boundary terminals"); + + Dict kwargs; + fill_kwargs(&kwargs, 11, 100, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + + std::vector injected(n); + std::vector expected(n); + std::vector expected_rewards(n), expected_terminals(n); + for (int i = 0; i < n; i++) { + injected[i] = {}; + injected[i].rng = (uint32_t)(1000 + i); + injected[i].state = 0x1234u; + injected[i].target = 0xbeefu; + injected[i].max_steps = 100; + injected[i].scramble_depth = 2; + injected[i].curriculum_depth = 2; + injected[i].target_distance = 2; + + expected[i].rng = injected[i].rng; + expected[i].state = injected[i].state; + expected[i].target = injected[i].target; + expected[i].max_steps = injected[i].max_steps; + expected[i].scramble_depth = injected[i].scramble_depth; + expected[i].curriculum_depth = injected[i].curriculum_depth; + expected[i].target_distance = injected[i].target_distance; + oracle_step(&expected[i], actions_under_test[i], &table, + 2, 16, 100, PERF_WEIGHTING_LINEAR, + &expected_rewards[i], &expected_terminals[i]); + } + check_cuda(cudaMemcpy(g_gpu.states, injected.data(), + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyHostToDevice), + "copy boundary states"); + check_cuda(cudaMemcpy(actions, actions_under_test, + sizeof(actions_under_test), cudaMemcpyHostToDevice), + "copy boundary actions"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "boundary step"); + + std::vector actual_states(n); + std::vector actual_envs(n); + std::vector actual_obs((size_t)n * OBS_SIZE); + std::vector actual_rewards(n), actual_terminals(n); + check_cuda(cudaMemcpy(actual_states.data(), g_gpu.states, + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), + "copy boundary results"); + check_cuda(cudaMemcpy(actual_envs.data(), envs, + (size_t)n * sizeof(Env), cudaMemcpyDeviceToHost), + "copy boundary logs"); + check_cuda(cudaMemcpy(actual_obs.data(), observations, + actual_obs.size() * sizeof(obs_t), cudaMemcpyDeviceToHost), + "copy boundary observations"); + check_cuda(cudaMemcpy(actual_rewards.data(), rewards, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost), + "copy boundary rewards"); + check_cuda(cudaMemcpy(actual_terminals.data(), terminals, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost), + "copy boundary terminals"); + + for (int i = 0; i < n; i++) { + bool invalid = !std::isfinite(actions_under_test[i]) + || actions_under_test[i] < 0.0f || actions_under_test[i] > 7.0f; + expect_state_equal(actual_states[i], expected[i]); + expect_log_equal(actual_envs[i].log, expected[i].log); + expect_observation_equal( + &actual_obs[(size_t)i * OBS_SIZE], expected[i]); + EXPECT_EQ(float_bits(actual_rewards[i]), + float_bits(expected_rewards[i])); + EXPECT_EQ(float_bits(actual_terminals[i]), + float_bits(expected_terminals[i])); + EXPECT_EQ(float_bits(actual_terminals[i]), + float_bits(invalid ? 1.0f : 0.0f)); + EXPECT_EQ(float_bits(actual_envs[i].log.n), + float_bits(invalid ? 1.0f : 0.0f)); + } + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree boundary observations"); + check_cuda(cudaFree(actions), "cudaFree boundary actions"); + check_cuda(cudaFree(rewards), "cudaFree boundary rewards"); + check_cuda(cudaFree(terminals), "cudaFree boundary terminals"); +} + +static const VisibleTargetRecord* find_solution_record( + const VisibleTargetTable* table, const OracleState& state) { + const VisibleTargetDepth* depth = oracle_depth(table, state.scramble_depth); + EXPECT_TRUE(depth != nullptr); + for (uint32_t i = 0; i < depth->stored_count; i++) { + const VisibleTargetRecord* record = + &table->records[depth->first_record + i]; + if (record->start == state.state && record->target == state.target) { + return record; + } + } + return nullptr; +} + +static void test_solution_curriculum_and_logs() { + constexpr int seed = 69; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, OBS_SIZE * sizeof(obs_t)), + "cudaMalloc solution observations"); + check_cuda(cudaMalloc(&actions, sizeof(float)), "cudaMalloc solution action"); + check_cuda(cudaMalloc(&rewards, sizeof(float)), "cudaMalloc solution reward"); + check_cuda(cudaMalloc(&terminals, sizeof(float)), "cudaMalloc solution terminal"); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + Dict kwargs; + fill_kwargs(&kwargs, seed, 0, PERF_WEIGHTING_QUADRATIC); + Env* envs = puf_vec_create(1, &kwargs, + observations, actions, rewards, terminals); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "solution reset"); + + OracleState oracle = {}; + unsigned int running_seed = seed; + oracle.rng = (uint32_t)rand_r(&running_seed); + oracle.curriculum_depth = 2; + oracle_reset_state(&oracle, &table, 0); + static const int expected_depths[] = {2, 4, 5, 6, 8, 16}; + for (int expected_depth : expected_depths) { + EXPECT_EQ(oracle.scramble_depth, expected_depth); + const VisibleTargetRecord* record = find_solution_record(&table, oracle); + EXPECT_TRUE(record != nullptr); + for (int move = 0; move < record->solution_length; move++) { + float action = (float)((record->packed_actions >> (3 * move)) & 7u); + check_cuda(cudaMemcpy(actions, &action, sizeof(float), + cudaMemcpyHostToDevice), "copy solution action"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "solution step"); + + float expected_reward = 0.0f; + float expected_terminal = 0.0f; + oracle_step(&oracle, action, &table, 2, 16, 0, + PERF_WEIGHTING_QUADRATIC, + &expected_reward, &expected_terminal); + GpuAffineLockState actual_state; + Env actual_env; + obs_t actual_obs[OBS_SIZE]; + float actual_reward = 0.0f; + float actual_terminal = 0.0f; + check_cuda(cudaMemcpy(&actual_state, g_gpu.states, + sizeof(actual_state), cudaMemcpyDeviceToHost), + "copy solution state"); + check_cuda(cudaMemcpy(&actual_env, envs, + sizeof(actual_env), cudaMemcpyDeviceToHost), + "copy solution log"); + check_cuda(cudaMemcpy(actual_obs, observations, + sizeof(actual_obs), cudaMemcpyDeviceToHost), + "copy solution observations"); + check_cuda(cudaMemcpy(&actual_reward, rewards, + sizeof(float), cudaMemcpyDeviceToHost), + "copy solution reward"); + check_cuda(cudaMemcpy(&actual_terminal, terminals, + sizeof(float), cudaMemcpyDeviceToHost), + "copy solution terminal"); + expect_state_equal(actual_state, oracle); + expect_log_equal(actual_env.log, oracle.log); + expect_observation_equal(actual_obs, oracle); + EXPECT_NEAR(actual_reward, expected_reward, 0.0f); + EXPECT_NEAR(actual_terminal, expected_terminal, 0.0f); + EXPECT_NEAR(actual_terminal, + move + 1 == record->solution_length ? 1.0f : 0.0f, 0.0f); + } + } + EXPECT_NEAR(oracle.log.solve_rate, 6.0f, 0.0f); + EXPECT_NEAR(oracle.log.d6_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(oracle.log.d8_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(oracle.log.d16_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(oracle.log.max_depth_solve, 1.0f, 0.0f); + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree solution observations"); + check_cuda(cudaFree(actions), "cudaFree solution action"); + check_cuda(cudaFree(rewards), "cudaFree solution reward"); + check_cuda(cudaFree(terminals), "cudaFree solution terminal"); +} + +static void test_linear_solve_scoring() { + constexpr int seed = 17; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, OBS_SIZE * sizeof(obs_t)), + "cudaMalloc linear observations"); + check_cuda(cudaMalloc(&actions, sizeof(float)), + "cudaMalloc linear action"); + check_cuda(cudaMalloc(&rewards, sizeof(float)), + "cudaMalloc linear reward"); + check_cuda(cudaMalloc(&terminals, sizeof(float)), + "cudaMalloc linear terminal"); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + Dict kwargs; + fill_kwargs(&kwargs, seed, 0, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(1, &kwargs, + observations, actions, rewards, terminals); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "linear reset"); + + OracleState oracle = {}; + unsigned int running_seed = seed; + oracle.rng = (uint32_t)rand_r(&running_seed); + oracle.curriculum_depth = 2; + oracle_reset_state(&oracle, &table, 0); + const VisibleTargetRecord* record = find_solution_record(&table, oracle); + EXPECT_TRUE(record != nullptr); + for (int move = 0; move < record->solution_length; move++) { + float action = (float)((record->packed_actions >> (3 * move)) & 7u); + check_cuda(cudaMemcpy(actions, &action, sizeof(action), + cudaMemcpyHostToDevice), "copy linear solution action"); + puf_step(envs); + } + check_cuda(cudaDeviceSynchronize(), "linear solution"); + + Env actual_env = {}; + float actual_reward = 0.0f; + float actual_terminal = 0.0f; + check_cuda(cudaMemcpy(&actual_env, envs, sizeof(actual_env), + cudaMemcpyDeviceToHost), "copy linear log"); + check_cuda(cudaMemcpy(&actual_reward, rewards, sizeof(actual_reward), + cudaMemcpyDeviceToHost), "copy linear reward"); + check_cuda(cudaMemcpy(&actual_terminal, terminals, sizeof(actual_terminal), + cudaMemcpyDeviceToHost), "copy linear terminal"); + EXPECT_NEAR(actual_env.log.perf, 2.0f / 16.0f, 0.0f); + EXPECT_NEAR(actual_env.log.score, 2.0f / 16.0f, 0.0f); + EXPECT_NEAR(actual_env.log.solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(actual_env.log.n, 1.0f, 0.0f); + EXPECT_NEAR(actual_reward, 1.0f, 0.0f); + EXPECT_NEAR(actual_terminal, 1.0f, 0.0f); + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree linear observations"); + check_cuda(cudaFree(actions), "cudaFree linear action"); + check_cuda(cudaFree(rewards), "cudaFree linear reward"); + check_cuda(cudaFree(terminals), "cudaFree linear terminal"); +} + +static void expect_device_canaries(const unsigned char* device_storage, + size_t prefix_bytes, size_t payload_bytes, size_t suffix_bytes, + unsigned char canary, const char* label) { + size_t total_bytes = prefix_bytes + payload_bytes + suffix_bytes; + std::vector host_storage(total_bytes); + check_cuda(cudaMemcpy(host_storage.data(), device_storage, total_bytes, + cudaMemcpyDeviceToHost), "copy canary storage"); + for (size_t i = 0; i < prefix_bytes; i++) { + if (host_storage[i] != canary) { + std::fprintf(stderr, + "%s prefix canary overwritten at byte %zu: 0x%02x != 0x%02x\n", + label, i, host_storage[i], canary); + std::exit(1); + } + } + size_t suffix_start = prefix_bytes + payload_bytes; + for (size_t i = suffix_start; i < total_bytes; i++) { + if (host_storage[i] != canary) { + std::fprintf(stderr, + "%s suffix canary overwritten at byte %zu: 0x%02x != 0x%02x\n", + label, i - suffix_start, host_storage[i], canary); + std::exit(1); + } + } +} + +static void run_io_canary_case(int n) { + constexpr size_t guard_bytes = 64; + constexpr size_t observation_prefix = guard_bytes + sizeof(obs_t); + constexpr unsigned char canary = 0xa5u; + size_t observation_bytes = (size_t)n * OBS_SIZE * sizeof(obs_t); + size_t scalar_bytes = (size_t)n * sizeof(float); + size_t observation_storage_bytes = observation_prefix + + observation_bytes + guard_bytes; + size_t scalar_storage_bytes = guard_bytes + scalar_bytes + guard_bytes; + + unsigned char* observation_storage = nullptr; + unsigned char* reward_storage = nullptr; + unsigned char* terminal_storage = nullptr; + float* actions = nullptr; + check_cuda(cudaMalloc(&observation_storage, observation_storage_bytes), + "cudaMalloc guarded observations"); + check_cuda(cudaMalloc(&reward_storage, scalar_storage_bytes), + "cudaMalloc guarded rewards"); + check_cuda(cudaMalloc(&terminal_storage, scalar_storage_bytes), + "cudaMalloc guarded terminals"); + check_cuda(cudaMalloc(&actions, scalar_bytes), + "cudaMalloc guarded actions"); + + obs_t* observations = reinterpret_cast( + observation_storage + observation_prefix); + float* rewards = reinterpret_cast(reward_storage + guard_bytes); + float* terminals = reinterpret_cast(terminal_storage + guard_bytes); + EXPECT_EQ((uintptr_t)observations & 1u, 0u); + EXPECT_EQ((uintptr_t)observations & 3u, 2u); + EXPECT_EQ((uintptr_t)rewards & 3u, 0u); + EXPECT_EQ((uintptr_t)terminals & 3u, 0u); + + check_cuda(cudaMemset(observation_storage, canary, + observation_storage_bytes), "initialize observation canaries"); + check_cuda(cudaMemset(reward_storage, canary, scalar_storage_bytes), + "initialize reward canaries"); + check_cuda(cudaMemset(terminal_storage, canary, scalar_storage_bytes), + "initialize terminal canaries"); + check_cuda(cudaMemset(actions, 0, scalar_bytes), + "initialize guarded actions"); + + Dict kwargs; + fill_kwargs(&kwargs, 31 + n, 1, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "guarded reset"); + expect_device_canaries(observation_storage, observation_prefix, + observation_bytes, guard_bytes, canary, "reset observations"); + expect_device_canaries(reward_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "reset rewards"); + expect_device_canaries(terminal_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "reset terminals"); + + check_cuda(cudaMemset(observation_storage, canary, + observation_storage_bytes), "reinitialize observation canaries"); + check_cuda(cudaMemset(reward_storage, canary, scalar_storage_bytes), + "reinitialize reward canaries"); + check_cuda(cudaMemset(terminal_storage, canary, scalar_storage_bytes), + "reinitialize terminal canaries"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "guarded step"); + expect_device_canaries(observation_storage, observation_prefix, + observation_bytes, guard_bytes, canary, "step observations"); + expect_device_canaries(reward_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "step rewards"); + expect_device_canaries(terminal_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "step terminals"); + + puf_close(envs); + dict_clear(&kwargs); + check_cuda(cudaFree(observation_storage), + "cudaFree guarded observations"); + check_cuda(cudaFree(reward_storage), "cudaFree guarded rewards"); + check_cuda(cudaFree(terminal_storage), "cudaFree guarded terminals"); + check_cuda(cudaFree(actions), "cudaFree guarded actions"); +} + +static void test_io_canaries_and_observation_alignment() { +#if AFFINE_LOCK_GPU_SHARED_OBS + constexpr int environments_per_block = AFFINE_LOCK_GPU_SHARED_BLOCK; +#else + constexpr int environments_per_block = + AFFINE_LOCK_GPU_BLOCK / AFFINE_LOCK_GPU_LANES; +#endif + run_io_canary_case(environments_per_block); + run_io_canary_case(environments_per_block + 1); +} + +static void test_nondefault_stream_and_cuda_graph() { + constexpr int n = 4099; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMalloc graph observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), + "cudaMalloc graph actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), + "cudaMalloc graph rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), + "cudaMalloc graph terminals"); + check_cuda(cudaMemset(actions, 0, (size_t)n * sizeof(float)), + "clear graph actions"); + + Dict kwargs; + fill_kwargs(&kwargs, 123, 3, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + cudaStream_t stream; + check_cuda(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), + "create graph stream"); + puf_bind_stream(stream); + puf_reset(envs); + check_cuda(cudaStreamSynchronize(stream), "graph reset"); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + std::vector oracle(n); + unsigned int running_seed = 123; + for (int i = 0; i < n; i++) { + oracle[i].rng = (uint32_t)rand_r(&running_seed); + oracle[i].curriculum_depth = 2; + oracle_reset_state(&oracle[i], &table, 3); + } + + GpuAffineLockState* actual_states = nullptr; + Env* actual_envs = nullptr; + obs_t* actual_obs = nullptr; + float* actual_rewards = nullptr; + float* actual_terminals = nullptr; + check_cuda(cudaMallocHost((void**)&actual_states, + (size_t)n * sizeof(GpuAffineLockState)), + "cudaMallocHost graph states"); + check_cuda(cudaMallocHost((void**)&actual_envs, + (size_t)n * sizeof(Env)), "cudaMallocHost graph envs"); + check_cuda(cudaMallocHost((void**)&actual_obs, + (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMallocHost graph observations"); + check_cuda(cudaMallocHost((void**)&actual_rewards, + (size_t)n * sizeof(float)), "cudaMallocHost graph rewards"); + check_cuda(cudaMallocHost((void**)&actual_terminals, + (size_t)n * sizeof(float)), "cudaMallocHost graph terminals"); + + cudaGraph_t graph; + cudaGraphExec_t graph_exec; + check_cuda(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal), + "begin graph capture"); + for (int i = 0; i < 3; i++) { + puf_step(envs); + } + check_cuda(cudaStreamEndCapture(stream, &graph), "end graph capture"); + check_cuda(cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0), + "instantiate graph"); + check_cuda(cudaGraphLaunch(graph_exec, stream), "launch graph first"); + check_cuda(cudaGraphLaunch(graph_exec, stream), "launch graph second"); + + check_cuda(cudaMemcpyAsync(actual_states, g_gpu.states, + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost, stream), + "queue graph states D2H"); + check_cuda(cudaMemcpyAsync(actual_envs, envs, + (size_t)n * sizeof(Env), cudaMemcpyDeviceToHost, stream), + "queue graph logs D2H"); + check_cuda(cudaMemcpyAsync(actual_obs, observations, + (size_t)n * OBS_SIZE * sizeof(obs_t), cudaMemcpyDeviceToHost, stream), + "queue graph observations D2H"); + check_cuda(cudaMemcpyAsync(actual_rewards, rewards, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost, stream), + "queue graph rewards D2H"); + check_cuda(cudaMemcpyAsync(actual_terminals, terminals, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost, stream), + "queue graph terminals D2H"); + check_cuda(cudaStreamSynchronize(stream), "synchronize graph and D2H"); + + std::vector expected_rewards(n), expected_terminals(n); + for (int step = 0; step < 6; step++) { + for (int i = 0; i < n; i++) { + oracle_step(&oracle[i], 0.0f, &table, + 2, 16, 3, PERF_WEIGHTING_LINEAR, + &expected_rewards[i], &expected_terminals[i]); + } + } + for (int i = 0; i < n; i++) { + expect_state_equal(actual_states[i], oracle[i]); + expect_log_equal(actual_envs[i].log, oracle[i].log); + expect_observation_equal( + &actual_obs[(size_t)i * OBS_SIZE], oracle[i]); + EXPECT_EQ(float_bits(actual_rewards[i]), + float_bits(expected_rewards[i])); + EXPECT_EQ(float_bits(actual_terminals[i]), + float_bits(expected_terminals[i])); + } + + check_cuda(cudaGraphExecDestroy(graph_exec), "destroy graph exec"); + check_cuda(cudaGraphDestroy(graph), "destroy graph"); + check_cuda(cudaFreeHost(actual_states), "cudaFreeHost graph states"); + check_cuda(cudaFreeHost(actual_envs), "cudaFreeHost graph envs"); + check_cuda(cudaFreeHost(actual_obs), "cudaFreeHost graph observations"); + check_cuda(cudaFreeHost(actual_rewards), "cudaFreeHost graph rewards"); + check_cuda(cudaFreeHost(actual_terminals), "cudaFreeHost graph terminals"); + check_cuda(cudaStreamDestroy(stream), "destroy graph stream"); + puf_bind_stream(nullptr); + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree graph observations"); + check_cuda(cudaFree(actions), "cudaFree graph actions"); + check_cuda(cudaFree(rewards), "cudaFree graph rewards"); + check_cuda(cudaFree(terminals), "cudaFree graph terminals"); +} + +int main() { + test_deterministic_reset_and_step_parity(); + test_reset_rejection_sampling(); + test_exhaustive_action_transforms(); + test_action_boundaries(); + test_solution_curriculum_and_logs(); + test_linear_solve_scoring(); + test_io_canaries_and_observation_alignment(); + test_nondefault_stream_and_cuda_graph(); + test_puf_log_exports_cpu_contract(); + std::puts("affine_lock CUDA tests passed"); + return 0; +}