diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index f594963bff..1629d4867e 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/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]; + } } } diff --git a/pufferlib/ocean/drive/drive.py b/pufferlib/ocean/drive/drive.py index 360ce4712c..e610d6967b 100644 --- a/pufferlib/ocean/drive/drive.py +++ b/pufferlib/ocean/drive/drive.py @@ -726,7 +726,7 @@ 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. @@ -734,7 +734,11 @@ def set_agent_states(self, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long): 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 ───────────────────────────────────── diff --git a/pufferlib/ocean/env_binding.h b/pufferlib/ocean/env_binding.h index 54e28ee8d8..79deb3ebe6 100644 --- a/pufferlib/ocean/env_binding.h +++ b/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); diff --git a/tests/drive/test_drive_cosim_setters.c b/tests/drive/test_drive_cosim_setters.c index 937d062a63..dc140d4d97 100644 --- a/tests/drive/test_drive_cosim_setters.c +++ b/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);