Add lane heading error

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

notebooks/notebook_utils.py

@@ -50,7 +50,7 @@
     "accel_long",
     "accel_lat",
     "lane_center",
-    "lane_align",
+    "lane_heading_error",
     "speed_limit",
     "stopped",
 ]

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

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;

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;

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}"
 

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);

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;