diff --git a/pufferlib/config/ocean/g2048.ini b/pufferlib/config/ocean/g2048.ini index 688cf0b8b2..cd0f00e80a 100644 --- a/pufferlib/config/ocean/g2048.ini +++ b/pufferlib/config/ocean/g2048.ini @@ -19,13 +19,7 @@ num_envs = 4 [env] num_envs = 4096 num_agents = 1 -can_go_over_65536 = False -reward_scaler = 0.67 -endgame_env_prob = 0.05 scaffolding_ratio = 0.67 -use_heuristic_rewards = True -snake_reward_weight = 0.0005 -use_sparse_reward = False [sweep.policy.hidden_size] distribution = uniform_pow2 @@ -62,7 +56,6 @@ max = 2 mean = 1 scale = auto - [train] # 512 hidden: https://wandb.ai/kywch/pufferlib/runs/5thsjr61?nw=nwuserkywch total_timesteps = 6_767_676_767 diff --git a/pufferlib/ocean/g2048/binding.c b/pufferlib/ocean/g2048/binding.c index 679703e220..e9c6952b89 100644 --- a/pufferlib/ocean/g2048/binding.c +++ b/pufferlib/ocean/g2048/binding.c @@ -4,13 +4,7 @@ #include "../env_binding.h" static int my_init(Env* env, PyObject* args, PyObject* kwargs) { - env->can_go_over_65536 = unpack(kwargs, "can_go_over_65536"); - env->reward_scaler = unpack(kwargs, "reward_scaler"); - env->endgame_env_prob = unpack(kwargs, "endgame_env_prob"); env->scaffolding_ratio = unpack(kwargs, "scaffolding_ratio"); - env->use_heuristic_rewards = unpack(kwargs, "use_heuristic_rewards"); - env->snake_reward_weight = unpack(kwargs, "snake_reward_weight"); - env->use_sparse_reward = unpack(kwargs, "use_sparse_reward"); init(env); return 0; } @@ -24,8 +18,6 @@ static int my_log(PyObject* dict, Log* log) { assign_to_dict(dict, "lifetime_max_tile", log->lifetime_max_tile); assign_to_dict(dict, "reached_32768", log->reached_32768); assign_to_dict(dict, "reached_65536", log->reached_65536); - assign_to_dict(dict, "monotonicity_reward", log->monotonicity_reward); - assign_to_dict(dict, "snake_state", log->snake_state); - assign_to_dict(dict, "snake_reward", log->snake_reward); + assign_to_dict(dict, "reached_131072", log->reached_131072); return 0; } \ No newline at end of file diff --git a/pufferlib/ocean/g2048/eval.py b/pufferlib/ocean/g2048/eval.py index 138c272c75..29a4355960 100644 --- a/pufferlib/ocean/g2048/eval.py +++ b/pufferlib/ocean/g2048/eval.py @@ -6,9 +6,7 @@ def evaluate(env_name, load_model_path): args['env']['num_envs'] = 4096 args['load_model_path'] = load_model_path # Turn off endgame_envs and scaffolding episodes, which do not report results - args['env']['endgame_env_prob'] = 0 args['env']['scaffolding_ratio'] = 0 - args['env']['can_go_over_65536'] = True vecenv = pufferl.load_env(env_name, args) policy = pufferl.load_policy(args, vecenv, env_name) @@ -33,6 +31,7 @@ def evaluate(env_name, load_model_path): merge_scores = sum(n * s for n, s in zip(stats['n'], stats['merge_score'])) / num_episodes reached_32768 = sum(n * s for n, s in zip(stats['n'], stats['reached_32768'])) / num_episodes reached_65536 = sum(n * s for n, s in zip(stats['n'], stats['reached_65536'])) / num_episodes + reached_131072 = sum(n * s for n, s in zip(stats['n'], stats['reached_131072'])) / num_episodes print(f"Num episodes: {int(num_episodes)}") print(f"Max tile avg: {max_tiles:.1f}") @@ -41,24 +40,9 @@ def evaluate(env_name, load_model_path): print(f"Merge score -- Avg: {merge_scores:.1f}, Max: {max(stats['merge_score']):.1f}") print(f"Reached 32768 prob: {reached_32768*100:.2f} %") print(f"Reached 65536 prob: {reached_65536*100:.2f} %") + print(f"Reached 131072 prob: {reached_131072*100:.2f} %") """ - # hidden 256: https://wandb.ai/kywch/pufferlib/runs/nvd0pfuj?nw=nwuserkywch - Num episodes: 154406 - Max tile avg: 22532.9 - Episode length -- Avg: 16667.2, Max: 26659.1 - Merge score -- Avg: 462797.9, Max: 744224.9 - Reached 32768 prob: 46.08 % - Reached 65536 prob: 3.53 % - - # hidden 512: https://wandb.ai/kywch/pufferlib/runs/2ch3my60?nw=nwuserkywch - Num episodes: 119243 - Max tile avg: 30662.2 - Episode length -- Avg: 21539.7, Max: 29680.3 - Merge score -- Avg: 618011.8, Max: 918755.8 - Reached 32768 prob: 68.25 % - Reached 65536 prob: 13.09 % - # hidden 512 (replication): https://wandb.ai/kywch/pufferlib/runs/5thsjr61?nw=nwuserkywch Num episodes: 115652 Max tile avg: 31773.2 @@ -66,18 +50,30 @@ def evaluate(env_name, load_model_path): Merge score -- Avg: 639395.6, Max: 909969.8 Reached 32768 prob: 71.22 % Reached 65536 prob: 14.75 % + + # embeddings: https://wandb.ai/thatguy11325/pufferlib/runs/g2f00pcm?nw=nwuserthatguy11325 + Num episodes: 192276 + Max tile avg: 33166.4 + Episode length -- Avg: 26950.7, Max: 44906.1 + Merge score -- Avg: 770645.8, Max: 1040367.2 + Reached 32768 prob: 85.32 % + Reached 65536 prob: 10.15 % + + # embeddings + new reward: https://wandb.ai/kywch/pufferlib/runs/1v5kls7l?nw=nwuserkywch + Num episodes: 95611 + Max tile avg: 40980.9 + Episode length -- Avg: 26792.1, Max: 37442.2 + Merge score -- Avg: 779238.6, Max: 997571.8 + Reached 32768 prob: 84.88 % + Reached 65536 prob: 33.96 % + Reached 131072 prob: 0.00 % """ def finetune(env_name, load_model_path): args = pufferl.load_config(env_name) args['load_model_path'] = load_model_path - # args['env']['use_sparse_reward'] = True args['env']['scaffolding_ratio'] = 0.85 - # args['policy']['hidden_size'] = 512 - # args['rnn']['input_size'] = 512 - # args['rnn']['hidden_size'] = 512 - args['train']['total_timesteps'] = 1_000_000_000 args['train']['learning_rate'] = 0.00005 args['train']['anneal_lr'] = False @@ -90,12 +86,15 @@ def finetune(env_name, load_model_path): if __name__ == '__main__': import os import wandb + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('--entity', type=str, default='kywch') + parser.add_argument('--run-id', type=str, default='1v5kls7l') - # https://wandb.ai/kywch/pufferlib/runs/5thsjr61?nw=nwuserkywch - wandb_run_id = '5thsjr61' - wandb.init(id=wandb_run_id, project='pufferlib', entity='kywch') + args = parser.parse_args() - artifact = wandb.use_artifact(f'{wandb_run_id}:latest') + wandb.init(id=args.run_id, project='pufferlib', entity=args.entity) + artifact = wandb.use_artifact(f'{args.run_id}:latest') data_dir = artifact.download() model_file = max(os.listdir(data_dir)) model_path = f'{data_dir}/{model_file}' diff --git a/pufferlib/ocean/g2048/g2048.c b/pufferlib/ocean/g2048/g2048.c index 1cd1cf89e0..2929de4382 100644 --- a/pufferlib/ocean/g2048/g2048.c +++ b/pufferlib/ocean/g2048/g2048.c @@ -1,7 +1,7 @@ #include "g2048.h" #include "g2048_net.h" -#define OBS_DIM 289 +#define OBS_DIM 16 #define HIDDEN_DIM 512 // Set NO_RENDER to true to run evals without the render @@ -11,12 +11,7 @@ int main() { srand(time(NULL)); Game env = { - .can_go_over_65536 = true, - .reward_scaler = 0.0, - .endgame_env_prob = 0.0, .scaffolding_ratio = 0.0, - .use_heuristic_rewards = false, - .snake_reward_weight = 0.0, }; init(&env); @@ -30,8 +25,8 @@ int main() { env.actions = actions; env.rewards = rewards; - Weights* weights = load_weights("resources/g2048/g2048_weights.bin", 3713541); - G2048Net* net = make_g2048net(weights, OBS_DIM, HIDDEN_DIM); + Weights* weights = load_weights("resources/g2048/g2048_weights.bin", 3466859); + G2048Net* net = make_g2048net(weights, HIDDEN_DIM); c_reset(&env); if (!NO_RENDER) c_render(&env); printf("Starting...\n"); diff --git a/pufferlib/ocean/g2048/g2048.h b/pufferlib/ocean/g2048/g2048.h index 490b2f8e8f..2592d5b2d5 100644 --- a/pufferlib/ocean/g2048/g2048.h +++ b/pufferlib/ocean/g2048/g2048.h @@ -17,24 +17,21 @@ static inline int max(int a, int b) { return a > b ? a : b; } #define RIGHT 4 #define BASE_MAX_TICKS 1000 -// These work well -#define MERGE_REWARD_WEIGHT 0.0625f +// Reward constants +#define MERGE_BASE_REWARD 0.05f +#define MERGE_REWARD_SCALE 0.03f #define INVALID_MOVE_PENALTY -0.05f #define GAME_OVER_PENALTY -1.0f -// These may need experimenting, but work for now -#define STATE_REWARD_WEIGHT 0.01f // Fixed, small reward for maintaining "desirable" states -#define MONOTONICITY_REWARD_WEIGHT 0.00003f - -// Features: 18 per cell -// 1. Normalized tile value (current_val / max_val) -// 2. One-hot for empty (1 if empty, 0 if occupied) -// 3-18. One-hot for tile values 2^1 to 2^16 (16 features) -#define NUM_FEATURES 18 +// Pow 1.5 lookup table for tiles 128+ (index = row[i] - 6) +// Index: 1=128, 2=256, 3=512, 4=1024, 5=2048, 6=4096, 7=8192, 8=16384, 9=32768, 10=65536, 11=131k +static const float pow15_table[12] = { + 0.0f, 1.0f, 2.83f, 5.20f, 8.0f, 11.18f, 14.70f, 18.52f, 22.63f, 27.0f, 31.62f, 36.48f, +}; static inline float calculate_perf(unsigned char max_tile) { - // Reaching 65k -> 1.0, 32k -> 0.8, 16k -> 0.4, 8k -> 0.2, 4k -> 0.1, 2k -> 0.05 - float perf = 0.8f * (float)(1 << max_tile) / 32768.0f; + // Reaching 131k -> 1.0, 65k -> 0.8, 32k -> 0.4, 16k -> 0.2, 8k -> 0.1 + float perf = 0.8f * (float)(1 << max_tile) / 65536.0f; if (perf > 1.0f) perf = 1.0f; return perf; } @@ -48,9 +45,7 @@ typedef struct Log { float lifetime_max_tile; float reached_32768; float reached_65536; - float snake_state; - float monotonicity_reward; - float snake_reward; + float reached_131072; float n; } Log; @@ -62,16 +57,8 @@ typedef struct Game { float* terminals; // Required int num_agents; // Required for env_binding - bool can_go_over_65536; // Set false for training, true for eval - float reward_scaler; // Pufferlib clips rew from -1 to 1, adjust the resulting rew accordingly - - float endgame_env_prob; // The prob of env being initialized as an endgame-only env - bool is_endgame_env; float scaffolding_ratio; // The ratio for "scaffolding" runs, in which higher blocks are spawned bool is_scaffolding_episode; - bool use_heuristic_rewards; - float snake_reward_weight; - bool use_sparse_reward; // Ignore all rewards and provide 1 for reaching 16k, 32k, 65k int score; int tick; @@ -79,13 +66,8 @@ typedef struct Game { unsigned char lifetime_max_tile; unsigned char max_tile; // Episode max tile float episode_reward; // Accumulate episode reward - float monotonicity_reward; - float snake_reward; int moves_made; int max_episode_ticks; // Dynamic max_ticks based on score - bool is_snake_state; - int snake_state_tick; - bool stop_at_65536; // Cached values to avoid recomputation int empty_count; @@ -119,11 +101,6 @@ static Color tile_colors[17] = { {6, 24, 24, 255}, // 65536+ (Invisible) }; -// Precomputed pow(x, 1.5) lookup table for x in [0, 19] to avoid expensive pow() calls. -static const unsigned char pow_1_5_lookup[20] = { - 0, 1, 2, 5, 8, 11, 14, 18, 22, 27, 31, 36, 41, 46, 52, 57, 64, 69, 75, 81 -}; - // --- Logging --- void add_log(Game* game); @@ -135,48 +112,16 @@ void c_close(Game* game); void init(Game* game) { game->lifetime_max_tile = 0; - game->is_endgame_env = (rand() / (float)RAND_MAX) < game->endgame_env_prob; + memset(game->grid, EMPTY, SIZE * SIZE); } void update_observations(Game* game) { - // Observation: 4x4 grid, 18 features per cell - // 1. Normalized tile value (current_val / max_val) - // 2. One-hot for empty (1 if empty, 0 if occupied) - // 3. One-hot for tile values 2^1 to 2^16 (16 features) - // 4. Additional obs: is_snake_state (1) - - int num_cell = SIZE * SIZE; - int num_additional_obs = 1; - memset(game->observations, 0, (num_cell * NUM_FEATURES + num_additional_obs) * sizeof(unsigned char)); - for (int i = 0; i < SIZE; i++) { - for (int j = 0; j < SIZE; j++) { - int feat1_idx = (i * SIZE + j); - int feat2_idx = num_cell + feat1_idx; - int feat3_idx = 2 * num_cell + 16 * feat1_idx; - unsigned char grid_val = game->grid[i][j]; - - // Feature 1: The original tile values ** 1.5, to make a bit superlinear within uint8 - game->observations[feat1_idx] = pow_1_5_lookup[grid_val]; - - // Feature 2: One-hot for empty - game->observations[feat2_idx] = (grid_val == EMPTY) ? 1 : 0; - - // Features 3-18: One-hot for tile values - // NOTE: If this ever gets close to 131072, revisit this - if (grid_val > 0) { - grid_val = min(grid_val, 16); - game->observations[feat3_idx + grid_val - 1] = 1; - } - } - } - // Additional obs - int offset = num_cell * NUM_FEATURES; - game->observations[offset] = game->is_snake_state; + memcpy(game->observations, game->grid, SIZE * SIZE); } void add_log(Game* game) { // Scaffolding runs will distort stats, so skip logging - if (game->is_endgame_env || game->is_scaffolding_episode) return; + if (game->is_scaffolding_episode) return; // Update the lifetime best if (game->max_tile > game->lifetime_max_tile) { @@ -191,9 +136,7 @@ void add_log(Game* game) { game->log.lifetime_max_tile += (float)(1 << game->lifetime_max_tile); game->log.reached_32768 += (game->max_tile >= 15); game->log.reached_65536 += (game->max_tile >= 16); - game->log.snake_state += (float)game->snake_state_tick / (float)game->tick; - game->log.monotonicity_reward += game->monotonicity_reward * MONOTONICITY_REWARD_WEIGHT * game->reward_scaler; - game->log.snake_reward += game->snake_reward * game->snake_reward_weight * game->reward_scaler; + game->log.reached_131072 += (game->max_tile >= 17); game->log.n += 1; } @@ -222,69 +165,32 @@ static inline void place_tile_at_random_cell(Game* game, unsigned char tile) { } void set_scaffolding_curriculum(Game* game) { - game->stop_at_65536 = true; - if (game->lifetime_max_tile < 14) { + // Spawn one high tile from 8192 to 65536 int curriculum = rand() % 5; - - // Spawn one high tiles from 8192, 16384, 32768, 65536 unsigned char high_tile = max(12 + curriculum, game->lifetime_max_tile); place_tile_at_random_cell(game, high_tile); - if (high_tile >= 16) game->stop_at_65536 = false; } else { - int curriculum = rand() % 8; - - if (curriculum < 2) { // curriculum 0, 1 - place_tile_at_random_cell(game, 14 + curriculum); // Spawn one of 16384 or 32768 - + // base=14 until 65536 reached, then base=15 for 131072 practice + // All random placement, 1-2 tiles max + unsigned char base = (game->lifetime_max_tile >= 16) ? 15 : 14; + int curriculum = rand() % 4; + + if (curriculum == 0) { + place_tile_at_random_cell(game, base); + } else if (curriculum == 1) { + place_tile_at_random_cell(game, base + 1); } else if (curriculum == 2) { - // Place the tiles in the second row, so that they can be moved up in the first move - unsigned char tiles[] = {14, 13}; - memcpy(game->grid[1], tiles, 2); - game->empty_count -= 2; - } else if (curriculum == 3) { // harder - game->grid[1][0] = 14; game->empty_count--; - place_tile_at_random_cell(game, 13); - - } else if (curriculum == 4) { - unsigned char tiles[] = {15, 14}; - memcpy(game->grid[1], tiles, 2); - game->empty_count -= 2; - } else if (curriculum == 5) { // harder - game->grid[1][0] = 15; game->empty_count--; - place_tile_at_random_cell(game, 14); - - } else if (curriculum == 6) { - unsigned char tiles[] = {15, 14, 13}; - memcpy(game->grid[1], tiles, 3); - game->empty_count -= 3; - } else if (curriculum == 7) { // harder - game->grid[1][0] = 15; game->empty_count--; - place_tile_at_random_cell(game, 14); - place_tile_at_random_cell(game, 13); + place_tile_at_random_cell(game, base); + place_tile_at_random_cell(game, base - 1); + } else { + place_tile_at_random_cell(game, base + 1); + place_tile_at_random_cell(game, base); } } } -void set_endgame_curriculum(Game* game) { - game->stop_at_65536 = true; - int curriculum = rand() % 4; - - // Place the tiles in the second-third rows, so that they can be moved up in the first move - unsigned char tiles[] = {15, 14, 13, 12}; - memcpy(game->grid[1], tiles, 4); - game->empty_count -= 4; - - if (curriculum >= 1) { game->grid[2][3] = 11; game->empty_count--; } - if (curriculum >= 2) { - game->grid[2][2] = 10; - game->grid[2][1] = 9; - game->grid[2][0] = 8; - game->empty_count -= 3; - } -} - void c_reset(Game* game) { memset(game->grid, EMPTY, SIZE * SIZE); game->score = 0; @@ -296,30 +202,17 @@ void c_reset(Game* game) { game->moves_made = 0; game->max_episode_ticks = BASE_MAX_TICKS; game->max_tile = 0; - game->snake_state_tick = 0; - game->monotonicity_reward = 0; - game->snake_reward = 0; - game->is_snake_state = false; - game->stop_at_65536 = game->can_go_over_65536; - - if (game->terminals) game->terminals[0] = 0; - // End game envs only do endgame curriculum - if (game->is_endgame_env) { - set_endgame_curriculum(game); + // Higher tiles are spawned in scaffolding episodes + // Having high tiles saves moves to get there, allowing agents to experience it faster + game->is_scaffolding_episode = (rand() / (float)RAND_MAX) < game->scaffolding_ratio; + if (game->is_scaffolding_episode) { + set_scaffolding_curriculum(game); } else { - // Higher tiles are spawned in scaffolding episodes - // Having high tiles saves moves to get there, allowing agents to experience it faster - game->is_scaffolding_episode = (rand() / (float)RAND_MAX) < game->scaffolding_ratio; - if (game->is_scaffolding_episode) { - set_scaffolding_curriculum(game); - - } else { - // Add two random tiles at the start - for (int i = 0; i < 2; i++) { - place_tile_at_random_cell(game, get_new_tile()); - } + // Add two random tiles at the start + for (int i = 0; i < 2; i++) { + place_tile_at_random_cell(game, get_new_tile()); } } @@ -347,7 +240,13 @@ static inline bool slide_and_merge(Game* game, unsigned char* row, float* reward for (int i = 0; i < SIZE - 1; i++) { if (row[i] != EMPTY && row[i] == row[i + 1]) { row[i]++; - *reward += ((float)row[i]) * MERGE_REWARD_WEIGHT; + // Tiles 2-64 (row[i] 1-6): base reward only + // Tiles 128+ (row[i] 7+): base + pow1.5 scaled bonus + if (row[i] <= 6) { + *reward += MERGE_BASE_REWARD; + } else { + *reward += MERGE_BASE_REWARD + pow15_table[row[i] - 6] * MERGE_REWARD_SCALE; + } *score_increase += (float)(1 << (int)row[i]); // Shift remaining elements left for (int j = i + 1; j < SIZE - 1; j++) { @@ -411,7 +310,7 @@ bool move(Game* game, int direction, float* reward, float* score_increase) { bool is_game_over(Game* game) { // Use cached result if grid hasn't changed - if (!game->grid_changed && game->game_over_cached) { + if (!game->grid_changed) { return game->game_over_cached; } @@ -444,122 +343,28 @@ bool is_game_over(Game* game) { return true; } -// Combined grid stats and heuristic calculation for performance -float update_stats_and_get_heuristic_rewards(Game* game) { +void update_stats(Game* game) { int empty_count = 0; - int top_row_count = 0; unsigned char max_tile = 0; - unsigned char second_max_tile = 0; - unsigned char max_tile_in_row234 = 0; - float heuristic_state_reward = 0.0f; - float monotonicity_reward = 0.0f; - float snake_reward = 0.0f; - game->is_snake_state = false; for (int i = 0; i < SIZE; i++) { for (int j = 0; j < SIZE; j++) { unsigned char val = game->grid[i][j]; - // Update empty count and max tile if (val == EMPTY) empty_count++; - - // Count filled cells in the top row - if (i == 0 && val != EMPTY) top_row_count++; - - // Allow max and the second max tile to be the same - if (val >= max_tile) { - second_max_tile = max_tile; + if (val > max_tile) { max_tile = val; - } else if (val > second_max_tile && val < max_tile) { - second_max_tile = val; } - - // Get the max tile in the 2nd, 3rd, 4th row - if (i > 0 && val > max_tile_in_row234) max_tile_in_row234 = val; } } game->empty_count = empty_count; game->max_tile = max_tile; - - /* Heuristic rewards */ - - // Filled top row reward: A simple nudge to keep the top row filled - if (top_row_count == SIZE) heuristic_state_reward += STATE_REWARD_WEIGHT; - - bool max_in_top_left = (game->grid[0][0] == max_tile); - - // Corner reward: A simple nudge to keep the max tiles horizontally in the top row, left corner. - // When agents learn to put the max tile on the other corners, or put max tiles vertically - // they miss out snake rew, and this does happen sometimes. - if (max_in_top_left && game->grid[0][1] == second_max_tile && max_tile > 4) { - heuristic_state_reward += STATE_REWARD_WEIGHT; - } - - // Snake reward: look for the snake pattern, only when the max tile is at top left - if (max_in_top_left) { - monotonicity_reward += pow_1_5_lookup[max_tile]; - int evidence_for_snake = 0; - - for (int i = 0; i < 2; i++) { - unsigned char row_min = 32; - unsigned char next_row_max = 0; - for (int j = 0; j < SIZE; j++) { - unsigned char val = game->grid[i][j]; - - // Check horizontal monotonicity (snake pattern) for top two rows only - if (j < SIZE - 1) { - unsigned char next_col = game->grid[i][j+1]; - if (val != EMPTY && next_col != EMPTY) { - // Row 0: Reward decreasing left to right, e.g., 12-11-10-9 - if (i == 0 && val > next_col) { - monotonicity_reward += pow_1_5_lookup[next_col]; - evidence_for_snake++; - } - // Row 1: Reward increasing left to right, e.g., 5-6-7-8 - else if (i == 1 && val < next_col) { - monotonicity_reward += pow_1_5_lookup[val]; - } - } - } - - // Vertical monotonicity: give score after row scanning for min/max is done - if (val != EMPTY && val < row_min) row_min = val; - unsigned char next_row = game->grid[i+1][j]; - if (next_row != EMPTY && next_row > next_row_max) next_row_max = next_row; - // // Small column-level vertical reward - if (val != EMPTY && next_row != EMPTY && val > next_row) monotonicity_reward += next_row; - } - // Large row-level vertical reward - if (i < 2 && row_min < 20 && next_row_max > 0 && row_min > next_row_max) { - monotonicity_reward += 4 * pow_1_5_lookup[row_min]; - if (i == 0) evidence_for_snake++; - } - } - - // Snake bonus: sorted top row + the max_tile_in_row234 in the second row right - // For example, top row: 14-13-12-11, second row: ()-()-()-10 - unsigned char snake_tail = game->grid[1][3]; - if (evidence_for_snake >= 4 && snake_tail == max_tile_in_row234) { - game->is_snake_state = true; - game->snake_state_tick++; - snake_reward = snake_tail * snake_tail; - } - } - - // Trained models need game->is_snake_state as obs - if (!game->use_heuristic_rewards) return 0.0f; - - game->monotonicity_reward += monotonicity_reward; - game->snake_reward += snake_reward; - - return heuristic_state_reward + monotonicity_reward * MONOTONICITY_REWARD_WEIGHT + snake_reward * game->snake_reward_weight; } void c_step(Game* game) { float reward = 0.0f; float score_add = 0.0f; - unsigned char prev_max_tile = game->max_tile; bool did_move = move(game, game->actions[0] + 1, &reward, &score_add); game->tick++; @@ -568,11 +373,10 @@ void c_step(Game* game) { place_tile_at_random_cell(game, get_new_tile()); game->score += score_add; - // Add heuristic rewards/penalties and update grid stats - reward += update_stats_and_get_heuristic_rewards(game); - reward *= game->reward_scaler; + update_stats(game); - update_observations(game); // Observations only change if the grid changes + // Observations only change if the grid changes + update_observations(game); // This is to limit infinite invalid moves during eval (happens for noob agents) // Don't need to be tight. Don't need to show to human player. @@ -586,17 +390,11 @@ void c_step(Game* game) { bool game_over = is_game_over(game); bool max_ticks_reached = game->tick >= game->max_episode_ticks; - bool max_level_reached = game->stop_at_65536 && game->max_tile >= 16; - game->terminals[0] = (game_over || max_ticks_reached || max_level_reached) ? 1 : 0; + game->terminals[0] = (game_over || max_ticks_reached) ? 1 : 0; // Game over penalty overrides other rewards if (game_over) { - reward = GAME_OVER_PENALTY; - } - - if (game->use_sparse_reward) { - reward = 0; // Ignore all previous reward - if (game->max_tile >= 14 && game->max_tile > prev_max_tile) reward = 1; + reward += GAME_OVER_PENALTY; } game->rewards[0] = reward; @@ -619,8 +417,9 @@ void step_without_reset(Game* game) { game->moves_made++; place_tile_at_random_cell(game, get_new_tile()); game->score += score_add; - update_stats_and_get_heuristic_rewards(game); // The reward is ignored. - update_observations(game); // Observations only change if the grid changes + update_stats(game); + // Observations only change if the grid changes + update_observations(game); } bool game_over = is_game_over(game); diff --git a/pufferlib/ocean/g2048/g2048.py b/pufferlib/ocean/g2048/g2048.py index b1dbb0d0ea..ac8b1a85ac 100644 --- a/pufferlib/ocean/g2048/g2048.py +++ b/pufferlib/ocean/g2048/g2048.py @@ -7,37 +7,23 @@ from pufferlib.ocean.g2048 import binding class G2048(pufferlib.PufferEnv): - def __init__(self, num_envs=1, num_agents=1, reward_scaler=1.0, - can_go_over_65536=False, endgame_env_prob=0.0, scaffolding_ratio=0.0, - use_heuristic_rewards=False, snake_reward_weight=0.0, use_sparse_reward=False, + def __init__(self, num_envs=1, num_agents=1, scaffolding_ratio=0.0, render_mode=None, log_interval=128, buf=None, seed=0): self.single_observation_space = gymnasium.spaces.Box( - low=0, high=100, shape=(16*18 + 1,), dtype=np.uint8 + low=0, high=18, shape=(16,), dtype=np.uint8 ) self.single_action_space = gymnasium.spaces.Discrete(4) self.render_mode = render_mode self.num_agents = num_envs self.log_interval = log_interval - self.can_go_over_65536 = can_go_over_65536 - self.reward_scaler = reward_scaler - self.endgame_env_prob = endgame_env_prob self.scaffolding_ratio = scaffolding_ratio - self.use_heuristic_rewards = use_heuristic_rewards - self.snake_reward_weight = snake_reward_weight - self.use_sparse_reward = use_sparse_reward super().__init__(buf) self.c_envs = binding.vec_init( self.observations, self.actions, self.rewards, self.terminals, self.truncations, num_envs, seed, - can_go_over_65536 = self.can_go_over_65536, - reward_scaler = self.reward_scaler, - endgame_env_prob = self.endgame_env_prob, scaffolding_ratio = self.scaffolding_ratio, - use_heuristic_rewards = self.use_heuristic_rewards, - snake_reward_weight = self.snake_reward_weight, - use_sparse_reward = self.use_sparse_reward ) def reset(self, seed=0): @@ -86,4 +72,11 @@ def close(self): print('2048 SPS:', int(steps / (time.time() - start))) + num_dones = 0 + while num_dones < 300: + _, _, dones, _, _ = env.step(actions[i % CACHE]) + num_dones += dones.sum() + i += 1 + print('2048 gives out non-zero dones.') + env.close() diff --git a/pufferlib/ocean/g2048/g2048_net.h b/pufferlib/ocean/g2048/g2048_net.h index e3dcf9a9bd..b7f9abb220 100644 --- a/pufferlib/ocean/g2048/g2048_net.h +++ b/pufferlib/ocean/g2048/g2048_net.h @@ -2,9 +2,18 @@ typedef struct G2048Net G2048Net; +#define G2048_EMBED_DIM 3 +#define G2048_NUM_GRID_CELLS 16 +#define G2048_NUM_TILE_VALUES 18 +#define G2048_NUM_OBS (G2048_NUM_GRID_CELLS * G2048_EMBED_DIM) + struct G2048Net { int hidden_dim; - float* obs; + int* value_indices; // [16] - tile values for embedding lookup + int* pos_indices; // [16] - constant position indices 0-15 + float* embedded_obs; // [48] - working buffer for embedded observations + Embedding* value_embed; // [18, 3] - tile value embeddings + Embedding* pos_embed; // [16, 3] - position embeddings Linear* layer1; GELU* gelu1; Linear* layer2; @@ -21,32 +30,33 @@ struct G2048Net { Multidiscrete* multidiscrete; }; -G2048Net* make_g2048net(Weights* weights, int input_dim, int hidden_dim) { +G2048Net* make_g2048net(Weights* weights, int hidden_dim) { G2048Net* net = calloc(1, sizeof(G2048Net)); const int num_agents = 1; const int num_actions = 1; const int atn_sum = 4; int logit_sizes[1] = {4}; - net->obs = calloc(num_agents*input_dim, sizeof(float)); net->hidden_dim = hidden_dim; - if (hidden_dim <= 256) { - net->layer1 = make_linear(weights, num_agents, input_dim, 512); - net->gelu1 = make_gelu(num_agents, 512); - net->layer2 = make_linear(weights, num_agents, 512, 256); - net->gelu2 = make_gelu(num_agents, 256); - net->layer3 = make_linear(weights, num_agents, 256, hidden_dim); - net->gelu3 = make_gelu(num_agents, hidden_dim); - } else { - net->layer1 = make_linear(weights, num_agents, input_dim, 2*hidden_dim); - net->gelu1 = make_gelu(num_agents, 2*hidden_dim); - net->layer2 = make_linear(weights, num_agents, 2*hidden_dim, hidden_dim); - net->gelu2 = make_gelu(num_agents, hidden_dim); - net->layer3 = make_linear(weights, num_agents, hidden_dim, hidden_dim); - net->gelu3 = make_gelu(num_agents, hidden_dim); + net->value_indices = calloc(G2048_NUM_GRID_CELLS, sizeof(int)); + net->pos_indices = calloc(G2048_NUM_GRID_CELLS, sizeof(int)); + net->embedded_obs = calloc(G2048_NUM_OBS, sizeof(float)); + + for (int i = 0; i < G2048_NUM_GRID_CELLS; i++) { + net->pos_indices[i] = i; } + net->value_embed = make_embedding(weights, G2048_NUM_GRID_CELLS, G2048_NUM_TILE_VALUES, G2048_EMBED_DIM); + net->pos_embed = make_embedding(weights, G2048_NUM_GRID_CELLS, G2048_NUM_GRID_CELLS, G2048_EMBED_DIM); + + net->layer1 = make_linear(weights, num_agents, G2048_NUM_OBS, 2 * hidden_dim); + net->gelu1 = make_gelu(num_agents, 2 * hidden_dim); + net->layer2 = make_linear(weights, num_agents, 2 * hidden_dim, hidden_dim); + net->gelu2 = make_gelu(num_agents, hidden_dim); + net->layer3 = make_linear(weights, num_agents, hidden_dim, hidden_dim); + net->gelu3 = make_gelu(num_agents, hidden_dim); + net->actor_hidden = make_linear(weights, num_agents, hidden_dim, hidden_dim); net->gelu_actor = make_gelu(num_agents, hidden_dim); net->actor_head = make_linear(weights, num_agents, hidden_dim, atn_sum); @@ -61,33 +71,50 @@ G2048Net* make_g2048net(Weights* weights, int input_dim, int hidden_dim) { } void free_g2048net(G2048Net* net) { - free(net->obs); + free(net->value_indices); + free(net->pos_indices); + free(net->embedded_obs); + free(net->value_embed); + free(net->pos_embed); free(net->layer1); free(net->gelu1); free(net->layer2); free(net->gelu2); free(net->layer3); free(net->gelu3); - free(net->actor_hidden); free(net->gelu_actor); free(net->actor_head); free(net->value_hidden); free(net->gelu_value); free(net->value_head); - free(net->lstm); free(net->multidiscrete); free(net); } void forward_g2048net(G2048Net* net, unsigned char* observations, int* actions) { - for (int i = 0; i < net->layer1->input_dim; i++) { - net->obs[i] = (float)observations[i]; - if (i < 16) net->obs[i] /= 100.0f; + // Convert observations to value indices for embedding lookup + for (int i = 0; i < G2048_NUM_GRID_CELLS; i++) { + net->value_indices[i] = (int)observations[i]; + } + + // Embed tile values: value_embed(observations) -> [16, 3] + embedding(net->value_embed, net->value_indices); + + // Embed positions: pos_embed([0,1,...,15]) -> [16, 3] + embedding(net->pos_embed, net->pos_indices); + + // Add value and position embeddings, flatten to [48] + // PyTorch: grid_obs = (value_obs + pos_obs).flatten(1) + for (int i = 0; i < G2048_NUM_GRID_CELLS; i++) { + for (int j = 0; j < G2048_EMBED_DIM; j++) { + int idx = i * G2048_EMBED_DIM + j; + net->embedded_obs[idx] = net->value_embed->output[idx] + net->pos_embed->output[idx]; + } } - linear(net->layer1, net->obs); + linear(net->layer1, net->embedded_obs); gelu(net->gelu1, net->layer1->output); linear(net->layer2, net->gelu1->output); gelu(net->gelu2, net->layer2->output); diff --git a/pufferlib/ocean/torch.py b/pufferlib/ocean/torch.py index 4f263dfd8c..23cfcb5d22 100644 --- a/pufferlib/ocean/torch.py +++ b/pufferlib/ocean/torch.py @@ -931,26 +931,22 @@ def __init__(self, env, hidden_size=128): self.hidden_size = hidden_size self.is_continuous = False - num_obs = np.prod(env.single_observation_space.shape) + self.embed_dim = int(np.ceil(33**0.25)) + self.num_grid_cell = 4*4 + self.num_obs = self.num_grid_cell * self.embed_dim - if hidden_size <= 256: - self.encoder = torch.nn.Sequential( - pufferlib.pytorch.layer_init(nn.Linear(num_obs, 512)), - nn.GELU(), - pufferlib.pytorch.layer_init(nn.Linear(512, 256)), - nn.GELU(), - pufferlib.pytorch.layer_init(nn.Linear(256, hidden_size)), - nn.GELU(), - ) - else: - self.encoder = torch.nn.Sequential( - pufferlib.pytorch.layer_init(nn.Linear(num_obs, 2*hidden_size)), - nn.GELU(), - pufferlib.pytorch.layer_init(nn.Linear(2*hidden_size, hidden_size)), - nn.GELU(), - pufferlib.pytorch.layer_init(nn.Linear(hidden_size, hidden_size)), - nn.GELU(), - ) + self.value_embed = torch.nn.Embedding(18, self.embed_dim) + self.pos_embed = torch.nn.Embedding(self.num_grid_cell, self.embed_dim) + + self.encoder = torch.nn.Sequential( + torch.nn.Flatten(), + pufferlib.pytorch.layer_init(nn.Linear(self.num_obs, 2 * hidden_size)), + nn.GELU(), + pufferlib.pytorch.layer_init(nn.Linear(2 * hidden_size, hidden_size)), + nn.GELU(), + pufferlib.pytorch.layer_init(nn.Linear(hidden_size, hidden_size)), + nn.GELU(), + ) num_atns = env.single_action_space.n self.decoder = torch.nn.Sequential( @@ -973,13 +969,10 @@ def forward(self, observations, state=None): return self.forward_eval(observations, state) def encode_observations(self, observations, state=None): - batch_size = observations.shape[0] - observations = observations.view(batch_size, -1).float() - - # Scale the feat 1 (tile**1.5) - observations[:, :16] = observations[:, :16] / 100.0 - - return self.encoder(observations) + value_obs = self.value_embed(observations.long()) + pos_obs = self.pos_embed.weight.expand(*value_obs.shape) + grid_obs = (value_obs + pos_obs).flatten(1) + return self.encoder(grid_obs) def decode_actions(self, hidden): logits = self.decoder(hidden) diff --git a/pufferlib/resources/g2048/g2048_weights.bin b/pufferlib/resources/g2048/g2048_weights.bin index 3ebb3e38a3..30cc3b7bf4 100644 Binary files a/pufferlib/resources/g2048/g2048_weights.bin and b/pufferlib/resources/g2048/g2048_weights.bin differ