added reward/entropy/discount conditioning

Author: m2kulkarni

pufferlib/config/ocean/drive.ini

@@ -35,6 +35,7 @@ init_steps = 0 # Determines which step of the trajectory to initialize the agent
 control_all_agents = False # this should be set to false unless you want to specifically want to override and control expert marked vehicles
 num_policy_controlled_agents = -1 # note: if you add this you likely need to set num_agents to a smaller number
 deterministic_agent_selection = False # if this is true it overrides vehicles marked as expert to be policy controlled
+condition_type = "none" # Options: "none", "reward", "entropy", "discount", "all"
 
 [train]
 total_timesteps = 2_000_000_000

pufferlib/extensions/cuda/pufferlib.cu

@@ -20,8 +20,8 @@ __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 dones, torch::Tensor importance, torch::Tensor advantages, 
+        torch::Tensor gammas, int num_steps, int horizon) {
 
     // Validate input tensors
     torch::Device device = values.device();
@@ -35,24 +35,30 @@ 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);
@@ -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

@@ -174,6 +174,22 @@ static int my_init(Env* env, PyObject* args, PyObject* kwargs) {
     env->control_all_agents = unpack(kwargs, "control_all_agents");
     env->deterministic_agent_selection = unpack(kwargs, "deterministic_agent_selection");
     env->control_non_vehicles = (int)unpack(kwargs, "control_non_vehicles");
+
+    // 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");
+
     int map_id = unpack(kwargs, "map_id");
     int max_agents = unpack(kwargs, "max_agents");
     int init_steps = unpack(kwargs, "init_steps");

pufferlib/ocean/drive/drive.c

@@ -11,6 +11,7 @@
 typedef struct DriveNet DriveNet;
 struct DriveNet {
     int num_agents;
+    int conditioning_dims;
     float* obs_self;
     float* obs_partner;
     float* obs_road;
@@ -42,13 +43,20 @@ struct DriveNet {
     Multidiscrete* multidiscrete;
 };
 
-DriveNet* init_drivenet(Weights* weights, int num_agents) {
+DriveNet* init_drivenet(Weights* weights, int num_agents, bool use_rc, bool use_ec, bool use_dc) {
     DriveNet* net = calloc(1, sizeof(DriveNet));
     int hidden_size = 256;
     int input_size = 64;
 
     net->num_agents = num_agents;
-    net->obs_self = calloc(num_agents*7, sizeof(float)); // 7 features
+    net->conditioning_dims = (use_rc ? 3 : 0) + (use_ec ? 1 : 0) + (use_dc ? 1 : 0);
+
+    int ego_obs_size = 7; // base features
+    if (use_rc) ego_obs_size += 3; // reward conditioning
+    if (use_ec) ego_obs_size += 1; // entropy conditioning
+    if (use_dc) ego_obs_size += 1; // discount conditioning
+
+    net->obs_self = calloc(num_agents*ego_obs_size, sizeof(float));
     net->obs_partner = calloc(num_agents*63*7, sizeof(float)); // 63 objects, 7 features
     net->obs_road = calloc(num_agents*200*13, sizeof(float)); // 200 objects, 13 features
     net->partner_linear_output = calloc(num_agents*63*input_size, sizeof(float));
@@ -57,7 +65,7 @@ DriveNet* init_drivenet(Weights* weights, int num_agents) {
     net->road_linear_output_two = calloc(num_agents*200*input_size, sizeof(float));
     net->partner_layernorm_output = calloc(num_agents*63*input_size, sizeof(float));
     net->road_layernorm_output = calloc(num_agents*200*input_size, sizeof(float));
-    net->ego_encoder = make_linear(weights, num_agents, 7, input_size);
+    net->ego_encoder = make_linear(weights, num_agents, ego_obs_size, input_size);
     net->ego_layernorm = make_layernorm(weights, num_agents, input_size);
     net->ego_encoder_two = make_linear(weights, num_agents, input_size, input_size);
     net->road_encoder = make_linear(weights, num_agents, 13, input_size);
@@ -117,23 +125,25 @@ void free_drivenet(DriveNet* net) {
 }
 
 void forward(DriveNet* net, float* observations, int* actions) {
+    int ego_obs_size = 7 + net->conditioning_dims;
+
     // Clear previous observations
-    memset(net->obs_self, 0, net->num_agents * 7 * sizeof(float));
+    memset(net->obs_self, 0, net->num_agents * ego_obs_size * sizeof(float));
     memset(net->obs_partner, 0, net->num_agents * 63 * 7 * sizeof(float));
     memset(net->obs_road, 0, net->num_agents * 200 * 13 * sizeof(float));
 
     // Reshape observations into 2D boards and additional features
-    float (*obs_self)[7] = (float (*)[7])net->obs_self;
+    float* obs_self = net->obs_self;
     float (*obs_partner)[63][7] = (float (*)[63][7])net->obs_partner;
     float (*obs_road)[200][13] = (float (*)[200][13])net->obs_road;
 
     for (int b = 0; b < net->num_agents; b++) {
-        int b_offset = b * (7 + 63*7 + 200*7);  // offset for each batch
-        int partner_offset = b_offset + 7;
-        int road_offset = b_offset + 7 + 63*7;
+        int b_offset = b * (ego_obs_size + 63*7 + 200*7);  // offset for each batch
+        int partner_offset = b_offset + ego_obs_size;
+        int road_offset = b_offset + ego_obs_size + 63*7;
         // Process self observation
-        for(int i = 0; i < 7; i++) {
-            obs_self[b][i] = observations[b_offset + i];
+        for(int i = 0; i < ego_obs_size; i++) {
+            obs_self[b*ego_obs_size + i] = observations[b_offset + i];
         }
 
         // Process partner observation