From aeac7e0affc67a442b2a72aba415cd763bf461c6 Mon Sep 17 00:00:00 2001 From: Valentin Charraut Date: Wed, 9 Sep 2026 10:39:57 +0200 Subject: [PATCH 1/5] Refactor observation and inference feature handling; update boundary feature count --- notebooks/01_observations.py | 6 +++--- notebooks/05_inference.py | 12 ++++++------ pufferlib/ocean/drive/constants.h | 2 +- pufferlib/ocean/drive/drive.h | 15 ++++++--------- pufferlib/ocean/drive/render.h | 8 ++++---- 5 files changed, 20 insertions(+), 23 deletions(-) diff --git a/notebooks/01_observations.py b/notebooks/01_observations.py index 567cabf463..92c8c99750 100644 --- a/notebooks/01_observations.py +++ b/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], diff --git a/notebooks/05_inference.py b/notebooks/05_inference.py index 48337a52d8..17122e6712 100644 --- a/notebooks/05_inference.py +++ b/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") diff --git a/pufferlib/ocean/drive/constants.h b/pufferlib/ocean/drive/constants.h index fda70a1d88..e41942904f 100644 --- a/pufferlib/ocean/drive/constants.h +++ b/pufferlib/ocean/drive/constants.h @@ -223,7 +223,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 diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index 80d8b68591..f27e1e6466 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -351,7 +351,7 @@ static const RewardBound REWARD_BOUNDS[NUM_REWARD_COEFS] = { {0.8f, 1.25f, 0}, // REWARD_COEF_THROTTLE C_throttle {0.8f, 1.25f, 0}, // REWARD_COEF_STEER C_steer {0.666f, 1.5f, 0}, // REWARD_COEF_ACC C_acc - {0.666f, 1.5f, 0}, // REWARD_COEF_SPEED C_vel + {0.666f, 1.5f, 0}, // REWARD_COEF_SPEED C_vel }; // Meaning of the values: [min_range, max_range, use_log_scale] @@ -373,7 +373,7 @@ static const RewardBound REWARD_BOUNDS_LOG[NUM_REWARD_COEFS] = { {0.8f, 1.25f, 0}, // REWARD_COEF_THROTTLE C_throttle {0.8f, 1.25f, 0}, // REWARD_COEF_STEER C_steer {0.666f, 1.5f, 0}, // REWARD_COEF_ACC C_acc - {0.666f, 1.5f, 0}, // REWARD_COEF_SPEED C_vel + {0.666f, 1.5f, 0}, // REWARD_COEF_SPEED C_vel }; // ======================================== @@ -4065,11 +4065,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]; @@ -4084,10 +4085,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; } } diff --git a/pufferlib/ocean/drive/render.h b/pufferlib/ocean/drive/render.h index 8f7977b783..c3632796cd 100644 --- a/pufferlib/ocean/drive/render.h +++ b/pufferlib/ocean/drive/render.h @@ -925,8 +925,8 @@ void draw_agent_obs(Drive *env, int agent_index, int mode, int obs_only, int las float x_middle = agent_obs[entity_idx] * env->obs_norm_xy_offset_m; float y_middle = agent_obs[entity_idx + 1] * env->obs_norm_xy_offset_m; float z_middle = agent_obs[entity_idx + 2] * Z_BUFFER; - float rel_angle_x = (agent_obs[entity_idx + 5]); - float rel_angle_y = (agent_obs[entity_idx + 6]); + float rel_angle_x = (agent_obs[entity_idx + 4]); + float rel_angle_y = (agent_obs[entity_idx + 5]); float rel_angle = atan2f(rel_angle_y, rel_angle_x); float segment_size = agent_obs[entity_idx + 3] * env->obs_norm_road_seg_length_m; // Calculate endpoint using the relative angle directly @@ -979,8 +979,8 @@ void draw_agent_obs(Drive *env, int agent_index, int mode, int obs_only, int las float x_middle = agent_obs[entity_idx] * env->obs_norm_xy_offset_m; float y_middle = agent_obs[entity_idx + 1] * env->obs_norm_xy_offset_m; float z_middle = agent_obs[entity_idx + 2] * Z_BUFFER; - float rel_angle_x = agent_obs[entity_idx + 5]; - float rel_angle_y = agent_obs[entity_idx + 6]; + float rel_angle_x = agent_obs[entity_idx + 4]; + float rel_angle_y = agent_obs[entity_idx + 5]; float rel_angle = atan2f(rel_angle_y, rel_angle_x); float segment_size = agent_obs[entity_idx + 3] * env->obs_norm_road_seg_length_m; float x_start = x_middle - segment_size * cosf(rel_angle); From 298359c50672656901ef175af2c72bb2aef9f588 Mon Sep 17 00:00:00 2001 From: Valentin Charraut Date: Wed, 9 Sep 2026 10:48:18 +0200 Subject: [PATCH 2/5] Normalize goal direction vectors in write_reward_target_obs function --- pufferlib/ocean/drive/drive.h | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index f27e1e6466..7c0bf691ae 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -3852,8 +3852,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 = sqrtf(rel_goal_x * rel_goal_x + rel_goal_y * rel_goal_y); + float goal_norm = fmaxf(env->obs_norm_goal_offset_m, goal_distance); + obs[obs_idx++] = rel_goal_x / goal_norm; + obs[obs_idx++] = rel_goal_y / goal_norm; obs[obs_idx++] = (ego->list_goal_z[goal_idx] - ego->sim_z) / env->obs_norm_z_m; } From 2a2f266b0a5eb8655bedebefb26c27c36433b539 Mon Sep 17 00:00:00 2001 From: Valentin Charraut Date: Wed, 9 Sep 2026 11:42:08 +0200 Subject: [PATCH 3/5] Refactor observation handling: update reward coefficient comments and normalize observation indices --- pufferlib/ocean/drive/drive.h | 12 ++++++------ pufferlib/ocean/drive/drive.py | 11 +++++++++++ pufferlib/ocean/drive/render.h | 4 ++-- pufferlib/pufferl.py | 10 +++++++--- pufferlib/viz.py | 6 +++--- 5 files changed, 29 insertions(+), 14 deletions(-) diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index 7c0bf691ae..cf7bc56be9 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -351,7 +351,7 @@ static const RewardBound REWARD_BOUNDS[NUM_REWARD_COEFS] = { {0.8f, 1.25f, 0}, // REWARD_COEF_THROTTLE C_throttle {0.8f, 1.25f, 0}, // REWARD_COEF_STEER C_steer {0.666f, 1.5f, 0}, // REWARD_COEF_ACC C_acc - {0.666f, 1.5f, 0}, // REWARD_COEF_SPEED C_vel + {0.666f, 1.5f, 0}, // REWARD_COEF_SPEED C_vel }; // Meaning of the values: [min_range, max_range, use_log_scale] @@ -373,7 +373,7 @@ static const RewardBound REWARD_BOUNDS_LOG[NUM_REWARD_COEFS] = { {0.8f, 1.25f, 0}, // REWARD_COEF_THROTTLE C_throttle {0.8f, 1.25f, 0}, // REWARD_COEF_STEER C_steer {0.666f, 1.5f, 0}, // REWARD_COEF_ACC C_acc - {0.666f, 1.5f, 0}, // REWARD_COEF_SPEED C_vel + {0.666f, 1.5f, 0}, // REWARD_COEF_SPEED C_vel }; // ======================================== @@ -3853,10 +3853,10 @@ static int write_reward_target_obs(Drive *env, Agent *ego, float *obs, int obs_i &rel_goal_x, &rel_goal_y); // Goals beyond obs_norm_goal_offset_m collapse to a unit direction vector (keeps obs in [-1, 1]) - float goal_distance = sqrtf(rel_goal_x * rel_goal_x + rel_goal_y * rel_goal_y); - float goal_norm = fmaxf(env->obs_norm_goal_offset_m, goal_distance); - obs[obs_idx++] = rel_goal_x / goal_norm; - obs[obs_idx++] = rel_goal_y / goal_norm; + 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; } diff --git a/pufferlib/ocean/drive/drive.py b/pufferlib/ocean/drive/drive.py index fb2af24fbf..b1c9b632b4 100644 --- a/pufferlib/ocean/drive/drive.py +++ b/pufferlib/ocean/drive/drive.py @@ -278,6 +278,17 @@ def __init__( self.single_observation_space = gymnasium.spaces.Box(low=-1, high=1, shape=(self.num_obs,), dtype=np.float32) + self.normalized_obs_mask = np.ones(self.num_obs, dtype=bool) + traffic_control_base = ( + self.num_obs + - self.obs_valid_count_features + - self.obs_slots_traffic_controls_n * self.traffic_control_features + ) + for slot_idx in range(self.obs_slots_traffic_controls_n): + slot_end = traffic_control_base + (slot_idx + 1) * self.traffic_control_features + self.normalized_obs_mask[slot_end - 2 : slot_end] = False + self.normalized_obs_mask[self.num_obs - self.obs_valid_count_features :] = False + self.init_step = init_step # Per C environment randomized start point. When on, each parallel environment # starts the episode at a randomized point. diff --git a/pufferlib/ocean/drive/render.h b/pufferlib/ocean/drive/render.h index c3632796cd..c83397e0ff 100644 --- a/pufferlib/ocean/drive/render.h +++ b/pufferlib/ocean/drive/render.h @@ -925,8 +925,8 @@ void draw_agent_obs(Drive *env, int agent_index, int mode, int obs_only, int las float x_middle = agent_obs[entity_idx] * env->obs_norm_xy_offset_m; float y_middle = agent_obs[entity_idx + 1] * env->obs_norm_xy_offset_m; float z_middle = agent_obs[entity_idx + 2] * Z_BUFFER; - float rel_angle_x = (agent_obs[entity_idx + 4]); - float rel_angle_y = (agent_obs[entity_idx + 5]); + float rel_angle_x = agent_obs[entity_idx + 4]; + float rel_angle_y = agent_obs[entity_idx + 5]; float rel_angle = atan2f(rel_angle_y, rel_angle_x); float segment_size = agent_obs[entity_idx + 3] * env->obs_norm_road_seg_length_m; // Calculate endpoint using the relative angle directly diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index b86b09d03f..6bbd7326ed 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -295,6 +295,9 @@ def __init__(self, config, vecenv, policy, logger=None): # Initializations self.config = config self.vecenv = vecenv + self.normalized_obs_idx = torch.as_tensor( + np.flatnonzero(vecenv.driver_env.normalized_obs_mask), device=config["device"] + ) self.epoch = 0 self.global_step = 0 self.agent_steps = 0 @@ -387,9 +390,10 @@ 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.normalized_obs_idx is None else o_device[..., self.normalized_obs_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: diff --git a/pufferlib/viz.py b/pufferlib/viz.py index 0af8935e25..ff8eb82336 100644 --- a/pufferlib/viz.py +++ b/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( @@ -1415,7 +1415,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 { const out=[]; for(let i=0;i { const out=[]; for(let i=0;i Date: Wed, 9 Sep 2026 12:09:46 +0200 Subject: [PATCH 4/5] Update golden JSON files: add stop sign components and adjust reward metrics --- .../data/drive_rollout_golden.json | 2 + .../smoke_tests/data/drive_smoke_golden.json | 72 ++++++++++--------- 2 files changed, 39 insertions(+), 35 deletions(-) diff --git a/tests/smoke_tests/data/drive_rollout_golden.json b/tests/smoke_tests/data/drive_rollout_golden.json index 15014e813c..32f759d65e 100644 --- a/tests/smoke_tests/data/drive_rollout_golden.json +++ b/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 diff --git a/tests/smoke_tests/data/drive_smoke_golden.json b/tests/smoke_tests/data/drive_smoke_golden.json index 196b514fe9..2089224092 100644 --- a/tests/smoke_tests/data/drive_smoke_golden.json +++ b/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, From c713d5451e55e19adb49724818a7a2cdabfad3d7 Mon Sep 17 00:00:00 2001 From: Valentin Charraut Date: Wed, 16 Sep 2026 12:16:37 +0200 Subject: [PATCH 5/5] Refactor observation masks to improve traffic control feature handling --- pufferlib/ocean/drive/drive.py | 18 ++++++++++-------- pufferlib/pufferl.py | 8 +++++--- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/pufferlib/ocean/drive/drive.py b/pufferlib/ocean/drive/drive.py index 3a92a8cf1f..c34e1d9b15 100644 --- a/pufferlib/ocean/drive/drive.py +++ b/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,16 +273,16 @@ def __init__( self.single_observation_space = gymnasium.spaces.Box(low=-1, high=1, shape=(self.num_obs,), dtype=np.float32) - self.normalized_obs_mask = np.ones(self.num_obs, dtype=bool) - traffic_control_base = ( - self.num_obs - - self.obs_valid_count_features - - self.obs_slots_traffic_controls_n * self.traffic_control_features + # 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 = traffic_control_base + (slot_idx + 1) * self.traffic_control_features - self.normalized_obs_mask[slot_end - 2 : slot_end] = False - self.normalized_obs_mask[self.num_obs - self.obs_valid_count_features :] = False + 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 diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index ad766b8442..a1ed9a6bc6 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -301,8 +301,8 @@ def __init__(self, config, vecenv, policy, logger=None): # Initializations self.config = config self.vecenv = vecenv - self.normalized_obs_idx = torch.as_tensor( - np.flatnonzero(vecenv.driver_env.normalized_obs_mask), device=config["device"] + 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 @@ -396,7 +396,9 @@ 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. - obs_stat_source = o_device if self.normalized_obs_idx is None else o_device[..., self.normalized_obs_idx] + 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())