Merge branch 'PufferAI:3.0' into vision_test

Author: jutskitting

pufferlib/config/ocean/breakout.ini

@@ -10,7 +10,21 @@ num_envs = 8
 [env]
 num_envs = 1024
 frameskip = 4
-
+width = 576
+height = 330
+paddle_width = 62
+paddle_height = 8
+ball_width = 32
+ball_height = 32
+brick_width = 32
+brick_height = 12
+brick_rows = 6
+brick_cols = 18
+initial_ball_speed = 256
+max_ball_speed = 448
+paddle_speed = 620
+continuous = 0
+ 
 [policy]
 hidden_size = 128
 

pufferlib/config/ocean/cartpole.ini

@@ -6,6 +6,12 @@ rnn_name = Recurrent
 
 [env]
 num_envs = 4096
+cart_mass = 1.0
+pole_mass = 0.1
+pole_length = 0.5
+gravity = 9.8
+force_mag = 10.0
+dt = 0.02
 
 [train]
 total_timesteps = 20_000_000

pufferlib/ocean/breakout/binding.c

@@ -7,14 +7,17 @@ static int my_init(Env* env, PyObject* args, PyObject* kwargs) {
     env->frameskip = unpack(kwargs, "frameskip");
     env->width = unpack(kwargs, "width");
     env->height = unpack(kwargs, "height");
-    env->paddle_width = unpack(kwargs, "paddle_width");
+    env->initial_paddle_width = unpack(kwargs, "paddle_width");
     env->paddle_height = unpack(kwargs, "paddle_height");
     env->ball_width = unpack(kwargs, "ball_width");
     env->ball_height = unpack(kwargs, "ball_height");
     env->brick_width = unpack(kwargs, "brick_width");
     env->brick_height = unpack(kwargs, "brick_height");
     env->brick_rows = unpack(kwargs, "brick_rows");
     env->brick_cols = unpack(kwargs, "brick_cols");
+    env->initial_ball_speed = unpack(kwargs, "initial_ball_speed");
+    env->max_ball_speed = unpack(kwargs, "max_ball_speed");
+    env->paddle_speed = unpack(kwargs, "paddle_speed");
     env->continuous = unpack(kwargs, "continuous");
     init(env);
     return 0;

pufferlib/ocean/breakout/breakout.c

@@ -19,6 +19,9 @@ void demo() {
         .brick_height = 12,
         .brick_rows = 6,
         .brick_cols = 18,
+        .initial_ball_speed = 256,
+        .max_ball_speed = 448,
+        .paddle_speed = 620,
         .continuous = 0,
     };
     allocate(&env);

pufferlib/ocean/breakout/breakout.h

@@ -9,7 +9,6 @@
 #define NOOP 0
 #define LEFT 1
 #define RIGHT 2
-#define MAX_BALL_SPEED 448
 #define HALF_PADDLE_WIDTH 31
 #define Y_OFFSET 50
 #define TICK_RATE 1.0f/60.0f
@@ -55,9 +54,13 @@ typedef struct Breakout {
     float* brick_y;
     float* brick_states;
     int balls_fired;
+    float initial_paddle_width;
     float paddle_width;
     float paddle_height;
+    float paddle_speed;
     float ball_speed;
+    float initial_ball_speed;
+    float max_ball_speed;
     int hits;
     int width;
     int height;
@@ -289,15 +292,13 @@ bool calc_paddle_ball_collisions(Breakout* env, CollisionInfo* collision_info) {
     collision_info->brick_index = BRICK_INDEX_PADDLE_COLLISION;
 
     env->hit_brick = false;
-    float relative_intersection = ((env->ball_x +
-                                    env->ball_width / 2) -
-                                   env->paddle_x) /
-                                  env->paddle_width;
+    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->hits += 1;
-    if (env->hits % 4 == 0 && env->ball_speed < MAX_BALL_SPEED) {
+    if (env->hits % 4 == 0 && env->ball_speed < env->max_ball_speed) {
         env->ball_speed += 64;
     }
     if (env->score == env->half_max_score) {
@@ -336,12 +337,16 @@ void calc_all_wall_collisions(Breakout* env, CollisionInfo* collision_info) {
 // With rare floating point conditions, the ball could escape the bounds.
 // Let's handle that explicitly.
 void check_wall_bounds(Breakout* env) {
-    if (env->ball_x < 0)
-        env->ball_x += MAX_BALL_SPEED * 1.1f * TICK_RATE;
-    if (env->ball_x > env->width)
-        env->ball_x -= MAX_BALL_SPEED * 1.1f * TICK_RATE;
-    if (env->ball_y < 0)
-        env->ball_y += MAX_BALL_SPEED * 1.1f * TICK_RATE;
+    float offset = env->max_ball_speed * 1.1f * TICK_RATE;
+    if (env->ball_x < 0) {
+        env->ball_x += offset;
+    }
+    if (env->ball_x > env->width) {
+        env->ball_x -= offset;
+    }
+    if (env->ball_y < 0) {
+        env->ball_y += offset;
+    }
 }
 
 void destroy_brick(Breakout* env, int brick_idx) {
@@ -353,7 +358,7 @@ void destroy_brick(Breakout* env, int brick_idx) {
     env->rewards[0] += gained_points;
 
     if (brick_idx / env->brick_cols < 3) {
-        env->ball_speed = MAX_BALL_SPEED;
+        env->ball_speed = env->max_ball_speed;
     }
 }
 
@@ -393,8 +398,8 @@ void reset_round(Breakout* env) {
     env->balls_fired = 0;
     env->hit_brick = false;
     env->hits = 0;
-    env->ball_speed = 256;
-    env->paddle_width = 2 * HALF_PADDLE_WIDTH;
+    env->ball_speed = env->initial_ball_speed;
+    env->paddle_width = env->initial_paddle_width;
 
     env->paddle_x = env->width / 2.0 - env->paddle_width / 2;
     env->paddle_y = env->height - env->paddle_height - 10;
@@ -437,7 +442,7 @@ void step_frame(Breakout* env, float action) {
     if (env->continuous){
         act = action;
     }
-    env->paddle_x += act * 620 * TICK_RATE;
+    env->paddle_x += act * env->paddle_speed * TICK_RATE;
     if (env->paddle_x <= 0){
         env->paddle_x = fmaxf(0, env->paddle_x);
     } else {

pufferlib/ocean/breakout/breakout.py

@@ -1,10 +1,3 @@
-'''High-perf Pong
-
-Inspired from https://gist.github.com/Yttrmin/18ecc3d2d68b407b4be1
-& https://jair.org/index.php/jair/article/view/10819/25823
-& https://www.youtube.com/watch?v=PSQt5KGv7Vk
-'''
-
 import numpy as np
 import gymnasium
 
@@ -17,7 +10,10 @@ def __init__(self, num_envs=1, render_mode=None,
             paddle_width=62, paddle_height=8,
             ball_width=32, ball_height=32,
             brick_width=32, brick_height=12,
-            brick_rows=6, brick_cols=18, continuous=False, log_interval=128,
+            brick_rows=6, brick_cols=18,
+            initial_ball_speed=256, max_ball_speed=448,
+            paddle_speed=620,
+            continuous=False, log_interval=128,
             buf=None, seed=0):
         self.single_observation_space = gymnasium.spaces.Box(low=0, high=1,
             shape=(10 + brick_rows*brick_cols,), dtype=np.float32)
@@ -40,10 +36,14 @@ def __init__(self, num_envs=1, render_mode=None,
             self.actions = self.actions.astype(np.float32)
 
         self.c_envs = binding.vec_init(self.observations, self.actions, self.rewards,
-            self.terminals, self.truncations, num_envs, seed, frameskip=frameskip, width=width, height=height,
-            paddle_width=paddle_width, paddle_height=paddle_height, ball_width=ball_width, ball_height=ball_height,
-            brick_width=brick_width, brick_height=brick_height, brick_rows=brick_rows,
-            brick_cols=brick_cols, continuous=continuous
+            self.terminals, self.truncations, num_envs, seed, frameskip=frameskip,
+            width=width, height=height, paddle_width=paddle_width,
+            paddle_height=paddle_height, ball_width=ball_width,
+            ball_height=ball_height, brick_width=brick_width,
+            brick_height=brick_height, brick_rows=brick_rows,
+            brick_cols=brick_cols, initial_ball_speed=initial_ball_speed,
+            max_ball_speed=max_ball_speed, paddle_speed=paddle_speed,
+            continuous=continuous
         )
 
     def reset(self, seed=0):

pufferlib/ocean/cartpole/binding.c

@@ -3,6 +3,12 @@
 #include "../env_binding.h"
 
 static int my_init(Env* env, PyObject* args, PyObject* kwargs) {   
+    env->cart_mass = unpack(kwargs, "cart_mass");
+    env->pole_mass = unpack(kwargs, "pole_mass");
+    env->pole_length = unpack(kwargs, "pole_length");
+    env->gravity = unpack(kwargs, "gravity");
+    env->force_mag = unpack(kwargs, "force_mag");
+    env->tau = unpack(kwargs, "dt");
     env->continuous = unpack(kwargs, "continuous");
     init(env);
     return 0;

pufferlib/ocean/cartpole/cartpole.h

@@ -6,15 +6,6 @@
 #include <time.h>
 #include "raylib.h"
 
-#define GRAVITY 9.8f
-#define MASSCART 1.0f
-#define MASSPOLE 0.1f
-#define TOTAL_MASS (MASSPOLE + MASSCART)
-#define LENGTH 0.5f // half pole length
-#define POLEMASS_LENGTH (MASSPOLE * LENGTH)
-#define FORCE_MAG 10.0f
-#define TAU 0.02f // timestep duration
-
 #define X_THRESHOLD 2.4f
 #define THETA_THRESHOLD_RADIANS (12 * 2 * M_PI / 360)
 #define MAX_STEPS 200
@@ -51,6 +42,12 @@ struct Cartpole {
     float theta;
     float theta_dot;
     int tick;
+    float cart_mass;
+    float pole_mass;
+    float pole_length;
+    float gravity;
+    float force_mag;
+    float tau;
     int continuous;
     float episode_return;
 };
@@ -172,25 +169,29 @@ void c_step(Cartpole* env) {
     }
     /* ========================================================== */
 
-    if (!isfinite(a))             a = 0.0f;
+    if (!isfinite(a)) {
+        a = 0.0f;
+    }
     a = fminf(fmaxf(a, -1.0f), 1.0f);
     env->actions[0] = a;
 
-    float force = env->continuous ? a * FORCE_MAG
-                                  : (a > 0.5f ? FORCE_MAG : -FORCE_MAG);
+    float force = env->continuous ? a * env->force_mag
+                                  : (a > 0.5f ? env->force_mag: -env->force_mag);
 
     float costheta = cosf(env->theta);
     float sintheta = sinf(env->theta);
 
-    float temp = (force + POLEMASS_LENGTH * env->theta_dot * env->theta_dot * sintheta) / TOTAL_MASS;
-    float thetaacc = (GRAVITY * sintheta - costheta * temp) / 
-                     (LENGTH * (4.0f / 3.0f - MASSPOLE * costheta * costheta / TOTAL_MASS));
-    float xacc = temp - POLEMASS_LENGTH * thetaacc * costheta / TOTAL_MASS;
-
-    env->x += TAU * env->x_dot;
-    env->x_dot += TAU * xacc;
-    env->theta += TAU * env->theta_dot;
-    env->theta_dot += TAU * thetaacc;
+    float total_mass = env->cart_mass + env->pole_mass;
+    float polemass_length = total_mass + env->pole_mass;
+    float temp = (force + polemass_length * env->theta_dot * env->theta_dot * sintheta) / total_mass;
+    float thetaacc = (env->gravity * sintheta - costheta * temp) / 
+                     (env->pole_length * (4.0f / 3.0f - total_mass * costheta * costheta / total_mass));
+    float xacc = temp - polemass_length * thetaacc * costheta / total_mass;
+
+    env->x += env->tau * env->x_dot;
+    env->x_dot += env->tau * xacc;
+    env->theta += env->tau * env->theta_dot;
+    env->theta_dot += env->tau * thetaacc;
 
     env->tick += 1;
 

pufferlib/ocean/cartpole/cartpole.py

@@ -4,7 +4,10 @@
 from pufferlib.ocean.cartpole import binding
 
 class Cartpole(pufferlib.PufferEnv):
-    def __init__(self, num_envs=1, render_mode='human', report_interval=1, continuous=False, buf=None, seed=0):
+    def __init__(self, num_envs=1, cart_mass=1.0, pole_mass=0.1,
+            pole_length=0.5, gravity=9.8, force_mag=10.0, dt=0.02,
+            render_mode='human', report_interval=1, continuous=False,
+            buf=None, seed=0):
         self.render_mode = render_mode
         self.num_agents = num_envs
         self.report_interval = report_interval
@@ -35,6 +38,12 @@ def __init__(self, num_envs=1, render_mode='human', report_interval=1, continuou
             self.truncations,
             num_envs,
             seed,
+            cart_mass=cart_mass,
+            pole_mass=pole_mass,
+            pole_length=pole_length,
+            gravity=gravity,
+            force_mag=force_mag,
+            dt=dt,
             continuous=int(self.continuous),
         )
 

pufferlib/ocean/drone_race/binding.c

@@ -13,6 +13,9 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) {
 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, "collision_rate", log->collision_rate);
+    assign_to_dict(dict, "oob", log->oob);
+    assign_to_dict(dict, "timeout", log->timeout);
     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);