feat(cosim): optional seconds_stopped in set_agent_states

pufferlib/ocean/drive/drive.h

@@ -3244,7 +3244,8 @@ void c_set_agent_states(
     const float *vx,
     const float *vy,
     const float *yaw_rate,
-    const float *accel_long) {
+    const float *accel_long,
+    const float *seconds_stopped) {
     for (int k = 0; k < count; k++) {
         int agent_idx = idx[k];
         if (agent_idx < 0 || agent_idx >= env->num_total_agents) {
@@ -3283,8 +3284,10 @@ void c_set_agent_states(
         agent->accel_lat = agent->sim_speed_signed * agent->yaw_rate;
         refresh_lane_association(env, agent); // current_lane_idx / lane-dist / lane-angle for the new pose
 
-        // seconds_stopped is intentionally left untouched: c_step already updates it once per tick, before this
-        // overwrite runs, so redoing it here would double-count.
+        // NULL keeps c_step's own accumulation (co-sim default); an array injects stopped-time as state.
+        if (seconds_stopped) {
+            agent->seconds_stopped = seconds_stopped[k];
+        }
     }
 }
 

pufferlib/ocean/drive/drive.py

@@ -726,15 +726,19 @@ def get_global_agent_state(self):
 
         return states
 
-    def set_agent_states(self, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long):
+    def set_agent_states(self, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long, seconds_stopped=None):
         """Co-sim: overwrite the sim state of agents at global indices `idx`
         (e.g. CARLA background) with world-frame pose/velocity. The C side
         subtracts world_mean, recaches heading trig and recomputes speed.
         `yaw_rate` (rad/s) and `accel_long` (m/s^2) must come from the external
         sim's own physics (e.g. CARLA's get_angular_velocity()/get_acceleration(),
         nuPlan's EgoState.dynamic_car_state) -- not finite-differenced here, since
         this agent's previous state may already have been overwritten this tick by
-        env.step()'s own dynamics before this call runs."""
+        env.step()'s own dynamics before this call runs.
+        `seconds_stopped` (seconds each agent's speed has been below the stopped
+        threshold) is optional: pass it to inject stopped-time as state when this
+        env when the main driving agent is not the one being simulated in pufferdrive, 
+        or leave it None to keep c_step's own per-tick accumulation."""
         binding.vec_set_agent_states(
             self.c_envs,
             np.ascontiguousarray(idx, dtype=np.int32),
@@ -746,6 +750,7 @@ def set_agent_states(self, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long):
             np.ascontiguousarray(vy, dtype=np.float32),
             np.ascontiguousarray(yaw_rate, dtype=np.float32),
             np.ascontiguousarray(accel_long, dtype=np.float32),
+            None if seconds_stopped is None else np.ascontiguousarray(seconds_stopped, dtype=np.float32),
         )
 
     # ── Co-simulation external-state setters ─────────────────────────────────────

pufferlib/ocean/env_binding.h

@@ -1025,8 +1025,8 @@ static PyObject *vec_get_global_agent_state(PyObject *self, PyObject *args) {
 // ── Co-simulation external-state setters (mirror vec_get_global_agent_state) ──
 // Co-sim runs a single env (num_envs == 1); these operate on vec->envs[0].
 static PyObject *vec_set_agent_states(PyObject *self, PyObject *args) {
-    if (PyTuple_Size(args) != 10) {
-        PyErr_SetString(PyExc_TypeError, "vec_set_agent_states requires 10 arguments");
+    if (PyTuple_Size(args) != 11) {
+        PyErr_SetString(PyExc_TypeError, "vec_set_agent_states requires 11 arguments");
         return NULL;
     }
     VecEnv *vec = unpack_vecenv(args);
@@ -1042,12 +1042,17 @@ static PyObject *vec_set_agent_states(PyObject *self, PyObject *args) {
     PyObject *vy_arr = PyTuple_GetItem(args, 7);
     PyObject *yaw_rate_arr = PyTuple_GetItem(args, 8);
     PyObject *accel_long_arr = PyTuple_GetItem(args, 9);
+    PyObject *seconds_stopped_arr = PyTuple_GetItem(args, 10);  // None to keep c_step's accumulation
     if (!PyArray_Check(idx_arr) || !PyArray_Check(x_arr) || !PyArray_Check(y_arr) || !PyArray_Check(z_arr)
         || !PyArray_Check(heading_arr) || !PyArray_Check(vx_arr) || !PyArray_Check(vy_arr)
         || !PyArray_Check(yaw_rate_arr) || !PyArray_Check(accel_long_arr)) {
         PyErr_SetString(PyExc_TypeError, "All arrays must be NumPy arrays");
         return NULL;
     }
+    if (seconds_stopped_arr != Py_None && !PyArray_Check(seconds_stopped_arr)) {
+        PyErr_SetString(PyExc_TypeError, "seconds_stopped must be a NumPy array or None");
+        return NULL;
+    }
     int *idx = (int *) PyArray_DATA((PyArrayObject *) idx_arr);
     float *x = (float *) PyArray_DATA((PyArrayObject *) x_arr);
     float *y = (float *) PyArray_DATA((PyArrayObject *) y_arr);
@@ -1057,8 +1062,11 @@ static PyObject *vec_set_agent_states(PyObject *self, PyObject *args) {
     float *vy = (float *) PyArray_DATA((PyArrayObject *) vy_arr);
     float *yaw_rate = (float *) PyArray_DATA((PyArrayObject *) yaw_rate_arr);
     float *accel_long = (float *) PyArray_DATA((PyArrayObject *) accel_long_arr);
+    float *seconds_stopped =
+        seconds_stopped_arr == Py_None ? NULL : (float *) PyArray_DATA((PyArrayObject *) seconds_stopped_arr);
     int count = (int) PyArray_SIZE((PyArrayObject *) idx_arr);
-    c_set_agent_states((Drive *) vec->envs[0], count, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long);
+    c_set_agent_states(
+        (Drive *) vec->envs[0], count, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long, seconds_stopped);
     Py_RETURN_NONE;
 }
 
@@ -1530,6 +1538,7 @@ PyMODINIT_FUNC PyInit_binding(void) {
     // post-step speed/accel intent back out of the observation row.
     PyModule_AddObject(m, "MAX_SPEED", PyFloat_FromDouble(MAX_SPEED));
     PyModule_AddObject(m, "ACCEL_LONG_NORM", PyFloat_FromDouble(fabsf(ACCEL_LONG_LIMIT[0])));
+    PyModule_AddObject(m, "AGENT_STOPPED_SPEED_THRESHOLD", PyFloat_FromDouble(AGENT_STOPPED_SPEED_THRESHOLD));
     PyModule_AddIntConstant(m, "ACTION_TYPE_DISCRETE", ACTION_TYPE_DISCRETE);
     PyModule_AddIntConstant(m, "ACTION_TYPE_CONTINUOUS", ACTION_TYPE_CONTINUOUS);
     PyModule_AddIntConstant(m, "DYNAMICS_MODEL_CLASSIC", DYNAMICS_MODEL_CLASSIC);

tests/drive/test_drive_cosim_setters.c

@@ -130,7 +130,7 @@ static int test_set_agent_states_teleport_resets_prev_pose(void) {
     int idx[1] = {0};
     float x[1] = {10.0f}, y[1] = {20.0f}, z[1] = {0.0f}, h[1] = {0.5f};
     float vx[1] = {1.0f}, vy[1] = {0.0f}, yr[1] = {0.0f}, al[1] = {0.0f};
-    c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al);
+    c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, NULL);
 
     EXPECT_NEAR(agent.prev_x, agent.sim_x, 1e-6f);
     EXPECT_NEAR(agent.prev_y, agent.sim_y, 1e-6f);
@@ -139,6 +139,28 @@ static int test_set_agent_states_teleport_resets_prev_pose(void) {
     return 0;
 }
 
+static int test_set_agent_states_seconds_stopped_injects_or_preserves(void) {
+    // seconds_stopped is optional: a NULL array leaves c_step's own accumulation untouched
+    // (the co-sim default CARLA/nuplan rely on), a non-NULL array injects it as state.
+    Drive env = {0};
+    Agent agent = drive_test_agent(0.0f, 0.0f, 0.0f);
+    agent.seconds_stopped = 12.5f;
+    env.agents = &agent;
+    env.num_total_agents = 1;
+
+    int idx[1] = {0};
+    float x[1] = {1.0f}, y[1] = {2.0f}, z[1] = {0.0f}, h[1] = {0.0f};
+    float vx[1] = {0.0f}, vy[1] = {0.0f}, yr[1] = {0.0f}, al[1] = {0.0f};
+
+    c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, NULL);
+    EXPECT_NEAR(agent.seconds_stopped, 12.5f, 1e-6f); // NULL: untouched
+
+    float seconds_stopped[1] = {3.0f};
+    c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, seconds_stopped);
+    EXPECT_NEAR(agent.seconds_stopped, 3.0f, 1e-6f); // array: overwritten
+    return 0;
+}
+
 // ---------------------------------------------------------------------------
 // c_set_agent_goals
 // ---------------------------------------------------------------------------
@@ -226,6 +248,7 @@ int main(void) {
     RUN_TEST(test_set_traffic_light_states_writes_current_timestep_for_lights_only);
     RUN_TEST(test_set_traffic_light_states_skips_out_of_range_or_missing_states);
     RUN_TEST(test_set_agent_states_teleport_resets_prev_pose);
+    RUN_TEST(test_set_agent_states_seconds_stopped_injects_or_preserves);
     RUN_TEST(test_set_agent_goals_sets_positions_lane_and_count);
     RUN_TEST(test_set_agent_goals_caps_at_max_goals);
     RUN_TEST(test_set_agent_goals_out_of_range_agent_idx_is_noop);