Local snake newbind pr

pufferlib/ocean/snake/binding.c

@@ -0,0 +1,29 @@
+#include "snake.h"
+
+#define Env CSnake
+#include "../env_binding.h"
+
+static int my_init(Env* env, PyObject* args, PyObject* kwargs) {   
+    env->width = unpack(kwargs, "width");
+    env->height = unpack(kwargs, "height");
+    env->num_snakes = unpack(kwargs, "num_snakes");
+    env->vision = unpack(kwargs, "vision");
+    env->leave_corpse_on_death = unpack(kwargs, "leave_corpse_on_death");
+    env->food = unpack(kwargs, "num_food");
+    env->reward_food = unpack(kwargs, "reward_food");
+    env->reward_corpse = unpack(kwargs, "reward_corpse");
+    env->reward_death = unpack(kwargs, "reward_death");
+    env->max_snake_length = unpack(kwargs, "max_snake_length");
+    env->cell_size = unpack(kwargs, "cell_size");    
+    init_csnake(env);
+    return 0;
+}
+
+static int my_log(PyObject* dict, Log* log) {
+    assign_to_dict(dict, "perf", log->perf);
+    assign_to_dict(dict, "score", log->score);
+    assign_to_dict(dict, "episode_return", log->episode_return);
+    assign_to_dict(dict, "episode_length", log->episode_length);
+    assign_to_dict(dict, "n", log->n);
+    return 0;
+}

pufferlib/ocean/snake/snake.c

@@ -20,7 +20,7 @@ int demo() {
 
     Weights* weights = load_weights("resources/snake_weights.bin", 148357);
     LinearLSTM* net = make_linearlstm(weights, env.num_snakes, env.obs_size, 4);
-    Client* client = make_client(2, env.width, env.height);
+    env.client = make_client(2, env.width, env.height);
 
     while (!WindowShouldClose()) {
         // User can take control of the first snake
@@ -36,11 +36,11 @@ int demo() {
             forward_linearlstm(net, net->obs, env.actions);
         }
         c_step(&env);
-        c_render(client, &env);
+        c_render(&env);
     }
     free_linearlstm(net);
     free(weights);
-    close_client(client);
+    close_client(env.client);
     free_csnake(&env);
     return 0;
 }
@@ -71,11 +71,12 @@ void test_performance(float test_time) {
         i++;
     }
     int end = time(NULL);
+    free_csnake(&env);
     printf("SPS: %f\n", (float)env.num_snakes*i / (end - start));
 }
 
 int main() {
     demo();
-    //test_performance(30);
+    // test_performance(30);
     return 0;
 }