Merging conditioning stuff first (#4)

* Seperates visualization code from the core drive.c

Changes
- Remove GIF generation code from drive.c
- Improved load_weights to auto-detect file size
- view flag
- random maps if map-name not passed
- policy-name flag to make videos for a particular policy. Saved in the policy directoryq

* merging from main

* fixed test_drive_render.py

* added reward/entropy/discount conditioning

* pre-commit fixes

* some build fixes

* removed output files

* readd weights

* fixed viz env

* run pre-commit

* pre-commit fix

* fixed issues with conditioning and viz

* changed respawn timstep to always be at obs[6]

* merged main

* fixed wrong jerk/classic obs values

* was freeing twice

Author: m2kulkarni

pufferlib/config/ocean/drive.ini

@@ -42,6 +42,7 @@ offroad_behavior = 0
 scenario_length = 91
 resample_frequency = 910
 num_maps = 1
+condition_type = "none" # Options: "none", "reward", "entropy", "discount", "all"
 ; Determines which step of the trajectory to initialize the agents at upon reset
 init_steps = 0
 ; Options: "control_vehicles", "control_agents", "control_tracks_to_predict"

pufferlib/extensions/cuda/pufferlib.cu

@@ -21,7 +21,7 @@ __host__ __device__ void puff_advantage_row_cuda(float* values, float* rewards,
 
 void vtrace_check_cuda(torch::Tensor values, torch::Tensor rewards,
         torch::Tensor dones, torch::Tensor importance, torch::Tensor advantages,
-        int num_steps, int horizon) {
+        torch::Tensor gammas, int num_steps, int horizon) {
 
     // Validate input tensors
     torch::Device device = values.device();
@@ -35,27 +35,33 @@ void vtrace_check_cuda(torch::Tensor values, torch::Tensor rewards,
             t.contiguous();
         }
     }
+    // Validate gammas tensor
+    TORCH_CHECK(gammas.dim() == 1, "Gammas must be 1D");
+    TORCH_CHECK(gammas.size(0) == num_steps, "Gammas size must match num_steps");
+    TORCH_CHECK(gammas.dtype() == torch::kFloat32, "Gammas must be float32");
+    TORCH_CHECK(gammas.is_cuda(), "Gammas must be on GPU");
+    TORCH_CHECK(gammas.is_contiguous(), "Gammas must be contiguous");
 }
 
  // [num_steps, horizon]
 __global__ void puff_advantage_kernel(float* values, float* rewards,
-        float* dones, float* importance, float* advantages, float gamma,
+        float* dones, float* importance, float* advantages, float* gammas,
         float lambda, float rho_clip, float c_clip, int num_steps, int horizon) {
     int row = blockIdx.x*blockDim.x + threadIdx.x;
     if (row >= num_steps) {
         return;
     }
     int offset = row*horizon;
     puff_advantage_row_cuda(values + offset, rewards + offset, dones + offset,
-        importance + offset, advantages + offset, gamma, lambda, rho_clip, c_clip, horizon);
+        importance + offset, advantages + offset, gammas[row], lambda, rho_clip, c_clip, horizon);
 }
 
 void compute_puff_advantage_cuda(torch::Tensor values, torch::Tensor rewards,
         torch::Tensor dones, torch::Tensor importance, torch::Tensor advantages,
-        double gamma, double lambda, double rho_clip, double c_clip) {
+        torch::Tensor gammas, double lambda, double rho_clip, double c_clip) {
     int num_steps = values.size(0);
     int horizon = values.size(1);
-    vtrace_check_cuda(values, rewards, dones, importance, advantages, num_steps, horizon);
+    vtrace_check_cuda(values, rewards, dones, importance, advantages, gammas, num_steps, horizon);
     TORCH_CHECK(values.is_cuda(), "All tensors must be on GPU");
 
     int threads_per_block = 256;
@@ -67,7 +73,7 @@ void compute_puff_advantage_cuda(torch::Tensor values, torch::Tensor rewards,
         dones.data_ptr<float>(),
         importance.data_ptr<float>(),
         advantages.data_ptr<float>(),
-        gamma,
+        gammas.data_ptr<float>(),
         lambda,
         rho_clip,
         c_clip,

pufferlib/extensions/pufferlib.cpp

@@ -42,7 +42,7 @@ void puff_advantage_row(float* values, float* rewards, float* dones,
 
 void vtrace_check(torch::Tensor values, torch::Tensor rewards,
         torch::Tensor dones, torch::Tensor importance, torch::Tensor advantages,
-        int num_steps, int horizon) {
+        torch::Tensor gammas,  int num_steps, int horizon) {
 
     // Validate input tensors
     torch::Device device = values.device();
@@ -56,36 +56,42 @@ void vtrace_check(torch::Tensor values, torch::Tensor rewards,
             t.contiguous();
         }
     }
+    // Validate gammas tensor
+    TORCH_CHECK(gammas.dim() == 1, "Gammas must be 1D");
+    TORCH_CHECK(gammas.size(0) == num_steps, "Gammas size must match num_steps");
+    TORCH_CHECK(gammas.dtype() == torch::kFloat32, "Gammas must be float32");
+    TORCH_CHECK(gammas.is_contiguous(), "Gammas must be contiguous");
 }
 
 
 // [num_steps, horizon]
 void puff_advantage(float* values, float* rewards, float* dones, float* importance,
-        float* advantages, float gamma, float lambda, float rho_clip, float c_clip,
+        float* advantages, float* gammas, float lambda, float rho_clip, float c_clip,
         int num_steps, const int horizon){
-    for (int offset = 0; offset < num_steps*horizon; offset+=horizon) {
+    for (int row = 0; row < num_steps; row++) {
+        int offset = row * horizon;
         puff_advantage_row(values + offset, rewards + offset,
             dones + offset, importance + offset, advantages + offset,
-            gamma, lambda, rho_clip, c_clip, horizon
+            gammas[row], lambda, rho_clip, c_clip, horizon
         );
     }
 }
 
 
 void compute_puff_advantage_cpu(torch::Tensor values, torch::Tensor rewards,
         torch::Tensor dones, torch::Tensor importance, torch::Tensor advantages,
-        double gamma, double lambda, double rho_clip, double c_clip) {
+        torch::Tensor gammas, double lambda, double rho_clip, double c_clip) {
     int num_steps = values.size(0);
     int horizon = values.size(1);
-    vtrace_check(values, rewards, dones, importance, advantages, num_steps, horizon);
+    vtrace_check(values, rewards, dones, importance, advantages, gammas, num_steps, horizon);
     puff_advantage(values.data_ptr<float>(), rewards.data_ptr<float>(),
         dones.data_ptr<float>(), importance.data_ptr<float>(), advantages.data_ptr<float>(),
-        gamma, lambda, rho_clip, c_clip, num_steps, horizon
+        gammas.data_ptr<float>(), lambda, rho_clip, c_clip, num_steps, horizon
     );
 }
 
 TORCH_LIBRARY(pufferlib, m) {
-   m.def("compute_puff_advantage(Tensor(a!) values, Tensor(b!) rewards, Tensor(c!) dones, Tensor(d!) importance, Tensor(e!) advantages, float gamma, float lambda, float rho_clip, float c_clip) -> ()");
+   m.def("compute_puff_advantage(Tensor(a!) values, Tensor(b!) rewards, Tensor(c!) dones, Tensor(d!) importance, Tensor(e!) advantages, Tensor gammas, float lambda, float rho_clip, float c_clip) -> ()");
  }
 
 TORCH_LIBRARY_IMPL(pufferlib, CPU, m) {

pufferlib/ocean/drive/binding.c

@@ -73,6 +73,7 @@ static PyObject* my_shared(PyObject* self, PyObject* args, PyObject* kwargs) {
     int init_mode = unpack(kwargs, "init_mode");
     int control_mode = unpack(kwargs, "control_mode");
     int init_steps = unpack(kwargs, "init_steps");
+    int max_controlled_agents = unpack(kwargs, "max_controlled_agents");
     clock_gettime(CLOCK_REALTIME, &ts);
     srand(ts.tv_nsec);
     int total_agent_count = 0;
@@ -89,6 +90,7 @@ static PyObject* my_shared(PyObject* self, PyObject* args, PyObject* kwargs) {
         env->init_mode = init_mode;
         env->control_mode = control_mode;
         env->init_steps = init_steps;
+        env->max_controlled_agents = max_controlled_agents;
         sprintf(map_file, "resources/drive/binaries/map_%03d.bin", map_id);
         env->entities = load_map_binary(map_file, env);
         set_active_agents(env);
@@ -176,6 +178,10 @@ static int my_init(Env* env, PyObject* args, PyObject* kwargs) {
     }
     env->action_type = conf.action_type;
     env->dynamics_model = conf.dynamics_model;
+    if (PyDict_GetItemString(kwargs, "dynamics_model")) {
+        char* dynamics_str = unpack_str(kwargs, "dynamics_model");
+        env->dynamics_model = (strcmp(dynamics_str, "jerk") == 0) ? JERK : CLASSIC;
+    }
     env->reward_vehicle_collision = conf.reward_vehicle_collision;
     env->reward_offroad_collision = conf.reward_offroad_collision;
     env->reward_goal = conf.reward_goal;
@@ -188,6 +194,22 @@ static int my_init(Env* env, PyObject* args, PyObject* kwargs) {
     env->offroad_behavior = conf.offroad_behavior;
     env->max_controlled_agents = unpack(kwargs, "max_controlled_agents");
     env->dt = conf.dt;
+
+    // Conditioning parameters
+    env->use_rc = (bool)unpack(kwargs, "use_rc");
+    env->use_ec = (bool)unpack(kwargs, "use_ec");
+    env->use_dc = (bool)unpack(kwargs, "use_dc");
+    env->collision_weight_lb = (float)unpack(kwargs, "collision_weight_lb");
+    env->collision_weight_ub = (float)unpack(kwargs, "collision_weight_ub");
+    env->offroad_weight_lb = (float)unpack(kwargs, "offroad_weight_lb");
+    env->offroad_weight_ub = (float)unpack(kwargs, "offroad_weight_ub");
+    env->goal_weight_lb = (float)unpack(kwargs, "goal_weight_lb");
+    env->goal_weight_ub = (float)unpack(kwargs, "goal_weight_ub");
+    env->entropy_weight_lb = (float)unpack(kwargs, "entropy_weight_lb");
+    env->entropy_weight_ub = (float)unpack(kwargs, "entropy_weight_ub");
+    env->discount_weight_lb = (float)unpack(kwargs, "discount_weight_lb");
+    env->discount_weight_ub = (float)unpack(kwargs, "discount_weight_ub");
+
     env->init_mode = (int)unpack(kwargs, "init_mode");
     env->control_mode = (int)unpack(kwargs, "control_mode");
     int map_id = unpack(kwargs, "map_id");

pufferlib/ocean/drive/drive.c

@@ -19,7 +19,7 @@ void test_drivenet() {
 
     //Weights* weights = load_weights("resources/drive/puffer_drive_weights.bin");
     Weights* weights = load_weights("puffer_drive_weights.bin");
-    DriveNet* net = init_drivenet(weights, num_agents);
+    DriveNet* net = init_drivenet(weights, num_agents, CLASSIC, false, false, false);
 
     forward(net, observations, actions);
     for (int i = 0; i < num_agents*num_actions; i++) {
@@ -50,7 +50,7 @@ void demo() {
         .reward_ade = conf.reward_ade,
         .goal_radius = conf.goal_radius,
         .dt = conf.dt,
-	    .map_name = "resources/drive/binaries/map_000.bin",
+        .map_name = "resources/drive/binaries/map_000.bin",
         .init_steps = conf.init_steps,
         .collision_behavior = conf.collision_behavior,
         .offroad_behavior = conf.offroad_behavior,
@@ -59,7 +59,7 @@ void demo() {
     c_reset(&env);
     c_render(&env);
     Weights* weights = load_weights("resources/drive/puffer_drive_weights.bin");
-    DriveNet* net = init_drivenet(weights, env.active_agent_count, env.dynamics_model);
+    DriveNet* net = init_drivenet(weights, env.active_agent_count, env.dynamics_model, false, false, false);
     //Client* client = make_client(&env);
     int accel_delta = 2;
     int steer_delta = 4;