diff --git a/pufferlib/ocean/drive/binding.c b/pufferlib/ocean/drive/binding.c index f178fedfc..2fe4ab096 100644 --- a/pufferlib/ocean/drive/binding.c +++ b/pufferlib/ocean/drive/binding.c @@ -335,6 +335,7 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) { env->collision_behavior = conf.collision_behavior; env->offroad_behavior = conf.offroad_behavior; env->dt = conf.dt; + env->randomize_respawn = (int)unpack(kwargs, "randomize_respawn"); env->init_mode = (int)unpack(kwargs, "init_mode"); env->control_mode = (int)unpack(kwargs, "control_mode"); env->goal_behavior = (int)unpack(kwargs, "goal_behavior"); diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index fc4aa2bfc..796008ad9 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -347,6 +347,7 @@ struct Drive { char scenario_id[SCENARIO_ID_STR_LENGTH]; int collision_behavior; int offroad_behavior; + int randomize_respawn; float observation_window_size; float polyline_reduction_threshold; float polyline_max_segment_length; @@ -2824,16 +2825,105 @@ void compute_observations(Drive *env) { } } +// Find a random collision-free position on a drivable lane for an existing agent. +// Returns true if a valid position was found and updates the agent's sim_x/y/z/heading. +static bool randomize_agent_position(Drive *env, int agent_idx) { + Agent *agent = &env->agents[agent_idx]; + + // Pre-compute drivable lanes + int drivable_lanes[env->num_roads]; + float lane_lengths[env->num_roads]; + int num_drivable = 0; + float total_lane_length = 0.0f; + for (int i = 0; i < env->num_roads; i++) { + if (env->road_elements[i].type == ROAD_LANE && env->road_elements[i].polyline_length > 0.0f) { + drivable_lanes[num_drivable] = i; + lane_lengths[num_drivable] = env->road_elements[i].polyline_length; + total_lane_length += lane_lengths[num_drivable]; + num_drivable++; + } + } + + if (num_drivable == 0) + return false; + + for (int attempt = 0; attempt < MAX_SPAWN_ATTEMPTS; attempt++) { + // Length-weighted lane selection + float r = ((float)rand() / (float)RAND_MAX) * total_lane_length; + float cumulative = 0.0f; + int selected = num_drivable - 1; + for (int k = 0; k < num_drivable; k++) { + cumulative += lane_lengths[k]; + if (r < cumulative) { + selected = k; + break; + } + } + RoadMapElement *lane = &env->road_elements[drivable_lanes[selected]]; + + float spawn_x, spawn_y, spawn_z, spawn_heading; + get_random_point_on_lane(lane, &spawn_x, &spawn_y, &spawn_z, &spawn_heading); + spawn_z += agent->sim_height / 2.0f; + + // Temporarily invalidate this agent so check_spawn_collision skips it + float saved_x = agent->sim_x; + agent->sim_x = INVALID_POSITION; + bool collision = check_spawn_collision(env, env->active_agent_count, spawn_x, spawn_y, spawn_z, spawn_heading, + agent->sim_length, agent->sim_width, agent->sim_height); + agent->sim_x = saved_x; + if (collision) + continue; + + // Check offroad + if (check_spawn_offroad(env, spawn_x, spawn_y, spawn_z, spawn_heading, agent->sim_length, agent->sim_width, + agent->sim_height)) + continue; + + agent->sim_x = spawn_x; + agent->sim_y = spawn_y; + agent->sim_z = spawn_z; + agent->sim_heading = spawn_heading; + agent->heading_x = cosf(spawn_heading); + agent->heading_y = sinf(spawn_heading); + // Update stored initial position so future non-random resets are consistent + agent->log_trajectory_x[0] = spawn_x; + agent->log_trajectory_y[0] = spawn_y; + agent->log_trajectory_z[0] = spawn_z; + agent->log_heading[0] = spawn_heading; + return true; + } + return false; +} + void respawn_agent(Drive *env, int agent_idx) { Agent *agent = &env->agents[agent_idx]; - agent->sim_x = agent->log_trajectory_x[0]; - agent->sim_y = agent->log_trajectory_y[0]; - agent->sim_z = agent->log_trajectory_z[0]; - agent->sim_heading = agent->log_heading[0]; - agent->heading_x = cosf(agent->sim_heading); - agent->heading_y = sinf(agent->sim_heading); - agent->sim_vx = agent->log_velocity_x[0]; - agent->sim_vy = agent->log_velocity_y[0]; + + if (env->randomize_respawn) { + if (!randomize_agent_position(env, agent_idx)) { + // Fallback to original position if no valid spawn found + agent->sim_x = agent->log_trajectory_x[0]; + agent->sim_y = agent->log_trajectory_y[0]; + agent->sim_z = agent->log_trajectory_z[0]; + agent->sim_heading = agent->log_heading[0]; + agent->heading_x = cosf(agent->sim_heading); + agent->heading_y = sinf(agent->sim_heading); + } + // Sample a new goal relative to the new position + sample_new_goal(env, agent_idx); + agent->sim_vx = 0.0f; + agent->sim_vy = 0.0f; + agent->sim_speed = 0.0f; + agent->sim_speed_signed = 0.0f; + } else { + agent->sim_x = agent->log_trajectory_x[0]; + agent->sim_y = agent->log_trajectory_y[0]; + agent->sim_z = agent->log_trajectory_z[0]; + agent->sim_heading = agent->log_heading[0]; + agent->heading_x = cosf(agent->sim_heading); + agent->heading_y = sinf(agent->sim_heading); + agent->sim_vx = agent->log_velocity_x[0]; + agent->sim_vy = agent->log_velocity_y[0]; + } agent->metrics_array[COLLISION_IDX] = 0.0f; agent->metrics_array[OFFROAD_IDX] = 0.0f; agent->metrics_array[REACHED_GOAL_IDX] = 0.0f; @@ -3097,8 +3187,21 @@ void move_dynamics(Drive *env, int action_idx, int agent_idx) { void c_reset(Drive *env) { env->timestep = env->init_steps; - set_start_position(env); - reset_goal_positions(env); + if (env->randomize_respawn) { + // Randomize all agent positions on reset + for (int x = 0; x < env->active_agent_count; x++) { + int agent_idx = env->active_agent_indices[x]; + randomize_agent_position(env, agent_idx); + } + // Sample new goals relative to new positions + for (int x = 0; x < env->active_agent_count; x++) { + int agent_idx = env->active_agent_indices[x]; + sample_new_goal(env, agent_idx); + } + } else { + set_start_position(env); + reset_goal_positions(env); + } for (int x = 0; x < env->active_agent_count; x++) { env->logs[x] = (Log){0}; int agent_idx = env->active_agent_indices[x]; @@ -3128,7 +3231,7 @@ void c_reset(Drive *env) { agent->prev_goal_z = agent->sim_z; generate_reward_coefs(env, agent); - if (env->goal_behavior == GOAL_GENERATE_NEW) { + if (env->goal_behavior == GOAL_GENERATE_NEW && env->init_mode != INIT_VARIABLE_AGENT_NUMBER) { agent->goal_position_x = agent->init_goal_x; agent->goal_position_y = agent->init_goal_y; agent->goal_position_z = agent->init_goal_z; diff --git a/pufferlib/ocean/drive/drive.py b/pufferlib/ocean/drive/drive.py index ec13480da..cef4e2117 100644 --- a/pufferlib/ocean/drive/drive.py +++ b/pufferlib/ocean/drive/drive.py @@ -95,6 +95,7 @@ def __init__( spawn_length_min=2.0, spawn_length_max=5.5, spawn_height=1.5, + randomize_respawn=1, ): # env self.dt = dt @@ -125,6 +126,7 @@ def __init__( self.episode_length = episode_length self.termination_mode = termination_mode self.resample_frequency = resample_frequency + self.randomize_respawn = randomize_respawn self.dynamics_model = dynamics_model # reward randomization bounds self.reward_bound_goal_radius_min = reward_bound_goal_radius_min @@ -416,6 +418,7 @@ def __init__( spawn_length_min=self.spawn_length_min, spawn_length_max=self.spawn_length_max, spawn_height=self.spawn_height, + randomize_respawn=self.randomize_respawn, ) env_ids.append(env_id) @@ -573,6 +576,7 @@ def resample_maps(self): spawn_length_min=self.spawn_length_min, spawn_length_max=self.spawn_length_max, spawn_height=self.spawn_height, + randomize_respawn=self.randomize_respawn, ) env_ids.append(env_id) self.c_envs = binding.vectorize(*env_ids) @@ -799,9 +803,8 @@ def save_map_binary(map_data, output_file, unique_map_id): sdc_track_index = metadata.get("sdc_track_index", -1) # -1 as default if not found tracks_to_predict = metadata.get("tracks_to_predict", []) - # Write original scenario_id with fallback to placeholder - scenario_id = map_data.get("scenario_id", f"map_{unique_map_id:03d}") - f.write(struct.pack("16s", scenario_id.encode("utf-8"))) + # Note: C load_map_binary does NOT read a scenario_id prefix. + # Do not write one here or the binary will be misaligned. # Write sdc_track_index f.write(struct.pack("i", sdc_track_index)) diff --git a/tests/test_randomize_respawn.py b/tests/test_randomize_respawn.py new file mode 100644 index 000000000..d7b569bf1 --- /dev/null +++ b/tests/test_randomize_respawn.py @@ -0,0 +1,118 @@ +"""Test that randomize_respawn produces different agent positions across resets. + +Run on the cluster with: + srun ... python -m pytest tests/test_randomize_respawn.py -v +""" + +import numpy as np +import pytest +from pufferlib.ocean.drive.drive import Drive + + +MAP_DIR = "pufferlib/resources/drive/binaries/carla_data" + + +def get_agent_positions(env): + """Extract current agent positions from observations.""" + # Ego obs starts at index 0: sim_x, sim_y are the first features + # But obs are in ego frame (normalized). Use the C env directly. + # The simplest proxy: just hash the full observation vector. + return env.observations.copy() + + +@pytest.fixture +def env_randomize(): + e = Drive( + num_agents=8, + num_maps=2, + map_dir=MAP_DIR, + dynamics_model="classic", + min_agents_per_env=1, + max_agents_per_env=8, + init_mode="init_variable_agent_number", + control_mode="control_vehicles", + episode_length=300, + resample_frequency=0, + randomize_respawn=1, + ) + e.reset() + yield e + e.close() + + +@pytest.fixture +def env_no_randomize(): + e = Drive( + num_agents=8, + num_maps=2, + map_dir=MAP_DIR, + dynamics_model="classic", + min_agents_per_env=1, + max_agents_per_env=8, + init_mode="init_variable_agent_number", + control_mode="control_vehicles", + episode_length=10, + resample_frequency=0, + randomize_respawn=0, + ) + e.reset() + yield e + e.close() + + +def test_randomize_respawn_produces_different_positions(env_randomize): + """With randomize_respawn=1, positions should differ after episode reset.""" + env = env_randomize + actions = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + + # Get initial observations + obs_before = env.observations.copy() + + # Step until episode resets (episode_length=300, or force via resample) + env.resample_maps() + obs_after = env.observations.copy() + + # Observations should differ (agents at different positions) + assert not np.allclose(obs_before, obs_after, atol=1e-6), ( + "Observations should differ after reset with randomize_respawn=1" + ) + + +def test_no_randomize_same_positions(env_no_randomize): + """With randomize_respawn=0, positions should be the same after episode reset.""" + env = env_no_randomize + actions = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + + # Get initial observations + obs_before = env.observations.copy() + + # Step through the full episode to trigger c_reset + for _ in range(15): + env.step(actions) + + # After reset, positions should return to initial state + # Note: obs won't be exactly the same due to metrics/counters, + # but the position-related features should match + obs_after = env.observations.copy() + + # With no randomization, the first few ego features (position-related) + # should be identical after reset + # ego features: speed, heading components, goal direction, etc. + # After a full reset with no randomization, agents return to log_trajectory[0] + assert np.allclose(obs_before[:, :5], obs_after[:, :5], atol=0.1), ( + "Position features should be similar after reset with randomize_respawn=0" + ) + + +def test_multiple_resets_produce_variety(env_randomize): + """Multiple resets with randomize_respawn should produce different positions each time.""" + env = env_randomize + observations = [] + + for _ in range(5): + env.resample_maps() + observations.append(env.observations[:, :10].copy()) + + # Check that not all resets produce the same observations + all_same = all(np.allclose(observations[0], obs, atol=1e-6) for obs in observations[1:]) + assert not all_same, "5 resets with randomize_respawn should not all produce identical observations"