Change observation layout: goal normalization and road inputs

notebooks/01_observations.py

@@ -216,7 +216,7 @@
 # ## Lane / boundary segments
 
 # %%
-road_labels = ["rel_x", "rel_y", "rel_z", "length", "width", "dir_cos", "dir_sin", "goal_dist_abs", "goal_dist_rel"]
+road_labels = ["rel_x", "rel_y", "rel_z", "length", "dir_cos", "dir_sin", "width", "goal_dist_abs", "goal_dist_rel"]
 
 lane_active = ~np.all(lanes == 0, axis=1)
 bound_active = ~np.all(boundaries == 0, axis=1)
@@ -228,7 +228,7 @@
 
 # Mirror the canonical road rendering in pufferlib.viz.plot_observation
 for seg in lanes[lane_active]:
-    x, y, z, length, width, dc, ds = seg[:7]
+    x, y, z, length, dc, ds = seg[:6]
     # seg[7] = goal_dist_abs (0 near goal lane -> 1 far); green->red colormap
     color = plt.cm.RdYlGn_r(float(seg[7])) if env.obs_goal_lane_distance else "lightgrey"
     ax.scatter(x, y, color=color, s=10, zorder=1)
@@ -241,7 +241,7 @@
     )
 
 for seg in boundaries[bound_active]:
-    x, y, z, length, width, dc, ds = seg[:7]
+    x, y, z, length, dc, ds = seg[:6]
     ax.scatter(x, y, color="black", s=10, zorder=1)
     ax.plot(
         [x + dc * length / 2, x - dc * length / 2],

notebooks/05_inference.py

@@ -272,8 +272,8 @@ def run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE, horizon=HORI
 # - **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
-# - **Lanes** (MAX_LANES x 7): rel_x, rel_y, rel_z, seg_length, seg_width, dir_cos, dir_sin
-# - **Boundaries** (MAX_BOUNDS x 7): same as lanes
+# - **Lanes** (MAX_LANES x 9): rel_x, rel_y, rel_z, seg_length, dir_cos, dir_sin, seg_width, goal_dist_abs, goal_dist_rel
+# - **Boundaries** (MAX_BOUNDS x 6): first 6 lane features
 # - **Traffic controls** (MAX_TRAFFIC x 7): rel_x1, rel_y1, rel_x2, rel_y2, rel_z, type, state
 
 # %%
@@ -598,7 +598,7 @@ def unpack_all_timesteps(bufs, agent_idx):
 for i in range(lanes.shape[0]):
     if np.allclose(lanes[i], 0):
         continue
-    rx, ry, rz, length, _, dc, ds = lanes[i][:7]
+    rx, ry, rz, length, dc, ds = lanes[i][:6]
     ax.plot(
         [rx - dc * length / 2, rx + dc * length / 2],
         [ry - ds * length / 2, ry + ds * length / 2],
@@ -620,7 +620,7 @@ def unpack_all_timesteps(bufs, agent_idx):
 for i in range(boundaries.shape[0]):
     if np.allclose(boundaries[i], 0):
         continue
-    rx, ry, rz, length, _, dc, ds = boundaries[i][:7]
+    rx, ry, rz, length, dc, ds = boundaries[i][:6]
     ax.plot(
         [rx - dc * length / 2, rx + dc * length / 2],
         [ry - ds * length / 2, ry + ds * length / 2],
@@ -865,7 +865,7 @@ def unpack_all_timesteps(bufs, agent_idx):
 
 # %%
 # Road per-feature distributions (lanes + boundaries)
-road_labels = ["rel_x", "rel_y", "rel_z", "seg_length", "seg_width", "dir_cos", "dir_sin"]
+road_labels = ["rel_x", "rel_y", "rel_z", "seg_length", "dir_cos", "dir_sin"]
 lf = env.lane_features
 bf = env.boundary_features
 max_lanes = env.obs_slots_lane_kept
@@ -891,7 +891,7 @@ def unpack_all_timesteps(bufs, agent_idx):
     f"({100 * len(vis_bounds) / (all_bounds.shape[0] * max_bounds):.1f}%)"
 )
 
-fig, axes = plt.subplots(2, 7, figsize=(28, 8))
+fig, axes = plt.subplots(2, 6, figsize=(24, 8))
 for i, label in enumerate(road_labels):
     # Lanes
     axes[0, i].hist(vis_lanes[:, i], bins=80, edgecolor="black", alpha=0.7, color="silver")

pufferlib/ocean/drive/constants.h

@@ -219,7 +219,7 @@ static const int ROAD_OFFSETS[25][2]
 
 #define EGO_FEATURES 10
 #define LANE_FEATURES 9
-#define BOUNDARY_FEATURES 9
+#define BOUNDARY_FEATURES 6
 #define PARTNER_FEATURES 9
 #define TRAFFIC_CONTROL_FEATURES 7
 #define GOAL_FEATURES 3

pufferlib/ocean/drive/drive.h

@@ -3807,8 +3807,11 @@ static int write_reward_target_obs(Drive *env, Agent *ego, float *obs, int obs_i
             ego->list_goal_y[goal_idx],
             &rel_goal_x,
             &rel_goal_y);
-        obs[obs_idx++] = rel_goal_x / env->obs_norm_goal_offset_m;
-        obs[obs_idx++] = rel_goal_y / env->obs_norm_goal_offset_m;
+        // Goals beyond obs_norm_goal_offset_m collapse to a unit direction vector (keeps obs in [-1, 1])
+        float goal_distance_m = sqrtf(rel_goal_x * rel_goal_x + rel_goal_y * rel_goal_y);
+        float goal_normalization_m = fmaxf(env->obs_norm_goal_offset_m, goal_distance_m);
+        obs[obs_idx++] = rel_goal_x / goal_normalization_m;
+        obs[obs_idx++] = rel_goal_y / goal_normalization_m;
         obs[obs_idx++] = (ego->list_goal_z[goal_idx] - ego->sim_z) / env->obs_norm_z_m;
     }
 
@@ -4020,11 +4023,12 @@ static int write_road_obs(Drive *env, Agent *ego, float *obs, int obs_idx, int *
         segment_dest[feature_base + 1] = rel_y / env->obs_norm_xy_offset_m;
         segment_dest[feature_base + 2] = rel_z / env->obs_norm_z_m;
         segment_dest[feature_base + 3] = seg_half_len / env->obs_norm_road_seg_length_m;
-        segment_dest[feature_base + 4] = LANE_WIDTH / env->obs_norm_road_seg_width_m;
-        segment_dest[feature_base + 5] = rel_seg_dir_x;
-        segment_dest[feature_base + 6] = rel_seg_dir_y;
+        segment_dest[feature_base + 4] = rel_seg_dir_x;
+        segment_dest[feature_base + 5] = rel_seg_dir_y;
         // Goal-distance features: absolute and relative to ego's lane->goal distance.
         if (is_lane) {
+            // Constant until the map format carries per-lane width
+            segment_dest[feature_base + 6] = LANE_WIDTH / env->obs_norm_road_seg_width_m;
             float goal_dist_abs = 0.0f, goal_dist_rel = 0.0f; // 0 when flag off / unresolved
             if (env->obs_goal_lane_distance && goal_graph_idx >= 0 && entity_idx < env->num_road_elements) {
                 int lane_graph_idx = env->lane_graph.lane_to_graph_idx[entity_idx];
@@ -4039,10 +4043,6 @@ static int write_road_obs(Drive *env, Agent *ego, float *obs, int obs_idx, int *
             }
             segment_dest[feature_base + 7] = goal_dist_abs;
             segment_dest[feature_base + 8] = goal_dist_rel;
-        } else {
-            // NOTE: Remove this with next model
-            segment_dest[feature_base + 7] = 0.0f;
-            segment_dest[feature_base + 8] = 0.0f;
         }
     }
 

pufferlib/ocean/drive/drive.py

@@ -6,6 +6,8 @@
 import pufferlib
 from pufferlib.ocean.drive import binding
 
+TRAFFIC_CONTROL_CATEGORICAL_FEATURE_COUNT = 2  # type and state
+
 
 def compute_effective_road_obs_count(max_count, dropout):
     if max_count <= 0:
@@ -271,6 +273,17 @@ def __init__(
 
         self.single_observation_space = gymnasium.spaces.Box(low=-1, high=1, shape=(self.num_obs,), dtype=np.float32)
 
+        # Observation distribution stats exclude raw traffic-control categories and valid-slot counts.
+        self.obs_stats_feature_mask = np.ones(self.num_obs, dtype=bool)
+        valid_counts_start_idx = self.num_obs - self.obs_valid_count_features
+        traffic_controls_start_idx = (
+            valid_counts_start_idx - self.obs_slots_traffic_controls_n * self.traffic_control_features
+        )
+        for slot_idx in range(self.obs_slots_traffic_controls_n):
+            slot_end_idx = traffic_controls_start_idx + (slot_idx + 1) * self.traffic_control_features
+            self.obs_stats_feature_mask[slot_end_idx - TRAFFIC_CONTROL_CATEGORICAL_FEATURE_COUNT : slot_end_idx] = False
+        self.obs_stats_feature_mask[valid_counts_start_idx:] = False
+
         self.init_step = init_step
         # Per C environment randomized start point. When on, each parallel environment
         # starts the episode at a randomized point.

pufferlib/pufferl.py

@@ -301,6 +301,9 @@ def __init__(self, config, vecenv, policy, logger=None):
         # Initializations
         self.config = config
         self.vecenv = vecenv
+        self.obs_stats_feature_idx = torch.as_tensor(
+            np.flatnonzero(vecenv.driver_env.obs_stats_feature_mask), device=config["device"]
+        )
         self.epoch = 0
         self.global_step = 0
         self.agent_steps = 0
@@ -393,9 +396,12 @@ def evaluate(self):
             # Obs distribution stats (max/min/mean across the batch and obs
             # dims, appended per env step). Surfaces clipping / unbounded
             # features / normalization regressions in wandb.
-            self.stats["obs/max"].append(o_device.max().item())
-            self.stats["obs/min"].append(o_device.min().item())
-            self.stats["obs/mean"].append(o_device.mean().item())
+            obs_stat_source = (
+                o_device if self.obs_stats_feature_idx is None else o_device[..., self.obs_stats_feature_idx]
+            )
+            self.stats["obs/max"].append(obs_stat_source.max().item())
+            self.stats["obs/min"].append(obs_stat_source.min().item())
+            self.stats["obs/mean"].append(obs_stat_source.mean().item())
 
             profile("eval_forward", epoch)
             with torch.no_grad(), self.amp_context:

pufferlib/viz.py

@@ -692,7 +692,7 @@ def plot_observation(
         count_lane += 1
         rel_x, rel_y = lane_obs[i][0], lane_obs[i][1]
         length = lane_obs[i][3] * rl2p
-        dir_cos, dir_sin = lane_obs[i][5], lane_obs[i][6]
+        dir_cos, dir_sin = lane_obs[i][4], lane_obs[i][5]
         # idx 7 = goal_dist_abs (0 near goal lane -> 1 far/unreachable); green->red colormap
         color = plt.cm.RdYlGn_r(float(lane_obs[i][7])) if obs_goal_lane_distance else "lightgrey"
         ax.scatter(rel_x, rel_y, color=color, s=10, zorder=1)
@@ -711,7 +711,7 @@ def plot_observation(
         count_boundary += 1
         rel_x, rel_y = boundary_obs[i][0], boundary_obs[i][1]
         length = boundary_obs[i][3] * rl2p
-        dir_cos, dir_sin = boundary_obs[i][5], boundary_obs[i][6]
+        dir_cos, dir_sin = boundary_obs[i][4], boundary_obs[i][5]
         color = "black"
         ax.scatter(rel_x, rel_y, color=color, s=10, zorder=1)
         ax.plot(
@@ -1453,7 +1453,7 @@ def _render_interactive_replay_payload(compressed_payload, filename):
             const trafficStart = p;
             const rot = (x,y) => [-y,x];
             const zero = (off,n) => { for(let i=0;i<n;i++) if(obs[off+i] !== 0) return false; return true; };
-            const roads = (start,count,poolName,feat) => { const out=[]; for(let i=0;i<count;i++){ const o=start+i*feat; if(zero(o,feat)) continue; let xy=rot(obs[o]*Q,obs[o+1]*Q), cs=rot(obs[o+5]*Q,obs[o+6]*Q); out.push([xy[0],xy[1],obs[o+3]*Q*H.scales.road_length_to_position,cs[0],cs[1],poolAt(poolName,frame,slot,i)]); } return out; };
+            const roads = (start,count,poolName,feat) => { const out=[]; for(let i=0;i<count;i++){ const o=start+i*feat; if(zero(o,feat)) continue; let xy=rot(obs[o]*Q,obs[o+1]*Q), cs=rot(obs[o+4]*Q,obs[o+5]*Q); out.push([xy[0],xy[1],obs[o+3]*Q*H.scales.road_length_to_position,cs[0],cs[1],poolAt(poolName,frame,slot,i)]); } return out; };
             const partners = []; for(let i=0;i<H.obs_slots_partners_n;i++){ const o=partnersStart+i*H.partner_features; if(zero(o,H.partner_features)) continue; let xy=rot(obs[o]*Q,obs[o+1]*Q), h=Math.atan2(obs[o+6],obs[o+5]); h = ((h + Math.PI/2 + Math.PI) % (2*Math.PI)) - Math.PI; partners.push({x:xy[0],y:xy[1],l:obs[o+3]*Q*H.scales.veh_len_to_position,w:obs[o+4]*Q*H.scales.veh_width_to_position,h:h,pool:poolAt("pool_partner",frame,slot,i)}); }
             const gps = []; for(let i=0;i<H.num_goals;i++){ const o=targetStart+i*H.goal_features; if(zero(o,H.goal_features)) continue; let scale=H.scales.goal_to_position*Q, xy=rot(obs[o]*scale, obs[o+1]*scale); gps.push(xy); }
             const controls = []; for(let i=0;i<H.traffic_obs_count;i++){ const o=trafficStart+i*TF; if(zero(o,TF)) continue; let a=rot(obs[o]*Q,obs[o+1]*Q), b=rot(obs[o+2]*Q,obs[o+3]*Q); controls.push({type:Math.round(obs[o+5]*Q), state:Math.round(obs[o+6]*Q), x1:a[0], y1:a[1], x2:b[0], y2:b[1], pool:poolAt("pool_traffic",frame,slot,i)}); }

tests/smoke_tests/data/drive_rollout_golden.json

@@ -22,9 +22,11 @@
     "reward_components/overspeed": 0.0,
     "reward_components/red_light": -0.0007146410954495271,
     "reward_components/reverse": -0.010871957583973805,
+    "reward_components/stop_sign": 0.0,
     "reward_components/timestep": -8.810036033537471e-05,
     "reward_components/velocity": 0.00020452999203068127,
     "score": 0.0,
+    "stop_sign_violation_rate": 0.0,
     "total_distance_travelled_sum": 131.91801204681397,
     "total_infraction_count": 10.55,
     "velocity_progress_sum": 0.019642803197105724

tests/smoke_tests/data/drive_smoke_golden.json

@@ -1,50 +1,52 @@
 {
   "env": {
-    "avg_distance_per_infraction": 13.043761351398219,
-    "avg_speed_per_agent": 1.367759719491005,
-    "collision_rate": 0.024596774019300938,
-    "comfort_violation_count": 0.6676511764526367,
-    "dnf_rate": 0.5364415310323238,
-    "episode_length": 14.842876672744751,
-    "episode_return": -1.2550823464989662,
-    "lane_center_rate": 0.6568814143538475,
+    "avg_distance_per_infraction": 11.98020388045401,
+    "avg_speed_per_agent": 1.3697761297225952,
+    "collision_rate": 0.04865591321140528,
+    "comfort_violation_count": 0.6441232562065125,
+    "dnf_rate": 0.5329301133751869,
+    "episode_length": 13.433198928833008,
+    "episode_return": -1.2031443566083908,
+    "lane_center_rate": 0.6850986406207085,
     "n": 28.75,
     "num_goals_reached": 0.008333333767950535,
-    "obs/max": 49.0,
-    "obs/mean": 0.24173322669230402,
-    "obs/min": -1.072199359536171,
-    "offroad_rate": 0.4389616884291172,
+    "obs/max": 1.0000000186264515,
+    "obs/mean": 0.11787314258981496,
+    "obs/min": -1.000000019557774,
+    "offroad_rate": 0.4142473079264164,
     "red_light_violation_rate": 0.0,
     "reward_components/ade": 0.0,
-    "reward_components/collision": -0.05138678662478924,
-    "reward_components/comfort": -0.5063945986330509,
+    "reward_components/collision": -0.07210881542414427,
+    "reward_components/comfort": -0.43151041865348816,
     "reward_components/goal": 0.004166666883975267,
-    "reward_components/lane_align": -0.029466886539012194,
-    "reward_components/lane_center": -0.005276555719319731,
-    "reward_components/offroad": -0.6542690135538578,
+    "reward_components/lane_align": -0.026959585840813816,
+    "reward_components/lane_center": -0.004400172387249768,
+    "reward_components/offroad": -0.6615230031311512,
     "reward_components/overspeed": 0.0,
     "reward_components/red_light": 0.0,
-    "reward_components/reverse": -0.012481134268455207,
-    "reward_components/timestep": -9.589604633220006e-05,
-    "reward_components/velocity": 0.00012186067851871485,
+    "reward_components/reverse": -0.01090477011166513,
+    "reward_components/stop_sign": 0.0,
+    "reward_components/timestep": -8.623538178653689e-05,
+    "reward_components/velocity": 0.00018188954209108488,
     "score": 0.0,
-    "total_distance_travelled": 1395.6824645996094,
-    "total_infractions": 107.0,
-    "velocity_progress_sum": 0.011142019531689584
+    "stop_sign_violation_rate": 0.0,
+    "total_distance_travelled": 1269.901611328125,
+    "total_infractions": 106.0,
+    "velocity_progress_sum": 0.018718332110438496
   },
   "losses": {
-    "approx_kl": 0.002203181851655245,
-    "clipfrac": 0.016433581709861755,
-    "ema_max": 2.8828256130218506,
-    "entropy": 2.474294900894165,
-    "explained_variance": 0.004730343818664551,
-    "filter_threshold": 0.028828256130218506,
-    "filtered_fraction": 0.07795229784758584,
-    "kept_fraction": 0.9220477021524142,
-    "masked_fraction": 0.16064453125,
-    "old_approx_kl": 0.0029104824643582106,
-    "policy_loss": -0.0030369183514267206,
-    "value_loss": 0.18146912753582
+    "approx_kl": 0.005168732721358538,
+    "clipfrac": 0.076636902987957,
+    "ema_max": 2.959818959236145,
+    "entropy": 2.4640564918518066,
+    "explained_variance": 0.055533766746520996,
+    "filter_threshold": 0.029598189592361452,
+    "filtered_fraction": 0.09143855562026793,
+    "kept_fraction": 0.9085614443797321,
+    "masked_fraction": 0.16162109375,
+    "old_approx_kl": 0.0028026883956044912,
+    "policy_loss": 0.006841260474175215,
+    "value_loss": 0.3042336106300354
   },
   "meta": {
     "bptt_horizon": 64,