diff --git a/notebooks/05_inference.py b/notebooks/05_inference.py index 17122e6712..4f0656c824 100644 --- a/notebooks/05_inference.py +++ b/notebooks/05_inference.py @@ -247,11 +247,21 @@ def run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE, horizon=HORI ego_ts = np.array(ego_features_over_time) if dyn_model == "jerk": - labels = ["speed", "width", "length", "steering", "accel_long", "accel_lat", "lcenter", "lalign", "speed_limit"] + labels = [ + "speed", + "width", + "length", + "steering", + "accel_long", + "accel_lat", + "lcenter", + "lane_heading_error", + "speed_limit", + ] plot_idxs = [0, 3, 4, 5] # speed, steering, accel_long, accel_lat else: - labels = ["speed", "width", "length", "lcenter", "lalign", "speed_limit"] - plot_idxs = [0, 3, 4, 5] # speed, lcenter, lalign, speed_limit + labels = ["speed", "width", "length", "lcenter", "lane_heading_error", "speed_limit"] + plot_idxs = [0, 3, 4, 5] # speed, lcenter, lane_heading_error, speed_limit fig, axes = plt.subplots(len(plot_idxs), 1, figsize=(14, 3 * len(plot_idxs)), sharex=True) for i, idx in enumerate(plot_idxs): @@ -268,7 +278,7 @@ def run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE, horizon=HORI # ## Observation layer breakdown # # Obs layout (all ego-centric, normalized): -# - **Ego**: speed, width, length, [jerk: steering, accel_long, accel_lat], lane_center_dist, lane_angle, speed_limit +# - **Ego**: speed, width, length, [jerk: steering, accel_long, accel_lat], lane_center_dist, lane_heading_error, speed_limit # - **Conditioning** (if enabled): 17 reward coefs (goal_radius, goal_speed, collision, offroad, comfort, lane_align, vel_align, lane_center, center_bias, velocity, reverse, stop_line, timestep, overspeed, throttle, steer, acc) + target waypoints # - **Target**: static=rel_x,rel_y,rel_z per waypoint; dynamic=rel_x,rel_y,rel_z,heading_cos,heading_sin per waypoint # - **Partners** (MAX_PARTNERS x 9): rel_x, rel_y, rel_z, length, width, heading_cos, heading_sin, sim_speed_signed, seconds_stopped diff --git a/notebooks/notebook_utils.py b/notebooks/notebook_utils.py index 05a7091ef3..fcda675959 100644 --- a/notebooks/notebook_utils.py +++ b/notebooks/notebook_utils.py @@ -50,7 +50,7 @@ "accel_long", "accel_lat", "lane_center", - "lane_align", + "lane_heading_error", "speed_limit", "stopped", ] diff --git a/pufferlib/ocean/drive/constants.h b/pufferlib/ocean/drive/constants.h index 1160a52a80..95f9b3e4c3 100644 --- a/pufferlib/ocean/drive/constants.h +++ b/pufferlib/ocean/drive/constants.h @@ -258,26 +258,27 @@ static const int ROAD_OFFSETS[25][2] // ===================================================================================== // Indices into Agent.metrics_array; NUM_METRICS must stay equal to the index count -#define NUM_METRICS 18 -#define COLLISION_IDX 0 -#define OFFROAD_IDX 1 -#define RED_LIGHT_IDX 2 -#define STOP_SIGN_IDX 3 -#define REACHED_GOAL_IDX 4 -#define LANE_DIST_IDX 5 -#define LANE_ANGLE_IDX 6 -#define COMFORT_VIOLATION_IDX 7 -#define VELOCITY_PROGRESS_IDX 8 -#define SPEED_LIMIT_IDX 9 -#define AVG_DISPLACEMENT_ERROR_IDX 10 -#define PROGRESSION_IDX 11 +#define NUM_METRICS 19 +#define COLLISION_IDX 0 // Vehicle collision flag, {0, 1} +#define OFFROAD_IDX 1 // Off-road flag, {0, 1} +#define RED_LIGHT_IDX 2 // Red-light violation flag, {0, 1} +#define STOP_SIGN_IDX 3 // Stop-sign violation flag, {0, 1} +#define REACHED_GOAL_IDX 4 // Goal reached this step, {0, 1} +#define LANE_DIST_IDX 5 // Signed lane-center offset, meters +#define LANE_HEADING_ERROR_RADIANS_IDX 6 // Signed heading error, radians +#define LANE_HEADING_COSINE_IDX 7 // cos(heading error), [-1, 1] +#define COMFORT_VIOLATION_IDX 8 // Acceleration and jerk violation count +#define VELOCITY_PROGRESS_IDX 9 // Forward lane alignment, [0, 1] +#define SPEED_LIMIT_IDX 10 // Speed-limit violation flag, {0, 1} +#define AVG_DISPLACEMENT_ERROR_IDX 11 // Running replay displacement error, meters +#define PROGRESSION_IDX 12 // Reserved route-progression metric // Evaluation metrics -#define AT_FAULT_COLLISION_IDX 12 -#define TTC_IDX 13 -#define DISTANCE_TO_COLLISION_IDX 14 -#define PROGRESS_RATIO_IDX 15 -#define MULTI_LANE_TIME_IDX 16 -#define MULTI_LANE_SCORE_IDX 17 +#define AT_FAULT_COLLISION_IDX 13 // At-fault collision flag, {0, 1} +#define TTC_IDX 14 // Minimum vehicle time-to-collision, seconds +#define DISTANCE_TO_COLLISION_IDX 15 // Minimum vehicle collision distance, meters +#define PROGRESS_RATIO_IDX 16 // Reserved episode-progress ratio +#define MULTI_LANE_TIME_IDX 17 // Cumulative multi-lane occupancy, seconds +#define MULTI_LANE_SCORE_IDX 18 // Multi-lane compliance score, {0, 0.5, 1} // Time to collision #define DEFAULT_TTC 5.0f // when no vehicle ahead diff --git a/pufferlib/ocean/drive/datatypes.h b/pufferlib/ocean/drive/datatypes.h index b58b5a35b6..f7350db1ee 100644 --- a/pufferlib/ocean/drive/datatypes.h +++ b/pufferlib/ocean/drive/datatypes.h @@ -80,9 +80,10 @@ struct Agent { int current_route_idx; // Tracks progress through route array // Metrics and status tracking (size must match NUM_METRICS in drive.h) - float metrics_array[NUM_METRICS]; // [collision, offroad, red_light, stop_sign, reached_goal, lane_dist, lane_angle, - // comfort_violation, velocity_progress, speed_limit, avg_displacement_error, - // progression, at_fault_collision, ttc, distance_to_collision, progress_ratio, + float metrics_array[NUM_METRICS]; // [collision, offroad, red_light, stop_sign, reached_goal, lane_dist, + // lane_heading_alignment, lane_heading_error, comfort_violation, + // velocity_progress, speed_limit, avg_displacement_error, progression, + // at_fault_collision, ttc, distance_to_collision, progress_ratio, // multi_lane_time, multi_lane_score] int current_lane_idx; int previous_lane_idx; diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index 50eec3d0d5..091f56705c 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -3274,10 +3274,14 @@ static void compute_metrics(Drive *env, int agent_idx, int log_idx) { reset_agent_metrics(env, agent_idx); if (agent->sim_x == INVALID_POSITION) { + agent->previous_lane_idx = agent->current_lane_idx; + agent->current_lane_idx = -1; return; // invalid agent position } if (get_grid_index(env, agent->sim_x, agent->sim_y) == -1) { // Current agent is offgrid, treat as offroad + agent->previous_lane_idx = agent->current_lane_idx; + agent->current_lane_idx = -1; agent->metrics_array[OFFROAD_IDX] = 1.0f; apply_infraction_behavior(agent, env->offroad_behavior); return; @@ -3469,15 +3473,17 @@ static void compute_metrics(Drive *env, int agent_idx, int log_idx) { if (env->compute_eval_metrics && edge_dist > MULTI_LANE_THRESHOLD && agent->sim_speed > 0.0f) { agent_log->multi_lane_time += env->dt; } - // theta_f = angle relative to lane heading + // theta_f = heading difference between agent and lane (left = negative, right = positive) float theta_f = compute_heading_diff(agent->sim_heading, lane_heading); - agent->metrics_array[LANE_ANGLE_IDX] = cosf(theta_f); // Store cos(θ_f) + agent->metrics_array[LANE_HEADING_COSINE_IDX] = cosf(theta_f); + agent->metrics_array[LANE_HEADING_ERROR_RADIANS_IDX] = theta_f; } else { // Agent not on any lane agent->previous_lane_idx = -1; agent->current_lane_idx = -1; agent->metrics_array[LANE_DIST_IDX] = LANE_DISTANCE_NORMALIZATION; // Max distance (far from lane) - agent->metrics_array[LANE_ANGLE_IDX] = 0.0f; // Perpendicular (no alignment) + agent->metrics_array[LANE_HEADING_COSINE_IDX] = 0.0f; + agent->metrics_array[LANE_HEADING_ERROR_RADIANS_IDX] = 0.0f; } // Update cumulative metrics @@ -3499,7 +3505,7 @@ static void compute_metrics(Drive *env, int agent_idx, int log_idx) { // Velocity metric - forward progress aligned with lane const float VELOCITY_MIN_SPEED = 2.5f; // m/s if (agent->sim_speed_signed > VELOCITY_MIN_SPEED && lane_idx != -1) { - float cos_theta = agent->metrics_array[LANE_ANGLE_IDX]; + float cos_theta = agent->metrics_array[LANE_HEADING_COSINE_IDX]; agent->metrics_array[VELOCITY_PROGRESS_IDX] = fmaxf(cos_theta, 0.0f); if (env->compute_eval_metrics && cos_theta < 0.0f) { agent_log->wrong_way_distance += agent->sim_speed_signed * env->dt; @@ -3637,9 +3643,11 @@ static void compute_rewards(Drive *env, int i) { agent_log->reward_goal += reward_goal; } - // Get lane angle metric: cos(θ_f) where θ_f = heading diff from lane - float cos_theta = agent->metrics_array[LANE_ANGLE_IDX]; - float theta_f = acosf(fminf(fmaxf(cos_theta, -1.0f), 1.0f)); // Get |θ_f| from cos + float theta_f = M_PI / 2.0f; + float cos_theta = agent->metrics_array[LANE_HEADING_COSINE_IDX]; + if (agent->current_lane_idx != -1) { + theta_f = fabsf(agent->metrics_array[LANE_HEADING_ERROR_RADIANS_IDX]); + } agent_log->lane_heading_aligned_rate += (cos_theta >= LANE_ALIGN_COS_THRESHOLD) ? 1.0f : 0.0f; // Rl-align: min(cos,0) + vel_align*min(cos*v,0) + 0.0025*(1-|θ|/(π/2)) @@ -3765,7 +3773,7 @@ static int write_ego_obs(Drive *env, Agent *ego, float *obs, int obs_idx) { obs[obs_idx++] = ego->accel_long / fabsf(ACCEL_LONG_LIMIT[0]); obs[obs_idx++] = ego->accel_lat / ACCEL_LAT_LIMIT[1]; obs[obs_idx++] = fmaxf(-1.0f, fminf(1.0f, ego->metrics_array[LANE_DIST_IDX] / LANE_DISTANCE_NORMALIZATION)); - obs[obs_idx++] = ego->metrics_array[LANE_ANGLE_IDX]; + obs[obs_idx++] = ego->metrics_array[LANE_HEADING_ERROR_RADIANS_IDX] / M_PI; float current_lane_speed_limit = (ego->current_lane_idx != -1) ? env->road_elements[ego->current_lane_idx].speed_limit : -1.0f; obs[obs_idx++] = current_lane_speed_limit / env->obs_norm_speed_mps; diff --git a/pufferlib/viz.py b/pufferlib/viz.py index e6f547f86b..06183036aa 100644 --- a/pufferlib/viz.py +++ b/pufferlib/viz.py @@ -62,7 +62,8 @@ "stop_sign", "reached_goal", "lane_dist", - "lane_angle", + "lane_heading_error_radians", + "lane_heading_cosine", "comfort_violation", "velocity_progress", "speed_limit", @@ -602,7 +603,18 @@ def plot_observation( ) target_position_scale = scales["goal_to_position"] - ego_speed, ego_width, ego_length, steering_angle, accel_long, accel_lat, lcenter, lalign, speed_limit, _ = ego_state + ( + ego_speed, + ego_width, + ego_length, + steering_angle, + accel_long, + accel_lat, + lcenter, + lane_heading_error, + speed_limit, + _, + ) = ego_state ego_width *= scales["veh_width_to_position"] ego_length *= scales["veh_len_to_position"] @@ -646,7 +658,7 @@ def plot_observation( ax.scatter(wp_x, wp_y, color=color, marker=marker, s=s, zorder=15) # Add dynamics info text for DYNAMICS_MODEL_JERK model - ego_info = f"Speed: {ego_speed:.2f}\nLane Centering: {lcenter:.2f}\nLane Align: {lalign:.2f}\nSpeed Limit: {speed_limit:.2f}" + ego_info = f"Speed: {ego_speed:.2f}\nLane Centering: {lcenter:.2f}\nLane Heading Error: {lane_heading_error:.2f}\nSpeed Limit: {speed_limit:.2f}" ego_info += f"\nSteering: {steering_angle:.3f}\naccel_long: {accel_long:.2f}\naccel_lat: {accel_lat:.2f}" diff --git a/tests/drive/test_drive_metrics.c b/tests/drive/test_drive_metrics.c index 2b6ba6f8b6..cd25c2d9a4 100644 --- a/tests/drive/test_drive_metrics.c +++ b/tests/drive/test_drive_metrics.c @@ -47,7 +47,9 @@ static int test_metric_on_road_lane_alignment(void) { compute_metrics(&env, agent_idx, 0); EXPECT_TRUE(agent->current_lane_idx != -1); - EXPECT_TRUE(agent->metrics_array[LANE_ANGLE_IDX] >= -1.0f && agent->metrics_array[LANE_ANGLE_IDX] <= 1.0f); + EXPECT_TRUE( + agent->metrics_array[LANE_HEADING_COSINE_IDX] >= -1.0f + && agent->metrics_array[LANE_HEADING_COSINE_IDX] <= 1.0f); EXPECT_NEAR(agent->metrics_array[OFFROAD_IDX], 0.0f, 1e-5f); free_allocated(&env); diff --git a/tests/drive/test_drive_observations_rewards.c b/tests/drive/test_drive_observations_rewards.c index 8910e3d412..62e2c61487 100644 --- a/tests/drive/test_drive_observations_rewards.c +++ b/tests/drive/test_drive_observations_rewards.c @@ -77,7 +77,7 @@ static void init_reward_env(Drive *env, Agent *agent, Log *log, int *active, flo agent->reward_coefs[REWARD_COEF_COLLISION] = 3.0f; agent->reward_coefs[REWARD_COEF_OFFROAD] = 4.0f; agent->reward_coefs[REWARD_COEF_STOP_LINE] = 5.0f; - agent->metrics_array[LANE_ANGLE_IDX] = 1.0f; + agent->metrics_array[LANE_HEADING_COSINE_IDX] = 1.0f; } static int test_reward_terminal_components(void) { @@ -163,7 +163,8 @@ static int test_reward_lane_align_wrong_way(void) { init_reward_env(&env, &agent, &log, active, reward); reward[0] = 0.0f; - agent.metrics_array[LANE_ANGLE_IDX] = -1.0f; // cos(θ_f) = -1 → driving against lane + agent.metrics_array[LANE_HEADING_COSINE_IDX] = -1.0f; // cos(θ_f) = -1 → driving against lane + agent.metrics_array[LANE_HEADING_ERROR_RADIANS_IDX] = M_PI; agent.sim_speed_signed = 5.0f; agent.reward_coefs[REWARD_COEF_LANE_ALIGN] = 1.0f; agent.reward_coefs[REWARD_COEF_VEL_ALIGN] = 1.0f;