blastar post-code-review. complete.

xinpw8 · Jan 26, 2025 · b08cf8c · b08cf8c
1 parent 598d741
commit b08cf8c
Show file tree

Hide file tree

Showing 6 changed files with 155 additions and 303 deletions.
diff --git a/config/ocean/blastar.ini b/config/ocean/blastar.ini
@@ -11,7 +11,7 @@ num_envs = 4096
 anneal_lr = False
 batch_size = 131072 # 65536
 bptt_horizon = 8
-checkpoint_interval = 600
+checkpoint_interval = 200
 clip_coef = 0.2
 clip_vloss = True
 compile = False

diff --git a/pufferlib/ocean/blastar/blastar.c b/pufferlib/ocean/blastar/blastar.c
@@ -8,70 +8,50 @@
 #include "puffernet.h"
 
 const char* WEIGHTS_PATH = "/home/daa/pufferlib_testbench/PufferLib/pufferlib/resources/blastar/blastar_weights.bin";
-#define OBSERVATIONS_SIZE 20
+#define OBSERVATIONS_SIZE 10
 #define ACTIONS_SIZE 6
-#define NUM_WEIGHTS 135687
+#define NUM_WEIGHTS 134407
 
 void get_input(BlastarEnv* env) {
-    // Left
     if (IsKeyDown(KEY_LEFT) || IsKeyDown(KEY_A)) {
-        env->actions[0] = 1;
-    }
-    // Right
-    else if (IsKeyDown(KEY_RIGHT) || IsKeyDown(KEY_D)) {
-        env->actions[0] = 2;
-    }
-    // Up
-    else if (IsKeyDown(KEY_UP) || IsKeyDown(KEY_W)) {
-        env->actions[0] = 3;
-    }
-    // Down
-    else if (IsKeyDown(KEY_DOWN) || IsKeyDown(KEY_S)) {
-        env->actions[0] = 4;
-    }
-    // Fire
-    else if (IsKeyDown(KEY_SPACE)) {
-        env->actions[0] = 5;
-    }
-    // Noop
-    else {
-        env->actions[0] = 0;
+        env->actions[0] = 1;  // Left
+    } else if (IsKeyDown(KEY_RIGHT) || IsKeyDown(KEY_D)) {
+        env->actions[0] = 2;  // Right
+    } else if (IsKeyDown(KEY_UP) || IsKeyDown(KEY_W)) {
+        env->actions[0] = 3;  // Up
+    } else if (IsKeyDown(KEY_DOWN) || IsKeyDown(KEY_S)) {
+        env->actions[0] = 4;  // Down
+    } else if (IsKeyDown(KEY_SPACE)) {
+        env->actions[0] = 5;  // Fire
+    } else {
+        env->actions[0] = 0;  // Noop
     }
 }
 
 int demo() {
     Weights* weights = load_weights(WEIGHTS_PATH, NUM_WEIGHTS);
     LinearLSTM* net = make_linearlstm(weights, 1, OBSERVATIONS_SIZE, ACTIONS_SIZE);
-
     BlastarEnv env = {
-        .player.x = SCREEN_WIDTH / 2,
-        .player.y = SCREEN_HEIGHT - PLAYER_HEIGHT,
         .num_obs = OBSERVATIONS_SIZE,
     };
-    allocate(&env);
-
+    allocate(&env, env.num_obs);
     Client* client = make_client(&env);
-
     unsigned int seed = 12345;
     srand(seed);
     c_reset(&env);
-
     int running = 1;
     while (running) {
         if (IsKeyDown(KEY_LEFT_SHIFT)) {
             get_input(&env);  // Human input
         } else {
             forward_linearlstm(net, env.observations, env.actions);  // AI input
         }
-
         c_step(&env);
         c_render(client, &env);
-
         if (WindowShouldClose() || env.game_over) {
             running = 0;
         }
     }
-
     free_linearlstm(net);
     free(weights);
     close_client(client);
@@ -80,29 +60,27 @@ int demo() {
 }
 
 void perftest(float test_time) {
-    BlastarEnv env;
-    allocate(&env);
-
+    BlastarEnv env = {
+        .num_obs = OBSERVATIONS_SIZE,
+    };
+    allocate(&env, env.num_obs);
     unsigned int seed = 12345;
     srand(seed);
     c_reset(&env);
-
     int start = time(NULL);
     int steps = 0;
     while (time(NULL) - start < test_time) {
         env.actions[0] = rand() % ACTIONS_SIZE;  // Random actions
         c_step(&env);
         steps++;
     }
-
     int end = time(NULL);
     printf("Steps per second: %f\n", steps / (float)(end - start));
-
     free_allocated(&env);
 }
 
 int main() {
     demo();
-    // perftest(30.0f);
+    // perftest(10.0f);
     return 0;
 }