Env vec speeds look reasonable in testing

Author: jsuarez5341

profile_kernels.cu

@@ -1755,6 +1755,9 @@ EnvSpeedArgs* create_envspeedargs(int total_agents, int num_buffers, int num_thr
         fprintf(stderr, "Failed to create environments\n");
         return nullptr;
     }
+    for (int i = 0; i < num_buffers; i++) {
+        cudaStreamCreateWithFlags(&vec->streams[i], cudaStreamNonBlocking);
+    }
 
     int num_envs = vec->size;
     printf("Created %d envs (%s) for %d total_agents\n", num_envs, TOSTRING(ENV_NAME), total_agents);
@@ -1794,25 +1797,23 @@ float profile_env_rollout(EnvSpeedArgs* args, const char* name) {
     cudaEventCreate(&stop);
 
     // Warmup
-    auto start_time = std::chrono::steady_clock::now();
-    for (int i = 0; i < 10; ++i) {
+    float start_time = get_time_sec();
+    for (int i = 0; i < 100; ++i) {
         run_env_rollout(args);
         cudaDeviceSynchronize();
-        auto now = std::chrono::steady_clock::now();
-        float elapsed = std::chrono::duration<float>(now - start_time).count();
+        auto elapsed = get_time_sec() - start_time;
         if (elapsed > TIMEOUT_SEC) break;
     }
 
-    start_time = std::chrono::steady_clock::now();
+    start_time = get_time_sec();
     cudaProfilerStart();
     if (name) nvtxRangePushA(name);
     cudaEventRecord(start);
     float completed = 0;
     for (int i = 0; i < 1000; ++i) {
         run_env_rollout(args);
         completed += 1;
-        auto now = std::chrono::steady_clock::now();
-        float elapsed = std::chrono::duration<float>(now - start_time).count();
+        float elapsed = get_time_sec() - start_time;
         if (elapsed > TIMEOUT_SEC) break;
     }
     cudaDeviceSynchronize();
@@ -1883,10 +1884,12 @@ int main(int argc, char** argv) {
     int buffers = BUF;
     int threads = 16;
     int horizon = T;
+    int total_agents = BR * buffers;
     for (int i = 2; i < argc - 1; i++) {
         if (strcmp(argv[i], "--buffers") == 0) buffers = atoi(argv[++i]);
         else if (strcmp(argv[i], "--threads") == 0) threads = atoi(argv[++i]);
         else if (strcmp(argv[i], "--horizon") == 0) horizon = atoi(argv[++i]);
+        else if (strcmp(argv[i], "--total-agents") == 0) total_agents = atoi(argv[++i]);
     }
 
     warmup_gpu();
@@ -1914,7 +1917,7 @@ int main(int argc, char** argv) {
 
 #ifdef USE_STATIC_ENV
     if (strcmp(profile, "envspeed") == 0 || strcmp(profile, "all") == 0) {
-        profile_envspeed(buffers * BR, buffers, threads, horizon);
+        profile_envspeed(total_agents, buffers, threads, horizon);
     }
 #endif
 

pufferlib/config/ocean/benchmark.ini

@@ -5,7 +5,7 @@ policy_name = Policy
 rnn_name = Recurrent
 
 [env]
-bandwidth = 1
+bandwidth = 512
 compute = 0
 
 [vec]

pufferlib/extensions/env_binding.c

@@ -89,6 +89,7 @@ static void* static_omp_threadmanager(void* arg) {
 
     Env* envs = (Env*)vec->envs;
 
+    printf("Num workers: %d\n", num_workers);
     while (true) {
         while (atomic_load(&buffer_states[buf]) != OMP_RUNNING) {
             if (atomic_load(&threading->shutdown)) {
@@ -107,9 +108,15 @@ static void* static_omp_threadmanager(void* arg) {
                 cudaMemcpyDeviceToHost, stream);
             cudaStreamSynchronize(stream);
 
-            #pragma omp parallel for schedule(static) num_threads(num_workers)
-            for (int i = env_start; i < env_start + env_count; i++) {
-                c_step(&envs[i]);
+            if (num_workers > 1) {
+                #pragma omp parallel for schedule(static) num_threads(num_workers)
+                for (int i = env_start; i < env_start + env_count; i++) {
+                    c_step(&envs[i]);
+                }
+            } else {
+                for (int i = env_start; i < env_start + env_count; i++) {
+                    c_step(&envs[i]);
+                }
             }
 
             cudaMemcpyAsync(

pufferlib/ocean/benchmark/benchmark.h

@@ -26,7 +26,7 @@ void c_step(Benchmark* env) {
         result = sinf(result + 0.1f);
     }
 
-    memset(env->observations, result, env->bandwidth);
+    //memset(env->observations, result, env->bandwidth);
 }
 
 void c_render(Benchmark* env) { }

pufferlib/ocean/benchmark/binding.h

@@ -1,5 +1,5 @@
 #include "benchmark.h"
-#define OBS_SIZE 1 // TODO: Current API forces you to edit this per obs size
+#define OBS_SIZE 512 // TODO: Current API forces you to edit this per obs size
 #define NUM_ATNS 1
 #define ACT_SIZES {2}
 #define OBS_TYPE UNSIGNED_CHAR

pufferlib/ocean/breakout/breakout.h

@@ -159,18 +159,16 @@ void compute_observations(Breakout* env) {
     env->observations[7] = env->score / 864.0f;
     env->observations[8] = env->num_balls / 5.0f;
     env->observations[9] = env->paddle_width / (2.0f * HALF_PADDLE_WIDTH);
-    for (int i = 0; i < env->num_bricks; i++) {
-        env->observations[10 + i] = env->brick_states[i];
-    }
+    memcpy(env->observations + 10, env->brick_states, sizeof(float) * env->num_bricks);
 }
 
 // Collision of a stationary vertical line segment (xw,yw) to (xw,yw+hw)
 // with a moving line segment (x+vx*t,y+vy*t) to (x+vx*t,y+vy*t+h).
 static inline bool calc_vline_collision(float xw, float yw, float hw, float x,
         float y, float vx, float vy, float h, CollisionInfo* col) {
     float t_new = (xw - x) / vx;
-    float topmost = fmin(yw + hw, y + h + vy * t_new);
-    float botmost = fmax(yw, y + vy * t_new);
+    float topmost = fminf(yw + hw, y + h + vy * t_new);
+    float botmost = fmaxf(yw, y + vy * t_new);
     float overlap_new = topmost - botmost;
 
     // Collision finds the smallest time of collision with the greatest overlap
@@ -248,26 +246,52 @@ static inline void calc_brick_collision(Breakout* env, int idx,
     }
 }
 static inline int column_index(Breakout* env, float x) {
-    return (int)(floorf(x / env->brick_width));
+    return (int)(x / env->brick_width);
 }
 static inline int row_index(Breakout* env, float y) {
-    return (int)(floorf((y - Y_OFFSET) / env->brick_height));
+    return (int)((y - Y_OFFSET) / env->brick_height);
 }
 
 void calc_all_brick_collisions(Breakout* env, CollisionInfo* collision_info) {
-    int column_from = column_index(env, fminf(env->ball_x + env->ball_vx, env->ball_x));
-    column_from = fmaxf(column_from, 0);
-    int column_to = column_index(env, fmaxf(env->ball_x + env->ball_width + env->ball_vx, env->ball_x + env->ball_width));
-    column_to = fminf(column_to, env->brick_cols - 1);
-    int row_from = row_index(env, fminf(env->ball_y + env->ball_vy, env->ball_y));
-    row_from = fmaxf(row_from, 0);
-    int row_to = row_index(env, fmaxf(env->ball_y + env->ball_height + env->ball_vy, env->ball_y + env->ball_height));
-    row_to = fminf(row_to, env->brick_rows - 1);
+    float ball_x = env->ball_x;
+    float ball_x_dst = ball_x + env->ball_vx;
+    float ball_y = env->ball_y;
+    float ball_y_dst = ball_y + env->ball_vy;
+    float ball_width = env->ball_width;
+    float ball_height = env->ball_height;
+
+    int row_from = row_index(env, ball_y < ball_y_dst ? ball_y : ball_y_dst);
+    if (row_from < 0) {
+        row_from = 0;
+    }
+
+    if (row_from > env->brick_rows) {
+        return;
+    }
+
+    int column_from = column_index(env, ball_x < ball_x_dst ? ball_x : ball_x_dst);
+    if (column_from < 0) {
+        column_from = 0;
+    }
+
+    float ball_x_end = ball_x + ball_width;
+    float ball_x_dst_end = ball_x_dst + ball_width;
+    int column_to = column_index(env, ball_x_dst_end > ball_x_end ? ball_x_dst_end : ball_x_end);
+    if (column_to >= env->brick_cols) {
+        column_to = env->brick_cols - 1;
+    }
+
+    float ball_y_end = ball_y + ball_height;
+    float ball_y_dst_end = ball_y_dst + ball_height;
+    int row_to = row_index(env, ball_y_dst_end > ball_y_end ? ball_y_dst_end : ball_y_end);
+    if (row_to >= env->brick_rows) {
+        row_to = env->brick_rows - 1;
+    }
 
     for (int row = row_from; row <= row_to; row++) {
         for (int column = column_from; column <= column_to; column++) {
             int brick_index = row * env->brick_cols + column;
-            if (env->brick_states[brick_index] == 0.0)
+            if (env->brick_states[brick_index] == 0.0f)
                 calc_brick_collision(env, brick_index, collision_info);
         }
     }
@@ -297,8 +321,8 @@ bool calc_paddle_ball_collisions(Breakout* env, CollisionInfo* collision_info) {
     float relative_intersection = (
         (env->ball_x + env->ball_width / 2) - env->paddle_x) / env->paddle_width;
     float angle = -base_angle + relative_intersection * 2 * base_angle;
-    env->ball_vx = sin(angle) * env->ball_speed * TICK_RATE;
-    env->ball_vy = -cos(angle) * env->ball_speed * TICK_RATE;
+    env->ball_vx = sinf(angle) * env->ball_speed * TICK_RATE;
+    env->ball_vy = -cosf(angle) * env->ball_speed * TICK_RATE;
     env->hits += 1;
     if (env->hits % 4 == 0 && env->ball_speed < env->max_ball_speed) {
         env->ball_speed += 64;
@@ -430,8 +454,8 @@ void step_frame(Breakout* env, float action) {
         env->balls_fired = 1;
         float direction = M_PI / 3.25f;
 
-        env->ball_vy = cos(direction) * env->ball_speed * TICK_RATE;
-        env->ball_vx = sin(direction) * env->ball_speed * TICK_RATE;
+        env->ball_vy = cosf(direction) * env->ball_speed * TICK_RATE;
+        env->ball_vx = sinf(direction) * env->ball_speed * TICK_RATE;
         if (rand() % 2 == 0) {
             env->ball_vx = -env->ball_vx;
         }