Gpudrive

config/ocean/gpudrive.ini

@@ -3,28 +3,30 @@ package = ocean
 env_name = puffer_gpudrive
 policy_name = GPUDrive
 rnn_name = Recurrent
-vec = multiprocessing
+vec = native
 
 [policy]
-input_size = 256
-hidden_size = 256
+input_size = 64
+hidden_size = 128
 
 [rnn]
-input_size = 256
-hidden_size = 256
+input_size = 64
+hidden_size = 128
 
 [env]
-num_envs = 256
-reward_vehicle_collision = -0.25
-reward_offroad_collision = -0.25
+num_envs = 75
+reward_vehicle_collision = 0
+reward_offroad_collision = 0
 
 [train]
-total_timesteps = 250_000_000
-learning_rate = 0.05
-minibatch_size = 32768
-num_workers = 2
-num_envs = 2
+total_timesteps = 150_000_000
+learning_rate = 0.005
+num_workers = 1
+num_envs = 1
 env_batch_size = 1
+anneal_lr = True
+
+
 
 [sweep.env.reward_vehicle_collision]
 distribution = uniform

gpudrive_test.c

@@ -0,0 +1,29 @@
+#include "gpudrive_test.h"
+#include <time.h>
+
+void demo(){
+    long test_time = 10;
+    test_struct env = {
+        .num_agents = 128,
+        .active_agents =15,
+    };
+    init(&env);
+    long start = time(NULL);
+    int i = 0;
+    int mean_obs = 4000;
+    int max_obs = 20000;
+    while (time(NULL) - start < test_time) {
+        step(&env);
+        // return;
+        i++;
+    }
+    long end = time(NULL);
+    printf("SPS: %ld\n", (i *env.active_agents) / (end - start));
+    free_initialized(&env);
+
+}
+
+int main() {
+    demo();
+    return 0;
+}

gpudrive_test.h

@@ -0,0 +1,126 @@
+#include <stdlib.h>
+#include <stdio.h>
+#include <math.h>
+#include <assert.h>
+#include <string.h>
+#define SELF_OBS 7
+#define OTHER_OBS 11
+#define MAP_OBS 3000
+
+static int new_map_obs[MAP_OBS];
+
+typedef struct test_struct test_struct;
+struct test_struct {
+    int* observations;
+    int* agent_states;
+    int num_agents;
+    int active_agents;
+};
+
+void print_obs(test_struct* env, int agent_idx) {
+    printf("agent %d obs: ", agent_idx);
+    int obs_size_per_agent = SELF_OBS + (env->active_agents-1) * (OTHER_OBS);
+    for (int i = 0; i < obs_size_per_agent; i++) {
+        printf("%d ", env->observations[agent_idx*obs_size_per_agent + i]);
+    }
+    printf("\n");
+}
+
+void print_map_obs(test_struct* env) {
+    printf("map obs: ");
+    int obs_size_per_agent = SELF_OBS + (env->active_agents-1) * (OTHER_OBS);
+    for (int i = 0; i < (MAP_OBS); i++) {
+        printf("%d ", env->observations[env->active_agents*obs_size_per_agent + i]);
+    }
+    printf("\n");
+}
+
+void add_obs(test_struct* env) {
+    int obs_idx = 0;
+    int obs_size_per_agent = SELF_OBS + (env->active_agents-1) * (OTHER_OBS);
+    for (int i = 0; i < env->active_agents * obs_size_per_agent; i++) {
+        if (i % obs_size_per_agent == 0 && i != 0) {
+            obs_idx++;
+        }
+        env->observations[i] = obs_idx;
+    }
+    for (int i = 0; i < (MAP_OBS); i++) {
+        env->observations[env->active_agents*obs_size_per_agent + i] = rand() % 100;
+    }
+}
+
+void set_agents(test_struct* env) {
+    int obs_size_per_agent = SELF_OBS + (env->active_agents-1) * (OTHER_OBS);
+    for(int i=0;i<env->num_agents * (obs_size_per_agent);i++) {
+        env->agent_states[i] = rand() % 100;
+    }
+}
+
+void init(test_struct* env) {
+    int obs_size_per_agent = SELF_OBS + (env->active_agents-1) * (OTHER_OBS);
+    env->agent_states = (int*)calloc(env->num_agents * obs_size_per_agent, sizeof(int));
+    env->observations = (int*)calloc(env->active_agents*obs_size_per_agent + MAP_OBS, sizeof(int));
+    add_obs(env);
+}
+
+void free_initialized(test_struct* env) {
+    free(env->observations);
+    free(env->agent_states);
+}
+
+void compute_observations(test_struct* env, int* rand_agents){
+    int obs_size_per_agent = SELF_OBS + (env->active_agents-1) * (OTHER_OBS);
+    for(int i=0;i<env->active_agents;i++) {
+        // printf("selected agents %d\n", rand_agents[i]);
+        // printf("old obs\n");
+        // print_obs(env, i);
+        memcpy(env->observations + i*obs_size_per_agent, env->agent_states + rand_agents[i]*obs_size_per_agent, obs_size_per_agent*sizeof(int));
+        // printf("new obs\n");
+        // print_obs(env, i);
+	memcpy(env->observations + env->active_agents*obs_size_per_agent, new_map_obs, (MAP_OBS)*sizeof(int));
+
+    }
+
+
+}
+void change_world(test_struct* env, int* rand_agents){
+    int obs_size_per_agent = SELF_OBS + (env->active_agents-1) * (OTHER_OBS);
+    for(int i=0;i<env->active_agents;i++) {
+        rand_agents[i] = rand() % env->num_agents;
+        for(int j = 0;j<SELF_OBS;j++) {
+            env->agent_states[rand_agents[i]*obs_size_per_agent + j] = 5;
+        }
+        for(int j = 0;j<OTHER_OBS*(env->active_agents-1);j++) {
+            env->agent_states[rand_agents[i]*obs_size_per_agent + SELF_OBS + j] = 2;
+        }
+
+    }
+    /*for (int j=0;j<MAP_OBS;j++) {
+            new_map_obs[j] = 2;
+    	}
+	*/
+}
+
+void copy_world(test_struct* env, int* rand_agents){
+	int obs_size_per_agent = SELF_OBS + (env->active_agents-1) * (OTHER_OBS);
+	int* blah = (int*)calloc(obs_size_per_agent, sizeof(int));
+	//int* randoworld = (int*)calloc(MAP_OBS, sizeof(int));
+	for(int i =0;i<env->active_agents;i++){
+		rand_agents[i] = rand() % env->num_agents;
+		memcpy(env->agent_states+rand_agents[i]*obs_size_per_agent, blah, obs_size_per_agent*sizeof(int));
+	}
+	//memcpy(new_map_obs, randoworld, MAP_OBS*sizeof(int));
+	free(blah);
+	//free(randoworld);
+}
+void step(test_struct* env) {
+    int rand_agents[env->active_agents];
+    // do something
+    copy_world(env, rand_agents);
+    //change_world(env,rand_agents);
+    // compute new obs  
+    compute_observations(env, rand_agents); 
+        // print_obs(env, rand_agents[0]);
+    // print_obs(env, rand_agents[1]);
+    // print_map_obs(env);
+}

pufferlib/ocean/gpudrive/cy_gpudrive.pyx

@@ -1,4 +1,5 @@
-from libc.stdlib cimport calloc, free
+from libc.stdlib cimport calloc, malloc, free
+from libc.string cimport strcpy
 import numpy as np
 cdef extern from "gpudrive.h":
     int LOG_BUFFER_SIZE
@@ -82,6 +83,8 @@ cdef extern from "gpudrive.h":
         float reward_offroad_collision;
         char* map_name;
         char* reached_goal_this_turn;
+        float world_mean_x;
+        float world_mean_y;
 
     ctypedef struct Client
 
@@ -179,47 +182,62 @@ cdef class CyGPUDrive:
 
         self.client = NULL
         self.num_envs = num_envs
-        self.envs = <GPUDrive*> calloc(num_envs, sizeof(GPUDrive))
+        cdef int num_clones
+        num_clones = 8
+        self.envs = <GPUDrive*> calloc(num_envs*num_clones, sizeof(GPUDrive))
         self.agent_offsets = <int*> calloc(num_envs + 1, sizeof(int))
         self.logs = allocate_logbuffer(LOG_BUFFER_SIZE)
         cdef int i
         for i in range(num_envs + 1):
             self.agent_offsets[i] = offsets[i]
         cdef int inc
-        for i in range(num_envs):
-            inc = self.agent_offsets[i]
-            print(inc)
-            map_file = f"resources/gpudrive/binaries/map_{i:03d}.bin".encode('utf-8')
+        cdef int index
+        cdef int total_envs
+        total_envs = num_envs * num_clones
+        cdef int total_agents
+        total_agents = self.agent_offsets[num_envs]
+        cdef char* c_map_file
+        for i in range(total_envs):
+            env_index = i % num_envs
+            clone_index = i // num_envs
+            inc = self.agent_offsets[env_index]
+            count = self.agent_offsets[env_index+1] - self.agent_offsets[env_index]
+            clone_agent_offset = clone_index * total_agents + inc
+            print("Env Index: ", env_index)
+            print("Increment: ", inc)
+            print("clone_agent_offset: ", clone_agent_offset)
+            map_file = f"resources/gpudrive/binaries/map_{env_index:03d}.bin".encode('utf-8')
+            c_map_file = <char*>malloc(len(map_file) + 1)
+            strcpy(c_map_file, map_file)
             print("cython map_name", map_file)
             self.envs[i] = GPUDrive(
-                observations=&observations[inc, 0],
-                actions=&actions[inc,0],
-                rewards=&rewards[inc],
-                masks=&masks[inc],
-                dones=&terminals[inc],
+                observations=&observations[clone_agent_offset, 0],
+                actions=&actions[clone_agent_offset,0],
+                rewards=&rewards[clone_agent_offset],
+                masks=&masks[clone_agent_offset],
+                dones=&terminals[clone_agent_offset],
                 log_buffer=self.logs,
                 human_agent_idx=human_agent_idx,
                 reward_vehicle_collision=reward_vehicle_collision,
                 reward_offroad_collision=reward_offroad_collision,
-                map_name = map_file
+                map_name = c_map_file
             )
-            print("init")
             init(&self.envs[i])
             self.client = NULL
 
 
     def reset(self):
         cdef int i
-        for i in range(self.num_envs):
+        for i in range(self.num_envs*8):
             c_reset(&self.envs[i])
 
     def step(self):
         cdef int i
-        for i in range(self.num_envs):
+        for i in range(self.num_envs*8):
             c_step(&self.envs[i])
 
     def render(self):
-        cdef GPUDrive* env = &self.envs[211]
+        cdef GPUDrive* env = &self.envs[11]
         if self.client == NULL:
             import os
             cwd = os.getcwd()

pufferlib/ocean/gpudrive/gpudrive.c

@@ -205,7 +205,7 @@ void performance_test() {
 }
 
 int main() {
-    demo();
-    //performance_test();
+    // demo();
+    performance_test();
     return 0;
 }