Early environment resets based on agents' respawn status. (#167)

* Added early termination parameter based on respawn status of all agents in an episode

* pre-commit fix

* fix test

* Apply precommit.

* Reduce variance in aggregate metrics by logging only if we have data for at least num_agents.

---------

Co-authored-by: Daphne Cornelisse <cor.daphne@gmail.com>

Author: riccardosavorgnan

pufferlib/config/ocean/drive.ini

@@ -41,6 +41,7 @@ offroad_behavior = 0
 ; Number of steps before reset
 scenario_length = 91
 resample_frequency = 910
+termination_mode = 1 # 0 - terminate at scenario_length, 1 - terminate after all agents have been reset
 map_dir = "resources/drive/binaries/training"
 num_maps = 10000
 ; Determines which step of the trajectory to initialize the agents at upon reset

pufferlib/ocean/drive/binding.c

@@ -186,6 +186,7 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) {
     env->reward_goal_post_respawn = conf.reward_goal_post_respawn;
     env->reward_ade = conf.reward_ade;
     env->scenario_length = conf.scenario_length;
+    env->termination_mode = conf.termination_mode;
     env->collision_behavior = conf.collision_behavior;
     env->offroad_behavior = conf.offroad_behavior;
     env->max_controlled_agents = unpack(kwargs, "max_controlled_agents");

pufferlib/ocean/drive/drive.h

@@ -324,6 +324,7 @@ struct Drive {
     GridMap *grid_map;
     int *neighbor_offsets;
     int scenario_length;
+    int termination_mode;
     float reward_vehicle_collision;
     float reward_offroad_collision;
     float reward_ade;
@@ -2215,7 +2216,18 @@ void c_step(Drive *env) {
     memset(env->rewards, 0, env->active_agent_count * sizeof(float));
     memset(env->terminals, 0, env->active_agent_count * sizeof(unsigned char));
     env->timestep++;
-    if (env->timestep == env->scenario_length) {
+
+    int originals_remaining = 0;
+    for (int i = 0; i < env->active_agent_count; i++) {
+        int agent_idx = env->active_agent_indices[i];
+        // Keep flag true if there is at least one agent that has not been respawned yet
+        if (env->entities[agent_idx].respawn_count == 0) {
+            originals_remaining = 1;
+            break;
+        }
+    }
+
+    if (env->timestep == env->scenario_length || (!originals_remaining && env->termination_mode == 1)) {
         add_log(env);
         c_reset(env);
         return;

pufferlib/ocean/drive/drive.py

@@ -28,6 +28,7 @@ def __init__(
         offroad_behavior=0,
         dt=0.1,
         scenario_length=None,
+        termination_mode=None,
         resample_frequency=91,
         num_maps=100,
         num_agents=512,
@@ -57,6 +58,7 @@ def __init__(
         self.reward_ade = reward_ade
         self.human_agent_idx = human_agent_idx
         self.scenario_length = scenario_length
+        self.termination_mode = termination_mode
         self.resample_frequency = resample_frequency
         self.dynamics_model = dynamics_model
 
@@ -179,6 +181,7 @@ def __init__(
                 offroad_behavior=self.offroad_behavior,
                 dt=dt,
                 scenario_length=(int(scenario_length) if scenario_length is not None else None),
+                termination_mode=(int(self.termination_mode) if self.termination_mode is not None else 0),
                 max_controlled_agents=self.max_controlled_agents,
                 map_id=map_ids[i],
                 max_agents=nxt - cur,
@@ -204,7 +207,7 @@ def step(self, actions):
         self.tick += 1
         info = []
         if self.tick % self.report_interval == 0:
-            log = binding.vec_log(self.c_envs)
+            log = binding.vec_log(self.c_envs, self.num_agents)
             if log:
                 info.append(log)
                 # print(log)

pufferlib/ocean/env_binding.h

@@ -556,30 +556,48 @@ static int assign_to_dict(PyObject *dict, char *key, float value) {
 }
 
 static PyObject *vec_log(PyObject *self, PyObject *args) {
+    if (PyTuple_Size(args) != 2) {
+        PyErr_SetString(PyExc_TypeError, "vec_log requires 2 arguments");
+        return NULL;
+    }
+
     VecEnv *vec = unpack_vecenv(args);
     if (!vec) {
         return NULL;
     }
 
     // Iterates over logs one float at a time. Will break
     // horribly if Log has non-float data.
+    PyObject *num_agents_arg = PyTuple_GetItem(args, 1);
+    float num_agents = (float)PyLong_AsLong(num_agents_arg);
+
     Log aggregate = {0};
     int num_keys = sizeof(Log) / sizeof(float);
     for (int i = 0; i < vec->num_envs; i++) {
         Env *env = vec->envs[i];
         for (int j = 0; j < num_keys; j++) {
             ((float *)&aggregate)[j] += ((float *)&env->log)[j];
-            ((float *)&env->log)[j] = 0.0f;
         }
     }
 
     PyObject *dict = PyDict_New();
-    if (aggregate.n == 0.0f) {
+
+    // Only log if we have at least num_agents worth of data
+    if (aggregate.n < num_agents) {
         return dict;
     }
 
-    // Average
+    // Got enough data. Reset logs and return metrics
+    for (int i = 0; i < vec->num_envs; i++) {
+        Env *env = vec->envs[i];
+        for (int j = 0; j < num_keys; j++) {
+            ((float *)&env->log)[j] = 0.0f;
+        }
+    }
+
     float n = aggregate.n;
+
+    // Average across agents
     for (int i = 0; i < num_keys; i++) {
         ((float *)&aggregate)[i] /= n;
     }

pufferlib/ocean/env_config.h

@@ -23,6 +23,7 @@ typedef struct {
     float dt;
     int goal_behavior;
     int scenario_length;
+    int termination_mode;
     int init_steps;
     int init_mode;
     int control_mode;
@@ -78,6 +79,8 @@ static int handler(void *config, const char *section, const char *name, const ch
         env_config->dt = atof(value);
     } else if (MATCH("env", "scenario_length")) {
         env_config->scenario_length = atoi(value);
+    } else if (MATCH("env", "termination_mode")) {
+        env_config->termination_mode = atoi(value);
     } else if (MATCH("env", "init_steps")) {
         env_config->init_steps = atoi(value);
     } else if (MATCH("env", "init_mode")) {