Merge pull request #397 from kywch/tetris

Add levels, garbage lines, 7-bag system to tetris

Author: jsuarez5341

pufferlib/config/ocean/tetris.ini

@@ -8,11 +8,24 @@ rnn_name = Recurrent
 num_envs = 8
 
 [env]
-num_envs = 1024
-deck_size = 3
+num_envs = 2048
+n_rows = 20
+n_cols = 10
+use_deck_obs = True
+n_init_garbage = 4
+# This is experimental. Sometimes it works.
+n_noise_obs = 0
+
+[policy]
+hidden_size = 256
+
+[rnn]
+input_size = 256
+hidden_size = 256
 
 [train]
-total_timesteps = 5_000_000_000
+# https://wandb.ai/kywch/pufferlib/runs/era6a8p6?nw=nwuserkywch
+total_timesteps = 3_000_000_000
 batch_size = auto
 bptt_horizon = 64
 minibatch_size = 65536

pufferlib/ocean/tetris/binding.c

@@ -6,7 +6,9 @@
 static int my_init(Env* env, PyObject* args, PyObject* kwargs) {
     env->n_rows = unpack(kwargs, "n_rows");
     env->n_cols = unpack(kwargs, "n_cols");
-    env->deck_size = unpack(kwargs, "deck_size");
+    env->use_deck_obs = unpack(kwargs, "use_deck_obs");
+    env->n_noise_obs = unpack(kwargs, "n_noise_obs");
+    env->n_init_garbage = unpack(kwargs, "n_init_garbage");
     init(env);
     return 0;
 }
@@ -18,9 +20,13 @@ static int my_log(PyObject* dict, Log* log) {
     assign_to_dict(dict, "ep_return", log->ep_return);
     assign_to_dict(dict, "avg_combo", log->avg_combo);
     assign_to_dict(dict, "lines_deleted", log->lines_deleted);
+    assign_to_dict(dict, "game_level", log->game_level);
+    assign_to_dict(dict, "ticks_per_line", log->ticks_per_line);
 
-    assign_to_dict(dict, "atn_frac_soft_drop", log->atn_frac_soft_drop);
+    // assign_to_dict(dict, "atn_frac_soft_drop", log->atn_frac_soft_drop);
     assign_to_dict(dict, "atn_frac_hard_drop", log->atn_frac_hard_drop);
     assign_to_dict(dict, "atn_frac_rotate", log->atn_frac_rotate);
+    assign_to_dict(dict, "atn_frac_hold", log->atn_frac_hold);
+
     return 0;
 }

pufferlib/ocean/tetris/tetris.c

@@ -4,48 +4,95 @@
 #define min(a, b) (((a) < (b)) ? (a) : (b))
 #define max(a, b) (((a) > (b)) ? (a) : (b))
 
+// Network with hidden size 256. Should go to puffernet
+LinearLSTM* make_linearlstm_256(Weights* weights, int num_agents, int input_dim, int logit_sizes[], int num_actions) {
+    LinearLSTM* net = calloc(1, sizeof(LinearLSTM));
+    net->num_agents = num_agents;
+    net->obs = calloc(num_agents*input_dim, sizeof(float));
+    int hidden_dim = 256;
+    net->encoder = make_linear(weights, num_agents, input_dim, hidden_dim);
+    net->gelu1 = make_gelu(num_agents, hidden_dim);
+    int atn_sum = 0;
+    for (int i = 0; i < num_actions; i++) {
+        atn_sum += logit_sizes[i];
+    }
+    net->actor = make_linear(weights, num_agents, hidden_dim, atn_sum);
+    net->value_fn = make_linear(weights, num_agents, hidden_dim, 1);
+    net->lstm = make_lstm(weights, num_agents, hidden_dim, hidden_dim);
+    net->multidiscrete = make_multidiscrete(num_agents, logit_sizes, num_actions);
+    return net;
+}
+
+
 void demo() {
     Tetris env = {
         .n_rows = 20,
         .n_cols = 10,
-        .deck_size=3,
+        .use_deck_obs = true,
+        .n_noise_obs = 0,
+        .n_init_garbage = 0,
     };
     allocate(&env);
     env.client = make_client(&env);
     c_reset(&env);
 
-    Weights* weights = load_weights("resources/tetris/tetris_weights.bin", 163208);
+    Weights* weights = load_weights("resources/tetris/tetris_weights.bin", 588552);
     int logit_sizes[1] = {7};
-    LinearLSTM* net = make_linearlstm(weights, 1, 234, logit_sizes, 1);
+    LinearLSTM* net = make_linearlstm_256(weights, 1, 234, logit_sizes, 1);
 
+    // State tracking for single-press actions to avoid using IsKeyPressed
+    // because IsKeyPressed doesn't work well in web browsers
+    static bool rotate_key_was_down = false;
+    static bool hard_drop_key_was_down = false;
+    static bool swap_key_was_down = false;
+
+    int frame = 0;
+    env.actions[0] = 0;
     while (!WindowShouldClose()) {
+        bool process_logic = true;
+        frame++;
+
         if (IsKeyDown(KEY_LEFT_SHIFT)) {
-            if (IsKeyDown(KEY_LEFT)  || IsKeyDown(KEY_A)){
-                env.actions[0] = 1;
-            }
-            if (IsKeyDown(KEY_RIGHT)  || IsKeyDown(KEY_D)){
-                env.actions[0] = 2;
-            }
-            if (IsKeyPressed(KEY_UP)  || IsKeyDown(KEY_W)) {
-                env.actions[0] = 3;
-            }
-            if (IsKeyDown(KEY_DOWN)  || IsKeyDown(KEY_S)) {
-                env.actions[0] = 4;
-            }
-            if (IsKeyPressed(KEY_SPACE)) {
-                env.actions[0] = 5;
-            }
-            if (IsKeyPressed(KEY_C)) {
-                env.actions[0] = 6;
+            if (frame % 3 != 0) {
+                // This effectively slows down the client by 3x
+                process_logic = false;
+            } else {
+                // Use KeyDown for left, right, down to allow continuous input
+                // Though, IsKeyDown can overshoot ...
+                if (IsKeyDown(KEY_LEFT) || IsKeyDown(KEY_A)) {
+                    env.actions[0] = 1;
+                } else if (IsKeyDown(KEY_RIGHT) || IsKeyDown(KEY_D)) {
+                    env.actions[0] = 2;
+                } else if (IsKeyDown(KEY_DOWN) || IsKeyDown(KEY_S)) {
+                    env.actions[0] = 4; // Soft drop
+                }
+                // Manual state tracking for single-press actions, mutually exclusive
+                else if ((IsKeyDown(KEY_UP) || IsKeyDown(KEY_W)) && !rotate_key_was_down) {
+                    env.actions[0] = 3; // Rotate
+                } else if (IsKeyDown(KEY_SPACE) && !hard_drop_key_was_down) {
+                    env.actions[0] = 5; // Hard drop
+                } else if (IsKeyDown(KEY_C) && !swap_key_was_down) {
+                    env.actions[0] = 6; // Swap
+                }
             }
         } else {
             forward_linearlstm(net, env.observations, env.actions);
         }
 
-        c_step(&env);
-        env.actions[0] = 0;
+        if (process_logic) {
+            // Update key state flags after processing actions for the frame
+            rotate_key_was_down = IsKeyDown(KEY_UP) || IsKeyDown(KEY_W);
+            hard_drop_key_was_down = IsKeyDown(KEY_SPACE);
+            swap_key_was_down = IsKeyDown(KEY_C);
+
+            c_step(&env);
+
+            env.actions[0] = 0;
+        }
+
         c_render(&env);
     }
+
     free_linearlstm(net);
     free_allocated(&env);
     close_client(env.client);