diff --git a/README.md b/README.md index 41fcf3f614..926df70402 100644 --- a/README.md +++ b/README.md @@ -266,7 +266,7 @@ To add a metric to the report: add its key to `TREND_METRICS` / | `reward_vel_align` | `1.0` | Speed matching road limit | | `reward_lane_center` | `0.0038` | Lane centering | | `reward_stop_line` | `1.0` | Stop line penalty | -| `reward_overspeed` | `0.05` | Speeding penalty | +| `reward_overspeed` | `0.005` | Speeding penalty | ### `train` — PPO diff --git a/data_utils/datasets.yaml b/data_utils/datasets.yaml index 724de0ea1f..c839a13496 100644 --- a/data_utils/datasets.yaml +++ b/data_utils/datasets.yaml @@ -62,3 +62,15 @@ nuplan_val: CC BY-NC-SA 4.0 (nuPlan, Motional) — non-commercial use only, attribution required, share-alike. Redistribute only under this same license and keep this notice with any copy. + +carla_128_affine_zones_stop: + s3_uri: s3://pufferdrive-bins/carla/0.3.2/128_affine_zones_stop/ + size: ~106 MB, 128 bins + public: true + description: >- + CARLA Towns 01-07/10HD as PufferDrive bins with 16 affine augmentations + each (flips, shears, scales), per-lane speed zones, stop signs and the + 0.1 polyline tolerance (123Drive mapforge at v0.3.2-19). The training set + behind the k_scaled_004x runs. + license: >- + CARLA assets are CC BY 4.0 — attribution required. diff --git a/notebooks/01_observations.py b/notebooks/01_observations.py index 92c8c99750..6420781e82 100644 --- a/notebooks/01_observations.py +++ b/notebooks/01_observations.py @@ -69,6 +69,8 @@ obs_slots_traffic_controls_n=env.obs_slots_traffic_controls_n, obs_dropout_lane=env.obs_dropout_lane, obs_dropout_boundary=env.obs_dropout_boundary, + obs_partner_relative_velocity=env.obs_partner_relative_velocity, + obs_lane_speed_limit=env.obs_lane_speed_limit, ) print(f"ego: {ego.shape} = {ego}") print(f"target: {target.shape}") @@ -283,6 +285,8 @@ obs_slots_traffic_controls_n=env.obs_slots_traffic_controls_n, obs_dropout_lane=env.obs_dropout_lane, obs_dropout_boundary=env.obs_dropout_boundary, + obs_partner_relative_velocity=env.obs_partner_relative_velocity, + obs_lane_speed_limit=env.obs_lane_speed_limit, obs_lane_stride=env.obs_lane_stride, obs_boundary_stride=env.obs_boundary_stride, obs_goal_lane_distance=env.obs_goal_lane_distance, diff --git a/notebooks/05_inference.py b/notebooks/05_inference.py index 17122e6712..25999cddc8 100644 --- a/notebooks/05_inference.py +++ b/notebooks/05_inference.py @@ -205,6 +205,8 @@ def run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE, horizon=HORI obs_slots_boundary_n=env.obs_slots_boundary_n, obs_dropout_lane=env.obs_dropout_lane, obs_dropout_boundary=env.obs_dropout_boundary, + obs_partner_relative_velocity=env.obs_partner_relative_velocity, + obs_lane_speed_limit=env.obs_lane_speed_limit, obs_lane_stride=env.obs_lane_stride, obs_boundary_stride=env.obs_boundary_stride, obs_slots_traffic_controls_n=env.obs_slots_traffic_controls_n, @@ -241,6 +243,8 @@ def run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE, horizon=HORI obs_slots_boundary_n=env.obs_slots_boundary_n, obs_dropout_lane=env.obs_dropout_lane, obs_dropout_boundary=env.obs_dropout_boundary, + obs_partner_relative_velocity=env.obs_partner_relative_velocity, + obs_lane_speed_limit=env.obs_lane_speed_limit, obs_slots_traffic_controls_n=env.obs_slots_traffic_controls_n, ) ego_features_over_time.append(ego) @@ -290,6 +294,8 @@ def run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE, horizon=HORI obs_slots_boundary_n=env.obs_slots_boundary_n, obs_dropout_lane=env.obs_dropout_lane, obs_dropout_boundary=env.obs_dropout_boundary, + obs_partner_relative_velocity=env.obs_partner_relative_velocity, + obs_lane_speed_limit=env.obs_lane_speed_limit, obs_slots_traffic_controls_n=env.obs_slots_traffic_controls_n, ) @@ -364,6 +370,8 @@ def layer_stats(name, arr): "sim_speed_signed", "seconds_stopped", ] +if env.obs_partner_relative_velocity: + partner_labels += ["rel_vx", "rel_vy"] for p in range(min(int(n_visible), 5)): vals = ", ".join(f"{partner_labels[j]}={partners[p, j]:.3f}" for j in range(env.partner_features)) print(f" [{p}] {vals}") @@ -477,6 +485,8 @@ def unpack_all_timesteps(bufs, agent_idx): obs_slots_boundary_n=env.obs_slots_boundary_n, obs_dropout_lane=env.obs_dropout_lane, obs_dropout_boundary=env.obs_dropout_boundary, + obs_partner_relative_velocity=env.obs_partner_relative_velocity, + obs_lane_speed_limit=env.obs_lane_speed_limit, obs_slots_traffic_controls_n=env.obs_slots_traffic_controls_n, ) egos.append(ego) @@ -554,6 +564,8 @@ def unpack_all_timesteps(bufs, agent_idx): obs_slots_boundary_n=env.obs_slots_boundary_n, obs_dropout_lane=env.obs_dropout_lane, obs_dropout_boundary=env.obs_dropout_boundary, + obs_partner_relative_velocity=env.obs_partner_relative_velocity, + obs_lane_speed_limit=env.obs_lane_speed_limit, obs_slots_traffic_controls_n=env.obs_slots_traffic_controls_n, ) dists = np.sqrt(part[:, 0] ** 2 + part[:, 1] ** 2) @@ -581,6 +593,8 @@ def unpack_all_timesteps(bufs, agent_idx): obs_slots_boundary_n=env.obs_slots_boundary_n, obs_dropout_lane=env.obs_dropout_lane, obs_dropout_boundary=env.obs_dropout_boundary, + obs_partner_relative_velocity=env.obs_partner_relative_velocity, + obs_lane_speed_limit=env.obs_lane_speed_limit, obs_slots_traffic_controls_n=env.obs_slots_traffic_controls_n, ) @@ -644,7 +658,7 @@ def unpack_all_timesteps(bufs, agent_idx): for i in range(partners.shape[0]): if np.allclose(partners[i], 0): continue - rx, ry, rz, length, width, hc, hs, speed, _ = partners[i] + rx, ry, rz, length, width, hc, hs, speed = partners[i][:8] heading = np.arctan2(hs, hc) rect = Rectangle( (-length / 2, -width / 2), length, width, facecolor="orange", edgecolor="black", alpha=0.6, zorder=9 @@ -791,6 +805,8 @@ def unpack_all_timesteps(bufs, agent_idx): "sim_speed_signed", "seconds_stopped", ] +if env.obs_partner_relative_velocity: + partner_labels += ["rel_vx", "rel_vy"] obs_slots_partners_n = env.obs_slots_partners_n pf = env.partner_features diff --git a/pufferlib/config/evaluation/benchmark.yaml b/pufferlib/config/evaluation/benchmark.yaml index 6925cd7971..fc28008bd1 100644 --- a/pufferlib/config/evaluation/benchmark.yaml +++ b/pufferlib/config/evaluation/benchmark.yaml @@ -12,7 +12,7 @@ env: reward_stop_line: 1.0 reward_ade: 0.0 reward_goal: 1.0 - reward_overspeed: 0.05 + reward_overspeed: 0.1 # training-range max (REWARD_BOUNDS) reward_comfort: 0.05 reward_velocity: 0.0025 reward_lane_align: 0.025 @@ -28,7 +28,9 @@ env: obs_dropout_lane: 0.0 obs_dropout_boundary: 0.0 partner_blindness_prob: 0.0 + partner_blindness_prob_min: 0.0 partner_blindness_trigger_prob: 0.0 + partner_blindness_time_fraction_max: 0.0 phantom_braking_prob: 0.0 phantom_braking_trigger_prob: 0.0 use_neighbor_cache: true diff --git a/pufferlib/config/puffer_drive.yaml b/pufferlib/config/puffer_drive.yaml index 8c8cc5d4bd..a435b22647 100644 --- a/pufferlib/config/puffer_drive.yaml +++ b/pufferlib/config/puffer_drive.yaml @@ -57,19 +57,19 @@ env: # Speed cap when reward_randomization is off (eval); null = base_max_speed_mps. Sets C_vel = max_speed_mps / base_max_speed_mps (training range [1/a, a] below) max_speed_mps: null # Training C_vel ~ X(a) = 0.5 U(1/a, 1) + 0.5 U(1, a): speed cap = base_max_speed_mps * [1/a, a]; 1.5 = paper (13.3-30 m/s), 2.0 = 10-40 m/s - conditioning_speed_scale: 1.5 + conditioning_speed_scale: 2.0 # Optional nonzero launch speed for gigaflow random spawns spawn_initial_speed: 0.0 # Gigaflow spawn pose noise: lateral offset as a fraction of half the lane width, heading offset in degrees spawn_lateral_offset_max_frac: 1.0 spawn_heading_max_deg: 30.0 # Gigaflow dynamics randomization: per-step brownian pose buffeting (std in meters / degrees); 0 disables - pose_noise_xy_m: 0.0 - pose_noise_yaw_deg: 0.0 + pose_noise_xy_m: 0.025 + pose_noise_yaw_deg: 0.25 # Speed-limit randomization (needs SPDZONE1 maps): per episode, with this probability every speed zone shifts its # posted limit by one uniform offset in +-delta, clipped to [min, max]; junction lanes inherit the min over entries. # m/s: 9.72 = 35 km/h, 1.39 = 5 km/h, 36.11 = 130 km/h. 0 disables. - speed_limit_random_prob: 0.0 + speed_limit_random_prob: 1.0 speed_limit_random_delta_mps: 9.72 speed_limit_random_min_mps: 1.39 speed_limit_random_max_mps: 36.11 @@ -83,7 +83,7 @@ env: stop_sign_behavior: stop # Traffic controls observed and enforced; stop-sign runs are detected as in CaRL's RunStopSign2 (reward_stop_line) traffic_lights_enabled: true - stop_signs_enabled: false + stop_signs_enabled: true yield_signs_enabled: false # Skip red light violation detection (metrics, reward, behavior) disable_red_light_infractions: false @@ -94,13 +94,13 @@ env: # Share static map geometry (roads/grid/lane-graph) across envs using the same map use_map_cache: true # Load all configured maps before fork so workers inherit shared geometry - preload_map_cache: false + preload_map_cache: true # Precompute neighboring road entities for observation queries use_neighbor_cache: true # Number of steps before reset; keep it off multiples of train.bptt_horizon so episode ends don't all land on the unused last buffer slot - scenario_length: 2500 + scenario_length: 1200 # Frequency of resampling scenario (in steps), 0 to disable - resample_frequency: 256000 + resample_frequency: 0 # End the episode early when too many agents are inactive; otherwise use scenario_length termination_mode: true # Ratio of inactive agents (stopped/removed) that triggers early reset @@ -152,9 +152,9 @@ env: # Maximum speed at final waypoint to count goal reward goal_speed: 3.0 # true: goal-speed coef ~ U(0, 20) per agent under reward_randomization; false: pinned to goal_speed (GIGAFLOW v_goal = 3) - goal_speed_randomization: true + goal_speed_randomization: false # true: the final goal is only consumed below goal speed (GIGAFLOW R_goal); false: consumed on entry, reward gated - goal_reach_requires_speed: false + goal_reach_requires_speed: true # Number of target waypoints in the observation (including goal) num_goals: 3 # Waypoint spacing range in meters when building path @@ -179,7 +179,7 @@ env: reward_velocity: 0.0025 reward_reverse: 0.005 reward_timestep: 2.5e-05 - reward_overspeed: 0.05 + reward_overspeed: 0.005 # m/s above the lane limit before the overspeed penalty fires overspeed_tolerance_mps: 0.0 reward_ade: 0.0 @@ -195,11 +195,11 @@ env: obs_slots_boundary_n: 50 obs_slots_partners_n: 20 # Append the partner's velocity relative to ego (ego frame, 2 floats) to every partner row (GIGAFLOW observes velocity) - obs_partner_relative_velocity: false + obs_partner_relative_velocity: true # Ego lane-angle slot: true = signed heading error theta_f / pi (GIGAFLOW B.1), false = cos(theta_f) - obs_lane_heading_signed: false + obs_lane_heading_signed: true # Append each lane segment's speed limit (/ obs_norm_speed_mps) to its lane row so limit changes ahead are visible - obs_lane_speed_limit: false + obs_lane_speed_limit: true obs_slots_traffic_controls_n: 4 # Fraction of segment observation slots to drop (reduces obs size) obs_dropout_lane: 0.5 @@ -207,6 +207,8 @@ env: # Stride for lane and boundary segments; 1 means use every segment, 2 means use every other segment, etc. obs_lane_stride: 2 obs_boundary_stride: 1 + # true: strided-out segments merge into one chord so polylines stay connected; false = pre-2026-10-12 gaps + obs_stride_merge_segments: true # --- Observation normalization --- obs_norm_speed_mps: 60.0 # num_goals * max_goal_spacing; goals beyond this collapse to a unit direction vector @@ -219,6 +221,8 @@ env: obs_norm_z_m: 10.0 # --- Eval-only --- eval_perceived_size_margin_m: 0.1 + # Eval-only: partners longer than this (m) are observed as ceil(length / this) touching pieces (training max 7.0). 0 disables. + obs_partner_split_length_m: 0.0 # Eval-only: while standing exactly still, |longitudinal jerk| below this (m/s^3) is zeroed # so leftover mean-action jerk cannot creep the car. 0 disables. eval_standstill_jerk_deadband_mps3: 0.0 @@ -230,12 +234,18 @@ env: obs_range_traffic_control_m: 200.0 # --- Robustness features--- # Erratic agents are selected once per episode; masked out of the PPO rollout buffer. - # Upper bound of the per-episode blind-agent rate; each episode draws its rate uniformly from [0, this] + # Upper bound of the per-episode blind-agent rate; each episode draws its rate uniformly from [partner_blindness_prob_min, this] partner_blindness_prob: 0.05 - # Per-step probability that a flagged blind-partner agent zeros out its partner observations - partner_blindness_trigger_prob: 0.03 + # Lower bound of the per-episode blind-agent rate + partner_blindness_prob_min: 0.05 + # Per-step probability that a flagged blind-partner agent restricts its partner observations + partner_blindness_trigger_prob: 1.0 # Duration in seconds of partner blindness once triggered partner_blindness_duration_seconds: 3.0 + # Blind agents keep the partners in their forward corridor (false: they see no partners at all, pre-2026-10-12 rule) + partner_blindness_corridor: false + # 0 disables; else each blind agent draws its share of time spent blind uniformly from [share set by trigger prob and duration, this]; 1.0 = blind every step + partner_blindness_time_fraction_max: 0.0 # Upper bound of the per-episode phantom-braker rate; each episode draws its rate uniformly from [0, this] phantom_braking_prob: 0.1 # Per-step probability that a flagged phantom-braker actually initiates a brake window @@ -255,7 +265,7 @@ policy: # Encoder activation - options: "relu", "tanh", "gelu" encoder_activation: relu encoder_layer_norm: true - mask_padded_features: false + mask_padded_features: true # Shared backbone layer backbone_hidden_size: 1024 backbone_num_layers: 3 @@ -271,7 +281,9 @@ policy: critic_num_layers: 1 critic_head_layer_norm: false # Run the actor/critic heads in float32 under bf16 autocast (backbone stays bf16) - fp32_heads: false + fp32_heads: true + # Fused Triton kernels for the per-slot encoders and their max-pool (same bf16 math, fewer activation bytes) + fused_slot_encoder: true # Dual or shared actor-critic backbone shared_network: false action_type: discrete @@ -297,7 +309,7 @@ train: anneal_lr: true precision: bfloat16 # false: pure float32 matmuls/convs (no TensorFloat-32); requires precision float32 - tf32: true + tf32: false rollout_dtype: float32 total_timesteps: 500000000000 learning_rate: 0.0005 @@ -333,7 +345,7 @@ train: # Accumulate gradients above this size max_minibatch_size: 65536 bptt_horizon: 128 - compile: false + compile: true compile_mode: default compile_fullgraph: false vtrace_rho_clip: 1 @@ -349,6 +361,10 @@ train: # Rendering options render: false render_interval: 1000 + # Training maps per render pass, a seeded random subset; null renders every map + render_num_maps: null + # Stop after this many epochs in this process (debug/render runs); null runs to total_timesteps + stop_after_epochs: null sweep: method: Protein metric: score diff --git a/pufferlib/config_schema.py b/pufferlib/config_schema.py index 72a2c13e11..a9802c2ecf 100644 --- a/pufferlib/config_schema.py +++ b/pufferlib/config_schema.py @@ -336,6 +336,7 @@ class DriveEnvConfig: obs_dropout_boundary: float = _constrained_field(PROBABILITY_CONSTRAINT) obs_lane_stride: int = _constrained_field(POSITIVE_INT_CONSTRAINT) obs_boundary_stride: int = _constrained_field(POSITIVE_INT_CONSTRAINT) + obs_stride_merge_segments: bool = MISSING obs_norm_speed_mps: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) obs_norm_goal_offset_m: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) obs_norm_xy_offset_m: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) @@ -345,6 +346,7 @@ class DriveEnvConfig: obs_norm_road_seg_width_m: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) obs_norm_z_m: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) eval_perceived_size_margin_m: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT) + obs_partner_split_length_m: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT) eval_standstill_jerk_deadband_mps3: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT) obs_range_road_front_m: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) obs_range_road_behind_m: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) @@ -352,8 +354,11 @@ class DriveEnvConfig: obs_range_partner_m: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) obs_range_traffic_control_m: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) partner_blindness_prob: float = _constrained_field(PROBABILITY_CONSTRAINT) + partner_blindness_prob_min: float = _constrained_field(PROBABILITY_CONSTRAINT) partner_blindness_trigger_prob: float = _constrained_field(PROBABILITY_CONSTRAINT) partner_blindness_duration_seconds: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) + partner_blindness_corridor: bool = MISSING + partner_blindness_time_fraction_max: float = _constrained_field(PROBABILITY_CONSTRAINT) phantom_braking_prob: float = _constrained_field(PROBABILITY_CONSTRAINT) phantom_braking_trigger_prob: float = _constrained_field(PROBABILITY_CONSTRAINT) phantom_braking_duration_seconds: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) @@ -386,6 +391,7 @@ class DrivePolicyConfig: critic_num_layers: int = _constrained_field(NONNEGATIVE_INT_CONSTRAINT) critic_head_layer_norm: bool = MISSING fp32_heads: bool = MISSING + fused_slot_encoder: bool = False shared_network: bool = MISSING action_type: ActionType = MISSING @@ -414,7 +420,7 @@ class TrainingConfig: tf32: bool = MISSING rollout_dtype: RolloutDtype = MISSING total_timesteps: int = _constrained_field(POSITIVE_INT_CONSTRAINT) - learning_rate: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) + learning_rate: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT) gamma: float = _constrained_field(PROBABILITY_CONSTRAINT) gae_lambda: float = _constrained_field(PROBABILITY_CONSTRAINT) update_epochs: int = _constrained_field(POSITIVE_INT_CONSTRAINT) @@ -454,6 +460,8 @@ class TrainingConfig: adv_filter_leak_fraction: float = _constrained_field(PROBABILITY_CONSTRAINT) render: bool = MISSING render_interval: int = _constrained_field(POSITIVE_INT_CONSTRAINT) + render_num_maps: int | None = _constrained_field(POSITIVE_INT_CONSTRAINT) + stop_after_epochs: int | None = _constrained_field(POSITIVE_INT_CONSTRAINT) # Derived by load_config from rnn_name and intentionally absent from YAML. use_rnn: bool = MISSING @@ -610,6 +618,8 @@ def _validate_cross_field_constraints(config, context): _raise_config_error(context, "env.spawn_initial_speed", "must not exceed env.base_max_speed_mps") if env["min_goal_spacing"] > env["max_goal_spacing"]: _raise_config_error(context, "env.min_goal_spacing", "must not exceed env.max_goal_spacing") + if env["partner_blindness_prob_min"] > env["partner_blindness_prob"]: + _raise_config_error(context, "env.partner_blindness_prob_min", "must not exceed env.partner_blindness_prob") if env["preload_map_cache"] and not env["use_map_cache"]: _raise_config_error(context, "env.preload_map_cache", "requires env.use_map_cache to be true") @@ -672,6 +682,8 @@ def _validate_cross_field_constraints(config, context): not isinstance(evaluation_benchmarks, str) or not evaluation_benchmarks.strip() ): _raise_config_error(context, "train.evaluation_benchmarks", "must be a non-empty string") + if train["render_num_maps"] is not None and train["render_num_maps"] > config["env"]["num_maps"]: + _raise_config_error(context, "train.render_num_maps", "must not exceed env.num_maps") if not context.startswith("evaluation") and context != "simulation profiling": for field_name in ("batch_size", "bptt_horizon"): if train[field_name] != "auto": diff --git a/pufferlib/ocean/cosim/arch.py b/pufferlib/ocean/cosim/arch.py index 9e21009bc9..1ab69f335b 100644 --- a/pufferlib/ocean/cosim/arch.py +++ b/pufferlib/ocean/cosim/arch.py @@ -23,7 +23,9 @@ "obs_dropout_lane": 0.0, "obs_dropout_boundary": 0.0, "partner_blindness_prob": 0.0, + "partner_blindness_prob_min": 0.0, "partner_blindness_trigger_prob": 0.0, + "partner_blindness_time_fraction_max": 0.0, "phantom_braking_prob": 0.0, "phantom_braking_trigger_prob": 0.0, "traffic_light_behavior": "stop", @@ -35,7 +37,7 @@ "reward_stop_line": 1.0, "reward_ade": 0.0, "reward_goal": 1.0, - "reward_overspeed": 0.05, + "reward_overspeed": 0.1, # training-range max (REWARD_BOUNDS) "reward_comfort": 0.05, "reward_velocity": 0.0025, "reward_lane_align": 0.025, diff --git a/pufferlib/ocean/cosim/carla/README.md b/pufferlib/ocean/cosim/carla/README.md index d80b7137df..6e72985eb8 100644 --- a/pufferlib/ocean/cosim/carla/README.md +++ b/pufferlib/ocean/cosim/carla/README.md @@ -8,10 +8,12 @@ Evaluate a PufferDrive-trained policy with **CaRL's unmodified `original_leaderb Zero changes to CaRL. The leaderboard imports any agent file via `--agent`, so the whole integration is this package: -- `leaderboard_agent.py` — `AutonomousAgent` subclass (entry point `PufferAgent`). - Loads a PufferDrive checkpoint, builds the shadow `Drive` env from the checkpoint's config, and every CARLA tick (the shadow env runs at the 0.05 s tick, no action repeat): overwrites ALL shadow agents (ego included) from CARLA ground truth (ego pose/velocity, nearest vehicles+walkers with true bounding boxes, traffic-light states, route goals from the leaderboard's dense global plan), recomputes observations, runs the policy, integrates one dt to produce a target (speed, yaw), and returns a `carla.VehicleControl` every tick. Get `sensor.camera.rgb` chase cam via `sensors()` — the leaderboard's own standard sensor mechanism, so this doesn't touch the read-only-w.r.t.-CARLA contract; frames stream straight to an mp4 writer. `COSIM_OBS_HTML` records the interactive observation replay (exact policy input/outputs). +- `shadow_ego.py` — `ShadowEgo`, the carla-free core shared with the Leaderboard 1.0 integration (`lb1/`). + Loads a PufferDrive checkpoint, builds the shadow `Drive` env from the checkpoint's config, calibrates the bin frame, maps lights/stop signs and route goals, and every CARLA tick (the shadow env runs at the 0.05 s tick, no action repeat): overwrites ALL shadow agents (ego included) from the CARLA-frame snapshot (ego pose/velocity, nearest vehicles+walkers with true bounding boxes, traffic-light states, route goals from the leaderboard's target points), recomputes observations, runs the policy and integrates one dt into the ego's next pose (or a target (speed, yaw) for the controller). All per-route logs (telemetry, world log, obs replay, obs dump) are written here. +- `leaderboard_agent.py` — `AutonomousAgent` subclass (entry point `PufferAgent`), the CARLA side: reads the ground truth into the snapshot, teleports the ego to the returned pose (resting it on the road mesh), runs CaRL's criteria for the `carla_*` telemetry columns, and returns a `carla.VehicleControl` every tick. Get `sensor.camera.rgb` chase cam via `sensors()` — the leaderboard's own standard sensor mechanism, so this doesn't touch the read-only-w.r.t.-CARLA contract; frames stream straight to an mp4 writer. `COSIM_OBS_HTML` records the interactive observation replay (exact policy input/outputs). - `controller.py` — `TrackingController`: convert that shadow-env kinematic target into actual throttle/brake/steer at CARLA's native tick rate +- `lb1/` — the same policy under carla_garage's Leaderboard 1.0 (CARLA 0.9.10.1, longest6 v1): a Python 3.7 client agent plus `policy_server.py` running `ShadowEgo` in the PufferDrive venv; see `lb1/README.md`. All shadow-env assumptions about the model (observation layout, dynamics model, dt, goal count/radius, reward conditioning) are adopted from the checkpoint's sibling `config.yaml` @@ -41,7 +43,9 @@ Optional env vars (see `leaderboard_agent.py` docstring): `COSIM_DEVICE`, `COSIM_DYNAMICS_SOURCE`, `COSIM_NUM_AGENTS`, `COSIM_MAX_SPEED_MPS` (ego speed cap via C_vel; default = the checkpoint's `base_max_speed_mps`, must stay inside the trained range), `COSIM_ZERO_PARTNER_STOPPED_TIME` (1 = partners' `seconds_stopped` obs feature held at 0, ego -untouched; ablation for the parked-car / red-light-queue avoidance), `COSIM_OBS_HTML`, +untouched; ablation for the parked-car / red-light-queue avoidance), `COSIM_VRU_SIZE_PAD_M` +(added to the length and width of every walker and two-wheeler box in the policy's obs; the +longest6 scripts set 0.8 m), `COSIM_OBS_HTML`, `COSIM_DEBUG_CARLA_VIEW`, `COSIM_RECORD_INFRACTIONS`. ### Ego dynamics source diff --git a/pufferlib/ocean/cosim/carla/controller.py b/pufferlib/ocean/cosim/carla/controller.py index d3b8ee7d7e..2bc8a20da2 100644 --- a/pufferlib/ocean/cosim/carla/controller.py +++ b/pufferlib/ocean/cosim/carla/controller.py @@ -1,9 +1,11 @@ -"""Kinematic target -> carla.VehicleControl tracking controller. +"""Kinematic target -> (steer, throttle, brake) tracking controller. The policy's jerk action is integrated by the shadow PufferDrive env into a target (speed, yaw) one policy-dt ahead; this controller makes CARLA's physics chase that target. CARLA (not the shadow env) moves the ego, so all leaderboard -collision/infraction machinery sees a normally-driven vehicle. +collision/infraction machinery sees a normally-driven vehicle. The simulator +side wraps the returned values in carla.VehicleControl; this module never +imports carla so it also runs in the PufferDrive venv behind lb1/policy_server.py. Longitudinal: PI on speed error -> throttle/brake. Lateral: desired yaw rate over the policy horizon -> inverse kinematic bicycle @@ -15,41 +17,54 @@ import math -import carla import numpy as np from pufferlib.ocean.cosim.carla_bridge import wrap_deg_180 +WHEEL_ROW_COLUMNS = 4 # max_steer_deg, position x/y/z (cm, CARLA's physics control units) +WHEEL_POSITION_CM_PER_M = 100.0 +DEFAULT_WHEELBASE_M = 2.85 # Lincoln MKZ-ish fallbacks for implausible physics controls +DEFAULT_MAX_STEER_DEG = 70.0 +WHEELBASE_RANGE_M = (1.5, 5.0) +MAX_STEER_RANGE_DEG = (10.0, 90.0) + + +def wheel_rows(vehicle): + """[[max_steer_deg, x_cm, y_cm, z_cm], ...] of a live CARLA vehicle's wheels (read-only physics control).""" + return [[w.max_steer_angle, w.position.x, w.position.y, w.position.z] for w in vehicle.get_physics_control().wheels] + + +def vehicle_geometry_from_wheels(rows): + """(wheelbase_m, max_steer_rad) from wheel_rows; defaults for anything implausible.""" + rows = np.asarray(rows, dtype=np.float64).reshape(-1, WHEEL_ROW_COLUMNS) + if len(rows) == 0: + return DEFAULT_WHEELBASE_M, math.radians(DEFAULT_MAX_STEER_DEG) + max_steer_deg = float(rows[:, 0].max()) + steering, fixed = rows[rows[:, 0] > 0.0], rows[rows[:, 0] <= 0.0] + wheelbase = DEFAULT_WHEELBASE_M + if len(steering) and len(fixed): + axle_offset_cm = steering[:, 1:].mean(axis=0) - fixed[:, 1:].mean(axis=0) + wheelbase = float(np.linalg.norm(axle_offset_cm)) / WHEEL_POSITION_CM_PER_M + if not (WHEELBASE_RANGE_M[0] <= wheelbase <= WHEELBASE_RANGE_M[1]): + wheelbase = DEFAULT_WHEELBASE_M + if not (MAX_STEER_RANGE_DEG[0] <= max_steer_deg <= MAX_STEER_RANGE_DEG[1]): + max_steer_deg = DEFAULT_MAX_STEER_DEG + return wheelbase, math.radians(max_steer_deg) + def read_vehicle_geometry(vehicle): - """(wheelbase_m, max_steer_rad) from the vehicle's physics parameters - (read-only). Falls back to Lincoln MKZ-ish defaults on any surprise.""" + """(wheelbase_m, max_steer_rad) of a live CARLA vehicle; defaults on any surprise.""" try: - phys = vehicle.get_physics_control() - wheels = phys.wheels - max_steer_deg = max(w.max_steer_angle for w in wheels) - steering = [w for w in wheels if w.max_steer_angle > 0.0] - fixed = [w for w in wheels if w.max_steer_angle <= 0.0] - if steering and fixed: - front = np.mean([[w.position.x, w.position.y, w.position.z] for w in steering], axis=0) - rear = np.mean([[w.position.x, w.position.y, w.position.z] for w in fixed], axis=0) - wheelbase = float(np.linalg.norm(front - rear)) / 100.0 # positions are in cm - else: - wheelbase = 2.85 - if not (1.5 <= wheelbase <= 5.0): - wheelbase = 2.85 - if not (10.0 <= max_steer_deg <= 90.0): - max_steer_deg = 70.0 - return wheelbase, math.radians(max_steer_deg) + return vehicle_geometry_from_wheels(wheel_rows(vehicle)) except Exception: - return 2.85, math.radians(70.0) + return DEFAULT_WHEELBASE_M, math.radians(DEFAULT_MAX_STEER_DEG) class TrackingController: def __init__( self, - wheelbase_m=2.85, - max_steer_rad=math.radians(70.0), + wheelbase_m=DEFAULT_WHEELBASE_M, + max_steer_rad=math.radians(DEFAULT_MAX_STEER_DEG), horizon_s=0.1, kp_speed=0.7, ki_speed=0.15, @@ -77,8 +92,8 @@ def reset(self): self._yaw_errors = [] def step(self, current_speed, current_yaw_deg, target_speed, target_yaw_deg, tick_dt): - """One control tick. current_* from CARLA ground truth, target_* from the - shadow env's integrated policy action (held between policy steps).""" + """One control tick -> (steer, throttle, brake). current_* from CARLA ground truth, target_* from + the shadow env's integrated policy action (held between policy steps).""" # --- longitudinal --- err = target_speed - current_speed self._speed_errors.append(err) @@ -103,7 +118,7 @@ def step(self, current_speed, current_yaw_deg, target_speed, target_yaw_deg, tic steer_angle = math.atan2(self.wheelbase * desired_yaw_rate, max(current_speed, 1.0)) steer = float(np.clip(steer_angle / self.max_steer, -1.0, 1.0)) - return carla.VehicleControl(steer=steer, throttle=throttle, brake=brake) + return steer, throttle, brake def stats(self): """Per-route tracking fidelity: mean/max absolute speed error (m/s) and diff --git a/pufferlib/ocean/cosim/carla/lb1/README.md b/pufferlib/ocean/cosim/carla/lb1/README.md new file mode 100644 index 0000000000..03e28e2059 --- /dev/null +++ b/pufferlib/ocean/cosim/carla/lb1/README.md @@ -0,0 +1,99 @@ +# PufferDrive policy under carla_garage's CARLA Leaderboard 1.0 (longest6 v1) + +## Goal + +Evaluate a PufferDrive-trained policy with **carla_garage's unmodified Leaderboard 1.0** pipeline +(`leaderboard_evaluator_local.py`, CARLA 0.9.10.1, `longest6.xml` + `eval_scenarios.json`, +`BENCHMARK=longest6`) — the evaluator, scenarios, scoring and `tools/result_parser.py` that produced +the TransFuser++ longest6 numbers. The Leaderboard 2.0 counterpart (CaRL, CARLA 0.9.15) is +`../leaderboard_agent.py`; both share the same shadow env core. + +## Design + +Zero changes to carla_garage. The leaderboard imports any agent file via `--agent`; the difference to +the Leaderboard 2.0 agent is the interpreter: the `garage` conda env is **Python 3.7** (the only Python +the CARLA 0.9.10 egg ships for) and PufferDrive needs 3.9+, so one process cannot hold both. The +integration is therefore a client/server pair around the shared core: + +- `../shadow_ego.py` — `ShadowEgo`, the carla-free core used by both leaderboards: loads the checkpoint + and its `config.yaml`, builds the shadow `Drive` env (clean-eval profile, `cosim/arch.py`), calibrates + the bin frame against CARLA's driving waypoints, maps CARLA lights/stop signs onto the bin, turns the + leaderboard's target points into route goals, and per tick overwrites ALL shadow agents from the + snapshot, runs the policy and integrates one dt. All inputs/outputs are CARLA-frame plain data. +- `leaderboard_agent.py` (this directory) — the `AutonomousAgent` for the garage evaluator (entry point + `PufferAgentLB1`), Python 3.7, imports nothing from `pufferlib`. On the first tick it reads the town, the + ego's box/wheels, the route (`set_global_plan` sparse target points + the dense 1 m route), the driving + waypoints, every traffic light's stop waypoints (from the trigger volume: 0.9.10 has no + `get_stop_waypoints`) and every `traffic.stop` trigger volume, and sends them to the server. Every tick + it ships ego/partner states and light states, and teleports the ego to the returned pose + (`COSIM_DYNAMICS_SOURCE=pufferdrive`) or applies the returned throttle/brake/steer (`carla`). Chase-cam + mp4 (`COSIM_DEBUG_CARLA_VIEW`) and infraction clips (`COSIM_RECORD_INFRACTIONS`) are written here. +- `policy_server.py` — started by the agent in the PufferDrive venv (`COSIM_SERVER_PYTHON`, default + `/.venv/bin/python`), one process per route, `ShadowEgo` behind `lb1_protocol.py` + (length-prefixed JSON over an inherited socket). Telemetry, world log, obs replay and obs dump come out + of this process (same `COSIM_*` variables and file names as the Leaderboard 2.0 agent). +- `ground_truth.py` — the Leaderboard 1.0 criteria (red light, stop sign, outside lanes, collision via a + collision sensor on the client) recomputed from the client's observations, for the `carla_*` telemetry + columns and infraction clips. The leaderboard's own `result.json` is the score. + +Verified 2026-10-05 locally with `k_scaled_0040_1000` on a CARLA 0.9.10.1 offscreen-OpenGL server, policy +on the CPU at the 20 Hz tick: Town01 route 0 scenario-free completed with DS 100 / RC 100 and no +infraction, 655 s of game time in 22 min wall-clock (0.5x real time, 256 background cars), telemetry, +world log, obs replay (13096 frames) and chase-cam mp4 written, and both `analyze_carla_cosim.py` and +`render_carla_obs_html.py` accept the run. + +## How to run + +Prerequisites (this machine): CARLA 0.9.10.1 at `~/ordnung/internal/CARLA_0.9.10`, carla_garage at +`~/ordnung/internal/carla_garage_1`, the `garage` conda env (`environment.yml`, Python 3.7), the repo's +`.venv` for the server. The 0.9.10 egg links `libtiff.so.5`, which newer distros no longer ship (Ubuntu +24.04: `libtiff.so.6` only). The launcher scripts test `import carla` in the garage env and, when it fails, +first put the env's own lib dir on `LD_LIBRARY_PATH` (`conda install -n garage "libtiff<4.5"` installs a +`.so.5` there; the egg does not find it without the path) and otherwise link the system `libtiff.so.6` +under that name next to the results (`LIBTIFF5_COMPAT=0` disables the link). On a distro that still ships +`libtiff.so.5`, nothing happens. + +One route, own CARLA server (the script starts and stops it): + +```bash +CKPT=experiments/k_scaled_0040_1000/final_model.pt \ +ROUTES=~/ordnung/internal/carla_garage_1/leaderboard/data/longest6_split/longest_weathers_0.xml \ +SCENARIOS=~/ordnung/internal/carla_garage_1/leaderboard/data/scenarios/no_scenarios.json \ +bash pufferlib/ocean/cosim/carla/lb1/run_leaderboard.sh +``` + +Against a server you started yourself (`DISPLAY= ./CarlaUE4.sh -opengl -carla-rpc-port=2000`), from VS Code: +launch config "Cosim: CARLA leaderboard 1.0 agent (garage)" runs `run_evaluator.py --model-dir +--routes ...` in the garage env with outputs under `/eval/carla_leaderboard_v1/`. It points +`LD_LIBRARY_PATH` at `experiments/compat_lib` (gitignored), which holds the `libtiff.so.5` link: +`mkdir -p experiments/compat_lib && ln -s /usr/lib/x86_64-linux-gnu/libtiff.so.6 experiments/compat_lib/libtiff.so.5`. + +All 36 longest6 routes on the cluster: `scripts/kesai/14_carla_longest6_v1.sh` (one CARLA 0.9.10 server ++ evaluator per GPU over the `longest6_split` files, retries, garage's `result_parser.py` aggregation, +the nuPlan-style HTML report and the obs replay gallery, like `11_carla_longest6_v2.sh`). + +The evaluator reads `BENCHMARK=longest6` (background traffic on every spawn point, stop-sign penalty +1.0) and `ROUTES` (names the route records) from the environment; `run_leaderboard.sh`, +`run_evaluator.py` and script 14 export them. + +## Differences to the Leaderboard 2.0 co-sim + +- Target points every **50 m** (`downsample_route(route, 50)`) instead of 200 m, so the sliding goal + window is denser; junction entries/exits and lane changes are target points in both. +- Scoring is Leaderboard 1.0's: penalties 0.50/0.60/0.65/0.70 (pedestrian/vehicle/layout/red light), + stop sign 1.0 under `BENCHMARK=longest6`, no min-speed or scenario-timeout infractions, blocked after + 180 s (`BLOCKED_THRESHOLD`), route timeout 0.8 s/m. `result.json` records carry nine infraction keys. +- The ego is `vehicle.lincoln.mkz2017`; routes carry weather. +- CARLA 0.9.10: no `World.cast_ray`, so the ego's body chord rests on the lane's OpenDRIVE height profile + and a downward semantic lidar probes the road mesh's offset from it (Town02's mesh sits 0.22 m above its + OpenDRIVE; the egg's `Rotation.get_up_vector()` points down, so the pitch comes from the profile too); + stop waypoints from the trigger volume. +- The `carla_*` telemetry columns come from `ground_truth.py` (straight-ahead probes in place of lane + waypoints for the stop sign, no same-lane check for the red light), not from CaRL's criteria. + +## Outputs + +As for the Leaderboard 2.0 agent: `result.json` (the leaderboard's), `telemetry/.csv`, +`world_log/.npz`, `obs_html/.replay.zlib` (rendered by +`scripts/eval/render_carla_obs_html.py`), `carla_view/.mp4`. `scripts/eval/analyze_carla_cosim.py` +builds the HTML report from a run folder holding `result.json`, `world_log/`, `telemetry/`, `carla_view/`. diff --git a/pufferlib/ocean/cosim/carla/lb1/__init__.py b/pufferlib/ocean/cosim/carla/lb1/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pufferlib/ocean/cosim/carla/lb1/ground_truth.py b/pufferlib/ocean/cosim/carla/lb1/ground_truth.py new file mode 100644 index 0000000000..73da099f39 --- /dev/null +++ b/pufferlib/ocean/cosim/carla/lb1/ground_truth.py @@ -0,0 +1,178 @@ +"""CARLA-side ground truth of the Leaderboard 1.0 criteria, recomputed in the policy server from what the +Python 3.7 client observes (it cannot run CaRL's criteria the Leaderboard 2.0 agent uses). Feeds the +carla_* telemetry columns and the infraction clips; the leaderboard's own result.json stays the score. + +Per-tick observation sent by lb1/leaderboard_agent.py (CARLA frame): + {"collision": [other actor type ids since the last tick], + "lane": [driving_distance_m, driving_lane_width_m, parking_distance_m (-1: no parking lane), + parking_lane_width_m, driving_lane_yaw_deg, driving_is_junction]} +plus the snapshot's ego state. Ports of scenario_runner's RunningRedLightTest (tail line crossing the +0.4-lane-width stop line of a red light within 15 m, once per light), RunningStopTest (entered a stop +trigger volume, left its 20 m influence without a standstill) and OutsideRouteLanesTest (farther than +half a lane plus 1.3 m from both the driving and the parking lane, or facing more than 120 deg away from +a non-junction lane). +""" + +import math + +from pufferlib.ocean.cosim.carla.shadow_ego import STATE_X, STATE_Y, STATE_YAW_DEG, ego_speed_from_state + + +RED_LIGHT_DISTANCE_M = 15.0 # RunningRedLightTest.DISTANCE_LIGHT +STOP_LINE_HALF_WIDTH_FRACTION = 0.4 # of the lane width, RunningRedLightTest +TAIL_CLOSE_FRACTION = 0.8 # tail line from 0.8 * extent to extent + 1 m behind the ego centre +TAIL_FAR_MARGIN_M = 1.0 +STOP_SIGN_SPEED_THRESHOLD_MPS = 0.1 # RunningStopTest.SPEED_THRESHOLD +STOP_SIGN_PROXIMITY_M = 50.0 # RunningStopTest.PROXIMITY_THRESHOLD +STOP_SIGN_LOOKAHEAD_STEPS = 20 # RunningStopTest: 20 waypoints, WAYPOINT_STEP 1 m +STOP_SIGN_LOOKAHEAD_STEP_M = 1.0 +OUTSIDE_LANE_ALLOWED_DISTANCE_M = 1.3 # OutsideRouteLanesTest.ALLOWED_OUT_DISTANCE +WRONG_LANE_MAX_ANGLE_DEG = 120.0 # OutsideRouteLanesTest.MAX_ALLOWED_VEHICLE_ANGLE +( + LANE_DRIVING_DISTANCE, + LANE_DRIVING_WIDTH, + LANE_PARKING_DISTANCE, + LANE_PARKING_WIDTH, + LANE_YAW_DEG, + LANE_IS_JUNCTION, +) = range(6) + + +def segments_intersect(a0, a1, b0, b1): + """Proper or touching intersection of segments a0-a1 and b0-b1 (2-D).""" + + def orientation(p, q, r): + return (q[0] - p[0]) * (r[1] - p[1]) - (q[1] - p[1]) * (r[0] - p[0]) + + def on_segment(p, q, r): + return min(p[0], q[0]) <= r[0] <= max(p[0], q[0]) and min(p[1], q[1]) <= r[1] <= max(p[1], q[1]) + + o1, o2 = orientation(a0, a1, b0), orientation(a0, a1, b1) + o3, o4 = orientation(b0, b1, a0), orientation(b0, b1, a1) + if (o1 > 0) != (o2 > 0) and (o3 > 0) != (o4 > 0) and o1 != 0 and o2 != 0 and o3 != 0 and o4 != 0: + return True + if o1 == 0 and on_segment(a0, a1, b0): + return True + if o2 == 0 and on_segment(a0, a1, b1): + return True + if o3 == 0 and on_segment(b0, b1, a0): + return True + return o4 == 0 and on_segment(b0, b1, a1) + + +def point_inside_box(point, center, extent): + """RunningStopTest.point_inside_boundingbox: axis-aligned test, the trigger volume's rotation ignored.""" + return abs(point[0] - center[0]) < extent[0] and abs(point[1] - center[1]) < extent[1] + + +class LeaderboardOneGroundTruth: + def __init__(self, lights, stop_signs, ego_extent_x): + """lights / stop_signs: the init_route geometry dicts (carla_bridge.light_geometry_from_carla / + stop_sign_geometry_from_carla layout), ego_extent_x: half length of the ego's bounding box.""" + self.stop_lines_by_light = {} + self.trigger_by_light = {} + for light in lights: + lines = [] + for stop_wp in light["stop_waypoints"]: + yaw = math.radians(stop_wp["yaw_deg"]) + half_width = STOP_LINE_HALF_WIDTH_FRACTION * stop_wp["lane_width"] + across = (-math.sin(yaw) * half_width, math.cos(yaw) * half_width) + left = (stop_wp["x"] + across[0], stop_wp["y"] + across[1]) + right = (stop_wp["x"] - across[0], stop_wp["y"] - across[1]) + lines.append(((math.cos(yaw), math.sin(yaw)), left, right)) + self.stop_lines_by_light[int(light["id"])] = lines + self.trigger_by_light[int(light["id"])] = (light["trigger_x"], light["trigger_y"]) + self.stop_signs = [(sign["center"][:2], sign["extent"][:2]) for sign in stop_signs] + self.ego_extent_x = float(ego_extent_x) + self.last_red_light_id = None + self.target_stop_sign = None + self.stop_completed = False + self.affected_by_stop = False + + def update(self, snapshot, observation): + """-> {'collision', 'offroad', 'red_light', 'stop_sign'} flags (1.0 fired this tick) from one snapshot.""" + ego = [float(v) for v in snapshot["ego"]] + light_states = {int(light_id): state_name for light_id, state_name in snapshot["light_states"]} + return { + "collision": float(bool(observation.get("collision"))), + "offroad": float(self._outside_lanes(ego, observation["lane"])), + "red_light": float(self._ran_red_light(ego, light_states)), + "stop_sign": float(self._ran_stop_sign(ego, observation["lane"])), + } + + def _ran_red_light(self, ego, light_states): + yaw = math.radians(ego[STATE_YAW_DEG]) + forward = (math.cos(yaw), math.sin(yaw)) + x, y = ego[STATE_X], ego[STATE_Y] + tail_close = ( + x - TAIL_CLOSE_FRACTION * self.ego_extent_x * forward[0], + y - TAIL_CLOSE_FRACTION * self.ego_extent_x * forward[1], + ) + tail_far_m = self.ego_extent_x + TAIL_FAR_MARGIN_M + tail_far = (x - tail_far_m * forward[0], y - tail_far_m * forward[1]) + for light_id, lines in self.stop_lines_by_light.items(): + if light_states.get(light_id) != "Red" or light_id == self.last_red_light_id: + continue + trigger_x, trigger_y = self.trigger_by_light[light_id] + if math.hypot(trigger_x - x, trigger_y - y) > RED_LIGHT_DISTANCE_M: + continue + for lane_forward, left, right in lines: + if forward[0] * lane_forward[0] + forward[1] * lane_forward[1] <= 0.0: + continue # this stop line does not face our lane + if segments_intersect(tail_close, tail_far, left, right): + self.last_red_light_id = light_id + return True + return False + + def _affected_by_stop_sign(self, ego, sign): + """RunningStopTest.is_actor_affected_by_stop with straight-ahead probes in place of lane waypoints.""" + center, extent = sign + x, y = ego[STATE_X], ego[STATE_Y] + if math.hypot(center[0] - x, center[1] - y) > STOP_SIGN_PROXIMITY_M: + return False + yaw = math.radians(ego[STATE_YAW_DEG]) + for k in range(STOP_SIGN_LOOKAHEAD_STEPS + 1): + probe = ( + x + k * STOP_SIGN_LOOKAHEAD_STEP_M * math.cos(yaw), + y + k * STOP_SIGN_LOOKAHEAD_STEP_M * math.sin(yaw), + ) + if point_inside_box(probe, center, extent): + return True + return False + + def _ran_stop_sign(self, ego, lane): + if self.target_stop_sign is None: + lane_yaw = math.radians(lane[LANE_YAW_DEG]) + yaw = math.radians(ego[STATE_YAW_DEG]) + if math.cos(yaw - lane_yaw) <= 0.0: + return False # ignore all when going in a wrong lane + for sign in self.stop_signs: + if self._affected_by_stop_sign(ego, sign): + self.target_stop_sign = sign + break + return False + if not self.stop_completed and ego_speed_from_state(ego) < STOP_SIGN_SPEED_THRESHOLD_MPS: + self.stop_completed = True + center, extent = self.target_stop_sign + if not self.affected_by_stop and point_inside_box((ego[STATE_X], ego[STATE_Y]), center, extent): + self.affected_by_stop = True + if self._affected_by_stop_sign(ego, self.target_stop_sign): + return False + ran = self.affected_by_stop and not self.stop_completed + self.target_stop_sign = None + self.stop_completed = False + self.affected_by_stop = False + return ran + + def _outside_lanes(self, ego, lane): + driving_distance, parking_distance = lane[LANE_DRIVING_DISTANCE], lane[LANE_PARKING_DISTANCE] + if parking_distance >= 0.0 and driving_distance >= parking_distance: + distance, lane_width = parking_distance, lane[LANE_PARKING_WIDTH] + else: + distance, lane_width = driving_distance, lane[LANE_DRIVING_WIDTH] + outside = distance > lane_width / 2.0 + OUTSIDE_LANE_ALLOWED_DISTANCE_M + if lane[LANE_IS_JUNCTION] > 0.0: + return outside # lanes and roads are too chaotic at junctions + angle = (lane[LANE_YAW_DEG] - ego[STATE_YAW_DEG]) % 360.0 + wrong_lane = WRONG_LANE_MAX_ANGLE_DEG <= angle <= 360.0 - WRONG_LANE_MAX_ANGLE_DEG + return outside or wrong_lane diff --git a/pufferlib/ocean/cosim/carla/lb1/lb1_protocol.py b/pufferlib/ocean/cosim/carla/lb1/lb1_protocol.py new file mode 100644 index 0000000000..184d9379ff --- /dev/null +++ b/pufferlib/ocean/cosim/carla/lb1/lb1_protocol.py @@ -0,0 +1,57 @@ +"""Length-prefixed JSON messages between lb1/leaderboard_agent.py (CARLA 0.9.10, Python 3.7) and +lb1/policy_server.py (PufferDrive venv). Standard library only: the agent imports this file as a +sibling module (the leaderboard puts the agent's directory on sys.path), the server as a package module. + +Handshake, then one request/reply pair per CARLA tick: + server -> client {"type": "config", "partner_slots", "partner_max_abs_dz_m", "light_probe_steps_m", + "dynamics_source", "infraction_flags"} + client -> server {"type": "init", "route": shadow_ego.init_route dict} -> {"type": "ready"} + client -> server {"type": "tick", "snapshot": shadow_ego.step dict, + "ground_truth": lb1/ground_truth.py observation or null} -> {"type": "control", + "motion": [...] | null, "control": [steer, throttle, brake] | null, + "pd_flags": {...} | null, "carla_flags": {...} | null} + client -> server {"type": "finish"} -> {"type": "done"} +""" + +import json +import struct + + +HEADER = struct.Struct("!I") +MAX_MESSAGE_BYTES = 256 * 1024 * 1024 +RECV_CHUNK_BYTES = 1 << 20 + + +def send_message(sock, message): + payload = json.dumps(message, separators=(",", ":")).encode("utf-8") + if len(payload) > MAX_MESSAGE_BYTES: + raise ValueError(f"message of {len(payload)} bytes exceeds MAX_MESSAGE_BYTES") + sock.sendall(HEADER.pack(len(payload)) + payload) + + +def recv_message(sock): + """The next message, or None once the peer closed the connection between messages.""" + header = _recv_exact(sock, HEADER.size) + if header is None: + return None + (length,) = HEADER.unpack(header) + if length > MAX_MESSAGE_BYTES: + raise ValueError(f"incoming message of {length} bytes exceeds MAX_MESSAGE_BYTES") + payload = _recv_exact(sock, length) + if payload is None: + raise ConnectionError("peer closed the connection inside a message") + return json.loads(payload.decode("utf-8")) + + +def _recv_exact(sock, byte_count): + chunks = [] + remaining = byte_count + while remaining > 0: # every pass consumes >= 1 byte or returns + chunk = sock.recv(min(remaining, RECV_CHUNK_BYTES)) + if not chunk: + if chunks: + raise ConnectionError("peer closed the connection inside a message") + return None + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) diff --git a/pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py b/pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py new file mode 100644 index 0000000000..aab4292083 --- /dev/null +++ b/pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py @@ -0,0 +1,656 @@ +"""PufferDrive policy as a CARLA Leaderboard 1.0 agent (carla_garage's leaderboard_evaluator_local.py, +CARLA 0.9.10.1, longest6 v1). + +The leaderboard runs in carla_garage's `garage` conda env: Python 3.7, the only Python the CARLA 0.9.10 +egg supports, and one that cannot import PufferDrive (Python >= 3.9). This file is therefore the CARLA +side only -- the counterpart of ../leaderboard_agent.py's CARLA code: it reads ground truth (ego, nearby +actors, lights, stop signs, the route), ships it over a socket to lb1/policy_server.py (started here in +the PufferDrive venv, one process per route, shadow_ego.ShadowEgo inside) and moves the ego the way the +reply says. Nothing from `pufferlib` is imported here; lb1_protocol.py is a sibling module. + +Usage (BENCHMARK=longest6 and ROUTES= exported, see run_leaderboard.sh / 14_carla_longest6_v1.sh): + + python $GARAGE_WORK_DIR/leaderboard/leaderboard/leaderboard_evaluator_local.py \ + --routes $GARAGE_WORK_DIR/leaderboard/data/longest6_split/longest_weathers_0.xml \ + --scenarios $GARAGE_WORK_DIR/leaderboard/data/scenarios/eval_scenarios.json \ + --agent /path/to/pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py \ + --agent-config /path/to/experiments//final_model.pt \ + --checkpoint /path/to/results/result.json --track MAP --timeout 600 + +Environment variables (this process; the shadow env's COSIM_DEVICE ... COSIM_WORLD_LOG are read by +policy_server.py from the inherited environment, see shadow_ego.py): + SCENARIO_RUNNER_ROOT REQUIRED: carla_garage's scenario_runner checkout + COSIM_SERVER_PYTHON python of the PufferDrive venv that runs policy_server.py + (default: /.venv/bin/python) + COSIM_DEBUG_CARLA_VIEW=/dir write a CARLA chase-camera mp4 per route (native tick rate) + COSIM_RECORD_INFRACTIONS=/dir write a short chase-cam clip (last ~5 s) per ego infraction flagged by + the shadow env (pd_*) or by the Leaderboard 1.0 criteria recomputed in + the server (carla_*, lb1/ground_truth.py) + +CARLA 0.9.10 differences handled here: no World.cast_ray (the ego's body chord rests on the lane's OpenDRIVE +height profile, lifted by the road mesh's offset from it as probed by a downward semantic lidar; the egg's +Rotation.get_up_vector() points down, so the pitch comes from the profile too); stop waypoints come from the +trigger volume with RunningRedLightTest's own recipe on both sides (carla_bridge.light_stop_waypoints). +""" + +import math +import os +import re +import socket +import subprocess +from collections import deque +from datetime import datetime +from pathlib import Path + +import carla +import lb1_protocol +import numpy as np +from leaderboard.autoagents import autonomous_agent +from srunner.scenariomanager.carla_data_provider import CarlaDataProvider + + +REPO_ROOT = Path(__file__).resolve().parents[5] +POLICY_SERVER = Path(__file__).resolve().with_name("policy_server.py") +DEFAULT_SERVER_PYTHON = REPO_ROOT / ".venv" / "bin" / "python" +SERVER_ENV_DROPPED = ("LD_LIBRARY_PATH", "PYTHONPATH", "PYTHONHOME") # the garage env's libs/paths break the venv +PROTOCOL_TIMEOUT_S = 600.0 # one reply per tick; the evaluator's own agent watchdog is --timeout (600 s) +SERVER_EXIT_TIMEOUT_S = 60.0 + +INFRACTION_CLIP_SECONDS = 5.0 +INFRACTION_MIN_SEPARATION_M = 10.0 +CARLA_VIEW_SENSOR_ID = "puffer_chase_cam" +CARLA_VIEW_WIDTH, CARLA_VIEW_HEIGHT, CARLA_VIEW_FOV = 960, 540, 90 +CARLA_VIEW_TRANSFORM = dict(x=-2.0, y=0.0, z=2.1, roll=0.0, pitch=-15.0, yaw=0.0) # within the 3 m sensor radius +MP4_FOURCC = "mp4v" + +OFFSET_CALIBRATION_SPACING_M = 10.0 # carla_bridge.driving_waypoint_samples +TRIGGER_SAMPLE_SPAN_FRACTION = 0.9 # RunningRedLightTest.get_traffic_light_waypoints: avoids adjacent lanes +TRIGGER_SAMPLE_STEP_M = 1.0 +JUNCTION_ADVANCE_STEP_M = 0.5 +MAX_JUNCTION_ADVANCE_STEPS = 400 # 200 m: a stop waypoint farther from its junction is a map error +NO_LIGHT_ID = -1 +NO_PARKING_LANE = -1.0 + +WHEEL_POSITION_CM_PER_M = 100.0 +ROAD_PROBE_MOUNT_Z_M = 0.5 # inside the hull, above any plausible penetration +ROAD_PROBE_RANGE_M = 20.0 +ROAD_PROBE_POINTS_PER_SECOND = 2000 +ROAD_SURFACE_TAGS = (6, 7, 14, 15, 16, 22) # CityObjectLabel RoadLines, Roads, Ground, Bridge, RailTrack, Terrain +REST_ORIGIN_ABOVE_ROAD_M = 0.033 # mkz2017 origin height at rest on a flat road (measured, 0.9.10) +MESH_LIFT_MARGIN_M = 0.05 +MESH_OFFSET_DEADBAND_M = 0.02 +MESH_OFFSET_MAX_STEP_M = 0.3 + + +def get_entry_point(): + return "PufferAgentLB1" + + +def wrap_deg_180(d): + return (d + 180.0) % 360.0 - 180.0 + + +def actor_state_row(actor): + """shadow_ego ACTOR_STATE of a live CARLA actor (CARLA frame).""" + tf = actor.get_transform() + v = actor.get_velocity() + av = actor.get_angular_velocity() + acc = actor.get_acceleration() + return [tf.location.x, tf.location.y, tf.location.z, tf.rotation.yaw, v.x, v.y, v.z, av.z, acc.x, acc.y] + + +def partner_row(actor): + """shadow_ego PARTNER_ROW of a live CARLA vehicle or walker.""" + box = actor.bounding_box + is_walker = 1.0 if "walker" in actor.type_id else 0.0 + is_two_wheeler = 1.0 if actor.attributes.get("number_of_wheels") == "2" else 0.0 + return actor_state_row(actor) + [ + box.location.z, + box.extent.x, + box.extent.y, + box.extent.z, + is_walker, + is_two_wheeler, + actor.id, + ] + + +def plan_xyz(plan): + """[(carla.Transform, RoadOption)] -> [[x, y, z], ...] CARLA-frame positions.""" + return [[t.location.x, t.location.y, t.location.z] for t, _ in plan] + + +def lane_z_at(wp, offset_m): + """Lane height offset_m along the lane from wp (negative = behind); None past the lane network.""" + if offset_m == 0.0: + return wp.transform.location.z + found = wp.next(offset_m) if offset_m > 0.0 else wp.previous(-offset_m) + return found[0].transform.location.z if found else None + + +def lane_chord_pose(wp, sample_offsets_m): + """(origin z, pitch_deg) of a rigid body chord resting on the lane's height profile around wp.""" + sample_z = [lane_z_at(wp, s) for s in sample_offsets_m] + if any(z is None for z in sample_z): + return wp.transform.location.z, wrap_deg_180(wp.transform.rotation.pitch) + slope = (sample_z[-1] - sample_z[0]) / (sample_offsets_m[-1] - sample_offsets_m[0]) + z = max(z_i - slope * s_i for z_i, s_i in zip(sample_z, sample_offsets_m)) + return z, math.degrees(math.atan(slope)) + + +def driving_waypoint_samples(carla_map): + """[[x, y, z, right_x, right_y], ...] of the non-junction driving waypoints (carla_bridge.driving_waypoint_samples).""" + rows = [] + for wp in carla_map.generate_waypoints(OFFSET_CALIBRATION_SPACING_M): + if wp.is_junction or str(wp.lane_type) != "Driving": + continue + loc, right = wp.transform.location, wp.transform.get_right_vector() + rows.append([loc.x, loc.y, loc.z, right.x, right.y]) + return rows + + +def light_stop_waypoints(light, carla_map): + """Stop waypoints of a traffic light from its trigger volume: one waypoint per lane the volume spans, + advanced to the junction entry (RunningRedLightTest.get_traffic_light_waypoints): the py3.7 copy of + carla_bridge.light_stop_waypoints. TrafficLight.get_stop_waypoints is NOT that list in Town01/02.""" + base_transform = light.get_transform() + trigger = light.trigger_volume + lane_waypoints = [] + for x in np.arange( + -TRIGGER_SAMPLE_SPAN_FRACTION * trigger.extent.x, + TRIGGER_SAMPLE_SPAN_FRACTION * trigger.extent.x, + TRIGGER_SAMPLE_STEP_M, + ): + point = base_transform.transform(trigger.location + carla.Location(x=float(x))) + wp = carla_map.get_waypoint(point) + if wp is None: + continue + if not lane_waypoints or lane_waypoints[-1].road_id != wp.road_id or lane_waypoints[-1].lane_id != wp.lane_id: + lane_waypoints.append(wp) + stop_waypoints = [] + for wp in lane_waypoints: + for _ in range(MAX_JUNCTION_ADVANCE_STEPS): + if wp.is_junction: + break + ahead = wp.next(JUNCTION_ADVANCE_STEP_M) + if not ahead or ahead[0].is_junction: + break + wp = ahead[0] + stop_waypoints.append(wp) + return stop_waypoints + + +def light_geometry(lights, carla_map, probe_steps_m): + """carla_bridge.light_geometry_from_carla layout, built with the 0.9.10 API (probe steps from the server's config).""" + geometry = [] + for light in lights: + location = light.get_location() + trigger = light.get_transform().transform(light.trigger_volume.location) + stop_waypoints = [] + for wp in light_stop_waypoints(light, carla_map): + backward, forward = [], [] + for step_m in probe_steps_m: + probes = [wp] if step_m == 0.0 else wp.previous(step_m) + backward.append( + None if not probes else [probes[0].transform.location.x, probes[0].transform.location.y] + ) + for step_m in probe_steps_m[1:]: + probes = wp.next(step_m) or [] + forward.append( + None + if not probes + else [ + probes[0].transform.location.x, + probes[0].transform.location.y, + probes[0].transform.rotation.yaw, + ] + ) + stop_waypoints.append( + { + "x": wp.transform.location.x, + "y": wp.transform.location.y, + "yaw_deg": wp.transform.rotation.yaw, + "lane_width": wp.lane_width, + "backward": backward, + "forward": forward, + } + ) + geometry.append( + { + "id": light.id, + "x": location.x, + "y": location.y, + "trigger_x": trigger.x, + "trigger_y": trigger.y, + "stop_waypoints": stop_waypoints, + } + ) + return geometry + + +def stop_sign_geometry(world, carla_map): + """carla_bridge.stop_sign_geometry_from_carla layout, built with the 0.9.10 API.""" + geometry = [] + for actor in world.get_actors().filter("traffic.stop"): + actor_transform = actor.get_transform() + trigger = actor.trigger_volume + center = actor_transform.transform(trigger.location) + waypoint = carla_map.get_waypoint(center, project_to_road=True, lane_type=carla.LaneType.Driving) + if waypoint is None: + continue + forward = waypoint.transform.get_forward_vector() + lane_location = waypoint.transform.location + geometry.append( + { + "id": actor.id, + "center": [center.x, center.y, center.z], + "extent": [trigger.extent.x, trigger.extent.y, trigger.extent.z], + "yaw_deg": actor_transform.rotation.yaw, + "lane": [ + lane_location.x, + lane_location.y, + lane_location.z, + math.degrees(math.atan2(forward.y, forward.x)), + ], + } + ) + return geometry + + +class Mp4Writer: + """Streams RGB uint8 frames to an mp4 via OpenCV (carla_cosim.Mp4Writer, importable here).""" + + def __init__(self, out_path, fps): + self.out_path = str(out_path) + self.fps = float(fps) + self.writer = None + + def append_data(self, frame): + import cv2 + + frame = np.ascontiguousarray(frame) + if self.writer is None: + height, width = frame.shape[:2] + self.writer = cv2.VideoWriter(self.out_path, cv2.VideoWriter_fourcc(*MP4_FOURCC), self.fps, (width, height)) + if not self.writer.isOpened(): + raise RuntimeError(f"Mp4Writer({self.out_path}): OpenCV could not open the video writer") + self.writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) + + def close(self): + if self.writer is not None: + self.writer.release() + self.writer = None + + +def write_mp4(out_path, frames, fps): + writer = Mp4Writer(out_path, fps) + for frame in frames: + writer.append_data(frame) + writer.close() + + +class RoadProbe: + """Semantic lidar firing straight down from inside the hull: 0.9.10's substitute for World.cast_ray.""" + + def __init__(self, world, vehicle, tick_dt): + blueprint = world.get_blueprint_library().find("sensor.lidar.ray_cast_semantic") + for name, value in ( + ("channels", "1"), + ("range", str(ROAD_PROBE_RANGE_M)), + ("points_per_second", str(ROAD_PROBE_POINTS_PER_SECOND)), + ("rotation_frequency", str(1.0 / tick_dt)), + ("upper_fov", "-90"), + ("lower_fov", "-90"), + ): + blueprint.set_attribute(name, value) + mount = carla.Transform(carla.Location(x=0.0, y=0.0, z=ROAD_PROBE_MOUNT_Z_M)) + self.sensor = world.spawn_actor(blueprint, mount, attach_to=vehicle) + self.sensor.listen(self._on_data) + self._frame = None + self._distance_m = None + self._consumed_frame = None + + def _on_data(self, data): + distances = sorted(-d.point.z for d in data if d.object_tag in ROAD_SURFACE_TAGS) + if distances: + self._frame, self._distance_m = data.frame, distances[len(distances) // 2] + + def fresh_distance_m(self): + """Median road distance along the body's down axis from the newest unread scan; None when nothing new.""" + if self._frame is None or self._frame == self._consumed_frame: + return None + self._consumed_frame = self._frame + return self._distance_m + + def destroy(self): + self.sensor.stop() + self.sensor.destroy() + + +class PolicyServerLink: + """policy_server.py as a child process in the PufferDrive venv, one request/reply per message.""" + + def __init__(self, checkpoint, route_tag): + python = os.environ.get("COSIM_SERVER_PYTHON", str(DEFAULT_SERVER_PYTHON)) + if not Path(python).is_file(): + raise FileNotFoundError(f"COSIM_SERVER_PYTHON={python} does not exist (the PufferDrive venv python)") + self.sock, child_sock = socket.socketpair() + self.sock.settimeout(PROTOCOL_TIMEOUT_S) + env = {k: v for k, v in os.environ.items() if k not in SERVER_ENV_DROPPED} + env["PYTHONPATH"] = str(REPO_ROOT) + env["PYTHONUNBUFFERED"] = "1" # the server shares the evaluator's log file + command = [ + python, + str(POLICY_SERVER), + "--checkpoint", + checkpoint, + "--socket-fd", + str(child_sock.fileno()), + "--route-tag", + route_tag, + ] + self.process = subprocess.Popen(command, cwd=str(REPO_ROOT), env=env, pass_fds=(child_sock.fileno(),)) + child_sock.close() + self.config = self._receive() + if self.config.get("type") != "config": + raise RuntimeError(f"policy_server.py handshake failed: {self.config!r}") + + def _receive(self): + try: + message = lb1_protocol.recv_message(self.sock) + except (ConnectionError, socket.timeout) as error: # noqa: UP041 (alias of TimeoutError only from 3.10) + raise RuntimeError(f"policy_server.py connection lost ({error}, exit code {self.process.poll()!r})") + if message is None: + raise RuntimeError(f"policy_server.py exited (code {self.process.poll()!r}) without replying") + return message + + def request(self, message): + lb1_protocol.send_message(self.sock, message) + return self._receive() + + def close(self): + self.sock.close() + try: + self.process.wait(timeout=SERVER_EXIT_TIMEOUT_S) + except subprocess.TimeoutExpired: + self.process.kill() + print("[puffer_agent_lb1] policy_server.py did not exit; killed", flush=True) + + +class PufferAgentLB1(autonomous_agent.AutonomousAgent): + def setup(self, path_to_conf_file, route_index=None): + if "SCENARIO_RUNNER_ROOT" not in os.environ: + raise RuntimeError("SCENARIO_RUNNER_ROOT is not set; export carla_garage's scenario_runner checkout") + self.track = autonomous_agent.Track.MAP + route_index = re.sub(r"[^\w.-]", "_", str(route_index)) if route_index else "route" + self.route_tag = f"{route_index}_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}" + self.debug_carla_view_dir = os.environ.get("COSIM_DEBUG_CARLA_VIEW", None) + self.record_infractions_dir = os.environ.get("COSIM_RECORD_INFRACTIONS", None) + self.ground_truth = bool(self.record_infractions_dir or os.environ.get("COSIM_TELEMETRY")) + self.server = PolicyServerLink(str(Path(path_to_conf_file).resolve()), self.route_tag) + self.partner_slots = int(self.server.config["partner_slots"]) + self.partner_max_abs_dz_m = float(self.server.config["partner_max_abs_dz_m"]) + self.light_probe_steps_m = [float(v) for v in self.server.config["light_probe_steps_m"]] + self.dense_global_plan_world_coord = None + self.step = -1 + self.initialized = False + self.carla_view_writer = None + self._collision_sensor = None + self._collision_events = [] + self.road_probe = None + self.body_chord_offsets_m = None + self.mesh_offset_m = 0.0 + self._previous_body_z = None + self._previous_lane_center_z = None + self._previous_pitch_rad = 0.0 + + def sensors(self): + if not self.debug_carla_view_dir and not self.record_infractions_dir: + return [] + sensor = { + "type": "sensor.camera.rgb", + "id": CARLA_VIEW_SENSOR_ID, + "width": CARLA_VIEW_WIDTH, + "height": CARLA_VIEW_HEIGHT, + "fov": CARLA_VIEW_FOV, + } + sensor.update(CARLA_VIEW_TRANSFORM) + return [sensor] + + def set_global_plan(self, global_plan_gps, global_plan_world_coord): + """The base class keeps only the 50 m target points; the dense 1 m route gives the goals their direction.""" + super().set_global_plan(global_plan_gps, global_plan_world_coord) + self.dense_global_plan_world_coord = global_plan_world_coord + + def _init_on_first_step(self): + """Deferred init (the ego and world only exist once the route runs).""" + self.vehicle = CarlaDataProvider.get_hero_actor() + self.world = self.vehicle.get_world() + if not self.sensors(): # no camera anywhere: skip the server's spectator rendering + settings = self.world.get_settings() + settings.no_rendering_mode = True + self.world.apply_settings(settings) + self.cmap = self.world.get_map() + town = CarlaDataProvider.get_map().name.split("/")[-1] + self.tick_dt = float(self.world.get_settings().fixed_delta_seconds) + self.lights = list(self.world.get_actors().filter("traffic.traffic_light")) + body = self.vehicle.bounding_box + wheels = [ + [w.max_steer_angle, w.position.x, w.position.y, w.position.z] + for w in self.vehicle.get_physics_control().wheels + ] + front_left, rear_left = wheels[0][1:], wheels[2][1:] + half_wheelbase_m = ( + 0.5 * math.sqrt(sum((a - b) ** 2 for a, b in zip(front_left, rear_left))) / WHEEL_POSITION_CM_PER_M + ) + body_rear_m, body_front_m = body.location.x - body.extent.x, body.location.x + body.extent.x + assert body_rear_m < -half_wheelbase_m < half_wheelbase_m < body_front_m + self.body_chord_offsets_m = (body_rear_m, -half_wheelbase_m, 0.0, half_wheelbase_m, body_front_m) + self.road_probe = RoadProbe(self.world, self.vehicle, self.tick_dt) + route = { + "town": town, + "tick_dt": self.tick_dt, + "ego": { + "extent": [body.extent.x, body.extent.y, body.extent.z], + "box_z": body.location.z, + "wheels": wheels, + "state": actor_state_row(self.vehicle), + }, + "sparse_plan": plan_xyz(self._global_plan_world_coord), + "dense_plan": plan_xyz(self.dense_global_plan_world_coord), + "driving_waypoints": driving_waypoint_samples(self.cmap), + "lights": light_geometry(self.lights, self.cmap, self.light_probe_steps_m), + "stop_signs": stop_sign_geometry(self.world, self.cmap), + "infraction_flags": bool(self.record_infractions_dir), + "ground_truth": self.ground_truth, + } + reply = self.server.request({"type": "init", "route": route}) + if reply.get("type") != "ready": + raise RuntimeError(f"policy_server.py did not become ready: {reply!r}") + + if self.debug_carla_view_dir: + Path(self.debug_carla_view_dir).mkdir(parents=True, exist_ok=True) + out = str(Path(self.debug_carla_view_dir) / f"{self.route_tag}.mp4") + self.carla_view_writer = Mp4Writer(out, fps=round(1.0 / self.tick_dt)) + if self.record_infractions_dir: + Path(self.record_infractions_dir).mkdir(parents=True, exist_ok=True) + self.infraction_buffer = deque(maxlen=int(INFRACTION_CLIP_SECONDS / self.tick_dt)) + self.infraction_counter = 0 + self.last_infraction_location = self.vehicle.get_location() + if self.ground_truth: + blueprint = self.world.get_blueprint_library().find("sensor.other.collision") + self._collision_sensor = self.world.spawn_actor(blueprint, carla.Transform(), attach_to=self.vehicle) + self._collision_sensor.listen(self._on_collision) + print( + f"[puffer_agent_lb1] town={town} tick_dt={self.tick_dt:.3f} lights={len(self.lights)} " + f"stop_signs={len(route['stop_signs'])} target_points={len(route['sparse_plan'])} " + f"dense={len(route['dense_plan'])}", + flush=True, + ) + self.initialized = True + + def _on_collision(self, event): + self._collision_events.append(event.other_actor.type_id) + + # --- CARLA ground truth for the shadow env (read-only w.r.t. CARLA) ---- + + def _nearby_actors(self): + ego_loc = self.vehicle.get_location() + candidates = [] + for a in self.world.get_actors(): + if a.id == self.vehicle.id or not ("vehicle" in a.type_id or "walker.pedestrian" in a.type_id): + continue + loc = a.get_location() + # scenario_runner parks pending scenario actors hundreds of metres underground at the ego's waypoint + if abs(loc.z - ego_loc.z) > self.partner_max_abs_dz_m: + continue + candidates.append((loc.distance(ego_loc), a)) + candidates.sort(key=lambda item: item[0]) + return [a for _, a in candidates[: self.partner_slots]] + + def _snapshot(self): + ego_light = self.vehicle.get_traffic_light() + return { + "ego": actor_state_row(self.vehicle), + "partners": [partner_row(a) for a in self._nearby_actors()], + "light_states": [[lt.id, lt.get_state().name] for lt in self.lights], + "ego_light_id": NO_LIGHT_ID if ego_light is None else ego_light.id, + } + + def _ground_truth_observation(self): + """lb1/ground_truth.py's per-tick input: collisions since the last tick and the ego's lane situation.""" + collisions, self._collision_events = self._collision_events, [] + location = self.vehicle.get_location() + driving = self.cmap.get_waypoint(location, project_to_road=True, lane_type=carla.LaneType.Driving) + parking = self.cmap.get_waypoint(location, project_to_road=True, lane_type=carla.LaneType.Parking) + parking_distance = NO_PARKING_LANE if parking is None else location.distance(parking.transform.location) + return { + "collision": collisions, + "lane": [ + location.distance(driving.transform.location), + driving.lane_width, + parking_distance, + 0.0 if parking is None else parking.lane_width, + driving.transform.rotation.yaw, + 1.0 if driving.is_junction else 0.0, + ], + } + + def _teleport_carla_ego(self, motion): + """Place the CARLA ego at the shadow env's post-step pose (shadow_ego.StepResult.motion): the body chord on + the lane's height profile, lifted by the probed road-mesh offset (0.9.10 has no road-mesh raycast).""" + _, _, yaw0_deg, _, x, y, yaw_deg, speed, carla_z = motion + yaw_delta_deg = wrap_deg_180(yaw_deg - yaw0_deg) + wp = self.cmap.get_waypoint(carla.Location(x=x, y=y, z=carla_z)) + chord_z, pitch_deg = (carla_z, 0.0) if wp is None else lane_chord_pose(wp, self.body_chord_offsets_m) + self._update_mesh_offset() + z = chord_z + self.mesh_offset_m + REST_ORIGIN_ABOVE_ROAD_M + MESH_LIFT_MARGIN_M + pitch_rad, yaw_rad = math.radians(pitch_deg), math.radians(yaw_deg) + self._previous_body_z = z + self._previous_lane_center_z = carla_z if wp is None else wp.transform.location.z + self._previous_pitch_rad = pitch_rad + forward = ( + math.cos(yaw_rad) * math.cos(pitch_rad), + math.sin(yaw_rad) * math.cos(pitch_rad), + math.sin(pitch_rad), + ) + # Zero momentum before the teleport: CARLA's collision resolver reacts violently to a + # physics body carrying velocity into a new pose (carla issue #8076). + # One ordered batch: the server reorders async actor commands, and a late zero left the ego reading 0 m/s on 3 % of ticks + zero = carla.Vector3D(x=0.0, y=0.0, z=0.0) + CarlaDataProvider.get_client().apply_batch_sync( + [ + carla.command.ApplyTargetVelocity(self.vehicle.id, zero), + carla.command.ApplyTargetAngularVelocity(self.vehicle.id, zero), + carla.command.ApplyTransform( + self.vehicle.id, + carla.Transform( + carla.Location(x=x, y=y, z=z), carla.Rotation(pitch=pitch_deg, yaw=yaw_deg, roll=0.0) + ), + ), + carla.command.ApplyTargetVelocity( + self.vehicle.id, carla.Vector3D(x=speed * forward[0], y=speed * forward[1], z=speed * forward[2]) + ), + carla.command.ApplyTargetAngularVelocity( + self.vehicle.id, carla.Vector3D(x=0.0, y=0.0, z=yaw_delta_deg / self.tick_dt) + ), + ], + False, + ) + + def _update_mesh_offset(self): + """Road mesh height above the lane's height profile under the body, from the probe's newest scan at the previous pose.""" + distance_m = self.road_probe.fresh_distance_m() + if distance_m is None or self._previous_body_z is None: + return + # the ray follows the pitched body's down axis; the mount offset is along that axis too + vertical_clearance_m = (distance_m - ROAD_PROBE_MOUNT_Z_M) * math.cos(self._previous_pitch_rad) + measured_m = self._previous_body_z - vertical_clearance_m - self._previous_lane_center_z + step_m = measured_m - self.mesh_offset_m + if abs(step_m) > MESH_OFFSET_DEADBAND_M: + self.mesh_offset_m += max(-MESH_OFFSET_MAX_STEP_M, min(MESH_OFFSET_MAX_STEP_M, step_m)) + + def _maybe_save_infraction_clip(self, pd_flags, carla_flags): + """One mp4 of the rolling chase-cam buffer per flagged ego infraction, at most once per + INFRACTION_MIN_SEPARATION_M of ego travel (as ../leaderboard_agent.py).""" + fired = [f"pd_{name}" for name, value in (pd_flags or {}).items() if value > 0.0] + fired += [f"carla_{name}" for name, value in (carla_flags or {}).items() if value > 0.0] + if not fired or not self.infraction_buffer: + return + location = self.vehicle.get_location() + if location.distance(self.last_infraction_location) <= INFRACTION_MIN_SEPARATION_M: + return + out = str( + Path(self.record_infractions_dir) / f"{self.route_tag}_{'_'.join(fired)}_{self.infraction_counter:02d}.mp4" + ) + write_mp4(out, list(self.infraction_buffer), fps=round(1.0 / self.tick_dt)) + print(f"[puffer_agent_lb1] infraction {fired} -> {out}", flush=True) + self.infraction_counter += 1 + self.last_infraction_location = location + + def run_step(self, input_data, timestamp, sensors=None): + self.step += 1 + if not self.initialized: + self._init_on_first_step() + return carla.VehicleControl(steer=0.0, throttle=0.0, brake=1.0) + + if CARLA_VIEW_SENSOR_ID in input_data: + _, bgra = input_data[CARLA_VIEW_SENSOR_ID] # (H, W, 4) uint8, leaderboard's CallBack format + rgb = bgra[:, :, [2, 1, 0]] + if self.carla_view_writer is not None: + self.carla_view_writer.append_data(rgb) + if self.record_infractions_dir: + self.infraction_buffer.append(rgb.copy()) + + message = {"type": "tick", "snapshot": self._snapshot(), "ground_truth": None} + if self.ground_truth: + message["ground_truth"] = self._ground_truth_observation() + reply = self.server.request(message) + if reply.get("motion") is not None: + self._teleport_carla_ego(reply["motion"]) + if self.record_infractions_dir: + self._maybe_save_infraction_clip(reply.get("pd_flags"), reply.get("carla_flags")) + if reply.get("control") is None: + return carla.VehicleControl() + steer, throttle, brake = reply["control"] + return carla.VehicleControl(steer=steer, throttle=throttle, brake=brake) + + def destroy(self, results=None): + if not hasattr(self, "server"): + return + if self.initialized: + reply = self.server.request({"type": "finish"}) + if reply.get("type") != "done": + print(f"[puffer_agent_lb1] unexpected finish reply {reply!r}", flush=True) + self.server.close() + if self.carla_view_writer is not None: + self.carla_view_writer.close() + print(f"[puffer_agent_lb1] wrote CARLA chase-cam video for route {self.route_tag}", flush=True) + if self.road_probe is not None: + self.road_probe.destroy() + self.road_probe = None + if self._collision_sensor is not None: + self._collision_sensor.stop() + self._collision_sensor.destroy() + self._collision_sensor = None + self.initialized = False diff --git a/pufferlib/ocean/cosim/carla/lb1/policy_server.py b/pufferlib/ocean/cosim/carla/lb1/policy_server.py new file mode 100644 index 0000000000..6aec7b741d --- /dev/null +++ b/pufferlib/ocean/cosim/carla/lb1/policy_server.py @@ -0,0 +1,99 @@ +"""Shadow env + policy process behind lb1/leaderboard_agent.py (CARLA Leaderboard 1.0, CARLA 0.9.10). + +Started by the agent in the PufferDrive venv, one process per route, talking lb1_protocol over an +inherited socket: `init` builds the ShadowEgo for the route, every `tick` runs one policy step and +answers with the ego command (teleport pose or control values) plus both infraction sources, +`finish` flushes the per-route outputs. Options come from the COSIM_* environment variables +(shadow_ego.py), inherited from the evaluator. + +usage: policy_server.py --checkpoint <.pt or run dir> --socket-fd --route-tag +""" + +import argparse +import socket +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[5] +sys.path.insert(0, str(REPO_ROOT)) + +from pufferlib.ocean.cosim import carla_bridge as cb +from pufferlib.ocean.cosim.carla.lb1 import lb1_protocol +from pufferlib.ocean.cosim.carla.lb1.ground_truth import LeaderboardOneGroundTruth +from pufferlib.ocean.cosim.carla.shadow_ego import MAX_POLICY_STEPS, PARTNER_MAX_ABS_DZ_M, ShadowEgo + + +MAX_MESSAGES = MAX_POLICY_STEPS + 16 # ticks plus the handshake; the shadow env caps the route before this + + +def reply_for_step(result, carla_flags): + return { + "type": "control", + "motion": None if result.motion is None else [float(v) for v in result.motion], + "control": None if result.control is None else [float(v) for v in result.control], + "pd_flags": result.pd_flags, + "carla_flags": carla_flags, + } + + +def serve(sock, shadow): + lb1_protocol.send_message( + sock, + { + "type": "config", + "partner_slots": shadow.partner_slots, + "partner_max_abs_dz_m": PARTNER_MAX_ABS_DZ_M, + "light_probe_steps_m": list(cb.LIGHT_PROBE_STEPS_M), + "dynamics_source": shadow.dynamics_source, + }, + ) + ground_truth = None + for _ in range(MAX_MESSAGES): + message = lb1_protocol.recv_message(sock) + if message is None: + print("[policy_server] client closed the connection", flush=True) + return + kind = message["type"] + if kind == "init": + route = message["route"] + shadow.init_route(route) + if route.get("ground_truth", False): + ground_truth = LeaderboardOneGroundTruth( + route["lights"], route["stop_signs"], route["ego"]["extent"][0] + ) + lb1_protocol.send_message(sock, {"type": "ready", "num_traffic": shadow.num_traffic}) + elif kind == "tick": + snapshot = message["snapshot"] + result = shadow.step(snapshot) + carla_flags = None + if ground_truth is not None and message.get("ground_truth") is not None: + carla_flags = ground_truth.update(snapshot, message["ground_truth"]) + shadow.telemetry(result, carla_flags) + lb1_protocol.send_message(sock, reply_for_step(result, carla_flags)) + elif kind == "finish": + shadow.finish() + lb1_protocol.send_message(sock, {"type": "done"}) + return + else: + raise ValueError(f"unknown message type {kind!r}") + raise RuntimeError(f"more than MAX_MESSAGES={MAX_MESSAGES} messages on one route") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--checkpoint", required=True, help="PufferDrive checkpoint .pt or run dir (config.yaml beside)" + ) + parser.add_argument("--socket-fd", type=int, required=True, help="inherited socket to the leaderboard agent") + parser.add_argument("--route-tag", required=True, help="per-route file stem for every output") + args = parser.parse_args() + sock = socket.socket(fileno=args.socket_fd) + try: + serve(sock, ShadowEgo(args.checkpoint, args.route_tag)) + finally: + sock.close() + + +if __name__ == "__main__": + main() diff --git a/pufferlib/ocean/cosim/carla/lb1/run_evaluator.py b/pufferlib/ocean/cosim/carla/lb1/run_evaluator.py new file mode 100644 index 0000000000..ec49b8ac12 --- /dev/null +++ b/pufferlib/ocean/cosim/carla/lb1/run_evaluator.py @@ -0,0 +1,90 @@ +"""Run carla_garage's leaderboard_evaluator_local.py (Leaderboard 1.0, CARLA 0.9.10) for one model directory, +with all outputs under its eval folder; the Leaderboard 1.0 twin of ../run_evaluator.py. Runs in the `garage` +conda env (Python 3.7), so nothing from pufferlib is imported here. + +usage: run_evaluator.py --model-dir [evaluator args...] + + holds final_model.pt + config.yaml (blank -> DEFAULT_MODEL_DIR). The launcher adds --agent-config +/final_model.pt, --checkpoint /eval/carla_leaderboard_v1/result.json (the evaluator opens the +result file before any agent code runs) and --scenarios eval_scenarios.json, points COSIM_TELEMETRY / +COSIM_WORLD_LOG / COSIM_OBS_HTML / COSIM_DEBUG_CARLA_VIEW into the eval folder unless already set, and +exports what the garage evaluator reads from the environment: ROUTES (the --routes file, it names the +route records after it) and BENCHMARK=longest6 (all spawn points of background traffic, stop-sign +penalty 1.0). The evaluator is located through GARAGE_WORK_DIR. +""" + +import os +import runpy +import sys +from pathlib import Path + + +DEFAULT_MODEL_DIR = "experiments/k_scaled_0040_1000" +EVAL_SUBDIR = "eval/carla_leaderboard_v1" +DEFAULT_BENCHMARK = "longest6" +DEFAULT_SCENARIOS = "leaderboard/data/scenarios/eval_scenarios.json" +OUTPUT_ENV_DIRS = { + "COSIM_TELEMETRY": "telemetry", + "COSIM_WORLD_LOG": "world_log", + "COSIM_OBS_HTML": "obs_html", + "COSIM_DEBUG_CARLA_VIEW": "carla_view", +} + + +def option_value(argv, name): + """Value of `--name value` or `--name=value` in argv, else None.""" + for k, arg in enumerate(argv): + if arg == name and k + 1 < len(argv): + return argv[k + 1] + if arg.startswith(name + "="): + return arg.split("=", 1)[1] + return None + + +def main(): + argv = sys.argv[1:] + model_dir = "" + if "--model-dir" in argv: + k = argv.index("--model-dir") + has_value = k + 1 < len(argv) and not argv[k + 1].startswith("--") + model_dir = argv[k + 1] if has_value else "" + del argv[k : k + 2 if has_value else k + 1] + if not model_dir: + model_dir = DEFAULT_MODEL_DIR + print(f"[run_evaluator] --model-dir blank, using {DEFAULT_MODEL_DIR}") + model_dir = Path(model_dir) + if not model_dir.is_absolute(): + model_dir = Path.cwd() / model_dir + checkpoint = model_dir / "final_model.pt" + if not checkpoint.is_file(): + raise SystemExit(f"run_evaluator.py: no checkpoint at {checkpoint}") + garage_work_dir = os.environ.get("GARAGE_WORK_DIR") + if not garage_work_dir: + raise SystemExit("run_evaluator.py: GARAGE_WORK_DIR is not set (the carla_garage checkout)") + evaluator = Path(garage_work_dir) / "leaderboard" / "leaderboard" / "leaderboard_evaluator_local.py" + if not evaluator.is_file(): + raise SystemExit(f"run_evaluator.py: evaluator not found at {evaluator}") + routes = option_value(argv, "--routes") + if routes is None: + raise SystemExit("run_evaluator.py: --routes is required") + eval_dir = model_dir / EVAL_SUBDIR + eval_dir.mkdir(parents=True, exist_ok=True) + if option_value(argv, "--agent-config") is None: + argv += ["--agent-config", str(checkpoint)] + if option_value(argv, "--checkpoint") is None: + argv += ["--checkpoint", str(eval_dir / "result.json")] + if option_value(argv, "--scenarios") is None: + argv += ["--scenarios", str(Path(garage_work_dir) / DEFAULT_SCENARIOS)] + for env_name, subdir in OUTPUT_ENV_DIRS.items(): + if env_name not in os.environ: + os.environ[env_name] = str(eval_dir / subdir) + os.environ["ROUTES"] = routes + os.environ.setdefault("BENCHMARK", DEFAULT_BENCHMARK) + os.environ.setdefault("SCENARIO_RUNNER_ROOT", str(Path(garage_work_dir) / "scenario_runner")) + print(f"[run_evaluator] model {model_dir} -> outputs {eval_dir} (BENCHMARK={os.environ['BENCHMARK']})") + sys.argv = [str(evaluator)] + argv + runpy.run_path(str(evaluator), run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/pufferlib/ocean/cosim/carla/lb1/run_leaderboard.sh b/pufferlib/ocean/cosim/carla/lb1/run_leaderboard.sh new file mode 100644 index 0000000000..c93d34f833 --- /dev/null +++ b/pufferlib/ocean/cosim/carla/lb1/run_leaderboard.sh @@ -0,0 +1,117 @@ +#!/bin/bash +# Run lb1/leaderboard_agent.py under carla_garage's leaderboard_evaluator_local.py (CARLA Leaderboard 1.0, +# CARLA 0.9.10.1): start a CARLA server, evaluate one routes xml, stop the server. One machine, one GPU; +# scripts/kesai/14_carla_longest6_v1.sh is the cluster version over all 36 longest6 routes. +# +# Usage: +# CKPT=experiments//final_model.pt bash pufferlib/ocean/cosim/carla/lb1/run_leaderboard.sh +# +# Required env: +# CKPT PufferDrive checkpoint .pt (config.yaml beside it or one level up) or an experiment dir +# Optional env: +# ROUTES routes xml. Default: longest6_split/longest_weathers_0.xml (one Town01 route). +# SCENARIOS scenario json. Default: eval_scenarios.json (the benchmark); no_scenarios.json = empty. +# REPETITIONS default 1 +# OUT result json. Default: runs/cosim_leaderboard_v1_/result.json +# CARLA_PORT default 2000 (distinct per concurrent server); TM_PORT default CARLA_PORT+6000 +# CARLA_WINDOW 1: rendering window on $DISPLAY; default 0: offscreen OpenGL (DISPLAY unset, -opengl) +# CARLA_ROOT CARLA 0.9.10.1 install. Default: ~/ordnung/internal/CARLA_0.9.10 +# GARAGE_WORK_DIR carla_garage checkout. Default: ~/ordnung/internal/carla_garage_1 +# GARAGE_PY python of the `garage` conda env. Default: ~/miniconda3/envs/garage/bin/python +# COSIM_SERVER_PYTHON python of the PufferDrive venv for policy_server.py. Default: $PD/.venv/bin/python +# PD PufferDrive checkout. Default: the one holding this script. +# LIBTIFF5_COMPAT 1 (default): the 0.9.10 egg links libtiff.so.5. When `import carla` fails, the garage env's +# own lib dir is tried first (`conda install -n garage "libtiff<4.5"` puts one there), then +# the system libtiff.so.6 is linked under that name inside $OUT's folder (0 = never link). +# COSIM_* forwarded to the agent / policy server as-is (see lb1/leaderboard_agent.py, shadow_ego.py); +# COSIM_DEBUG_CARLA_VIEW defaults to "$(dirname "$OUT")/carla_view" (empty disables). +set -u + +: "${CKPT:?set CKPT=/path/to/final_model.pt}" +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +PD=${PD:-$(cd "$SCRIPT_DIR/../../../../.." && pwd)} +CARLA_ROOT=${CARLA_ROOT:-$HOME/ordnung/internal/CARLA_0.9.10} +GARAGE_WORK_DIR=${GARAGE_WORK_DIR:-$HOME/ordnung/internal/carla_garage_1} +GARAGE_PY=${GARAGE_PY:-$HOME/miniconda3/envs/garage/bin/python} +export COSIM_SERVER_PYTHON=${COSIM_SERVER_PYTHON:-$PD/.venv/bin/python} +CARLA_PORT=${CARLA_PORT:-2000} +TM_PORT=${TM_PORT:-$((CARLA_PORT + 6000))} +CARLA_WINDOW=${CARLA_WINDOW:-0} +ROUTES=${ROUTES:-$GARAGE_WORK_DIR/leaderboard/data/longest6_split/longest_weathers_0.xml} +SCENARIOS=${SCENARIOS:-$GARAGE_WORK_DIR/leaderboard/data/scenarios/eval_scenarios.json} +REPETITIONS=${REPETITIONS:-1} +TIMESTAMP=$(date +"%Y%m%d_%H%M%S") +OUT=${OUT:-$PD/runs/cosim_leaderboard_v1_${TIMESTAMP}/result.json} +OUT_DIR=$(dirname "$OUT") +mkdir -p "$OUT_DIR" +export COSIM_DEBUG_CARLA_VIEW=${COSIM_DEBUG_CARLA_VIEW-"$OUT_DIR/carla_view"} +EVALUATOR=$GARAGE_WORK_DIR/leaderboard/leaderboard/leaderboard_evaluator_local.py +CARLA_EGG=$CARLA_ROOT/PythonAPI/carla/dist/carla-0.9.10-py3.7-linux-x86_64.egg +for path in "$CKPT" "$GARAGE_PY" "$COSIM_SERVER_PYTHON" "$CARLA_ROOT/CarlaUE4.sh" "$CARLA_EGG" "$EVALUATOR" "$ROUTES" "$SCENARIOS"; do + [ -e "$path" ] || { echo "missing $path"; exit 1; } +done + +carla_imports() { PYTHONPATH="$CARLA_EGG" "$GARAGE_PY" -c "import carla" >/dev/null 2>&1; } +# the 0.9.10 egg links libtiff.so.5, which newer distros no longer ship (Ubuntu 24.04: libtiff.so.6 only) +if ! carla_imports; then + GARAGE_LIB=$(cd "$(dirname "$GARAGE_PY")/../lib" 2>/dev/null && pwd) + if [ -n "$GARAGE_LIB" ] && [ -e "$GARAGE_LIB/libtiff.so.5" ]; then + export LD_LIBRARY_PATH="$GARAGE_LIB${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + echo "[run_leaderboard_v1] libtiff.so.5 from the garage env: $GARAGE_LIB" + fi +fi +if ! carla_imports && [ "${LIBTIFF5_COMPAT:-1}" = "1" ]; then + SYSTEM_LIBTIFF=$( (ldconfig -p 2>/dev/null || /sbin/ldconfig -p 2>/dev/null) | awk '/libtiff.so.6 /{print $NF; exit}') + if [ -n "$SYSTEM_LIBTIFF" ]; then + mkdir -p "$OUT_DIR/compat_lib" && ln -sf "$SYSTEM_LIBTIFF" "$OUT_DIR/compat_lib/libtiff.so.5" + export LD_LIBRARY_PATH="$OUT_DIR/compat_lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + echo "[run_leaderboard_v1] libtiff.so.5 missing: using $SYSTEM_LIBTIFF via $OUT_DIR/compat_lib" + fi +fi +carla_imports || { echo "ERROR: $GARAGE_PY cannot import the CARLA egg:"; PYTHONPATH="$CARLA_EGG" "$GARAGE_PY" -c "import carla"; exit 1; } + +export CARLA_ROOT GARAGE_WORK_DIR +export SCENARIO_RUNNER_ROOT=$GARAGE_WORK_DIR/scenario_runner +export LEADERBOARD_ROOT=$GARAGE_WORK_DIR/leaderboard +export PYTHONPATH="$CARLA_ROOT/PythonAPI/carla:$CARLA_EGG:$SCENARIO_RUNNER_ROOT:$LEADERBOARD_ROOT:${PYTHONPATH:-}" +# the garage evaluator names the route records after $ROUTES; BENCHMARK=longest6 fills every spawn point +export ROUTES BENCHMARK=${BENCHMARK:-longest6} DATAGEN=0 + +if [ "$CARLA_WINDOW" = "1" ]; then + "$CARLA_ROOT/CarlaUE4.sh" -opengl -nosound -carla-rpc-port="$CARLA_PORT" -carla-streaming-port=$((CARLA_PORT + 1)) \ + > "$OUT_DIR/carla_server.log" 2>&1 & +else + env -u WAYLAND_DISPLAY DISPLAY= SDL_VIDEODRIVER=offscreen SDL_HINT_CUDA_DEVICE=0 "$CARLA_ROOT/CarlaUE4.sh" -opengl -nosound -carla-rpc-port="$CARLA_PORT" -carla-streaming-port=$((CARLA_PORT + 1)) \ + > "$OUT_DIR/carla_server.log" 2>&1 & +fi +SERVER_PID=$! +up=1 +for _ in $(seq 1 60); do + if ! kill -0 $SERVER_PID 2>/dev/null; then + echo "CARLA server process died"; break + fi + "$GARAGE_PY" -c " +import carla, sys +try: + c = carla.Client('localhost', $CARLA_PORT); c.set_timeout(10.0) + c.get_world().get_map() +except Exception: + sys.exit(1)" 2>/dev/null && up=0 && break + sleep 5 +done +if [ $up -ne 0 ]; then + echo "FAIL: CARLA server not ready on port $CARLA_PORT (see $OUT_DIR/carla_server.log)" + pkill -9 -P $SERVER_PID 2>/dev/null; kill -9 $SERVER_PID 2>/dev/null + exit 1 +fi +echo "[run_leaderboard_v1] CARLA 0.9.10 server ready (pid $SERVER_PID, port $CARLA_PORT)" + +echo "[run_leaderboard_v1] routes=$ROUTES scenarios=$SCENARIOS ckpt=$CKPT out=$OUT" +"$GARAGE_PY" -u "$EVALUATOR" --routes "$ROUTES" --scenarios "$SCENARIOS" --repetitions "$REPETITIONS" \ + --agent "$SCRIPT_DIR/leaderboard_agent.py" --agent-config "$CKPT" --checkpoint "$OUT" --track MAP \ + --port "$CARLA_PORT" --trafficManagerPort "$TM_PORT" --timeout 600 +RC=$? +# children first: killing the wrapper shell first reparents the UE4 binary and leaves it running +pkill -9 -P $SERVER_PID 2>/dev/null; kill -9 $SERVER_PID 2>/dev/null +echo "[run_leaderboard_v1] evaluator exit $RC; result -> $OUT" +exit $RC diff --git a/pufferlib/ocean/cosim/carla/leaderboard_agent.py b/pufferlib/ocean/cosim/carla/leaderboard_agent.py index fda865955c..75905a15ea 100644 --- a/pufferlib/ocean/cosim/carla/leaderboard_agent.py +++ b/pufferlib/ocean/cosim/carla/leaderboard_agent.py @@ -1,4 +1,4 @@ -"""PufferDrive policy as a CARLA leaderboard agent (CaRL original_leaderboard). +"""PufferDrive policy as a CARLA leaderboard agent (CaRL original_leaderboard, CARLA 0.9.15). Usage (CaRL env vars/paths as in CaRL/CARLA/README.md, plus PufferDrive repo root on PYTHONPATH for `pufferlib` and `data_utils`): @@ -9,66 +9,27 @@ --agent-config /path/to/experiments/puffer_drive_xxx/models/model_xxx.pt \ --checkpoint /path/to/results/result.json --track MAP -Everything about the shadow env (obs layout, dynamics, dt, goal count, agent -pool size, ...) comes from the checkpoint's sibling config.yaml, with the -clean-eval profile applied on top (cosim/arch.py CLEAN_EVAL_OVERRIDES); only -the structural co-sim keys (map, pool wiring) are set here. +This file is the CARLA side only: it reads ground truth (ego, nearby actors, lights, stop signs, the +route), hands it to shadow_ego.ShadowEgo (the shadow Drive env + policy, built from the checkpoint's +config.yaml) and moves the ego the way the result says. The Leaderboard 1.0 integration (lb1/) runs +the same core behind a socket. -Environment variables: +Environment variables (the shadow env's own, COSIM_DEVICE ... COSIM_WORLD_LOG, are listed in +shadow_ego.py): SCENARIO_RUNNER_ROOT REQUIRED: path to the scenario_runner checkout (route scenarios silently fail to load without it) CARL_WORK_DIR REQUIRED if COSIM_TELEMETRY or COSIM_RECORD_INFRACTIONS is set: path to the CaRL/CARLA checkout, to import its reward.criteria collision/offroad detectors - COSIM_DEVICE=cpu torch device for the policy - COSIM_DYNAMICS_SOURCE=pufferdrive - "pufferdrive" (default): PufferDrive's own - dynamics (the ones the policy trained on) move - the ego; the CARLA actor is teleported to match - every policy step. - "carla": TrackingController converts the shadow - env's target (speed, yaw) into throttle/brake/ - steer and CARLA's own vehicle physics moves the - ego (subject to dynamics-mismatch tracking lag). - COSIM_MAX_SPEED_MPS=20 ego speed cap; sets the shadow env's C_vel = - cap / base_max_speed_mps. Default: the - checkpoint's base_max_speed_mps (C_vel = 1). - Must lie in the trained range - base * [1/conditioning_speed_scale, scale] - (0036: 13.3-30 m/s, 0038+: 10-40 m/s). - COSIM_ZERO_PARTNER_STOPPED_TIME=0 - 1: every partner's seconds_stopped observation - feature is held at 0 (the ego's own counter is - untouched). Ablates the stopped-time cue behind - parked-car / red-light-queue avoidance. - COSIM_PEDESTRIAN_MIN_SIZE_M=0.0 - walker partner boxes below this on either axis - grow to it (nuPlan's pedestrian_min_size_m; the - training spawn floor is 0.8 m, CARLA walkers - are 0.4-0.5 m) COSIM_DEBUG_CARLA_VIEW=/dir write a CARLA chase-camera mp4 per route (native tick rate, streamed to disk frame-by-frame) COSIM_RECORD_INFRACTIONS=/dir write a short chase-cam clip (last ~5 s) per ego infraction (collision/offroad/red-light/stop-sign), from either the shadow pufferdrive env's own - model (_ego_infractions) or real CARLA ground - truth (_carla_infractions) - COSIM_TELEMETRY=/dir write a per-policy-step CSV per route, with - both infraction sources above as columns - COSIM_OBS_HTML=/dir write an interactive pufferlib.viz replay per - route (the exact obs + policy outputs the ego saw) - COSIM_OBS_HTML_RENDER=1 0: save only the compact .replay.zlib per route; - scripts/eval/render_carla_obs_html.py renders the - pages of a whole run into one gallery afterwards - COSIM_DUMP_OBS=/dir write the raw ego obs + action per policy step - per route (.npz: step, obs, action) - COSIM_WORLD_LOG=/dir write the bin-frame world state per policy step - per route (ego, partners, light states, route); - input of scripts/eval/analyze_carla_cosim.py + model (pd_*) or real CARLA ground truth (carla_*) """ -import json import math import os import re @@ -82,11 +43,9 @@ from leaderboard.autoagents import autonomous_agent from srunner.scenariomanager.carla_data_provider import CarlaDataProvider -from pufferlib.ocean.drive.drive import Drive -from pufferlib.ocean.cosim.goals import RouteGoalWindow, route_goals_from_target_points from pufferlib.ocean.cosim import carla_bridge as cb -from pufferlib.ocean.cosim.arch import checkpoint_config_path, shadow_env_kwargs -from pufferlib.ocean.cosim.carla.controller import TrackingController, read_vehicle_geometry +from pufferlib.ocean.cosim.carla.controller import vehicle_geometry_from_wheels, wheel_rows +from pufferlib.ocean.cosim.carla.shadow_ego import PARTNER_MAX_ABS_DZ_M, ShadowEgo # Rolling chase-cam window kept for COSIM_RECORD_INFRACTIONS clips, and the # minimum ego travel between two logged infractions (suppresses re-triggering @@ -94,18 +53,10 @@ INFRACTION_MIN_SEPARATION_M = 10.0 RED_LIGHT_CHECK_DISTANCE_M = 30.0 # matches CaRL RunRedLight's own distance_light default -FAR_AWAY = 1.0e6 # park unused shadow-env agent slots out of observation range -PARTNER_MAX_ABS_DZ_M = 20.0 # CARLA actors farther above/below the ego are hidden scenario props, not traffic -DEFAULT_BASE_MAX_SPEED_MPS = 20.0 # Drive() defaults, for checkpoints whose config predates the keys -DEFAULT_CONDITIONING_SPEED_SCALE = 1.5 ROAD_RAY_HALF_SPAN_M = ( 5.0 # vertical ray around the waypoint z; the ego roof and underpass roads are filtered by label/nearest ) ROAD_RAY_LABELS = (carla.CityObjectLabel.Roads, carla.CityObjectLabel.RoadLines, carla.CityObjectLabel.Bridge) -MAX_POLICY_STEPS = 100_000 # shadow-env episode cap (~1.4 h at the 0.05 s CARLA tick); sizes the light state buffers -SCENARIO_LENGTH_MARGIN_STEPS = 2 -# Shadow-env metrics_array indices (datatypes.h): collision/offroad/red-light flags. -EGO_INFRACTION_METRICS = {"collision": 0, "offroad": 1, "red_light": 2, "stop_sign": 3} CARLA_VIEW_SENSOR_ID = "puffer_chase_cam" CARLA_VIEW_WIDTH, CARLA_VIEW_HEIGHT, CARLA_VIEW_FOV = 960, 540, 90 @@ -119,35 +70,6 @@ def get_entry_point(): return "PufferAgent" -def clean_policy_state_dict(state_dict): - """Strip torch.compile / DDP prefixes (wandb, neptune, the _C kernel, ...) - from the leaderboard's evaluation environment.""" - - def clean(key): - while key.startswith(("module.", "_orig_mod.")): - key = key.split(".", 1)[1] - return key - - return {clean(k): v for k, v in state_dict.items()} - - -def resolve_checkpoint(path_to_conf_file): - """(checkpoint_path, config_dict) from a .pt file or an experiment dir.""" - import yaml - - p = Path(path_to_conf_file).resolve() - if p.is_file(): - ckpt, cfg_path = p, checkpoint_config_path(p) - else: - models = sorted((p / "models").glob("*.pt")) or sorted(p.glob("*.pt")) - if not models: - raise FileNotFoundError(f"no .pt checkpoint under {p}") - ckpt, cfg_path = models[-1], p / "config.yaml" - with open(cfg_path) as f: - cfg = yaml.safe_load(f) - return str(ckpt), cfg - - def road_mesh_z(world, x, y, reference_z): """Height of the road mesh at (x, y) from a vertical ray, the road hit nearest reference_z; None when none is hit.""" top = carla.Location(x=x, y=y, z=reference_z + ROAD_RAY_HALF_SPAN_M) @@ -171,29 +93,36 @@ def road_aligned_attitude(road_up, yaw_deg): return pitch_deg, roll_deg, forward -def cosim_max_speed_mps(env_cfg): - """COSIM_MAX_SPEED_MPS as a float inside the checkpoint's trained C_vel range, or None when unset.""" - raw = os.environ.get("COSIM_MAX_SPEED_MPS") - if raw is None: - return None - max_speed_mps = float(raw) - base_max_speed_mps = float(env_cfg.get("base_max_speed_mps") or DEFAULT_BASE_MAX_SPEED_MPS) - scale = float(env_cfg.get("conditioning_speed_scale") or DEFAULT_CONDITIONING_SPEED_SCALE) - low_mps, high_mps = base_max_speed_mps / scale, base_max_speed_mps * scale - if not (low_mps <= max_speed_mps <= high_mps): - raise ValueError( - f"COSIM_MAX_SPEED_MPS={max_speed_mps} is outside the trained C_vel range " - f"[{low_mps:.1f}, {high_mps:.1f}] m/s (base_max_speed_mps={base_max_speed_mps}, " - f"conditioning_speed_scale={scale})" - ) - return max_speed_mps - - def plan_xyz(plan): """[(carla.Transform, RoadOption)] -> (N, 3) CARLA-frame positions.""" return np.array([[t.location.x, t.location.y, t.location.z] for t, _ in plan], dtype=np.float64).reshape(-1, 3) +def actor_state_row(actor): + """shadow_ego ACTOR_STATE of a live CARLA actor (CARLA frame).""" + tf = actor.get_transform() + v = actor.get_velocity() + av = actor.get_angular_velocity() # deg/s, world frame (CARLA's rotation convention) + acc = actor.get_acceleration() # m/s^2, world frame + return [tf.location.x, tf.location.y, tf.location.z, tf.rotation.yaw, v.x, v.y, v.z, av.z, acc.x, acc.y] + + +def partner_row(actor): + """shadow_ego PARTNER_ROW of a live CARLA vehicle or walker.""" + box = actor.bounding_box + is_walker = 1.0 if "walker" in actor.type_id else 0.0 + is_two_wheeler = 1.0 if actor.attributes.get("number_of_wheels") == "2" else 0.0 + return actor_state_row(actor) + [ + box.location.z, + box.extent.x, + box.extent.y, + box.extent.z, + is_walker, + is_two_wheeler, + actor.id, + ] + + class PufferAgent(autonomous_agent.AutonomousAgent): def setup(self, path_to_conf_file, route_index=None): if "SCENARIO_RUNNER_ROOT" not in os.environ: @@ -206,64 +135,16 @@ def setup(self, path_to_conf_file, route_index=None): # CaRL's `route_index` is route_date_string = Path(ROUTES).stem, fixed # once per evaluator process unless --collect-dataset is passed. self.video_tag = f"{self.route_index}_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}" - self.checkpoint, self.cfg = resolve_checkpoint(path_to_conf_file) - - self.device = os.environ.get("COSIM_DEVICE", "cpu") - self.dynamics_source = os.environ.get("COSIM_DYNAMICS_SOURCE", "pufferdrive") - if self.dynamics_source not in ("carla", "pufferdrive"): - raise ValueError(f"COSIM_DYNAMICS_SOURCE must be 'carla' or 'pufferdrive', got {self.dynamics_source!r}") - env_cfg = self.cfg["env"] - self.max_speed_mps = cosim_max_speed_mps(env_cfg) - zero_stopped_raw = os.environ.get("COSIM_ZERO_PARTNER_STOPPED_TIME", "0") - if zero_stopped_raw not in ("0", "1"): - raise ValueError(f"COSIM_ZERO_PARTNER_STOPPED_TIME must be '0' or '1', got {zero_stopped_raw!r}") - self.zero_partner_stopped_time = zero_stopped_raw == "1" - self.pedestrian_min_size_m = float(os.environ.get("COSIM_PEDESTRIAN_MIN_SIZE_M", "0.0")) - if not (self.pedestrian_min_size_m >= 0.0): - raise ValueError(f"COSIM_PEDESTRIAN_MIN_SIZE_M must be >= 0, got {self.pedestrian_min_size_m}") - # Shadow agent pool == the training per-env cap - self.num_agents = int(env_cfg["max_agents_per_env"]) + self.shadow = ShadowEgo(path_to_conf_file, self.video_tag) self.debug_carla_view_dir = os.environ.get("COSIM_DEBUG_CARLA_VIEW", None) self.record_infractions_dir = os.environ.get("COSIM_RECORD_INFRACTIONS", None) - self.telemetry_dir = os.environ.get("COSIM_TELEMETRY", None) - self.telemetry_file = None - self.obs_html_dir = os.environ.get("COSIM_OBS_HTML", None) - self.obs_dump_dir = os.environ.get("COSIM_DUMP_OBS", None) # raw ego obs per policy step (.npz) - self._obs_dump = [] - self.world_log_dir = os.environ.get("COSIM_WORLD_LOG", None) # bin-frame world state per policy step (.npz) - self._world_log = {"ego": [], "partners": [], "lights": []} - self._last_partners = None - self._partner_stopped_s = {} # CARLA actor id -> seconds at standstill - from pufferlib.ocean.drive import binding - - self._partner_stopped_speed_threshold = float(binding.AGENT_STOPPED_SPEED_THRESHOLD) - self.obs_html_max_steps = int(os.environ.get("COSIM_OBS_HTML_MAX_STEPS", "12000")) - obs_html_render_raw = os.environ.get("COSIM_OBS_HTML_RENDER", "1") - if obs_html_render_raw not in ("0", "1"): - raise ValueError(f"COSIM_OBS_HTML_RENDER must be '0' or '1', got {obs_html_render_raw!r}") - self.obs_html_render = obs_html_render_raw == "1" - self._obs_html = None - # Per-policy-step CSV of the ego's current goal vs the lane - # find_goal_lane snapped it to (goal_lane_idx from env.get_state()): - # lane direction, lane-to-goal distance, dot vs the route direction - # used at snap time -- diagnoses whether the goal feature is guiding - # toward a real, correctly-directed lane. Off unless set. - self.debug_goal_lane_dir = os.environ.get("COSIM_DEBUG_GOAL_LANE", None) - self._goal_lane_debug_file = None - self._road_data = None self.step = -1 self.initialized = False - self.policy = None self.carla_view_writer = None - self.goal_window = None - self.target = (0.0, 0.0) # (target_speed, target_yaw_deg), held between policy steps - # (x, y, yaw_deg, speed) of the shadow ego before and after the current policy step, CARLA frame - self._motion = None self._carla_collision = None # set by _init_carla_infraction_detectors when needed self._carla_stop_sign = None - self._stop_sign_lines = np.zeros((0, 6), np.float32) # CARLA trigger volumes as bin-frame stop lines def sensors(self): if not self.debug_carla_view_dir and not self.record_infractions_dir: @@ -279,67 +160,6 @@ def sensors(self): } ] - def _load_policy_and_env(self, town_bin): - arch = shadow_env_kwargs( - self.cfg, - overrides=dict( - map_dir=town_bin, - num_maps=1, - # One policy agent (the ego); every other slot is a static partner streamed from CARLA. - num_agents=1, - min_agents_per_env=1, - cosim_partner_slots=self.num_agents - 1, - goal_source="external", - dt=self.dt, - scenario_length=MAX_POLICY_STEPS + SCENARIO_LENGTH_MARGIN_STEPS, - resample_frequency=0, - termination_mode=0, - # Enforcement off (detection flags still fire): in one endless - # episode a "stop" latch is permanent and fires only AFTER - # CARLA scored the infraction (route 5 froze 320 s -> DNF). - collision_behavior="ignore", - offroad_behavior="ignore", - traffic_light_behavior="ignore", - **({} if self.max_speed_mps is None else {"max_speed_mps": self.max_speed_mps}), - ), - ) - self.dynamics_model = arch.get("dynamics_model", "classic") - # jerk dynamics build accel across steps while sync() re-seeds accel - self.speed_intent_extension_s = 2.0 * self.dt if self.dynamics_model == "jerk" else 0.0 - - self.env = Drive(**arch) - self.env.reset() - - import torch - import pufferlib.ocean.torch as drive_torch - - policy_state_dict = clean_policy_state_dict( - torch.load(self.checkpoint, map_location=self.device, weights_only=False) - ) - env_for_policy = self.env - self._boundary_feat_ckpt = None - bw = policy_state_dict.get("actor_backbone.boundary_encoder.0.weight") - if bw is not None and bw.shape[1] < env_for_policy.boundary_features: - # Older checkpoints emit fewer boundary features than this drive.h - # (e.g. 7 vs 9 -- two zero-padded GPS columns were added later). - import copy - - self._boundary_feat_ckpt = int(bw.shape[1]) - env_for_policy = copy.copy(self.env) - env_for_policy.boundary_features = self._boundary_feat_ckpt - print( - f"[puffer_agent] checkpoint uses {self._boundary_feat_ckpt} boundary features " - f"(env emits {self.env.boundary_features}); stripping zero-padded columns" - ) - - policy_cls = getattr(drive_torch, self.cfg.get("policy_name", "Drive")) - # pre-3.0 checkpoints keep action_type only in the env section - self.cfg["policy"].setdefault("action_type", self.cfg["env"]["action_type"]) - self.policy = policy_cls(env_for_policy, **self.cfg["policy"]).to(self.device) - self.policy.load_state_dict(policy_state_dict) - self.policy.eval() - print(f"[puffer_agent] loaded policy from {self.checkpoint}") - def _init_on_first_step(self): """Deferred init (the ego and world only exist once the route runs) — same pattern as CaRL's eval_agent.agent_init. Everything here is @@ -349,89 +169,31 @@ def _init_on_first_step(self): self.cmap = self.world.get_map() town = CarlaDataProvider.get_map().name.split("/")[-1] self.tick_dt = float(self.world.get_settings().fixed_delta_seconds) # 0.05 @ 20 Hz - self.dt = self.tick_dt # the policy runs every CARLA tick, like the leaderboard's sensor agents - - self.town_bin = cb.bin_path_for_town(town) - self._load_policy_and_env(self.town_bin) - - self.transform = cb.CarlaTransform(town, offset=cb.town_offset(self.town_bin)) - offset, z_offset, residual_before, residual_after = cb.calibrate_town_offset( - self.cmap, self.transform, self.town_bin - ) - print( - f"[puffer_agent] bin offset calibrated against CARLA lanes: shift " - f"({offset[0] - self.transform.tx:+.2f}, {offset[1] - self.transform.ty:+.2f}) m, z {z_offset:+.2f} m, " - f"median lane residual {residual_before:.2f} -> {residual_after:.2f} m" - ) - self.transform = cb.CarlaTransform(town, offset=offset, z_offset=z_offset) - if self.dynamics_source == "pufferdrive": - self._sync_ego_from_carla(zero_velocity=True) - ego_loc = self.vehicle.get_location() - target_points = plan_xyz(self._global_plan_world_coord) - dense_points = plan_xyz(self.dense_global_plan_world_coord) - target_points[:, 2] = self.transform.z_to_bin(target_points[:, 2]) - dense_points[:, 2] = self.transform.z_to_bin(dense_points[:, 2]) - self.route_goals = route_goals_from_target_points( - target_points, - dense_points, - self.transform.loc_to_bin, - self.transform.loc_to_bin(ego_loc.x, ego_loc.y), - self.env.goal_radius, - ) - # the leaderboard's target points, next num_goals of them at all times (fewer only at the route end) - self.goal_window = RouteGoalWindow(self.env, self.route_goals, sliding=True) - - import data_utils.mirror_map_bin as mbin - - bin_data = mbin.read_bin(Path(self.town_bin)) - bin_traffic = bin_data["traffic"] - self.stop_line_centers = np.array( - [ - [0.5 * (t["stop_line"][0] + t["stop_line"][3]), 0.5 * (t["stop_line"][1] + t["stop_line"][4])] - for t in bin_traffic - ] - ).reshape(-1, 2) - - if self.debug_goal_lane_dir: - # goal_lane_idx comes from road_elements[] in the C file; bin_data["roads"] - # is the same bin read in the same order, so it's usable as-is. - self._road_data = bin_data["roads"] - Path(self.debug_goal_lane_dir).mkdir(parents=True, exist_ok=True) - self._goal_lane_debug_file = open(Path(self.debug_goal_lane_dir) / f"{self.video_tag}.csv", "w") - self._goal_lane_debug_file.write( - "step,ego_x,ego_y,goal_x,goal_y,gdir_x,gdir_y,goal_lane_idx,lane_dir_dot_route,lane_dist_to_goal_m\n" - ) - self.lights = list(self.world.get_actors().filter("traffic.traffic_light")) - light_map, bin_num_traffic = cb.map_lights_to_bin(self.lights, self.transform, self.town_bin) - # the shadow env runs on CARLA's trigger volumes: the bin's exported stop lines sit up to 9 m off and miss a few - self._stop_sign_lines, stop_sign_headings = cb.stop_signs_from_carla(self.world, self.cmap, self.transform) - self.num_traffic = self.env.set_stop_signs(self._stop_sign_lines, stop_sign_headings) - print( - f"[puffer_agent] stop signs: {len(stop_sign_headings)} CARLA trigger volumes replace the bin's " - f"(traffic elements {bin_num_traffic} -> {self.num_traffic})" - ) - # One id -> bin-indices table, reused by both passes in - # _read_light_states, so they can't disagree on which bin element a - # given light maps to (a prior independent geometric lookup for the - # ego-governance pass picked the wrong element ~5% of the time). - self.light_bin_indices_by_id = {lt.id: light_map[li] for li, lt in enumerate(self.lights)} - self.last_light_states = np.zeros(self.num_traffic, np.int32) - wheelbase, max_steer = read_vehicle_geometry(self.vehicle) - self.wheelbase_m = wheelbase body = self.vehicle.bounding_box self.body_front_m = body.location.x + body.extent.x self.body_rear_m = body.location.x - body.extent.x - self.controller = TrackingController(wheelbase_m=wheelbase, max_steer_rad=max_steer, horizon_s=self.dt) - - if self.obs_html_dir: - from pufferlib.ocean.cosim.obs_replay import ObsReplayCapture - - Path(self.obs_html_dir).mkdir(parents=True, exist_ok=True) - self._obs_html = ObsReplayCapture( - self.env, self.policy, Path(self.obs_html_dir) / self.video_tag, max_steps=self.obs_html_max_steps - ) + wheels = wheel_rows(self.vehicle) + self.wheelbase_m, _ = vehicle_geometry_from_wheels(wheels) + self.shadow.init_route( + { + "town": town, + "tick_dt": self.tick_dt, + "ego": { + "extent": [body.extent.x, body.extent.y, body.extent.z], + "box_z": body.location.z, + "wheels": wheels, + "state": actor_state_row(self.vehicle), + }, + "sparse_plan": plan_xyz(self._global_plan_world_coord), + "dense_plan": plan_xyz(self.dense_global_plan_world_coord), + "driving_waypoints": cb.driving_waypoint_samples(self.cmap), + "lights": cb.light_geometry_from_carla(self.lights, self.cmap), + "stop_signs": cb.stop_sign_geometry_from_carla(self.world, self.cmap), + "infraction_flags": bool(self.record_infractions_dir), + } + ) if self.debug_carla_view_dir: from pufferlib.ocean.cosim.carla_cosim import Mp4Writer @@ -440,16 +202,6 @@ def _init_on_first_step(self): out = str(Path(self.debug_carla_view_dir) / f"{self.video_tag}.mp4") self.carla_view_writer = Mp4Writer(out, fps=round(1.0 / self.tick_dt)) - if self.telemetry_dir: - Path(self.telemetry_dir).mkdir(parents=True, exist_ok=True) - self.telemetry_file = open(Path(self.telemetry_dir) / f"{self.video_tag}.csv", "w") - self.telemetry_file.write( - "step,current_speed,target_speed,ego_action,goal_cursor,goal_dist_m," - "near_light_dist_m,near_light_state," - "pd_infr_collision,pd_infr_offroad,pd_infr_red,pd_infr_stop," - "carla_infr_collision,carla_infr_offroad,carla_infr_red,carla_infr_stop\n" - ) - if self.record_infractions_dir: from collections import deque @@ -459,18 +211,11 @@ def _init_on_first_step(self): self.infraction_counter = 0 self.last_infraction_location = self.vehicle.get_location() - if self.telemetry_dir or self.record_infractions_dir: + if self.shadow.telemetry_dir or self.record_infractions_dir: self._init_carla_infraction_detectors() - - print( - f"[puffer_agent] town={town} tick_dt={self.tick_dt} dt={self.dt} route_goals={len(self.route_goals)} " - f"max_speed_mps={self.env.max_speed_mps:.1f} (C_vel={self.env.max_speed_mps / self.env.base_max_speed_mps:.2f}) " - f"zero_partner_stopped_time={int(self.zero_partner_stopped_time)} " - f"pedestrian_min_size_m={self.pedestrian_min_size_m:g}" - ) self.initialized = True - # --- shadow-env sync (read-only w.r.t. CARLA) -------------------------- + # --- CARLA ground truth for the shadow env (read-only w.r.t. CARLA) ---- def _nearby_actors(self): ego_loc = self.vehicle.get_location() @@ -484,196 +229,25 @@ def _nearby_actors(self): continue candidates.append((loc.distance(ego_loc), a)) candidates.sort(key=lambda item: item[0]) - return [a for _, a in candidates[: self.num_agents - 1]] + return [a for _, a in candidates[: self.shadow.partner_slots]] - def _read_states(self, actors): - idx, x, y, z, h, vx, vy, yaw_rate, accel_long, stopped_s = [], [], [], [], [], [], [], [], [], [] - stopped_by_id = {} - for j, a in enumerate(actors): - idx.append(1 + j) # agent 0 = ego; others fill 1..M - bx, by, bz, bh, bvx, bvy, byr, bal = self.transform.actor_state_to_bin(a) - x.append(bx) - y.append(by) - z.append(bz) - h.append(bh) - vx.append(bvx) - vy.append(bvy) - yaw_rate.append(byr) - accel_long.append(bal) - # stopped time follows the CARLA actor, not the shadow slot (slots reshuffle by distance every step) - is_stopped = math.hypot(bvx, bvy) <= self._partner_stopped_speed_threshold - stopped_by_id[a.id] = self._partner_stopped_s.get(a.id, 0.0) + self.dt if is_stopped else 0.0 - stopped_s.append(stopped_by_id[a.id]) - self._partner_stopped_s = stopped_by_id - if self.zero_partner_stopped_time: - stopped_s = [0.0] * len(actors) - return ( - np.array(idx, np.int32), - np.array(x, np.float32), - np.array(y, np.float32), - np.array(z, np.float32), - np.array(h, np.float32), - np.array(vx, np.float32), - np.array(vy, np.float32), - np.array(yaw_rate, np.float32), - np.array(accel_long, np.float32), - np.array(stopped_s, np.float32), - ) - - def _read_sizes(self, actors): - idx, length, width = [], [], [] - for j, a in enumerate(actors): - idx.append(1 + j) - ext = a.bounding_box.extent - walker_floor_m = self.pedestrian_min_size_m if "walker" in a.type_id else 0.0 - length.append(max(2.0 * ext.x, 0.1, walker_floor_m)) - width.append(max(2.0 * ext.y, 0.1, walker_floor_m)) - return (np.array(idx, np.int32), np.array(length, np.float32), np.array(width, np.float32)) - - def _read_light_states(self): - """Ground truth for every mapped bin traffic element, via the - precomputed light_bin_indices_by_id table (see _init_on_first_step).""" - # unmapped light elements read OFF (a training state), never UNKNOWN - states = np.full(self.num_traffic, cb.TRAFFIC_LIGHT_STATE_OFF, np.int32) - for lt in self.lights: - state = cb.carla_light_to_puffer(lt.get_state()) - for j in self.light_bin_indices_by_id[lt.id]: - if 0 <= j < self.num_traffic: - states[j] = state + def _snapshot(self): ego_light = self.vehicle.get_traffic_light() - if ego_light is not None and ego_light.id in self.light_bin_indices_by_id: - state = cb.carla_light_to_puffer(ego_light.get_state()) - for j in self.light_bin_indices_by_id[ego_light.id]: - if 0 <= j < self.num_traffic: - states[j] = state - return states - - def _sync_ego_from_carla(self, zero_velocity=False): - """Overwrite agent 0 (ego) from CARLA's ground-truth pose/size.""" - ego_state = self.transform.actor_state_to_bin(self.vehicle) - if zero_velocity: - x, y, z, heading, *_ = ego_state - ego_state = (x, y, z, heading, 0.0, 0.0, 0.0, 0.0) - self.env.set_agent_states(np.array([0], np.int32), *[np.array([v], np.float32) for v in ego_state]) - ego_ext = self.vehicle.bounding_box.extent - self.env.set_agent_sizes( - np.array([0], np.int32), np.array([2.0 * ego_ext.x], np.float32), np.array([2.0 * ego_ext.y], np.float32) - ) - - def _sync_carla(self): - """Overwrite the shadow env's background agents/lights (+ ego, only in - dynamics_source='carla' mode) from CARLA ground truth and return the - recomputed observation array (num_agents, obs_dim). Row 0 = ego.""" - if self.dynamics_source == "carla": - self._sync_ego_from_carla() - ego = self.env.get_global_agent_state() - ebx, eby = float(ego["x"][0]), float(ego["y"][0]) - - actors = self._nearby_actors() - if actors: - states = self._read_states(actors) - sizes = self._read_sizes(actors) - self.env.set_agent_states(*states) - self.env.set_agent_sizes(*sizes) - is_walker = np.array(["walker" in a.type_id for a in actors], np.float32) - self._last_partners = (states[1], states[2], states[4], states[5], states[6], sizes[1], sizes[2], is_walker) - else: - self._last_partners = None - n_used = 1 + len(actors) - if n_used < self.num_agents: - sp = np.arange(n_used, self.num_agents, dtype=np.int32) - zf = np.full(len(sp), FAR_AWAY, np.float32) - zz = np.zeros_like(zf) - self.env.set_agent_states(sp, zf, zf, zf, zz, zz, zz, zz, zz) - - self.last_light_states = self._read_light_states() - self.env.set_traffic_light_states(self.last_light_states) - - self.goal_window.sync(ebx, eby, float(ego["heading"][0])) - if self._goal_lane_debug_file is not None: - self._write_goal_lane_debug_row(ebx, eby, self.route_goals[self.goal_window.current_index]) - return np.asarray(self.env.recompute_observations()) - - def _write_goal_lane_debug_row(self, ebx, eby, current_goal): - """One CSV row: the CURRENT goal (window slot 0, current_goal_idx after - the reset in c_set_agent_goals) vs the lane find_goal_lane snapped it - to -- that lane's direction (nearest segment to the goal, mirroring - find_goal_lane's own point-to-segment search) dotted against the route - direction used at snap time, and the snap distance. dot < 0 means the - snapped lane runs opposite the route (oncoming lane); a large snap - distance means no nearby lane passed find_goal_lane's gates at all.""" - gx, gy, _, gdir_x, gdir_y = current_goal - state = self.env.get_state() - scenario = state[0] if isinstance(state, list) else state - ego_agent = (scenario.get("agents") or [{}])[0] - lane_idx = int(ego_agent.get("goal_lane_idx", -1)) - dot, dist = float("nan"), float("nan") - if lane_idx >= 0 and self._road_data is not None and lane_idx < len(self._road_data): - road = self._road_data[lane_idx] - xs, ys = np.asarray(road["x"], dtype=np.float64), np.asarray(road["y"], dtype=np.float64) - if len(xs) >= 2: - seg_dx, seg_dy = xs[1:] - xs[:-1], ys[1:] - ys[:-1] - seg_len = np.hypot(seg_dx, seg_dy) - seg_len_safe = np.where(seg_len > 1e-6, seg_len, 1.0) - t = np.clip(((gx - xs[:-1]) * seg_dx + (gy - ys[:-1]) * seg_dy) / (seg_len_safe**2), 0.0, 1.0) - px, py = xs[:-1] + t * seg_dx, ys[:-1] + t * seg_dy - d = np.hypot(gx - px, gy - py) - k = int(np.argmin(d)) - dist = float(d[k]) - route_norm = math.hypot(gdir_x, gdir_y) - if seg_len[k] > 1e-6 and route_norm > 1e-6: - dot = float((gdir_x * seg_dx[k] + gdir_y * seg_dy[k]) / (seg_len[k] * route_norm)) - self._goal_lane_debug_file.write( - f"{self.step},{ebx:.2f},{eby:.2f},{gx:.2f},{gy:.2f},{gdir_x:.3f},{gdir_y:.3f}," - f"{lane_idx},{dot:.3f},{dist:.2f}\n" - ) - self._goal_lane_debug_file.flush() - - def _carla_integrate(self, actions): - """Step the shadow env ONCE (one policy step == one shadow tick). - - dynamics_source='carla': return the ego's target kinematic state for - the tracking controller - - dynamics_source='pufferdrive': PufferDrive's own jerk/classic dynamics - (the ones the policy trained on) ARE the ego's motion -- no target to - chase. The CARLA ego is teleported to the post-step pose (see - _teleport_carla_ego), where the next tick's observation reads it. - """ - if self.step >= MAX_POLICY_STEPS: - raise RuntimeError(f"route exceeded MAX_POLICY_STEPS={MAX_POLICY_STEPS} policy steps") - before = self.env.get_global_agent_state() - speed_before = self._ego_speed() - self.env.step(actions) - after = self.env.get_global_agent_state() - target_yaw_deg = self.transform.bin_heading_to_yaw(float(after["heading"][0])) - speed_after = self._ego_speed() - accel_after = float(np.asarray(self.env.observations)[0][4]) * self._accel_long_norm() - - if self.dynamics_source == "pufferdrive": - x0, y0 = self.transform.bin_to_loc(float(before["x"][0]), float(before["y"][0])) - x1, y1 = self.transform.bin_to_loc(float(after["x"][0]), float(after["y"][0])) - yaw0_deg = self.transform.bin_heading_to_yaw(float(before["heading"][0])) - self._motion = (x0, y0, yaw0_deg, speed_before, x1, y1, target_yaw_deg, speed_after, float(after["z"][0])) - self._teleport_carla_ego(1.0) - return speed_after, target_yaw_deg # logged, not chased (see run_step) - - target_speed = speed_after + max(accel_after, 0.0) * self.speed_intent_extension_s - return target_speed, target_yaw_deg + return { + "ego": actor_state_row(self.vehicle), + "partners": [partner_row(a) for a in self._nearby_actors()], + "light_states": [[lt.id, lt.get_state().name] for lt in self.lights], + "ego_light_id": -1 if ego_light is None else ego_light.id, + } - def _teleport_carla_ego(self, fraction): - """Place the CARLA ego at `fraction` (0 = pre-step, 1 = post-step) of the current policy step's motion.""" - x0, y0, yaw0_deg, speed0, x1, y1, yaw1_deg, speed1, sim_z = self._motion - yaw_delta_deg = cb.wrap_deg_180(yaw1_deg - yaw0_deg) - x = x0 + fraction * (x1 - x0) - y = y0 + fraction * (y1 - y0) - yaw_deg = yaw0_deg + fraction * yaw_delta_deg - speed = speed0 + fraction * (speed1 - speed0) + def _teleport_carla_ego(self, motion): + """Place the CARLA ego at the shadow env's post-step pose (StepResult.motion).""" + _, _, yaw0_deg, _, x, y, yaw_deg, speed, carla_z = motion + yaw_delta_deg = cb.wrap_deg_180(yaw_deg - yaw0_deg) # CARLA's live road-mesh height, not the shadow env's lane-averaged sim_z: on graded # multi-level roads the average can land the body mid-structure (measured: 20-100+ # road collisions per Town03/04 route). sim_z only off the drivable network. # get_waypoint is nearest-in-3D: without sim_z the z=0 default snaps an overpass ego to the road beneath - carla_z = self.transform.z_to_carla(sim_z) wp = self.cmap.get_waypoint(carla.Location(x=x, y=y, z=carla_z)) z = wp.transform.location.z if wp is not None else carla_z road_up = wp.transform.rotation.get_up_vector() if wp is not None else carla.Vector3D(x=0.0, y=0.0, z=1.0) @@ -692,16 +266,27 @@ def _teleport_carla_ego(self, fraction): ) # Zero momentum before the teleport: CARLA's collision resolver reacts violently to a # physics body carrying velocity into a new pose (carla issue #8076). + # One ordered batch: the server reorders async actor commands, and a late zero left the ego reading 0 m/s on 3 % of ticks zero = carla.Vector3D(x=0.0, y=0.0, z=0.0) - self.vehicle.set_target_velocity(zero) - self.vehicle.set_target_angular_velocity(zero) - self.vehicle.set_transform( - carla.Transform(carla.Location(x=x, y=y, z=z), carla.Rotation(pitch=pitch_deg, yaw=yaw_deg, roll=roll_deg)) - ) - self.vehicle.set_target_velocity( - carla.Vector3D(x=speed * forward[0], y=speed * forward[1], z=speed * forward[2]) + CarlaDataProvider.get_client().apply_batch_sync( + [ + carla.command.ApplyTargetVelocity(self.vehicle.id, zero), + carla.command.ApplyTargetAngularVelocity(self.vehicle.id, zero), + carla.command.ApplyTransform( + self.vehicle.id, + carla.Transform( + carla.Location(x=x, y=y, z=z), carla.Rotation(pitch=pitch_deg, yaw=yaw_deg, roll=roll_deg) + ), + ), + carla.command.ApplyTargetVelocity( + self.vehicle.id, carla.Vector3D(x=speed * forward[0], y=speed * forward[1], z=speed * forward[2]) + ), + carla.command.ApplyTargetAngularVelocity( + self.vehicle.id, carla.Vector3D(x=0.0, y=0.0, z=yaw_delta_deg / self.shadow.dt) + ), + ], + False, ) - self.vehicle.set_target_angular_velocity(carla.Vector3D(x=0.0, y=0.0, z=yaw_delta_deg / self.dt)) def _road_surface(self, x, y, reference_z, yaw_deg): """(z, pitch_deg) resting the body on the road mesh: bumper-to-bumper chord, lifted until no sample is below the mesh; None off the mesh.""" @@ -716,32 +301,7 @@ def _road_surface(self, x, y, reference_z, yaw_deg): z = max(z_i - slope * s_i for z_i, s_i in zip(sample_z, sample_offsets_m)) return z, math.degrees(math.atan(slope)) - def _ego_speed(self): - return float(np.asarray(self.env.observations)[0][0]) * self._max_speed() - - def _max_speed(self): - return self.env.obs_norm_speed_mps - - def _accel_long_norm(self): - from pufferlib.ocean.drive import binding - - return binding.ACCEL_LONG_NORM - - def _ego_infractions(self): - """{'collision': f, 'offroad': f, 'red_light': f, 'stop_sign': f} from the shadow - pufferdrive env's own model of the ego (compute_metrics, refreshed by - the last integrate()) -- what the policy would have caused according - to PufferDrive's own dynamics, not necessarily what the real CARLA - vehicle did (see _carla_infractions).""" - state = self.env.get_state() - scenario = state[0] if isinstance(state, list) else state - agents = scenario.get("agents") or [] - if not agents: - return {name: 0.0 for name in EGO_INFRACTION_METRICS} - metrics = agents[0].get("metrics_array") or [] - return { - name: float(metrics[idx]) if idx < len(metrics) else 0.0 for name, idx in EGO_INFRACTION_METRICS.items() - } + # --- CARLA ground-truth infraction detectors (CaRL's criteria) ----------- def _init_carla_infraction_detectors(self): """Real CARLA ground truth for collision/offroad/red-light/stop-sign, reusing @@ -820,7 +380,7 @@ def _poll_carla_stop_sign(self): def _carla_infractions(self): """{'collision': f, 'offroad': f, 'red_light': f, 'stop_sign': f} from real CARLA ground truth (_init_carla_infraction_detectors) -- contrast with - _ego_infractions, which is the shadow env's model of the ego.""" + StepResult.pd_flags, which is the shadow env's model of the ego.""" collision, self._carla_collision_pending = self._carla_collision_pending, False red_light, self._ran_red_light_pending = self._ran_red_light_pending, False stop_sign, self._ran_stop_sign_pending = self._ran_stop_sign_pending, False @@ -833,137 +393,6 @@ def _carla_infractions(self): "stop_sign": float(stop_sign), } - # --- policy + capture --------------------------------------------------- - - def _adapt_obs_for_policy(self, obs): - """Strip the trailing zero-padded boundary columns when the checkpoint - predates them (see _load_policy_and_env). No-op otherwise.""" - if self._boundary_feat_ckpt is None: - return obs - env = self.env - n_slots = env.obs_slots_boundary_kept - feat_env, feat_ckpt = env.boundary_features, self._boundary_feat_ckpt - b0 = ( - env.ego_features - + env.num_reward_coefs - + env.goal_dim - + env.obs_slots_partners_n * env.partner_features - + env.obs_slots_lane_kept * env.lane_features - ) - b1 = b0 + n_slots * feat_env - obs = np.asarray(obs) - boundary = obs[:, b0:b1].reshape(obs.shape[0], n_slots, feat_env)[:, :, :feat_ckpt] - return np.concatenate([obs[:, :b0], boundary.reshape(obs.shape[0], -1), obs[:, b1:]], axis=1) - - def _policy_actions(self, obs): - """-> (actions (num_agents, act_dim) int32, aux dict). aux carries - per-agent value/entropy/action-probs/pool for the obs_html capture""" - import torch - import pufferlib.pytorch - - import pufferlib.spaces - - obs = self._adapt_obs_for_policy(obs) - # A discrete policy head on a continuous env needs the bin->continuous - # mapping (pufferl.py feeds cont_action to the env, never the bin indices). - env_continuous = isinstance(self.env.single_action_space, pufferlib.spaces.Box) - action_selection = ( - pufferlib.pytorch.ACTION_SELECT_MEAN - if env_continuous and not self.policy.is_continuous - else pufferlib.pytorch.ACTION_SELECT_MODE - ) - with torch.no_grad(): - logits, value = self.policy.forward_eval(torch.as_tensor(obs).to(self.device)) - action, _, entropy, cont_action = pufferlib.pytorch.sample_logits( - logits, action_selection=action_selection, env_continuous=env_continuous, policy=self.policy - ) - env_action = cont_action if env_continuous and cont_action is not None else action - actions = env_action.cpu().numpy().reshape(1, -1) - if not env_continuous: - actions = actions.astype(np.int32) - aux = {} - if self._obs_html is not None: - aux = self._obs_html.policy_outputs(torch.as_tensor(obs).to(self.device), logits, value, entropy) - aux["action_index"] = action.cpu().numpy().reshape(-1) if not self.policy.is_continuous else None - return actions, aux - - def _record_world_state(self, obs): - """Bin-frame snapshot at sync time (pre-step ego, partners as streamed, light states).""" - ego = self.env.get_global_agent_state() - accel_long = float(obs[0][4]) * self._accel_long_norm() - step_idx = len(self._world_log["ego"]) - self._world_log["ego"].append( - ( - self.step, - float(ego["x"][0]), - float(ego["y"][0]), - float(ego["heading"][0]), - self._ego_speed(), - accel_long, - float(self.goal_window.current_index), - ) - ) - if self._last_partners is not None: - x, y, h, vx, vy, length, width, is_walker = self._last_partners - speed = np.hypot(vx, vy) - rows = np.column_stack([np.full(len(x), step_idx, np.float32), x, y, h, length, width, speed, is_walker]) - self._world_log["partners"].append(rows.astype(np.float32)) - self._world_log["lights"].append(self.last_light_states.astype(np.int8)) - - def _write_world_log(self): - Path(self.world_log_dir).mkdir(parents=True, exist_ok=True) - dense = np.array( - [self.transform.loc_to_bin(t.location.x, t.location.y) for t, _ in self.dense_global_plan_world_coord], - np.float32, - ) - partners = self._world_log["partners"] - np.savez_compressed( - Path(self.world_log_dir) / f"{self.video_tag}.npz", - ego=np.array(self._world_log["ego"], np.float32), - partners=np.concatenate(partners) if partners else np.zeros((0, 8), np.float32), - lights=np.stack(self._world_log["lights"]) if self._world_log["lights"] else np.zeros((0, 0), np.int8), - route_goals=self.route_goals, - stop_signs=self._stop_sign_lines, - dense_route=dense, - meta=json.dumps( - { - "town": Path(self.town_bin).stem.split("__")[-1], - "town_bin": self.town_bin, - "dt": self.dt, - "tick_dt": self.tick_dt, - "offset": [self.transform.tx, self.transform.ty], - "z_offset": self.transform.tz, - } - ), - ) - - def _write_telemetry_row(self, ego_action, pd_flags, carla_flags): - """One CSV row per policy step: what the loop commanded vs achieved, - plus the nearest mapped light's state and both infraction sources - (pd_* = shadow pufferdrive env, carla_* = real CARLA ground truth).""" - ego = self.env.get_global_agent_state() - ex, ey = float(ego["x"][0]), float(ego["y"][0]) - cur = min(self.goal_window.current_index, len(self.route_goals) - 1) - goal_dist = float(np.hypot(self.route_goals[cur, 0] - ex, self.route_goals[cur, 1] - ey)) - near_dist, near_state = -1.0, -1 - if len(self.stop_line_centers): - d2 = (self.stop_line_centers[:, 0] - ex) ** 2 + (self.stop_line_centers[:, 1] - ey) ** 2 - j = int(d2.argmin()) - near_dist, near_state = float(np.sqrt(d2[j])), int(self.last_light_states[j]) - current_speed = ( - float(np.asarray(self.env.observations)[0, 0]) * self._max_speed() - if self.dynamics_source == "pufferdrive" - else self.vehicle.get_velocity().length() - ) - self.telemetry_file.write( - f"{self.step},{current_speed:.3f},{self.target[0]:.3f}," - f"{ego_action},{self.goal_window.current_index},{goal_dist:.1f},{near_dist:.1f},{near_state}," - f"{pd_flags['collision']:.0f},{pd_flags['offroad']:.0f},{pd_flags['red_light']:.0f}," - f"{pd_flags['stop_sign']:.0f}," - f"{carla_flags['collision']:.0f},{carla_flags['offroad']:.0f},{carla_flags['red_light']:.0f}," - f"{carla_flags['stop_sign']:.0f}\n" - ) - def _maybe_save_infraction_clip(self, pd_flags, carla_flags): """Dump the rolling chase-cam buffer as one mp4 when either infraction source flags an ego infraction (collision/offroad/red-light/stop-sign), at most @@ -1005,61 +434,27 @@ def run_step(self, input_data, timestamp, sensors=None): self._poll_carla_red_light() # filters single-tick get_traffic_light() noise self._poll_carla_stop_sign() - obs = self._sync_carla() # shadow env <- CARLA ground truth - actions, aux = self._policy_actions(obs) - if self.obs_dump_dir: - self._obs_dump.append((self.step, obs[0].astype(np.float32), actions[0].astype(np.float32))) - if self.world_log_dir: - self._record_world_state(obs) - if self._obs_html is not None: - self._obs_html.capture(obs, actions, aux, aux.get("action_index")) - self.target = self._carla_integrate(actions) # policy intent, one dt ahead - if self.telemetry_file is not None or self.record_infractions_dir: - # _init_carla_infraction_detectors ran for either dir being set - pd_flags, carla_flags = self._ego_infractions(), self._carla_infractions() - if self.telemetry_file is not None: - self._write_telemetry_row(float(actions[0, 0]), pd_flags, carla_flags) + result = self.shadow.step(self._snapshot()) # shadow env <- CARLA ground truth, policy, one dt + if result.motion is not None: + self._teleport_carla_ego(result.motion) + if self._carla_collision is not None: + carla_flags = self._carla_infractions() + self.shadow.telemetry(result, carla_flags) if self.record_infractions_dir: - self._maybe_save_infraction_clip(pd_flags, carla_flags) + self._maybe_save_infraction_clip(result.pd_flags, carla_flags) - if self.dynamics_source == "pufferdrive": + if result.control is None: return carla.VehicleControl() - - # Controller runs every tick against the latest CARLA state, chasing the - # target held from the last policy step. - current_speed = self.vehicle.get_velocity().length() - current_yaw = self.vehicle.get_transform().rotation.yaw - target_speed, target_yaw = self.target - return self.controller.step(current_speed, current_yaw, target_speed, target_yaw, self.tick_dt) + steer, throttle, brake = result.control + return carla.VehicleControl(steer=steer, throttle=throttle, brake=brake) def destroy(self, results=None): if not self.initialized: return - print( - f"[puffer_agent] route done: goals {self.goal_window.current_index + 1}/{len(self.route_goals)}, " - f"tracking {self.controller.stats()}" - ) + self.shadow.finish() if self.carla_view_writer is not None: self.carla_view_writer.close() print(f"[puffer_agent] wrote CARLA chase-cam video for route {self.video_tag}") - if self.telemetry_file is not None: - self.telemetry_file.close() if self._carla_collision is not None: self._carla_collision.clean() # stop+destroy the spawned collision sensor actor - if self._goal_lane_debug_file is not None: - self._goal_lane_debug_file.close() - if self._obs_html is not None: - written = self._obs_html.write(render_html=self.obs_html_render) - print(f"[puffer_agent] wrote obs replay ({len(self._obs_html)} frames) -> {written}") - if self.world_log_dir and self._world_log["ego"]: - self._write_world_log() - if self.obs_dump_dir and self._obs_dump: - Path(self.obs_dump_dir).mkdir(parents=True, exist_ok=True) - np.savez_compressed( - Path(self.obs_dump_dir) / f"{self.video_tag}.npz", - step=np.array([s for s, _, _ in self._obs_dump], np.int32), - obs=np.stack([o for _, o, _ in self._obs_dump]), - action=np.stack([a for _, _, a in self._obs_dump]), - ) - self.env.close() self.initialized = False diff --git a/pufferlib/ocean/cosim/carla/shadow_ego.py b/pufferlib/ocean/cosim/carla/shadow_ego.py new file mode 100644 index 0000000000..ea4ae7f950 --- /dev/null +++ b/pufferlib/ocean/cosim/carla/shadow_ego.py @@ -0,0 +1,837 @@ +"""PufferDrive shadow env + policy driving one ego from another simulator's ground truth. + +Shared by the CARLA leaderboard integrations: CaRL's Leaderboard 2.0 agent (leaderboard_agent.py, +CARLA 0.9.15, in-process) and the Leaderboard 1.0 client/server pair (lb1/, CARLA 0.9.10, whose +Python 3.7 cannot import PufferDrive). The simulator side only reads CARLA and moves the ego; +everything the policy needs happens here and never imports carla. Inputs and outputs are in the +CARLA frame (metres, degrees, left-handed y) as plain lists/dicts, so they also travel as JSON. + +Everything about the shadow env (obs layout, dynamics, dt, goal count, agent pool size, ...) comes +from the checkpoint's sibling config.yaml with the clean-eval profile applied on top (cosim/arch.py +CLEAN_EVAL_OVERRIDES); only the structural co-sim keys (map, pool wiring) are set here. + +Actor state rows (CARLA frame, physics quantities from CARLA's own physics): + ACTOR_STATE [x, y, z, yaw_deg, vx, vy, vz, angular_velocity_z_deg_s, accel_x, accel_y] + PARTNER_ROW ACTOR_STATE + [box_z, extent_x, extent_y, extent_z, is_walker, is_two_wheeler, actor_id] + +init_route(route) once the ego exists, route = { + town, tick_dt, + ego: {extent: [x, y, z], box_z, wheels: controller.wheel_rows, state: ACTOR_STATE}, + sparse_plan: (N, 3) leaderboard target points, dense_plan: (M, 3) the dense route, + driving_waypoints: carla_bridge.driving_waypoint_samples rows, + lights: carla_bridge.light_geometry_from_carla dicts, stop_signs: carla_bridge.stop_sign_geometry_from_carla dicts, + infraction_flags: bool, also compute the shadow env's own infraction flags when no telemetry is written } +step(snapshot) every simulator tick, snapshot = { + ego: ACTOR_STATE, partners: PARTNER_ROW list (nearest first, at most partner_slots, |dz| <= PARTNER_MAX_ABS_DZ_M), + light_states: [[light_id, carla.TrafficLightState name], ...], ego_light_id: id reported for the ego or -1 } + +Environment variables: + COSIM_DEVICE=cpu torch device for the policy + COSIM_DYNAMICS_SOURCE=pufferdrive + "pufferdrive" (default): PufferDrive's own dynamics (the ones the policy + trained on) move the ego; StepResult.motion is the pose the simulator + teleports the ego to every step. + "carla": TrackingController converts the shadow env's target (speed, yaw) + into StepResult.control = (steer, throttle, brake) and CARLA's own vehicle + physics moves the ego (subject to dynamics-mismatch tracking lag). + COSIM_MAX_SPEED_MPS=20 ego speed cap; sets the shadow env's C_vel = cap / base_max_speed_mps. + Default: the checkpoint's base_max_speed_mps (C_vel = 1). Must lie in the + trained range base * [1/conditioning_speed_scale, scale] + (0036: 13.3-30 m/s, 0038+: 10-40 m/s). + COSIM_ZERO_PARTNER_STOPPED_TIME=0 + 1: every partner's seconds_stopped observation feature is held at 0 (the + ego's own counter is untouched). Ablates the stopped-time cue behind + parked-car / red-light-queue avoidance. + COSIM_VRU_SIZE_PAD_M=0.0 added to the length and width of every walker and two-wheeler partner box + (nuPlan's vru_size_pad_m; the training spawn floor is 0.8 m, CARLA walkers + are 0.4-0.5 m and bicycles 0.4 m wide, so 0.8 lifts them strictly above + the smallest training vehicle) + COSIM_TELEMETRY=/dir write a per-policy-step CSV per route, with the shadow env's and the + simulator's infraction flags as columns + COSIM_OBS_HTML=/dir write an interactive pufferlib.viz replay per route (the exact obs + + policy outputs the ego saw) + COSIM_OBS_HTML_RENDER=1 0: save only the compact .replay.zlib per route; + scripts/eval/render_carla_obs_html.py renders the pages of a whole run + into one gallery afterwards + COSIM_OBS_HTML_MAX_STEPS=12000 frames kept per route replay + COSIM_DUMP_OBS=/dir write the raw ego obs + action per policy step per route (.npz) + COSIM_WORLD_LOG=/dir write the bin-frame world state per policy step per route (ego, + partners, light states, route); input of scripts/eval/analyze_carla_cosim.py + COSIM_DEBUG_GOAL_LANE=/dir per-step CSV of the current goal vs the lane find_goal_lane snapped it to +""" + +import json +import math +import os +from pathlib import Path + +import numpy as np + +from pufferlib.ocean.cosim import carla_bridge as cb +from pufferlib.ocean.cosim.arch import checkpoint_config_path, shadow_env_kwargs +from pufferlib.ocean.cosim.carla.controller import TrackingController, vehicle_geometry_from_wheels +from pufferlib.ocean.cosim.goals import RouteGoalWindow, route_goals_from_target_points +from pufferlib.ocean.drive.drive import Drive + + +ACTOR_STATE_COLUMNS = 10 +PARTNER_ROW_COLUMNS = 17 +( + STATE_X, + STATE_Y, + STATE_Z, + STATE_YAW_DEG, + STATE_VX, + STATE_VY, + STATE_VZ, + STATE_ANGULAR_VELOCITY_Z, + STATE_ACCEL_X, + STATE_ACCEL_Y, +) = range(ACTOR_STATE_COLUMNS) +( + PARTNER_BOX_Z, + PARTNER_EXTENT_X, + PARTNER_EXTENT_Y, + PARTNER_EXTENT_Z, + PARTNER_IS_WALKER, + PARTNER_IS_TWO_WHEELER, + PARTNER_ACTOR_ID, +) = range(ACTOR_STATE_COLUMNS, PARTNER_ROW_COLUMNS) + +FAR_AWAY = 1.0e6 # park unused shadow-env agent slots out of observation range +PARTNER_MAX_ABS_DZ_M = 20.0 # CARLA actors farther above/below the ego are hidden scenario props, not traffic +DEFAULT_BASE_MAX_SPEED_MPS = 20.0 # Drive() defaults, for checkpoints whose config predates the keys +DEFAULT_CONDITIONING_SPEED_SCALE = 1.5 +MAX_POLICY_STEPS = 100_000 # shadow-env episode cap (~1.4 h at the 0.05 s CARLA tick); sizes the light state buffers +SCENARIO_LENGTH_MARGIN_STEPS = 2 +MIN_PARTNER_SIZE_M = 0.1 +# Shadow-env metrics_array indices (datatypes.h): collision/offroad/red-light/stop-sign flags. +EGO_INFRACTION_METRICS = {"collision": 0, "offroad": 1, "red_light": 2, "stop_sign": 3} +DYNAMICS_SOURCES = ("carla", "pufferdrive") +TELEMETRY_HEADER = ( + "step,current_speed,target_speed,ego_action,goal_cursor,goal_dist_m," + "near_light_dist_m,near_light_state," + "pd_infr_collision,pd_infr_offroad,pd_infr_red,pd_infr_stop," + "carla_infr_collision,carla_infr_offroad,carla_infr_red,carla_infr_stop\n" +) +GOAL_LANE_DEBUG_HEADER = ( + "step,ego_x,ego_y,goal_x,goal_y,gdir_x,gdir_y,goal_lane_idx,lane_dir_dot_route,lane_dist_to_goal_m\n" +) + + +def clean_policy_state_dict(state_dict): + """Strip torch.compile / DDP prefixes (wandb, neptune, the _C kernel, ...) + from the leaderboard's evaluation environment.""" + + def clean(key): + while key.startswith(("module.", "_orig_mod.")): + key = key.split(".", 1)[1] + return key + + return {clean(k): v for k, v in state_dict.items()} + + +def resolve_checkpoint(path_to_conf_file): + """(checkpoint_path, config_dict) from a .pt file or an experiment dir.""" + import yaml + + p = Path(path_to_conf_file).resolve() + if p.is_file(): + ckpt, cfg_path = p, checkpoint_config_path(p) + else: + models = sorted((p / "models").glob("*.pt")) or sorted(p.glob("*.pt")) + if not models: + raise FileNotFoundError(f"no .pt checkpoint under {p}") + ckpt, cfg_path = models[-1], p / "config.yaml" + with open(cfg_path) as f: + cfg = yaml.safe_load(f) + return str(ckpt), cfg + + +def cosim_max_speed_mps(env_cfg): + """COSIM_MAX_SPEED_MPS as a float inside the checkpoint's trained C_vel range, or None when unset.""" + raw = os.environ.get("COSIM_MAX_SPEED_MPS") + if raw is None: + return None + max_speed_mps = float(raw) + base_max_speed_mps = float(env_cfg.get("base_max_speed_mps") or DEFAULT_BASE_MAX_SPEED_MPS) + scale = float(env_cfg.get("conditioning_speed_scale") or DEFAULT_CONDITIONING_SPEED_SCALE) + low_mps, high_mps = base_max_speed_mps / scale, base_max_speed_mps * scale + if not (low_mps <= max_speed_mps <= high_mps): + raise ValueError( + f"COSIM_MAX_SPEED_MPS={max_speed_mps} is outside the trained C_vel range " + f"[{low_mps:.1f}, {high_mps:.1f}] m/s (base_max_speed_mps={base_max_speed_mps}, " + f"conditioning_speed_scale={scale})" + ) + return max_speed_mps + + +def env_flag(name, default="0"): + raw = os.environ.get(name, default) + if raw not in ("0", "1"): + raise ValueError(f"{name} must be '0' or '1', got {raw!r}") + return raw == "1" + + +def ego_speed_from_state(state): + """CARLA's Vector3D.length() of the ego velocity (3-D, as the leaderboard criteria read it).""" + return math.sqrt(state[STATE_VX] ** 2 + state[STATE_VY] ** 2 + state[STATE_VZ] ** 2) + + +class StepResult: + """One policy step: `actions` (num_agents, act_dim), the shadow env's infraction flags (None unless + requested), and the ego command -- `motion` = [x0, y0, yaw0_deg, speed0, x1, y1, yaw1_deg, speed1, z1] + (CARLA frame, pre- and post-step pose to teleport along) in pufferdrive mode, `control` = (steer, + throttle, brake) in carla mode; `target` = (speed_mps, yaw_deg) the policy intends one dt ahead.""" + + def __init__(self, actions, target, motion=None, control=None, pd_flags=None): + self.actions = actions + self.target = target + self.motion = motion + self.control = control + self.pd_flags = pd_flags + + +class ShadowEgo: + def __init__(self, path_to_conf_file, route_tag): + self.route_tag = route_tag + self.checkpoint, self.cfg = resolve_checkpoint(path_to_conf_file) + self.device = os.environ.get("COSIM_DEVICE", "cpu") + self.dynamics_source = os.environ.get("COSIM_DYNAMICS_SOURCE", "pufferdrive") + if self.dynamics_source not in DYNAMICS_SOURCES: + raise ValueError(f"COSIM_DYNAMICS_SOURCE must be one of {DYNAMICS_SOURCES}, got {self.dynamics_source!r}") + env_cfg = self.cfg["env"] + self.max_speed_mps = cosim_max_speed_mps(env_cfg) + self.zero_partner_stopped_time = env_flag("COSIM_ZERO_PARTNER_STOPPED_TIME") + self.vru_size_pad_m = float(os.environ.get("COSIM_VRU_SIZE_PAD_M", "0.0")) + if not (self.vru_size_pad_m >= 0.0): + raise ValueError(f"COSIM_VRU_SIZE_PAD_M must be >= 0, got {self.vru_size_pad_m}") + # Shadow agent pool == the training per-env cap + self.num_agents = int(env_cfg["max_agents_per_env"]) + + self.telemetry_dir = os.environ.get("COSIM_TELEMETRY", None) + self.telemetry_file = None + self.obs_html_dir = os.environ.get("COSIM_OBS_HTML", None) + self.obs_dump_dir = os.environ.get("COSIM_DUMP_OBS", None) + self._obs_dump = [] + self.world_log_dir = os.environ.get("COSIM_WORLD_LOG", None) + self._world_log = {"ego": [], "partners": [], "lights": []} + self._last_partners = None + self._partner_stopped_s = {} # CARLA actor id -> seconds at standstill + from pufferlib.ocean.drive import binding + + self._partner_stopped_speed_threshold = float(binding.AGENT_STOPPED_SPEED_THRESHOLD) + self.obs_html_max_steps = int(os.environ.get("COSIM_OBS_HTML_MAX_STEPS", "12000")) + self.obs_html_render = env_flag("COSIM_OBS_HTML_RENDER", "1") + self._obs_html = None + self.debug_goal_lane_dir = os.environ.get("COSIM_DEBUG_GOAL_LANE", None) + self._goal_lane_debug_file = None + self._road_data = None + + self.step_count = 0 + self.initialized = False + self.infraction_flags = False + self.policy = None + self.env = None + self.goal_window = None + self.target = (0.0, 0.0) # (target_speed, target_yaw_deg), held between policy steps + self._ego_speed_carla = 0.0 # CARLA's own ego speed at the last sync (carla dynamics telemetry) + self._stop_sign_lines = np.zeros((0, 6), np.float32) # CARLA trigger volumes as bin-frame stop lines + + @property + def partner_slots(self): + return self.num_agents - 1 + + def _load_policy_and_env(self, town_bin): + arch = shadow_env_kwargs( + self.cfg, + overrides=dict( + map_dir=town_bin, + num_maps=1, + # One policy agent (the ego); every other slot is a static partner streamed from CARLA. + num_agents=1, + min_agents_per_env=1, + cosim_partner_slots=self.partner_slots, + goal_source="external", + dt=self.dt, + scenario_length=MAX_POLICY_STEPS + SCENARIO_LENGTH_MARGIN_STEPS, + resample_frequency=0, + termination_mode=0, + # Enforcement off (detection flags still fire): in one endless + # episode a "stop" latch is permanent and fires only AFTER + # CARLA scored the infraction (route 5 froze 320 s -> DNF). + collision_behavior="ignore", + offroad_behavior="ignore", + traffic_light_behavior="ignore", + **({} if self.max_speed_mps is None else {"max_speed_mps": self.max_speed_mps}), + ), + ) + self.dynamics_model = arch.get("dynamics_model", "classic") + # jerk dynamics build accel across steps while sync() re-seeds accel + self.speed_intent_extension_s = 2.0 * self.dt if self.dynamics_model == "jerk" else 0.0 + + self.env = Drive(**arch) + self.env.reset() + + import torch + + import pufferlib.ocean.torch as drive_torch + + policy_state_dict = clean_policy_state_dict( + torch.load(self.checkpoint, map_location=self.device, weights_only=False) + ) + env_for_policy = self.env + self._boundary_feat_ckpt = None + bw = policy_state_dict.get("actor_backbone.boundary_encoder.0.weight") + if bw is not None and bw.shape[1] < env_for_policy.boundary_features: + # Older checkpoints emit fewer boundary features than this drive.h + # (e.g. 7 vs 9 -- two zero-padded GPS columns were added later). + import copy + + self._boundary_feat_ckpt = int(bw.shape[1]) + env_for_policy = copy.copy(self.env) + env_for_policy.boundary_features = self._boundary_feat_ckpt + print( + f"[puffer_agent] checkpoint uses {self._boundary_feat_ckpt} boundary features " + f"(env emits {self.env.boundary_features}); stripping zero-padded columns" + ) + + policy_cls = getattr(drive_torch, self.cfg.get("policy_name", "Drive")) + # pre-3.0 checkpoints keep action_type only in the env section + self.cfg["policy"].setdefault("action_type", self.cfg["env"]["action_type"]) + self.policy = policy_cls(env_for_policy, **self.cfg["policy"]).to(self.device) + self.policy.load_state_dict(policy_state_dict) + self.policy.eval() + print(f"[puffer_agent] loaded policy from {self.checkpoint}") + + def init_route(self, route): + """Build the shadow env for this route's town and seed it from the simulator's ground truth (see + the module docstring for the `route` dict). Read-only with respect to the simulator.""" + town = route["town"] + self.tick_dt = float(route["tick_dt"]) + self.dt = self.tick_dt # the policy runs every CARLA tick, like the leaderboard's sensor agents + self.infraction_flags = bool(route.get("infraction_flags", False)) + ego = route["ego"] + self.ego_extent = [float(v) for v in ego["extent"]] + self.ego_box_z = float(ego["box_z"]) + + self.town_bin = cb.bin_path_for_town(town) + self._load_policy_and_env(self.town_bin) + + self.transform = cb.CarlaTransform(town, offset=cb.town_offset(self.town_bin)) + offset, z_offset, residual_before, residual_after = cb.calibrate_town_offset( + route["driving_waypoints"], self.transform, self.town_bin + ) + print( + f"[puffer_agent] bin offset calibrated against CARLA lanes: shift " + f"({offset[0] - self.transform.tx:+.2f}, {offset[1] - self.transform.ty:+.2f}) m, z {z_offset:+.2f} m, " + f"median lane residual {residual_before:.2f} -> {residual_after:.2f} m" + ) + self.transform = cb.CarlaTransform(town, offset=offset, z_offset=z_offset) + ego_state = [float(v) for v in ego["state"]] + if self.dynamics_source == "pufferdrive": + self._sync_ego(ego_state, zero_velocity=True) + target_points = np.asarray(route["sparse_plan"], dtype=np.float64).reshape(-1, 3).copy() + dense_points = np.asarray(route["dense_plan"], dtype=np.float64).reshape(-1, 3).copy() + self.dense_plan_xy = dense_points[:, :2].copy() # CARLA frame, for the world log + target_points[:, 2] = self.transform.z_to_bin(target_points[:, 2]) + dense_points[:, 2] = self.transform.z_to_bin(dense_points[:, 2]) + self.route_goals = route_goals_from_target_points( + target_points, + dense_points, + self.transform.loc_to_bin, + self.transform.loc_to_bin(ego_state[STATE_X], ego_state[STATE_Y]), + self.env.goal_radius, + ) + # the leaderboard's target points, next num_goals of them at all times (fewer only at the route end) + self.goal_window = RouteGoalWindow(self.env, self.route_goals, sliding=True) + + import data_utils.mirror_map_bin as mbin + + bin_data = mbin.read_bin(Path(self.town_bin)) + bin_traffic = bin_data["traffic"] + self.stop_line_centers = np.array( + [ + [0.5 * (t["stop_line"][0] + t["stop_line"][3]), 0.5 * (t["stop_line"][1] + t["stop_line"][4])] + for t in bin_traffic + ] + ).reshape(-1, 2) + + if self.debug_goal_lane_dir: + # goal_lane_idx comes from road_elements[] in the C file; bin_data["roads"] + # is the same bin read in the same order, so it's usable as-is. + self._road_data = bin_data["roads"] + Path(self.debug_goal_lane_dir).mkdir(parents=True, exist_ok=True) + self._goal_lane_debug_file = open(Path(self.debug_goal_lane_dir) / f"{self.route_tag}.csv", "w") + self._goal_lane_debug_file.write(GOAL_LANE_DEBUG_HEADER) + + lights = route["lights"] + light_map, bin_num_traffic = cb.map_lights_to_bin(lights, self.transform, self.town_bin) + dense_route_bin_xy = np.array( + [self.transform.loc_to_bin(x, y) for x, y in self.dense_plan_xy], np.float64 + ).reshape(-1, 2) + light_line_indices, light_lines = cb.light_stop_line_overrides( + lights, light_map, self.transform, self.town_bin, route_xy=dense_route_bin_xy + ) + if len(light_line_indices): + moved_centers = 0.5 * (light_lines[:, 0:2] + light_lines[:, 3:5]) + max_shift_m = float(np.hypot(*(moved_centers - self.stop_line_centers[light_line_indices]).T).max()) + self.env.set_traffic_light_lines(light_line_indices, light_lines) + self.stop_line_centers[light_line_indices] = moved_centers + print( + f"[puffer_agent] light stop lines: {len(light_line_indices)} on the route moved onto the leaderboard's " + f"red-light lines (max {max_shift_m:.1f} m)" + ) + # the shadow env runs on CARLA's trigger volumes: the bin's exported stop lines sit up to 9 m off and miss a few + self._stop_sign_lines, stop_sign_headings = cb.stop_sign_lines(route["stop_signs"], self.transform) + self.num_traffic = self.env.set_stop_signs(self._stop_sign_lines, stop_sign_headings) + print( + f"[puffer_agent] stop signs: {len(stop_sign_headings)} CARLA trigger volumes replace the bin's " + f"(traffic elements {bin_num_traffic} -> {self.num_traffic})" + ) + # One id -> bin-indices table, reused by both passes in + # _read_light_states, so they can't disagree on which bin element a + # given light maps to (a prior independent geometric lookup for the + # ego-governance pass picked the wrong element ~5% of the time). + self.light_bin_indices_by_id = {int(light["id"]): light_map[i] for i, light in enumerate(lights)} + self.last_light_states = np.zeros(self.num_traffic, np.int32) + + wheelbase, max_steer = vehicle_geometry_from_wheels(ego["wheels"]) + self.controller = TrackingController(wheelbase_m=wheelbase, max_steer_rad=max_steer, horizon_s=self.dt) + + if self.obs_html_dir: + from pufferlib.ocean.cosim.obs_replay import ObsReplayCapture + + Path(self.obs_html_dir).mkdir(parents=True, exist_ok=True) + self._obs_html = ObsReplayCapture( + self.env, self.policy, Path(self.obs_html_dir) / self.route_tag, max_steps=self.obs_html_max_steps + ) + + if self.telemetry_dir: + Path(self.telemetry_dir).mkdir(parents=True, exist_ok=True) + self.telemetry_file = open(Path(self.telemetry_dir) / f"{self.route_tag}.csv", "w") + self.telemetry_file.write(TELEMETRY_HEADER) + + print( + f"[puffer_agent] town={town} tick_dt={self.tick_dt} dt={self.dt} route_goals={len(self.route_goals)} " + f"max_speed_mps={self.env.max_speed_mps:.1f} (C_vel={self.env.max_speed_mps / self.env.base_max_speed_mps:.2f}) " + f"zero_partner_stopped_time={int(self.zero_partner_stopped_time)} " + f"vru_size_pad_m={self.vru_size_pad_m:g}" + ) + self.initialized = True + + # --- shadow-env sync (read-only w.r.t. the simulator) ------------------- + + def _read_states(self, partners): + idx, x, y, z, h, vx, vy, yaw_rate, accel_long, stopped_s = [], [], [], [], [], [], [], [], [], [] + stopped_by_id = {} + for j, row in enumerate(partners): + idx.append(1 + j) # agent 0 = ego; others fill 1..M + bx, by, bz, bh, bvx, bvy, byr, bal = self.transform.state_to_bin( + row[STATE_X], + row[STATE_Y], + row[STATE_Z], + row[STATE_YAW_DEG], + row[STATE_VX], + row[STATE_VY], + row[STATE_ANGULAR_VELOCITY_Z], + row[STATE_ACCEL_X], + row[STATE_ACCEL_Y], + row[PARTNER_BOX_Z], + row[PARTNER_EXTENT_Z], + ) + x.append(bx) + y.append(by) + z.append(bz) + h.append(bh) + vx.append(bvx) + vy.append(bvy) + yaw_rate.append(byr) + accel_long.append(bal) + # stopped time follows the CARLA actor, not the shadow slot (slots reshuffle by distance every step) + actor_id = int(row[PARTNER_ACTOR_ID]) + is_stopped = math.hypot(bvx, bvy) <= self._partner_stopped_speed_threshold + stopped_by_id[actor_id] = self._partner_stopped_s.get(actor_id, 0.0) + self.dt if is_stopped else 0.0 + stopped_s.append(stopped_by_id[actor_id]) + self._partner_stopped_s = stopped_by_id + if self.zero_partner_stopped_time: + stopped_s = [0.0] * len(partners) + return ( + np.array(idx, np.int32), + np.array(x, np.float32), + np.array(y, np.float32), + np.array(z, np.float32), + np.array(h, np.float32), + np.array(vx, np.float32), + np.array(vy, np.float32), + np.array(yaw_rate, np.float32), + np.array(accel_long, np.float32), + np.array(stopped_s, np.float32), + ) + + def _read_sizes(self, partners): + idx, length, width = [], [], [] + for j, row in enumerate(partners): + idx.append(1 + j) + is_vru = row[PARTNER_IS_WALKER] > 0.0 or row[PARTNER_IS_TWO_WHEELER] > 0.0 + vru_pad_m = self.vru_size_pad_m if is_vru else 0.0 + length.append(max(2.0 * row[PARTNER_EXTENT_X], MIN_PARTNER_SIZE_M) + vru_pad_m) + width.append(max(2.0 * row[PARTNER_EXTENT_Y], MIN_PARTNER_SIZE_M) + vru_pad_m) + return (np.array(idx, np.int32), np.array(length, np.float32), np.array(width, np.float32)) + + def _read_light_states(self, light_states, ego_light_id): + """Ground truth for every mapped bin traffic element, via the + precomputed light_bin_indices_by_id table (see init_route).""" + # unmapped light elements read OFF (a training state), never UNKNOWN + states = np.full(self.num_traffic, cb.TRAFFIC_LIGHT_STATE_OFF, np.int32) + state_by_id = {} + for light_id, state_name in light_states: + state = cb.light_state_from_name(state_name) + state_by_id[int(light_id)] = state + for j in self.light_bin_indices_by_id[int(light_id)]: + if 0 <= j < self.num_traffic: + states[j] = state + ego_light_id = int(ego_light_id) + if ego_light_id in self.light_bin_indices_by_id: + state = state_by_id[ego_light_id] + for j in self.light_bin_indices_by_id[ego_light_id]: + if 0 <= j < self.num_traffic: + states[j] = state + return states + + def _sync_ego(self, ego_state, zero_velocity=False): + """Overwrite agent 0 (ego) from the simulator's ground-truth pose/size.""" + bin_state = self.transform.state_to_bin( + ego_state[STATE_X], + ego_state[STATE_Y], + ego_state[STATE_Z], + ego_state[STATE_YAW_DEG], + ego_state[STATE_VX], + ego_state[STATE_VY], + ego_state[STATE_ANGULAR_VELOCITY_Z], + ego_state[STATE_ACCEL_X], + ego_state[STATE_ACCEL_Y], + self.ego_box_z, + self.ego_extent[2], + ) + if zero_velocity: + x, y, z, heading, *_ = bin_state + bin_state = (x, y, z, heading, 0.0, 0.0, 0.0, 0.0) + self.env.set_agent_states(np.array([0], np.int32), *[np.array([v], np.float32) for v in bin_state]) + self.env.set_agent_sizes( + np.array([0], np.int32), + np.array([2.0 * self.ego_extent[0]], np.float32), + np.array([2.0 * self.ego_extent[1]], np.float32), + ) + + def _sync(self, snapshot): + """Overwrite the shadow env's background agents/lights (+ ego, only in + dynamics_source='carla' mode) from the simulator's ground truth and return the + recomputed observation array (num_agents, obs_dim). Row 0 = ego.""" + ego_state = [float(v) for v in snapshot["ego"]] + self._ego_speed_carla = ego_speed_from_state(ego_state) + if self.dynamics_source == "carla": + self._sync_ego(ego_state) + ego = self.env.get_global_agent_state() + ebx, eby = float(ego["x"][0]), float(ego["y"][0]) + + partners = np.asarray(snapshot["partners"], dtype=np.float64).reshape(-1, PARTNER_ROW_COLUMNS) + if len(partners): + states = self._read_states(partners) + sizes = self._read_sizes(partners) + self.env.set_agent_states(*states) + self.env.set_agent_sizes(*sizes) + is_walker = (partners[:, PARTNER_IS_WALKER] > 0.0).astype(np.float32) + self._last_partners = (states[1], states[2], states[4], states[5], states[6], sizes[1], sizes[2], is_walker) + else: + self._last_partners = None + n_used = 1 + len(partners) + if n_used < self.num_agents: + sp = np.arange(n_used, self.num_agents, dtype=np.int32) + zf = np.full(len(sp), FAR_AWAY, np.float32) + zz = np.zeros_like(zf) + self.env.set_agent_states(sp, zf, zf, zf, zz, zz, zz, zz, zz) + + self.last_light_states = self._read_light_states(snapshot["light_states"], snapshot["ego_light_id"]) + self.env.set_traffic_light_states(self.last_light_states) + + self.goal_window.sync(ebx, eby, float(ego["heading"][0])) + if self._goal_lane_debug_file is not None: + self._write_goal_lane_debug_row(ebx, eby, self.route_goals[self.goal_window.current_index]) + return np.asarray(self.env.recompute_observations()) + + def _write_goal_lane_debug_row(self, ebx, eby, current_goal): + """One CSV row: the CURRENT goal (window slot 0, current_goal_idx after + the reset in c_set_agent_goals) vs the lane find_goal_lane snapped it + to -- that lane's direction (nearest segment to the goal, mirroring + find_goal_lane's own point-to-segment search) dotted against the route + direction used at snap time, and the snap distance. dot < 0 means the + snapped lane runs opposite the route (oncoming lane); a large snap + distance means no nearby lane passed find_goal_lane's gates at all.""" + gx, gy, _, gdir_x, gdir_y = current_goal + state = self.env.get_state() + scenario = state[0] if isinstance(state, list) else state + ego_agent = (scenario.get("agents") or [{}])[0] + lane_idx = int(ego_agent.get("goal_lane_idx", -1)) + dot, dist = float("nan"), float("nan") + if lane_idx >= 0 and self._road_data is not None and lane_idx < len(self._road_data): + road = self._road_data[lane_idx] + xs, ys = np.asarray(road["x"], dtype=np.float64), np.asarray(road["y"], dtype=np.float64) + if len(xs) >= 2: + seg_dx, seg_dy = xs[1:] - xs[:-1], ys[1:] - ys[:-1] + seg_len = np.hypot(seg_dx, seg_dy) + seg_len_safe = np.where(seg_len > 1e-6, seg_len, 1.0) + t = np.clip(((gx - xs[:-1]) * seg_dx + (gy - ys[:-1]) * seg_dy) / (seg_len_safe**2), 0.0, 1.0) + px, py = xs[:-1] + t * seg_dx, ys[:-1] + t * seg_dy + d = np.hypot(gx - px, gy - py) + k = int(np.argmin(d)) + dist = float(d[k]) + route_norm = math.hypot(gdir_x, gdir_y) + if seg_len[k] > 1e-6 and route_norm > 1e-6: + dot = float((gdir_x * seg_dx[k] + gdir_y * seg_dy[k]) / (seg_len[k] * route_norm)) + self._goal_lane_debug_file.write( + f"{self.step_count},{ebx:.2f},{eby:.2f},{gx:.2f},{gy:.2f},{gdir_x:.3f},{gdir_y:.3f}," + f"{lane_idx},{dot:.3f},{dist:.2f}\n" + ) + self._goal_lane_debug_file.flush() + + def _integrate(self, actions, ego_state): + """Step the shadow env ONCE (one policy step == one shadow tick). + + dynamics_source='carla': the tracking controller chases the ego's target kinematic state + with CARLA's physics -> StepResult.control. + + dynamics_source='pufferdrive': PufferDrive's own jerk/classic dynamics (the ones the policy + trained on) ARE the ego's motion -- no target to chase. StepResult.motion is the pre/post-step + pose the simulator teleports the ego along, where the next tick's observation reads it. + """ + if self.step_count >= MAX_POLICY_STEPS: + raise RuntimeError(f"route exceeded MAX_POLICY_STEPS={MAX_POLICY_STEPS} policy steps") + before = self.env.get_global_agent_state() + speed_before = self._ego_speed() + self.env.step(actions) + after = self.env.get_global_agent_state() + target_yaw_deg = self.transform.bin_heading_to_yaw(float(after["heading"][0])) + speed_after = self._ego_speed() + accel_after = float(np.asarray(self.env.observations)[0][4]) * self._accel_long_norm() + + if self.dynamics_source == "pufferdrive": + x0, y0 = self.transform.bin_to_loc(float(before["x"][0]), float(before["y"][0])) + x1, y1 = self.transform.bin_to_loc(float(after["x"][0]), float(after["y"][0])) + yaw0_deg = self.transform.bin_heading_to_yaw(float(before["heading"][0])) + z1 = self.transform.z_to_carla(float(after["z"][0])) + self.target = (speed_after, target_yaw_deg) # logged, not chased + return StepResult( + actions, self.target, motion=[x0, y0, yaw0_deg, speed_before, x1, y1, target_yaw_deg, speed_after, z1] + ) + + target_speed = speed_after + max(accel_after, 0.0) * self.speed_intent_extension_s + self.target = (target_speed, target_yaw_deg) + # Controller runs every tick against the latest CARLA state, chasing the target of this policy step. + control = self.controller.step( + self._ego_speed_carla, ego_state[STATE_YAW_DEG], target_speed, target_yaw_deg, self.tick_dt + ) + return StepResult(actions, self.target, control=control) + + def step(self, snapshot): + """Shadow env <- simulator ground truth, policy, one dt of dynamics -> StepResult.""" + self.step_count += 1 + obs = self._sync(snapshot) + actions, aux = self._policy_actions(obs) + if self.obs_dump_dir: + self._obs_dump.append((self.step_count, obs[0].astype(np.float32), actions[0].astype(np.float32))) + if self.world_log_dir: + self._record_world_state(obs) + if self._obs_html is not None: + self._obs_html.capture(obs, actions, aux, aux.get("action_index")) + result = self._integrate(actions, [float(v) for v in snapshot["ego"]]) + if self.telemetry_file is not None or self.infraction_flags: + result.pd_flags = self._ego_infractions() + return result + + def _ego_speed(self): + return float(np.asarray(self.env.observations)[0][0]) * self._max_speed() + + def _max_speed(self): + return self.env.obs_norm_speed_mps + + def _accel_long_norm(self): + from pufferlib.ocean.drive import binding + + return binding.ACCEL_LONG_NORM + + def _ego_infractions(self): + """{'collision': f, 'offroad': f, 'red_light': f, 'stop_sign': f} from the shadow + pufferdrive env's own model of the ego (compute_metrics, refreshed by + the last integrate()) -- what the policy would have caused according + to PufferDrive's own dynamics, not necessarily what the real CARLA + vehicle did (the simulator side's own detectors tell that).""" + state = self.env.get_state() + scenario = state[0] if isinstance(state, list) else state + agents = scenario.get("agents") or [] + if not agents: + return {name: 0.0 for name in EGO_INFRACTION_METRICS} + metrics = agents[0].get("metrics_array") or [] + return { + name: float(metrics[idx]) if idx < len(metrics) else 0.0 for name, idx in EGO_INFRACTION_METRICS.items() + } + + # --- policy + capture --------------------------------------------------- + + def _adapt_obs_for_policy(self, obs): + """Strip the trailing zero-padded boundary columns when the checkpoint + predates them (see _load_policy_and_env). No-op otherwise.""" + if self._boundary_feat_ckpt is None: + return obs + env = self.env + n_slots = env.obs_slots_boundary_kept + feat_env, feat_ckpt = env.boundary_features, self._boundary_feat_ckpt + b0 = ( + env.ego_features + + env.num_reward_coefs + + env.goal_dim + + env.obs_slots_partners_n * env.partner_features + + env.obs_slots_lane_kept * env.lane_features + ) + b1 = b0 + n_slots * feat_env + obs = np.asarray(obs) + boundary = obs[:, b0:b1].reshape(obs.shape[0], n_slots, feat_env)[:, :, :feat_ckpt] + return np.concatenate([obs[:, :b0], boundary.reshape(obs.shape[0], -1), obs[:, b1:]], axis=1) + + def _policy_actions(self, obs): + """-> (actions (num_agents, act_dim) int32, aux dict). aux carries + per-agent value/entropy/action-probs/pool for the obs_html capture""" + import torch + + import pufferlib.pytorch + import pufferlib.spaces + + obs = self._adapt_obs_for_policy(obs) + # A discrete policy head on a continuous env needs the bin->continuous + # mapping (pufferl.py feeds cont_action to the env, never the bin indices). + env_continuous = isinstance(self.env.single_action_space, pufferlib.spaces.Box) + action_selection = ( + pufferlib.pytorch.ACTION_SELECT_MEAN + if env_continuous and not self.policy.is_continuous + else pufferlib.pytorch.ACTION_SELECT_MODE + ) + with torch.no_grad(): + logits, value = self.policy.forward_eval(torch.as_tensor(obs).to(self.device)) + action, _, entropy, cont_action = pufferlib.pytorch.sample_logits( + logits, action_selection=action_selection, env_continuous=env_continuous, policy=self.policy + ) + env_action = cont_action if env_continuous and cont_action is not None else action + actions = env_action.cpu().numpy().reshape(1, -1) + if not env_continuous: + actions = actions.astype(np.int32) + aux = {} + if self._obs_html is not None: + aux = self._obs_html.policy_outputs(torch.as_tensor(obs).to(self.device), logits, value, entropy) + aux["action_index"] = action.cpu().numpy().reshape(-1) if not self.policy.is_continuous else None + return actions, aux + + def _record_world_state(self, obs): + """Bin-frame snapshot at sync time (pre-step ego, partners as streamed, light states).""" + ego = self.env.get_global_agent_state() + accel_long = float(obs[0][4]) * self._accel_long_norm() + step_idx = len(self._world_log["ego"]) + self._world_log["ego"].append( + ( + self.step_count, + float(ego["x"][0]), + float(ego["y"][0]), + float(ego["heading"][0]), + self._ego_speed(), + accel_long, + float(self.goal_window.current_index), + ) + ) + if self._last_partners is not None: + x, y, h, vx, vy, length, width, is_walker = self._last_partners + speed = np.hypot(vx, vy) + rows = np.column_stack([np.full(len(x), step_idx, np.float32), x, y, h, length, width, speed, is_walker]) + self._world_log["partners"].append(rows.astype(np.float32)) + self._world_log["lights"].append(self.last_light_states.astype(np.int8)) + + def _write_world_log(self): + Path(self.world_log_dir).mkdir(parents=True, exist_ok=True) + dense = np.array([self.transform.loc_to_bin(x, y) for x, y in self.dense_plan_xy], np.float32).reshape(-1, 2) + partners = self._world_log["partners"] + np.savez_compressed( + Path(self.world_log_dir) / f"{self.route_tag}.npz", + ego=np.array(self._world_log["ego"], np.float32), + partners=np.concatenate(partners) if partners else np.zeros((0, 8), np.float32), + lights=np.stack(self._world_log["lights"]) if self._world_log["lights"] else np.zeros((0, 0), np.int8), + route_goals=self.route_goals, + stop_signs=self._stop_sign_lines, + dense_route=dense, + meta=json.dumps( + { + "town": Path(self.town_bin).stem.split("__")[-1], + "town_bin": self.town_bin, + "dt": self.dt, + "tick_dt": self.tick_dt, + "offset": [self.transform.tx, self.transform.ty], + "z_offset": self.transform.tz, + } + ), + ) + + def telemetry(self, result, carla_flags): + """One CSV row per policy step (when COSIM_TELEMETRY is set): what the loop commanded vs + achieved, the nearest mapped light's state and both infraction sources (pd_* = shadow pufferdrive + env, carla_* = the simulator side's own ground truth, None -> all 0).""" + if self.telemetry_file is None: + return + pd_flags = result.pd_flags + carla_flags = carla_flags or {name: 0.0 for name in EGO_INFRACTION_METRICS} + ego = self.env.get_global_agent_state() + ex, ey = float(ego["x"][0]), float(ego["y"][0]) + cur = min(self.goal_window.current_index, len(self.route_goals) - 1) + goal_dist = float(np.hypot(self.route_goals[cur, 0] - ex, self.route_goals[cur, 1] - ey)) + near_dist, near_state = -1.0, -1 + if len(self.stop_line_centers): + d2 = (self.stop_line_centers[:, 0] - ex) ** 2 + (self.stop_line_centers[:, 1] - ey) ** 2 + j = int(d2.argmin()) + near_dist, near_state = float(np.sqrt(d2[j])), int(self.last_light_states[j]) + current_speed = ( + float(np.asarray(self.env.observations)[0, 0]) * self._max_speed() + if self.dynamics_source == "pufferdrive" + else self._ego_speed_carla + ) + self.telemetry_file.write( + f"{self.step_count},{current_speed:.3f},{self.target[0]:.3f}," + f"{float(result.actions[0, 0])},{self.goal_window.current_index},{goal_dist:.1f},{near_dist:.1f},{near_state}," + f"{pd_flags['collision']:.0f},{pd_flags['offroad']:.0f},{pd_flags['red_light']:.0f}," + f"{pd_flags['stop_sign']:.0f}," + f"{carla_flags['collision']:.0f},{carla_flags['offroad']:.0f},{carla_flags['red_light']:.0f}," + f"{carla_flags['stop_sign']:.0f}\n" + ) + + def finish(self): + """Route done: flush every per-route output and release the shadow env.""" + if not self.initialized: + return + print( + f"[puffer_agent] route done: goals {self.goal_window.current_index + 1}/{len(self.route_goals)}, " + f"tracking {self.controller.stats()}" + ) + if self.telemetry_file is not None: + self.telemetry_file.close() + if self._goal_lane_debug_file is not None: + self._goal_lane_debug_file.close() + if self._obs_html is not None: + written = self._obs_html.write(render_html=self.obs_html_render) + print(f"[puffer_agent] wrote obs replay ({len(self._obs_html)} frames) -> {written}") + if self.world_log_dir and self._world_log["ego"]: + self._write_world_log() + if self.obs_dump_dir and self._obs_dump: + Path(self.obs_dump_dir).mkdir(parents=True, exist_ok=True) + np.savez_compressed( + Path(self.obs_dump_dir) / f"{self.route_tag}.npz", + step=np.array([s for s, _, _ in self._obs_dump], np.int32), + obs=np.stack([o for _, o, _ in self._obs_dump]), + action=np.stack([a for _, _, a in self._obs_dump]), + ) + self.env.close() + self.initialized = False diff --git a/pufferlib/ocean/cosim/carla_bridge.py b/pufferlib/ocean/cosim/carla_bridge.py index 8931f3dae2..7ab3b37fc2 100644 --- a/pufferlib/ocean/cosim/carla_bridge.py +++ b/pufferlib/ocean/cosim/carla_bridge.py @@ -93,17 +93,33 @@ def town_offset(bin_path: str): OFFSET_CALIBRATION_MAX_Z_OFFSET_M = 5.0 # exported bins sit 0-1.7 m below CARLA's ground; more means a frame mix-up -def calibrate_town_offset(carla_map, transform, town_bin): +DRIVING_WAYPOINT_COLUMNS = 5 # x, y, z, right_x, right_y in the CARLA frame + + +def driving_waypoint_samples(carla_map): + """(N, 5) [x, y, z, right_x, right_y] of CARLA's non-junction driving waypoints, calibrate_town_offset's input.""" + rows = [] + for wp in carla_map.generate_waypoints(OFFSET_CALIBRATION_SPACING_M): + if wp.is_junction or str(wp.lane_type) != "Driving": + continue + loc, right = wp.transform.location, wp.transform.get_right_vector() + rows.append((loc.x, loc.y, loc.z, right.x, right.y)) + return np.array(rows, dtype=np.float64).reshape(-1, DRIVING_WAYPOINT_COLUMNS) + + +def calibrate_town_offset(driving_waypoints, transform, town_bin): """(tx, ty) refined so the bin's lane polylines coincide with CARLA's lane centres. - The exported bins sit a constant fraction of a metre off the CARLA frame (measured 2026-09-23 with - the stored offsets: Town01 0.27 m, Town02 0.6 m), and the ego plus every streamed actor inherited it, - driving 0.44 m right of CARLA's lane centre on Town02 and clipping curbs the shadow env could not see. - Least squares of the translation over the lane-normal residuals of CARLA's driving waypoints against - the nearest bin lane segment, two passes (re-matched after the first shift), mismatches beyond - OFFSET_CALIBRATION_MATCH_MAX_M dropped. The bins also sit a constant per-town height below CARLA's - ground (Town03 0.77 m, Town04 1.24 m, Town05 1.64 m, Town06 0), so the same matches give the z offset as - the median of bin lane z minus waypoint z. Returns (offset, z_offset_m, residual_before_m, residual_after_m).""" + driving_waypoints: (N, 5) CARLA-frame [x, y, z, right_x, right_y] samples of the driving lanes + (driving_waypoint_samples). The exported bins sit a constant fraction of a metre off the CARLA frame + (measured 2026-09-23 with the stored offsets: Town01 0.27 m, Town02 0.6 m), and the ego plus every + streamed actor inherited it, driving 0.44 m right of CARLA's lane centre on Town02 and clipping curbs + the shadow env could not see. Least squares of the translation over the lane-normal residuals of the + waypoints against the nearest bin lane segment, two passes (re-matched after the first shift), + mismatches beyond OFFSET_CALIBRATION_MATCH_MAX_M dropped. The bins also sit a constant per-town height + below CARLA's ground (Town03 0.77 m, Town04 1.24 m, Town05 1.64 m, Town06 0), so the same matches give + the z offset as the median of bin lane z minus waypoint z. + Returns (offset, z_offset_m, residual_before_m, residual_after_m).""" data = _mbin.read_bin(Path(town_bin)) seg_start, seg_end, seg_z_start, seg_z_end = [], [], [], [] for road in data["roads"]: @@ -121,18 +137,11 @@ def calibrate_town_offset(carla_map, transform, town_bin): seg_z_end = np.concatenate(seg_z_end) seg_dir = seg_end - seg_start seg_len_sq = np.maximum((seg_dir**2).sum(1), 1e-9) - points, normals, waypoint_z = [], [], [] - for wp in carla_map.generate_waypoints(OFFSET_CALIBRATION_SPACING_M): - if wp.is_junction or str(wp.lane_type) != "Driving": - continue - loc, right = wp.transform.location, wp.transform.get_right_vector() - points.append(transform.loc_to_bin(loc.x, loc.y)) - normals.append((right.x, -right.y)) # y flips into the bin frame - waypoint_z.append(loc.z) - points = np.array(points, dtype=np.float64).reshape(-1, 2) - normals = np.array(normals, dtype=np.float64).reshape(-1, 2) + samples = np.asarray(driving_waypoints, dtype=np.float64).reshape(-1, DRIVING_WAYPOINT_COLUMNS) + points = np.array([transform.loc_to_bin(x, y) for x, y in samples[:, :2]], dtype=np.float64).reshape(-1, 2) + normals = np.column_stack([samples[:, 3], -samples[:, 4]]) # y flips into the bin frame normals /= np.maximum(np.hypot(normals[:, 0], normals[:, 1]), 1e-9)[:, None] - waypoint_z = np.array(waypoint_z, dtype=np.float64) + waypoint_z = samples[:, 2] def normal_residuals(shift): residuals = np.empty(len(points)) @@ -204,24 +213,41 @@ def yaw_to_bin(self, yaw_deg): def vel_to_bin(self, vx, vy): return vx, -vy + def state_to_bin(self, x, y, z, yaw_deg, vx, vy, angular_velocity_z_deg_s, accel_x, accel_y, box_z, box_extent_z): + """(x, y, z, heading, vx, vy, yaw_rate, accel_long) in the bin frame from a CARLA-frame actor state. + Angular velocity (deg/s) and acceleration (m/s^2, world frame) must come from CARLA's own physics, not + finite-differenced from a previous PufferDrive state: env.step() may have already overwritten this + actor's previous state with a throwaway dummy-action rollout. box_z/box_extent_z: bounding-box pivot + offset and half height (vehicles pivot at the ground, walkers at the capsule centre); bin z = ground.""" + bx, by = self.loc_to_bin(x, y) + heading = self.yaw_to_bin(yaw_deg) + yaw_rate = -math.radians( + angular_velocity_z_deg_s + ) # mirrored y flips rotation sense, like yaw_to_bin's negation + accel_long = accel_x * math.cos(heading) + (-accel_y) * math.sin(heading) # mirror y, like vel_to_bin + ground_z = z + box_z - box_extent_z + return (bx, by, self.z_to_bin(ground_z), heading, vx, -vy, yaw_rate, accel_long) + def actor_state_to_bin(self, actor): - """Return (x, y, z, heading, vx, vy, yaw_rate, accel_long) in the bin frame for a CARLA actor. - yaw_rate/accel_long come from CARLA's own physics (get_angular_velocity/get_acceleration), not - finite-differenced from a previous PufferDrive state: env.step() may have already overwritten - this actor's previous state with a throwaway dummy-action rollout by the time PufferDrive would - otherwise read "previous" state itself.""" + """state_to_bin of a live CARLA actor.""" tf = actor.get_transform() v = actor.get_velocity() - av = actor.get_angular_velocity() # deg/s, world frame (CARLA's rotation convention) - acc = actor.get_acceleration() # m/s^2, world frame - bx, by = self.loc_to_bin(tf.location.x, tf.location.y) - heading = self.yaw_to_bin(tf.rotation.yaw) - yaw_rate = -math.radians(av.z) # mirrored y flips rotation sense, like yaw_to_bin's negation - accel_x, accel_y = acc.x, -acc.y # mirror y, like vel_to_bin - accel_long = accel_x * math.cos(heading) + accel_y * math.sin(heading) - box = actor.bounding_box # pivot differs per asset (vehicles: ground, walkers: capsule centre); bin z = ground - ground_z = tf.location.z + box.location.z - box.extent.z - return (bx, by, self.z_to_bin(ground_z), heading, v.x, -v.y, yaw_rate, accel_long) + av = actor.get_angular_velocity() + acc = actor.get_acceleration() + box = actor.bounding_box + return self.state_to_bin( + tf.location.x, + tf.location.y, + tf.location.z, + tf.rotation.yaw, + v.x, + v.y, + av.z, + acc.x, + acc.y, + box.location.z, + box.extent.z, + ) # --- bin frame -> CARLA (to teleport the ego back into CARLA) --- def bin_to_loc(self, bx, by): @@ -234,23 +260,22 @@ def bin_heading_to_yaw(self, heading_rad): return -math.degrees(heading_rad) -# CARLA traffic-light state -> PufferDrive enum (datatypes.h:61-67) +# CARLA traffic-light state name -> PufferDrive enum (datatypes.h:61-67) # UNKNOWN=0 RED=1 YELLOW=2 GREEN=3 OFF=4 TRAFFIC_LIGHT_STATE_OFF = 4 +LIGHT_STATE_BY_NAME = {"Red": 1, "Yellow": 2, "Green": 3, "Off": 4, "Unknown": 0} -def carla_light_to_puffer(state) -> int: - import carla +def light_state_from_name(state_name) -> int: + """carla.TrafficLightState. (as sent over a co-sim protocol) -> PufferDrive enum.""" + return LIGHT_STATE_BY_NAME.get(state_name, 0) + - return { - carla.TrafficLightState.Red: 1, - carla.TrafficLightState.Yellow: 2, - carla.TrafficLightState.Green: 3, - carla.TrafficLightState.Off: 4, - carla.TrafficLightState.Unknown: 0, - }.get(state, 0) +def carla_light_to_puffer(state) -> int: + return light_state_from_name(state.name) +LIGHT_PROBE_STEPS_M = (0.0, 3.0, 6.0, 10.0, 15.0) # walk-back (all) / forward (all but 0) probes of map_lights_to_bin LIGHT_LANE_MATCH_MAX_DIST_M = 12.0 # measured Town01 controlled-lane stub distance (see below) LIGHT_LANE_FORWARD_MATCH_MAX_DIST_M = ( 3.0 # forward fallback lands ON the connector; a lane width away is the neighbour's @@ -258,6 +283,13 @@ def carla_light_to_puffer(state) -> int: LIGHT_LANE_ALIGN_COS_MIN = ( 0.5 # forward fallback: the lane must run along the waypoint's heading (no crossing connectors) ) +LIGHT_STOP_LINE_TOLERANCE_M = 1.0 # exported bin lines within this of the light's stop waypoint stay (median 0.5 m) +ROUTE_STOP_LINE_MARGIN_M = 1.0 # a route crossing a stop line extended by this at both ends drives that line's lane +ROUTE_STOP_LINE_ALIGN_COS_MIN = 0.5 # ... in the element's travel direction (opposing traffic crosses the same line) +TRIGGER_SAMPLE_SPAN_FRACTION = 0.9 # RunningRedLightTest.get_traffic_light_waypoints: avoids adjacent lanes +TRIGGER_SAMPLE_STEP_M = 1.0 +JUNCTION_ADVANCE_STEP_M = 0.5 +MAX_JUNCTION_ADVANCE_STEPS = 400 # 200 m: a trigger volume farther from its junction is a map error JUNCTION_CLUSTER_M = 60.0 # stop lines this close are treated as one junction (outlier repair, see below) # An assigned light farther from its junction cluster-mates' lights than this # multiple of the cluster-mates' own spread is a wrong-junction match. @@ -288,7 +320,7 @@ def find(k): return clusters -def _drop_junction_outliers(mapping, lights, transform, stop_centers): +def _drop_junction_outliers(mapping, light_geometry, transform, stop_centers): """(light_idx, element_idx) pairs to remove from `mapping`: elements whose assigned light sits far outside where its junction cluster-mates' lights actually are -- see map_lights_to_bin's docstring.""" @@ -309,8 +341,8 @@ def _drop_junction_outliers(mapping, lights, transform, stop_centers): continue light_pos = {} for j in elements: - loc = lights[element_light[j]].get_location() - light_pos[j] = np.array(transform.loc_to_bin(loc.x, loc.y)) + light = light_geometry[element_light[j]] + light_pos[j] = np.array(transform.loc_to_bin(light["x"], light["y"])) for j in elements: others = np.array([light_pos[j2] for j2 in elements if j2 != j]) centroid = others.mean(axis=0) @@ -324,8 +356,92 @@ def _drop_junction_outliers(mapping, lights, transform, stop_centers): return to_drop -def map_lights_to_bin(lights, transform, town_bin): - """mapping[i] = list of bin traffic-element indices controlled by lights[i]. +def light_stop_waypoints(light, carla_map): + """Stop waypoints of a traffic light from its trigger volume: one waypoint per lane the volume spans, + advanced to the junction entry, the line RunningRedLightTest scores (its get_traffic_light_waypoints). + Not TrafficLight.get_stop_waypoints: in Town01/02 those sit 10-14 m before the junction entry.""" + import carla + + base_transform = light.get_transform() + trigger = light.trigger_volume + lane_waypoints = [] + for x in np.arange( + -TRIGGER_SAMPLE_SPAN_FRACTION * trigger.extent.x, + TRIGGER_SAMPLE_SPAN_FRACTION * trigger.extent.x, + TRIGGER_SAMPLE_STEP_M, + ): + wp = carla_map.get_waypoint(base_transform.transform(trigger.location + carla.Location(x=float(x)))) + if not lane_waypoints or lane_waypoints[-1].road_id != wp.road_id or lane_waypoints[-1].lane_id != wp.lane_id: + lane_waypoints.append(wp) + stop_waypoints = [] + for wp in lane_waypoints: + for _ in range(MAX_JUNCTION_ADVANCE_STEPS): + if wp.is_junction: + break + ahead = wp.next(JUNCTION_ADVANCE_STEP_M) + if not ahead or ahead[0].is_junction: + break + wp = ahead[0] + stop_waypoints.append(wp) + return stop_waypoints + + +def light_geometry_from_carla(lights, carla_map): + """Plain-data geometry of live CARLA traffic lights for map_lights_to_bin (CARLA frame). + + One dict per light: id, x/y of the light actor, trigger_x/trigger_y of its trigger volume and one + entry per stop waypoint (light_stop_waypoints) with x, y, yaw_deg, lane_width, + `backward` = [x, y] of the first waypoint LIGHT_PROBE_STEPS_M behind it (None where the lane ends) + and `forward` = [x, y, yaw_deg] of the first waypoint each non-zero step ahead of it.""" + geometry = [] + for light in lights: + location = light.get_location() + trigger = light.get_transform().transform(light.trigger_volume.location) + stop_waypoints = [] + for wp in light_stop_waypoints(light, carla_map): + backward, forward = [], [] + for step_m in LIGHT_PROBE_STEPS_M: + probes = [wp] if step_m == 0.0 else wp.previous(step_m) + backward.append( + None if not probes else [probes[0].transform.location.x, probes[0].transform.location.y] + ) + for step_m in LIGHT_PROBE_STEPS_M[1:]: + probes = wp.next(step_m) or [] + forward.append( + None + if not probes + else [ + probes[0].transform.location.x, + probes[0].transform.location.y, + probes[0].transform.rotation.yaw, + ] + ) + stop_waypoints.append( + { + "x": wp.transform.location.x, + "y": wp.transform.location.y, + "yaw_deg": wp.transform.rotation.yaw, + "lane_width": wp.lane_width, + "backward": backward, + "forward": forward, + } + ) + geometry.append( + { + "id": light.id, + "x": location.x, + "y": location.y, + "trigger_x": trigger.x, + "trigger_y": trigger.y, + "stop_waypoints": stop_waypoints, + } + ) + return geometry + + +def map_lights_to_bin(light_geometry, transform, town_bin): + """mapping[i] = list of bin traffic-element indices controlled by light_geometry[i] + (light_geometry_from_carla, or the same dicts built by a CARLA 0.9.10 co-sim client). Semantic matching: each of the light's STOP WAYPOINTS is snapped to its bin LANE (nearest drivable-lane segment -- stop lines lie ON lane segments, so @@ -406,26 +522,26 @@ def controlling_elements_near(px, py, max_dist_m, heading=None): return controlling return [] - probe_steps_m = (0.0, 3.0, 6.0, 10.0, 15.0) mapping = [] - for lt in lights: + for light in light_geometry: element_indices = [] - for wp in lt.get_stop_waypoints(): + for stop_waypoint in light["stop_waypoints"]: if not len(seg_lane): continue controlling = [] - for step_m in probe_steps_m: - probes = [wp] if step_m == 0.0 else wp.previous(step_m) - for probe in probes[:1]: - bx, by = transform.loc_to_bin(probe.transform.location.x, probe.transform.location.y) - controlling = controlling_elements_near(bx, by, LIGHT_LANE_MATCH_MAX_DIST_M) + for probe in stop_waypoint["backward"]: + if probe is None: + continue + bx, by = transform.loc_to_bin(probe[0], probe[1]) + controlling = controlling_elements_near(bx, by, LIGHT_LANE_MATCH_MAX_DIST_M) if controlling: break - for step_m in probe_steps_m[1:] if not controlling else (): # fallback: forward into the junction - for probe in (wp.next(step_m) or [])[:1]: - bx, by = transform.loc_to_bin(probe.transform.location.x, probe.transform.location.y) - heading = transform.yaw_to_bin(probe.transform.rotation.yaw) - controlling = controlling_elements_near(bx, by, LIGHT_LANE_FORWARD_MATCH_MAX_DIST_M, heading) + for probe in stop_waypoint["forward"] if not controlling else (): # fallback: forward into the junction + if probe is None: + continue + bx, by = transform.loc_to_bin(probe[0], probe[1]) + heading = transform.yaw_to_bin(probe[2]) + controlling = controlling_elements_near(bx, by, LIGHT_LANE_FORWARD_MATCH_MAX_DIST_M, heading) if controlling: break element_indices.extend(controlling) @@ -437,22 +553,92 @@ def controlling_elements_near(px, py, max_dist_m, heading=None): for t in data["traffic"] ] ).reshape(-1, 2) - for light_idx, element_idx in _drop_junction_outliers(mapping, lights, transform, stop_centers): + for light_idx, element_idx in _drop_junction_outliers(mapping, light_geometry, transform, stop_centers): mapping[light_idx].remove(element_idx) return mapping, len(data["traffic"]) -def stop_signs_from_carla(world, carla_map, transform): - """(lines (K, 6), headings (K,)) in the bin frame, one per CARLA `traffic.stop` actor: the trigger - volume's centre projected onto its driving lane, spanning the trigger's width across that lane, with - the lane's travel direction as heading and the lane waypoint's height (a few trigger volumes sit up - to 1.5 m above or below their road). Feed to Drive.set_stop_signs so the shadow env runs on the - volumes the leaderboard scores against instead of the bin's exported stop lines (which sit up to - 9 m away on a few Town03/05 approaches and miss four signs).""" +def route_crosses_stop_line(route_xy, stop_line, heading): + """True when the polyline route_xy (N, 2) crosses stop_line [x1, y1, z1, x2, y2, z2] (extended by + ROUTE_STOP_LINE_MARGIN_M at both ends) while running along `heading`, the element's travel direction.""" + route = np.asarray(route_xy, dtype=np.float64).reshape(-1, 2) + if len(route) < 2: + return False + x1, y1, _, x2, y2, _ = stop_line + line_len = math.hypot(x2 - x1, y2 - y1) + if line_len <= 0.0: + return False + ux, uy = (x2 - x1) / line_len, (y2 - y1) / line_len + ax, ay = x1 - ROUTE_STOP_LINE_MARGIN_M * ux, y1 - ROUTE_STOP_LINE_MARGIN_M * uy + bx, by = x2 + ROUTE_STOP_LINE_MARGIN_M * ux, y2 + ROUTE_STOP_LINE_MARGIN_M * uy + p, q = route[:-1], route[1:] + d = q - p + seg_len = np.maximum(np.hypot(d[:, 0], d[:, 1]), 1e-9) + along = (d[:, 0] * math.cos(heading) + d[:, 1] * math.sin(heading)) / seg_len + line_side_p = (bx - ax) * (p[:, 1] - ay) - (by - ay) * (p[:, 0] - ax) + line_side_q = (bx - ax) * (q[:, 1] - ay) - (by - ay) * (q[:, 0] - ax) + seg_side_a = d[:, 0] * (ay - p[:, 1]) - d[:, 1] * (ax - p[:, 0]) + seg_side_b = d[:, 0] * (by - p[:, 1]) - d[:, 1] * (bx - p[:, 0]) + crossing = ((line_side_p < 0.0) != (line_side_q < 0.0)) & ((seg_side_a < 0.0) != (seg_side_b < 0.0)) + return bool(np.any(crossing & (along >= ROUTE_STOP_LINE_ALIGN_COS_MIN))) + + +def light_stop_line_overrides(light_geometry, mapping, transform, town_bin, route_xy=None): + """(element indices (K,), lines (K, 6)) for the bin light elements whose exported stop line sits more than + LIGHT_STOP_LINE_TOLERANCE_M along the lane from the light's stop waypoints (light_stop_waypoints: the + line RunningRedLightTest scores against; Town04: four approaches whose bin line is 8 m past it, so an + ego stopping at the bin line has already run the light). Each line is the bin's own translated along + the element's travel direction onto its most upstream waypoint (a line spanning several lanes then + precedes every lane's CARLA line); a light's waypoints go to the nearest of its elements. + route_xy (N, 2), bin frame: only lines the route crosses in their travel direction move (the ones the + leaderboard scores); the rest keep the exported position the policy trained with, since a moved + neighbouring line beside the ego's own red line is out of distribution (Town04 junctions 483/1249 sit + 3.8 m apart in CARLA's model and the policy ran that red). None moves every line.""" + import data_utils.mirror_map_bin as mbin + + bin_traffic = mbin.read_bin(Path(town_bin))["traffic"] + along_offsets_by_element = {} + for light, elements in zip(light_geometry, mapping): + if not elements: + continue + mids = { + j: ( + 0.5 * (bin_traffic[j]["stop_line"][0] + bin_traffic[j]["stop_line"][3]), + 0.5 * (bin_traffic[j]["stop_line"][1] + bin_traffic[j]["stop_line"][4]), + ) + for j in elements + } + for waypoint in light["stop_waypoints"]: + wx, wy = transform.loc_to_bin(waypoint["x"], waypoint["y"]) + j = min(elements, key=lambda e: math.hypot(mids[e][0] - wx, mids[e][1] - wy)) + heading = bin_traffic[j]["heading"] + along = (wx - mids[j][0]) * math.cos(heading) + (wy - mids[j][1]) * math.sin(heading) + along_offsets_by_element.setdefault(j, []).append(along) + indices, lines = [], [] + for j in sorted(along_offsets_by_element): + shift_m = min(along_offsets_by_element[j]) + if abs(shift_m) <= LIGHT_STOP_LINE_TOLERANCE_M: + continue + if route_xy is not None and not route_crosses_stop_line( + route_xy, bin_traffic[j]["stop_line"], bin_traffic[j]["heading"] + ): + continue + heading = bin_traffic[j]["heading"] + dx, dy = shift_m * math.cos(heading), shift_m * math.sin(heading) + x1, y1, z1, x2, y2, z2 = bin_traffic[j]["stop_line"] + indices.append(j) + lines.append([x1 + dx, y1 + dy, z1, x2 + dx, y2 + dy, z2]) + return np.array(indices, np.int32), np.array(lines, np.float32).reshape(-1, 6) + + +def stop_sign_geometry_from_carla(world, carla_map): + """Plain-data geometry of every CARLA `traffic.stop` actor for stop_sign_lines (CARLA frame): id, the + trigger volume's centre/extent/yaw_deg and `lane` = [x, y, z, yaw_deg] of the driving-lane waypoint + under the centre. Signs whose centre projects onto no driving lane are left out.""" import carla - lines, headings = [], [] + geometry = [] for actor in world.get_actors().filter("traffic.stop"): actor_transform = actor.get_transform() trigger = actor.trigger_volume @@ -461,14 +647,44 @@ def stop_signs_from_carla(world, carla_map, transform): if waypoint is None: continue forward = waypoint.transform.get_forward_vector() - lane_yaw = math.atan2(forward.y, forward.x) - box_yaw = math.radians(actor_transform.rotation.yaw) + lane_location = waypoint.transform.location + geometry.append( + { + "id": actor.id, + "center": [center.x, center.y, center.z], + "extent": [trigger.extent.x, trigger.extent.y, trigger.extent.z], + "yaw_deg": actor_transform.rotation.yaw, + "lane": [ + lane_location.x, + lane_location.y, + lane_location.z, + math.degrees(math.atan2(forward.y, forward.x)), + ], + } + ) + return geometry + + +def stop_sign_lines(stop_sign_geometry, transform): + """(lines (K, 6), headings (K,)) in the bin frame, one per stop sign of stop_sign_geometry_from_carla: the + trigger volume's centre projected onto its driving lane, spanning the trigger's width across that lane, + with the lane's travel direction as heading and the lane waypoint's height (a few trigger volumes sit up + to 1.5 m above or below their road). Feed to Drive.set_stop_signs so the shadow env runs on the volumes + the leaderboard scores against instead of the bin's exported stop lines (which sit up to 9 m away on a + few Town03/05 approaches and miss four signs).""" + lines, headings = [], [] + for sign in stop_sign_geometry: + center_x, center_y, _ = sign["center"] + extent_x, extent_y, _ = sign["extent"] + _, _, lane_z, lane_yaw_deg = sign["lane"] + lane_yaw = math.radians(lane_yaw_deg) + box_yaw = math.radians(sign["yaw_deg"]) box_x_along_lane = abs(math.cos(box_yaw - lane_yaw)) >= math.cos(math.pi / 4.0) - half_width = trigger.extent.y if box_x_along_lane else trigger.extent.x + half_width = extent_y if box_x_along_lane else extent_x across_x, across_y = -math.sin(lane_yaw), math.cos(lane_yaw) - left = transform.loc_to_bin(center.x - half_width * across_x, center.y - half_width * across_y) - right = transform.loc_to_bin(center.x + half_width * across_x, center.y + half_width * across_y) - line_z = transform.z_to_bin(waypoint.transform.location.z) + left = transform.loc_to_bin(center_x - half_width * across_x, center_y - half_width * across_y) + right = transform.loc_to_bin(center_x + half_width * across_x, center_y + half_width * across_y) + line_z = transform.z_to_bin(lane_z) lines.append([left[0], left[1], line_z, right[0], right[1], line_z]) - headings.append(transform.yaw_to_bin(math.degrees(lane_yaw))) + headings.append(transform.yaw_to_bin(lane_yaw_deg)) return np.array(lines, np.float32).reshape(-1, 6), np.array(headings, np.float32) diff --git a/pufferlib/ocean/cosim/carla_cosim.py b/pufferlib/ocean/cosim/carla_cosim.py index 1909f0b9a2..0fbb07fd4c 100644 --- a/pufferlib/ocean/cosim/carla_cosim.py +++ b/pufferlib/ocean/cosim/carla_cosim.py @@ -414,13 +414,13 @@ def main(): world, tm, ego, bg, lights = build_carla(client, town, route_wps, args.num_background, 0.1) cmap = world.get_map() transform = cb.CarlaTransform(town, offset=cb.town_offset(town_bin)) - offset, z_offset, _, _ = cb.calibrate_town_offset(cmap, transform, town_bin) + offset, z_offset, _, _ = cb.calibrate_town_offset(cb.driving_waypoint_samples(cmap), transform, town_bin) transform = cb.CarlaTransform(town, offset=offset, z_offset=z_offset) _bin_lanes = cb._bin_lane_points(town_bin) # bin lane points (== global frame) for diagnostics dense_route = densify_route(route_wps) # fine-sampled route for lane-centered goal placement route_goals = build_route_goals(dense_route, transform, cmap) # fixed 20-m lane-centered goal sequence goal_window = RouteGoalWindow(env, route_goals) - light_map, num_traffic = cb.map_lights_to_bin(lights, transform, town_bin) + light_map, num_traffic = cb.map_lights_to_bin(cb.light_geometry_from_carla(lights, cmap), transform, town_bin) print( f"[cosim] carla: ego + {len(bg)} background + {len(lights)} lights; " f"offset={transform.tx:.1f},{transform.ty:.1f} z={transform.tz:+.2f}" diff --git a/pufferlib/ocean/cosim/nuplan/README.md b/pufferlib/ocean/cosim/nuplan/README.md index a035df5e5e..4384784678 100644 --- a/pufferlib/ocean/cosim/nuplan/README.md +++ b/pufferlib/ocean/cosim/nuplan/README.md @@ -30,11 +30,12 @@ Zero changes to nuplan-devkit or CaRL's `carl_nuplan`. Hydra loads any planner v change only where the lane graph forces one), so multi-lane roads stop pulling the ego toward the road centre. All come from the challenge's route roadblock ids only; the planner - never reads the logged ego. Eval hack `pedestrian_min_size_m` (env - `PEDESTRIAN_MIN_SIZE_M`): pedestrian and bicycle partner boxes are grown to at least - this size in the policy's observation (training spawns nothing below 0.8 x 0.8 m, - nuPlan pedestrians are 0.4-0.8 m); nuPlan still scores the true boxes. Off by - default, script 9 sets 0.8 m. The shadow env consumes the goals of its + never reads the logged ego. Eval hack `vru_size_pad_m` (env `VRU_SIZE_PAD_M`): this + much is added to every pedestrian and bicycle partner box's length and width in the + policy's observation (training spawns nothing below 0.8 x 0.8 m, nuPlan pedestrians + are 0.4-0.8 m, bicycles ~0.5 m wide), so they read strictly larger than the smallest + training vehicle; nuPlan still scores the true boxes. Off by default, script 9 sets + 0.8 m. The shadow env consumes the goals of its `num_goals` window itself exactly like training (`goal_regen_mode: finite`, consumed slots zeroed in the obs); the planner pushes the next window only when the current one is exhausted or its current goal is clearly behind the diff --git a/pufferlib/ocean/cosim/nuplan/config/planner/pufferdrive_planner.yaml b/pufferlib/ocean/cosim/nuplan/config/planner/pufferdrive_planner.yaml index 6bef40984c..5e9dc697ab 100644 --- a/pufferlib/ocean/cosim/nuplan/config/planner/pufferdrive_planner.yaml +++ b/pufferlib/ocean/cosim/nuplan/config/planner/pufferdrive_planner.yaml @@ -21,8 +21,8 @@ pufferdrive_planner: startup_jerk_cap_seconds: 1.5 startup_accel_jerk_cap_mps3: 0.0 startup_brake_jerk_cap_mps3: 0.0 - # eval hack: grow pedestrian/bicycle partner boxes to at least this size [m] in the policy's obs (training min 0.8); 0 = off - pedestrian_min_size_m: 0.0 + # eval hack: grow pedestrian and bicycle partner boxes by this much [m] in length and width in the policy's obs (training min 0.8 x 0.8); 0 = off + vru_size_pad_m: 0.0 env_overrides: obs_slots_partners_n: 40 # GIGAFLOW eval value (Table A4); slots are max-pooled, so any checkpoint transfers stop_signs_enabled: false # nuPlan city bins export TURN_STOP zones as stop signs; keep them out of the obs diff --git a/pufferlib/ocean/cosim/nuplan/planner.py b/pufferlib/ocean/cosim/nuplan/planner.py index 512f590d45..dd23c1a4d3 100644 --- a/pufferlib/ocean/cosim/nuplan/planner.py +++ b/pufferlib/ocean/cosim/nuplan/planner.py @@ -93,7 +93,7 @@ def __init__( startup_jerk_cap_seconds: float = 1.5, startup_accel_jerk_cap_mps3: float = 0.0, startup_brake_jerk_cap_mps3: float = 0.0, - pedestrian_min_size_m: float = 0.0, + vru_size_pad_m: float = 0.0, env_overrides: Optional[Dict] = None, obs_html_dir: Optional[str] = None, obs_html_max_steps: int = 800, @@ -125,9 +125,10 @@ def __init__( during the start window [m/s³]; 0 = uncapped. nuPlan's Savitzky-Golay jerk extrapolates at the trajectory edges, so a ramp that starts at t=0 reads ~2x its raw jerk. :param startup_brake_jerk_cap_mps3: same cap for braking jerk; 0 = uncapped. - :param pedestrian_min_size_m: eval hack: pedestrian and bicycle partner boxes are grown to at - least this length and width [m] before the policy sees them (training spawns nothing below - 0.8 x 0.8 m, nuPlan pedestrians are 0.4-0.8 m); nuPlan keeps scoring the true boxes. 0 = off. + :param vru_size_pad_m: eval hack: pedestrian and bicycle partner boxes grow by this much in length + and width [m] before the policy sees them, so they read strictly larger than the smallest + training vehicle (0.8 x 0.8 m; nuPlan pedestrians are 0.4-0.8 m, bicycles ~0.5 m wide); nuPlan + keeps scoring the true boxes. 0 = off. :param env_overrides: Drive env kwargs overriding DEFAULT_ARCH. :param obs_html_dir: write the interactive pufferlib.viz observation replay per scenario (the exact obs the policy received, its outputs and encoder pool winners). @@ -157,9 +158,9 @@ def __init__( < 0 ): raise ValueError("startup_jerk_cap_seconds and the startup jerk caps must be >= 0") - self._pedestrian_min_size_m = float(pedestrian_min_size_m) - if self._pedestrian_min_size_m < 0: - raise ValueError("pedestrian_min_size_m must be >= 0") + self._vru_size_pad_m = float(vru_size_pad_m) + if self._vru_size_pad_m < 0: + raise ValueError("vru_size_pad_m must be >= 0") self._startup_jerk_cap_steps = 0 # resolved in _build from the scenario dt self._startup_action_bounds = None # (lo, hi) on the continuous jerk action while the cap is live self._env_overrides = env_overrides or {} @@ -509,8 +510,8 @@ def _sync(self, ego_state: EgoState, detections: DetectionsTracks, traffic_light objs = objs[: self._num_agents - 1] if objs: idx, x, y, z, hh, vx, vy, tp, ln, wd = nb.tracked_objects_to_arrays(objs, tf) - if self._pedestrian_min_size_m > 0: - ln, wd = nb.floor_vru_partner_sizes(tp, ln, wd, self._pedestrian_min_size_m) + if self._vru_size_pad_m > 0: + ln, wd = nb.pad_vru_partner_sizes(tp, ln, wd, self._vru_size_pad_m) env.set_agent_states(idx, x, y, z, hh, vx, vy, np.zeros_like(vx), np.zeros_like(vx)) env.set_agent_sizes(idx, ln, wd) surplus = np.arange(1 + len(objs), self._num_agents, dtype=np.int32) diff --git a/pufferlib/ocean/cosim/nuplan/run_nuplan_planner.sh b/pufferlib/ocean/cosim/nuplan/run_nuplan_planner.sh index bc21e10f92..6656e6bce2 100755 --- a/pufferlib/ocean/cosim/nuplan/run_nuplan_planner.sh +++ b/pufferlib/ocean/cosim/nuplan/run_nuplan_planner.sh @@ -37,10 +37,12 @@ # SLIDING_GOAL_WINDOW=true|false refill the goal window after every consumed goal (default false). # STARTUP_ACCEL_JERK_CAP / STARTUP_BRAKE_JERK_CAP [m/s^3] ego jerk caps for the first # STARTUP_JERK_CAP_SECONDS (default 1.5) of every scenario; unset/0 = off. -# PEDESTRIAN_MIN_SIZE_M [m] grow pedestrian/bicycle partner boxes to at least this size in the -# policy's observation (training never spawns below 0.8 m); unset/0 = off. +# VRU_SIZE_PAD_M [m] add this to pedestrian and bicycle partner boxes' length and width in the policy's +# observation (training never spawns below 0.8 m; 0.8 = strictly above the smallest training +# vehicle); unset/0 = off. # OBS_SLOTS_PARTNERS_N shadow-env partner obs slots (unset = planner yaml, 40). # EVAL_PERCEIVED_SIZE_MARGIN_M [m per side] ego box inflation in the policy's obs (unset = arch.py, 0.2). +# REWARD_OVERSPEED overspeed conditioning alpha fed to the policy (unset = arch.py, 0.1). # COSIM_OBS_HTML=all|failures|infractions|0 interactive observation replay (exact policy # input/outputs per step, pufferlib.viz HTML) -> $GROUP/obs_html. all: every # scenario. failures (default): only scenarios scoring below @@ -85,9 +87,10 @@ EXTRA_ARGS+=("planner.pufferdrive_planner.city_bin_dir=$CITY_BIN_DIR") [ -n "${STARTUP_JERK_CAP_SECONDS:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.startup_jerk_cap_seconds=$STARTUP_JERK_CAP_SECONDS") [ -n "${STARTUP_ACCEL_JERK_CAP:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.startup_accel_jerk_cap_mps3=$STARTUP_ACCEL_JERK_CAP") [ -n "${STARTUP_BRAKE_JERK_CAP:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.startup_brake_jerk_cap_mps3=$STARTUP_BRAKE_JERK_CAP") -[ -n "${PEDESTRIAN_MIN_SIZE_M:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.pedestrian_min_size_m=$PEDESTRIAN_MIN_SIZE_M") +[ -n "${VRU_SIZE_PAD_M:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.vru_size_pad_m=$VRU_SIZE_PAD_M") [ -n "${OBS_SLOTS_PARTNERS_N:-}" ] && EXTRA_ARGS+=("++planner.pufferdrive_planner.env_overrides.obs_slots_partners_n=$OBS_SLOTS_PARTNERS_N") [ -n "${EVAL_PERCEIVED_SIZE_MARGIN_M:-}" ] && EXTRA_ARGS+=("++planner.pufferdrive_planner.env_overrides.eval_perceived_size_margin_m=$EVAL_PERCEIVED_SIZE_MARGIN_M") +[ -n "${REWARD_OVERSPEED:-}" ] && EXTRA_ARGS+=("++planner.pufferdrive_planner.env_overrides.reward_overspeed=$REWARD_OVERSPEED") [ -n "${WORKER:-}" ] && EXTRA_ARGS+=("worker=$WORKER") [ -n "${THREADS_PER_NODE:-}" ] && EXTRA_ARGS+=("worker.threads_per_node=$THREADS_PER_NODE") [ -n "${LIMIT_TOTAL_SCENARIOS:-}" ] && EXTRA_ARGS+=("scenario_filter.limit_total_scenarios=$LIMIT_TOTAL_SCENARIOS") diff --git a/pufferlib/ocean/cosim/nuplan_bridge.py b/pufferlib/ocean/cosim/nuplan_bridge.py index 599ca2fdc7..b07288adb9 100644 --- a/pufferlib/ocean/cosim/nuplan_bridge.py +++ b/pufferlib/ocean/cosim/nuplan_bridge.py @@ -611,12 +611,12 @@ def partner_tracked_objects(tracked_objects, map_api, static_on_lane: dict): VRU_PARTNER_TYPE_IDS = (_NUPLAN_AGENT_TYPE["PEDESTRIAN"], _NUPLAN_AGENT_TYPE["BICYCLE"]) -def floor_vru_partner_sizes(types, lengths, widths, min_size_m: float): - """Pedestrian/bicycle boxes below min_size_m on either axis grow to min_size_m (training never - spawns agents under 0.8 x 0.8 m; nuPlan pedestrians are mostly 0.4-0.8 m). Other types untouched.""" +def pad_vru_partner_sizes(types, lengths, widths, pad_m: float): + """Pedestrian and bicycle boxes grow by pad_m in length and width (0.8 puts every nuPlan pedestrian + and bicycle strictly above the 0.8 x 0.8 m training minimum). Other types untouched.""" vru = np.isin(types, VRU_PARTNER_TYPE_IDS) - lengths = np.where(vru, np.maximum(lengths, min_size_m), lengths).astype(np.float32) - widths = np.where(vru, np.maximum(widths, min_size_m), widths).astype(np.float32) + lengths = np.where(vru, lengths + pad_m, lengths).astype(np.float32) + widths = np.where(vru, widths + pad_m, widths).astype(np.float32) return lengths, widths diff --git a/pufferlib/ocean/cuda_slot_encoder.py b/pufferlib/ocean/cuda_slot_encoder.py new file mode 100644 index 0000000000..565cf3e7b4 --- /dev/null +++ b/pufferlib/ocean/cuda_slot_encoder.py @@ -0,0 +1,988 @@ +"""CUDA (NVRTC) backend for the fused slot encoder: the same Linear -> LayerNorm -> ReLU -> Linear -> masked +max-pool forward and its backward as the Triton kernels in fused_slot_encoder.py, written with mma.sync tensor-core +instructions, smem-staged LayerNorm and cp.async weight streaming. Selected with PUFFER_SLOT_ENCODER_BACKEND=cuda. + +Measured on an RTX 5090 at 65536 rows (2026-10-15): forward 1.08-1.14x faster than Triton with activation stores +(training), 1.24-1.35x without (rollout); backward on par (0.98-1.00x, bandwidth-bound dataflow). +""" + +import functools + +import numpy as np +import torch + +try: + from cuda.bindings import nvrtc, driver as cu +except ImportError: # cuda-python < 12.6 kept the modules at the package root + try: + from cuda import nvrtc, cuda as cu + except ImportError as import_error: + raise ImportError( + "PUFFER_SLOT_ENCODER_BACKEND=cuda needs the cuda-bindings package (NVRTC + driver API); " + "install it or unset the variable to use the Triton kernels" + ) from import_error + +CUDA_SRC_FWD = r""" +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef long long int64_t; +#define NEG_INF __int_as_float(0xff800000) +#define BLOCK_B 64 +#define WARP_M 32 +#define WARP_N 64 +#define NWARPS_M 2 + +__device__ __forceinline__ uint32_t smem_u32(const void *p) { return (uint32_t) __cvta_generic_to_shared(p); } +__device__ __forceinline__ uint32_t f32_to_bf16_bits(float f) { uint16_t r; asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(r) : "f"(f)); return r; } +__device__ __forceinline__ float bf16_bits_to_f32(uint32_t b) { return __uint_as_float(b << 16); } +__device__ __forceinline__ float bf16_round(float f) { return bf16_bits_to_f32(f32_to_bf16_bits(f)); } +__device__ __forceinline__ uint32_t pack_bf16x2(float lo, float hi) { uint32_t r; asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(r) : "f"(hi), "f"(lo)); return r; } +__device__ __forceinline__ void mma_16816(float c[4], const uint32_t a[4], uint32_t b0, uint32_t b1) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} +__device__ __forceinline__ void ldmatrix_x4(uint32_t r[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(addr)); +} +__device__ __forceinline__ void cp_async16(uint32_t smem_addr, const void *gptr) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" :: "r"(smem_addr), "l"(gptr)); +} +__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;"); } +template __device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;" :: "n"(N)); } + +// D: encoder width (128/256), F_PAD: padded input features (16/32), STORE: write per-slot activations +template +__global__ void __launch_bounds__(32 * NWARPS_M * (D / WARP_N), 1) +fused_slot_encoder_fwd_cuda( + const float *__restrict__ x, const int64_t *__restrict__ counts, + const uint16_t *__restrict__ w1, const float *__restrict__ b1, const float *__restrict__ gamma, const float *__restrict__ beta, + const uint16_t *__restrict__ w2, const float *__restrict__ b2, + uint16_t *__restrict__ normed_out, uint16_t *__restrict__ pooled_out, uint8_t *__restrict__ argmax_out, + int B, int S, int F, int64_t stride_xb, int64_t stride_xs, int64_t stride_xf) +{ + constexpr int NWARPS_N = D / WARP_N; + constexpr int THREADS = 32 * NWARPS_M * NWARPS_N; + constexpr int LDX = F_PAD + 8; // bf16 elements per sX / sW1 row + constexpr int KCHUNK = 32; + constexpr int LDW2 = KCHUNK + 8; // bf16 per sW2 row + constexpr int NCHUNK = D / KCHUNK; + constexpr int LDA = D + 8; // bf16 per sAct row + constexpr int TPR = D / 32; // LayerNorm threads per row (32 columns each) + constexpr int ROWS_PER_PASS = THREADS / TPR; + constexpr int X_ELEMS = BLOCK_B * F_PAD / THREADS; + + extern __shared__ __align__(16) unsigned char smem_raw[]; + uint16_t *sX = (uint16_t *) smem_raw; // [BLOCK_B][LDX] + uint16_t *sW1 = sX + BLOCK_B * LDX; // [D][LDX] + uint16_t *sW2 = sW1 + D * LDX; // [2][D][LDW2] + uint16_t *sAct = sW2 + 2 * D * LDW2; // [BLOCK_B][LDA] + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int wm = warp / NWARPS_N; + const int wn = warp % NWARPS_N; + const int b0 = blockIdx.x * BLOCK_B; + + // W1 ([D][F_PAD] bf16) resident in smem, zero-padded columns + for (int i = tid; i < D * (F_PAD / 8); i += THREADS) { + int n = i / (F_PAD / 8), q = i % (F_PAD / 8); + *(uint4 *) &sW1[n * LDX + q * 8] = *(const uint4 *) &w1[n * F_PAD + q * 8]; + } + // counts for the 4 rows this lane owns in the mma C layout: (wm*32 + i*16 + lane/4 + {0, 8}) + int my_counts[4]; + #pragma unroll + for (int i = 0; i < 2; i++) { + #pragma unroll + for (int h = 0; h < 2; h++) { + int row = wm * WARP_M + i * 16 + (lane >> 2) + h * 8; + int b = b0 + row; + my_counts[i * 2 + h] = (b < B) ? (int) counts[b] : 0; + } + } + // W2 chunk streaming: chunk c = columns [c*32, +32) of W2[n][k]; 16-byte pieces + auto issue_w2_chunk = [&](int chunk, int stage) { + const uint16_t *src = w2 + chunk * KCHUNK; + uint16_t *dst = sW2 + stage * D * LDW2; + #pragma unroll + for (int i = 0; i < (D * 4) / THREADS; i++) { + int e = tid + i * THREADS; + int n = e >> 2, q = e & 3; + cp_async16(smem_u32(&dst[n * LDW2 + q * 8]), src + n * D + q * 8); + } + cp_async_commit(); + }; + issue_w2_chunk(0, 0); + issue_w2_chunk(1, 1); + + // first slot's x tile -> registers + float xreg[X_ELEMS]; + auto load_x = [&](int slot) { + #pragma unroll + for (int q = 0; q < X_ELEMS; q++) { + int e = tid + q * THREADS; + int row = e / F_PAD, f = e % F_PAD; + int b = b0 + row; + xreg[q] = (b < B && f < F) ? x[b * stride_xb + slot * stride_xs + f * stride_xf] : 0.0f; + } + }; + auto store_x = [&]() { + #pragma unroll + for (int q = 0; q < X_ELEMS; q++) { + int e = tid + q * THREADS; + int row = e / F_PAD, f = e % F_PAD; + sX[row * LDX + f] = (uint16_t) f32_to_bf16_bits(xreg[q]); + } + }; + load_x(0); + store_x(); + __syncthreads(); + + float rmax[2][8][4]; + uint32_t amax[2][8]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) { + amax[i][j] = 0u; + #pragma unroll + for (int e = 0; e < 4; e++) rmax[i][j][e] = NEG_INF; + } + + const uint32_t sAct_base = smem_u32(sAct); + const uint32_t sX_base = smem_u32(sX); + const uint32_t sW1_base = smem_u32(sW1); + const uint32_t sW2_base = smem_u32(sW2); + + for (int slot = 0; slot < S; slot++) { + if (slot + 1 < S) load_x(slot + 1); + // ---- GEMM1: pre_norm = bf16(x @ W1^T + b1) -> sAct (bf16) + { + float acc[2][8][4]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) + #pragma unroll + for (int e = 0; e < 4; e++) acc[i][j][e] = 0.0f; + #pragma unroll + for (int kk = 0; kk < F_PAD / 16; kk++) { + uint32_t afrag[2][4]; + #pragma unroll + for (int i = 0; i < 2; i++) { + int row = wm * WARP_M + i * 16 + (lane & 15); + int col = kk * 16 + (lane >> 4) * 8; + ldmatrix_x4(afrag[i], sX_base + (row * LDX + col) * 2); + } + #pragma unroll + for (int jj = 0; jj < 4; jj++) { + uint32_t bfrag[4]; + int m = lane >> 3, r = lane & 7; + int n = wn * WARP_N + jj * 16 + (m >> 1) * 8 + r; + int k = kk * 16 + (m & 1) * 8; + ldmatrix_x4(bfrag, sW1_base + (n * LDX + k) * 2); + #pragma unroll + for (int i = 0; i < 2; i++) { + mma_16816(acc[i][jj * 2], afrag[i], bfrag[0], bfrag[1]); + mma_16816(acc[i][jj * 2 + 1], afrag[i], bfrag[2], bfrag[3]); + } + } + } + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) { + int col = wn * WARP_N + j * 8 + (lane & 3) * 2; + float bb0 = b1[col], bb1 = b1[col + 1]; + #pragma unroll + for (int h = 0; h < 2; h++) { + int row = wm * WARP_M + i * 16 + (lane >> 2) + h * 8; + float v0 = bf16_round(acc[i][j][h * 2] + bb0); + float v1 = bf16_round(acc[i][j][h * 2 + 1] + bb1); + *(uint32_t *) &sAct[row * LDA + col] = pack_bf16x2(v0, v1); + } + } + } + __syncthreads(); + // ---- LayerNorm + ReLU in place (bf16), optional global store of normed_relu + #pragma unroll + for (int pass = 0; pass < BLOCK_B / ROWS_PER_PASS; pass++) { + int row = pass * ROWS_PER_PASS + tid / TPR; + int c = tid % TPR; + float v[32]; + #pragma unroll + for (int j = 0; j < 4; j++) { + uint4 raw = *(const uint4 *) &sAct[row * LDA + 8 * (c + TPR * j)]; + uint32_t w[4] = {raw.x, raw.y, raw.z, raw.w}; + #pragma unroll + for (int t = 0; t < 4; t++) { + v[j * 8 + t * 2] = bf16_bits_to_f32(w[t] & 0xFFFFu); + v[j * 8 + t * 2 + 1] = bf16_bits_to_f32(w[t] >> 16); + } + } + float s = 0.0f; + #pragma unroll + for (int t = 0; t < 32; t++) s += v[t]; + #pragma unroll + for (int off = TPR / 2; off > 0; off >>= 1) s += __shfl_xor_sync(0xffffffffu, s, off); + float mean = s / (float) D; + float ss = 0.0f; + #pragma unroll + for (int t = 0; t < 32; t++) { v[t] -= mean; ss += v[t] * v[t]; } + #pragma unroll + for (int off = TPR / 2; off > 0; off >>= 1) ss += __shfl_xor_sync(0xffffffffu, ss, off); + float rstd = rsqrtf(ss / (float) D + 1e-5f); + int b = b0 + row; + #pragma unroll + for (int j = 0; j < 4; j++) { + int col = 8 * (c + TPR * j); + uint32_t w[4]; + #pragma unroll + for (int t = 0; t < 4; t++) { + float g0 = gamma[col + t * 2], g1 = gamma[col + t * 2 + 1]; + float be0 = beta[col + t * 2], be1 = beta[col + t * 2 + 1]; + float n0 = fmaxf(v[j * 8 + t * 2] * rstd * g0 + be0, 0.0f); + float n1 = fmaxf(v[j * 8 + t * 2 + 1] * rstd * g1 + be1, 0.0f); + w[t] = pack_bf16x2(n0, n1); + } + uint4 packed = make_uint4(w[0], w[1], w[2], w[3]); + *(uint4 *) &sAct[row * LDA + col] = packed; + if (STORE && b < B) { + *(uint4 *) &normed_out[((int64_t) b * S + slot) * D + col] = packed; + } + } + } + __syncthreads(); + // ---- GEMM2: encoded = normed @ W2^T, W2 streamed in K chunks + float acc[2][8][4]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) + #pragma unroll + for (int e = 0; e < 4; e++) acc[i][j][e] = 0.0f; + for (int chunk = 0; chunk < NCHUNK; chunk++) { + cp_async_wait<1>(); + __syncthreads(); + const uint32_t stage_base = sW2_base + (chunk & 1) * D * LDW2 * 2; + #pragma unroll + for (int kk = 0; kk < KCHUNK / 16; kk++) { + uint32_t afrag[2][4]; + #pragma unroll + for (int i = 0; i < 2; i++) { + int row = wm * WARP_M + i * 16 + (lane & 15); + int col = chunk * KCHUNK + kk * 16 + (lane >> 4) * 8; + ldmatrix_x4(afrag[i], sAct_base + (row * LDA + col) * 2); + } + #pragma unroll + for (int jj = 0; jj < 4; jj++) { + uint32_t bfrag[4]; + int m = lane >> 3, r = lane & 7; + int n = wn * WARP_N + jj * 16 + (m >> 1) * 8 + r; + int k = kk * 16 + (m & 1) * 8; + ldmatrix_x4(bfrag, stage_base + (n * LDW2 + k) * 2); + #pragma unroll + for (int i = 0; i < 2; i++) { + mma_16816(acc[i][jj * 2], afrag[i], bfrag[0], bfrag[1]); + mma_16816(acc[i][jj * 2 + 1], afrag[i], bfrag[2], bfrag[3]); + } + } + } + __syncthreads(); + int next = chunk + 2; + if (next < NCHUNK) issue_w2_chunk(next, chunk & 1); + else issue_w2_chunk(next - NCHUNK, chunk & 1); // prefetch next slot's first chunks + } + // ---- epilogue: bias, bf16 round, running masked max; STORE keeps the first arg-max slot per feature + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) { + int col = wn * WARP_N + j * 8 + (lane & 3) * 2; + float bb0 = b2[col], bb1 = b2[col + 1]; + uint32_t a4 = amax[i][j]; + #pragma unroll + for (int h = 0; h < 2; h++) { + // invalid slots compete as -inf via bit masking: no select for the compiler to turn into branches + uint32_t keep = (slot < my_counts[i * 2 + h]) ? 0xFFFFFFFFu : 0u; + float v0 = bf16_round(acc[i][j][h * 2] + bb0); + float v1 = bf16_round(acc[i][j][h * 2 + 1] + bb1); + #pragma unroll + for (int e = 0; e < 2; e++) { + float v = __uint_as_float((__float_as_uint(e ? v1 : v0) & keep) | (0xff800000u & ~keep)); + float &m = rmax[i][j][h * 2 + e]; + if (STORE) { + int sh = (h * 2 + e) * 8; + uint32_t updated = (a4 & ~(0xFFu << sh)) | ((uint32_t) slot << sh); + a4 = (v > m) ? updated : a4; + } + m = fmaxf(m, v); + } + } + amax[i][j] = a4; + } + __syncthreads(); // keeps the next slot's x loads issued before the epilogue instead of sunk to store_x + if (slot + 1 < S) store_x(); + __syncthreads(); + } + // ---- pooled = counts > 0 ? running max : 0 ; arg-max slots + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) { + int col = wn * WARP_N + j * 8 + (lane & 3) * 2; + #pragma unroll + for (int h = 0; h < 2; h++) { + int row = wm * WARP_M + i * 16 + (lane >> 2) + h * 8; + int b = b0 + row; + if (b >= B) continue; + bool any = my_counts[i * 2 + h] > 0; + float p0 = any ? rmax[i][j][h * 2] : 0.0f; + float p1 = any ? rmax[i][j][h * 2 + 1] : 0.0f; + *(uint32_t *) &pooled_out[(int64_t) b * D + col] = pack_bf16x2(p0, p1); + if (STORE) { + uint32_t a4 = amax[i][j]; + uint16_t ap = (uint16_t) (((a4 >> (h * 16)) & 0xFFu) | (((a4 >> (h * 16 + 8)) & 0xFFu) << 8)); + *(uint16_t *) &argmax_out[(int64_t) b * D + col] = ap; + } + } + } +} +""" + + +def _check(res): + err, *rest = res + if int(err) != 0: + raise RuntimeError(f"CUDA error {err}") + return rest[0] if len(rest) == 1 else rest + + +def _instantiate_name(D, F_PAD, store): + return f"fused_slot_encoder_fwd_cuda<{D},{F_PAD},{'true' if store else 'false'}>" + + +@functools.lru_cache(maxsize=None) +def _compile(D, F_PAD, store, device_index): + major, minor = torch.cuda.get_device_capability(device_index) + name = _instantiate_name(D, F_PAD, store) + prog = _check(nvrtc.nvrtcCreateProgram(CUDA_SRC_FWD.encode(), b"fused_slot_encoder.cu", 0, [], [])) + _check(nvrtc.nvrtcAddNameExpression(prog, name.encode())) + opts = [ + f"--gpu-architecture=sm_{major}{minor}".encode(), + b"--std=c++17", + b"-default-device", + b"--extra-device-vectorization", + ] + (err,) = nvrtc.nvrtcCompileProgram(prog, len(opts), opts) + log_size = _check(nvrtc.nvrtcGetProgramLogSize(prog)) + log = b" " * log_size + nvrtc.nvrtcGetProgramLog(prog, log) + if int(err) != 0: + raise RuntimeError("NVRTC failed:\n" + log.decode()) + lowered = _check(nvrtc.nvrtcGetLoweredName(prog, name.encode())) + cubin_size = _check(nvrtc.nvrtcGetCUBINSize(prog)) + cubin = b" " * cubin_size + _check(nvrtc.nvrtcGetCUBIN(prog, cubin)) + torch.cuda.init() + torch.zeros(1, device=f"cuda:{device_index}") + module = _check(cu.cuModuleLoadData(cubin)) + func = _check(cu.cuModuleGetFunction(module, lowered)) + threads = 32 * 2 * (D // 64) + smem = (64 * (F_PAD + 8) + D * (F_PAD + 8) + 2 * D * 40 + 64 * (D + 8)) * 2 + _check(cu.cuFuncSetAttribute(func, cu.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem)) + regs = _check(cu.cuFuncGetAttribute(cu.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, func)) + local = _check(cu.cuFuncGetAttribute(cu.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, func)) + return func, threads, smem, regs, local, log.decode() + + +MAX_SLOTS = 255 + + +def feature_pad(F): + return 16 if F <= 16 else 32 + + +def cuda_slot_encoder_fwd(x, counts, w1_padded, b1, gamma, beta, w2, b2, store_acts): + """Mirror of fused_slot_encoder_fwd; w2 is [D][D] (not transposed). Returns pooled, argmax (uint8, STORE only), normed_relu.""" + B, S, F = x.shape + D, F_PAD = w1_padded.shape + assert x.dtype == torch.float32 and w1_padded.dtype == torch.bfloat16 and w2.dtype == torch.bfloat16 + assert w1_padded.is_contiguous() and w2.is_contiguous() and x.stride(2) == 1 + func, threads, smem, _, _, _ = _compile(D, F_PAD, bool(store_acts), x.device.index or 0) + if S > MAX_SLOTS: + raise ValueError(f"fused slot encoder stores the arg-max slot as uint8: {S} slots > {MAX_SLOTS}") + rows = B * S if store_acts else 0 + normed = torch.empty(rows, D, dtype=torch.bfloat16, device=x.device) + pooled = torch.empty(B, D, dtype=torch.bfloat16, device=x.device) + argmax = torch.empty(B if store_acts else 0, D, dtype=torch.uint8, device=x.device) + counts = counts.contiguous() + assert counts.dtype == torch.int64 + args = [ + np.array([x.data_ptr()], dtype=np.uint64), + np.array([counts.data_ptr()], dtype=np.uint64), + np.array([w1_padded.data_ptr()], dtype=np.uint64), + np.array([b1.data_ptr()], dtype=np.uint64), + np.array([gamma.data_ptr()], dtype=np.uint64), + np.array([beta.data_ptr()], dtype=np.uint64), + np.array([w2.data_ptr()], dtype=np.uint64), + np.array([b2.data_ptr()], dtype=np.uint64), + np.array([normed.data_ptr()], dtype=np.uint64), + np.array([pooled.data_ptr()], dtype=np.uint64), + np.array([argmax.data_ptr()], dtype=np.uint64), + np.array([B], dtype=np.int32), + np.array([S], dtype=np.int32), + np.array([F], dtype=np.int32), + np.array([x.stride(0)], dtype=np.int64), + np.array([x.stride(1)], dtype=np.int64), + np.array([x.stride(2)], dtype=np.int64), + ] + arg_ptrs = np.array([a.ctypes.data for a in args], dtype=np.uint64) + grid = (B + 63) // 64 + stream = torch.cuda.current_stream(x.device).cuda_stream + _check(cu.cuLaunchKernel(func, grid, 1, 1, threads, 1, 1, smem, stream, arg_ptrs.ctypes.data, 0)) + return pooled, argmax, normed + + +TILE_ROWS = 64 +MAX_ROWS_PER_BLOCK = 1024 +BLOCKS_PER_SM_TARGET = 4 + +CUDA_SRC_BWD = r""" +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef long long int64_t; +#define TILE_ROWS 64 +#define WARP_M 32 +#define WARP_N 64 +#define NWARPS_M 2 + +__device__ __forceinline__ uint32_t smem_u32(const void *p) { return (uint32_t) __cvta_generic_to_shared(p); } +__device__ __forceinline__ uint32_t f32_to_bf16_bits(float f) { uint16_t r; asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(r) : "f"(f)); return r; } +__device__ __forceinline__ float bf16_bits_to_f32(uint32_t b) { return __uint_as_float(b << 16); } +__device__ __forceinline__ float bf16_round(float f) { return bf16_bits_to_f32(f32_to_bf16_bits(f)); } +__device__ __forceinline__ uint32_t pack_bf16x2(float lo, float hi) { uint32_t r; asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(r) : "f"(hi), "f"(lo)); return r; } +__device__ __forceinline__ void mma_16816(float c[4], const uint32_t a[4], uint32_t b0, uint32_t b1) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} +__device__ __forceinline__ void ldmatrix_x4(uint32_t r[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(addr)); +} +__device__ __forceinline__ void ldmatrix_x4_trans(uint32_t r[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];" : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(addr)); +} +__device__ __forceinline__ void cp_async16(uint32_t smem_addr, const void *gptr) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" :: "r"(smem_addr), "l"(gptr)); +} +__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;"); } +template __device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;" :: "n"(N)); } + +// Row-wise sum over the mma C layout: lane partial over its 16 columns -> 4-lane shuffle -> cross-warp smem. +// vals[i][j][e]; result per (i, h) row in row_sum[i][h]. sRed: [NWARPS_N][TILE_ROWS] floats. +template +__device__ __forceinline__ void row_reduce_c_layout(const float vals[2][8][4], float row_sum[2][2], float *sRed, int wm, int wn, int lane) { + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int h = 0; h < 2; h++) { + float s = 0.0f; + #pragma unroll + for (int j = 0; j < 8; j++) s += vals[i][j][h * 2] + vals[i][j][h * 2 + 1]; + s += __shfl_xor_sync(0xffffffffu, s, 1); + s += __shfl_xor_sync(0xffffffffu, s, 2); + if ((lane & 3) == 0) { + int row = wm * WARP_M + i * 16 + (lane >> 2) + h * 8; + sRed[wn * TILE_ROWS + row] = s; + } + } + __syncthreads(); + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int h = 0; h < 2; h++) { + int row = wm * WARP_M + i * 16 + (lane >> 2) + h * 8; + float s = 0.0f; + #pragma unroll + for (int w = 0; w < NWARPS_N; w++) s += sRed[w * TILE_ROWS + row]; + row_sum[i][h] = s; + } + __syncthreads(); +} + +template +__global__ void __launch_bounds__(32 * NWARPS_M * (D / WARP_N), 1) +fused_slot_encoder_bwd_cuda( + const float *__restrict__ x, const int64_t *__restrict__ counts, + const uint16_t *__restrict__ w1, const float *__restrict__ b1, const float *__restrict__ gamma, const float *__restrict__ beta, + const uint16_t *__restrict__ w2, + const uint8_t *__restrict__ argmax, const uint16_t *__restrict__ grad_pooled, + uint16_t *__restrict__ grad_encoded_out, float *__restrict__ grad_w1_partial, + float *__restrict__ grad_b2_partial, float *__restrict__ grad_gamma_partial, float *__restrict__ grad_beta_partial, float *__restrict__ grad_b1_partial, + int B, int S, int F, int64_t stride_xb, int64_t stride_xs, int64_t stride_xf, int64_t stride_grad_row, int rows_per_block) +{ + constexpr int NWARPS_N = D / WARP_N; + constexpr int THREADS = 32 * NWARPS_M * NWARPS_N; + constexpr int LDX = F_PAD + 8; + constexpr int KCHUNK = 32; + constexpr int LDW2 = D + 8; // W2 chunk rows are K (= n of the forward), columns N (= k) + constexpr int NCHUNK = D / KCHUNK; + constexpr int LDA = D + 8; + constexpr int PIECES_PER_ROW = D / 8; + constexpr int ROWS_PER_PASS = THREADS / PIECES_PER_ROW; + constexpr int X_ELEMS = TILE_ROWS * F_PAD / THREADS; + const int64_t total_rows = (int64_t) B * S; + + extern __shared__ __align__(16) unsigned char smem_raw[]; + uint16_t *sX = (uint16_t *) smem_raw; // [TILE_ROWS][LDX] + uint16_t *sW1 = sX + TILE_ROWS * LDX; // [D][LDX] + uint16_t *sW2 = sW1 + D * LDX; // [2][KCHUNK][LDW2] + uint16_t *sA = sW2 + 2 * KCHUNK * LDW2; // [TILE_ROWS][LDA] + float *sRed = (float *) (sA + TILE_ROWS * LDA); // [NWARPS_N][TILE_ROWS] + float *sCol = sRed + NWARPS_N * TILE_ROWS; // [NWARPS_M * 8][D] column partial scratch + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int wm = warp / NWARPS_N; + const int wn = warp % NWARPS_N; + + for (int i = tid; i < D * (F_PAD / 8); i += THREADS) { + int n = i / (F_PAD / 8), q = i % (F_PAD / 8); + *(uint4 *) &sW1[n * LDX + q * 8] = *(const uint4 *) &w1[n * F_PAD + q * 8]; + } + float *sAccW1 = sCol; // [D][F_PAD] W1 weight gradient of this block + for (int i = tid; i < D * F_PAD; i += THREADS) sAccW1[i] = 0.0f; + // W2 chunk c = rows [c*32, +32) of W2[n][k] (K = n for this GEMM), 16-byte pieces + auto issue_w2_chunk = [&](int chunk, int stage) { + const uint16_t *src = w2 + (int64_t) chunk * KCHUNK * D; + uint16_t *dst = sW2 + stage * KCHUNK * LDW2; + #pragma unroll + for (int i = 0; i < (KCHUNK * PIECES_PER_ROW) / THREADS; i++) { + int e = tid + i * THREADS; + int r = e / PIECES_PER_ROW, q = e % PIECES_PER_ROW; + cp_async16(smem_u32(&dst[r * LDW2 + q * 8]), src + r * D + q * 8); + } + cp_async_commit(); + }; + issue_w2_chunk(0, 0); + issue_w2_chunk(1, 1); + + float acc_gamma[2][8][2], acc_beta[2][8][2], acc_b1[2][8][2]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) + #pragma unroll + for (int e = 0; e < 2; e++) { acc_gamma[i][j][e] = 0.0f; acc_beta[i][j][e] = 0.0f; acc_b1[i][j][e] = 0.0f; } + float acc_b2[8]; + #pragma unroll + for (int t = 0; t < 8; t++) acc_b2[t] = 0.0f; + + const uint32_t sA_base = smem_u32(sA); + const uint32_t sX_base = smem_u32(sX); + const uint32_t sW1_base = smem_u32(sW1); + const uint32_t sW2_base = smem_u32(sW2); + const int64_t block_row0 = (int64_t) blockIdx.x * rows_per_block; + __syncthreads(); + + for (int tile = 0; tile < rows_per_block / TILE_ROWS; tile++) { + const int64_t row0 = block_row0 + (int64_t) tile * TILE_ROWS; + if (row0 >= total_rows) break; + // ---- x tile -> sX (bf16), pool backward -> grad_encoded (global + sA) + #pragma unroll + for (int q = 0; q < X_ELEMS; q++) { + int e = tid + q * THREADS; + int r = e / F_PAD, f = e % F_PAD; + int64_t row = row0 + r; + float v = 0.0f; + if (row < total_rows && f < F) { + int64_t b = row / S, s = row - b * S; + v = x[b * stride_xb + s * stride_xs + f * stride_xf]; + } + sX[r * LDX + f] = (uint16_t) f32_to_bf16_bits(v); + } + #pragma unroll + for (int pass = 0; pass < TILE_ROWS / ROWS_PER_PASS; pass++) { + int r = pass * ROWS_PER_PASS + tid / PIECES_PER_ROW; + int p = tid % PIECES_PER_ROW; + int64_t row = row0 + r; + uint32_t packed[4] = {0u, 0u, 0u, 0u}; + if (row < total_rows) { + int64_t b = row / S, s = row - b * S; + bool valid = s < counts[b]; + uint2 am = *(const uint2 *) &argmax[b * D + p * 8]; + uint4 gpl = *(const uint4 *) &grad_pooled[b * stride_grad_row + p * 8]; + uint32_t aw[2] = {am.x, am.y}; + uint32_t gw[4] = {gpl.x, gpl.y, gpl.z, gpl.w}; + #pragma unroll + for (int t = 0; t < 4; t++) { + float out[2]; + #pragma unroll + for (int e = 0; e < 2; e++) { + int64_t winner = (aw[t >> 1] >> (8 * ((t & 1) * 2 + e))) & 0xFFu; + float gv = bf16_bits_to_f32((gw[t] >> (16 * e)) & 0xFFFFu); + out[e] = (valid && winner == s) ? gv : 0.0f; + acc_b2[t * 2 + e] += out[e]; + } + packed[t] = pack_bf16x2(out[0], out[1]); + } + *(uint4 *) &grad_encoded_out[row * D + p * 8] = make_uint4(packed[0], packed[1], packed[2], packed[3]); + } + *(uint4 *) &sA[r * LDA + p * 8] = make_uint4(packed[0], packed[1], packed[2], packed[3]); + } + __syncthreads(); + // ---- GEMM1 recompute: pre_norm = bf16(x @ W1^T + b1), kept in registers (C layout) + float pre[2][8][4]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) + #pragma unroll + for (int e = 0; e < 4; e++) pre[i][j][e] = 0.0f; + #pragma unroll + for (int kk = 0; kk < F_PAD / 16; kk++) { + uint32_t afrag[2][4]; + #pragma unroll + for (int i = 0; i < 2; i++) { + int row = wm * WARP_M + i * 16 + (lane & 15); + int col = kk * 16 + (lane >> 4) * 8; + ldmatrix_x4(afrag[i], sX_base + (row * LDX + col) * 2); + } + #pragma unroll + for (int jj = 0; jj < 4; jj++) { + uint32_t bfrag[4]; + int m = lane >> 3, r = lane & 7; + int n = wn * WARP_N + jj * 16 + (m >> 1) * 8 + r; + int k = kk * 16 + (m & 1) * 8; + ldmatrix_x4(bfrag, sW1_base + (n * LDX + k) * 2); + #pragma unroll + for (int i = 0; i < 2; i++) { + mma_16816(pre[i][jj * 2], afrag[i], bfrag[0], bfrag[1]); + mma_16816(pre[i][jj * 2 + 1], afrag[i], bfrag[2], bfrag[3]); + } + } + } + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) { + int col = wn * WARP_N + j * 8 + (lane & 3) * 2; + float bb0 = b1[col], bb1 = b1[col + 1]; + #pragma unroll + for (int h = 0; h < 2; h++) { + pre[i][j][h * 2] = bf16_round(pre[i][j][h * 2] + bb0); + pre[i][j][h * 2 + 1] = bf16_round(pre[i][j][h * 2 + 1] + bb1); + } + } + // ---- GEMM3: grad_normed_relu = bf16(grad_encoded @ W2), W2 rows streamed as K chunks (ldmatrix.trans) + float gnr[2][8][4]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) + #pragma unroll + for (int e = 0; e < 4; e++) gnr[i][j][e] = 0.0f; + for (int chunk = 0; chunk < NCHUNK; chunk++) { + cp_async_wait<1>(); + __syncthreads(); + const uint32_t stage_base = sW2_base + (chunk & 1) * KCHUNK * LDW2 * 2; + #pragma unroll + for (int kk = 0; kk < KCHUNK / 16; kk++) { + uint32_t afrag[2][4]; + #pragma unroll + for (int i = 0; i < 2; i++) { + int row = wm * WARP_M + i * 16 + (lane & 15); + int col = chunk * KCHUNK + kk * 16 + (lane >> 4) * 8; + ldmatrix_x4(afrag[i], sA_base + (row * LDA + col) * 2); + } + #pragma unroll + for (int jj = 0; jj < 4; jj++) { + // B[k][n] row-major tile: matrices (k 0-7 | n-tile j), (k 8-15 | j), (k 0-7 | j+1), (k 8-15 | j+1) + uint32_t bfrag[4]; + int m = lane >> 3, r = lane & 7; + int k = kk * 16 + (m & 1) * 8 + r; + int n = wn * WARP_N + jj * 16 + (m >> 1) * 8; + ldmatrix_x4_trans(bfrag, stage_base + (k * LDW2 + n) * 2); + #pragma unroll + for (int i = 0; i < 2; i++) { + mma_16816(gnr[i][jj * 2], afrag[i], bfrag[0], bfrag[1]); + mma_16816(gnr[i][jj * 2 + 1], afrag[i], bfrag[2], bfrag[3]); + } + } + } + __syncthreads(); + int next = chunk + 2; + issue_w2_chunk(next < NCHUNK ? next : next - NCHUNK, chunk & 1); + } + // ---- LayerNorm forward stats + ReLU/LayerNorm backward in the C layout + float row_sum[2][2]; + row_reduce_c_layout(pre, row_sum, sRed, wm, wn, lane); + float mean[2][2]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int h = 0; h < 2; h++) mean[i][h] = row_sum[i][h] / (float) D; + float sq[2][8][4]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) + #pragma unroll + for (int e = 0; e < 4; e++) { + float c = pre[i][j][e] - mean[i][e >> 1]; + pre[i][j][e] = c; // centered + sq[i][j][e] = c * c; + } + row_reduce_c_layout(sq, row_sum, sRed, wm, wn, lane); + float rstd[2][2]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int h = 0; h < 2; h++) rstd[i][h] = rsqrtf(row_sum[i][h] / (float) D + 1e-5f); + // pre -> normalized; gnr -> grad_scaled (= grad_normed * gamma); sq -> grad_scaled * normalized + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) { + int col = wn * WARP_N + j * 8 + (lane & 3) * 2; + float g[2] = {gamma[col], gamma[col + 1]}; + float be[2] = {beta[col], beta[col + 1]}; + #pragma unroll + for (int e = 0; e < 4; e++) { + float normalized = pre[i][j][e] * rstd[i][e >> 1]; + float normed = normalized * g[e & 1] + be[e & 1]; + float grad_normed = normed > 0.0f ? bf16_round(gnr[i][j][e]) : 0.0f; + acc_gamma[i][j][e & 1] += grad_normed * normalized; + acc_beta[i][j][e & 1] += grad_normed; + float grad_scaled = grad_normed * g[e & 1]; + pre[i][j][e] = normalized; + gnr[i][j][e] = grad_scaled; + sq[i][j][e] = grad_scaled * normalized; + } + } + float mean_grad[2][2], mean_grad_normalized[2][2]; + row_reduce_c_layout(gnr, row_sum, sRed, wm, wn, lane); + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int h = 0; h < 2; h++) mean_grad[i][h] = row_sum[i][h] / (float) D; + row_reduce_c_layout(sq, row_sum, sRed, wm, wn, lane); + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int h = 0; h < 2; h++) mean_grad_normalized[i][h] = row_sum[i][h] / (float) D; + // grad_pre_norm = bf16(rstd * (grad_scaled - mean_grad - normalized * mean_grad_normalized)) -> sA + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) { + int col = wn * WARP_N + j * 8 + (lane & 3) * 2; + #pragma unroll + for (int h = 0; h < 2; h++) { + float v[2]; + #pragma unroll + for (int e = 0; e < 2; e++) { + int idx = h * 2 + e; + v[e] = bf16_round(rstd[i][h] * (gnr[i][j][idx] - mean_grad[i][h] - pre[i][j][idx] * mean_grad_normalized[i][h])); + acc_b1[i][j][e] += v[e]; + } + int row = wm * WARP_M + i * 16 + (lane >> 2) + h * 8; + *(uint32_t *) &sA[row * LDA + col] = pack_bf16x2(v[0], v[1]); + } + } + __syncthreads(); + // ---- grad_w1 += grad_pre_norm^T @ x over this tile: A = sA read transposed, B = sX, K = the 64 lanes + { + float acc_w1[2][F_PAD / 8][4]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int jn = 0; jn < F_PAD / 8; jn++) + #pragma unroll + for (int e = 0; e < 4; e++) acc_w1[i][jn][e] = 0.0f; + const int m = lane >> 3, r = lane & 7; + #pragma unroll + for (int ks = 0; ks < TILE_ROWS / 16; ks++) { + uint32_t bfrag[F_PAD / 16][4]; + #pragma unroll + for (int jj = 0; jj < F_PAD / 16; jj++) { + int k = ks * 16 + (m & 1) * 8 + r; + int n = jj * 16 + (m >> 1) * 8; + ldmatrix_x4_trans(bfrag[jj], sX_base + (k * LDX + n) * 2); + } + #pragma unroll + for (int i = 0; i < 2; i++) { + uint32_t afrag[4]; + int k = ks * 16 + (m >> 1) * 8 + r; + int d = warp * 32 + i * 16 + (m & 1) * 8; + ldmatrix_x4_trans(afrag, sA_base + (k * LDA + d) * 2); + #pragma unroll + for (int jj = 0; jj < F_PAD / 16; jj++) { + mma_16816(acc_w1[i][jj * 2], afrag, bfrag[jj][0], bfrag[jj][1]); + mma_16816(acc_w1[i][jj * 2 + 1], afrag, bfrag[jj][2], bfrag[jj][3]); + } + } + } + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int jn = 0; jn < F_PAD / 8; jn++) + #pragma unroll + for (int h = 0; h < 2; h++) { + int d = warp * 32 + i * 16 + (lane >> 2) + h * 8; + int f = jn * 8 + (lane & 3) * 2; + sAccW1[d * F_PAD + f] += acc_w1[i][jn][h * 2]; + sAccW1[d * F_PAD + f + 1] += acc_w1[i][jn][h * 2 + 1]; + } + } + __syncthreads(); + } + for (int i = tid; i < D * F_PAD; i += THREADS) grad_w1_partial[(int64_t) blockIdx.x * D * F_PAD + i] = sAccW1[i]; + __syncthreads(); + // ---- partial sums: columns owned per lane (C layout) -> reduce over the 8 lanes sharing lane&3 and over wm + auto store_c_layout_partial = [&](float (*acc)[8][2], float *dst) { + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 8; j++) + #pragma unroll + for (int e = 0; e < 2; e++) { + float s = acc[i][j][e]; + s += __shfl_xor_sync(0xffffffffu, s, 4); + s += __shfl_xor_sync(0xffffffffu, s, 8); + s += __shfl_xor_sync(0xffffffffu, s, 16); + if ((lane >> 2) == 0) { + int col = wn * WARP_N + j * 8 + (lane & 3) * 2 + e; + sCol[(wm * 2 + i) * D + col] = s; + } + } + __syncthreads(); + for (int col = tid; col < D; col += THREADS) { + float s = 0.0f; + #pragma unroll + for (int g = 0; g < NWARPS_M * 2; g++) s += sCol[g * D + col]; + dst[(int64_t) blockIdx.x * D + col] = s; + } + __syncthreads(); + }; + store_c_layout_partial(acc_gamma, grad_gamma_partial); + store_c_layout_partial(acc_beta, grad_beta_partial); + store_c_layout_partial(acc_b1, grad_b1_partial); + // grad_b2: thread owns cols [8*(tid % PIECES_PER_ROW), +8); reduce over the ROWS_PER_PASS row groups + #pragma unroll + for (int t = 0; t < 8; t++) { + sCol[(tid / PIECES_PER_ROW) * D + (tid % PIECES_PER_ROW) * 8 + t] = acc_b2[t]; + } + __syncthreads(); + for (int col = tid; col < D; col += THREADS) { + float s = 0.0f; + #pragma unroll + for (int g = 0; g < ROWS_PER_PASS; g++) s += sCol[g * D + col]; + grad_b2_partial[(int64_t) blockIdx.x * D + col] = s; + } +} +""" + + +@functools.lru_cache(maxsize=None) +def _compile_bwd(D, F_PAD, device_index): + major, minor = torch.cuda.get_device_capability(device_index) + name = f"fused_slot_encoder_bwd_cuda<{D},{F_PAD}>" + prog = _check(nvrtc.nvrtcCreateProgram(CUDA_SRC_BWD.encode(), b"fused_slot_encoder_bwd.cu", 0, [], [])) + _check(nvrtc.nvrtcAddNameExpression(prog, name.encode())) + opts = [ + f"--gpu-architecture=sm_{major}{minor}".encode(), + b"--std=c++17", + b"-default-device", + b"--extra-device-vectorization", + ] + (err,) = nvrtc.nvrtcCompileProgram(prog, len(opts), opts) + log_size = _check(nvrtc.nvrtcGetProgramLogSize(prog)) + log = b" " * log_size + nvrtc.nvrtcGetProgramLog(prog, log) + if int(err) != 0: + raise RuntimeError("NVRTC failed:\n" + log.decode()) + lowered = _check(nvrtc.nvrtcGetLoweredName(prog, name.encode())) + cubin_size = _check(nvrtc.nvrtcGetCUBINSize(prog)) + cubin = b" " * cubin_size + _check(nvrtc.nvrtcGetCUBIN(prog, cubin)) + torch.cuda.init() + torch.zeros(1, device=f"cuda:{device_index}") + module = _check(cu.cuModuleLoadData(cubin)) + func = _check(cu.cuModuleGetFunction(module, lowered)) + nwarps_n = D // 64 + threads = 32 * 2 * nwarps_n + rows_per_pass = threads // (D // 8) + smem = (64 * (F_PAD + 8) + D * (F_PAD + 8) + 2 * 32 * (D + 8) + 64 * (D + 8)) * 2 + ( + nwarps_n * 64 + max(4, rows_per_pass, F_PAD) * D + ) * 4 + _check(cu.cuFuncSetAttribute(func, cu.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem)) + regs = _check(cu.cuFuncGetAttribute(cu.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, func)) + local = _check(cu.cuFuncGetAttribute(cu.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, func)) + return func, threads, smem, regs, local + + +@functools.lru_cache(maxsize=None) +def _sm_count(device_index): + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def bwd_rows_per_block(rows, device_index): + """Lanes per block: enough blocks for BLOCKS_PER_SM_TARGET waves, capped so the W1 partials stay small.""" + per_block = rows // (BLOCKS_PER_SM_TARGET * _sm_count(device_index)) // TILE_ROWS * TILE_ROWS + return max(TILE_ROWS, min(MAX_ROWS_PER_BLOCK, per_block)) + + +def cuda_slot_encoder_bwd(x, counts, w1_padded, b1, gamma, beta, w2, normed_relu, argmax, grad_pooled): + """Mirror of fused_slot_encoder_bwd: returns grad_w1, grad_b1, grad_gamma, grad_beta, grad_w2, grad_b2.""" + B, S, F = x.shape + D, F_PAD = w1_padded.shape + func, threads, smem, _, _ = _compile_bwd(D, F_PAD, x.device.index or 0) + rows = B * S + if grad_pooled.stride(1) != 1: + grad_pooled = grad_pooled.contiguous() + counts = counts.contiguous() + rows_per_block = bwd_rows_per_block(rows, x.device.index or 0) + blocks = (rows + rows_per_block - 1) // rows_per_block + grad_encoded = torch.empty(rows, D, dtype=torch.bfloat16, device=x.device) + grad_w1_partial = torch.empty(blocks, D, F_PAD, dtype=torch.float32, device=x.device) + partials = torch.empty(4, blocks, D, dtype=torch.float32, device=x.device) + p = lambda t: np.array([t.data_ptr()], dtype=np.uint64) + args = [ + p(x), + p(counts), + p(w1_padded), + p(b1), + p(gamma), + p(beta), + p(w2), + p(argmax), + p(grad_pooled), + p(grad_encoded), + p(grad_w1_partial), + p(partials[0]), + p(partials[1]), + p(partials[2]), + p(partials[3]), + np.array([B], dtype=np.int32), + np.array([S], dtype=np.int32), + np.array([F], dtype=np.int32), + np.array([x.stride(0)], dtype=np.int64), + np.array([x.stride(1)], dtype=np.int64), + np.array([x.stride(2)], dtype=np.int64), + np.array([grad_pooled.stride(0)], dtype=np.int64), + np.array([rows_per_block], dtype=np.int32), + ] + arg_ptrs = np.array([a.ctypes.data for a in args], dtype=np.uint64) + stream = torch.cuda.current_stream(x.device).cuda_stream + _check(cu.cuLaunchKernel(func, blocks, 1, 1, threads, 1, 1, smem, stream, arg_ptrs.ctypes.data, 0)) + grad_w2 = (grad_encoded.t() @ normed_relu).float() + grad_w1 = grad_w1_partial[:, :, :F].sum(0).to(torch.bfloat16).float() + grad_b2 = partials[0].sum(0).to(torch.bfloat16).float() + grad_gamma = partials[1].sum(0) + grad_beta = partials[2].sum(0) + grad_b1 = partials[3].sum(0).to(torch.bfloat16).float() + return grad_w1, grad_b1, grad_gamma, grad_beta, grad_w2, grad_b2 diff --git a/pufferlib/ocean/drive/binding.c b/pufferlib/ocean/drive/binding.c index 0c73e69a79..d21d15200b 100644 --- a/pufferlib/ocean/drive/binding.c +++ b/pufferlib/ocean/drive/binding.c @@ -1926,6 +1926,7 @@ static int unpack_map_cache_args(PyObject *kwargs, Drive *config, const char *** config->use_neighbor_cache = (int) unpack(kwargs, "use_neighbor_cache"); config->obs_lane_stride = (int) unpack(kwargs, "obs_lane_stride"); config->obs_boundary_stride = (int) unpack(kwargs, "obs_boundary_stride"); + config->obs_stride_merge_segments = (int) unpack(kwargs, "obs_stride_merge_segments"); config->obs_range_road_front_m = (float) unpack(kwargs, "obs_range_road_front_m"); config->obs_range_road_behind_m = (float) unpack(kwargs, "obs_range_road_behind_m"); config->obs_range_road_side_m = (float) unpack(kwargs, "obs_range_road_side_m"); @@ -1933,7 +1934,8 @@ static int unpack_map_cache_args(PyObject *kwargs, Drive *config, const char *** return -1; } if (config->use_neighbor_cache < 0 || config->use_neighbor_cache > 1 || config->obs_lane_stride < 1 - || config->obs_boundary_stride < 1 || !isfinite(config->obs_range_road_front_m) + || config->obs_boundary_stride < 1 || config->obs_stride_merge_segments < 0 + || config->obs_stride_merge_segments > 1 || !isfinite(config->obs_range_road_front_m) || !isfinite(config->obs_range_road_behind_m) || !isfinite(config->obs_range_road_side_m) || config->obs_range_road_front_m < 0.0f || config->obs_range_road_behind_m < 0.0f || config->obs_range_road_side_m < 0.0f) { @@ -2063,6 +2065,7 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) { env->yield_signs_enabled = (bool) unpack(kwargs, "yield_signs_enabled"); env->obs_lane_stride = (int) unpack(kwargs, "obs_lane_stride"); env->obs_boundary_stride = (int) unpack(kwargs, "obs_boundary_stride"); + env->obs_stride_merge_segments = (int) unpack(kwargs, "obs_stride_merge_segments"); env->dt = (float) unpack(kwargs, "dt"); env->base_max_speed_mps = (float) unpack(kwargs, "base_max_speed_mps"); env->max_speed_mps = (float) unpack(kwargs, "max_speed_mps"); @@ -2145,6 +2148,11 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) { env->obs_norm_road_seg_width_m = (float) unpack(kwargs, "obs_norm_road_seg_width_m"); env->obs_norm_z_m = (float) unpack(kwargs, "obs_norm_z_m"); env->eval_perceived_size_margin_m = (float) unpack(kwargs, "eval_perceived_size_margin_m"); + env->obs_partner_split_length_m = (float) unpack(kwargs, "obs_partner_split_length_m"); + if (!(env->obs_partner_split_length_m >= 0.0f) || !isfinite(env->obs_partner_split_length_m)) { + PyErr_SetString(PyExc_ValueError, "obs_partner_split_length_m must be finite and >= 0"); + return -1; + } env->eval_standstill_jerk_deadband_mps3 = (float) unpack(kwargs, "eval_standstill_jerk_deadband_mps3"); env->obs_range_traffic_control_m = (float) unpack(kwargs, "obs_range_traffic_control_m"); env->obs_range_partner_m = (float) unpack(kwargs, "obs_range_partner_m"); @@ -2154,9 +2162,20 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) { env->obs_slots_lane_kept = (int) unpack(kwargs, "obs_slots_lane_kept"); env->obs_slots_boundary_kept = (int) unpack(kwargs, "obs_slots_boundary_kept"); env->partner_blindness_prob = (float) unpack(kwargs, "partner_blindness_prob"); + env->partner_blindness_prob_min = (float) unpack(kwargs, "partner_blindness_prob_min"); + if (!(env->partner_blindness_prob_min >= 0.0f && env->partner_blindness_prob_min <= env->partner_blindness_prob)) { + PyErr_SetString(PyExc_ValueError, "partner_blindness_prob_min must be within [0, partner_blindness_prob]"); + return -1; + } env->partner_blindness_trigger_prob = (float) unpack(kwargs, "partner_blindness_trigger_prob"); env->partner_blindness_duration = (int) ceilf((float) unpack(kwargs, "partner_blindness_duration_seconds") / env->dt); + env->partner_blindness_corridor = (int) unpack(kwargs, "partner_blindness_corridor"); + env->partner_blindness_time_fraction_max = (float) unpack(kwargs, "partner_blindness_time_fraction_max"); + if (!(env->partner_blindness_time_fraction_max >= 0.0f && env->partner_blindness_time_fraction_max <= 1.0f)) { + PyErr_SetString(PyExc_ValueError, "partner_blindness_time_fraction_max must be within [0, 1]"); + return -1; + } env->phantom_braking_prob = (float) unpack(kwargs, "phantom_braking_prob"); env->phantom_braking_trigger_prob = (float) unpack(kwargs, "phantom_braking_trigger_prob"); env->phantom_braking_duration = (int) ceilf((float) unpack(kwargs, "phantom_braking_duration_seconds") / env->dt); diff --git a/pufferlib/ocean/drive/constants.h b/pufferlib/ocean/drive/constants.h index 9977ff483d..02f1c8f908 100644 --- a/pufferlib/ocean/drive/constants.h +++ b/pufferlib/ocean/drive/constants.h @@ -156,9 +156,9 @@ static const float STEERING_VALUES[NUM_STEERING_ACTIONS] #define LANE_DISTANCE_NORMALIZATION 4.0f #define LANE_SWITCH_THRESHOLD 0.05f // Hysteresis: new lane must be 5% better to switch #define LANE_ALIGN_COS_THRESHOLD 0.5f -#define LANE_CURVATURE_NORM 0.2f // 1/m; a 5 m turn radius saturates the obs -#define MAX_CHECKED_LANES 32 -#define LANE_WIDTH 3.7f // fallback when the map carries no per-point widths +#define LANE_CURVATURE_NORM 0.2f // 1/m; a 5 m turn radius saturates the obs +#define LANE_SCAN_GUARD_MARGIN_M 0.01f // float slack so an unlisted segment never ties the cached nearest one +#define LANE_WIDTH 3.7f // fallback when the map carries no per-point widths #define MAX_LANE_WIDTH_M 50.0f #define LANE_MARGIN 0.2f @@ -214,8 +214,9 @@ static const float STEERING_VALUES[NUM_STEERING_ACTIONS] // Depends on resolution of data Formula: 3 * (2 + GRID_CELL_SIZE*sqrt(2)/resolution) // => For each entity type in gridmap, diagonal poly-lines -> sqrt(2), include diagonal ends -> 2 #define MAX_ENTITIES_PER_CELL 64 -#define ROAD_QUERY_ENTITY_COUNT (MAX_ENTITIES_PER_CELL * 25) // 5x5 cell neighborhood -#define OBS_ENTITY_NONE 0 // grid entity not sampled for road observations +#define ROAD_QUERY_CELL_COUNT 25 // 5x5 cell neighborhood +#define ROAD_QUERY_ENTITY_COUNT (MAX_ENTITIES_PER_CELL * ROAD_QUERY_CELL_COUNT) +#define OBS_ENTITY_NONE 0 // grid entity not sampled for road observations #define OBS_ENTITY_LANE 1 #define OBS_ENTITY_EDGE 2 @@ -274,6 +275,8 @@ static const int ROAD_OFFSETS[25][2] #define OBS_VALID_COUNT_FEATURES 4 // Heading deviation since last kept point that forces a keep when obs stride > 1 (~15 degrees). #define OBS_STRIDE_HEADING_THRESHOLD 0.2618f +// Lateral slack beyond the ego's half width on each side of a blind partner's forward corridor +#define PARTNER_BLINDNESS_CORRIDOR_MARGIN_M 1.0f // ===================================================================================== // 9. REWARDS diff --git a/pufferlib/ocean/drive/datatypes.h b/pufferlib/ocean/drive/datatypes.h index 0daefdfd39..16e7a35a90 100644 --- a/pufferlib/ocean/drive/datatypes.h +++ b/pufferlib/ocean/drive/datatypes.h @@ -133,10 +133,11 @@ struct Agent { // Reward conditioning coefficients (per-agent, randomized at spawn) float reward_coefs[NUM_REWARD_COEFS]; - int phantom_braking_counter; // >0 means currently phantom braking - int partner_blindness_counter; // >0 means currently blind to partners - unsigned char is_blind_partner; // episode-level flag: agent sees no other agents - unsigned char is_phantom_braker; // episode-level flag: agent may phantom-brake + int phantom_braking_counter; // >0 means currently phantom braking + int partner_blindness_counter; // steps left in the current blind interval, 0 = sees partners normally + float partner_blindness_trigger_prob; // per-step chance of starting a blind interval while not blind + unsigned char is_blind_partner; // episode-level flag: agent may restrict partner visibility + unsigned char is_phantom_braker; // episode-level flag: agent may phantom-brake }; typedef struct { @@ -147,6 +148,7 @@ typedef struct { float direction_x; float direction_y; float width_m; + int observed; } RoadObservationSegment; struct RoadMapElement { @@ -169,6 +171,7 @@ struct RoadMapElement { float *cum_lengths; float *widths; // per-point lane width, meters int speed_zone_idx; // lanes sharing one posted limit; -1 = none (junction lane or map without zones) + float max_segment_half_length_m; }; struct TrafficControlElement { @@ -187,7 +190,7 @@ struct TrafficControlElement { struct LaneGraph { int n_lanes; int *lane_ids; - float *distances; // n_lanes * n_lanes row-major (row = from, col = to) + float *distances; // [goal * n_lanes + from]: row = goal lane, so an ego's lane obs read one contiguous row int *lane_to_graph_idx; // road-element idx -> graph idx (-1 if lane absent from graph), sized num_road_elements }; diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index 88b83bc039..7599963e0c 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -101,8 +101,16 @@ struct GridMapEntity { int entity_idx; // Index into the road_elements array int geometry_idx; // Index into element's geometry array int obs_kind; // OBS_ENTITY_NONE / LANE / EDGE: how road observations sample this entity + int pool_idx; // Position in entity_pool and segment_pool }; +// One road-query candidate of a cell's 5x5 window: a road-edge segment, or a drivable lane with its listed segments. +typedef struct { + int entity_idx; + int seg_first; + int seg_last; +} LaneQueryEntry; + struct GridMap { float top_left_x; float top_left_y; @@ -117,8 +125,12 @@ struct GridMap { int *grid_index_drivable; int num_drivable_grid_cell; int total_entities; - GridMapEntity *entity_pool; // one block; cells[i] point into it + GridMapEntity *entity_pool; // one block; cells[i] point into it + RoadObservationSegment *segment_pool; // parallel to entity_pool GridMapEntity **cells; + LaneQueryEntry *lane_query_pool; // per cell, window order, lanes deduplicated at first appearance + int *lane_query_offset; // cell count + 1 + int *window_entity_count; // every grid entity in the cell's 5x5 window uint16_t **neighbor_cache_pool_idx; // per cell: entity_pool indices of its spiral neighborhood }; @@ -139,6 +151,7 @@ struct SharedMapData { struct LaneGraph lane_graph; int obs_lane_stride; int obs_boundary_stride; + int obs_stride_merge_segments; int ref_count; pid_t owner_pid; }; @@ -288,6 +301,7 @@ struct Drive { int yield_signs_enabled; int obs_lane_stride; int obs_boundary_stride; + int obs_stride_merge_segments; int obs_slots_lane_kept; int obs_slots_boundary_kept; int road_dropout_enabled; @@ -300,7 +314,8 @@ struct Drive { float obs_norm_road_seg_width_m; float obs_norm_z_m; float eval_perceived_size_margin_m; - int cosim_eval_semantics; // co-sim: apply the eval-only obs/action tweaks without native eval batching + float obs_partner_split_length_m; // 0 disables; longer partners are observed as ceil(length / this) equal pieces + int cosim_eval_semantics; // co-sim: apply the eval-only obs/action tweaks without native eval batching float eval_standstill_jerk_deadband_mps3; // 0 disables; eval-only standstill deadband on j_long float obs_range_traffic_control_m; float obs_range_partner_m; @@ -309,8 +324,11 @@ struct Drive { float obs_range_road_side_m; // Robustness float partner_blindness_prob; + float partner_blindness_prob_min; float partner_blindness_trigger_prob; int partner_blindness_duration; + int partner_blindness_corridor; + float partner_blindness_time_fraction_max; float phantom_braking_prob; float phantom_braking_trigger_prob; int phantom_braking_duration; @@ -354,7 +372,7 @@ static const RewardBound REWARD_BOUNDS[NUM_REWARD_COEFS] = { {2.5e-4f, 7.5e-3f, 0}, // REWARD_COEF_REVERSE α_reverse ~ U(2.5e-4, 7.5e-3) {0.0f, 1.0f, 0}, // REWARD_COEF_STOP_LINE α_stop-line ~ U(0, 1) {0.0f, 5e-5f, 0}, // REWARD_COEF_TIMESTEP α_timestep ~ U(0, 5e-5f) - {0.0f, 1.0f, 0}, // REWARD_COEF_OVERSPEED α_overspeed ~ U(0, 1) + {0.0f, 0.1f, 0}, // REWARD_COEF_OVERSPEED α_overspeed ~ U(0, 0.1) {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 @@ -376,7 +394,7 @@ static const RewardBound REWARD_BOUNDS_LOG[NUM_REWARD_COEFS] = { {2.5e-4f, 7.5e-3f, 1}, // REWARD_COEF_REVERSE α_reverse ~ logU(2.5e-4, 7.5e-3) {0.0f, 1.0f, 0}, // REWARD_COEF_STOP_LINE α_stop-line ~ U(0, 1) {0.0f, 5e-5f, 0}, // REWARD_COEF_TIMESTEP α_timestep ~ U(0, 5e-5f) - {0.0f, 1.0f, 0}, // REWARD_COEF_OVERSPEED α_overspeed ~ U(0, 1) + {0.0f, 0.1f, 0}, // REWARD_COEF_OVERSPEED α_overspeed ~ U(0, 0.1) {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 @@ -417,7 +435,9 @@ static float clip(float value, float min, float max) { } static float normalize_heading(float heading) { - heading = fmodf(heading, 2.0f * M_PI); + if (!(fabsf(heading) < 2.0f * M_PI)) { // fmodf is the identity below one period + heading = fmodf(heading, 2.0f * M_PI); + } if (heading > M_PI) { heading -= 2.0f * M_PI; } else if (heading < -M_PI) { @@ -481,8 +501,9 @@ static void begin_episode_rng(Drive *env) { env->episode_seed = rng_next(&env->seed_stream_rng) >> 1; } rng_seed(&env->rng_state, env->episode_seed); - // Per-episode perturbation rates: uniform in [0, configured max] - env->episode_partner_blindness_prob = sample_uniform(&env->rng_state, 0.0f, env->partner_blindness_prob); + // Per-episode perturbation rates: uniform in [floor, configured max]; phantom braking has no floor knob + env->episode_partner_blindness_prob + = sample_uniform(&env->rng_state, env->partner_blindness_prob_min, env->partner_blindness_prob); env->episode_phantom_braking_prob = sample_uniform(&env->rng_state, 0.0f, env->phantom_braking_prob); } @@ -1954,14 +1975,17 @@ static bool update_stop_sign_state(Drive *env, int agent_idx) { agent->stop_sign_target_idx = -1; agent->stop_sign_stop_completed = 0; } - if (agent->stop_sign_target_idx < 0) { + // an uncompleted target keeps competing: a curved approach latches the neighbour lane's sign first + if (agent->stop_sign_target_idx < 0 || !agent->stop_sign_stop_completed) { int standstill_idx = find_stop_sign_standstill(env, agent); if (standstill_idx >= 0) { agent->stop_sign_standstill_idx = standstill_idx; } - agent->stop_sign_target_idx = find_stop_sign_target(env, agent); - agent->stop_sign_stop_completed - = agent->stop_sign_target_idx >= 0 && agent->stop_sign_target_idx == agent->stop_sign_standstill_idx; + int target_idx = find_stop_sign_target(env, agent); + if (target_idx >= 0 && target_idx != agent->stop_sign_target_idx) { + agent->stop_sign_target_idx = target_idx; + agent->stop_sign_stop_completed = target_idx == agent->stop_sign_standstill_idx; + } if (agent->stop_sign_target_idx < 0) { return false; } @@ -2541,6 +2565,16 @@ static inline void sample_erratic_flags(Drive *env, Agent *agent) { && sample_uniform(&env->rng_state, 0.0f, 1.0f) < env->episode_partner_blindness_prob) ? 1 : 0; + agent->partner_blindness_trigger_prob = env->partner_blindness_trigger_prob; + if (agent->is_blind_partner && env->partner_blindness_time_fraction_max > 0.0f) { + // blind-time fraction of a blind/sighted cycle = duration / (duration + mean gap), mean gap = (1 - p) / p + float duration_steps = (float) env->partner_blindness_duration; + float base_trigger_prob = env->partner_blindness_trigger_prob; + float base_fraction + = duration_steps * base_trigger_prob / (duration_steps * base_trigger_prob + 1.0f - base_trigger_prob); + float fraction = sample_uniform(&env->rng_state, base_fraction, env->partner_blindness_time_fraction_max); + agent->partner_blindness_trigger_prob = fraction / (fraction + duration_steps * (1.0f - fraction)); + } agent->is_phantom_braker = (env->episode_phantom_braking_prob > 0.0f && sample_uniform(&env->rng_state, 0.0f, 1.0f) < env->episode_phantom_braking_prob) ? 1 @@ -2770,7 +2804,6 @@ static void sample_zone_speed_limits(Drive *env) { } static void generate_traffic_light_states(Drive *env) { - int steps = env->scenario_length; int training_mode = !env->eval_mode || env->eval_training_render; // 20% chance: disable ALL lights for this episode @@ -2786,10 +2819,7 @@ static void generate_traffic_light_states(Drive *env) { continue; } - int fill_steps = steps; - if (tc->state_size < fill_steps) { - fill_steps = tc->state_size; - } + int fill_steps = tc->state_size; if (disable_all) { for (int t = 0; t < fill_steps; t++) { @@ -3655,7 +3685,8 @@ void init(Drive *env) { begin_episode_rng(env); sample_zone_speed_limits(env); if (env->simulation_mode == SIMULATION_MODE_GIGAFLOW) { - int steps = env->scenario_length; + int steps + = env->scenario_length + 1; // the terminal step's obs and red-light check read states[scenario_length] if (steps > 0) { for (int i = 0; i < env->num_traffic_elements; i++) { TrafficControlElement *traffic = &env->traffic_elements[i]; @@ -3800,6 +3831,106 @@ void c_get_global_agent_state( // pose. No side effects (no agent_log, no metrics_array, no infraction handling) -- shared by // compute_metrics (which owns everything side-effecting: infraction application, goal advancement, // eval-stat accumulation) and refresh_lane_association (cosim-only snapshot refresh) below. +// Squared distance from the agent to lane segment seg_idx; cross_out carries the left/right sign. +static inline float lane_segment_distance_sq( + const RoadMapElement *element, + int seg_idx, + float agent_x, + float agent_y, + float *cross_out) { + float seg_start_x = element->x[seg_idx]; + float seg_start_y = element->y[seg_idx]; + float seg_end_x = element->x[seg_idx + 1]; + float seg_end_y = element->y[seg_idx + 1]; + float seg_dx = seg_end_x - seg_start_x; + float seg_dy = seg_end_y - seg_start_y; + float seg_length_sq = seg_dx * seg_dx + seg_dy * seg_dy; + float to_agent_x = agent_x - seg_start_x; + float to_agent_y = agent_y - seg_start_y; + float cross = seg_dx * to_agent_y - seg_dy * to_agent_x; + *cross_out = cross; + if (seg_length_sq <= 1e-6f) { + return to_agent_x * to_agent_x + to_agent_y * to_agent_y; + } + float t = (to_agent_x * seg_dx + to_agent_y * seg_dy) / seg_length_sq; + if (t <= 0.0f) { + return to_agent_x * to_agent_x + to_agent_y * to_agent_y; + } + if (t >= 1.0f) { + float dxe = agent_x - seg_end_x; + float dye = agent_y - seg_end_y; + return dxe * dxe + dye * dye; + } + return (cross * cross) / seg_length_sq; +} + +// Closest segment in [seg_first, seg_last], lowest index on ties; min_dist_sq stays 1e18 if none passes the z filter. +static void scan_lane_segment_range( + const RoadMapElement *element, + const Agent *agent, + int seg_first, + int seg_last, + int *closest_seg_idx, + float *min_dist_sq, + float *closest_cross) { + *closest_seg_idx = 0; + *min_dist_sq = 1e18f; + *closest_cross = 0.0f; + for (int seg_idx = seg_first; seg_idx <= seg_last; seg_idx++) { + if (fabsf(element->z[seg_idx] - agent->sim_z) > Z_BUFFER) { + continue; + } + float cross; + float dist_sq = lane_segment_distance_sq(element, seg_idx, agent->sim_x, agent->sim_y, &cross); + if (dist_sq < *min_dist_sq) { + *min_dist_sq = dist_sq; + *closest_seg_idx = seg_idx; + *closest_cross = cross; + } + } +} + +// Distance from the agent to the nearest outer edge of its 5x5 road query window, meters. +static float road_query_window_margin_m(const Drive *env, const Agent *agent) { + float rel_x = agent->sim_x - env->grid_map->top_left_x; + float rel_y = agent->sim_y - env->grid_map->bottom_right_y; + float cell_x = (float) (int) (rel_x / GRID_CELL_SIZE); + float cell_y = (float) (int) (rel_y / GRID_CELL_SIZE); + float margin_x = fminf(rel_x - (cell_x - 2.0f) * GRID_CELL_SIZE, (cell_x + 3.0f) * GRID_CELL_SIZE - rel_x); + float margin_y = fminf(rel_y - (cell_y - 2.0f) * GRID_CELL_SIZE, (cell_y + 3.0f) * GRID_CELL_SIZE - rel_y); + return fminf(margin_x, margin_y); +} + +// Unlisted segments lie beyond window_margin - half_length; returns false when all segments exceed max_distance_m. +static bool find_closest_lane_segment( + const RoadMapElement *element, + const Agent *agent, + const LaneQueryEntry *entry, + float window_margin_m, + float max_distance_m, + int *closest_seg_idx, + float *min_dist_sq, + float *closest_cross) { + scan_lane_segment_range( + element, + agent, + entry->seg_first, + entry->seg_last, + closest_seg_idx, + min_dist_sq, + closest_cross); + float unlisted_reach_m = window_margin_m - element->max_segment_half_length_m - LANE_SCAN_GUARD_MARGIN_M; + float listed_dist_m = *min_dist_sq < 1e18f ? sqrtf(*min_dist_sq) : 1e9f; + if (listed_dist_m < unlisted_reach_m) { + return true; + } + if (unlisted_reach_m > max_distance_m && listed_dist_m > max_distance_m) { + return false; + } + scan_lane_segment_range(element, agent, 0, element->segment_size - 2, closest_seg_idx, min_dist_sq, closest_cross); + return true; +} + static void find_lane_and_offroad( Drive *env, Agent *agent, @@ -3816,38 +3947,30 @@ static void find_lane_and_offroad( int lane_seg_idx = 0; float signed_lane_distance = 0.0f, lane_heading = 0.0f; - GridMapEntity entity_list[ROAD_QUERY_ENTITY_COUNT]; - int list_size = get_neighbors_entities( - env, - agent->sim_x, - agent->sim_y, - entity_list, - ROAD_QUERY_ENTITY_COUNT, - ROAD_OFFSETS, - 25); - - if (list_size <= 0) { + GridMap *grid = env->grid_map; + int cell_idx = get_grid_index(env, agent->sim_x, agent->sim_y); + const LaneQueryEntry *entries = NULL; + int entry_count = 0; + float window_margin_m = 0.0f; + if (cell_idx == -1 || grid->window_entity_count[cell_idx] <= 0) { is_offroad = true; + } else { + entries = &grid->lane_query_pool[grid->lane_query_offset[cell_idx]]; + entry_count = grid->lane_query_offset[cell_idx + 1] - grid->lane_query_offset[cell_idx]; + window_margin_m = road_query_window_margin_m(env, agent); } // Vehicle-width based distance threshold (3x width) float max_distance_threshold = 3.0f * agent->sim_width; - int checked_lanes[MAX_CHECKED_LANES]; - int num_checked_lanes = 0; - - // Loop through road entities and compute associated metrics (offroad, lane alignment) - for (int i = 0; i < list_size; i++) { - if (entity_list[i].entity_idx == -1) { - continue; - } - - int entity_idx = entity_list[i].entity_idx; - int geometry_idx = entity_list[i].geometry_idx; + // Window candidates in grid order: road edges (offroad crossing) and drivable lanes (association) + for (int i = 0; i < entry_count; i++) { + const LaneQueryEntry *entry = &entries[i]; + int entity_idx = entry->entity_idx; RoadMapElement *element = &env->road_elements[entity_idx]; - // Check for offroad crossing with road edges if (is_road_edge(element->type)) { + int geometry_idx = entry->seg_first; float abs_dz = fabsf(element->z[geometry_idx] - agent->sim_z); if (abs_dz > Z_BUFFER) { continue; @@ -3858,71 +3981,29 @@ static void find_lane_and_offroad( element->x[geometry_idx + 1], element->y[geometry_idx + 1], agent); - } - - if (is_offroad) { - break; - } - - if (!is_drivable_road_lane(element->type)) { - continue; - } - - int already_checked = 0; - for (int c = 0; c < num_checked_lanes; c++) { - if (checked_lanes[c] == entity_idx) { - already_checked = 1; + if (is_offroad) { break; } - } - if (already_checked) { continue; } - if (num_checked_lanes < MAX_CHECKED_LANES) { - checked_lanes[num_checked_lanes++] = entity_idx; - } // Find closest segment on this lane (signed distance: left = negative, right = positive) int closest_seg_idx = 0; float signed_dist = 1e9f; int num_segments = element->segment_size - 1; if (num_segments >= 1) { - float min_dist_sq = 1e18f; - float closest_cross = 0.0f; - for (int seg_idx = 0; seg_idx < num_segments; seg_idx++) { - if (fabsf(element->z[seg_idx] - agent->sim_z) > Z_BUFFER) { - continue; - } - float seg_start_x = element->x[seg_idx]; - float seg_start_y = element->y[seg_idx]; - float seg_end_x = element->x[seg_idx + 1]; - float seg_end_y = element->y[seg_idx + 1]; - float seg_dx = seg_end_x - seg_start_x; - float seg_dy = seg_end_y - seg_start_y; - float seg_length_sq = seg_dx * seg_dx + seg_dy * seg_dy; - float to_agent_x = agent->sim_x - seg_start_x; - float to_agent_y = agent->sim_y - seg_start_y; - float cross = seg_dx * to_agent_y - seg_dy * to_agent_x; - float dist_sq; - if (seg_length_sq > 1e-6f) { - float t = (to_agent_x * seg_dx + to_agent_y * seg_dy) / seg_length_sq; - if (t <= 0.0f) { - dist_sq = to_agent_x * to_agent_x + to_agent_y * to_agent_y; - } else if (t >= 1.0f) { - float dxe = agent->sim_x - seg_end_x; - float dye = agent->sim_y - seg_end_y; - dist_sq = dxe * dxe + dye * dye; - } else { - dist_sq = (cross * cross) / seg_length_sq; - } - } else { - dist_sq = to_agent_x * to_agent_x + to_agent_y * to_agent_y; - } - if (dist_sq < min_dist_sq) { - min_dist_sq = dist_sq; - closest_seg_idx = seg_idx; - closest_cross = cross; - } + float min_dist_sq, closest_cross; + bool within_reach = find_closest_lane_segment( + element, + agent, + entry, + window_margin_m, + max_distance_threshold, + &closest_seg_idx, + &min_dist_sq, + &closest_cross); + if (!within_reach) { + continue; } float abs_dist_val = sqrtf(min_dist_sq); signed_dist = (closest_cross >= 0.0f) ? -abs_dist_val : abs_dist_val; @@ -4034,7 +4115,8 @@ int c_set_agent_states( const float *vy, const float *yaw_rate, const float *accel_long, - const float *seconds_stopped) { + const float *seconds_stopped, + const float *steering) { for (int k = 0; k < count; k++) { int agent_idx = idx[k]; if (agent_idx < 0 || agent_idx >= env->num_total_agents) { @@ -4046,6 +4128,9 @@ int c_set_agent_states( if (seconds_stopped && !(seconds_stopped[k] >= 0.0f)) { return -1; } + if (steering && !isfinite(steering[k])) { + return -1; + } Agent *agent = &env->agents[agent_idx]; agent->sim_x = x[k] - env->world_mean_x; agent->sim_y = y[k] - env->world_mean_y; @@ -4072,11 +4157,17 @@ int c_set_agent_states( // called, so "the agent's previous state" at this point is a throwaway dummy-action rollout, // not the true previous external state. agent->yaw_rate = yaw_rate[k]; - float speed_for_steering = fmaxf(fabsf(agent->sim_speed_signed), 1.0f); - float steering = atanf(agent->yaw_rate * agent->wheelbase / speed_for_steering); - agent->steering_angle = clip(steering, -STEERING_LIMIT, STEERING_LIMIT); + if (steering) { + agent->steering_angle = clip(steering[k], -STEERING_LIMIT, STEERING_LIMIT); + agent->accel_lat + = agent->sim_speed_signed * agent->sim_speed_signed * tanf(agent->steering_angle) / agent->wheelbase; + } else { + float speed_for_steering = fmaxf(fabsf(agent->sim_speed_signed), 1.0f); + float derived_steering = atanf(agent->yaw_rate * agent->wheelbase / speed_for_steering); + agent->steering_angle = clip(derived_steering, -STEERING_LIMIT, STEERING_LIMIT); + agent->accel_lat = agent->sim_speed_signed * agent->yaw_rate; + } agent->accel_long = accel_long[k]; - agent->accel_lat = agent->sim_speed_signed * agent->yaw_rate; refresh_lane_association(env, agent); // current_lane_idx / lane-dist / lane-angle for the new pose // NULL keeps c_step's own accumulation (co-sim default); an array injects stopped-time as state. @@ -4123,6 +4214,10 @@ int c_set_traffic_light_states(Drive *env, const int *states) { return -1; } traffic->states[ts] = states[i]; + // holds until the next sync: c_step's metrics and obs read timestep + 1 before that sync arrives + if (ts + 1 < traffic->state_size) { + traffic->states[ts + 1] = states[i]; + } } return 0; } @@ -4184,6 +4279,41 @@ int c_set_stop_signs(Drive *env, int count, const float *lines, const float *hea return 0; } +// Co-sim: move traffic-light stop lines onto the external sim's own (the lines its red-light judge scores against), +// indices = element index per line, lines = count x [x1, y1, z1, x2, y2, z2] in world coordinates. +int c_set_traffic_light_lines(Drive *env, int count, const int *indices, const float *lines) { + if (count < 0 || (count > 0 && (indices == NULL || lines == NULL))) { + return -1; + } + for (int k = 0; k < count; k++) { + int idx = indices[k]; + if (idx < 0 || idx >= env->num_traffic_elements + || env->traffic_elements[idx].type != TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT) { + return -1; + } + for (int v = 0; v < 6; v++) { + if (!isfinite(lines[k * 6 + v])) { + return -1; + } + } + float line_dx = lines[k * 6 + 3] - lines[k * 6]; + float line_dy = lines[k * 6 + 4] - lines[k * 6 + 1]; + if (line_dx * line_dx + line_dy * line_dy <= 0.0f) { + return -1; + } + } + for (int k = 0; k < count; k++) { + TrafficControlElement *tc = &env->traffic_elements[indices[k]]; + tc->stop_line[0] = lines[k * 6] - env->world_mean_x; + tc->stop_line[1] = lines[k * 6 + 1] - env->world_mean_y; + tc->stop_line[2] = lines[k * 6 + 2]; + tc->stop_line[3] = lines[k * 6 + 3] - env->world_mean_x; + tc->stop_line[4] = lines[k * 6 + 4] - env->world_mean_y; + tc->stop_line[5] = lines[k * 6 + 5]; + } + return 0; +} + void c_get_agent_goal_progress(Drive *env, int agent_idx, int *current_goal_idx_out, int *goal_count_out) { Agent *agent = &env->agents[agent_idx]; *current_goal_idx_out = agent->current_goal_idx; @@ -4285,27 +4415,28 @@ void c_get_road_edge_polylines(Drive *env, float *x_out, float *y_out, int *leng // Noise & Robustness Functions // ======================================== -static void subsample_road_observation_rows( - Rng *rng_state, - float *buffer, - int collected_count, - int keep_count, - int feature_count) { +// Partial Fisher-Yates over row indices; gathering rows in this order equals swapping the rows themselves. +static void subsample_road_observation_rows(Rng *rng_state, int *row_order, int collected_count, int keep_count) { if (keep_count <= 0 || collected_count <= keep_count) { return; } - float tmp[feature_count]; for (int sample_idx = 0; sample_idx < keep_count; sample_idx++) { int remaining = collected_count - sample_idx; int swap_idx = (remaining > 1) ? sample_idx + rng_below(rng_state, remaining) : sample_idx; - if (swap_idx == sample_idx) { - continue; - } - float *a = &buffer[sample_idx * feature_count]; - float *b = &buffer[swap_idx * feature_count]; - memcpy(tmp, a, sizeof(tmp)); - memcpy(a, b, sizeof(tmp)); - memcpy(b, tmp, sizeof(tmp)); + int sampled_row = row_order[sample_idx]; + row_order[sample_idx] = row_order[swap_idx]; + row_order[swap_idx] = sampled_row; + } +} + +static void gather_road_observation_rows( + float *dest, + const float *buffer, + const int *row_order, + int keep_count, + int feature_count) { + for (int k = 0; k < keep_count; k++) { + memcpy(&dest[k * feature_count], &buffer[row_order[k] * feature_count], feature_count * sizeof(float)); } } @@ -4595,7 +4726,7 @@ static void compute_rewards(Drive *env, int i) { } // Speed limit reward - float speed_reward = -agent->reward_coefs[REWARD_COEF_OVERSPEED] * agent->metrics_array[SPEED_LIMIT_IDX]; + float speed_reward = -agent->reward_coefs[REWARD_COEF_OVERSPEED] * env->dt * agent->metrics_array[SPEED_LIMIT_IDX]; env->rewards[i] += speed_reward; agent_log->reward_overspeed += speed_reward; @@ -4734,21 +4865,37 @@ static int write_reward_target_obs(Drive *env, Agent *ego, float *obs, int obs_i return obs_idx; } -static int write_partner_obs(Drive *env, Agent *ego, int agent_idx, float *obs, int obs_idx, int *partner_count) { - // Partner blindness: zero partner obs for the configured duration once triggered +typedef struct { + float x; + float y; + float z; + int agent_idx; +} PartnerScanEntry; + +static int write_partner_obs( + Drive *env, + Agent *ego, + int ego_agent_idx, + const PartnerScanEntry *partner_scan, + float *obs, + int obs_idx, + int *partner_count) { + // Blind steps hide all partners, or with partner_blindness_corridor only those outside the forward corridor if (ego->partner_blindness_counter > 0) { ego->partner_blindness_counter--; } - if (ego->partner_blindness_counter == 0 && ego->is_blind_partner && env->partner_blindness_trigger_prob > 0.0f - && sample_uniform(&env->rng_state, 0.0f, 1.0f) < env->partner_blindness_trigger_prob) { + if (ego->partner_blindness_counter == 0 && ego->is_blind_partner && ego->partner_blindness_trigger_prob > 0.0f + && sample_uniform(&env->rng_state, 0.0f, 1.0f) < ego->partner_blindness_trigger_prob) { ego->partner_blindness_counter = env->partner_blindness_duration; } - if (ego->partner_blindness_counter > 0) { + int blind = ego->partner_blindness_counter > 0; + if (blind && !env->partner_blindness_corridor) { int partner_obs_stride = env->obs_slots_partners_n * partner_feature_count(env); memset(&obs[obs_idx], 0, partner_obs_stride * sizeof(float)); *partner_count = 0; return obs_idx + partner_obs_stride; } + int corridor_only = blind; typedef struct { int index; @@ -4757,36 +4904,38 @@ static int write_partner_obs(Drive *env, Agent *ego, int agent_idx, float *obs, } AgentDistance; AgentDistance nearby_agents[env->num_agents]; int nearby_count = 0; + float range_sq = env->obs_range_partner_m * env->obs_range_partner_m; for (int j = 0; j < env->num_agents; j++) { - int index = -1; - if (j < env->active_agent_count) { - index = env->active_agent_indices[j]; - } else if (j < env->num_agents) { - index = env->static_agent_indices[j - env->active_agent_count]; - } - if (index == env->active_agent_indices[agent_idx]) { - continue; // Skip self, but don't increment obs_idx - } - Agent *other = &env->agents[index]; - float dx = other->sim_x - ego->sim_x; - float dy = other->sim_y - ego->sim_y; - float dz = other->sim_z - ego->sim_z; + float dx = partner_scan[j].x - ego->sim_x; + float dy = partner_scan[j].y - ego->sim_y; + float dz = partner_scan[j].z - ego->sim_z; float dist_sq = dx * dx + dy * dy + dz * dz; - if (dist_sq > env->obs_range_partner_m * env->obs_range_partner_m) { - continue; - } - nearby_agents[nearby_count].index = index; + int partner_visible = 1; + if (corridor_only) { + Agent *other = &env->agents[partner_scan[j].agent_idx]; + float rel_x, rel_y, rel_heading_x, rel_heading_y; + project_point_to_ego_frame(ego, other->sim_x, other->sim_y, &rel_x, &rel_y); + project_vector_to_ego_frame(ego, other->cos_heading, other->sin_heading, &rel_heading_x, &rel_heading_y); + float lateral_half_extent_m + = 0.5f * (fabsf(rel_heading_y) * other->sim_length + fabsf(rel_heading_x) * other->sim_width); + partner_visible = rel_x > 0.0f + && fabsf(rel_y) <= lateral_half_extent_m + 0.5f * ego->sim_width + PARTNER_BLINDNESS_CORRIDOR_MARGIN_M; + } + nearby_agents[nearby_count].index = partner_scan[j].agent_idx; nearby_agents[nearby_count].dist_sq = dist_sq; nearby_agents[nearby_count].dz = dz; - nearby_count++; + // branchless append: self, out-of-range and corridor-hidden candidates are overwritten by the next one + nearby_count += (partner_scan[j].agent_idx != ego_agent_idx) & !(dist_sq > range_sq) & partner_visible; } int partners_written = 0; int partners_to_write = (nearby_count < env->obs_slots_partners_n) ? nearby_count : env->obs_slots_partners_n; for (int k = 0; k < partners_to_write; k++) { int nearest_idx = k; + float nearest_dist_sq = nearby_agents[k].dist_sq; for (int j = k + 1; j < nearby_count; j++) { - if (nearby_agents[j].dist_sq < nearby_agents[nearest_idx].dist_sq) { + if (nearby_agents[j].dist_sq < nearest_dist_sq) { + nearest_dist_sq = nearby_agents[j].dist_sq; nearest_idx = j; } } @@ -4797,7 +4946,8 @@ static int write_partner_obs(Drive *env, Agent *ego, int agent_idx, float *obs, } } - for (int j = 0; j < partners_to_write; j++) { + int row_features = partner_feature_count(env); + for (int j = 0; j < partners_to_write && partners_written < env->obs_slots_partners_n; j++) { Agent *other = &env->agents[nearby_agents[j].index]; float rel_x, rel_y, rel_heading_x, rel_heading_y, rel_vx, rel_vy; project_point_to_ego_frame(ego, other->sim_x, other->sim_y, &rel_x, &rel_y); @@ -4811,21 +4961,37 @@ static int write_partner_obs(Drive *env, Agent *ego, int agent_idx, float *obs, &rel_vx, &rel_vy); project_vector_to_ego_frame(ego, other->cos_heading, other->sin_heading, &rel_heading_x, &rel_heading_y); - obs[obs_idx++] = rel_x / env->obs_norm_xy_offset_m; - obs[obs_idx++] = rel_y / env->obs_norm_xy_offset_m; + // float count: agent lengths are not range-checked at load, an int cast could overflow + float piece_count = 1.0f; + if (env->obs_partner_split_length_m > 0.0f && other->sim_length > env->obs_partner_split_length_m) { + piece_count = ceilf(other->sim_length / env->obs_partner_split_length_m); + } + float piece_length_m = other->sim_length / piece_count; + float rear_piece_offset_m = 0.5f * (piece_length_m - other->sim_length); + int row_start = obs_idx; + obs[obs_idx++] = (rel_x + rear_piece_offset_m * rel_heading_x) / env->obs_norm_xy_offset_m; + obs[obs_idx++] = (rel_y + rear_piece_offset_m * rel_heading_y) / env->obs_norm_xy_offset_m; obs[obs_idx++] = nearby_agents[j].dz / env->obs_norm_z_m; - obs[obs_idx++] = other->sim_length / env->obs_norm_veh_length_m; + obs[obs_idx++] = piece_length_m / env->obs_norm_veh_length_m; obs[obs_idx++] = other->sim_width / env->obs_norm_veh_width_m; obs[obs_idx++] = rel_heading_x; obs[obs_idx++] = rel_heading_y; obs[obs_idx++] = other->sim_speed_signed / env->obs_norm_speed_mps; - // TODO(hack): partner seconds_stopped is a temporary feature; remove later. obs[obs_idx++] = fminf(1.0f, other->seconds_stopped / MAX_STOPPED_SECONDS); if (env->obs_partner_relative_velocity) { obs[obs_idx++] = rel_vx / env->obs_norm_speed_mps; obs[obs_idx++] = rel_vy / env->obs_norm_speed_mps; } partners_written++; + for (int piece_idx = 1; (float) piece_idx < piece_count && partners_written < env->obs_slots_partners_n; + piece_idx++) { + float piece_offset_m = rear_piece_offset_m + (float) piece_idx * piece_length_m; + memcpy(&obs[obs_idx], &obs[row_start], row_features * sizeof(float)); + obs[obs_idx] = (rel_x + piece_offset_m * rel_heading_x) / env->obs_norm_xy_offset_m; + obs[obs_idx + 1] = (rel_y + piece_offset_m * rel_heading_y) / env->obs_norm_xy_offset_m; + obs_idx += row_features; + partners_written++; + } } *partner_count = partners_written; @@ -4841,7 +5007,7 @@ static float lane_graph_distance_to_goal_m(Drive *env, int lane_idx, int goal_gr if (lane_graph_idx < 0) { return -1.0f; } - float distance_m = env->lane_graph.distances[lane_graph_idx * env->lane_graph.n_lanes + goal_graph_idx]; + float distance_m = env->lane_graph.distances[goal_graph_idx * env->lane_graph.n_lanes + lane_graph_idx]; return (isfinite(distance_m) && distance_m >= 0.0f) ? distance_m : INFINITY; } @@ -4884,10 +5050,7 @@ static int write_road_obs_rows( entity != NULL && rows_written < slot_cap; entity = neighbor_cursor_next(env, neighbor_cursor)) { int entity_idx = entity->entity_idx; - int geometry_idx = entity->geometry_idx; - RoadMapElement *road_element = &env->road_elements[entity_idx]; - - const RoadObservationSegment *segment = &road_element->observation_segments[geometry_idx]; + const RoadObservationSegment *segment = &env->grid_map->segment_pool[entity->pool_idx]; float rel_x, rel_y; float rel_z = segment->mid_z_m - ego->sim_z; project_point_to_ego_frame(ego, segment->mid_x_m, segment->mid_y_m, &rel_x, &rel_y); @@ -4973,19 +5136,27 @@ static int write_road_obs(Drive *env, Agent *ego, float *obs, int obs_idx, int * = (boundaries_found < env->obs_slots_boundary_kept) ? boundaries_found : env->obs_slots_boundary_kept; *lane_count = lanes_to_copy; *boundary_count = boundaries_to_copy; - subsample_road_observation_rows(&env->rng_state, lanes_buffer, lanes_found, lanes_to_copy, lane_row_features); - subsample_road_observation_rows( - &env->rng_state, - boundaries_buffer, - boundaries_found, - boundaries_to_copy, - BOUNDARY_FEATURES); - memcpy(&obs[lane_obs_idx], lanes_buffer, lanes_to_copy * lane_row_features * sizeof(float)); + int lane_order[env->obs_slots_lane_n]; + int boundary_order[env->obs_slots_boundary_n]; + for (int k = 0; k < lanes_found; k++) { + lane_order[k] = k; + } + for (int k = 0; k < boundaries_found; k++) { + boundary_order[k] = k; + } + subsample_road_observation_rows(&env->rng_state, lane_order, lanes_found, lanes_to_copy); + subsample_road_observation_rows(&env->rng_state, boundary_order, boundaries_found, boundaries_to_copy); + gather_road_observation_rows(&obs[lane_obs_idx], lanes_buffer, lane_order, lanes_to_copy, lane_row_features); memset( &obs[lane_obs_idx + lanes_to_copy * lane_row_features], 0, (env->obs_slots_lane_kept - lanes_to_copy) * lane_row_features * sizeof(float)); - memcpy(&obs[boundary_obs_idx], boundaries_buffer, boundaries_to_copy * BOUNDARY_FEATURES * sizeof(float)); + gather_road_observation_rows( + &obs[boundary_obs_idx], + boundaries_buffer, + boundary_order, + boundaries_to_copy, + BOUNDARY_FEATURES); memset( &obs[boundary_obs_idx + boundaries_to_copy * BOUNDARY_FEATURES], 0, @@ -5041,8 +5212,10 @@ static int write_traffic_control_obs(Drive *env, Agent *ego, float *obs, int obs = (visible_count < env->obs_slots_traffic_controls_n) ? visible_count : env->obs_slots_traffic_controls_n; for (int k = 0; k < controls_to_observe; k++) { int nearest_idx = k; + float nearest_dist_sq = visible_controls[k].dist_sq; for (int j = k + 1; j < visible_count; j++) { - if (visible_controls[j].dist_sq < visible_controls[nearest_idx].dist_sq) { + if (visible_controls[j].dist_sq < nearest_dist_sq) { + nearest_dist_sq = visible_controls[j].dist_sq; nearest_idx = j; } } @@ -5088,6 +5261,13 @@ static void compute_observations(Drive *env) { int obs_per_agent = compute_observation_size(env); memset(env->observations, 0, obs_per_agent * env->active_agent_count * sizeof(float)); float (*obs_matrix)[obs_per_agent] = (float (*)[obs_per_agent]) env->observations; + PartnerScanEntry partner_scan[env->num_agents]; + for (int j = 0; j < env->num_agents; j++) { + int index = (j < env->active_agent_count) ? env->active_agent_indices[j] + : env->static_agent_indices[j - env->active_agent_count]; + Agent *agent = &env->agents[index]; + partner_scan[j] = (PartnerScanEntry) {agent->sim_x, agent->sim_y, agent->sim_z, index}; + } for (int i = 0; i < env->active_agent_count; i++) { float *obs = &obs_matrix[i][0]; int agent_idx = env->active_agent_indices[i]; @@ -5100,7 +5280,7 @@ static void compute_observations(Drive *env) { obs_idx = write_ego_obs(env, ego, obs, obs_idx); obs_idx = write_reward_target_obs(env, ego, obs, obs_idx); - obs_idx = write_partner_obs(env, ego, i, obs, obs_idx, &partner_count); + obs_idx = write_partner_obs(env, ego, agent_idx, partner_scan, obs, obs_idx, &partner_count); obs_idx = write_road_obs(env, ego, obs, obs_idx, &lane_count, &boundary_count); obs_idx = write_traffic_control_obs(env, ego, obs, obs_idx, &traffic_control_count); obs[obs_idx++] = (float) lane_count; @@ -5229,7 +5409,7 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { int j_lat_idx = action_val % num_lat; j_long = JERK_LONG[j_long_idx]; j_lat = JERK_LAT[j_lat_idx]; - } else if (env->action_type == ACTION_TYPE_CONTINUOUS) { + } else { float (*action_array_f)[2] = (float (*)[2]) env->actions; // Asymmetric scaling for longitudinal jerk to match discrete action space // Discrete: JERK_LONG = [-15, -4, 0, 4] (more braking than acceleration) @@ -5269,6 +5449,8 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { } else { a_long_new = clip(a_long_new, ACCEL_LONG_LIMIT[0], ACCEL_LONG_LIMIT[1] * c_acc); } + // Comfort scores the driver's jerk, so take it before the stop resets below zero the stored acceleration + float jerk_long_before_stop_reset = (a_long_new - agent->accel_long) / env->dt; // Calculate new lateral acceleration from jerk (Eq. 2 in paper) float a_lat_new = agent->accel_lat + c_steer * j_lat * env->dt; @@ -5356,7 +5538,7 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { update_agent_speed(agent); // Update jerk and acceleration - agent->jerk_long = (a_long_new - agent->accel_long) / env->dt; + agent->jerk_long = jerk_long_before_stop_reset; agent->jerk_lat = (a_lat_new - agent->accel_lat) / env->dt; agent->accel_long = a_long_new; agent->accel_lat = a_lat_new; diff --git a/pufferlib/ocean/drive/drive.py b/pufferlib/ocean/drive/drive.py index a0783b0670..0992f5980f 100644 --- a/pufferlib/ocean/drive/drive.py +++ b/pufferlib/ocean/drive/drive.py @@ -33,7 +33,7 @@ def __init__( reward_reverse=0.005, reward_stop_line=1.0, reward_timestep=0.000025, - reward_overspeed=0.05, + reward_overspeed=0.005, overspeed_tolerance_mps=0.0, reward_ade=0.0, min_goal_spacing=20.0, @@ -118,6 +118,7 @@ def __init__( obs_slots_boundary_n=32, obs_lane_stride=1, obs_boundary_stride=1, + obs_stride_merge_segments=False, obs_slots_partners_n=16, obs_partner_relative_velocity=False, obs_lane_heading_signed=False, @@ -136,6 +137,7 @@ def __init__( obs_norm_road_seg_width_m=5.0, obs_norm_z_m=10.0, eval_perceived_size_margin_m=0.1, + obs_partner_split_length_m=0.0, eval_standstill_jerk_deadband_mps3=0.0, obs_range_traffic_control_m=100.0, obs_range_partner_m=100.0, @@ -145,8 +147,11 @@ def __init__( obs_dropout_lane=0.0, obs_dropout_boundary=0.0, partner_blindness_prob=0.0, + partner_blindness_prob_min=0.0, partner_blindness_trigger_prob=0.1, partner_blindness_duration_seconds=1.0, + partner_blindness_corridor=True, + partner_blindness_time_fraction_max=0.0, phantom_braking_prob=0.0, phantom_braking_trigger_prob=0.0, phantom_braking_duration_seconds=1.0, @@ -304,6 +309,7 @@ def __init__( self.obs_slots_boundary_n = obs_slots_boundary_n self.obs_lane_stride = obs_lane_stride self.obs_boundary_stride = obs_boundary_stride + self.obs_stride_merge_segments = int(bool(obs_stride_merge_segments)) self.obs_slots_partners_n = obs_slots_partners_n self.obs_partner_relative_velocity = int(bool(obs_partner_relative_velocity)) self.obs_lane_heading_signed = int(bool(obs_lane_heading_signed)) @@ -323,6 +329,9 @@ def __init__( self.obs_norm_road_seg_width_m = float(obs_norm_road_seg_width_m) self.obs_norm_z_m = float(obs_norm_z_m) self.eval_perceived_size_margin_m = float(eval_perceived_size_margin_m) + self.obs_partner_split_length_m = float(obs_partner_split_length_m) + if not np.isfinite(self.obs_partner_split_length_m) or self.obs_partner_split_length_m < 0.0: + raise ValueError(f"obs_partner_split_length_m must be finite and >= 0. Got: {obs_partner_split_length_m}") self.eval_standstill_jerk_deadband_mps3 = float(eval_standstill_jerk_deadband_mps3) if self.eval_standstill_jerk_deadband_mps3 < 0: raise ValueError( @@ -344,8 +353,11 @@ def __init__( self.obs_dropout_boundary, ) self.partner_blindness_prob = float(partner_blindness_prob) + self.partner_blindness_prob_min = float(partner_blindness_prob_min) self.partner_blindness_trigger_prob = float(partner_blindness_trigger_prob) self.partner_blindness_duration_seconds = float(partner_blindness_duration_seconds) + self.partner_blindness_corridor = int(bool(partner_blindness_corridor)) + self.partner_blindness_time_fraction_max = float(partner_blindness_time_fraction_max) self.phantom_braking_prob = float(phantom_braking_prob) self.phantom_braking_trigger_prob = float(phantom_braking_trigger_prob) self.phantom_braking_duration_seconds = float(phantom_braking_duration_seconds) @@ -524,6 +536,7 @@ def preload_shared_resources(self): "use_neighbor_cache": self.use_neighbor_cache, "obs_lane_stride": self.obs_lane_stride, "obs_boundary_stride": self.obs_boundary_stride, + "obs_stride_merge_segments": self.obs_stride_merge_segments, "obs_range_road_front_m": self.obs_range_road_front_m, "obs_range_road_behind_m": self.obs_range_road_behind_m, "obs_range_road_side_m": self.obs_range_road_side_m, @@ -574,6 +587,7 @@ def _env_init_kwargs(self, map_file, max_agents): "obs_slots_boundary_n": self.obs_slots_boundary_n, "obs_lane_stride": self.obs_lane_stride, "obs_boundary_stride": self.obs_boundary_stride, + "obs_stride_merge_segments": self.obs_stride_merge_segments, "obs_slots_partners_n": self.obs_slots_partners_n, "obs_partner_relative_velocity": self.obs_partner_relative_velocity, "obs_lane_heading_signed": self.obs_lane_heading_signed, @@ -634,6 +648,7 @@ def _env_init_kwargs(self, map_file, max_agents): "obs_norm_road_seg_width_m": self.obs_norm_road_seg_width_m, "obs_norm_z_m": self.obs_norm_z_m, "eval_perceived_size_margin_m": self.eval_perceived_size_margin_m, + "obs_partner_split_length_m": self.obs_partner_split_length_m, "eval_standstill_jerk_deadband_mps3": self.eval_standstill_jerk_deadband_mps3, "obs_range_traffic_control_m": self.obs_range_traffic_control_m, "obs_range_partner_m": self.obs_range_partner_m, @@ -643,8 +658,11 @@ def _env_init_kwargs(self, map_file, max_agents): "obs_slots_lane_kept": self.obs_slots_lane_kept, "obs_slots_boundary_kept": self.obs_slots_boundary_kept, "partner_blindness_prob": self.partner_blindness_prob, + "partner_blindness_prob_min": self.partner_blindness_prob_min, "partner_blindness_trigger_prob": self.partner_blindness_trigger_prob, "partner_blindness_duration_seconds": self.partner_blindness_duration_seconds, + "partner_blindness_corridor": self.partner_blindness_corridor, + "partner_blindness_time_fraction_max": self.partner_blindness_time_fraction_max, "phantom_braking_prob": self.phantom_braking_prob, "phantom_braking_trigger_prob": self.phantom_braking_trigger_prob, "phantom_braking_duration_seconds": self.phantom_braking_duration_seconds, @@ -832,7 +850,9 @@ def get_global_agent_state(self, include_static=False): return states - def set_agent_states(self, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long, seconds_stopped=None): + def set_agent_states( + self, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long, seconds_stopped=None, steering=None + ): """Co-sim: overwrite the sim state of agents at global indices `idx` (e.g. CARLA background) with world-frame pose/velocity. The C side subtracts world_mean, recaches heading trig and recomputes speed. @@ -858,6 +878,7 @@ def set_agent_states(self, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long, np.ascontiguousarray(yaw_rate, dtype=np.float32), np.ascontiguousarray(accel_long, dtype=np.float32), None if seconds_stopped is None else np.ascontiguousarray(seconds_stopped, dtype=np.float32), + None if steering is None else np.ascontiguousarray(steering, dtype=np.float32), ) # ── Co-simulation external-state setters ───────────────────────────────────── @@ -893,6 +914,14 @@ def set_stop_signs(self, lines, headings): headings = np.ascontiguousarray(np.asarray(headings, dtype=np.float32).reshape(-1)) return int(binding.vec_set_stop_signs(self.c_envs, lines, headings)) + def set_traffic_light_lines(self, indices, lines): + """move traffic-light elements' stop lines onto the external sim's own + (co-sim): indices (K,) element indices, lines (K, 6) world-frame + endpoints [x1, y1, z1, x2, y2, z2].""" + indices = np.ascontiguousarray(np.asarray(indices, dtype=np.int32).reshape(-1)) + lines = np.ascontiguousarray(np.asarray(lines, dtype=np.float32).reshape(-1, 6)) + binding.vec_set_traffic_light_lines(self.c_envs, indices, lines) + def set_agent_goals(self, agent_idx, gx, gy, gz, gdir_x=None, gdir_y=None): """set an agent's goal waypoints (e.g. the ego's route) in world coords (C subtracts world_mean). Each waypoint is snapped to its diff --git a/pufferlib/ocean/drive/map_data.h b/pufferlib/ocean/drive/map_data.h index 8254dc4fbb..3b79c415ca 100644 --- a/pufferlib/ocean/drive/map_data.h +++ b/pufferlib/ocean/drive/map_data.h @@ -35,10 +35,6 @@ static void add_entity_to_grid( int geometry_idx, int obs_kind, int *cell_entities_insert_index) { - if (grid_index == -1) { - return; - } - int count = cell_entities_insert_index[grid_index]; if (count >= env->grid_map->cell_entities_count[grid_index]) { printf( @@ -50,12 +46,174 @@ static void add_entity_to_grid( return; } + int pool_idx = (int) (env->grid_map->cells[grid_index] - env->grid_map->entity_pool) + count; env->grid_map->cells[grid_index][count].entity_idx = entity_idx; env->grid_map->cells[grid_index][count].geometry_idx = geometry_idx; env->grid_map->cells[grid_index][count].obs_kind = obs_kind; + env->grid_map->cells[grid_index][count].pool_idx = pool_idx; + env->grid_map->segment_pool[pool_idx] = env->road_elements[entity_idx].observation_segments[geometry_idx]; cell_entities_insert_index[grid_index] = count + 1; } +// Appends one window candidate; drivable lanes already present only widen their listed segment range. +static int append_lane_query_entry( + LaneQueryEntry *entries, + int count, + int first_entry, + int entity_idx, + int seg_idx, + int is_lane) { + if (is_lane) { + for (int e = first_entry; e < count; e++) { + if (entries[e].entity_idx != entity_idx) { + continue; + } + entries[e].seg_first = seg_idx < entries[e].seg_first ? seg_idx : entries[e].seg_first; + entries[e].seg_last = seg_idx > entries[e].seg_last ? seg_idx : entries[e].seg_last; + return count; + } + } + entries[count] = (LaneQueryEntry) {entity_idx, seg_idx, seg_idx}; + return count + 1; +} + +static int road_query_window_entity_count(const GridMap *grid, int cell_idx) { + int cell_x = cell_idx % grid->grid_cols; + int cell_y = cell_idx / grid->grid_cols; + int total = 0; + for (int o = 0; o < ROAD_QUERY_CELL_COUNT; o++) { + int nx = cell_x + ROAD_OFFSETS[o][0]; + int ny = cell_y + ROAD_OFFSETS[o][1]; + if (nx < 0 || nx >= grid->grid_cols || ny < 0 || ny >= grid->grid_rows) { + continue; + } + total += grid->cell_entities_count[ny * grid->grid_cols + nx]; + } + return total; +} + +// Appends the cell's window candidates in grid order starting at entries[first_entry]; returns the new count. +static int append_cell_lane_query_entries(const Drive *env, int cell_idx, LaneQueryEntry *entries, int first_entry) { + const GridMap *grid = env->grid_map; + int cell_x = cell_idx % grid->grid_cols; + int cell_y = cell_idx / grid->grid_cols; + int count = first_entry; + for (int o = 0; o < ROAD_QUERY_CELL_COUNT; o++) { + int nx = cell_x + ROAD_OFFSETS[o][0]; + int ny = cell_y + ROAD_OFFSETS[o][1]; + if (nx < 0 || nx >= grid->grid_cols || ny < 0 || ny >= grid->grid_rows) { + continue; + } + int neighbor_idx = ny * grid->grid_cols + nx; + for (int e = 0; e < grid->cell_entities_count[neighbor_idx]; e++) { + const GridMapEntity *entity = &grid->cells[neighbor_idx][e]; + int type = env->road_elements[entity->entity_idx].type; + if (!is_road_edge(type) && !is_drivable_road_lane(type)) { + continue; + } + count = append_lane_query_entry( + entries, + count, + first_entry, + entity->entity_idx, + entity->geometry_idx, + is_drivable_road_lane(type)); + } + } + return count; +} + +// Per cell, the road-query candidates of its 5x5 window in grid order, built once per map. +static void cache_lane_query_lists(Drive *env) { + GridMap *grid = env->grid_map; + int cell_count = grid->grid_cols * grid->grid_rows; + int capacity = grid->total_entities > 0 ? grid->total_entities : 1; + LaneQueryEntry *entries = (LaneQueryEntry *) malloc(capacity * sizeof(LaneQueryEntry)); + grid->lane_query_offset = (int *) malloc((cell_count + 1) * sizeof(int)); + grid->window_entity_count = (int *) malloc(cell_count * sizeof(int)); + if (entries == NULL || grid->lane_query_offset == NULL || grid->window_entity_count == NULL) { + raise_error_with_message(ERROR_MEMORY_ALLOCATION, "Failed to allocate the road query cache"); + } + int count = 0; + for (int cell_idx = 0; cell_idx < cell_count; cell_idx++) { + grid->lane_query_offset[cell_idx] = count; + int window_count = road_query_window_entity_count(grid, cell_idx); + grid->window_entity_count[cell_idx] = window_count; + if (count + window_count > capacity) { + capacity = 2 * (count + window_count); + entries = (LaneQueryEntry *) realloc(entries, capacity * sizeof(LaneQueryEntry)); + if (entries == NULL) { + raise_error_with_message(ERROR_MEMORY_ALLOCATION, "Failed to grow the road query cache"); + } + } + count = append_cell_lane_query_entries(env, cell_idx, entries, count); + } + grid->lane_query_offset[cell_count] = count; + grid->lane_query_pool = entries; +} + +static void set_observation_segment_geometry( + RoadMapElement *road, + int start_point_idx, + int end_point_idx, + int carrier_idx) { + RoadObservationSegment *segment = &road->observation_segments[carrier_idx]; + // Preserve the observation formulas and float rounding, including zero-length segments. + segment->mid_x_m = (road->x[start_point_idx] + road->x[end_point_idx]) / 2.0f; + segment->mid_y_m = (road->y[start_point_idx] + road->y[end_point_idx]) / 2.0f; + segment->mid_z_m = (road->z[start_point_idx] + road->z[end_point_idx]) / 2.0f; + float half_dx_m = road->x[end_point_idx] - segment->mid_x_m; + float half_dy_m = road->y[end_point_idx] - segment->mid_y_m; + segment->half_length_m = sqrtf(half_dx_m * half_dx_m + half_dy_m * half_dy_m); + segment->direction_x = segment->half_length_m > 0.0f ? half_dx_m / segment->half_length_m : half_dx_m; + segment->direction_y = segment->half_length_m > 0.0f ? half_dy_m / segment->half_length_m : half_dy_m; + segment->width_m + = is_road_lane(road->type) ? 0.5f * (road->widths[start_point_idx] + road->widths[end_point_idx]) : 0.0f; +} + +// Legacy stride: every obs_stride-th segment is observed with its own extent, so the skipped ones leave gaps. +static void mark_strided_observation_segments(RoadMapElement *road, int obs_stride) { + int last_kept_idx = 0; + for (int segment_idx = 0; segment_idx < road->segment_size - 1; segment_idx++) { + int observed = 1; + if (obs_stride > 1 && segment_idx > 0) { + float heading_dev_rad + = fabsf(normalize_heading(road->headings[segment_idx] - road->headings[last_kept_idx])); + observed = segment_idx - last_kept_idx >= obs_stride || heading_dev_rad > OBS_STRIDE_HEADING_THRESHOLD; + } + if (observed) { + last_kept_idx = segment_idx; + } + road->observation_segments[segment_idx].observed = observed; + } +} + +// One chord from the run's first point to its last point, carried by the run's middle segment. +static void close_observation_segment_run(RoadMapElement *road, int first_segment_idx, int last_segment_idx) { + int carrier_idx = (first_segment_idx + last_segment_idx) / 2; + set_observation_segment_geometry(road, first_segment_idx, last_segment_idx + 1, carrier_idx); + road->observation_segments[carrier_idx].observed = 1; +} + +// Runs of up to obs_stride near-collinear segments become one chord, so the polyline stays connected. +static void merge_strided_observation_segments(RoadMapElement *road, int obs_stride) { + int segment_count = road->segment_size - 1; + int run_start_idx = 0; + for (int segment_idx = 0; segment_idx < segment_count; segment_idx++) { + if (segment_idx > run_start_idx + && fabsf(normalize_heading(road->headings[segment_idx] - road->headings[run_start_idx])) + > OBS_STRIDE_HEADING_THRESHOLD) { + close_observation_segment_run(road, run_start_idx, segment_idx - 1); + run_start_idx = segment_idx; + } + if (segment_idx - run_start_idx + 1 < obs_stride && segment_idx < segment_count - 1) { + continue; + } + close_observation_segment_run(road, run_start_idx, segment_idx); + run_start_idx = segment_idx + 1; + } +} + // Build after loading lane widths and before sharing the immutable road elements. static void cache_road_observation_geometry(Drive *env) { for (int road_idx = 0; road_idx < env->num_road_elements; road_idx++) { @@ -77,20 +235,14 @@ static void cache_road_observation_geometry(Drive *env) { if (road->observation_segments == NULL) { raise_error_with_message(ERROR_MEMORY_ALLOCATION, "Failed to cache geometry for road %d", road_idx); } - int has_lane_width = is_road_lane(road->type); for (int segment_idx = 0; segment_idx < road->segment_size - 1; segment_idx++) { - RoadObservationSegment *segment = &road->observation_segments[segment_idx]; - // Preserve the observation formulas and float rounding, including zero-length segments. - segment->mid_x_m = (road->x[segment_idx] + road->x[segment_idx + 1]) / 2.0f; - segment->mid_y_m = (road->y[segment_idx] + road->y[segment_idx + 1]) / 2.0f; - segment->mid_z_m = (road->z[segment_idx] + road->z[segment_idx + 1]) / 2.0f; - float half_dx_m = road->x[segment_idx + 1] - segment->mid_x_m; - float half_dy_m = road->y[segment_idx + 1] - segment->mid_y_m; - segment->half_length_m = sqrtf(half_dx_m * half_dx_m + half_dy_m * half_dy_m); - segment->direction_x = segment->half_length_m > 0.0f ? half_dx_m / segment->half_length_m : half_dx_m; - segment->direction_y = segment->half_length_m > 0.0f ? half_dy_m / segment->half_length_m : half_dy_m; - segment->width_m - = has_lane_width ? 0.5f * (road->widths[segment_idx] + road->widths[segment_idx + 1]) : 0.0f; + set_observation_segment_geometry(road, segment_idx, segment_idx + 1, segment_idx); + } + int obs_stride = is_road_lane(road->type) ? env->obs_lane_stride : env->obs_boundary_stride; + if (env->obs_stride_merge_segments) { + merge_strided_observation_segments(road, obs_stride); + } else { + mark_strided_observation_segments(road, obs_stride); } } } @@ -143,8 +295,9 @@ static int init_grid_map(Drive *env) { env->grid_map = NULL; return 1; } - env->grid_map->grid_cols = ceil(grid_width / GRID_CELL_SIZE); - env->grid_map->grid_rows = ceil(grid_height / GRID_CELL_SIZE); + // +1 so a point on the far edge indexes the last column/row instead of falling outside + env->grid_map->grid_cols = (int) (grid_width / GRID_CELL_SIZE) + 1; + env->grid_map->grid_rows = (int) (grid_height / GRID_CELL_SIZE) + 1; long long grid_cell_count_wide = (long long) env->grid_map->grid_cols * (long long) env->grid_map->grid_rows; if (env->grid_map->grid_cols < 1 || env->grid_map->grid_rows < 1 || grid_cell_count_wide > MAX_GRID_CELL_COUNT) { fprintf( @@ -167,12 +320,22 @@ static int init_grid_map(Drive *env) { continue; } RoadMapElement *element = &env->road_elements[i]; + element->max_segment_half_length_m = 0.0f; for (int j = 0; j < element->segment_size - 1; j++) { + float half_dx = (element->x[j + 1] - element->x[j]) * 0.5f; + float half_dy = (element->y[j + 1] - element->y[j]) * 0.5f; + float half_length_m = sqrtf(half_dx * half_dx + half_dy * half_dy); + element->max_segment_half_length_m = fmaxf(element->max_segment_half_length_m, half_length_m); float x_center = (element->x[j] + element->x[j + 1]) / 2; float y_center = (element->y[j] + element->y[j + 1]) / 2; int grid_index = get_grid_index(env, x_center, y_center); if (grid_index == -1) { - continue; + fprintf(stderr, "[ERROR] -> Road %d segment %d lies outside the grid extent\n", i, j); + free(env->grid_map->cells); + free(env->grid_map->cell_entities_count); + free(env->grid_map); + env->grid_map = NULL; + return 1; } env->grid_map->cell_entities_count[grid_index]++; } @@ -195,6 +358,8 @@ static int init_grid_map(Drive *env) { env->grid_map->total_entities += env->grid_map->cell_entities_count[grid_index]; } env->grid_map->entity_pool = (GridMapEntity *) calloc(env->grid_map->total_entities + 1, sizeof(GridMapEntity)); + env->grid_map->segment_pool + = (RoadObservationSegment *) calloc(env->grid_map->total_entities + 1, sizeof(RoadObservationSegment)); int pool_offset = 0; for (int grid_index = 0; grid_index < grid_cell_count; grid_index++) { int count = env->grid_map->cell_entities_count[grid_index]; @@ -209,32 +374,14 @@ static int init_grid_map(Drive *env) { } RoadMapElement *element = &env->road_elements[i]; int is_lane = is_road_lane(element->type); - int obs_stride = 1; - if (is_lane) { - obs_stride = env->obs_lane_stride; - } else if (is_road_edge(element->type)) { - obs_stride = env->obs_boundary_stride; - } - int last_kept_idx = 0; for (int j = 0; j < element->segment_size - 1; j++) { - // Keep a point every obs_stride points, plus wherever heading deviates enough - // since the last kept point (densifies curves/intersections) - int valid_for_obs = 1; - if (obs_stride > 1 && j > 0) { - float heading_dev = fabsf(normalize_heading(element->headings[j] - element->headings[last_kept_idx])); - valid_for_obs = j - last_kept_idx >= obs_stride || heading_dev > OBS_STRIDE_HEADING_THRESHOLD; - } int obs_kind = OBS_ENTITY_NONE; - if (valid_for_obs) { - last_kept_idx = j; + if (element->observation_segments[j].observed) { obs_kind = is_lane ? OBS_ENTITY_LANE : OBS_ENTITY_EDGE; } float x_center = (element->x[j] + element->x[j + 1]) / 2; float y_center = (element->y[j] + element->y[j + 1]) / 2; int grid_index = get_grid_index(env, x_center, y_center); - if (grid_index == -1) { - continue; - } add_entity_to_grid(env, grid_index, i, j, obs_kind, cell_entities_insert_index); if (is_drivable_road_lane(element->type) && !drivable_grid_seen[grid_index]) { drivable_grid_seen[grid_index] = true; @@ -252,6 +399,7 @@ static int init_grid_map(Drive *env) { } free(drivable_grid_seen); free(cell_entities_insert_index); + cache_lane_query_lists(env); return 0; } @@ -775,6 +923,14 @@ int load_map_binary(const char *filename, Drive *drive) { fclose(file); return -1; } + for (int p = 0; p < slen; p++) { + if (!isfinite(road->x[p]) || !isfinite(road->y[p]) || !isfinite(road->z[p]) + || !isfinite(road->headings[p])) { + printf("[ERROR] -> Road %d point %d has a non-finite coordinate or heading.\n", i, p); + fclose(file); + return -1; + } + } if (is_road_lane(road->type)) { if (fread(&road->num_entries, sizeof(int), 1, file) != 1) { @@ -921,6 +1077,15 @@ int load_map_binary(const char *filename, Drive *drive) { fclose(file); return -1; } + // File stores [from][to]; the lane observations read one goal row per ego, so keep [to][from]. + for (int from = 0; from < n_lanes_graph; from++) { + for (int to = from + 1; to < n_lanes_graph; to++) { + float from_to = drive->lane_graph.distances[from * n_lanes_graph + to]; + drive->lane_graph.distances[from * n_lanes_graph + to] + = drive->lane_graph.distances[to * n_lanes_graph + from]; + drive->lane_graph.distances[to * n_lanes_graph + from] = from_to; + } + } // Build reverse lookup road-element idx -> graph idx int num_roads = drive->num_road_elements; @@ -1009,7 +1174,8 @@ static struct SharedMapData *map_cache_lookup(Drive *env) { for (int i = 0; i < g_map_cache_count; i++) { if (g_map_cache[i] != NULL && strcmp(g_map_cache[i]->map_name, env->map_name) == 0 && g_map_cache[i]->obs_lane_stride == env->obs_lane_stride - && g_map_cache[i]->obs_boundary_stride == env->obs_boundary_stride) { + && g_map_cache[i]->obs_boundary_stride == env->obs_boundary_stride + && g_map_cache[i]->obs_stride_merge_segments == env->obs_stride_merge_segments) { return g_map_cache[i]; } } @@ -1039,6 +1205,7 @@ static struct SharedMapData *map_cache_store(Drive *env) { entry->lane_graph = env->lane_graph; entry->obs_lane_stride = env->obs_lane_stride; entry->obs_boundary_stride = env->obs_boundary_stride; + entry->obs_stride_merge_segments = env->obs_stride_merge_segments; entry->ref_count = 1; entry->owner_pid = getpid(); map_cache_insert(entry); @@ -1051,6 +1218,10 @@ static void free_grid_map(GridMap *grid_map) { } int grid_cell_count = grid_map->grid_cols * grid_map->grid_rows; free(grid_map->entity_pool); + free(grid_map->segment_pool); + free(grid_map->lane_query_pool); + free(grid_map->lane_query_offset); + free(grid_map->window_entity_count); free(grid_map->cells); free(grid_map->cell_entities_count); free(grid_map->grid_index_drivable); diff --git a/pufferlib/ocean/env_binding.h b/pufferlib/ocean/env_binding.h index cedb4c1215..0c6b91f9f4 100644 --- a/pufferlib/ocean/env_binding.h +++ b/pufferlib/ocean/env_binding.h @@ -987,8 +987,8 @@ static PyObject *vec_get_global_agent_state(PyObject *self, PyObject *args) { // ── Co-simulation external-state setters (mirror vec_get_global_agent_state) ── // Co-sim runs a single env (num_envs == 1); these operate on vec->envs[0]. static PyObject *vec_set_agent_states(PyObject *self, PyObject *args) { - if (PyTuple_Size(args) != 11) { - PyErr_SetString(PyExc_TypeError, "vec_set_agent_states requires 11 arguments"); + if (PyTuple_Size(args) != 12) { + PyErr_SetString(PyExc_TypeError, "vec_set_agent_states requires 12 arguments"); return NULL; } VecEnv *vec = unpack_vecenv(args); @@ -1005,6 +1005,7 @@ static PyObject *vec_set_agent_states(PyObject *self, PyObject *args) { PyObject *yaw_rate_arr = PyTuple_GetItem(args, 8); PyObject *accel_long_arr = PyTuple_GetItem(args, 9); PyObject *seconds_stopped_arr = PyTuple_GetItem(args, 10); // None keeps c_step's accumulation + PyObject *steering_arr = PyTuple_GetItem(args, 11); // None derives steering from yaw_rate if (!PyArray_Check(idx_arr) || !PyArray_Check(x_arr) || !PyArray_Check(y_arr) || !PyArray_Check(z_arr) || !PyArray_Check(heading_arr) || !PyArray_Check(vx_arr) || !PyArray_Check(vy_arr) || !PyArray_Check(yaw_rate_arr) || !PyArray_Check(accel_long_arr)) { @@ -1015,6 +1016,10 @@ static PyObject *vec_set_agent_states(PyObject *self, PyObject *args) { PyErr_SetString(PyExc_TypeError, "seconds_stopped must be a NumPy array or None"); return NULL; } + if (steering_arr != Py_None && !PyArray_Check(steering_arr)) { + PyErr_SetString(PyExc_TypeError, "steering must be a NumPy array or None"); + return NULL; + } int *idx = (int *) PyArray_DATA((PyArrayObject *) idx_arr); float *x = (float *) PyArray_DATA((PyArrayObject *) x_arr); float *y = (float *) PyArray_DATA((PyArrayObject *) y_arr); @@ -1026,10 +1031,20 @@ static PyObject *vec_set_agent_states(PyObject *self, PyObject *args) { float *accel_long = (float *) PyArray_DATA((PyArrayObject *) accel_long_arr); float *seconds_stopped = seconds_stopped_arr == Py_None ? NULL : (float *) PyArray_DATA((PyArrayObject *) seconds_stopped_arr); + float *steering = steering_arr == Py_None ? NULL : (float *) PyArray_DATA((PyArrayObject *) steering_arr); int count = (int) PyArray_SIZE((PyArrayObject *) idx_arr); - PyObject *value_arrs[9] - = {x_arr, y_arr, z_arr, heading_arr, vx_arr, vy_arr, yaw_rate_arr, accel_long_arr, seconds_stopped_arr}; - for (int k = 0; k < 9; k++) { + PyObject *value_arrs[10] + = {x_arr, + y_arr, + z_arr, + heading_arr, + vx_arr, + vy_arr, + yaw_rate_arr, + accel_long_arr, + seconds_stopped_arr, + steering_arr}; + for (int k = 0; k < 10; k++) { if (value_arrs[k] == Py_None) { continue; } @@ -1050,11 +1065,13 @@ static PyObject *vec_set_agent_states(PyObject *self, PyObject *args) { vy, yaw_rate, accel_long, - seconds_stopped) + seconds_stopped, + steering) != 0) { PyErr_SetString( PyExc_ValueError, - "vec_set_agent_states: agent index out of range, non-finite state or negative seconds_stopped"); + "vec_set_agent_states: agent index out of range, non-finite state, negative seconds_stopped or " + "non-finite steering"); return NULL; } Py_RETURN_NONE; @@ -1166,6 +1183,38 @@ static PyObject *vec_set_stop_signs(PyObject *self, PyObject *args) { return PyLong_FromLong(drive->num_traffic_elements); } +static PyObject *vec_set_traffic_light_lines(PyObject *self, PyObject *args) { + if (PyTuple_Size(args) != 3) { + PyErr_SetString(PyExc_TypeError, "vec_set_traffic_light_lines requires 3 arguments (vec_env, indices, lines)"); + return NULL; + } + VecEnv *vec = unpack_vecenv(args); + if (!vec) { + return NULL; + } + PyObject *indices_arr = PyTuple_GetItem(args, 1); + PyObject *lines_arr = PyTuple_GetItem(args, 2); + if (!PyArray_Check(indices_arr) || !PyArray_Check(lines_arr)) { + PyErr_SetString(PyExc_TypeError, "vec_set_traffic_light_lines: indices and lines must be NumPy arrays"); + return NULL; + } + int count = (int) PyArray_SIZE((PyArrayObject *) indices_arr); + if ((int) PyArray_SIZE((PyArrayObject *) lines_arr) != count * 6) { + PyErr_SetString(PyExc_ValueError, "vec_set_traffic_light_lines: lines must hold 6 floats per index"); + return NULL; + } + Drive *drive = (Drive *) vec->envs[0]; + const int *indices = (const int *) PyArray_DATA((PyArrayObject *) indices_arr); + const float *lines = (const float *) PyArray_DATA((PyArrayObject *) lines_arr); + if (c_set_traffic_light_lines(drive, count, indices, lines) != 0) { + PyErr_SetString( + PyExc_ValueError, + "vec_set_traffic_light_lines: index is not a traffic light, non-finite value or zero-length line"); + return NULL; + } + Py_RETURN_NONE; +} + static PyObject *vec_get_agent_goal_progress(PyObject *self, PyObject *args) { if (PyTuple_Size(args) != 2) { PyErr_SetString(PyExc_TypeError, "vec_get_agent_goal_progress requires 2 arguments"); @@ -1513,6 +1562,10 @@ static PyMethodDef methods[] vec_set_stop_signs, METH_VARARGS, "Replace the stop signs with the external sim's own (co-sim)"}, + {"vec_set_traffic_light_lines", + vec_set_traffic_light_lines, + METH_VARARGS, + "Move traffic-light stop lines onto the external sim's own (co-sim)"}, {"vec_set_agent_goals", vec_set_agent_goals, METH_VARARGS, "Set an agent's goal waypoints (co-sim)"}, {"vec_get_agent_goal_progress", vec_get_agent_goal_progress, diff --git a/pufferlib/ocean/fused_slot_encoder.py b/pufferlib/ocean/fused_slot_encoder.py new file mode 100644 index 0000000000..769318293b --- /dev/null +++ b/pufferlib/ocean/fused_slot_encoder.py @@ -0,0 +1,666 @@ +"""Fused per-slot encoder: Linear -> LayerNorm -> ReLU -> Linear -> masked max-pool over slots, in two Triton kernels. + +Reproduces the bf16-autocast numerics of DriveBackbone._encode_and_pool (bf16 matmul inputs and outputs, fp32 +accumulation, fp32 LayerNorm, gradient shared evenly between tied maxima) without materializing the 256-wide +per-slot activations between ops. The forward walks the slots one at a time for a block of batch rows, so the +max-pool is an elementwise running max with no padded slots. +""" + +import os +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + + +def cuda_backend_active(): + return os.environ.get("PUFFER_SLOT_ENCODER_BACKEND", "triton") == "cuda" + + +if cuda_backend_active(): + import pufferlib.ocean.cuda_slot_encoder # noqa: F401 fail at import time, not in the first policy call + + +@dataclass(frozen=True) +class KernelConfig: + fwd_block_batch: int + fwd_block_k: int + fwd_num_warps: int + fwd_num_stages: int + bwd_block_rows: int + bwd_block_k: int + bwd_num_warps: int + bwd_num_stages: int + pool_block_rows: int + pool_num_warps: int + + +# Timed on an RTX 5090 (GB202, compute capability 12.0). +GB202_CONFIG = KernelConfig(32, 32, 4, 2, 64, 64, 8, 2, 32, 4) +# Timed on an RTX PRO 6000 Blackwell Server at 131072 rows (scripts/kesai/15_tune_fused_slot_encoder_slurm.sh, kesai job 34). +RTX_PRO_6000_CONFIG = KernelConfig(32, 32, 4, 2, 64, 64, 8, 3, 64, 8) +# Timed on an H100 80GB HBM3 at 131072 rows (scripts/kesai/15_tune_fused_slot_encoder_slurm.sh, job 30133). +H100_CONFIG = KernelConfig(32, 32, 4, 3, 32, 64, 4, 3, 32, 4) +# Untimed placeholder: replace with the KernelConfig printed by scripts/tune_fused_slot_encoder.py on that GPU. +B200_CONFIG = GB202_CONFIG +B300_CONFIG = B200_CONFIG # Blackwell Ultra shares the B200 SM design +CONFIGS_BY_CAPABILITY = {(12, 0): GB202_CONFIG, (9, 0): H100_CONFIG, (10, 0): B200_CONFIG, (10, 3): B300_CONFIG} +# Same capability, different chip speeds: a device-name prefix wins over the capability table. +CONFIGS_BY_DEVICE_NAME = {"NVIDIA RTX PRO 6000 Blackwell": RTX_PRO_6000_CONFIG} +CONFIG_OVERRIDE = None +ROWS_PER_PARTIAL_SUM = 256 +MIN_FEATURE_PAD = 16 +MAX_FEATURE_PAD = 64 +LAYER_NORM_EPS_VALUE = 1e-5 +LAYER_NORM_EPS = tl.constexpr(LAYER_NORM_EPS_VALUE) + + +def kernel_config(device): + if CONFIG_OVERRIDE is not None: + return CONFIG_OVERRIDE + device_name = torch.cuda.get_device_name(device) + for name_prefix, config in CONFIGS_BY_DEVICE_NAME.items(): + if device_name.startswith(name_prefix): + return config + return CONFIGS_BY_CAPABILITY.get(torch.cuda.get_device_capability(device), GB202_CONFIG) + + +@triton.jit +def _bf16_round(x): + return x.to(tl.bfloat16).to(tl.float32) + + +@triton.jit +def _pre_norm_chunk(x_bf16, k_offs, w1_ptr, b1_ptr, F_PAD: tl.constexpr): + f_offs = tl.arange(0, F_PAD) + w1t_chunk = tl.load(w1_ptr + k_offs[None, :] * F_PAD + f_offs[:, None]) + b1_chunk = tl.load(b1_ptr + k_offs) + return _bf16_round(tl.dot(x_bf16, w1t_chunk) + b1_chunk[None, :]) + + +@triton.jit +def _row_stats_full(x_bf16, w1_ptr, b1_ptr, D: tl.constexpr, F_PAD: tl.constexpr): + pre_norm = _pre_norm_chunk(x_bf16, tl.arange(0, D), w1_ptr, b1_ptr, F_PAD) + mean = tl.sum(pre_norm, axis=1) / D + centered = pre_norm - mean[:, None] + rstd = tl.rsqrt(tl.sum(centered * centered, axis=1) / D + LAYER_NORM_EPS) + return mean, rstd + + +@triton.jit +def _normed_relu_chunk(x_bf16, mean, rstd, k_offs, w1_ptr, b1_ptr, gamma_ptr, beta_ptr, F_PAD: tl.constexpr): + gamma_chunk = tl.load(gamma_ptr + k_offs) + beta_chunk = tl.load(beta_ptr + k_offs) + pre_norm_chunk = _pre_norm_chunk(x_bf16, k_offs, w1_ptr, b1_ptr, F_PAD) + normed_chunk = (pre_norm_chunk - mean[:, None]) * rstd[:, None] * gamma_chunk[None, :] + beta_chunk[None, :] + return tl.maximum(normed_chunk, 0.0).to(tl.bfloat16) + + +@triton.jit +def _encode_columns( + x_bf16, + mean, + rstd, + row_mask, + flat_rows, + n_offs, + w1_ptr, + b1_ptr, + gamma_ptr, + beta_ptr, + w2t_ptr, + b2_ptr, + normed_relu_ptr, + D: tl.constexpr, + F_PAD: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_N: tl.constexpr, + STORE_NORMED: tl.constexpr, +): + # columns recomputed per K chunk: register tensors cannot be sliced + encoded = tl.zeros((x_bf16.shape[0], BLOCK_N), tl.float32) + for k0 in range(0, D, BLOCK_K): + k_offs = k0 + tl.arange(0, BLOCK_K) + normed_relu_chunk = _normed_relu_chunk(x_bf16, mean, rstd, k_offs, w1_ptr, b1_ptr, gamma_ptr, beta_ptr, F_PAD) + if STORE_NORMED: + tl.store( + normed_relu_ptr + flat_rows[:, None] * D + k_offs[None, :], normed_relu_chunk, mask=row_mask[:, None] + ) + w2t_chunk = tl.load(w2t_ptr + k_offs[:, None] * D + n_offs[None, :]) + encoded = tl.dot(normed_relu_chunk, w2t_chunk, encoded) + b2 = tl.load(b2_ptr + n_offs) + return _bf16_round(encoded + b2[None, :]) + + +@triton.jit +def fused_slot_encoder_fwd_kernel( + x_ptr, + counts_ptr, + w1_ptr, + b1_ptr, + gamma_ptr, + beta_ptr, + w2t_ptr, + b2_ptr, + normed_relu_ptr, + pooled_ptr, + argmax_ptr, + B, + S, + F, + stride_xb, + stride_xs, + stride_xf, + BLOCK_B: tl.constexpr, + D: tl.constexpr, + F_PAD: tl.constexpr, + BLOCK_K: tl.constexpr, + STORE_ACTS: tl.constexpr, +): + pid = tl.program_id(0) + batch_offs = pid * BLOCK_B + tl.arange(0, BLOCK_B) + batch_mask = batch_offs < B + counts = tl.load(counts_ptr + batch_offs, mask=batch_mask, other=0) + d_offs = tl.arange(0, D) + f_offs = tl.arange(0, F_PAD) + running_max = tl.full((BLOCK_B, D), float("-inf"), tl.float32) + argmax = tl.zeros((BLOCK_B, D), tl.int32) + for slot in range(S): + x = tl.load( + x_ptr + batch_offs[:, None] * stride_xb + slot * stride_xs + f_offs[None, :] * stride_xf, + mask=batch_mask[:, None] & (f_offs[None, :] < F), + other=0.0, + ) + x_bf16 = x.to(tl.bfloat16) + flat_rows = batch_offs * S + slot + mean, rstd = _row_stats_full(x_bf16, w1_ptr, b1_ptr, D, F_PAD) + encoded = _encode_columns( + x_bf16, + mean, + rstd, + batch_mask, + flat_rows, + d_offs, + w1_ptr, + b1_ptr, + gamma_ptr, + beta_ptr, + w2t_ptr, + b2_ptr, + normed_relu_ptr, + D, + F_PAD, + BLOCK_K, + D, + STORE_ACTS, + ) + valid = (batch_mask & (slot < counts))[:, None] + greater = valid & (encoded > running_max) + if STORE_ACTS: + argmax = tl.where(greater, slot, argmax) + running_max = tl.where(greater, encoded, running_max) + pooled = tl.where((counts > 0)[:, None], running_max, 0.0) + out_offs = batch_offs[:, None] * D + d_offs[None, :] + tl.store(pooled_ptr + out_offs, pooled.to(tl.bfloat16), mask=batch_mask[:, None]) + if STORE_ACTS: + tl.store(argmax_ptr + out_offs, argmax.to(tl.uint8), mask=batch_mask[:, None]) + + +@triton.jit +def fused_slot_pool_bwd_kernel( + argmax_ptr, + counts_ptr, + grad_pooled_ptr, + grad_encoded_ptr, + grad_b2_partial_ptr, + B, + S, + stride_grad_row, + BLOCK_R: tl.constexpr, + ITERS: tl.constexpr, + D: tl.constexpr, +): + pid = tl.program_id(0) + d_offs = tl.arange(0, D) + grad_b2 = tl.zeros((D,), tl.float32) + for it in range(ITERS): + rows = (pid * ITERS + it) * BLOCK_R + tl.arange(0, BLOCK_R) + row_mask = rows < B * S + batch = rows // S + slots = rows % S + counts = tl.load(counts_ptr + batch, mask=row_mask, other=0) + valid = row_mask & (slots < counts) + row_offs = rows[:, None] * D + d_offs[None, :] + batch_offs = batch[:, None] * D + d_offs[None, :] + grad_offs = batch[:, None] * stride_grad_row + d_offs[None, :] + argmax = tl.load(argmax_ptr + batch_offs, mask=row_mask[:, None], other=255).to(tl.int32) + grad_pooled = tl.load(grad_pooled_ptr + grad_offs, mask=row_mask[:, None], other=0.0).to(tl.float32) + grad_encoded = tl.where(valid[:, None] & (argmax == slots[:, None]), grad_pooled, 0.0) + tl.store(grad_encoded_ptr + row_offs, grad_encoded.to(tl.bfloat16), mask=row_mask[:, None]) + grad_b2 += tl.sum(grad_encoded, axis=0) + tl.store(grad_b2_partial_ptr + pid * D + d_offs, grad_b2) + + +@triton.jit +def _grad_normed_relu_columns( + grad_encoded_ptr, + rows, + row_mask, + n_offs, + w2_ptr, + D: tl.constexpr, + BLOCK_R: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_N: tl.constexpr, +): + grad_normed_relu = tl.zeros((BLOCK_R, BLOCK_N), tl.float32) + for k0 in range(0, D, BLOCK_K): + k_offs = k0 + tl.arange(0, BLOCK_K) + grad_chunk = tl.load(grad_encoded_ptr + rows[:, None] * D + k_offs[None, :], mask=row_mask[:, None], other=0.0) + w2_chunk = tl.load(w2_ptr + k_offs[:, None] * D + n_offs[None, :]) + grad_normed_relu = tl.dot(grad_chunk, w2_chunk, grad_normed_relu) + return _bf16_round(grad_normed_relu) + + +@triton.jit +def _layer_norm_grad_columns( + x_bf16, mean, rstd, grad_normed_relu, n_offs, w1_ptr, b1_ptr, gamma_ptr, beta_ptr, F_PAD: tl.constexpr +): + gamma_chunk = tl.load(gamma_ptr + n_offs) + beta_chunk = tl.load(beta_ptr + n_offs) + pre_norm = _pre_norm_chunk(x_bf16, n_offs, w1_ptr, b1_ptr, F_PAD) + normalized = (pre_norm - mean[:, None]) * rstd[:, None] + normed = normalized * gamma_chunk[None, :] + beta_chunk[None, :] + grad_normed = tl.where(normed > 0.0, grad_normed_relu, 0.0) + return normalized, grad_normed, grad_normed * gamma_chunk[None, :] + + +@triton.jit +def _backward_rows_full( + x_bf16, + rows, + row_mask, + grad_encoded_ptr, + w1t, + b1, + gamma, + beta, + w2_ptr, + grad_pre_norm_ptr, + grad_gamma, + grad_beta, + grad_b1, + D: tl.constexpr, + BLOCK_R: tl.constexpr, + BLOCK_K: tl.constexpr, +): + d_offs = tl.arange(0, D) + grad_normed_relu = _grad_normed_relu_columns( + grad_encoded_ptr, rows, row_mask, d_offs, w2_ptr, D, BLOCK_R, BLOCK_K, D + ) + pre_norm = _bf16_round(tl.dot(x_bf16, w1t) + b1[None, :]) + mean = tl.sum(pre_norm, axis=1) / D + centered = pre_norm - mean[:, None] + rstd = tl.rsqrt(tl.sum(centered * centered, axis=1) / D + LAYER_NORM_EPS) + normalized = centered * rstd[:, None] + normed = normalized * gamma[None, :] + beta[None, :] + grad_normed = tl.where(normed > 0.0, grad_normed_relu, 0.0) + grad_scaled = grad_normed * gamma[None, :] + mean_grad = tl.sum(grad_scaled, axis=1) / D + mean_grad_normalized = tl.sum(grad_scaled * normalized, axis=1) / D + grad_pre_norm = _bf16_round( + rstd[:, None] * (grad_scaled - mean_grad[:, None] - normalized * mean_grad_normalized[:, None]) + ) + tl.store( + grad_pre_norm_ptr + rows[:, None] * D + d_offs[None, :], + grad_pre_norm.to(tl.bfloat16), + mask=row_mask[:, None], + ) + grad_gamma += tl.sum(grad_normed * normalized, axis=0) + grad_beta += tl.sum(grad_normed, axis=0) + grad_b1 += tl.sum(grad_pre_norm, axis=0) + return grad_gamma, grad_beta, grad_b1 + + +@triton.jit +def fused_slot_encoder_bwd_kernel( + x_ptr, + grad_encoded_ptr, + w1_ptr, + b1_ptr, + gamma_ptr, + beta_ptr, + w2_ptr, + grad_pre_norm_ptr, + grad_gamma_partial_ptr, + grad_beta_partial_ptr, + grad_b1_partial_ptr, + B, + S, + F, + stride_xb, + stride_xs, + stride_xf, + BLOCK_R: tl.constexpr, + ITERS: tl.constexpr, + D: tl.constexpr, + F_PAD: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid = tl.program_id(0) + d_offs = tl.arange(0, D) + f_offs = tl.arange(0, F_PAD) + w1t = tl.load(w1_ptr + d_offs[None, :] * F_PAD + f_offs[:, None]) + b1 = tl.load(b1_ptr + d_offs) + gamma = tl.load(gamma_ptr + d_offs) + beta = tl.load(beta_ptr + d_offs) + grad_gamma = tl.zeros((D,), tl.float32) + grad_beta = tl.zeros((D,), tl.float32) + grad_b1 = tl.zeros((D,), tl.float32) + for it in range(ITERS): + rows = (pid * ITERS + it) * BLOCK_R + tl.arange(0, BLOCK_R) + row_mask = rows < B * S + batch = rows // S + slots = rows % S + x = tl.load( + x_ptr + batch[:, None] * stride_xb + slots[:, None] * stride_xs + f_offs[None, :] * stride_xf, + mask=row_mask[:, None] & (f_offs[None, :] < F), + other=0.0, + ) + grad_gamma, grad_beta, grad_b1 = _backward_rows_full( + x.to(tl.bfloat16), + rows, + row_mask, + grad_encoded_ptr, + w1t, + b1, + gamma, + beta, + w2_ptr, + grad_pre_norm_ptr, + grad_gamma, + grad_beta, + grad_b1, + D, + BLOCK_R, + BLOCK_K, + ) + tl.store(grad_gamma_partial_ptr + pid * D + d_offs, grad_gamma) + tl.store(grad_beta_partial_ptr + pid * D + d_offs, grad_beta) + tl.store(grad_b1_partial_ptr + pid * D + d_offs, grad_b1) + + +def _feature_pad(in_features): + return max(MIN_FEATURE_PAD, triton.next_power_of_2(in_features)) + + +def pad_weight(w1_bf16, feature_pad): + out_features, in_features = w1_bf16.shape + padded = torch.zeros(out_features, feature_pad, dtype=torch.bfloat16, device=w1_bf16.device) + padded[:, :in_features] = w1_bf16 + return padded + + +@torch.library.custom_op("pufferdrive::fused_slot_encoder_fwd", mutates_args=()) +def fused_slot_encoder_fwd( + x: torch.Tensor, + counts: torch.Tensor, + w1_padded: torch.Tensor, + b1: torch.Tensor, + gamma: torch.Tensor, + beta: torch.Tensor, + w2t: torch.Tensor, + b2: torch.Tensor, + store_acts: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, S, F = x.shape + D, feature_pad = w1_padded.shape + config = kernel_config(x.device) + rows = B * S if store_acts else 0 + normed_relu = torch.empty(rows, D, dtype=torch.bfloat16, device=x.device) + pooled = torch.empty(B, D, dtype=torch.bfloat16, device=x.device) + argmax = torch.empty(B if store_acts else 0, D, dtype=torch.uint8, device=x.device) + grid = (triton.cdiv(B, config.fwd_block_batch),) + fused_slot_encoder_fwd_kernel[grid]( + x, + counts, + w1_padded, + b1, + gamma, + beta, + w2t, + b2, + normed_relu, + pooled, + argmax, + B, + S, + F, + x.stride(0), + x.stride(1), + x.stride(2), + BLOCK_B=config.fwd_block_batch, + D=D, + F_PAD=feature_pad, + BLOCK_K=config.fwd_block_k, + STORE_ACTS=store_acts, + num_warps=config.fwd_num_warps, + num_stages=config.fwd_num_stages, + ) + return pooled, argmax, normed_relu + + +@fused_slot_encoder_fwd.register_fake +def _(x, counts, w1_padded, b1, gamma, beta, w2t, b2, store_acts): + B, S, _ = x.shape + D = w1_padded.shape[0] + rows = B * S if store_acts else 0 + return ( + torch.empty(B, D, dtype=torch.bfloat16, device=x.device), + torch.empty(B, D, dtype=torch.uint8, device=x.device), + torch.empty(rows, D, dtype=torch.bfloat16, device=x.device), + torch.empty(rows, D, dtype=torch.bfloat16, device=x.device), + ) + + +@torch.library.custom_op("pufferdrive::fused_slot_encoder_bwd", mutates_args=()) +def fused_slot_encoder_bwd( + x: torch.Tensor, + counts: torch.Tensor, + w1_padded: torch.Tensor, + b1: torch.Tensor, + gamma: torch.Tensor, + beta: torch.Tensor, + w2: torch.Tensor, + normed_relu: torch.Tensor, + argmax: torch.Tensor, + grad_pooled: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if cuda_backend_active(): + from pufferlib.ocean.cuda_slot_encoder import cuda_slot_encoder_bwd + + return cuda_slot_encoder_bwd(x, counts, w1_padded, b1, gamma, beta, w2, normed_relu, argmax, grad_pooled) + B, S, F = x.shape + D, feature_pad = w1_padded.shape + config = kernel_config(x.device) + rows = B * S + partial_count = triton.cdiv(rows, ROWS_PER_PARTIAL_SUM) + # The cat backward hands over a column slice of the feature gradient; only its row stride is non-trivial. + if grad_pooled.stride(1) != 1: + grad_pooled = grad_pooled.contiguous() + grad_encoded = torch.empty(rows, D, dtype=torch.bfloat16, device=x.device) + grad_b2_partial = torch.empty(partial_count, D, dtype=torch.float32, device=x.device) + fused_slot_pool_bwd_kernel[(partial_count,)]( + argmax, + counts, + grad_pooled, + grad_encoded, + grad_b2_partial, + B, + S, + grad_pooled.stride(0), + BLOCK_R=config.pool_block_rows, + ITERS=ROWS_PER_PARTIAL_SUM // config.pool_block_rows, + D=D, + num_warps=config.pool_num_warps, + ) + grad_pre_norm = torch.empty(rows, D, dtype=torch.bfloat16, device=x.device) + grad_gamma_partial = torch.empty(partial_count, D, dtype=torch.float32, device=x.device) + grad_beta_partial = torch.empty(partial_count, D, dtype=torch.float32, device=x.device) + grad_b1_partial = torch.empty(partial_count, D, dtype=torch.float32, device=x.device) + fused_slot_encoder_bwd_kernel[(partial_count,)]( + x, + grad_encoded, + w1_padded, + b1, + gamma, + beta, + w2, + grad_pre_norm, + grad_gamma_partial, + grad_beta_partial, + grad_b1_partial, + B, + S, + F, + x.stride(0), + x.stride(1), + x.stride(2), + BLOCK_R=config.bwd_block_rows, + ITERS=ROWS_PER_PARTIAL_SUM // config.bwd_block_rows, + D=D, + F_PAD=feature_pad, + BLOCK_K=config.bwd_block_k, + num_warps=config.bwd_num_warps, + num_stages=config.bwd_num_stages, + ) + # Weight gradients stay bf16 cuBLAS GEMMs with fp32 accumulation, as in the autocast reference. + grad_w2 = (grad_encoded.t() @ normed_relu).float() + grad_w1 = (grad_pre_norm.t() @ x.reshape(rows, F).to(torch.bfloat16)).float() + grad_b2 = grad_b2_partial.sum(0).to(torch.bfloat16).float() + grad_b1 = grad_b1_partial.sum(0).to(torch.bfloat16).float() + return grad_w1, grad_b1, grad_gamma_partial.sum(0), grad_beta_partial.sum(0), grad_w2, grad_b2 + + +@fused_slot_encoder_bwd.register_fake +def _(x, counts, w1_padded, b1, gamma, beta, w2, normed_relu, argmax, grad_pooled): + F = x.shape[2] + D = w1_padded.shape[0] + f32 = {"dtype": torch.float32, "device": x.device} + return ( + torch.empty(D, F, **f32), + torch.empty(D, **f32), + torch.empty(D, **f32), + torch.empty(D, **f32), + torch.empty(D, D, **f32), + torch.empty(D, **f32), + ) + + +@torch.library.custom_op("pufferdrive::fused_slot_encoder", mutates_args=()) +def fused_slot_encoder( + x: torch.Tensor, + counts: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + gamma: torch.Tensor, + beta: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, + store_acts: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + w1_padded = pad_weight(w1.to(torch.bfloat16), _feature_pad(w1.shape[1])) + b1_bf16_values = b1.to(torch.bfloat16).float() + w2_bf16 = w2.to(torch.bfloat16).contiguous() + b2_bf16_values = b2.to(torch.bfloat16).float() + if cuda_backend_active(): + from pufferlib.ocean.cuda_slot_encoder import cuda_slot_encoder_fwd + + pooled, argmax, normed_relu = cuda_slot_encoder_fwd( + x, counts, w1_padded, b1_bf16_values, gamma, beta, w2_bf16, b2_bf16_values, store_acts + ) + return pooled, argmax, normed_relu, w1_padded, b1_bf16_values, w2_bf16 + pooled, argmax, normed_relu = fused_slot_encoder_fwd( + x, counts, w1_padded, b1_bf16_values, gamma, beta, w2_bf16.t().contiguous(), b2_bf16_values, store_acts + ) + return pooled, argmax, normed_relu, w1_padded, b1_bf16_values, w2_bf16 + + +@fused_slot_encoder.register_fake +def _(x, counts, w1, b1, gamma, beta, w2, b2, store_acts): + B, S, F = x.shape + D = w1.shape[0] + rows = B * S if store_acts else 0 + bf16 = {"dtype": torch.bfloat16, "device": x.device} + return ( + torch.empty(B, D, **bf16), + torch.empty(B if store_acts else 0, D, dtype=torch.uint8, device=x.device), + torch.empty(rows, D, **bf16), + torch.empty(D, _feature_pad(F), **bf16), + torch.empty(D, dtype=torch.float32, device=x.device), + torch.empty(D, D, **bf16), + ) + + +def _setup_context(ctx, inputs, output): + x, counts, _, _, gamma, beta, _, _, _ = inputs + pooled, argmax, normed_relu, w1_padded, b1_bf16_values, w2_bf16 = output + # only pooled receives a gradient; without this autograd zero-fills rows x D tensors for the other outputs + ctx.set_materialize_grads(False) + ctx.save_for_backward(x, counts, w1_padded, b1_bf16_values, gamma, beta, w2_bf16, normed_relu, argmax) + + +def _backward(ctx, grad_pooled, grad_argmax, grad_normed_relu, grad_w1_padded, grad_b1_values, grad_w2_bf16): + if grad_pooled is None: + return None, None, None, None, None, None, None, None, None + x, counts, w1_padded, b1_bf16_values, gamma, beta, w2_bf16, normed_relu, argmax = ctx.saved_tensors + grad_w1, grad_b1, grad_gamma, grad_beta, grad_w2, grad_b2 = fused_slot_encoder_bwd( + x, + counts, + w1_padded, + b1_bf16_values, + gamma, + beta, + w2_bf16, + normed_relu, + argmax, + grad_pooled.to(torch.bfloat16), + ) + return None, None, grad_w1, grad_b1, grad_gamma, grad_beta, grad_w2, grad_b2, None + + +fused_slot_encoder.register_autograd(_backward, setup_context=_setup_context) + + +def validate_fused_encoder(encoder): + linear_in, norm, activation, linear_out = encoder + if not (isinstance(norm, torch.nn.LayerNorm) and isinstance(activation, torch.nn.ReLU)): + raise TypeError("fused_slot_encoder requires encoder_layer_norm=true and encoder_activation=relu") + width = linear_in.out_features + if width != triton.next_power_of_2(width) or width % 64 != 0 or width > 1024: + raise ValueError(f"fused_slot_encoder needs a power-of-two encoder width that is a multiple of 64, got {width}") + if linear_in.in_features > MAX_FEATURE_PAD: + raise ValueError(f"fused_slot_encoder supports at most {MAX_FEATURE_PAD} input features per slot") + if norm.eps != LAYER_NORM_EPS_VALUE or linear_out.in_features != width or linear_out.out_features != width: + raise ValueError("fused_slot_encoder expects LayerNorm eps 1e-5 and a square second linear layer") + + +def fused_encoder_active(objects): + return objects.is_cuda and torch.is_autocast_enabled("cuda") and torch.get_autocast_dtype("cuda") == torch.bfloat16 + + +def fused_encode_and_pool(objects, valid_counts, encoder): + linear_in, norm, _, linear_out = encoder + store_acts = torch.is_grad_enabled() and linear_in.weight.requires_grad + return fused_slot_encoder( + objects, + valid_counts.contiguous(), + linear_in.weight, + linear_in.bias, + norm.weight, + norm.bias, + linear_out.weight, + linear_out.bias, + store_acts, + )[0] diff --git a/pufferlib/ocean/torch.py b/pufferlib/ocean/torch.py index bf23ae2cc2..ae6f12d2e4 100644 --- a/pufferlib/ocean/torch.py +++ b/pufferlib/ocean/torch.py @@ -8,6 +8,13 @@ import pufferlib.models from pufferlib.ocean.drive import binding +try: + from pufferlib.ocean.fused_slot_encoder import fused_encode_and_pool, fused_encoder_active, validate_fused_encoder +except ModuleNotFoundError as import_error: # triton ships only with CUDA torch wheels + if import_error.name != "triton": + raise + fused_encode_and_pool = fused_encoder_active = validate_fused_encoder = None + from pufferlib.models import Default as Policy # noqa: F401 from pufferlib.models import Convolutional as Conv # noqa: F401 @@ -23,6 +30,8 @@ class DriveBackbone(nn.Module): - Split Actor/Critic (configurable) """ + fused_slot_encoder = False # class default: tests build instances without __init__ + def _create_encoder(self, in_features, out_size): layers = [pufferlib.pytorch.layer_init(nn.Linear(in_features, out_size))] if self.encoder_layer_norm: @@ -35,11 +44,14 @@ def _encode_and_pool(self, objects, valid_counts, encoder): if not self.mask_padded_features: return encoder(objects).max(dim=1).values + if self.fused_slot_encoder and fused_encoder_active(objects): + return fused_encode_and_pool(objects, valid_counts, encoder) + # Dense encode + masked fill keeps shapes static for torch.compile; padded rows are rare. valid_mask = torch.arange(objects.shape[1], device=objects.device) < valid_counts.unsqueeze(1) encoded_objects = encoder(objects) masked_objects = encoded_objects.masked_fill(~valid_mask.unsqueeze(2), torch.finfo(encoded_objects.dtype).min) - pooled = masked_objects.amax(dim=1) + pooled = masked_objects.max(dim=1).values return torch.where(valid_counts.unsqueeze(1) == 0, pooled.new_zeros(()), pooled) def __init__( @@ -59,6 +71,7 @@ def __init__( backbone_activation, backbone_layer_norm, mask_padded_features, + fused_slot_encoder, ): super().__init__() self.encoder_act_cls = ACTIVATIONS[encoder_activation] @@ -115,6 +128,16 @@ def __init__( self.context_encoder = self._create_encoder(self.context_dim, context_input_size) encoders_out += context_input_size + self.fused_slot_encoder = fused_slot_encoder + if fused_slot_encoder and validate_fused_encoder is None: + raise ImportError("fused_slot_encoder=true requires triton") + if fused_slot_encoder and not mask_padded_features: + raise ValueError("fused_slot_encoder requires mask_padded_features=true") + slot_encoder_names = ("lane_encoder", "boundary_encoder", "partner_encoder", "traffic_control_encoder") + for encoder_name in slot_encoder_names: + if fused_slot_encoder and hasattr(self, encoder_name): + validate_fused_encoder(getattr(self, encoder_name)) + # 2. Main Backbone MLP backbone_act_cls = ACTIVATIONS[backbone_activation] backbone_layers = [] @@ -321,6 +344,7 @@ def __init__( actor_head_layer_norm: bool = False, critic_head_layer_norm: bool = False, fp32_heads: bool = False, + fused_slot_encoder: bool = False, ): super().__init__() @@ -378,6 +402,7 @@ def __init__( "backbone_activation": backbone_activation, "backbone_layer_norm": backbone_layer_norm, "mask_padded_features": mask_padded_features, + "fused_slot_encoder": fused_slot_encoder, } # Instantiate backbones diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index 2704ab32aa..0e93120696 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -796,10 +796,13 @@ def _train_ppo_transition(self, losses, profile, epoch): self.optimizer.zero_grad() total_minibatches = 0 retained_transition_count = keep_idx.numel() + # A lone trailing transition has no sample std to normalize its advantage with (NaN), so it sits out this update + if retained_transition_count % self.minibatch_size == 1: + retained_transition_count -= 1 optimizer_batch_size = self.minibatch_size * self.accumulate_minibatches for _ in range(config["update_epochs"]): - permutation = keep_idx[torch.randperm(retained_transition_count, device=keep_idx.device)] + permutation = keep_idx[torch.randperm(keep_idx.numel(), device=keep_idx.device)] for start in range(0, retained_transition_count, self.minibatch_size): profile("train_copy", epoch) mb_idx = permutation[start : start + self.minibatch_size] @@ -1482,6 +1485,48 @@ def _latest_checkpoint_path(run_dir): return max(model_files, key=os.path.getctime) if model_files else None +def parse_cpulist(cpulist): + """sysfs cpulist ("0-119,256-375") to a set of CPU ids.""" + cpus = set() + for part in cpulist.strip().split(","): + if not part: + continue + lo, _, hi = part.partition("-") + cpus.update(range(int(lo), int(hi or lo) + 1)) + return cpus + + +def apply_matmul_env_flags(): + """PUFFER_BF16_REDUCED_REDUCTION=1 lets cuBLAS pick split-K kernels whose partial sums reduce in bf16.""" + if os.environ.get("PUFFER_BF16_REDUCED_REDUCTION") == "1": + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = True + print("cuBLAS bf16 reduced-precision reductions enabled") + + +def bind_process_to_gpu_numa_node(device_index): + """Pin this process (and the env workers it forks) to the CPUs of the GPU's NUMA node; no-op without NUMA or sysfs.""" + if os.environ.get("PUFFER_NUMA_BIND", "1") == "0" or not torch.cuda.is_available(): + return None + props = torch.cuda.get_device_properties(device_index) + pci_dir = f"/sys/bus/pci/devices/{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:{props.pci_device_id:02x}.0" + try: + numa_node = int(open(f"{pci_dir}/numa_node").read()) + if numa_node < 0: + return None + node_cpus = parse_cpulist(open(f"/sys/devices/system/node/node{numa_node}/cpulist").read()) + except (OSError, ValueError): + return None + cpus = node_cpus & os.sched_getaffinity(0) + if not cpus: + print( + f"NUMA bind skipped: GPU {device_index} is on node {numa_node} but none of its CPUs are in this process's affinity mask" + ) + return None + os.sched_setaffinity(0, cpus) + print(f"GPU {device_index} on NUMA node {numa_node}: process and env workers bound to {len(cpus)} CPUs") + return numa_node + + def train(env_name, args=None, vecenv=None, policy=None, logger=None, early_stop_fn=None): args = args or load_config(env_name) @@ -1507,6 +1552,7 @@ def train(env_name, args=None, vecenv=None, policy=None, logger=None, early_stop "obs_slots_lane_n", "obs_boundary_stride", "obs_lane_stride", + "obs_stride_merge_segments", "obs_dropout_boundary", "obs_dropout_lane", "obs_slots_partners_n", @@ -1542,6 +1588,7 @@ def train(env_name, args=None, vecenv=None, policy=None, logger=None, early_stop validate_puffer_drive_resources(args, "training") training_evaluation_scheduled = drive_benchmark.validate_training_evaluation_config(args) + apply_matmul_env_flags() # Assume TorchRun DDP is used if LOCAL_RANK is set world_size = int(os.environ.get("WORLD_SIZE", 1)) global_rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", 0))) @@ -1551,6 +1598,9 @@ def train(env_name, args=None, vecenv=None, policy=None, logger=None, early_stop local_rank = int(os.environ["LOCAL_RANK"]) print(f"rank: {global_rank} (local {local_rank}), MASTER_ADDR={master_addr}, MASTER_PORT={master_port}") torch.cuda.set_device(local_rank) + bind_process_to_gpu_numa_node(local_rank) + elif is_cuda_device(args["train"]["device"]) and torch.cuda.is_available(): + bind_process_to_gpu_numa_node(torch.device(args["train"]["device"]).index or 0) train_seed = args["train"]["seed"] if train_seed is None: @@ -1625,6 +1675,10 @@ def train(env_name, args=None, vecenv=None, policy=None, logger=None, early_stop # numpy scalars of the saved RNG state, not just tensors. tstate = torch.load(trainer_state_path, map_location=train_config["device"], weights_only=False) pufferl.optimizer.load_state_dict(tstate["optimizer_state_dict"]) + # The saved groups carry the old run's annealed lr; train.learning_rate is this run's base. + for param_group in pufferl.optimizer.param_groups: + param_group["lr"] = train_config["learning_rate"] + param_group["initial_lr"] = train_config["learning_rate"] pufferl.global_step = tstate.get("global_step", pufferl.global_step) pufferl.epoch = tstate.get("update", pufferl.epoch) # Fast-forward the LR scheduler to the resumed epoch so the cosine @@ -1653,8 +1707,10 @@ def train(env_name, args=None, vecenv=None, policy=None, logger=None, early_stop logging_threshold = min(0.20 * train_config["total_timesteps"], 100_000_000) all_logs = [] last_training_evaluation_epoch = None + stop_after_epochs = args["train"]["stop_after_epochs"] + stop_epoch = float("inf") if stop_after_epochs is None else pufferl.epoch + stop_after_epochs - while pufferl.global_step < train_config["total_timesteps"]: + while pufferl.global_step < train_config["total_timesteps"] and pufferl.epoch < stop_epoch: if is_cuda_device(train_config["device"]): torch.compiler.cudagraph_mark_step_begin() try: @@ -1783,6 +1839,9 @@ def eval( scenario_offset = eval_config["scenario_offset"] configured_output_subdir = eval_config["output_subdir"] eval_training_render = args["env"]["eval_training_render"] + eval_device = torch.device(args["train"]["device"]) + if eval_device.type == "cuda": + bind_process_to_gpu_numa_node(eval_device.index or 0) report_to_wandb = bool(args["wandb"]) and not use_training_config and failure_replay_csv is None environment_config, benchmarks = drive_benchmark.load_benchmark_config( @@ -2496,10 +2555,9 @@ def _run_eval_rollout( def render_training_replays(env_name, args, policy, epoch, global_step, run_dir): - """Render one HTML replay per training map with the live policy, observations included.""" + """Render one HTML replay per selected training map with the live policy, observations included.""" run_args = copy.deepcopy(args) env_config = run_args["env"] - num_scenarios = env_config["num_maps"] # Fixed-length episodes: eval-mode summaries flush on the resample boundary. env_config["termination_mode"] = False env_config["compute_eval_metrics"] = True @@ -2510,9 +2568,19 @@ def render_training_replays(env_name, args, policy, epoch, global_step, run_dir) capture_observations = run_args["eval"]["capture_observations"] # Observation replays render at multi-GB RSS each; the obs writer count bounds that too. max_renderer_count = run_args["eval"]["observation_replay_writer_count"] if capture_observations else None + # Same maps and episode seeds at every render of a run, so renders across epochs compare like for like. + render_seed = run_args["train"]["seed"] if run_args["train"]["seed"] is not None else 0 + render_rng = np.random.default_rng(render_seed) + render_num_maps = run_args["train"]["render_num_maps"] + if render_num_maps is None: + map_indices = list(range(env_config["num_maps"])) + else: + map_indices = sorted(render_rng.choice(env_config["num_maps"], size=render_num_maps, replace=False).tolist()) + map_seed_pairs = [(map_idx, int(render_rng.integers(0, 2**63, dtype=np.int64))) for map_idx in map_indices] + num_scenarios = len(map_seed_pairs) num_workers = min(run_args["vec"]["num_envs"], num_scenarios) - worker_env_kwargs, total_steps = drive_benchmark._plan_benchmark_eval_workers( - run_args, num_scenarios, num_workers, env_config["scenario_length"], capture_replay=True + worker_env_kwargs, total_steps = drive_benchmark._plan_failure_replay_workers( + run_args, map_seed_pairs, num_workers, env_config["scenario_length"] ) output_dir = os.path.join(run_dir, "renders", f"epoch_{epoch:06d}_step_{global_step}") @@ -2527,7 +2595,7 @@ def render_training_replays(env_name, args, policy, epoch, global_step, run_dir) "Rendering training replays", num_scenarios, policy=policy, - replay_output_dir=os.path.join(output_dir, "replays"), + replay_output_dir=os.path.join(output_dir, drive_eval_replay.ZLIB_REPLAY_DIR_NAME), capture_observations=capture_observations, ) drive_eval_replay._render_eval_replays( diff --git a/pufferlib/resources/drive/binaries/carla/opendrive__Town01.bin b/pufferlib/resources/drive/binaries/carla/opendrive__Town01.bin index b0b3f0402f..db51075617 100644 Binary files a/pufferlib/resources/drive/binaries/carla/opendrive__Town01.bin and b/pufferlib/resources/drive/binaries/carla/opendrive__Town01.bin differ diff --git a/pufferlib/resources/drive/binaries/carla/opendrive__Town02.bin b/pufferlib/resources/drive/binaries/carla/opendrive__Town02.bin index 64aec35867..3184cf7578 100644 Binary files a/pufferlib/resources/drive/binaries/carla/opendrive__Town02.bin and b/pufferlib/resources/drive/binaries/carla/opendrive__Town02.bin differ diff --git a/pufferlib/resources/drive/binaries/carla/opendrive__Town03.bin b/pufferlib/resources/drive/binaries/carla/opendrive__Town03.bin index 0612a6ca66..9ba4881637 100644 Binary files a/pufferlib/resources/drive/binaries/carla/opendrive__Town03.bin and b/pufferlib/resources/drive/binaries/carla/opendrive__Town03.bin differ diff --git a/pufferlib/resources/drive/binaries/carla/opendrive__Town04.bin b/pufferlib/resources/drive/binaries/carla/opendrive__Town04.bin index 7a57d217ac..c50074700b 100644 Binary files a/pufferlib/resources/drive/binaries/carla/opendrive__Town04.bin and b/pufferlib/resources/drive/binaries/carla/opendrive__Town04.bin differ diff --git a/pufferlib/resources/drive/binaries/carla/opendrive__Town05.bin b/pufferlib/resources/drive/binaries/carla/opendrive__Town05.bin index d2bc49c454..f413079330 100644 Binary files a/pufferlib/resources/drive/binaries/carla/opendrive__Town05.bin and b/pufferlib/resources/drive/binaries/carla/opendrive__Town05.bin differ diff --git a/pufferlib/resources/drive/binaries/carla/opendrive__Town06.bin b/pufferlib/resources/drive/binaries/carla/opendrive__Town06.bin index 5831d8531e..04a1b9542f 100644 Binary files a/pufferlib/resources/drive/binaries/carla/opendrive__Town06.bin and b/pufferlib/resources/drive/binaries/carla/opendrive__Town06.bin differ diff --git a/pufferlib/resources/drive/binaries/carla/opendrive__Town07.bin b/pufferlib/resources/drive/binaries/carla/opendrive__Town07.bin index 2573c12d15..2a30a55a92 100644 Binary files a/pufferlib/resources/drive/binaries/carla/opendrive__Town07.bin and b/pufferlib/resources/drive/binaries/carla/opendrive__Town07.bin differ diff --git a/pufferlib/resources/drive/binaries/carla/opendrive__Town10HD.bin b/pufferlib/resources/drive/binaries/carla/opendrive__Town10HD.bin index bba155c792..c6587b07bc 100644 Binary files a/pufferlib/resources/drive/binaries/carla/opendrive__Town10HD.bin and b/pufferlib/resources/drive/binaries/carla/opendrive__Town10HD.bin differ diff --git a/pufferlib/vector.py b/pufferlib/vector.py index 2d9254fa57..a40397adf3 100644 --- a/pufferlib/vector.py +++ b/pufferlib/vector.py @@ -6,6 +6,7 @@ import time import psutil from multiprocessing import get_all_start_methods, get_context +from multiprocessing.connection import wait from pufferlib import PufferEnv, set_buffers import pufferlib.spaces @@ -15,8 +16,7 @@ SEND = 2 RECV = 3 CLOSE = 4 -MAIN = 5 -INFO = 6 +WORKER_POLL_TIMEOUT_SECONDS = 5.0 def recv_precheck(vecenv): @@ -253,36 +253,24 @@ def _worker_process( else: envs = Serial(env_creators, env_args, env_kwargs, num_envs, buf=buf, seed=seed) - semaphores = np.ndarray(num_workers, dtype=np.uint8, buffer=shm["semaphores"]) notify = np.ndarray(num_workers, dtype=bool, buffer=shm["notify"]) - start = time.time() while True: + command = recv_pipe.recv() if notify[worker_idx]: envs.notify() notify[worker_idx] = False - sem = semaphores[worker_idx] - if sem >= MAIN: - if time.time() - start > 0.5: - time.sleep(0.01) - continue - - start = time.time() - if sem == RESET: + if command == RESET: seed = recv_pipe.recv() _, infos = envs.reset(seed=seed) - elif sem == STEP: + elif command == STEP: _, _, _, _, infos = envs.step(atn_arr) - elif sem == CLOSE: + elif command == CLOSE: envs.close() send_pipe.send(None) break - if infos: - semaphores[worker_idx] = INFO - send_pipe.send(infos) - else: - semaphores[worker_idx] = MAIN + send_pipe.send(infos) class Multiprocessing: @@ -319,14 +307,14 @@ def __init__( import psutil - cpu_cores = psutil.cpu_count(logical=False) - if num_workers > cpu_cores and not overwork: + cpu_threads = psutil.cpu_count(logical=True) + if num_workers > cpu_threads and not overwork: raise pufferlib.APIUsageError( " ".join( [ - f"num_workers ({num_workers}) > hardware cores ({cpu_cores}) is disallowed by default.", - "PufferLib multiprocessing is heavily optimized for 1 process per hardware core.", - "If you really want to do this, set overwork=True (--vec-overwork in our demo.py).", + f"num_workers ({num_workers}) > hardware threads ({cpu_threads}) is disallowed by default.", + "Idle workers block instead of spinning, so one process per hardware thread is the limit.", + "If you really want to do this, set overwork=True.", ] ) ) @@ -391,7 +379,6 @@ def __init__( terminals=process_context.RawArray("b", num_agents), truncateds=process_context.RawArray("b", num_agents), masks=process_context.RawArray("b", num_agents), - semaphores=process_context.RawArray("c", num_workers), notify=process_context.RawArray("b", num_workers), ) shape = (num_workers, agents_per_worker) @@ -404,10 +391,8 @@ def __init__( terminals=np.ndarray(shape, dtype=bool, buffer=self.shm["terminals"]), truncations=np.ndarray(shape, dtype=bool, buffer=self.shm["truncateds"]), masks=np.ndarray(shape, dtype=bool, buffer=self.shm["masks"]), - semaphores=np.ndarray(num_workers, dtype=np.uint8, buffer=self.shm["semaphores"]), notify=np.ndarray(num_workers, dtype=bool, buffer=self.shm["notify"]), ) - self.buf["semaphores"][:] = MAIN self.send_pipes, w_recv_pipes = zip(*[process_context.Pipe() for _ in range(num_workers)]) w_send_pipes, self.recv_pipes = zip(*[process_context.Pipe() for _ in range(num_workers)]) @@ -452,28 +437,49 @@ def __init__( self.ready_workers = [] self.waiting_workers = [] + self.batch_workers = [] + self.pending_result = np.zeros(num_workers, dtype=bool) + + def _check_worker_alive(self, worker): + process = self.processes[worker] + if not process.is_alive(): + raise RuntimeError(f"Vectorization worker {worker} exited with code {process.exitcode}") + + def _recv_worker_result(self, worker): + pipe = self.recv_pipes[worker] + while not pipe.poll(WORKER_POLL_TIMEOUT_SECONDS): + self._check_worker_alive(worker) + + self.pending_result[worker] = False + infos = pipe.recv() + if infos: + self.infos[worker] = infos + + def _collect_ready_workers(self): + # sync_traj consumes workers in send order so batches stay aligned with the experience buffer + if self.sync_traj: + worker = self.waiting_workers.pop(0) + self._recv_worker_result(worker) + self.ready_workers.append(worker) + return + + pipes = [self.recv_pipes[worker] for worker in self.waiting_workers] + ready_pipes = wait(pipes, WORKER_POLL_TIMEOUT_SECONDS) + if not ready_pipes: + for worker in self.waiting_workers: + self._check_worker_alive(worker) + return + + for pipe in ready_pipes: + worker = self.recv_pipe_dict[pipe] + self._recv_worker_result(worker) + self.waiting_workers.remove(worker) + self.ready_workers.append(worker) def recv(self): recv_precheck(self) while True: - # Bandaid patch for new experience buffer desync - if self.sync_traj: - worker = self.waiting_workers[0] - sem = self.buf["semaphores"][worker] - if sem >= MAIN: - self.waiting_workers.pop(0) - self.ready_workers.append(worker) - else: - worker = self.waiting_workers.pop(0) - sem = self.buf["semaphores"][worker] - if sem >= MAIN: - self.ready_workers.append(worker) - else: - self.waiting_workers.append(worker) - - if sem == INFO: - self.infos[worker] = self.recv_pipes[worker].recv() - + self._collect_ready_workers() if not self.ready_workers: continue @@ -523,6 +529,7 @@ def recv(self): break self.w_slice = w_slice + self.batch_workers = s_range buf = self.buf o = buf["observations"][w_slice].reshape(self.obs_batch_shape) @@ -546,21 +553,14 @@ def send(self, actions): actions = send_precheck(self, actions).reshape(self.atn_batch_shape) # TODO: What shape? - idxs = self.w_slice - self.actions[idxs] = actions - self.buf["semaphores"][idxs] = STEP + self.actions[self.w_slice] = actions + for worker in self.batch_workers: + self.send_pipes[worker].send(STEP) + self.pending_result[self.batch_workers] = True def async_reset(self, seed=0): - # Flush any waiting workers - while self.waiting_workers: - worker = self.waiting_workers.pop(0) - sem = self.buf["semaphores"][worker] - if sem >= MAIN: - self.ready_workers.append(worker) - if sem == INFO: - self.recv_pipes[worker].recv() - else: - self.waiting_workers.append(worker) + for worker in np.flatnonzero(self.pending_result): + self._recv_worker_result(worker) self.flag = RECV self.prev_env_id = [] @@ -571,13 +571,14 @@ def async_reset(self, seed=0): self.waiting_workers = list(range(self.num_workers)) self.infos = [[] for _ in range(self.num_workers)] - self.buf["semaphores"][:] = RESET worker_ss = ( np.random.SeedSequence(seed).spawn(self.num_workers) if seed is not None else [None] * self.num_workers ) for i in range(self.num_workers): s = int(worker_ss[i].generate_state(1)[0]) & 0x7FFFFFFF if worker_ss[i] is not None else None + self.send_pipes[i].send(RESET) self.send_pipes[i].send(s) + self.pending_result[:] = True def notify(self): self.buf["notify"][:] = True diff --git a/pufferlib/viz.py b/pufferlib/viz.py index 8461f02df5..7b8aa1a255 100644 --- a/pufferlib/viz.py +++ b/pufferlib/viz.py @@ -79,6 +79,8 @@ PAYLOAD_CHUNK_SIZE = 4 * 1024 * 1024 OBS_QUANTIZE_FRAME_CHUNK = 64 REPLAY_COMPRESSION_LEVEL = 1 +REPLAY_HEADER_PROBE_BYTES = 1 << 16 +REPLAY_INFLATE_MAX_ROUNDS = 1 << 12 SIMULATOR_FIGSIZE = (20.0, 20.0) SIMULATOR_DPI = 100 SIMULATOR_GOAL_RADIUS_METERS = 2.0 @@ -731,7 +733,7 @@ def plot_observation( continue count_lane += 1 rel_x, rel_y = lane_obs[i][0], lane_obs[i][1] - length = lane_obs[i][3] * rl2p + length = 2 * lane_obs[i][3] * rl2p 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" @@ -757,7 +759,7 @@ def plot_observation( continue count_boundary += 1 rel_x, rel_y = boundary_obs[i][0], boundary_obs[i][1] - length = boundary_obs[i][3] * rl2p + length = 2 * boundary_obs[i][3] * rl2p 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) @@ -873,6 +875,7 @@ def encode_interactive_replay(scenario, replay): elem_type = int(elem.get("type", 0)) xs = elem.get("x") or [] ys = elem.get("y") or [] + zs = elem.get("z") or [] if not xs or not ys: continue if 1 <= elem_type <= 3: @@ -888,7 +891,7 @@ def encode_interactive_replay(scenario, replay): road_types.append(draw_type) road_ids.append(int(elem.get("id", road_idx))) for i in range(count): - road_points.append((float(xs[i]), float(ys[i]))) + road_points.append((float(xs[i]), float(ys[i]), float(zs[i]) if i < len(zs) else 0.0)) traffic_stop_lines = [] traffic_types = [] @@ -930,7 +933,7 @@ def encode_interactive_replay(scenario, replay): np.subtract(agent_quantized[:, 1:], agent_quantized[:, :-1], out=agent_delta[:, 1:]) chunks = { - "road_points": np.asarray(road_points or [(0.0, 0.0)], dtype=np.float32), + "road_points": np.asarray(road_points or [(0.0, 0.0, 0.0)], dtype=np.float32), "road_lengths": np.asarray(road_lengths or [0], dtype=np.int32), "road_types": np.asarray(road_types or [0], dtype=np.int16), "road_ids": np.asarray(road_ids or [0], dtype=np.int32), @@ -1100,6 +1103,7 @@ def _render_interactive_replay_payload(compressed_payload, filename): .perturbation-line.blindness { border-color:#6d28d9; border-top-style:dashed; } .perturbation-line.phantom-braking { border-color:#b45309; } .perturbation-line.human-log { border-color:#ff0000; } + .perturbation-line.collision-onset { border-color:currentColor; border-top-style:dashed; } /* Agent panel: dark instrument-cluster surface in both themes — scoped variable overrides restyle all children. */ #hud-telemetry { --border:#4a5468; --muted:#aab3c5; --field:rgba(255,255,255,.07); --accent:#7cbcff; position:absolute; top:14px; right:14px; width:372px; max-height:calc(100vh - 90px); padding:12px 14px; overflow-y:auto; display:none; background:rgba(54,62,77,.95); color:#eef1f6; } [data-theme="dark"] #hud-telemetry { background:rgba(48,56,70,.95); } @@ -1128,6 +1132,11 @@ def _render_interactive_replay_payload(compressed_payload, filename): .toggle-header.is-collapsed span:last-child { transform:rotate(-90deg); } .toggle-body.is-collapsed { display:none; } #controls { position:absolute; left:50%; bottom:16px; transform:translateX(-50%); padding:9px 12px; display:flex; gap:12px; align-items:center; } + #view-bar { position:absolute; right:14px; bottom:16px; padding:7px 10px; display:flex; gap:8px; align-items:center; font-size:11px; } + #view-bar .btn.on { border-color:var(--accent); color:var(--accent); } + #view-bar .zx { display:flex; gap:4px; align-items:center; color:var(--muted); } + #view-bar input[type=range] { width:70px; accent-color:var(--accent); } + #viewHint { color:var(--muted); max-width:330px; } .step-counter { font-size:12px; min-width:90px; text-align:center; } #obs-container { position:absolute; left:14px; bottom:16px; width:390px; height:390px; min-width:250px; min-height:250px; max-width:92vw; max-height:86vh; display:none; overflow:hidden; resize:both; border-radius:10px; } #obs-title { position:absolute; top:0; left:0; right:0; z-index:2; display:flex; gap:6px; align-items:center; padding:6px 10px; background:var(--surface); border-bottom:1px solid var(--border); color:var(--muted); font-size:10px; font-weight:600; letter-spacing:.1em; text-transform:uppercase; cursor:grab; } @@ -1150,6 +1159,7 @@ def _render_interactive_replay_payload(compressed_payload, filename):
Active perturbations
Partner blindness
Phantom braking
+
Collision: sharp box = sim footprint, dashed = pose one step before the onset
@@ -1194,6 +1204,13 @@ def _render_interactive_replay_payload(compressed_payload, filename): +
+ + + + + +
__PAYLOAD_CHUNKS__ @@ -1223,7 +1240,9 @@ def _render_interactive_replay_payload(compressed_payload, filename): const obsC = document.getElementById('obs-canvas'), obsCtx = obsC.getContext('2d'); const dpr = window.devicePixelRatio || 1; let step = 0, play = false, speed = 4, lastTick = 0; - let cam = {x:0,y:0,z:5,drag:false,lx:0,ly:0}; + let cam = {x:0,y:0,z:5,drag:false,lx:0,ly:0,cz:0,yaw:0,pitch:0,zx:1}; + let viewMode = '2d', rotateMode = false, dragMode = null, lastTarget = null, mapLayer = null, mapLayerKey = '', mapZMean = 0, RP = 2, roadRangesById = new Map(); + const FOV_DEG = 50, NEAR_FRACTION = 0.05, MAX_PITCH = 80 * Math.PI / 180, DEFAULT_PITCH = 55 * Math.PI / 180, ORBIT_RAD_PER_PX = 0.008, KEY_ROTATE_STEP = 5 * Math.PI / 180, AGENT_PICK_PX = 18; let followedId = null, isEgoCam = false, darkMode = false, showGhost = false, ghostPaths = []; const OBS_ZOOM_DEFAULT = 2.2, OBS_ZOOM_MIN = .45, OBS_ZOOM_MAX = 8; let obsZoom = OBS_ZOOM_DEFAULT, obsExpanded = false, obsMode = 2, obsShowLimits = true; @@ -1353,21 +1372,25 @@ def _render_interactive_replay_payload(compressed_payload, filename): function buildMapPaths() { paths = {0:new Path2D(),1:new Path2D(),2:new Path2D()}; - roadPathsById = new Map(); - let p = 0; + roadPathsById = new Map(); roadRangesById = new Map(); + RP = (H.chunks.road_points && H.chunks.road_points.shape && H.chunks.road_points.shape[1]) || 2; + let p = 0, zSum = 0, zCount = 0; for (let i=0;i 2) for (let j=0;j 0)) return -1; + let onset = frame; + while (onset > 0 && C.metrics_f32[base - stride] > 0) { onset--; base -= stride; } + return onset; + } + function drawCollisionFootprint(a, f, outline) { + // Sprites round their corners; the simulator collides on the sharp box drawn here. + const onset = infractionOnset(f, a.idx, 0); + if (onset < 0) return; + ctx.save(); + ctx.strokeStyle = outline; ctx.lineWidth = 1.5 / cam.z; + ctx.strokeRect(-a.l/2, -a.w/2, a.l, a.w); + ctx.restore(); + if (onset !== f || f === 0) return; + const prev = agentAt(f - 1, a.idx); + if (!prev) return; + // The swept test can flag a contact made between frames: show where the car was one step earlier. + ctx.save(); + ctx.rotate(-a.h); ctx.translate(prev.x - a.x, prev.y - a.y); ctx.rotate(prev.h); + ctx.strokeStyle = outline; ctx.lineWidth = 1.5 / cam.z; ctx.setLineDash([4 / cam.z, 3 / cam.z]); + ctx.strokeRect(-prev.l/2, -prev.w/2, prev.l, prev.w); + ctx.restore(); + } function drawPerturbationOutlines(a) { if (a.phantomBrakingActive) { drawPerturbationOutline(a, PHANTOM_BRAKING_OUTLINE_COLOR, [], 3); @@ -1442,7 +1492,7 @@ def _render_interactive_replay_payload(compressed_payload, filename): const goalRadiusField = H.agent_goal_radius_field; const defaultGoalRadius = H.default_goal_radius_meters || DEFAULT_GOAL_RADIUS_METERS; const goalRadius = Number.isInteger(goalRadiusField) && goalRadiusField < F.af ? C.agent_f32[fb+goalRadiusField] : defaultGoalRadius; - return {idx:idx, id:C.agent_i32[ib], type:agentType, cl:C.agent_i32[ib+6], slot:C.agent_i32[ib+7], partnerBlindnessActive:C.agent_i32[ib+8] === 1, phantomBrakingActive:C.agent_i32[ib+9] === 1, x:C.agent_f32[fb], y:C.agent_f32[fb+1], h:C.agent_f32[fb+3], l:C.agent_f32[fb+4], w:C.agent_f32[fb+5], s:C.agent_f32[fb+6], st:C.agent_f32[fb+7], al:C.agent_f32[fb+8], alat:C.agent_f32[fb+9], jl:C.agent_f32[fb+10], jlat:C.agent_f32[fb+11], goalRadius:goalRadius, c:agentColor}; + return {idx:idx, id:C.agent_i32[ib], type:agentType, cl:C.agent_i32[ib+6], slot:C.agent_i32[ib+7], partnerBlindnessActive:C.agent_i32[ib+8] === 1, phantomBrakingActive:C.agent_i32[ib+9] === 1, x:C.agent_f32[fb], y:C.agent_f32[fb+1], z:C.agent_f32[fb+2], h:C.agent_f32[fb+3], l:C.agent_f32[fb+4], w:C.agent_f32[fb+5], s:C.agent_f32[fb+6], st:C.agent_f32[fb+7], al:C.agent_f32[fb+8], alat:C.agent_f32[fb+9], jl:C.agent_f32[fb+10], jlat:C.agent_f32[fb+11], goalRadius:goalRadius, c:agentColor}; } function getFrameAgents(frame) { const out = []; for (let i=0;i { if(!H || e.target.tagName === 'INPUT') return; if(e.code === 'Space'){ toggle(); e.preventDefault(); } if(e.code === 'ArrowRight'){ play=false; updateBtn(); step=Math.min(step+1,frameMax()); draw(true); } if(e.code === 'ArrowLeft'){ play=false; updateBtn(); step=Math.max(step-1,0); draw(true); } if(e.code === 'Escape'){ followedId=null; isEgoCam=false; updateUI(); draw(true); } if(e.code === 'KeyG'){ showGhost=!showGhost; draw(true); } }); - c.onwheel = e => { e.preventDefault(); cam.z *= Math.exp(-e.deltaY * .001); draw(true); }; - c.onmousedown = e => { if(!H) return; const r=c.getBoundingClientRect(), wx=(e.clientX-r.left-c.width/2)/cam.z+cam.x, wy=(e.clientY-r.top-c.height/2)/-cam.z+cam.y; let hit=null, agents=getFrameAgents(Math.floor(step)); if(!isEgoCam) for(const a of agents) if(Math.hypot(wx-a.x, wy-a.y) < Math.max(a.l,3)){ hit=a.id; break; } if(hit !== null){ followedId=hit; cam.drag=false; } else { followedId=null; isEgoCam=false; cam.drag=true; cam.lx=e.clientX; cam.ly=e.clientY; } draw(true); }; - window.onmouseup = () => cam.drag = false; - c.onmousemove = e => { if(cam.drag && !isEgoCam){ cam.x -= (e.clientX-cam.lx)/cam.z; cam.y -= (e.clientY-cam.ly)/-cam.z; cam.lx=e.clientX; cam.ly=e.clientY; draw(true); } }; + document.addEventListener('keydown', e => { if(!H || e.target.tagName === 'INPUT') return; if(e.code === 'Space'){ toggle(); e.preventDefault(); } if(e.code === 'ArrowRight'){ play=false; updateBtn(); step=Math.min(step+1,frameMax()); draw(true); } if(e.code === 'ArrowLeft'){ play=false; updateBtn(); step=Math.max(step-1,0); draw(true); } if(e.code === 'Escape'){ followedId=null; isEgoCam=false; updateUI(); draw(true); } if(e.code === 'KeyG'){ showGhost=!showGhost; draw(true); } if(e.code === 'KeyV'){ toggleView(); } if(e.code === 'KeyR'){ toggleRotateMode(); } if(e.code === 'KeyT'){ topView(); } if(is3D() && (e.code === 'KeyQ' || e.code === 'KeyE')){ cam.yaw += (e.code === 'KeyQ' ? -1 : 1) * KEY_ROTATE_STEP; updateViewBar(); draw(true); } if(is3D() && (e.code === 'KeyW' || e.code === 'KeyS')){ cam.pitch = Math.max(0, Math.min(MAX_PITCH, cam.pitch + (e.code === 'KeyW' ? 1 : -1) * KEY_ROTATE_STEP)); updateViewBar(); draw(true); } }); + c.addEventListener('contextmenu', e => e.preventDefault()); + c.onwheel = e => { e.preventDefault(); const factor = Math.exp(-e.deltaY * .001); if (is3D()) { const r=c.getBoundingClientRect(), sx=e.clientX-r.left, sy=e.clientY-r.top; setProjection(lastTarget); const before = groundOffset(sx, sy); cam.z *= factor; const after = groundOffset(sx, sy); if (before && after && followedId === null) { cam.x += before[0]-after[0]; cam.y += before[1]-after[1]; } } else cam.z *= factor; draw(true); }; + c.onmousedown = e => { if(!H) return; const r=c.getBoundingClientRect(), sx=e.clientX-r.left, sy=e.clientY-r.top; const orbit = e.button === 1 || e.button === 2 || e.shiftKey || (rotateMode && e.button === 0); if (orbit) { e.preventDefault(); if (!is3D()) enter3D(); dragMode='orbit'; cam.drag=true; cam.lx=e.clientX; cam.ly=e.clientY; draw(true); return; } if (e.button !== 0) return; const hit = pickAgent(sx, sy); if(hit !== null){ followedId=hit; cam.drag=false; } else { followedId=null; isEgoCam=false; dragMode='pan'; cam.drag=true; cam.lx=e.clientX; cam.ly=e.clientY; if (is3D()) { setProjection(null); cam.g0 = groundOffset(sx, sy); cam.t0 = [cam.x, cam.y]; } } draw(true); }; + window.onmouseup = () => { cam.drag = false; dragMode = null; }; + c.onmousemove = e => { if(!cam.drag) return; if (dragMode === 'orbit') { const dx=e.clientX-cam.lx, dy=e.clientY-cam.ly; cam.lx=e.clientX; cam.ly=e.clientY; if (!(isEgoCam && lastTarget)) cam.yaw += dx*ORBIT_RAD_PER_PX; cam.pitch = Math.max(0, Math.min(MAX_PITCH, cam.pitch + dy*ORBIT_RAD_PER_PX)); updateViewBar(); draw(true); return; } if (isEgoCam) return; if (is3D()) { const r=c.getBoundingClientRect(); setProjection(null); const g = groundOffset(e.clientX-r.left, e.clientY-r.top); if (g && cam.g0) { cam.x = cam.t0[0] + cam.g0[0] - g[0]; cam.y = cam.t0[1] + cam.g0[1] - g[1]; } } else { cam.x -= (e.clientX-cam.lx)/cam.z; cam.y -= (e.clientY-cam.ly)/-cam.z; cam.lx=e.clientX; cam.ly=e.clientY; } draw(true); }; obsC.addEventListener('wheel', e => { e.preventDefault(); obsZoom = Math.max(OBS_ZOOM_MIN, Math.min(OBS_ZOOM_MAX, obsZoom * Math.exp(-e.deltaY * .001))); draw(true); }, {passive:false}); function dragPanel(handleId, panelId) { const h=document.getElementById(handleId), p=document.getElementById(panelId); let on=false,sx=0,sy=0,sl=0,st=0; h.addEventListener('mousedown', e => { if(e.target.closest('button')) return; on=true; sx=e.clientX; sy=e.clientY; const r=p.getBoundingClientRect(); sl=r.left; st=r.top; p.style.right='auto'; p.style.bottom='auto'; p.style.left=sl+'px'; p.style.top=st+'px'; }); window.addEventListener('mousemove', e => { if(on){ p.style.left=(sl+e.clientX-sx)+'px'; p.style.top=(st+e.clientY-sy)+'px'; }}); window.addEventListener('mouseup', () => on=false); } dragPanel('obs-title','obs-container'); @@ -1599,7 +1650,7 @@ def _render_interactive_replay_payload(compressed_payload, filename): const rot = (x,y) => [-y,x]; const zero = (off,n) => { for(let i=0;i { const out=[]; for(let i=0;i= 0 ? v(o+limitIdx)*speedScale : null]); } return out; }; + const roads = (start,count,poolName,feat) => { const out=[]; for(let i=0;i= 0 ? v(o+limitIdx)*speedScale : null]); } return out; }; const partners = []; for(let i=0;i { + const onset = infractionOnset(f, agent.idx, metricIdx); + if (onset < 0) return; + warnings.push(onset < f ? `${name} since step ${onset}` : name); + }); const warnKey = warnings.join('|'), warnRow = document.getElementById('warn-row'); if (warnKey !== lastWarnKey) { lastWarnKey = warnKey; warnRow.style.display = warnings.length ? 'flex' : 'none'; warnRow.innerHTML = warnings.map(w=>`${w}`).join(''); } const obs = decodeObs(f, agent.slot); if (obs) { obsBox.style.display='block'; drawObs(obs, {l:agent.l*H.scales.meters_to_position, w:agent.w*H.scales.meters_to_position}); } else obsBox.style.display='none'; } + // ── 3D view: perspective orbit camera around (cam.x, cam.y, cam.cz); cam.z stays the px-per-metre scale at the target ── + let P_ca = 1, P_sa = 0, P_cb = 1, P_sb = 0, P_D = 1; + function is3D() { return viewMode === '3d'; } + function camDist() { return (c.height / cam.z) / (2 * Math.tan(FOV_DEG * Math.PI / 360)); } + function setProjection(target) { const a = (isEgoCam && target) ? (Math.PI / 2 - target.h) : cam.yaw; P_ca = Math.cos(a); P_sa = Math.sin(a); P_cb = Math.cos(cam.pitch); P_sb = Math.sin(cam.pitch); P_D = camDist(); } + function proj(x, y, z) { + const rx = x - cam.x, ry = y - cam.y, rz = (z - cam.cz) * cam.zx; + const x1 = rx * P_ca - ry * P_sa, y1 = rx * P_sa + ry * P_ca; + const depth = P_D + y1 * P_sb - rz * P_cb; + if (depth < NEAR_FRACTION * P_D) return null; + const f = cam.z * P_D / depth; + return [c.width / 2 + f * x1, c.height / 2 - f * (y1 * P_cb + rz * P_sb)]; + } + function unproj(sx, sy, z) { + const rz = (z - cam.cz) * cam.zx, u = (sx - c.width / 2) / (cam.z * P_D), w = (c.height / 2 - sy) / (cam.z * P_D); + const denom = w * P_sb - P_cb; + if (Math.abs(denom) < 1e-9) return null; + const y1 = (rz * (P_sb + w * P_cb) - w * P_D) / denom, depth = P_D + y1 * P_sb - rz * P_cb; + if (depth <= 0) return null; + const x1 = u * depth; + return [cam.x + x1 * P_ca + y1 * P_sa, cam.y - x1 * P_sa + y1 * P_ca, z]; + } + function groundOffset(sx, sy) { const g = unproj(sx, sy, cam.cz); return g ? [g[0] - cam.x, g[1] - cam.y] : null; } + function projPoint(idx) { const o = idx * RP; return proj(C.road_points[o], C.road_points[o + 1], RP > 2 ? C.road_points[o + 2] : 0); } + function tracePolyline(g, points) { let pen = false; for (const s of points) { if (!s) { pen = false; continue; } if (pen) g.lineTo(s[0], s[1]); else g.moveTo(s[0], s[1]); pen = true; } } + function traceRoads(g, type) { + let p = 0; + for (let i = 0; i < H.road_polyline_count; i++) { + const len = C.road_lengths[i]; + if (C.road_types[i] === type) { let pen = false; for (let j = 0; j < len; j++) { const s = projPoint(p + j); if (!s) { pen = false; continue; } if (pen) g.lineTo(s[0], s[1]); else g.moveTo(s[0], s[1]); pen = true; } } + p += len; + } + } + function mapLayerFor(colors) { + const key = [c.width, c.height, cam.x, cam.y, cam.cz, cam.z, cam.zx, P_ca, P_sa, P_cb, P_sb, darkMode].join(','); + if (mapLayer && mapLayerKey === key) return mapLayer; + if (!mapLayer || mapLayer.width !== c.width || mapLayer.height !== c.height) { mapLayer = document.createElement('canvas'); mapLayer.width = c.width; mapLayer.height = c.height; } + const g = mapLayer.getContext('2d'); g.clearRect(0, 0, c.width, c.height); g.lineCap = 'round'; + const styles = [[0, colors.road, .5, []], [1, colors.line, .5, [1, 1]], [2, colors.edge, .8, []]]; + for (const [type, color, widthM, dash] of styles) { g.strokeStyle = color; g.lineWidth = Math.max(1, widthM * cam.z); g.setLineDash(dash.map(d => Math.max(2, d * cam.z))); g.beginPath(); traceRoads(g, type); g.stroke(); } + g.setLineDash([]); mapLayerKey = key; return mapLayer; + } + function strokeRoad3D(roadId, color, widthPx) { const range = roadRangesById.get(roadId); if (!range) return; const pts = []; for (let j = 0; j < range[1]; j++) pts.push(projPoint(range[0] + j)); ctx.strokeStyle = color; ctx.lineWidth = widthPx; ctx.beginPath(); tracePolyline(ctx, pts); ctx.stroke(); } + function boxCorners(x, y, z, h, l, w, pad) { const ch = Math.cos(h), sh = Math.sin(h), hl = l / 2 + pad, hw = w / 2 + pad; return [[-hl, -hw], [hl, -hw], [hl, hw], [-hl, hw]].map(([u, v]) => proj(x + u * ch - v * sh, y + u * sh + v * ch, z)); } + function circlePoints(x, y, z, r) { const out = []; for (let k = 0; k < 16; k++) { const t = k / 16 * 2 * Math.PI; out.push(proj(x + r * Math.cos(t), y + r * Math.sin(t), z)); } return out; } + function tracePolygon(pts) { if (pts.some(s => !s)) return false; ctx.beginPath(); pts.forEach((s, i) => i ? ctx.lineTo(s[0], s[1]) : ctx.moveTo(s[0], s[1])); ctx.closePath(); return true; } + function drawAgent3D(a, f, colors) { + const outline = darkMode ? '#fff' : '#111'; + const body = a.type === 2 ? circlePoints(a.x, a.y, a.z, Math.max(a.w, 0.7) / 2) : boxCorners(a.x, a.y, a.z, a.h, a.l, a.w, 0); + if (!tracePolygon(body)) return; + ctx.fillStyle = a.c; ctx.fill(); ctx.strokeStyle = outline; ctx.lineWidth = 1; ctx.setLineDash([]); ctx.stroke(); + const cen = proj(a.x, a.y, a.z); + if (a.type !== 2 && cen) { const nose = proj(a.x + Math.cos(a.h) * a.l / 2, a.y + Math.sin(a.h) * a.l / 2, a.z); if (nose) { ctx.beginPath(); ctx.moveTo(cen[0], cen[1]); ctx.lineTo(nose[0], nose[1]); ctx.stroke(); } } + if (a.phantomBrakingActive && tracePolygon(boxCorners(a.x, a.y, a.z, a.h, a.l, a.w, 0.5))) { ctx.strokeStyle = PHANTOM_BRAKING_OUTLINE_COLOR; ctx.lineWidth = 2.5; ctx.stroke(); } + if (a.partnerBlindnessActive && tracePolygon(boxCorners(a.x, a.y, a.z, a.h, a.l, a.w, a.phantomBrakingActive ? 1 : 0.5))) { ctx.strokeStyle = PARTNER_BLINDNESS_OUTLINE_COLOR; ctx.lineWidth = 2.5; ctx.setLineDash([6, 4]); ctx.stroke(); ctx.setLineDash([]); } + const onset = infractionOnset(f, a.idx, 0); + if (onset === f && f > 0 && a.type !== 2) { const prev = agentAt(f - 1, a.idx); if (prev && tracePolygon(boxCorners(prev.x, prev.y, a.z, prev.h, prev.l, prev.w, 0))) { ctx.strokeStyle = outline; ctx.lineWidth = 1.5; ctx.setLineDash([4, 3]); ctx.stroke(); ctx.setLineDash([]); } } + if (cen) { ctx.fillStyle = colors.text; ctx.font = '600 12px system-ui'; ctx.textAlign = 'center'; ctx.fillText(a.id, cen[0], cen[1] - 9); } + if (a.id === followedId && tracePolygon(circlePoints(a.x, a.y, a.z, Math.max(a.l, a.w) * 1.2))) { ctx.strokeStyle = colors.accent; ctx.lineWidth = 3; ctx.stroke(); } + } + function drawGhosts3D(f) { + if (!showGhost || !C.ghost_f32) return; + const N = H.chunks.ghost_f32.shape[1]; + ctx.strokeStyle = 'rgba(255,0,0,.55)'; ctx.lineWidth = 1.5; ctx.setLineDash([]); + for (let j = 0; j < N; j++) { + const live = agentAt(f, j), z = live ? live.z : cam.cz; + ctx.beginPath(); let pen = false; + for (let k = 0; k < H.frames; k++) { const b = (k * N + j) * 5; if (C.ghost_f32[b + 4] <= 0) { pen = false; continue; } const s = proj(C.ghost_f32[b], C.ghost_f32[b + 1], z); if (!s) { pen = false; continue; } if (pen) ctx.lineTo(s[0], s[1]); else ctx.moveTo(s[0], s[1]); pen = true; } + ctx.stroke(); + } + ctx.strokeStyle = '#ff0000'; ctx.fillStyle = 'rgba(255,0,0,.22)'; ctx.lineWidth = 1; ctx.setLineDash([3, 2]); + for (let j = 0; j < N; j++) { const b = (f * N + j) * 5, w = C.ghost_f32[b + 4]; if (w <= 0) continue; const live = agentAt(f, j); if (tracePolygon(boxCorners(C.ghost_f32[b], C.ghost_f32[b + 1], live ? live.z : cam.cz, C.ghost_f32[b + 2], C.ghost_f32[b + 3], w, 0))) { ctx.fill(); ctx.stroke(); } } + ctx.setLineDash([]); + } + function draw3D(f, target, colors) { + setProjection(target); + ctx.setTransform(1, 0, 0, 1, 0, 0); + ctx.fillStyle = colors.bg; ctx.fillRect(0, 0, c.width, c.height); + ctx.drawImage(mapLayerFor(colors), 0, 0); + ctx.lineCap = 'round'; + if (target && target.cl >= 0) { strokeRoad3D(target.cl, 'rgba(255,255,255,.35)', 5); strokeRoad3D(target.cl, colors.accent, 4); } + drawGhosts3D(f); + for (const a of getFrameAgents(f)) drawAgent3D(a, f, colors); + ctx.lineCap = 'butt'; + for (let i = 0; i < H.traffic_static_count; i++) { + const t = trafficAt(f, i); if (!t) continue; + const sl = t.stop_line, s0 = proj(sl[0], sl[1], sl[2]), s1 = proj(sl[3], sl[4], sl[5]); if (!s0 || !s1) continue; + if (t.type === 1) { ctx.strokeStyle = trafficColor(t); ctx.lineWidth = 3; ctx.setLineDash([]); } else { ctx.strokeStyle = t.type === 2 ? '#ff0000' : '#ffd700'; ctx.lineWidth = 2.5; ctx.setLineDash([6, 4]); } + ctx.beginPath(); ctx.moveTo(s0[0], s0[1]); ctx.lineTo(s1[0], s1[1]); ctx.stroke(); ctx.setLineDash([]); + } + if (target) { for (const g of selectedGoals(f, target)) { if (tracePolygon(circlePoints(g.x, g.y, target.z, g.radius))) { ctx.fillStyle = 'rgba(56,189,248,.22)'; ctx.fill(); ctx.strokeStyle = '#38bdf8'; ctx.lineWidth = 2.5; ctx.stroke(); } } } + } + function pickAgent(sx, sy) { + if (isEgoCam) return null; + const agents = getFrameAgents(Math.floor(step)); + if (is3D()) { setProjection(lastTarget); for (const a of agents) { const s = proj(a.x, a.y, a.z); if (s && Math.hypot(sx - s[0], sy - s[1]) < AGENT_PICK_PX) return a.id; } return null; } + const wx = (sx - c.width / 2) / cam.z + cam.x, wy = (sy - c.height / 2) / -cam.z + cam.y; + for (const a of agents) if (Math.hypot(wx - a.x, wy - a.y) < Math.max(a.l, 3)) return a.id; + return null; + } + function enter3D() { viewMode = '3d'; if (cam.pitch === 0) cam.pitch = DEFAULT_PITCH; cam.cz = lastTarget ? lastTarget.z : mapZMean; updateViewBar(); } + function exit3D() { viewMode = '2d'; cam.pitch = 0; cam.yaw = 0; updateViewBar(); } + function toggleView() { if (is3D()) exit3D(); else enter3D(); draw(true); } + function topView() { if (!is3D()) return; cam.pitch = 0; cam.yaw = 0; updateViewBar(); draw(true); } + function toggleRotateMode() { rotateMode = !rotateMode; updateViewBar(); } + function setZExaggeration(v) { cam.zx = Math.max(1, Math.min(10, Number(v) || 1)); updateViewBar(); draw(true); } + function updateViewBar() { + const b3 = document.getElementById('btn3d'), br = document.getElementById('btnRotate'), hint = document.getElementById('viewHint'), zx = document.getElementById('zxVal'); + if (!b3) return; + const deg = r => Math.round(r * 180 / Math.PI); + b3.classList.toggle('on', is3D()); b3.textContent = is3D() ? '3D' : '2D'; + br.classList.toggle('on', rotateMode); + zx.textContent = cam.zx + 'x'; + hint.textContent = is3D() ? `tilt ${deg(cam.pitch)}° yaw ${deg(cam.yaw)}° · rotate: right/middle-drag, R+drag or Q/E W/S · T top · V 2D` : 'V or right-drag: 3D'; + } function draw(force=false) { if(!H) return; const f = Math.max(0, Math.min(frameMax(), Math.floor(step))); if(!force && f === lastDrawn) return; const target = followedId !== null ? findAgent(f, followedId) : null; - if (target) { cam.x = target.x; cam.y = target.y; } + lastTarget = target; + if (target) { cam.x = target.x; cam.y = target.y; if (is3D()) cam.cz = target.z; } updateUI(target); + if (is3D()) { draw3D(f, target, getColors()); lastDrawn = f; return; } const colors = getColors(); ctx.fillStyle = colors.bg; ctx.fillRect(0,0,c.width,c.height); ctx.save(); ctx.translate(c.width/2,c.height/2); ctx.scale(cam.z,-cam.z); if(isEgoCam && target) ctx.rotate(Math.PI/2 - target.h); ctx.translate(-cam.x,-cam.y); ctx.lineCap='round'; ctx.strokeStyle=colors.road; ctx.lineWidth=.5; ctx.stroke(paths[0]); ctx.strokeStyle=colors.line; ctx.setLineDash([1,1]); ctx.stroke(paths[1]); ctx.setLineDash([]); ctx.strokeStyle=colors.edge; ctx.lineWidth=.8; ctx.stroke(paths[2]); drawCurrentLane(target, colors); drawGhosts(f); - for(const a of getFrameAgents(f)){ ctx.save(); ctx.translate(a.x,a.y); ctx.rotate(a.h); drawAgentBody(a, darkMode?'#fff':'#111'); drawPerturbationOutlines(a); ctx.restore(); ctx.save(); ctx.translate(a.x,a.y); if(isEgoCam && target) ctx.rotate(-Math.PI/2 + target.h); else ctx.scale(1,-1); ctx.fillStyle=colors.text; ctx.font='600 '+(14/cam.z)+'px system-ui'; ctx.textAlign='center'; ctx.fillText(a.id,0,(isEgoCam && target)?a.w/2+.5:-a.w/2-.5); ctx.restore(); if(a.id === followedId){ ctx.save(); ctx.translate(a.x,a.y); ctx.strokeStyle=colors.accent; ctx.lineWidth=3/cam.z; ctx.beginPath(); ctx.arc(0,0,Math.max(a.l,a.w)*1.2,0,7); ctx.stroke(); ctx.restore(); } } + for(const a of getFrameAgents(f)){ ctx.save(); ctx.translate(a.x,a.y); ctx.rotate(a.h); drawAgentBody(a, darkMode?'#fff':'#111'); drawPerturbationOutlines(a); drawCollisionFootprint(a, f, darkMode?'#fff':'#111'); ctx.restore(); ctx.save(); ctx.translate(a.x,a.y); if(isEgoCam && target) ctx.rotate(-Math.PI/2 + target.h); else ctx.scale(1,-1); ctx.fillStyle=colors.text; ctx.font='600 '+(14/cam.z)+'px system-ui'; ctx.textAlign='center'; ctx.fillText(a.id,0,(isEgoCam && target)?a.w/2+.5:-a.w/2-.5); ctx.restore(); if(a.id === followedId){ ctx.save(); ctx.translate(a.x,a.y); ctx.strokeStyle=colors.accent; ctx.lineWidth=3/cam.z; ctx.beginPath(); ctx.arc(0,0,Math.max(a.l,a.w)*1.2,0,7); ctx.stroke(); ctx.restore(); } } for(let i=0;i (header dict, chunk name -> ndarray view) of a saved .replay.zlib.""" +def replay_step_link(step, seconds=None): + """Link for a panel injected into an interactive replay page: clicking seeks the viewer to `step`.""" + label = f"step {int(step)}" + ("" if seconds is None else f" ({seconds:.1f} s)") + return ( + f'{label}' + ) + + +def _inflate_through(decompressor, raw, needed_bytes): + rounds = 0 + while len(raw) < needed_bytes and decompressor.unconsumed_tail and rounds < REPLAY_INFLATE_MAX_ROUNDS: + raw += decompressor.decompress(decompressor.unconsumed_tail, needed_bytes - len(raw)) + rounds += 1 + if len(raw) < needed_bytes: + raise ValueError(f"truncated replay: {len(raw)} of {needed_bytes} bytes") + return raw + + +def _inflate_replay_prefix(compressed, chunk_names): + """Inflate a .replay.zlib only through the last of chunk_names (a long route inflates to GBs).""" + decompressor = zlib.decompressobj() + raw = bytearray(decompressor.decompress(compressed, REPLAY_HEADER_PROBE_BYTES)) + header_end = 4 + struct.unpack(" (header dict, chunk name -> ndarray view) of a saved .replay.zlib; chunk_names limits the chunks + returned and the inflated prefix of the stream.""" with open(replay_path, "rb") as replay_file: - raw = zlib.decompress(replay_file.read()) + compressed = replay_file.read() + raw = zlib.decompress(compressed) if chunk_names is None else _inflate_replay_prefix(compressed, chunk_names) header_len = struct.unpack("/renders//replays [dt] +(needs eval.keep_zlib_replays=true on the render) +""" + +import glob +import os +import sys +from collections import Counter + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from pufferlib.viz import read_replay_zlib # noqa: E402 + +COLLISION_IDX, OFFROAD_IDX, RED_LIGHT_IDX, STOP_SIGN_IDX, LANE_DIST_IDX = 0, 1, 2, 3, 5 +I_TYPE, I_ACTIVE, I_STOPPED, I_REMOVED, I_SLOT, I_BLIND, I_PHANTOM = 1, 3, 4, 5, 7, 8, 9 +F_X, F_Y, F_HEADING, F_LENGTH, F_WIDTH, F_SPEED = 0, 1, 3, 4, 5, 6 +PERTURBATION_LOOKBACK_FRAMES = 10 # 3 s bouts at dt 0.3 +SPAWN_FRAMES = 10 + + +def first_frame(flags): + hits = np.flatnonzero(flags) + return int(hits[0]) if hits.size else None + + +def classify_collision(f32, i32, m, t, a): + pos = f32[t][:, [F_X, F_Y]] + valid = i32[t, :, 2] == 1 + valid[a] = False + dist = np.linalg.norm(pos - pos[a], axis=1) + dist[~valid] = np.inf + # Prefer the agent whose collision flag fires in the same frame; fall back to the nearest body. + co_flagged = valid & (m[t, :, COLLISION_IDX] == 1) + if co_flagged.any(): + dist_co = np.where(co_flagged, dist, np.inf) + p = int(np.argmin(dist_co)) + else: + p = int(np.argmin(dist)) + gap = dist[p] - 0.5 * (f32[t, a, F_LENGTH] + f32[t, p, F_LENGTH]) + if not np.isfinite(dist[p]) or gap > 2.0: + return "no partner found", None + t0 = max(0, t - PERTURBATION_LOOKBACK_FRAMES) + tp = t - 1 if t > 0 else t + tags = [] + if i32[t0 : t + 1, p, I_BLIND].any(): + tags.append("partner blind now") + if i32[t0 : t + 1, p, I_PHANTOM].any(): + tags.append("partner phantom-braking now") + if not tags and (i32[:, p, I_BLIND].any() or i32[:, p, I_PHANTOM].any()): + tags.append("partner is a designated perturber (bout elsewhere)") + if t > 0 and i32[t - 1, p, I_STOPPED] == 1: + tags.append("partner parked by earlier infraction") + elif f32[tp, p, F_SPEED] < 0.3: + tags.append("partner stationary") + if i32[t, p, I_ACTIVE] == 0: + tags.append("partner not policy-controlled") + if not tags: + if m[t, p, COLLISION_IDX] == 1 and first_frame(m[:, p, COLLISION_IDX]) == t: + tags.append("both moving, mutual first collision") + else: + tags.append("both moving, other") + dh = abs((f32[t, a, F_HEADING] - f32[t, p, F_HEADING] + np.pi) % (2 * np.pi) - np.pi) + rel = pos[p] - pos[a] + ahead = np.cos(f32[t, a, F_HEADING]) * rel[0] + np.sin(f32[t, a, F_HEADING]) * rel[1] > 0 + if dh < np.radians(30): + geometry = "rear-end, ego behind" if ahead else "rear-end, ego in front" + elif dh > np.radians(150): + geometry = "head-on" + else: + geometry = "crossing/side" + detail = f"ego {f32[tp, a, F_SPEED]:.1f} m/s, partner {p} {f32[tp, p, F_SPEED]:.1f} m/s, {geometry}" + return " + ".join(tags), detail + + +def analyze(path, dt): + _, c = read_replay_zlib(path, chunk_names=["agent_f32", "agent_i32", "metrics_f32"]) + f32, i32, m = c["agent_f32"], c["agent_i32"], c["metrics_f32"] + frames, agents = m.shape[:2] + rows = [] + logged = 0 + for a in range(agents): + if i32[0, a, I_ACTIVE] != 1: + continue + perturbed = bool(i32[:, a, I_BLIND].any() or i32[:, a, I_PHANTOM].any()) + if not perturbed: + logged += 1 + events = { + "collision": first_frame(m[:, a, COLLISION_IDX]), + "offroad": first_frame(m[:, a, OFFROAD_IDX]), + "red_light": first_frame(m[:, a, RED_LIGHT_IDX]), + "stop_sign": first_frame(m[:, a, STOP_SIGN_IDX]), + } + events = {k: v for k, v in events.items() if v is not None} + if not events: + continue + kind = min(events, key=events.get) + t = events[kind] + tp = max(t - 1, 0) + speed = f32[tp, a, F_SPEED] + if kind == "collision": + tag, detail = classify_collision(f32, i32, m, t, a) + elif kind == "offroad": + tag = ( + "offroad within spawn frames" + if t < SPAWN_FRAMES + else ( + "offroad at creep speed (<1 m/s)" + if speed < 1.0 + else ("offroad slow (1-4 m/s)" if speed < 4.0 else "offroad fast (>4 m/s)") + ) + ) + detail = f"{speed:.1f} m/s, lane dist {m[tp, a, LANE_DIST_IDX]:.1f} m" + else: + tag = kind + detail = f"{speed:.1f} m/s" + rows.append((kind, tag, detail, t * dt, a, perturbed)) + return frames, agents, logged, rows + + +def main(): + replays_dir = sys.argv[1] + dt = float(sys.argv[2]) if len(sys.argv) > 2 else 0.3 + paths = sorted(glob.glob(os.path.join(replays_dir, "*.replay.zlib"))) + totals = Counter() + logged_total = 0 + for path in paths: + frames, agents, logged, rows = analyze(path, dt) + logged_total += logged + print(f"\n== {os.path.basename(path)}: {frames} frames, {agents} agents ({logged} logged policy agents)") + for kind, tag, detail, t_s, a, perturbed in sorted(rows, key=lambda r: r[3]): + flag = " [ego perturbed, not logged]" if perturbed else "" + print(f" t={t_s:6.1f}s agent {a:3d} {kind:9s} {tag}: {detail}{flag}") + if not perturbed: + totals[(kind, tag)] += 1 + n_inf = sum(totals.values()) + print( + f"\n== totals over {len(paths)} episodes, {logged_total} logged agent-episodes, {n_inf} first infractions" + f" ({n_inf / max(logged_total, 1):.3f} per agent-episode)" + ) + for (kind, tag), n in sorted(totals.items(), key=lambda kv: -kv[1]): + print(f" {n:3d} {kind:9s} {tag}") + + +if __name__ == "__main__": + main() diff --git a/scripts/baseline_10G/config-reward-dense.yaml b/scripts/baseline_10G/config-reward-dense.yaml index 0c1386b24e..1ae8495d20 100644 --- a/scripts/baseline_10G/config-reward-dense.yaml +++ b/scripts/baseline_10G/config-reward-dense.yaml @@ -51,6 +51,7 @@ env: obs_slots_traffic_controls_n: 4 offroad_behavior: 1 partner_blindness_prob: 0.0 + partner_blindness_prob_min: 0.0 partner_blindness_trigger_prob: 0.0 phantom_braking_duration: 10 phantom_braking_prob: 0.0 diff --git a/scripts/baseline_10G/config-reward-sparse.yaml b/scripts/baseline_10G/config-reward-sparse.yaml index cda5e9d4bc..a3f21b8cd7 100644 --- a/scripts/baseline_10G/config-reward-sparse.yaml +++ b/scripts/baseline_10G/config-reward-sparse.yaml @@ -51,6 +51,7 @@ env: obs_slots_traffic_controls_n: 4 offroad_behavior: 1 partner_blindness_prob: 0.0 + partner_blindness_prob_min: 0.0 partner_blindness_trigger_prob: 0.0 phantom_braking_duration: 10 phantom_braking_prob: 0.0 diff --git a/scripts/cluster_configs/single_agent_speed_run.yaml b/scripts/cluster_configs/single_agent_speed_run.yaml index b6dd076258..f1fea24a89 100644 --- a/scripts/cluster_configs/single_agent_speed_run.yaml +++ b/scripts/cluster_configs/single_agent_speed_run.yaml @@ -57,6 +57,7 @@ env.reward_overspeed: 0.0 # Behavioral perturbations off (single agent, no partners to perturb) env.phantom_braking_prob: 0 env.partner_blindness_prob: 0 +env.partner_blindness_prob_min: 0 # Policy: smaller shared-backbone net than the default split actor/critic policy.partner_input_size: 128 diff --git a/scripts/config-reward-dense.yaml b/scripts/config-reward-dense.yaml index 0c1386b24e..1ae8495d20 100644 --- a/scripts/config-reward-dense.yaml +++ b/scripts/config-reward-dense.yaml @@ -51,6 +51,7 @@ env: obs_slots_traffic_controls_n: 4 offroad_behavior: 1 partner_blindness_prob: 0.0 + partner_blindness_prob_min: 0.0 partner_blindness_trigger_prob: 0.0 phantom_braking_duration: 10 phantom_braking_prob: 0.0 diff --git a/scripts/config-reward-sparse.yaml b/scripts/config-reward-sparse.yaml index cda5e9d4bc..a3f21b8cd7 100644 --- a/scripts/config-reward-sparse.yaml +++ b/scripts/config-reward-sparse.yaml @@ -51,6 +51,7 @@ env: obs_slots_traffic_controls_n: 4 offroad_behavior: 1 partner_blindness_prob: 0.0 + partner_blindness_prob_min: 0.0 partner_blindness_trigger_prob: 0.0 phantom_braking_duration: 10 phantom_braking_prob: 0.0 diff --git a/scripts/eval/analyze_carla_cosim.py b/scripts/eval/analyze_carla_cosim.py index 7e8a008c22..fc69a275a7 100644 --- a/scripts/eval/analyze_carla_cosim.py +++ b/scripts/eval/analyze_carla_cosim.py @@ -1,9 +1,10 @@ """Failure analysis + visualizations for CARLA leaderboard runs of the PufferDrive co-sim agent (pufferlib/ocean/cosim/carla/leaderboard_agent.py run with COSIM_WORLD_LOG and COSIM_TELEMETRY set). Same report shape as analyze_nuplan_cosim.py: score table worst first, failure-category counts, -per-route diagnosis, six-frame top-down strip + speed plot + chase-cam / top-down video. +per-route diagnosis, six-frame top-down strip + speed plot + chase-cam video. The matplotlib top-down mp4 +is opt-in (--video): it costs about a minute per 1000 logged steps, the obs replay pages cover the same ground. -usage: python scripts/eval/analyze_carla_cosim.py [ ...] [--max-inline N] [--no-video] +usage: python scripts/eval/analyze_carla_cosim.py [ ...] [--max-inline N] [--video] [--obs-html-dir D] holds result.json (leaderboard_evaluator --checkpoint), world_log/*.npz, telemetry/*.csv and @@ -38,6 +39,7 @@ from render_carla_obs_html import page_name import data_utils.mirror_map_bin as mbin +import pufferlib from pufferlib.ocean.cosim.carla_cosim import write_mp4 @@ -97,7 +99,12 @@ def __init__(self, npz_path): self.meta = json.loads(str(d["meta"])) self.tick_dt = float(self.meta["tick_dt"]) self.t = self.ego[:, 0] * self.tick_dt - bin_data = mbin.read_bin(Path(self.meta["town_bin"])) + town_bin = Path(self.meta["town_bin"]) + if not town_bin.exists(): # log written on another machine: this checkout's copy of the same town bin + town_bin = ( + Path(pufferlib.__file__).resolve().parent / "resources" / "drive" / "binaries" / "carla" / town_bin.name + ) + bin_data = mbin.read_bin(town_bin) lo = self.ego[:, 1:3].min(axis=0) - MAP_CROP_MARGIN_M hi = self.ego[:, 1:3].max(axis=0) + MAP_CROP_MARGIN_M self.lanes, self.edges = [], [] @@ -367,7 +374,7 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("paths", nargs="+", help="run dirs followed by the report dir") parser.add_argument("--max-inline", type=int, default=12) - parser.add_argument("--no-video", action="store_true") + parser.add_argument("--video", action="store_true", help="also render the matplotlib top-down mp4 per route (slow)") parser.add_argument("--obs-html-dir", help="the run's obs replay gallery dir, linked per route") args = parser.parse_args() run_dirs, report_dir = args.paths[:-1], Path(args.paths[-1]) @@ -377,7 +384,7 @@ def main(): rows = [] for run_dir in run_dirs: for record, log_path, chase_mp4, tag in collect_routes(run_dir): - rows.append(render_route(record, RouteLog(log_path), chase_mp4, report_dir, tag, not args.no_video)) + rows.append(render_route(record, RouteLog(log_path), chase_mp4, report_dir, tag, args.video)) print(f"{tag}: {rows[-1]['categories']} | {rows[-1]['diagnosis']}") df = pd.DataFrame(rows).sort_values("driving_score") df.to_csv(report_dir / "routes.csv", index=False) diff --git a/scripts/eval/nuplan_log_remap.py b/scripts/eval/nuplan_log_remap.py index 2f6496a044..4de095b72c 100644 --- a/scripts/eval/nuplan_log_remap.py +++ b/scripts/eval/nuplan_log_remap.py @@ -6,7 +6,8 @@ when splits/val is absent) and provides `load_log`, which unpickles with placeholder classes for any module that cannot be imported here. -Env: CLUSTER_NUPLAN_ROOT (default /home/shared/data/nuplan), NUPLAN_DATA_ROOT, NUPLAN_MAPS_ROOT. +Env: CLUSTER_NUPLAN_ROOT (the root baked into the logs: default /home/shared/data/nuplan for scaleway, +/cephfs/data/nuplan for logs from kesai), NUPLAN_DATA_ROOT, NUPLAN_MAPS_ROOT. """ import io diff --git a/scripts/eval/render_alpasim_obs_html.py b/scripts/eval/render_alpasim_obs_html.py new file mode 100644 index 0000000000..377803d051 --- /dev/null +++ b/scripts/eval/render_alpasim_obs_html.py @@ -0,0 +1,287 @@ +"""Render the interactive observation replays of an AlpaSim run (scripts/kesai/12_alpasim.sh: +/obs_replays/.replay.zlib) into one page per scene plus an index.html gallery (same +navigator as the nuPlan, CARLA and self-play eval replays). + +usage: python scripts/eval/render_alpasim_obs_html.py [--prune] [--workers N] + +Pages are named s__.html so the gallery lists the worst first (flags: COL at-fault +collision, OFF offroad, CORR left the corridor laterally, PROG progress score < 1); each page carries the +scene's AlpaSim score criteria and per-metric values as a panel, every violated per-step metric linking to the +replay step of its first violation. The page the run wrote (.html) is replaced by the renamed one. +--prune deletes each .replay.zlib once its page is written (the page embeds the same bytes). +Reads the tidy_alpasim_run.py layout (/alpasim/*) and the untidied AlpaSim trees of older runs. +""" + +import argparse +import html +import json +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np +import pandas as pd +import yaml + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) +from pufferlib import viz # noqa: E402 + +FLAG_METRICS = ( + ("COL", "collision_at_fault"), + ("OFF", "offroad"), + ("CORR", "left_corridor_laterally"), +) +PROGRESS_FLAG = "PROG" +SCORED_METRICS = ( + ("collision_at_fault", "at-fault collision"), + ("offroad", "offroad"), + ("left_corridor_laterally", "left corridor laterally"), +) +UNSCORED_SERIES = ( + ("collision_any", "collision (any)"), + ("collision_front", "collision front"), + ("collision_lateral", "collision lateral"), + ("collision_rear", "collision rear"), + ("wrong_lane", "wrong lane"), +) +PER_STEP_SERIES = ( + "collision_any", + "collision_front", + "collision_lateral", + "collision_rear", + "offroad", + "left_corridor_laterally", + "wrong_lane", + "eval_relevant", + "dist_to_gt_trajectory", +) +ALPASIM_STEP_SECONDS = 0.1 + + +def run_layout(run_dir): + """-> (obs replay dir, results-summary.json, metrics_unprocessed.parquet, eval-config.yaml) of the tidied + layout, else of an AlpaSim wizard tree (2026-10-12 layout, then the older alpagym/ nesting).""" + obs_dirs = [base / "obs_replays" for base in (run_dir, run_dir / "alpagym") if (base / "obs_replays").is_dir()] + if not obs_dirs: + raise SystemExit(f"no obs_replays/ under {run_dir} or {run_dir / 'alpagym'}") + tidied = run_dir / "alpasim" + candidates = [ + (tidied / "results-summary.json", tidied / "metrics_unprocessed.parquet", tidied / "eval-config.yaml") + ] + for wizard_dir in (run_dir / "alpasim" / "wizard_0", run_dir / "alpagym" / "alpasim" / "wizard_0"): + aggregate = wizard_dir / "aggregate" + candidates.append( + ( + aggregate / "results-summary.json", + aggregate / "metrics_unprocessed.parquet", + wizard_dir / "eval-config.yaml", + ) + ) + for summary_path, metrics_path, eval_config_path in candidates: + if summary_path.is_file(): + return obs_dirs[0], summary_path, metrics_path, eval_config_path + raise SystemExit(f"no results-summary.json under {run_dir}: the run has not been scored") + + +def scene_rows(summary_path): + """clipgt -> results-summary rollout row (score, status, failure_reason, metrics, score_metrics).""" + summary = json.loads(summary_path.read_text()) + rows = {} + for row in summary["rollouts"]: + if row["clipgt_id"] in rows: + raise SystemExit(f"{row['clipgt_id']} has several rollouts; the scene's single replay cannot be paired") + rows[row["clipgt_id"]] = row + return rows, summary.get("score_criteria", {}) + + +def max_gt_dist_m(eval_config_path): + return float(yaml.safe_load(eval_config_path.read_text())["aggregation_modifiers"]["max_dist_to_gt_trajectory"]) + + +def scene_series(metrics_path, rows): + """clipgt -> its rollout's per-step metrics (PER_STEP_SERIES rows of metrics_unprocessed.parquet).""" + df = pd.read_parquet(metrics_path, columns=["name", "timestamps_us", "values", "clipgt_id", "rollout_id"]) + df = df[df["name"].isin(PER_STEP_SERIES)] + series = {} + for (clipgt, rollout_id), scene_df in df.groupby(["clipgt_id", "rollout_id"]): + if clipgt in rows and rows[clipgt]["rollout_id"] == rollout_id: + series[clipgt] = scene_df + return series + + +def scoring_timeline(df, frames, stop_gt_dist_m): + """AlpaSim scores from the first eval_relevant step through the first collision, offroad or stop_gt_dist_m + from the gt trajectory. -> {scored: series -> first replay step > 0 inside that window, after: first one past it, + stop_step, stop_reasons}. The replay ends with the per-step series, so index i is replay step i - (length - frames).""" + wide = df.pivot(index="timestamps_us", columns="name", values="values").sort_index() + wide["collision_at_fault"] = ((wide["collision_front"] > 0) | (wide["collision_lateral"] > 0)).astype(np.float64) + offset = len(wide) - frames + relevant = np.flatnonzero(wide["eval_relevant"].to_numpy() > 0) + start = int(relevant[0]) if relevant.size else len(wide) + stop_events = { + "collision": wide["collision_any"].to_numpy() > 0, + "offroad": wide["offroad"].to_numpy() > 0, + f">= {stop_gt_dist_m:g} m from the gt trajectory": wide["dist_to_gt_trajectory"].to_numpy() >= stop_gt_dist_m, + } + stop_hits = np.flatnonzero(np.logical_or.reduce(list(stop_events.values())) & (np.arange(len(wide)) >= start)) + stop = int(stop_hits[0]) if stop_hits.size else len(wide) - 1 + timeline = {"scored": {}, "after": {}, "stop_step": None, "stop_reasons": []} + if stop_hits.size: + timeline["stop_step"] = stop - offset + timeline["stop_reasons"] = [reason for reason, hits in stop_events.items() if hits[stop]] + for name in [metric for metric, _ in SCORED_METRICS] + [series for series, _ in UNSCORED_SERIES]: + hits = np.flatnonzero(wide[name].to_numpy() > 0) + inside, later = hits[(hits >= start) & (hits <= stop)], hits[hits > stop] + if inside.size: + timeline["scored"][name] = int(inside[0]) - offset + if later.size: + timeline["after"][name] = int(later[0]) - offset + return timeline + + +def step_html(step): + if step < 0: + return f"{-step * ALPASIM_STEP_SECONDS:.1f} s before the handover" + return viz.replay_step_link(step, step * ALPASIM_STEP_SECONDS) + + +def metric_line(label, value, failed, name, timeline): + color = "#ff6b6b" if failed else "#7ed491" + mark = "✗" if failed else "✓" + detail = "" + if failed and name in timeline["scored"]: + detail = f" · {step_html(timeline['scored'][name])}" + elif not failed and name in timeline["after"]: + detail = ( + f' · {step_html(timeline["after"][name])}, after scoring stopped' + ) + return ( + f'
{mark} {label}' + f'{value}{detail}
' + ) + + +def score_panel(row, criteria, timeline): + """Fixed panel injected into the viewer: AlpaSim score, every score criterion, the unscored per-step metrics, + each violation linking to its first step in the scored window.""" + metrics, score_metrics = row["metrics"], row["score_metrics"] + lines = [] + for metric, label in SCORED_METRICS: + value = float(score_metrics[metric]) + criterion = html.escape(criteria.get(metric, "")) + lines.append( + metric_line( + label, f'{value:.0f} ({criterion})', value > 0, metric, timeline + ) + ) + progress_score = float(score_metrics["progress_score"]) + lines.append( + metric_line( + "progress score", + f'{progress_score:.2f} (clipped rel {float(score_metrics["progress_clipped_rel"]):.2f}, ' + f"full at 0.8)", + progress_score < 1.0, + "progress", + timeline, + ) + ) + if timeline["stop_step"] is not None: + lines.append( + f'
scoring stopped at {step_html(timeline["stop_step"])}: ' + f"{', '.join(timeline['stop_reasons'])}
" + ) + for series, label in UNSCORED_SERIES: + value = float(metrics[series]) + lines.append( + metric_line(f'{label}', f"{value:.0f}", value > 0, series, timeline) + ) + lines.append( + f'
driven {float(metrics["dist_traveled_m"]):.0f} m of gt {float(metrics["gt_dist_traveled_m"]):.0f} m · ' + f"gt trajectory dist {float(metrics['dist_to_gt_trajectory']):.1f} m (lateral {float(metrics['lateral_dist_to_gt_trajectory']):.1f} m)
" + ) + failure = f" · {html.escape(row['failure_reason'])}" if row["failure_reason"] else "" + return ( + '
' + f'AlpaSim score {float(row["score"]):.2f} · ' + f"{html.escape(row['status'])}{failure} · {html.escape(row['clipgt_id'])}" + + "".join(lines) + + "
" + ) + + +def page_name(clipgt, row): + if row is None: + return f"{clipgt}.html" + flags = [flag for flag, metric in FLAG_METRICS if float(row["score_metrics"][metric]) > 0] + if float(row["score_metrics"]["progress_score"]) < 1.0: + flags.append(PROGRESS_FLAG) + return f"s{float(row['score']):.2f}_" + "".join(f"{flag}_" for flag in flags) + f"{clipgt}.html" + + +def gallery_metrics(row): + """Failure flags in the names the self-play gallery filters on (rate > 0 = flagged).""" + if row is None: + return {} + return { + "score": float(row["score"]), + "collision_rate": float(row["metrics"]["collision_any"]), + "at_fault_collision_rate": float(row["score_metrics"]["collision_at_fault"]), + "offroad_rate": float(row["score_metrics"]["offroad"]), + } + + +def render_page(zlib_path, html_path, panel): + viz.render_interactive_replay_zlib(str(zlib_path), str(html_path)) + if panel: + page = Path(html_path).read_text() + Path(html_path).write_text(page.replace("", "" + panel, 1)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("run_dir", help="output dir of 12_alpasim.sh (holds obs_replays/ and alpasim/)") + ap.add_argument("--prune", action="store_true", help="delete each .replay.zlib once its page is written") + ap.add_argument("--workers", type=int, default=min(16, os.cpu_count() or 1)) + args = ap.parse_args() + obs_dir, summary_path, metrics_path, eval_config_path = run_layout(Path(args.run_dir)) + replays = {p.name[: -len(".replay.zlib")]: p for p in sorted(obs_dir.glob("*.replay.zlib"))} + if not replays: + raise SystemExit(f"no {obs_dir}/*.replay.zlib") + rows, criteria = scene_rows(summary_path) + series = scene_series(metrics_path, rows) + stop_gt_dist_m = max_gt_dist_m(eval_config_path) + jobs = [(clipgt, replay, obs_dir / page_name(clipgt, rows.get(clipgt))) for clipgt, replay in replays.items()] + + def render(job): + clipgt, replay, html_path = job + row = rows.get(clipgt) + panel = "" + if row is not None: + frames = int(viz.read_replay_zlib(replay, chunk_names=("metrics_f32",))[0]["frames"]) + panel = score_panel(row, criteria, scoring_timeline(series[clipgt], frames, stop_gt_dist_m)) + render_page(replay, html_path, panel) + run_page = obs_dir / f"{clipgt}.html" + if run_page != html_path and run_page.exists(): + run_page.unlink() + + with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool: + list(pool.map(render, jobs)) + file_metrics = {html_path.name: gallery_metrics(rows.get(clipgt)) for clipgt, _, html_path in jobs} + viz.build_gallery_index(str(obs_dir), file_metrics=file_metrics) + if args.prune: + for replay in replays.values(): + replay.unlink() + unscored = sum(clipgt not in rows for clipgt in replays) + print( + f"[render_alpasim_obs_html] {len(jobs)} scene pages rendered" + f"{f', {unscored} without a score (no panel)' if unscored else ''}{', replays pruned' if args.prune else ''}" + f" -> {obs_dir / 'index.html'}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/eval/render_carla_obs_html.py b/scripts/eval/render_carla_obs_html.py index 37b47f1efc..7911c902c7 100644 --- a/scripts/eval/render_carla_obs_html.py +++ b/scripts/eval/render_carla_obs_html.py @@ -1,5 +1,5 @@ -"""Render the interactive observation replays of a CARLA longest6 run (scripts/kesai/11_carla_longest6.sh: -/routes/route_XX/obs_html/*.replay.zlib) into /obs_html/, one page per route plus an +"""Render the interactive observation replays of a CARLA longest6 run (scripts/kesai/11_carla_longest6_v2.sh: +/routes/route_XX[_s]/obs_html/*.replay.zlib) into /obs_html/, one page per route plus an index.html gallery (same navigator as the nuPlan and self-play eval replays). usage: python scripts/eval/render_carla_obs_html.py [--routes ] [--workers N] @@ -7,7 +7,8 @@ Pages are named ds___route.html so the gallery lists the worst first (flags: COL collision, RED red light, STOP stop sign, OFF outside route lanes, DEV route deviation, BLOCK blocked, TIMEOUT route/scenario timeout); each page carries the route's leaderboard scores and -infractions as a panel. --routes (the evaluated xml) names the town of every route, else the world log does. +infractions as a panel, every positioned infraction linking to the replay step where the ego was nearest. +--routes (the evaluated xml) names the town of every route, else the world log does. Retried routes leave the files of crashed attempts behind: the newest files pair with the records. """ @@ -28,6 +29,7 @@ REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT)) from pufferlib import viz +from pufferlib.ocean.cosim.carla_bridge import TOWN_OFFSETS COLLISION_KEYS = ("collisions_pedestrian", "collisions_vehicle", "collisions_layout") @@ -55,7 +57,10 @@ ("yield_emergency_vehicle_infractions", "emergency vehicles"), ) MESSAGE_MAX_CHARS = 120 +POSITION_PATTERN = re.compile(r"x=([-\d.]+), y=([-\d.]+)") +NEAREST_STEP_MAX_M = 5.0 ROUTE_ID_PATTERN = re.compile(r"RouteScenario_(\d+)(?:_rep(\d+))?") +SEED_DIR_PATTERN = re.compile(r"_s(\d+)$") # route_08_s1: seed 1 of route 8 (the eval scripts drive one pass per seed) def infraction_messages(record, key): @@ -70,18 +75,18 @@ def route_index_rep(record): return int(match.group(1)), int(match.group(2) or 0) -def page_name(record, town): +def page_name(record, town, seed=None): route_idx, rep = route_index_rep(record) parts = [f"ds{float(record['scores']['score_composed']):05.1f}"] parts += [flag for flag, keys in FLAG_INFRACTIONS if any(infraction_messages(record, key) for key in keys)] if town: parts.append(town) - parts.append(f"route{route_idx:02d}" + (f"_rep{rep}" if rep else "")) + parts.append(f"route{route_idx:02d}" + (f"_rep{rep}" if rep else "") + (f"_s{seed}" if seed is not None else "")) return "_".join(parts) + ".html" -def score_panel(record, town): - """Fixed panel injected into the viewer: leaderboard scores, every infraction category, the messages.""" +def score_panel(record, town, steps, seed=None): + """Fixed panel injected into the viewer: leaderboard scores, every infraction category, the messages + steps.""" scores, meta = record["scores"], record.get("meta", {}) route_idx, rep = route_index_rep(record) lines = [] @@ -94,15 +99,21 @@ def score_panel(record, town): f'{len(messages)}' ) for message in messages: + step_html = f" · {steps[message]}" if message in steps else "" lines.append( - f'
↳ {html.escape(message[:MESSAGE_MAX_CHARS])}
' + f'
↳ {html.escape(message[:MESSAGE_MAX_CHARS])}' + f"{step_html}
" ) if meta: lines.append( f'
route {float(meta.get("route_length", 0)):.0f} m · ' f"game {float(meta.get('duration_game', 0)):.0f} s · wall {float(meta.get('duration_system', 0)):.0f} s
" ) - title = f"{town + ' ' if town else ''}route {route_idx}" + (f" rep {rep}" if rep else "") + title = ( + f"{town + ' ' if town else ''}route {route_idx}" + + (f" rep {rep}" if rep else "") + + (f" seed {seed}" if seed is not None else "") + ) return ( '
seek link of every positioned infraction: the replay step where the ego was nearest to it. + CARLA -> bin offset from the route's world log (calibrated), else the town default.""" + offset = log_meta.get("offset") or TOWN_OFFSETS.get(town) + if offset is None: + return {} + tick_dt = log_meta.get("tick_dt") + _, chunks = viz.read_replay_zlib(replay_path, chunk_names=("agent_f32",)) + ego_xy = chunks["agent_f32"][:, 0, :2] + messages = [message for key, _ in PANEL_INFRACTIONS for message in infraction_messages(record, key)] + steps = {} + for message in messages: + match = POSITION_PATTERN.search(message) + if match is None: + continue + bin_x, bin_y = float(match.group(1)) + offset[0], -float(match.group(2)) + offset[1] + distance_m = np.hypot(ego_xy[:, 0] - bin_x, ego_xy[:, 1] - bin_y) + step = int(np.argmin(distance_m)) + if distance_m[step] > NEAREST_STEP_MAX_M: + steps[message] = "beyond the replay" + continue + steps[message] = viz.replay_step_link(step, None if tick_dt is None else step * tick_dt) + return steps + + def render_page(zlib_path, html_path, panel): viz.render_interactive_replay_zlib(str(zlib_path), str(html_path)) page = Path(html_path).read_text() @@ -137,8 +173,8 @@ def route_towns(routes_xml): } -def world_log_town(npz_path): - return json.loads(str(np.load(npz_path)["meta"])).get("town", "") +def world_log_meta(npz_path): + return json.loads(str(np.load(npz_path)["meta"])) def newest_files(route_dir, subdir, suffix, count): @@ -148,9 +184,11 @@ def newest_files(route_dir, subdir, suffix, count): def collect_jobs(run_dir, towns): - """-> [(record, replay zlib path, town)] over routes/route_*, the newest replay per route record.""" + """-> [(record, replay zlib path, town, world log meta, seed or None)] over routes/route_*, the newest replay per record.""" jobs = [] for route_dir in sorted(Path(run_dir, "routes").glob("route_*")): + seed_match = SEED_DIR_PATTERN.search(route_dir.name) + seed = seed_match.group(1) if seed_match else None result_path = route_dir / "result.json" if not result_path.is_file(): print(f"[{route_dir.name}] no result.json, skipped") @@ -165,16 +203,15 @@ def collect_jobs(run_dir, towns): records = records[-len(replays) :] world_logs = newest_files(route_dir, "world_log", ".npz", len(replays)) for k, (record, replay) in enumerate(zip(records, replays)): - town = towns.get(route_index_rep(record)[0]) - if town is None and k < len(world_logs): - town = world_log_town(world_logs[k]) - jobs.append((record, replay, town or "")) + log_meta = world_log_meta(world_logs[k]) if k < len(world_logs) else {} + town = towns.get(route_index_rep(record)[0]) or log_meta.get("town", "") + jobs.append((record, replay, town, log_meta, seed)) return jobs def main(): ap = argparse.ArgumentParser() - ap.add_argument("run_dir", help="output dir of 11_carla_longest6.sh (holds routes/route_*/)") + ap.add_argument("run_dir", help="output dir of 11_carla_longest6_v2.sh (holds routes/route_*/)") ap.add_argument("--routes", help="evaluated route xml, names the town of every route") ap.add_argument("--workers", type=int, default=min(16, os.cpu_count() or 1)) args = ap.parse_args() @@ -185,10 +222,20 @@ def main(): raise SystemExit(f"no routes/route_*/obs_html/*.replay.zlib with a result.json under {run_dir}") obs_dir = run_dir / "obs_html" obs_dir.mkdir(exist_ok=True) - pages = [(obs_dir / page_name(record, town), record, replay, town) for record, replay, town in jobs] + pages = [ + (obs_dir / page_name(record, town, seed), record, replay, town, log_meta, seed) + for record, replay, town, log_meta, seed in jobs + ] + + def render(page): + html_path, record, replay, town, log_meta, seed = page + render_page( + replay, html_path, score_panel(record, town, infraction_steps(record, replay, town, log_meta), seed) + ) + with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool: - list(pool.map(lambda job: render_page(job[2], job[0], score_panel(job[1], job[3])), pages)) - file_metrics = {page.name: gallery_metrics(record) for page, record, _, _ in pages} + list(pool.map(render, pages)) + file_metrics = {page.name: gallery_metrics(record) for page, record, *_ in pages} report = run_dir / "report" / "index.html" links = [("Analysis report", os.path.relpath(report, obs_dir))] if report.is_file() else [] viz.build_gallery_index(str(obs_dir), file_metrics=file_metrics, links=links) diff --git a/scripts/eval/render_obs_html.py b/scripts/eval/render_obs_html.py index 7cad4b7b31..26130fb559 100644 --- a/scripts/eval/render_obs_html.py +++ b/scripts/eval/render_obs_html.py @@ -7,7 +7,8 @@ Selection: aggregated score < --max-score OR any --metric < 1 OR explicitly listed --tokens. Pages are named s___.html so the gallery lists the worst first -(flags: COL at-fault collision, OFF drivable area, DIR driving direction, STALL no progress). +(flags: COL at-fault collision, OFF drivable area, DIR driving direction, STALL no progress). Every violated +metric links to the replay step of its first violation (collisions: the shadow env's own flag). --prune deletes the .replay.zlib of every scenario that was not selected. """ @@ -18,6 +19,7 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path +import numpy as np import pandas as pd REPO_ROOT = Path(__file__).resolve().parents[2] @@ -50,14 +52,25 @@ "ego_yaw_acceleration", "ego_yaw_rate", ) -COMFORT_BOUNDS = { - "ego_lon_acceleration": "-4.05..2.40 m/s^2", - "ego_lon_jerk": "|j| <= 4.13 m/s^3", - "ego_jerk": "|j| <= 8.37 m/s^3", - "ego_lat_acceleration": "|a| <= 4.89 m/s^2", - "ego_yaw_acceleration": "|a| <= 1.93 rad/s^2", - "ego_yaw_rate": "|w| <= 0.95 rad/s", +COMFORT_BOUNDS = { # (lower, upper, label): nuPlan's within-bound thresholds + "ego_lon_acceleration": (-4.05, 2.40, "-4.05..2.40 m/s^2"), + "ego_lon_jerk": (-4.13, 4.13, "|j| <= 4.13 m/s^3"), + "ego_jerk": (-8.37, 8.37, "|j| <= 8.37 m/s^3"), + "ego_lat_acceleration": (-4.89, 4.89, "|a| <= 4.89 m/s^2"), + "ego_yaw_acceleration": (-1.93, 1.93, "|a| <= 1.93 rad/s^2"), + "ego_yaw_rate": (-0.95, 0.95, "|w| <= 0.95 rad/s"), } +NUPLAN_STEP_SECONDS = 0.1 +DRIVING_DIRECTION_COMPLIANCE_THRESHOLD_M = 2.0 +TTC_LEAST_MIN_S = 0.95 +COLLISION_METRIC = "no_ego_at_fault_collisions" +COLLISION_METRIC_IDX = viz.METRIC_LABELS.index("collision") +VIOLATION_SERIES = ( # scored metric, parquet with its per-step series, (lower, upper) of a compliant value + ("drivable_area_compliance", "corners_in_drivable_area", 1.0, None), + ("driving_direction_compliance", "driving_direction_compliance", -DRIVING_DIRECTION_COMPLIANCE_THRESHOLD_M, None), + ("time_to_collision_within_bound", "time_to_collision_within_bound", TTC_LEAST_MIN_S, None), + ("speed_limit_compliance", "speed_limit_compliance", None, 0.0), +) def scenario_rows(group_dir): @@ -71,6 +84,48 @@ def scenario_rows(group_dir): return rows +def first_out_of_bounds(values, lower=None, upper=None): + """-> first index with value < lower or value > upper, None when every value is within bounds.""" + values = np.asarray(values, dtype=np.float64) + outside = np.zeros(values.shape, dtype=bool) + if lower is not None: + outside |= values < lower + if upper is not None: + outside |= values > upper + hits = np.flatnonzero(outside) + return int(hits[0]) if hits.size else None + + +def step_link(step): + return viz.replay_step_link(step, step * NUPLAN_STEP_SECONDS) + + +def violation_steps(group_dir): + """token -> {scored metric: first violating step} from the per-step time series nuPlan writes per metric.""" + series_files = [ + (metric, Path(sim_dir) / "metrics" / f"{parquet_name}.parquet", lower, upper) + for sim_dir in glob.glob(f"{group_dir}/simulation/*/20*") + for metric, parquet_name, lower, upper in VIOLATION_SERIES + ] + out = {} + for metric, parquet_path, lower, upper in series_files: + if not parquet_path.exists(): + continue + for _, r in pd.read_parquet(parquet_path).iterrows(): + values = r["time_series_values"] + step = None if values is None else first_out_of_bounds(values, lower, upper) + if step is not None: + out.setdefault(str(r["scenario_name"]), {})[metric] = step + return out + + +def replay_collision_step(replay_path): + """-> first replay step where the shadow env flagged an ego collision (at fault or not), None when never.""" + _, chunks = viz.read_replay_zlib(replay_path, chunk_names=("metrics_f32",)) + hits = np.flatnonzero(chunks["metrics_f32"][:, 0, COLLISION_METRIC_IDX] > 0) + return int(hits[0]) if hits.size else None + + def comfort_failures(group_dir): """token -> ['ego_lon_jerk -2.7..5.0', ...]: nuPlan's own out-of-bound comfort sub-metrics.""" out = {} @@ -83,15 +138,20 @@ def comfort_failures(group_dir): flag = [c for c in df.columns if c.endswith("within_bounds_stat_value")] if not flag: continue + lower, upper, bound_label = COMFORT_BOUNDS[name] for _, r in df[~df[flag[0]].astype(bool)].iterrows(): + step = first_out_of_bounds(r["time_series_values"], lower, upper) + step_html = "" if step is None else f" · {step_link(step)}" out.setdefault(str(r["scenario_name"]), []).append( - f"{name} {float(r[f'min_{name}_stat_value']):.2f}..{float(r[f'max_{name}_stat_value']):.2f} (bound {COMFORT_BOUNDS[name]})" + f"{name} {float(r[f'min_{name}_stat_value']):.2f}..{float(r[f'max_{name}_stat_value']):.2f} " + f"(bound {bound_label}){step_html}" ) return out -def score_panel(token, row, comfort): - """Fixed panel injected into the viewer: nuPlan score, every scored sub-metric, violated ones highlighted.""" +def score_panel(token, row, comfort, steps): + """Fixed panel injected into the viewer: nuPlan score, every scored sub-metric, violated ones highlighted + with the replay step of their first violation.""" if row is None: return "" lines = [] @@ -99,9 +159,13 @@ def score_panel(token, row, comfort): value = float(row[metric]) color = "#ff6b6b" if value < 1.0 else "#7ed491" mark = "✗" if value < 1.0 else "✓" + step_html = "" + if value < 1.0 and metric in steps: + source = ' (shadow env flag)' if metric == COLLISION_METRIC else "" + step_html = f" · {step_link(steps[metric])}{source}" lines.append( f'
{mark} {label}' - f'{value:.2f} ({role})
' + f'{value:.2f} ({role}){step_html}' ) for detail in comfort: lines.append(f'
↳ {detail}
') @@ -164,19 +228,26 @@ def main(): raise SystemExit(f"no obs_html/*.replay.zlib under {args.group_dir}") rows = scenario_rows(args.group_dir) comfort = comfort_failures(args.group_dir) + steps = violation_steps(args.group_dir) wanted = selected_tokens(rows, args.max_score, args.metric) | set(args.tokens) jobs = [ (token, replays[token], obs_dir / page_name(token, rows.get(token))) for token in sorted(wanted) if token in replays ] + + def render(job): + token, replay, html_path = job + row = rows.get(token) + token_steps = dict(steps.get(token, {})) + collision_failed = row is not None and float(row[COLLISION_METRIC]) < 1.0 + collision_step = replay_collision_step(replay) if collision_failed else None + if collision_step is not None: + token_steps[COLLISION_METRIC] = collision_step + render_page(replay, html_path, score_panel(token, row, comfort.get(token, []), token_steps)) + with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool: - list( - pool.map( - lambda job: render_page(job[1], job[2], score_panel(job[0], rows.get(job[0]), comfort.get(job[0], []))), - jobs, - ) - ) + list(pool.map(render, jobs)) file_metrics = {html.name: gallery_metrics(rows.get(token)) for token, _, html in jobs} pruned = 0 if args.prune: diff --git a/scripts/eval/tidy_alpasim_run.py b/scripts/eval/tidy_alpasim_run.py new file mode 100644 index 0000000000..77884751c0 --- /dev/null +++ b/scripts/eval/tidy_alpasim_run.py @@ -0,0 +1,154 @@ +"""Tidy the output dir of scripts/kesai/12_alpasim.sh down to what is read after the run: + + results.txt official score + per-metric means (written by 12_alpasim.sh) + obs_replays/ teacher observation replays (render_alpasim_obs_html.py turns them into the gallery) + logs/ evaluate.log, evaluation.log, evaluation.yaml, config.json, run_info.txt, txt-logs/ (enroot and pyxis + runs: one log per service container) + alpasim/ results-summary.json, metrics_unprocessed.parquet (every rollout's per-step metrics), eval-config.yaml, + wizard-config.yaml + videos/ the in-run mp4s, only with RENDER_VIDEO=1 + +usage: python scripts/eval/tidy_alpasim_run.py [--keep_rollout_logs none] [--keep_container_info] + +Everything else AlpaSim writes is deleted: per-rollout metrics.parquet (concatenated in metrics_unprocessed.parquet) and +_complete markers, controller csvs, runtime worker logs (copies of evaluate.log's runtime lines), generated configs, +docker-compose.yaml, prometheus configs, aggregate plots. Kept rollout.asl stay at alpasim/wizard_0/rollouts///, +where AlpaSim's and alpasim-training's offline tools look. evaluate.log loses its container INFO lines unless +--keep_container_info: AlpaSim's runtime workers and the renderer log every sim event at INFO whatever wizard.log_level says. +A run without results-summary.json (never aggregated) keeps its AlpaSim tree. Runs in the 2026-10-12 layout +(alpasim/wizard_0/ and logs/ already in place) are tidied the same way. +""" + +import argparse +import json +import re +from pathlib import Path + +KEEP_RULES = { + "all": lambda row: True, + "imperfect": lambda row: row["score"] < 1, + "failed": lambda row: row["status"] == "fail", + "none": lambda row: False, +} +RUN_LOGS = ("evaluate.log", "run_info.txt") +ALPAGYM_LOGS = ("config.json", "evaluation.yaml", "logs/evaluation.log") +WIZARD_KEPT = ( + "aggregate/results-summary.json", + "aggregate/metrics_unprocessed.parquet", + "eval-config.yaml", + "wizard-config.yaml", +) +CONTAINER_INFO_LINE = re.compile(rb"^[a-z0-9_-]+-[0-9]+-[0-9]+ +\| .*( INFO:|\[INFO\])") # docker compose prefix +SERVICE_INFO_LINE = re.compile(rb"^.*( INFO:|\[INFO\])") # txt-logs/out-*-log.txt of enroot and pyxis runs + + +def move(src, dst_dir): + if not src.exists(): + return + dst = dst_dir / src.name + if dst.exists(): + raise SystemExit(f"{dst} already exists, not overwriting it with {src}") + dst_dir.mkdir(exist_ok=True) + src.rename(dst) + + +def drop_info_lines(log_path, info_line): + """-> bytes removed from log_path by dropping the lines matching info_line, AlpaSim's container INFO lines.""" + size_before = log_path.stat().st_size + filtered_path = log_path.with_suffix(".filtered") + with open(log_path, "rb") as source, open(filtered_path, "wb") as target: + for line in source: + if info_line.match(line) is None: + target.write(line) + filtered_path.replace(log_path) + return size_before - log_path.stat().st_size + + +def remove_empty_dirs(root): + for directory in sorted((p for p in root.rglob("*") if p.is_dir()), key=lambda p: len(p.parts), reverse=True): + if not any(directory.iterdir()): + directory.rmdir() + if root.is_dir() and not any(root.iterdir()): + root.rmdir() + + +def sweep_wizard(wizard_dir, rows, keep_rule): + """Delete every file of the AlpaSim tree except the rollout.asl of the summary rows keep_rule keeps. + -> (files removed, bytes removed, rollout logs kept).""" + kept_logs = { + wizard_dir / "rollouts" / row["clipgt_id"] / row["rollout_id"] / "rollout.asl" for row in rows if keep_rule(row) + } + removed_count = removed_bytes = 0 + for path in sorted(wizard_dir.rglob("*")): + if not path.is_file() or path in kept_logs: + continue + removed_bytes += path.stat().st_size + path.unlink() + removed_count += 1 + remove_empty_dirs(wizard_dir) + return removed_count, removed_bytes, sum(path.is_file() for path in kept_logs) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("run_dir", help="output dir of 12_alpasim.sh") + ap.add_argument("--keep_rollout_logs", choices=tuple(KEEP_RULES), default="none") + ap.add_argument( + "--keep_container_info", action="store_true", help="keep the containers' INFO lines in evaluate.log" + ) + args = ap.parse_args() + run_dir = Path(args.run_dir) + alpagym_dir, alpasim_dir, logs_dir = run_dir / "alpagym", run_dir / "alpasim", run_dir / "logs" + + move(alpagym_dir / "obs_replays", run_dir) + for name in RUN_LOGS: + move(run_dir / name, logs_dir) + for name in ALPAGYM_LOGS: + move(alpagym_dir / name, logs_dir) + if (alpagym_dir / "alpasim" / "wizard_0").is_dir(): + alpasim_dir.mkdir(exist_ok=True) + move(alpagym_dir / "alpasim" / "wizard_0", alpasim_dir) + remove_empty_dirs(alpagym_dir) + if alpagym_dir.exists(): + print( + f"[tidy_alpasim_run] left in place, unknown: {sorted(str(p) for p in alpagym_dir.rglob('*') if p.is_file())}" + ) + + removed_count = removed_bytes = 0 + for stale in sorted(logs_dir.glob("runtime_worker_*.log")) + sorted( + logs_dir.glob("controller/alpasim_controller_*.csv") + ): + removed_bytes += stale.stat().st_size + stale.unlink() + removed_count += 1 + remove_empty_dirs(logs_dir / "controller") + evaluate_log = logs_dir / "evaluate.log" + if evaluate_log.is_file() and not args.keep_container_info: + removed_bytes += drop_info_lines(evaluate_log, CONTAINER_INFO_LINE) + + wizard_dir = alpasim_dir / "wizard_0" + summary_path = wizard_dir / "aggregate" / "results-summary.json" + if not summary_path.is_file(): + tidied = (alpasim_dir / "results-summary.json").is_file() + print(f"[tidy_alpasim_run] {'already tidied' if tidied else f'no {summary_path}: AlpaSim tree kept as is'}") + return + rows = json.loads(summary_path.read_text())["rollouts"] + for mp4 in sorted(wizard_dir.glob("rollouts/*/*/*.mp4")): + move(mp4, run_dir / "videos") + for name in WIZARD_KEPT: + move(wizard_dir / name, alpasim_dir) + move(wizard_dir / "txt-logs", logs_dir) + if not args.keep_container_info: + for service_log in sorted(logs_dir.glob("txt-logs/out-*-log.txt")): + removed_bytes += drop_info_lines(service_log, SERVICE_INFO_LINE) + wizard_count, wizard_bytes, kept_logs = sweep_wizard(wizard_dir, rows, KEEP_RULES[args.keep_rollout_logs]) + removed_count += wizard_count + removed_bytes += wizard_bytes + print( + f"[tidy_alpasim_run] removed {removed_count} files ({removed_bytes / 1e9:.2f} GB), kept {kept_logs} rollout logs " + f"({args.keep_rollout_logs}) -> {run_dir}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/kesai/0_build_puffer.sh b/scripts/kesai/0_build_puffer.sh index 4f68a5d85b..8c37ced632 100644 --- a/scripts/kesai/0_build_puffer.sh +++ b/scripts/kesai/0_build_puffer.sh @@ -6,10 +6,12 @@ #SBATCH --gres=gpu:1 #SBATCH --mem=64G #SBATCH --cpus-per-task=24 -#SBATCH --output=/home/bjaeger/logs/eval_server_%a_%A.out -#SBATCH --error=/home/bjaeger/logs/eval_server_%a_%A.err +#SBATCH --output=experiments/logs/build_puffer_%j.out +#SBATCH --error=experiments/logs/build_puffer_%j.err #SBATCH --partition=dev +source scripts/kesai/cluster_env.sh || exit 1 + # print info about current job echo "START TIME: $(date)" start=`date +%s` diff --git a/scripts/kesai/10_diag_reset_loop_8GPU_slurm.sh b/scripts/kesai/10_diag_reset_loop_8GPU_slurm.sh index 38f4f575a9..a0544638d9 100755 --- a/scripts/kesai/10_diag_reset_loop_8GPU_slurm.sh +++ b/scripts/kesai/10_diag_reset_loop_8GPU_slurm.sh @@ -4,13 +4,15 @@ #SBATCH --nodes 1 #SBATCH --time 0-02:00 #SBATCH --gres gpu:8 -#SBATCH --mem=1007G -#SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.err +#SBATCH --mem=0 +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --output experiments/logs/diag_reset_%a_%A.out +#SBATCH --error experiments/logs/diag_reset_%a_%A.err #SBATCH --partition dev #SBATCH --array=0-1 +source scripts/kesai/cluster_env.sh || exit 1 + # Short 1-node run to catch sub-envs that early-reset every tick right after # creation (see environment/early_reset_short, environment/spawn_* and the # [DRIVE DIAG] lines in the .err). Task 0 = training spawn noise, task 1 = noise off. @@ -26,7 +28,7 @@ else SPAWN_NOISE_ARGS="env.spawn_lateral_offset_max_frac=0.0 env.spawn_heading_max_deg=0.0" fi echo ${RUN_NAME} -export DATA_DIR=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} +export DATA_DIR=$PD/experiments/${RUN_NAME} echo ${DATA_DIR} export NUMEXPR_NUM_THREADS=1 @@ -38,9 +40,9 @@ export OMP_NUM_THREADS=1 source .venv/bin/activate # Only task 0 builds; concurrent in-place builds race on shared NFS build files. -BUILD_STATUS_FILE=/home/bjaeger/PufferDrive/experiments/logs/build_status_${SLURM_ARRAY_JOB_ID} +BUILD_STATUS_FILE=$PD/experiments/logs/build_status_${SLURM_ARRAY_JOB_ID} if [ "${SLURM_ARRAY_TASK_ID}" -eq 0 ]; then - bash scripts/kesai/build_ext_if_changed.sh /home/bjaeger/PufferDrive + bash scripts/kesai/build_ext_if_changed.sh $PD BUILD_STATUS=$? echo ${BUILD_STATUS} > ${BUILD_STATUS_FILE} else @@ -68,7 +70,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met wandb_project=nightly-multi-long \ wandb_group=diag_reset_loop \ train.data_dir=${DATA_DIR} \ - env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla_128_affine_zones \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla_128_affine_zones \ env.num_maps=128 \ train.name=${RUN_NAME} \ run_name=${RUN_NAME} \ @@ -94,7 +96,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met tb=True echo "Training exited with status $?" -echo "Summarize with: python scripts/kesai/summarize_reset_diag.py /home/bjaeger/PufferDrive/experiments/logs/log_${SLURM_ARRAY_TASK_ID}_${SLURM_ARRAY_JOB_ID}.err" +echo "Summarize with: python scripts/kesai/summarize_reset_diag.py $PD/experiments/logs/diag_reset_${SLURM_ARRAY_TASK_ID}_${SLURM_ARRAY_JOB_ID}.err" end=$(date +%s) echo "END TIME: $(date)" echo "Runtime: $((end-start))" diff --git a/scripts/kesai/11_carla_longest6.sh b/scripts/kesai/11_carla_longest6.sh deleted file mode 100644 index 3f589db93e..0000000000 --- a/scripts/kesai/11_carla_longest6.sh +++ /dev/null @@ -1,216 +0,0 @@ -#!/bin/bash -#SBATCH --job-name eval_longest6 -#SBATCH --ntasks 1 -#SBATCH --nodes 1 -#SBATCH --time 1-00:00 -#SBATCH --gres gpu:8 -#SBATCH --mem=1007G -#SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/log_%j.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/log_%j.err -#SBATCH --partition dev -# CARLA longest6 evaluation of a PufferDrive checkpoint through CaRL's unmodified original_leaderboard -# (pufferlib/ocean/cosim/carla/leaderboard_agent.py), the CARLA counterpart of 9_nuPlan.sh. -# -# One CARLA server + one evaluator per GPU (1-8: whatever the allocation holds, e.g. `sbatch --gres gpu:2`); -# the 36 routes are dealt round-robin over the GPUs, every route gets its own server (restarted per route, -# CaRL's own recipe), crashed routes are retried, and the per-route result jsons are aggregated at the end -# with CaRL's tools/result_parser.py (results.csv) plus the nuPlan-style HTML report ($OUT/report/index.html) -# and the obs replay gallery of all routes ($OUT/obs_html/index.html, as in 9_nuPlan.sh). -# -# Overridable env: RUN_DIR, ROUTES, SCENARIOS (0 = pufferlib's scenario-free longest6), REPETITIONS, -# ROUTE_SUBSET ("0 5 17": only these route ids), NUM_GPUS, PORT_BASE (CARLA rpc port of gpu 0; gpu w uses -# PORT_BASE+50w, TM PORT_BASE+6000+50w), LOGGING (1, default: telemetry, world log and the HTML report per -# route; 0: scores only), CARLA_VIEW (0, default: CARLA runs without rendering (-nullrhi; collision sensors -# and telemetry work, an RGB camera segfaults the server); 1: chase-cam mp4 per route on a rendering server, -# needs LOGGING=1), OBS_HTML (1 = also the interactive obs replay per route, rendered into the one gallery -# folder $OUT/obs_html at the end, large; needs LOGGING=1), REPORT (0 = skip the HTML report), -# MAX_ATTEMPTS, CARLA_ROOT, CARL_WORK_DIR, PY, PD, COSIM_MAX_SPEED_MPS (ego speed cap, default 30), -# COSIM_ZERO_PARTNER_STOPPED_TIME (default 1: partners' stopped-time obs held at 0; 0 = real stopped times), -# COSIM_PEDESTRIAN_MIN_SIZE_M (default 0: true CARLA walker boxes; 0.8 = the training spawn floor). -set -u - -export PD=${PD:-/home/bjaeger/PufferDrive} -export PY=${PY:-/home/bjaeger/miniconda3/envs/carl/bin/python} # cp310: CARLA 0.9.15 ships no newer wheel -export CARLA_ROOT=${CARLA_ROOT:-/home/bjaeger/CARLA_0.9.15} -export CARL_WORK_DIR=${CARL_WORK_DIR:-/home/bjaeger/CaRL/CARLA} -RUN_DIR=${RUN_DIR:-/home/bjaeger/PufferDrive/experiments/k_scaled_0041_1000} -# the agent finds config.yaml next to final_model.pt (or one level above a models/*.pt) -export CKPT=$RUN_DIR/final_model.pt -SCENARIOS=${SCENARIOS:-1} -if [ "$SCENARIOS" = "1" ]; then - ROUTES=${ROUTES:-$CARL_WORK_DIR/custom_leaderboard/leaderboard/data/longest6.xml} -else - ROUTES=${ROUTES:-$PD/pufferlib/ocean/cosim/carla/routes/longest6_no_scenarios.xml} -fi -REPETITIONS=${REPETITIONS:-1} -ROUTE_SUBSET=${ROUTE_SUBSET:-} -PORT_BASE=${PORT_BASE:-2000} -MAX_ATTEMPTS=${MAX_ATTEMPTS:-3} -LOGGING=${LOGGING:-1} -CARLA_VIEW=${CARLA_VIEW:-0} -[ "$CARLA_VIEW" = "1" ] && [ "$LOGGING" != "1" ] && { echo "CARLA_VIEW=1 needs LOGGING=1"; exit 1; } -OBS_HTML=${OBS_HTML:-1} -REPORT=${REPORT:-$LOGGING} -EVALUATOR=$CARL_WORK_DIR/original_leaderboard/leaderboard/leaderboard/leaderboard_evaluator.py -for path in "$CKPT" "$PY" "$CARLA_ROOT/CarlaUE4.sh" "$EVALUATOR" "$ROUTES" "$CARL_WORK_DIR/tools/result_parser.py"; do - [ -e "$path" ] || { echo "missing $path"; exit 1; } -done -for path in "$PY" "$CARLA_ROOT/CarlaUE4.sh"; do - [ -x "$path" ] || { echo "not executable: $path (chmod +x, or the filesystem is mounted noexec)"; exit 1; } -done -echo "Evaluating checkpoint: $CKPT" - -# GPUs of this allocation; one CARLA server + one evaluator per GPU -NUM_GPUS=${NUM_GPUS:-$(nvidia-smi -L 2>/dev/null | grep -c "^GPU")} -if [ "$NUM_GPUS" -lt 1 ] || [ "$NUM_GPUS" -gt 8 ]; then echo "NUM_GPUS=$NUM_GPUS must be 1..8"; exit 1; fi - -# PD first, and drop inherited entries from other PufferDrive checkouts (see 9_nuPlan.sh) -INHERITED_PYTHONPATH=$(echo "${PYTHONPATH:-}" | tr ':' '\n' | grep -v "cosim_Puffer\|/PufferDrive" | paste -sd: -) -export PYTHONPATH="$PD:$CARL_WORK_DIR/original_leaderboard/leaderboard:$CARL_WORK_DIR/original_leaderboard/scenario_runner:$CARLA_ROOT/PythonAPI/carla${INHERITED_PYTHONPATH:+:$INHERITED_PYTHONPATH}" -export SCENARIO_RUNNER_ROOT=$CARL_WORK_DIR/original_leaderboard/scenario_runner -cd "$PD" || exit 1 -export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 NUMEXPR_NUM_THREADS=4 -export COSIM_DEVICE=${COSIM_DEVICE:-cpu} -export COSIM_DYNAMICS_SOURCE=pufferdrive -export COSIM_MAX_SPEED_MPS=${COSIM_MAX_SPEED_MPS:-30} -export COSIM_ZERO_PARTNER_STOPPED_TIME=${COSIM_ZERO_PARTNER_STOPPED_TIME:-1} -export COSIM_PEDESTRIAN_MIN_SIZE_M=${COSIM_PEDESTRIAN_MIN_SIZE_M:-0} -export COSIM_OBS_HTML_MAX_STEPS=${COSIM_OBS_HTML_MAX_STEPS:-20000} -export COSIM_OBS_HTML_RENDER=0 # routes save the compact replay only; render_carla_obs_html.py renders all pages into $OUT/obs_html - -TAG=longest6$([ "$SCENARIOS" = "1" ] || echo "_noscen")$([ "$REPETITIONS" = "1" ] || echo "_rep$REPETITIONS") -# results live in the model's own eval folder, next to the PufferDrive benchmark evals -OUT=$RUN_DIR/eval/carla_${TAG}_$(date +%Y%m%d_%H%M%S)_${SLURM_JOB_ID:-local} -mkdir -p "$OUT/routes" "$OUT/aggregate" -{ - echo "checkpoint $CKPT"; echo "routes $ROUTES"; echo "repetitions $REPETITIONS"; echo "gpus $NUM_GPUS" - echo "git $(git -C "$PD" rev-parse --short HEAD 2>/dev/null) $(git -C "$PD" diff --quiet 2>/dev/null || echo dirty)" -} > "$OUT/run_info.txt" -echo "Results -> $OUT" - -mkdir -p "$PD/experiments/logs" # the build helper keeps its source hash there -PATH="$(dirname "$PY"):$PATH" bash "$PD/scripts/kesai/build_ext_if_changed.sh" "$PD" || exit 1 -IMPORTED_AGENT=$("$PY" -c "import pufferlib.ocean.cosim.carla.leaderboard_agent as a; print(a.__file__)") || { echo "ERROR: agent import failed (see traceback above)"; exit 1; } -case "$IMPORTED_AGENT" in - "$PD"/*) echo "agent: $IMPORTED_AGENT" ;; - *) echo "ERROR: agent imported from '$IMPORTED_AGENT', expected $PD"; exit 1 ;; -esac - -# route ids of the xml, dealt round-robin: gpu w runs ids w, w+N, w+2N, ... -mapfile -t ROUTE_IDS < <("$PY" -c " -import sys, xml.etree.ElementTree as ET -print('\n'.join(r.attrib['id'] for r in ET.parse(sys.argv[1]).getroot().iter('route')))" "$ROUTES") -[ -n "$ROUTE_SUBSET" ] && read -r -a ROUTE_IDS <<< "$ROUTE_SUBSET" -echo "${#ROUTE_IDS[@]} routes over $NUM_GPUS GPUs" - -route_done() { # 0 when the route json holds a record that is not one of the crash states CaRL resubmits - "$PY" - "$1" <<'EOF' -import json, sys -try: - records = json.load(open(sys.argv[1]))["_checkpoint"]["records"] -except Exception: - sys.exit(1) -crashed = ("Failed - Agent couldn't be set up", "Failed", "Failed - Simulation crashed", "Failed - Agent crashed") -sys.exit(0 if records and all(r["status"] not in crashed for r in records) else 1) -EOF -} - -run_route() { # $1 gpu, $2 route id, $3 attempt: one CARLA server + one evaluator run, into $OUT/routes/route_ - local gpu=$1 route_id=$2 attempt=$3 - local route_dir=$OUT/routes/route_$(printf "%02d" "$route_id") - local port=$((PORT_BASE + gpu * 50)) tm_port=$((PORT_BASE + 6000 + gpu * 50)) - local server_log=$route_dir/carla_server_attempt$attempt.log evaluator_log=$route_dir/evaluator_attempt$attempt.log - mkdir -p "$route_dir" - rm -f "$route_dir/result.json" - # only the chase camera needs the renderer; collision sensors and telemetry work under -nullrhi - local render_args=(-RenderOffScreen -nullrhi) - [ "$CARLA_VIEW" = "1" ] && render_args=(-RenderOffScreen -graphicsadapter="$gpu") - "$CARLA_ROOT/CarlaUE4.sh" "${render_args[@]}" -nosound \ - -carla-rpc-port="$port" -carla-streaming-port=$((port + 1)) > "$server_log" 2>&1 & - local server_pid=$! up=1 - echo "$server_pid" >> "$OUT/server_pids" - for _ in $(seq 1 60); do - kill -0 "$server_pid" 2>/dev/null || { echo "[gpu$gpu route$route_id] CARLA server died"; break; } - "$PY" -c " -import carla, sys -try: - c = carla.Client('localhost', $port); c.set_timeout(10.0); c.get_world().get_map() -except Exception: - sys.exit(1)" 2>/dev/null && up=0 && break - sleep 5 - done - if [ $up -ne 0 ]; then - pkill -9 -P "$server_pid" 2>/dev/null; kill -9 "$server_pid" 2>/dev/null - echo "[gpu$gpu route$route_id] last lines of $server_log:"; tail -n 20 "$server_log" - return 1 - fi - local log_env=() - if [ "$LOGGING" = "1" ]; then - log_env=(COSIM_TELEMETRY="$route_dir/telemetry" COSIM_WORLD_LOG="$route_dir/world_log") - [ "$CARLA_VIEW" = "1" ] && log_env+=(COSIM_DEBUG_CARLA_VIEW="$route_dir/carla_view") - [ "$OBS_HTML" = "1" ] && log_env+=(COSIM_OBS_HTML="$route_dir/obs_html") - fi - env CUDA_VISIBLE_DEVICES="$gpu" "${log_env[@]}" \ - "$PY" -u "$EVALUATOR" --routes "$ROUTES" --routes-subset "$route_id" --repetitions "$REPETITIONS" \ - --agent "$PD/pufferlib/ocean/cosim/carla/leaderboard_agent.py" --agent-config "$CKPT" \ - --checkpoint "$route_dir/result.json" --track MAP --port "$port" --traffic-manager-port "$tm_port" \ - > "$evaluator_log" 2>&1 - # children first: killing the wrapper shell first reparents the UE4 binary and leaves it running - pkill -9 -P "$server_pid" 2>/dev/null; kill -9 "$server_pid" 2>/dev/null; wait "$server_pid" 2>/dev/null - route_done "$route_dir/result.json" && return 0 - echo "[gpu$gpu route$route_id] no valid route record; last lines of $evaluator_log:"; tail -n 20 "$evaluator_log" - return 1 -} - -run_worker() { # $1 gpu: its share of the routes, sequentially, each retried up to MAX_ATTEMPTS - local gpu=$1 - for ((k = gpu; k < ${#ROUTE_IDS[@]}; k += NUM_GPUS)); do - local route_id=${ROUTE_IDS[k]} attempt - for ((attempt = 1; attempt <= MAX_ATTEMPTS; attempt++)); do - echo "[gpu$gpu] route $route_id attempt $attempt $(date +%H:%M:%S)" - run_route "$gpu" "$route_id" "$attempt" && break - echo "[gpu$gpu] route $route_id attempt $attempt failed" - done - done -} - -# only the servers this job started (children first: the UE4 binary outlives its wrapper shell otherwise) -trap 'for p in $(cat "$OUT/server_pids" 2>/dev/null); do pkill -9 -P "$p" 2>/dev/null; kill -9 "$p" 2>/dev/null; done' EXIT -WORKER_PIDS=() -for ((gpu = 0; gpu < NUM_GPUS; gpu++)); do - run_worker "$gpu" > "$OUT/worker_gpu$gpu.log" 2>&1 & - WORKER_PIDS+=($!) -done -wait "${WORKER_PIDS[@]}" -cat "$OUT"/worker_gpu*.log - -# aggregate: CaRL's parser over the per-route jsons -> $OUT/aggregate/results.csv, plus a per-route table -for f in "$OUT"/routes/route_*/result.json; do cp "$f" "$OUT/aggregate/$(basename "$(dirname "$f")").json"; done -"$PY" "$CARL_WORK_DIR/tools/result_parser.py" --xml "$ROUTES" --results "$OUT/aggregate" -"$PY" - "$OUT/aggregate" "${#ROUTE_IDS[@]}" <<'EOF' -import glob, json, sys -rows = [] -for f in sorted(glob.glob(sys.argv[1] + "/route_*.json")): - for r in json.load(open(f))["_checkpoint"]["records"]: - rows.append((r["route_id"], r["status"], r["scores"]["score_composed"], r["scores"]["score_route"], r["scores"]["score_penalty"])) -print(f"\n{len(rows)} route records of {sys.argv[2]} routes") -for route_id, status, ds, rc, ip in rows: - print(f" {route_id:22s} {status:34s} DS {ds:6.2f} RC {rc:6.2f} IP {ip:5.3f}") -if rows: - n = len(rows) - print(f"mean DS {sum(r[2] for r in rows)/n:.2f} RC {sum(r[3] for r in rows)/n:.2f} IP {sum(r[4] for r in rows)/n:.3f}") -EOF - -REPORT_ARGS=() -[ "$OBS_HTML" = "1" ] && REPORT_ARGS+=(--obs-html-dir "$OUT/obs_html") -if [ "$REPORT" = "1" ] && [ "$LOGGING" = "1" ]; then - "$PY" "$PD/scripts/eval/analyze_carla_cosim.py" "$OUT"/routes/route_* "$OUT/report" "${REPORT_ARGS[@]}" || echo "HTML report failed (results above are unaffected)" -fi -# after the report: the gallery links to it when it exists -if [ "$OBS_HTML" = "1" ] && [ "$LOGGING" = "1" ]; then - "$PY" "$PD/scripts/eval/render_carla_obs_html.py" "$OUT" --routes "$ROUTES" || echo "obs replay gallery failed (results above are unaffected)" -fi -[ -f "$OUT/report/index.html" ] && echo "report -> $OUT/report/index.html" -[ -f "$OUT/obs_html/index.html" ] && echo "obs replays -> $OUT/obs_html/index.html" -echo "Done -> $OUT" diff --git a/scripts/kesai/11_carla_longest6_v2.sh b/scripts/kesai/11_carla_longest6_v2.sh new file mode 100644 index 0000000000..b0d9865e2f --- /dev/null +++ b/scripts/kesai/11_carla_longest6_v2.sh @@ -0,0 +1,320 @@ +#!/bin/bash +#SBATCH --job-name eval_longest6_v2 +#SBATCH --ntasks 1 +#SBATCH --nodes 1 +#SBATCH --time 1-00:00 +#SBATCH --gres gpu:8 +#SBATCH --mem=0 +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --output experiments/logs/carla_longest6_v2_%j.out +#SBATCH --error experiments/logs/carla_longest6_v2_%j.err +#SBATCH --partition dev +# CARLA longest6 evaluation of a PufferDrive checkpoint through CaRL's unmodified original_leaderboard +# (pufferlib/ocean/cosim/carla/leaderboard_agent.py), the CARLA counterpart of 9_nuPlan.sh. +# +# WORKERS parallel workers (default 16), each one CARLA server + one evaluator at a time, spread round-robin over the +# GPUs of the allocation (1-8, e.g. `sbatch --gres gpu:2`): under -nullrhi a server needs ~2 cores and no GPU, so +# several share one. The 36 routes x SEEDS passes (route runs, $OUT/routes/route__s/) go longest-first to the +# least loaded worker (route length from the xml), every route run gets its own server (restarted per route, CaRL's own +# recipe), worker launches are staggered, a crashed route run is retried on the next GPU, and the result jsons are +# aggregated at the end with CaRL's tools/result_parser.py over $OUT/aggregate/s/ (results.csv) plus the nuPlan-style HTML report ($OUT/report/index.html) and the obs replay +# gallery of all routes ($OUT/obs_html/index.html, as in 9_nuPlan.sh). Speed: wall and game time of every route go +# to $OUT/timing.csv, the run totals to run_info.txt and stdout. +# +# Overridable env: RUN_DIR, ROUTES, SCENARIOS (0 = pufferlib's scenario-free longest6), SEEDS ("0 1 2" by default: +# traffic-manager seeds, one full pass over the routes per seed, the protocol of CaRL's and carla_garage's +# evaluate_routes_slurm.py with --num_repetitions 3; results.csv then carries mean and std over the seeds per route and +# per town, the stdout summary the DS of every seed and mean +- std over the seeds; "0" = one pass with the evaluator's +# default seed), REPETITIONS (leaderboard-internal repetitions within a pass, 1), ROUTE_SUBSET ("0 5 17": only these +# route ids), WORKERS, NUM_GPUS, STAGGER_S (15: seconds between worker launches), +# PORT_BASE (CARLA rpc port of worker 0; worker w uses PORT_BASE+50w, TM PORT_BASE+6000+50w), OMP_NUM_THREADS (policy +# threads; 2 from 16 workers on, else 4), LOGGING (1, default: telemetry, world log and the HTML report per +# route; 0: scores only), CARLA_VIEW (0, default: CARLA runs without rendering (-nullrhi; collision sensors +# and telemetry work, an RGB camera segfaults the server); 1: chase-cam mp4 per route on a rendering server, +# needs LOGGING=1), OBS_HTML (1 = also the interactive obs replay per route, rendered into the one gallery +# folder $OUT/obs_html at the end, large; needs LOGGING=1), REPORT (0 = skip the HTML report), +# MAX_ATTEMPTS, CARLA_ROOT, CARL_WORK_DIR, PY, PD, COSIM_MAX_SPEED_MPS (ego speed cap, default 30), +# COSIM_ZERO_PARTNER_STOPPED_TIME (default 1: partners' stopped-time obs held at 0; 0 = real stopped times), +# COSIM_VRU_SIZE_PAD_M (default 0.8: added to the length and width of every walker and two-wheeler box, so they +# read strictly larger than the smallest training vehicle; 0 = true CARLA boxes). +set -u +source scripts/kesai/cluster_env.sh || exit 1 + +export PY=${PY:-$HOME/miniconda3/envs/carl/bin/python} # cp310: CARLA 0.9.15 ships no newer wheel +export CARLA_ROOT=${CARLA_ROOT:-$HOME/CARLA_0.9.15} +export CARL_WORK_DIR=${CARL_WORK_DIR:-$HOME/CaRL/CARLA} +RUN_DIR=${RUN_DIR:-$PD/experiments/k_scaled_0045_1000} +# the agent finds config.yaml next to final_model.pt (or one level above a models/*.pt) +export CKPT=$RUN_DIR/final_model.pt +SCENARIOS=${SCENARIOS:-1} +if [ "$SCENARIOS" = "1" ]; then + ROUTES=${ROUTES:-$CARL_WORK_DIR/custom_leaderboard/leaderboard/data/longest6.xml} +else + ROUTES=${ROUTES:-$PD/pufferlib/ocean/cosim/carla/routes/longest6_no_scenarios.xml} +fi +REPETITIONS=${REPETITIONS:-1} +SEEDS=${SEEDS:-0 1 2} # traffic-manager seeds, one pass over the routes each (the CaRL / carla_garage protocol) +for seed in $SEEDS; do [[ "$seed" =~ ^[0-9]+$ ]] || { echo "SEEDS=\"$SEEDS\" must be integers"; exit 1; }; done +read -r -a SEED_LIST <<< "$SEEDS" +[ "${#SEED_LIST[@]}" -gt 0 ] || { echo "SEEDS must name at least one seed"; exit 1; } +ROUTE_SUBSET=${ROUTE_SUBSET:-} +PORT_BASE=${PORT_BASE:-2000} +MAX_ATTEMPTS=${MAX_ATTEMPTS:-3} +WORKERS=${WORKERS:-16} +STAGGER_S=${STAGGER_S:-15} +LOGGING=${LOGGING:-1} +CARLA_VIEW=${CARLA_VIEW:-0} +[ "$CARLA_VIEW" = "1" ] && [ "$LOGGING" != "1" ] && { echo "CARLA_VIEW=1 needs LOGGING=1"; exit 1; } +OBS_HTML=${OBS_HTML:-1} +REPORT=${REPORT:-$LOGGING} +EVALUATOR=$CARL_WORK_DIR/original_leaderboard/leaderboard/leaderboard/leaderboard_evaluator.py +for path in "$CKPT" "$PY" "$CARLA_ROOT/CarlaUE4.sh" "$EVALUATOR" "$ROUTES" "$CARL_WORK_DIR/tools/result_parser.py"; do + [ -e "$path" ] || { echo "missing $path"; exit 1; } +done +for path in "$PY" "$CARLA_ROOT/CarlaUE4.sh"; do + [ -x "$path" ] || { echo "not executable: $path (chmod +x, or the filesystem is mounted noexec)"; exit 1; } +done +echo "Evaluating checkpoint: $CKPT" + +# GPUs of this allocation; the workers share them round-robin +NUM_GPUS=${NUM_GPUS:-$(nvidia-smi -L 2>/dev/null | grep -c "^GPU")} +if [ "$NUM_GPUS" -lt 1 ] || [ "$NUM_GPUS" -gt 8 ]; then echo "NUM_GPUS=$NUM_GPUS must be 1..8"; exit 1; fi +if [ "$WORKERS" -lt 1 ] || [ "$WORKERS" -gt 64 ]; then echo "WORKERS=$WORKERS must be 1..64"; exit 1; fi + +# PD first, and drop inherited entries from other PufferDrive checkouts (see 9_nuPlan.sh) +INHERITED_PYTHONPATH=$(echo "${PYTHONPATH:-}" | tr ':' '\n' | grep -v "cosim_Puffer\|/PufferDrive" | paste -sd: -) +export PYTHONPATH="$PD:$CARL_WORK_DIR/original_leaderboard/leaderboard:$CARL_WORK_DIR/original_leaderboard/scenario_runner:$CARLA_ROOT/PythonAPI/carla${INHERITED_PYTHONPATH:+:$INHERITED_PYTHONPATH}" +export SCENARIO_RUNNER_ROOT=$CARL_WORK_DIR/original_leaderboard/scenario_runner +cd "$PD" || exit 1 +OMP_THREADS=${OMP_NUM_THREADS:-$([ "$WORKERS" -ge 16 ] && echo 2 || echo 4)} +export OMP_NUM_THREADS=$OMP_THREADS MKL_NUM_THREADS=$OMP_THREADS OPENBLAS_NUM_THREADS=$OMP_THREADS NUMEXPR_NUM_THREADS=$OMP_THREADS +export COSIM_DEVICE=${COSIM_DEVICE:-cpu} +export COSIM_DYNAMICS_SOURCE=pufferdrive +export COSIM_MAX_SPEED_MPS=${COSIM_MAX_SPEED_MPS:-30} +export COSIM_ZERO_PARTNER_STOPPED_TIME=${COSIM_ZERO_PARTNER_STOPPED_TIME:-1} +export COSIM_VRU_SIZE_PAD_M=${COSIM_VRU_SIZE_PAD_M:-0.8} +export COSIM_OBS_HTML_MAX_STEPS=${COSIM_OBS_HTML_MAX_STEPS:-20000} +export COSIM_OBS_HTML_RENDER=0 # routes save the compact replay only; render_carla_obs_html.py renders all pages into $OUT/obs_html + +TAG=longest6_v2$([ "$SCENARIOS" = "1" ] || echo "_noscen")$([ "$REPETITIONS" = "1" ] || echo "_rep$REPETITIONS")_seeds$(echo "$SEEDS" | tr ' ' '-') +TAG=$TAG$([ "$COSIM_VRU_SIZE_PAD_M" = "0" ] || echo "_vrupad$COSIM_VRU_SIZE_PAD_M") +# results live in the model's own eval folder, next to the PufferDrive benchmark evals +OUT=$RUN_DIR/eval/carla_${TAG}_$(date +%Y%m%d_%H%M%S)_${SLURM_JOB_ID:-local} +mkdir -p "$OUT/routes" "$OUT/aggregate" +{ + echo "checkpoint $CKPT"; echo "routes $ROUTES"; echo "seeds $SEEDS"; echo "repetitions $REPETITIONS"; echo "gpus $NUM_GPUS" + echo "vru_size_pad_m $COSIM_VRU_SIZE_PAD_M" + echo "git $(git -C "$PD" rev-parse --short HEAD 2>/dev/null) $(git -C "$PD" diff --quiet 2>/dev/null || echo dirty)" +} > "$OUT/run_info.txt" +echo "Results -> $OUT" + +mkdir -p "$PD/experiments/logs" # the build helper keeps its source hash there +PATH="$(dirname "$PY"):$PATH" bash "$PD/scripts/kesai/build_ext_if_changed.sh" "$PD" || exit 1 +IMPORTED_AGENT=$("$PY" -c "import pufferlib.ocean.cosim.carla.leaderboard_agent as a; print(a.__file__)") || { echo "ERROR: agent import failed (see traceback above)"; exit 1; } +case "$IMPORTED_AGENT" in + "$PD"/*) echo "agent: $IMPORTED_AGENT" ;; + *) echo "ERROR: agent imported from '$IMPORTED_AGENT', expected $PD"; exit 1 ;; +esac + +# route ids of the xml +mapfile -t ROUTE_IDS < <("$PY" -c " +import sys, xml.etree.ElementTree as ET +print('\n'.join(r.attrib['id'] for r in ET.parse(sys.argv[1]).getroot().iter('route')))" "$ROUTES") +[ -n "$ROUTE_SUBSET" ] && read -r -a ROUTE_IDS <<< "$ROUTE_SUBSET" +# work units = route x seed (":"), each its own server run and result folder +UNITS=() +for seed in "${SEED_LIST[@]}"; do for route_id in "${ROUTE_IDS[@]}"; do UNITS+=("$route_id:$seed"); done; done +[ "$WORKERS" -gt "${#UNITS[@]}" ] && WORKERS=${#UNITS[@]} +# longest route first to the least loaded worker; length = the xml's waypoint polyline (an earlier run's duration would +# read a route that failed early as short) +mapfile -t ASSIGNMENT < <("$PY" - "$ROUTES" "$WORKERS" "${UNITS[@]}" <<'EOF' +import math, sys, xml.etree.ElementTree as ET +routes_xml, workers, units = sys.argv[1], int(sys.argv[2]), sys.argv[3:] # unit = ":" +SETUP_M = 300.0 # server start + world load per route, in route-length-equivalent metres +length = {} +for route in ET.parse(routes_xml).getroot().iter("route"): + pts = [(float(e.get("x")), float(e.get("y"))) for e in route.iter() if e.tag in ("waypoint", "position")] + length[route.get("id")] = sum(math.dist(a, b) for a, b in zip(pts, pts[1:])) +loads = [0.0] * workers; assign = [[] for _ in range(workers)] +for unit in sorted(units, key=lambda u: -length[u.split(":")[0]]): + rid = unit.split(":")[0] + w = min(range(workers), key=lambda i: loads[i]); assign[w].append(unit); loads[w] += length[rid] + SETUP_M +for w in range(workers): + print(w, " ".join(assign[w]), f"# {loads[w] / 1000:.1f} km") +EOF +) +echo "${#ROUTE_IDS[@]} routes x ${#SEED_LIST[@]} seeds = ${#UNITS[@]} route runs over $WORKERS workers on $NUM_GPUS GPUs, policy threads $OMP_THREADS" +printf " worker %s\n" "${ASSIGNMENT[@]}" +echo "route,worker,gpu,attempt,start,wall_s,game_s,pace,status,driving_score,seed" > "$OUT/timing.csv" + +route_done() { # 0 when the route json holds a record that is not one of the crash states CaRL resubmits + "$PY" - "$1" <<'EOF' +import json, sys +try: + records = json.load(open(sys.argv[1]))["_checkpoint"]["records"] +except Exception: + sys.exit(1) +crashed = ("Failed - Agent couldn't be set up", "Failed", "Failed - Simulation crashed", "Failed - Agent crashed") +sys.exit(0 if records and all(r["status"] not in crashed for r in records) else 1) +EOF +} + +log_route_timing() { # $1 worker, $2 gpu, $3 route id, $4 attempt, $5 start epoch, $6 result json, $7 seed + "$PY" - "$@" "$(date +%s)" "$OUT/timing.csv" <<'EOF' +import json, sys, time +w, gpu, route_id, attempt, t0, path, seed, t1, csv_path = sys.argv[1:10] +rec = json.load(open(path))["_checkpoint"]["records"][0] +wall = int(t1) - int(t0); game = float(rec["meta"]["duration_game"]); ds = float(rec["scores"]["score_composed"]) +pace = game / wall if wall > 0 else 0.0 +with open(csv_path, "a") as f: + f.write(f"{route_id},{w},{gpu},{attempt},{time.strftime('%H:%M:%S', time.localtime(int(t0)))},{wall},{game:.1f},{pace:.3f},{rec['status']},{ds:.2f},{seed}\n") +print(f"[w{w} gpu{gpu} route{route_id}{' s' + seed if seed else ''}] done: {wall} s wall, {game:.0f} s game, {pace:.2f}x real time, DS {ds:.1f}, {rec['status']}") +EOF +} + +run_route() { # $1 worker, $2 gpu, $3 route id, $4 attempt, $5 seed: one CARLA server + one evaluator run, into $OUT/routes/route__s + local w=$1 gpu=$2 route_id=$3 attempt=$4 seed=$5 + local route_dir=$OUT/routes/route_$(printf "%02d" "$route_id")_s$seed + local port=$((PORT_BASE + w * 50)) tm_port=$((PORT_BASE + 6000 + w * 50)) + local server_log=$route_dir/carla_server_attempt$attempt.log evaluator_log=$route_dir/evaluator_attempt$attempt.log + local t_start + t_start=$(date +%s) + mkdir -p "$route_dir" + rm -f "$route_dir/result.json" + # only the chase camera needs the renderer; collision sensors and telemetry work under -nullrhi + local pin + pin=$(numa_pin_for_gpu "$gpu") + local render_args=(-RenderOffScreen -nullrhi) + [ "$CARLA_VIEW" = "1" ] && render_args=(-RenderOffScreen -graphicsadapter="$gpu") + $pin "$CARLA_ROOT/CarlaUE4.sh" "${render_args[@]}" -nosound \ + -carla-rpc-port="$port" -carla-streaming-port=$((port + 1)) > "$server_log" 2>&1 & + local server_pid=$! up=1 + echo "$server_pid" >> "$OUT/server_pids" + for _ in $(seq 1 60); do + kill -0 "$server_pid" 2>/dev/null || { echo "[w$w gpu$gpu route$route_id s$seed] CARLA server died"; break; } + "$PY" -c " +import carla, sys +try: + c = carla.Client('localhost', $port); c.set_timeout(10.0); c.get_world().get_map() +except Exception: + sys.exit(1)" 2>/dev/null && up=0 && break + sleep 5 + done + if [ $up -ne 0 ]; then + pkill -9 -P "$server_pid" 2>/dev/null; kill -9 "$server_pid" 2>/dev/null + echo "[w$w gpu$gpu route$route_id s$seed] last lines of $server_log:"; tail -n 20 "$server_log" + return 1 + fi + local log_env=() + if [ "$LOGGING" = "1" ]; then + log_env=(COSIM_TELEMETRY="$route_dir/telemetry" COSIM_WORLD_LOG="$route_dir/world_log") + [ "$CARLA_VIEW" = "1" ] && log_env+=(COSIM_DEBUG_CARLA_VIEW="$route_dir/carla_view") + [ "$OBS_HTML" = "1" ] && log_env+=(COSIM_OBS_HTML="$route_dir/obs_html") + fi + $pin env CUDA_VISIBLE_DEVICES="$gpu" "${log_env[@]}" \ + "$PY" -u "$EVALUATOR" --routes "$ROUTES" --routes-subset "$route_id" --repetitions "$REPETITIONS" \ + --agent "$PD/pufferlib/ocean/cosim/carla/leaderboard_agent.py" --agent-config "$CKPT" \ + --checkpoint "$route_dir/result.json" --track MAP --port "$port" --traffic-manager-port "$tm_port" \ + --traffic-manager-seed "$seed" \ + > "$evaluator_log" 2>&1 + # children first: killing the wrapper shell first reparents the UE4 binary and leaves it running + pkill -9 -P "$server_pid" 2>/dev/null; kill -9 "$server_pid" 2>/dev/null; wait "$server_pid" 2>/dev/null + if ! route_done "$route_dir/result.json"; then + echo "[w$w gpu$gpu route$route_id s$seed] no valid route record; last lines of $evaluator_log:"; tail -n 20 "$evaluator_log" + return 1 + fi + log_route_timing "$w" "$gpu" "$route_id" "$attempt" "$t_start" "$route_dir/result.json" "$seed" + return 0 +} + +run_worker() { # $1 worker: its assigned route runs, sequentially, each retried up to MAX_ATTEMPTS on the next GPU + local w=$1 units unit route_id seed attempt gpu + units=${ASSIGNMENT[w]#* } + units=${units%%#*} + for unit in $units; do + route_id=${unit%%:*}; seed=${unit##*:} + for ((attempt = 1; attempt <= MAX_ATTEMPTS; attempt++)); do + gpu=$(( (w + attempt - 1) % NUM_GPUS )) + echo "[w$w] route $route_id${seed:+ seed $seed} attempt $attempt gpu $gpu $(date +%H:%M:%S)" + run_route "$w" "$gpu" "$route_id" "$attempt" "$seed" && break + echo "[w$w] route $route_id${seed:+ seed $seed} attempt $attempt failed" + done + done +} + +# only the servers this job started (children first: the UE4 binary outlives its wrapper shell otherwise) +trap 'for p in $(cat "$OUT/server_pids" 2>/dev/null); do pkill -9 -P "$p" 2>/dev/null; kill -9 "$p" 2>/dev/null; done' EXIT +WORKER_PIDS=() +ROUTES_T0=$(date +%s) +for ((w = 0; w < WORKERS; w++)); do + [ "$w" -gt 0 ] && sleep "$STAGGER_S" # simultaneous engine starts are the one thing the crashed v1 runs shared + run_worker "$w" > "$OUT/worker_w$w.log" 2>&1 & + WORKER_PIDS+=($!) +done +wait "${WORKER_PIDS[@]}" +ROUTES_WALL_S=$(( $(date +%s) - ROUTES_T0 )) +cat "$OUT"/worker_w*.log +"$PY" - "$OUT/timing.csv" "$ROUTES_WALL_S" "$WORKERS" "$NUM_GPUS" "${#UNITS[@]}" <<'EOF' | tee -a "$OUT/run_info.txt" +import csv, sys +rows = list(csv.DictReader(open(sys.argv[1]))); wall, workers, gpus, n_routes = (int(x) for x in sys.argv[2:6]) +if not rows: + print(f"speed: no route finished in {wall / 3600:.2f} h wall ({workers} workers on {gpus} GPUs)"); sys.exit(0) +game = sum(float(r["game_s"]) for r in rows); route_wall = sum(int(r["wall_s"]) for r in rows) +per_worker = {} +for r in rows: + per_worker[r["worker"]] = per_worker.get(r["worker"], 0) + int(r["wall_s"]) +print(f"speed: {len(rows)}/{n_routes} route runs in {wall / 3600:.2f} h wall, {workers} workers on {gpus} GPUs; game time {game / 3600:.2f} h; " + f"per-route pace {game / route_wall:.2f}x real time; busiest worker {max(per_worker.values()) / 3600:.2f} h, " + f"idlest {min(per_worker.values()) / 3600:.2f} h; route runs/h {len(rows) / (wall / 3600):.1f}") +EOF + +# aggregate: CaRL's parser over the per-route jsons, one subfolder per seed (the parser walks $OUT/aggregate and reports +# mean/std over the seeds per route and town, the layout of CaRL's own evaluation) -> $OUT/aggregate/results.csv, plus +# a per-route table +for f in "$OUT"/routes/route_*/result.json; do + name=$(basename "$(dirname "$f")") + mkdir -p "$OUT/aggregate/s${name##*_s}" && cp "$f" "$OUT/aggregate/s${name##*_s}/${name%_s*}.json" +done +"$PY" "$CARL_WORK_DIR/tools/result_parser.py" --xml "$ROUTES" --results "$OUT/aggregate" +"$PY" - "$OUT/aggregate" "${#UNITS[@]}" "**/route_*.json" <<'EOF' +import glob, json, os, re, statistics, sys +rows = [] # (seed, route_id, status, DS, RC, IP); seed "" = a run from before the seeds layout +for f in glob.glob(os.path.join(sys.argv[1], sys.argv[3]), recursive=True): + parent = os.path.basename(os.path.dirname(f)) + seed = parent[1:] if parent.startswith("s") and parent[1:].isdigit() else "" + for r in json.load(open(f))["_checkpoint"]["records"]: + rows.append((seed, r["route_id"], r["status"], r["scores"]["score_composed"], r["scores"]["score_route"], r["scores"]["score_penalty"])) +def route_index(route_id): + m = re.search(r"_(\d+)", route_id) + return int(m.group(1)) if m else 0 +rows.sort(key=lambda r: (int(r[0]) if r[0] else -1, route_index(r[1]))) +print(f"\n{len(rows)} route records of {sys.argv[2]} route runs") +for seed, route_id, status, ds, rc, ip in rows: + print(f" {('s' + seed) if seed else '-':>3s} {route_id:22s} {status:34s} DS {ds:6.2f} RC {rc:6.2f} IP {ip:5.3f}") +if rows: + n = len(rows) + print(f"mean DS {sum(r[3] for r in rows)/n:.2f} RC {sum(r[4] for r in rows)/n:.2f} IP {sum(r[5] for r in rows)/n:.3f} (all {n} route runs)") + by_seed = {} + for r in rows: + by_seed.setdefault(r[0], []).append(r) + if len(by_seed) > 1: # the paper-style number: the benchmark DS per seed, mean +- std over the seeds + ds = {s: sum(r[3] for r in v) / len(v) for s, v in by_seed.items()} + rc = {s: sum(r[4] for r in v) / len(v) for s, v in by_seed.items()} + for s in sorted(by_seed, key=lambda s: int(s) if s else -1): + print(f"seed {s}: DS {ds[s]:.2f} RC {rc[s]:.2f} ({len(by_seed[s])} routes)") + print(f"over seeds: DS {statistics.mean(ds.values()):.2f} +- {statistics.pstdev(ds.values()):.2f} " + f"RC {statistics.mean(rc.values()):.2f} +- {statistics.pstdev(rc.values()):.2f} (population std over {len(by_seed)} seeds, as in results.csv)") +EOF + +REPORT_ARGS=() +[ "$OBS_HTML" = "1" ] && REPORT_ARGS+=(--obs-html-dir "$OUT/obs_html") +if [ "$REPORT" = "1" ] && [ "$LOGGING" = "1" ]; then + "$PY" "$PD/scripts/eval/analyze_carla_cosim.py" "$OUT"/routes/route_* "$OUT/report" "${REPORT_ARGS[@]}" || echo "HTML report failed (results above are unaffected)" +fi +# after the report: the gallery links to it when it exists +if [ "$OBS_HTML" = "1" ] && [ "$LOGGING" = "1" ]; then + "$PY" "$PD/scripts/eval/render_carla_obs_html.py" "$OUT" --routes "$ROUTES" || echo "obs replay gallery failed (results above are unaffected)" +fi +[ -f "$OUT/report/index.html" ] && echo "report -> $OUT/report/index.html" +[ -f "$OUT/obs_html/index.html" ] && echo "obs replays -> $OUT/obs_html/index.html" +echo "Done -> $OUT" diff --git a/scripts/kesai/12_alpasim.sh b/scripts/kesai/12_alpasim.sh index dcc2882388..fadf133703 100644 --- a/scripts/kesai/12_alpasim.sh +++ b/scripts/kesai/12_alpasim.sh @@ -4,76 +4,131 @@ #SBATCH --nodes 1 #SBATCH --time 1-00:00 #SBATCH --gres gpu:8 -#SBATCH --mem=1007G -#SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/log_%j.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/log_%j.err +#SBATCH --mem=400G # --gres, --mem and --cpus-per-task are overridable on the sbatch line +#SBATCH --cpus-per-task 96 +#SBATCH --output experiments/logs/alpasim_%j.out +#SBATCH --error experiments/logs/alpasim_%j.err #SBATCH --partition dev # AlpaSim closed-loop evaluation of a PufferDrive checkpoint on NuRec scenes, the AlpaSim counterpart of 9_nuPlan.sh: # alpasim-training's teacher driver runs the policy in a shadow env of this checkout (PUFFERDRIVE_ROOT=$PD) on the -# scenes converted to PufferDrive bins; AlpaSim renders, simulates the ego and replays the logged traffic. +# scenes converted to PufferDrive bins; AlpaSim renders, simulates the ego and replays the logged traffic. AlpaSim's +# services run as containers: docker compose on scaleway; on kesai, which has no docker, as enroot containers started +# straight from squashfs images inside this job (CONTAINER_RUNTIME=enroot) or as pyxis job steps (pyxis; every step first +# unpacks its image to node-local disk, ~300 GB per run with this topology). # -# One-time setup: alpasim-training at ALPAGYM_ROOT (`uv sync --all-packages`), the AlpaSim fork at ALPASIM_ROOT, the -# NuRec scenes and converted bins under NUREC_ROOT (alpasim-training README steps 3-4; all-usdzs/ must hold real files), -# and the AlpaSim base image archived to ALPASIM_IMAGE_TAR (REPO=$ALPASIM_ROOT OUT_DIR= sbatch -# --output= --error= $ALPAGYM_ROOT/scripts/build_alpasim_image.sbatch). +# One-time setup: alpasim-training at ALPAGYM_ROOT (`uv sync --all-packages`), the AlpaSim fork at ALPASIM_ROOT, the NuRec +# scenes as real files at USDZ_DIR, their bins with 130 km/h lane limits at BINS (alpasim-training README step 4) and the +# AlpaSim base image plus the NRE renderer image: +# docker: the base image archived in ALPASIM_IMAGE_DIR (cd $ALPAGYM_ROOT && ALPASIM_ROOT=$ALPASIM_ROOT +# OUT_DIR=$ALPASIM_IMAGE_DIR sbatch scripts/build_alpasim_image.sbatch; ~3 h on a dev node, ~31 GB archive), readable +# by this user; the renderer is pulled from nvcr.io. +# enroot/pyxis: both as squashfs files in SQSH_DIR, named the way the wizard looks them up (image basename with ':' +# and '-' replaced by '_'): +# enroot import -o $SQSH_DIR/nre_ga_26.04.sqsh docker://nvcr.io#nvidia/nre/nre-ga:26.04 +# python scripts/kesai/docker_archive_to_sqsh.py alpasim-base-.tar.zst $SQSH_DIR/alpasim_base_.sqsh +# the archive being build_alpasim_image.sbatch's, copied over from a docker host (~100 GB of scratch, see the script). # -# Overridable env: RUN_DIR, CKPT (default RUN_DIR/final_model.pt, else the newest RUN_DIR/models/model_*.pt), PD, -# ALPAGYM_ROOT, ALPASIM_ROOT, NUREC_ROOT, BINS, USDZ_DIR, ALPASIM_IMAGE_TAR, TOPOLOGY (8gpu_64rollouts = whole node; -# 1gpu for a quick check), PARALLEL (scenes in flight: 32 on 8gpu, at most 4 on 1gpu), N_SCENES (0 = every scene that -# passed the conversion checks), SCENE_LIST (file of scene ids), GOAL_MODE (gt | route), CAM_W/CAM_H (16:9 render size; -# the teacher never sees the cameras), RENDER_VIDEO (1 = one mp4 per scene), SIM_TIMEOUT_S (per batch of PARALLEL scenes). +# Settings are the plain assignments below. Only RUN_DIR, TOPOLOGY and N_SCENES are also read from the environment: 3_train +# and 7_eval chain this eval after their own by exporting RUN_DIR. The topology follows the GPUs the job holds (8, 2 or 1), +# so a run on two GPUs of a shared node (12 scenes in flight instead of 32) is +# sbatch --gres gpu:2 --mem 200G --cpus-per-task 48 scripts/kesai/12_alpasim.sh +# and a smoke test is +# N_SCENES=2 sbatch --gres gpu:1 --mem 200G --cpus-per-task 48 scripts/kesai/12_alpasim.sh set -u +source scripts/kesai/cluster_env.sh || exit 1 -export PD=${PD:-/home/bjaeger/PufferDrive} -ALPAGYM_ROOT=${ALPAGYM_ROOT:-/home/bjaeger/alpasim-training} -export ALPASIM_ROOT=${ALPASIM_ROOT:-/home/bjaeger/alpasim} -NUREC_ROOT=${NUREC_ROOT:-/home/shared/data/nurec} -BINS=${BINS:-$NUREC_ROOT/bins} -USDZ_DIR=${USDZ_DIR:-$NUREC_ROOT/all-usdzs} -RUN_DIR=${RUN_DIR:-/home/bjaeger/PufferDrive/experiments/k_scaled_0043_1000} -TOPOLOGY=${TOPOLOGY:-8gpu_64rollouts} -case "$TOPOLOGY" in - 8gpu*) PARALLEL=${PARALLEL:-32} ;; - *) PARALLEL=${PARALLEL:-4} ;; +TOPOLOGY=${TOPOLOGY:-} # empty = from the GPU count below: 8gpu_64rollouts, 2gpu or 1gpu +N_SCENES=${N_SCENES:-0} # 0 = the whole SCENE_LIST +RUN_DIR=${RUN_DIR:-$PD/experiments/k_scaled_0059_1000} +CKPT=$RUN_DIR/final_model.pt +[ -f "$CKPT" ] || CKPT=$(ls "$RUN_DIR"/models/model_*.pt 2>/dev/null | sort | tail -n 1) +GOAL_MODE=gt_time:5 # gt_time: | gt | route +GOAL_RADIUS_M=2 # goal arrival radius Damiano's reference evals use; the teacher's own gt_time default is 6 +VRU_SIZE_PAD_M=0.8 # added to the length and width of every pedestrian and two-wheeler box the teacher sees; 0 = true boxes +PARTNER_SPLIT_LENGTH_M=7 # actors longer than this (training max 7 m) reach the teacher as ceil(length / this) touching pieces; 0 = one box +RENDER_VIDEO=0 # one mp4 per scene; 1 also streams the teacher BEV and makes every policy step ~5x slower +OBS_REPLAY=1 # PufferDrive observation replay html per scene -> $OUT/obs_replays +CAM_W=640 # 16:9 render size; the teacher never sees the cameras +CAM_H=360 +SIM_TIMEOUT_S=3600 # per batch of PARALLEL scenes +ALPASIM_LOG_LEVEL=WARNING # AlpaSim services; INFO also keeps the per-event INFO lines runtime workers and renderer print anyway (~110 MB per 400 scenes) +PROMETHEUS=0 # AlpaSim performance telemetry (~80 MB per run); only alpagym_runtime.perf reads it +KEEP_ROLLOUT_LOGS=none # rollout.asl (~36 MB per scene, offline re-scoring/failure analysis only): all | imperfect (score < 1) | failed | none (never written) +ALPAGYM_ROOT=$HOME/alpasim-training +# held-out split Damiano evaluates on; empty = every converted scene that passed the checks (2149 on the cluster, 1584 of them train) +SCENE_LIST=$ALPAGYM_ROOT/scene_splits/eval_405.txt +export ALPASIM_ROOT=$HOME/alpasim +# shoulder lanes off-road (py123d 1_bernhard_dev) +BINS=$SHARED_DATA/nurec/bins/eval405_130kmh_ndshoulders +case "$CLUSTER" in + kesai) + CONTAINER_RUNTIME=enroot # enroot | pyxis + USDZ_DIR=$SHARED_DATA/nurec/usdz + SQSH_DIR=$SHARED_DATA/nurec/sqsh + FUSE_DIR=$HOME/.cache/alpasim-fuse-tools # squashfuse + fuse-overlayfs the wizard downloads once for enroot + ;; + *) + CONTAINER_RUNTIME=docker + # USDZ_DIR stays Damiano's, the shared all-usdzs symlinks dangle inside the containers + USDZ_DIR=$HOME_BASE/ddacol/code/alpasim-training/data/distill/usdz + ALPASIM_IMAGE_DIR=$HOME_BASE/ddacol/code/alpasim-training/data/images + ;; esac -N_SCENES=${N_SCENES:-0} -SCENE_LIST=${SCENE_LIST:-} -GOAL_MODE=${GOAL_MODE:-gt} -CAM_W=${CAM_W:-640} -CAM_H=${CAM_H:-360} -RENDER_VIDEO=${RENDER_VIDEO:-1} -SIM_TIMEOUT_S=${SIM_TIMEOUT_S:-3600} ALPAGYM_PY=$ALPAGYM_ROOT/.venv/bin/python # uv: the host syncs the AlpaSim checkout's venv before it starts the Wizard export PATH="$HOME/.local/bin:$PATH" -if [ -z "${CKPT:-}" ]; then - CKPT=$RUN_DIR/final_model.pt - [ -f "$CKPT" ] || CKPT=$(ls "$RUN_DIR"/models/model_*.pt 2>/dev/null | sort | tail -n 1) -fi [ -n "$CKPT" ] || { echo "no checkpoint in $RUN_DIR (final_model.pt or models/model_*.pt)"; exit 1; } -for path in "$CKPT" "$ALPAGYM_PY" "$ALPASIM_ROOT/pyproject.toml" "$BINS/scene_manifest.jsonl" "$USDZ_DIR"; do - [ -e "$path" ] || { echo "missing $path"; exit 1; } +case "$KEEP_ROLLOUT_LOGS" in all|imperfect|failed|none) ;; *) echo "KEEP_ROLLOUT_LOGS must be all, imperfect, failed or none, got $KEEP_ROLLOUT_LOGS"; exit 1 ;; esac +for path in "$CKPT" "$ALPAGYM_PY" "$ALPASIM_ROOT/pyproject.toml" "$BINS/scene_manifest.jsonl" "$USDZ_DIR" ${SCENE_LIST:+"$SCENE_LIST"}; do + [ -r "$path" ] || { echo "missing or not readable: $path"; exit 1; } done [ -f "$(dirname "$CKPT")/config.yaml" ] || [ -f "$(dirname "$(dirname "$CKPT")")/config.yaml" ] \ || { echo "no config.yaml next to or one level above $CKPT"; exit 1; } -[ -z "$(find "$USDZ_DIR" -maxdepth 1 -type l -print -quit)" ] \ +[ -z "$(find "$USDZ_DIR" -type l -print -quit)" ] \ || { echo "$USDZ_DIR holds symlinks, which dangle inside the renderer container; point USDZ_DIR at the real files"; exit 1; } -for tool in docker nvidia-smi uv; do +case "$CONTAINER_RUNTIME" in + docker) RUNTIME_TOOLS="docker" ;; + enroot) RUNTIME_TOOLS="enroot /usr/bin/fusermount" ;; # fusermount: what the wizard's fusermount3 symlink points at + pyxis) RUNTIME_TOOLS="srun" ;; + *) echo "CONTAINER_RUNTIME must be docker, enroot or pyxis, got $CONTAINER_RUNTIME"; exit 1 ;; +esac +for tool in $RUNTIME_TOOLS nvidia-smi uv; do command -v "$tool" > /dev/null || { echo "missing tool: $tool"; exit 1; } done +[ "$CONTAINER_RUNTIME" != pyxis ] || srun --help 2>&1 | grep -q -- --container-image \ + || { echo "srun has no --container-image option: pyxis is not installed on $(hostname)"; exit 1; } +NUM_GPUS=$(nvidia-smi -L | grep -c "^GPU") # SLURM shows the job only its own GPUs +if [ -z "$TOPOLOGY" ]; then + case "$NUM_GPUS" in + 8) TOPOLOGY=8gpu_64rollouts ;; + 2) TOPOLOGY=2gpu ;; + 1) TOPOLOGY=1gpu ;; + *) echo "no topology for $NUM_GPUS GPUs: request 8, 2 or 1 (--gres) or set TOPOLOGY"; exit 1 ;; + esac +fi case "$TOPOLOGY" in 8gpu*) - NUM_GPUS=$(nvidia-smi -L | grep -c "^GPU") [ "$NUM_GPUS" -eq 8 ] || { echo "$TOPOLOGY pins its services to GPUs 0-7 and needs all 8, got $NUM_GPUS"; exit 1; } - ;; + PARALLEL=64 ;; # scenes in flight = the topology's capacity; the runtime fills renderers in order, 32 left two idle + 2gpu) PARALLEL=12 ;; + *) PARALLEL=4 ;; esac ALPASIM_VERSION=$(grep -m1 '^version' "$ALPASIM_ROOT/pyproject.toml" | sed -E 's/.*"(.*)".*/\1/') BASE_TAG=alpasim-base:$ALPASIM_VERSION -ALPASIM_IMAGE_TAR=${ALPASIM_IMAGE_TAR:-/home/bjaeger/docker_images/alpasim-base-$ALPASIM_VERSION.tar.zst} NRE_TAG=$(grep -m1 -oE 'nvcr.io/nvidia/nre/[^[:space:]"]+' "$ALPASIM_ROOT/src/wizard/configs/base_config.yaml") [ -n "$NRE_TAG" ] || { echo "no NRE renderer image in $ALPASIM_ROOT/src/wizard/configs/base_config.yaml"; exit 1; } +if [ "$CONTAINER_RUNTIME" = docker ]; then + ALPASIM_IMAGE_TAR=$ALPASIM_IMAGE_DIR/alpasim-base-$ALPASIM_VERSION.tar.zst + [ -r "$ALPASIM_IMAGE_TAR" ] || docker image inspect "$BASE_TAG" > /dev/null 2>&1 \ + || { echo "$ALPASIM_IMAGE_TAR is missing or not readable and $BASE_TAG is not on $(hostname)"; exit 1; } +else + BASE_SQSH=$SQSH_DIR/$(basename "$BASE_TAG" | sed 's/[:-]/_/g').sqsh + NRE_SQSH=$SQSH_DIR/$(basename "$NRE_TAG" | sed 's/[:-]/_/g').sqsh + for sqsh in "$BASE_SQSH" "$NRE_SQSH"; do + [ -r "$sqsh" ] || { echo "missing image $sqsh; create it as described at the top of this script"; exit 1; } + done +fi echo "Evaluating checkpoint: $CKPT" # drop inherited entries from other PufferDrive checkouts: the teacher would import their pufferlib instead of $PD's @@ -93,7 +148,13 @@ mapfile -t SCENES < <("$ALPAGYM_PY" - "$BINS/scene_manifest.jsonl" "$N_SCENES" " import json, sys manifest, n_scenes, scene_list = sys.argv[1], int(sys.argv[2]), sys.argv[3] rows = [json.loads(line) for line in open(manifest).read().splitlines()[1:]] -scenes = [r["alpasim_scene_id"] for r in rows if r["source_parity_pass"] and r["gt_replay_pass"]] +def passed_checks(row): + if "gt_replay_pass" in row: + return row["source_parity_pass"] and row["gt_replay_pass"] + if "offroad_verdict" in row: # manifests built before alpasim-training d29bdd9 (2026-08-31) + return row["source_parity_pass"] and row["offroad_verdict"] == "both_pass" + sys.exit(f"{manifest}: unknown scene row schema, keys: {sorted(row)}") +scenes = [r["alpasim_scene_id"] for r in rows if passed_checks(r)] print(f"{len(scenes)} of {len(rows)} converted scenes passed the conversion checks", file=sys.stderr) if scene_list: usable = set(scenes) @@ -106,58 +167,135 @@ EOF [ ${#SCENES[@]} -gt 0 ] || { echo "no scenes to evaluate"; exit 1; } # results live in the model's own eval folder, next to the PufferDrive benchmark evals -OUT=$RUN_DIR/eval/alpasim_${GOAL_MODE}_$(basename "$CKPT" .pt)_$(date +%Y%m%d_%H%M%S)_${SLURM_JOB_ID:-local} +OUT=$RUN_DIR/eval/alpasim_${GOAL_MODE//:/}_r${GOAL_RADIUS_M}$([ "$VRU_SIZE_PAD_M" = "0" ] || echo "_vrupad$VRU_SIZE_PAD_M")$([ "$PARTNER_SPLIT_LENGTH_M" = "0" ] || echo "_split$PARTNER_SPLIT_LENGTH_M")_$(basename "$CKPT" .pt)_$(date +%Y%m%d_%H%M%S)_${SLURM_JOB_ID:-local} mkdir -p "$OUT" +# the trainer deletes every older models/model_*.pt when it saves the next one: evaluate a copy that outlives the run +CKPT_SOURCE=$CKPT +CKPT_CONFIG=$(dirname "$CKPT")/config.yaml +[ -f "$CKPT_CONFIG" ] || CKPT_CONFIG=$(dirname "$(dirname "$CKPT")")/config.yaml +cp "$CKPT" "$OUT/" && cp "$CKPT_CONFIG" "$OUT/config.yaml" || { echo "could not copy $CKPT and $CKPT_CONFIG into $OUT"; exit 1; } +CKPT=$OUT/$(basename "$CKPT") { - echo "checkpoint $CKPT"; echo "scenes ${#SCENES[@]}"; echo "topology $TOPOLOGY parallel $PARALLEL goal_mode $GOAL_MODE" + echo "checkpoint $CKPT_SOURCE (evaluated copy: $CKPT)"; echo "scenes ${#SCENES[@]}"; echo "topology $TOPOLOGY parallel $PARALLEL goal_mode $GOAL_MODE goal_radius_m $GOAL_RADIUS_M" + echo "vru_size_pad_m $VRU_SIZE_PAD_M" + echo "partner_split_length_m $PARTNER_SPLIT_LENGTH_M" + echo "container_runtime $CONTAINER_RUNTIME" for repo in "$PD" "$ALPAGYM_ROOT" "$ALPASIM_ROOT"; do echo "git $repo $(git -C "$repo" rev-parse --short HEAD 2>/dev/null) $(git -C "$repo" diff --quiet 2>/dev/null || echo dirty)" done } > "$OUT/run_info.txt" echo "Results -> $OUT" -docker info > /dev/null 2>&1 || { echo "docker is not usable on $(hostname)"; exit 1; } -# own compose project and port window per job, so evals sharing a node never touch each other's services -export COMPOSE_PROJECT_NAME=alpasim_${SLURM_JOB_ID:-$$} +# own port window per job, so evals sharing a node never touch each other's services BASEPORT=$((6000 + (${SLURM_JOB_ID:-$$} % 90) * 20)) +if [ "$CONTAINER_RUNTIME" = docker ]; then + docker info > /dev/null 2>&1 || { echo "docker is not usable on $(hostname)"; exit 1; } + export COMPOSE_PROJECT_NAME=alpasim_${SLURM_JOB_ID:-$$} +fi LOADED_BASE=0 PULLED_NRE=0 cleanup() { - local compose_file=$OUT/alpagym/alpasim/wizard_0/docker-compose.yaml - [ -f "$compose_file" ] && docker compose -f "$compose_file" down --volumes --remove-orphans > /dev/null 2>&1 - for container in $(docker ps -aq --filter "label=com.docker.compose.project=$COMPOSE_PROJECT_NAME"); do - docker rm -f "$container" > /dev/null 2>&1 - done - # /var/lib/docker is tmpfs on the compute nodes: an image left behind keeps holding ~40 GB of node RAM - [ "$LOADED_BASE" = 1 ] && docker rmi "$BASE_TAG" > /dev/null 2>&1 - [ "$PULLED_NRE" = 1 ] && docker rmi "$NRE_TAG" > /dev/null 2>&1 + # the service containers outlive the wizard when the host stops it before they exit + case "$CONTAINER_RUNTIME" in + docker) + local compose_file=$OUT/alpagym/alpasim/wizard_0/docker-compose.yaml + [ -f "$compose_file" ] && docker compose -f "$compose_file" down --volumes --remove-orphans > /dev/null 2>&1 + for container in $(docker ps -aq --filter "label=com.docker.compose.project=$COMPOSE_PROJECT_NAME"); do + docker rm -f "$container" > /dev/null 2>&1 + done + # /var/lib/docker is tmpfs on the compute nodes: an image left behind keeps holding ~40 GB of node RAM + [ "$LOADED_BASE" = 1 ] && docker rmi "$BASE_TAG" > /dev/null 2>&1 + [ "$PULLED_NRE" = 1 ] && docker rmi "$NRE_TAG" > /dev/null 2>&1 + ;; + enroot) # the node is ours alone: every enroot container of this user on it belongs to this eval + pkill -TERM -u "$(id -un)" -f "enroot start --rw" > /dev/null 2>&1 + ;; + pyxis) + squeue -h -s -j "${SLURM_JOB_ID:-0}" -o "%i %j" 2>/dev/null | while read -r step_id step_name; do + case "$step_name" in alpasim-*) scancel --signal=TERM "$step_id" > /dev/null 2>&1 ;; esac + done + ;; + esac } trap cleanup EXIT -if ! docker image inspect "$BASE_TAG" > /dev/null 2>&1; then - [ -f "$ALPASIM_IMAGE_TAR" ] || { echo "$BASE_TAG is not on $(hostname) and $ALPASIM_IMAGE_TAR is missing"; exit 1; } - echo "Loading $BASE_TAG from $ALPASIM_IMAGE_TAR" - LOADED_BASE=1 - zstd -dc "$ALPASIM_IMAGE_TAR" | docker load || exit 1 -fi -if ! docker image inspect "$NRE_TAG" > /dev/null 2>&1; then - echo "Pulling $NRE_TAG" - PULLED_NRE=1 - docker pull "$NRE_TAG" || exit 1 -fi +case "$CONTAINER_RUNTIME" in + docker) + if ! docker image inspect "$BASE_TAG" > /dev/null 2>&1; then + [ -r "$ALPASIM_IMAGE_TAR" ] || { echo "$BASE_TAG is not on $(hostname) and $ALPASIM_IMAGE_TAR is missing or not readable"; exit 1; } + echo "Loading $BASE_TAG from $ALPASIM_IMAGE_TAR" + LOADED_BASE=1 + zstd -dc "$ALPASIM_IMAGE_TAR" | docker load || exit 1 + fi + if ! docker image inspect "$NRE_TAG" > /dev/null 2>&1; then + echo "Pulling $NRE_TAG" + PULLED_NRE=1 + docker pull "$NRE_TAG" || exit 1 + fi + WIZARD_RUN_METHOD=DOCKER_COMPOSE + ;; + enroot) + echo "enroot $(enroot version 2>&1 | head -n 1) on $(hostname)" + # the host creates $ALPASIM_ROOT/.venv (uv sync) on the first run; until then the wizard fetches the tools itself + [ ! -x "$ALPASIM_ROOT/.venv/bin/python" ] \ + || "$ALPASIM_ROOT/.venv/bin/python" -c "from alpasim_wizard.deployment.fuse_tools import ensure_fuse_tools; ensure_fuse_tools('$FUSE_DIR')" \ + || { echo "no squashfuse + fuse-overlayfs in $FUSE_DIR: the wizard downloads them from github.com and archive.ubuntu.com (needs dpkg-deb)"; exit 1; } + # the wizard's squashfuse is an Ubuntu 18.04 build: nodes without libfuse2 need its libraries copied into $FUSE_DIR/lib + MISSING_FUSE_LIBS=$([ -x "$FUSE_DIR/bin/squashfuse" ] && LD_LIBRARY_PATH="$FUSE_DIR/lib" ldd "$FUSE_DIR/bin/squashfuse" | grep "not found") + [ -z "$MISSING_FUSE_LIBS" ] || { echo "$FUSE_DIR/bin/squashfuse cannot run on $(hostname):"; echo "$MISSING_FUSE_LIBS"; \ + echo "apt-get download the packages (libfuse2t64, liblzo2-2), dpkg-deb -x them and copy the .so files into $FUSE_DIR/lib"; exit 1; } + WIZARD_RUN_METHOD=SLURM_ENROOT + ;; + pyxis) + WIZARD_RUN_METHOD=SLURM + ;; +esac # the teacher is an external driver, but the 8gpu topologies re-assert driver.skip=false and abort at startup -OVERRIDES="wizard.baseport=$BASEPORT runtime.endpoints.driver.skip=true" -[ "$RENDER_VIDEO" = "1" ] || OVERRIDES="$OVERRIDES eval.video.render_video=false" -echo "${#SCENES[@]} scenes, $TOPOLOGY with $PARALLEL in flight; full log: $OUT/evaluate.log" +[ "$PROMETHEUS" = "1" ] && PROMETHEUS_FLAG=true || PROMETHEUS_FLAG=false +[ "$KEEP_ROLLOUT_LOGS" = "none" ] && SAVE_ROLLOUT_LOG_FLAG=false || SAVE_ROLLOUT_LOG_FLAG=true +OVERRIDES="wizard.baseport=$BASEPORT runtime.endpoints.driver.skip=true wizard.log_level=$ALPASIM_LOG_LEVEL wizard.prometheus.start_prometheus=$PROMETHEUS_FLAG" +OVERRIDES="$OVERRIDES runtime.simulation_config.save_rollout_log=$SAVE_ROLLOUT_LOG_FLAG" +if [ "$CONTAINER_RUNTIME" != docker ]; then + # root inside the containers as with docker (the image's venv and caches are root's), started in the image's /repo: + # pyxis would otherwise start the base-image services in the host cwd, which does not exist inside them + OVERRIDES="$OVERRIDES defines.base_image=$BASE_TAG wizard.sqshcaches=[$SQSH_DIR] wizard.slurm_cpu_bind_none=true" + for service in renderer physics controller trafficsim runtime prometheus; do + OVERRIDES="$OVERRIDES services.$service.remap_root=true" + done + for service in physics controller trafficsim runtime prometheus; do + OVERRIDES="$OVERRIDES services.$service.workdir=/repo" + done + # enroot's import leaves the NRE entrypoint's $0.runfiles (/app/run.runfiles) empty: name the real Bazel runfiles tree + NRE_RUNFILES=/app/internal/scripts/pycena/runtime/pycena_nrm_full.runfiles + OVERRIDES="$OVERRIDES 'services.renderer.environments=[\"HOME=/tmp\",\"XDG_CACHE_HOME=/tmp/.cache\",\"OMP_NUM_THREADS=1\",\"RUNFILES_DIR=$NRE_RUNFILES\"]'" + OVERRIDES="$OVERRIDES services.renderer.workdir=$NRE_RUNFILES/_main" + if [ "$CONTAINER_RUNTIME" = enroot ]; then + # enroot runs the NRE image's entrypoint (the app itself) in front of the command; /etc/rc sources rc.local first + printf 'exec "$@"\n' > "$OUT/nre_rc.local" + OVERRIDES="$OVERRIDES 'services.renderer.volumes=[\"$USDZ_DIR:/mnt/nre-data\",\"$ALPASIM_ROOT/data/nre-artifacts/ego-hoods:/mnt/ego-hoods\",\"$OUT/nre_rc.local:/etc/rc.local\"]'" + OVERRIDES="$OVERRIDES wizard.fuse_dir=$FUSE_DIR" + fi +fi +[ "$RENDER_VIDEO" = "1" ] && RENDER_VIDEO_FLAG=true || RENDER_VIDEO_FLAG=false +[ "$OBS_REPLAY" = "1" ] && OBS_REPLAY_FLAG=true || OBS_REPLAY_FLAG=false +echo "${#SCENES[@]} scenes, $TOPOLOGY with $PARALLEL in flight on $CONTAINER_RUNTIME; full log: $OUT/evaluate.log (logs/ once the run is tidied)" +# the teacher serves 2 x PARALLEL sessions from one process: one BLAS/OpenMP thread each, or torch spins every core of the node +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 cd "$ALPAGYM_ROOT" || exit 1 "$ALPAGYM_PY" -m alpagym_host.evaluate \ --model_kind teacher \ --teacher_checkpoint "$CKPT" \ --teacher_bins_dir "$BINS" \ --teacher_goal_mode "$GOAL_MODE" \ + --teacher_goal_radius "$GOAL_RADIUS_M" \ + --teacher_vru_size_pad_m "$VRU_SIZE_PAD_M" \ + --teacher_partner_split_length_m "$PARTNER_SPLIT_LENGTH_M" \ + --render_video "$RENDER_VIDEO_FLAG" \ + --teacher_obs_replay "$OBS_REPLAY_FLAG" \ --local_scene_dir "$USDZ_DIR" \ --scene_ids "${SCENES[@]}" \ --topology "$TOPOLOGY" \ + --wizard_run_method "$WIZARD_RUN_METHOD" \ --num_parallel_scenes "$PARALLEL" \ --camera_width "$CAM_W" --camera_height "$CAM_H" \ --simulation_timeout_s "$SIM_TIMEOUT_S" \ @@ -166,24 +304,60 @@ cd "$ALPAGYM_ROOT" || exit 1 > "$OUT/evaluate.log" 2>&1 EVAL_STATUS=$? +# results.txt: the official score, AlpaSim's km between incidents (all and at-fault) and the mean of every per-scene metric if [ -f "$OUT/alpagym/evaluation.yaml" ]; then - "$ALPAGYM_PY" - "$OUT/alpagym" <<'EOF' -import sys + "$ALPAGYM_PY" - "$OUT/alpagym/evaluation.yaml" "$OUT/alpagym/alpasim/wizard_0/aggregate/results-summary.json" "$OUT/results.txt" <<'EOF' +import json, statistics, sys from pathlib import Path import yaml -run = Path(sys.argv[1]) -summary = yaml.safe_load((run / "evaluation.yaml").read_text()) -print(f"\n{len(summary['scenes'])} scenes evaluated, complete: {summary['complete']}") -aggregate = run / "metrics" / "aggregate.yaml" -if aggregate.exists(): - for name, stats in yaml.safe_load(aggregate.read_text())["metrics"].items(): - print(f" {name:34s} mean {stats['mean']:9.4f} median {stats['median']:9.4f} n {stats['num_finite']}") +summary = yaml.safe_load(Path(sys.argv[1]).read_text()) +aggregate_path = Path(sys.argv[2]) +scenes = summary["scenes"] +official = summary.get("official_score") +if official is None: + lines = [f"official_score n/a (incomplete: {len(scenes)} scenes finished, {len(summary['failed_scenes'])} failed)"] +else: + lines = [f"official_score {official['score']:.4f} scenes {official['num_scored']} failed {official['num_failed']}"] +rows = [] +runs = json.loads(aggregate_path.read_text())["metrics_results"] if aggregate_path.exists() else [] +for run in runs: + prefix = f"{run['run_name']} " if len(runs) > 1 else "" + for name in ("avg_dist_between_incidents", "avg_dist_between_incidents_at_fault"): + km_per_incident = run.get(name) + if km_per_incident is None: + rows.append((prefix + name, "n/a (no incidents)")) + continue + per_km = f"{1.0 / km_per_incident:.4f} incidents/km" if km_per_incident > 0 else "inf incidents/km" + rows.append((prefix + name, f"{km_per_incident:9.4f} km ({per_km})")) +for name in sorted({name for scene in scenes for name in scene["metrics"]}): + values = [float(scene["metrics"][name]) for scene in scenes if name in scene["metrics"]] + rows.append((name, f"mean {statistics.mean(values):9.4f} n {len(values)}")) +width = max((len(name) for name, _ in rows), default=0) +lines += [f"{name:{width}s} {value}" for name, value in rows] +text = "".join(line + "\n" for line in lines) +Path(sys.argv[3]).write_text(text) +sys.stdout.write("\n" + text) EOF fi if [ $EVAL_STATUS -ne 0 ]; then echo "AlpaSim evaluation failed (exit $EVAL_STATUS); last host lines of $OUT/evaluate.log:" grep -vE '^[a-z0-9_-]+-[0-9]+-[0-9]+ +\|' "$OUT/evaluate.log" | tail -n 30 + for service_log in "$OUT"/alpagym/alpasim/wizard_0/txt-logs/out-*-log.txt; do + [ -f "$service_log" ] || continue + echo "last lines of $service_log:" + tail -n 5 "$service_log" + done +fi +# the stack is done: take it down while docker-compose.yaml exists, the tidy deletes it +cleanup +# results.txt, obs_replays/, logs/, alpasim/ (+ videos/); container INFO lines stay in evaluate.log for debugging and failed runs +TIDY_FLAGS="" +{ [ "$ALPASIM_LOG_LEVEL" = "INFO" ] || [ $EVAL_STATUS -ne 0 ]; } && TIDY_FLAGS="--keep_container_info" +"$ALPAGYM_PY" "$PD/scripts/eval/tidy_alpasim_run.py" "$OUT" --keep_rollout_logs "$KEEP_ROLLOUT_LOGS" $TIDY_FLAGS || echo "tidy failed" +# score-named pages with the AlpaSim score panel + index.html gallery, as 9_nuPlan.sh and 11_carla_longest6_v2.sh +if [ -d "$OUT/obs_replays" ] && [ -f "$OUT/alpasim/results-summary.json" ]; then + "$ALPAGYM_PY" "$PD/scripts/eval/render_alpasim_obs_html.py" "$OUT" --prune || echo "obs replay gallery failed" fi -echo "per-scene metrics -> $OUT/alpagym/metrics/per_scene.csv; videos -> $OUT/alpagym/alpasim/wizard_0/rollouts" +echo "results -> $OUT/results.txt; obs replays -> $OUT/obs_replays/index.html; logs -> $OUT/logs; per-scene scores + per-step metrics -> $OUT/alpasim" echo "Done -> $OUT" exit $EVAL_STATUS diff --git a/scripts/kesai/13_nuPlan_nonreactive.sh b/scripts/kesai/13_nuPlan_nonreactive.sh index 18d75f0c9b..3762752d3b 100644 --- a/scripts/kesai/13_nuPlan_nonreactive.sh +++ b/scripts/kesai/13_nuPlan_nonreactive.sh @@ -4,13 +4,13 @@ #SBATCH --nodes 1 #SBATCH --time 1-00:00 #SBATCH --gres gpu:8 -#SBATCH --mem=1007G -#SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/log_%j.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/log_%j.err +#SBATCH --mem=0 +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --output experiments/logs/nuPlan_nonreactive_%j.out +#SBATCH --error experiments/logs/nuPlan_nonreactive_%j.err #SBATCH --partition dev # nuPlan Val14 closed-loop evaluation against non-reactive log-replay agents (CLS-NR): 9_nuPlan.sh with the non-reactive # challenge, so the reactive and non-reactive scores always come from the same planner settings. Overridable env: as in # 9_nuPlan.sh; results go to $RUN_DIR/eval/nuplan_val14_nonreactive_*. -PD=${PD:-/home/bjaeger/PufferDrive} +source scripts/kesai/cluster_env.sh || exit 1 CHALLENGES=closed_loop_nonreactive_agents_pufferdrive exec bash "$PD/scripts/kesai/9_nuPlan.sh" diff --git a/scripts/kesai/14_carla_longest6_v1.sh b/scripts/kesai/14_carla_longest6_v1.sh new file mode 100755 index 0000000000..ef36e977fc --- /dev/null +++ b/scripts/kesai/14_carla_longest6_v1.sh @@ -0,0 +1,433 @@ +#!/bin/bash +#SBATCH --job-name eval_longest6_v1 +#SBATCH --ntasks 1 +#SBATCH --nodes 1 +#SBATCH --time 1-00:00 +#SBATCH --gres gpu:8 +#SBATCH --mem=0 +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --output experiments/logs/carla_longest6_v1_%j.out +#SBATCH --error experiments/logs/carla_longest6_v1_%j.err +#SBATCH --partition dev +# CARLA longest6 v1 evaluation of a PufferDrive checkpoint through carla_garage's unmodified Leaderboard 1.0 +# (CARLA 0.9.10.1, pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py), the Leaderboard 1.0 counterpart of +# 11_carla_longest6_v2.sh. The garage evaluator runs in the `garage` conda env (Python 3.7, the CARLA 0.9.10 egg); +# the agent starts pufferlib/ocean/cosim/carla/lb1/policy_server.py in a PufferDrive python ($SERVER_PY) per route. +# +# WORKERS parallel workers (default 16), each one CARLA server + one evaluator at a time, spread round-robin over the +# GPUs of the allocation (1-8, e.g. `sbatch --gres gpu:2`): a server needs ~2 cores and no GPU once the agent has +# turned the spectator render off, so several share a GPU. The 36 routes (one split xml each) x SEEDS passes (route +# runs, $OUT/routes/route__s/) go longest-first to the least loaded worker (route length from the xml), every +# route run gets its own server (restarted per route), worker launches are staggered, a crashed route run is retried on +# the next GPU, and the result jsons are aggregated at the end with garage's tools/result_parser.py over +# $OUT/aggregate/s/ (results.csv) plus the nuPlan-style HTML report +# ($OUT/report/index.html) and the obs replay gallery of all routes ($OUT/obs_html/index.html, as in +# 11_carla_longest6_v2.sh). Speed: wall and game time of every route go to $OUT/timing.csv, the run totals to +# run_info.txt and stdout. +# +# Overridable env: RUN_DIR, ROUTE_DIR (the split xmls), ROUTES_XML (the combined longest6.xml, for the parser +# and the gallery), SCENARIOS (1 = eval_scenarios.json, 0 = no_scenarios.json), SEEDS ("0 1 2" by default: traffic-manager +# seeds, one full pass over the routes per seed, the protocol of carla_garage's and CaRL's evaluate_routes_slurm.py with +# --num_repetitions 3; results.csv then carries mean and std over the seeds per route and per town, the stdout summary the +# DS of every seed and mean +- std over the seeds; SEEDS= (empty) = one pass with the evaluator's default seed 100 and the +# folder layout of the runs before 26 Oct 2026), REPETITIONS (leaderboard-internal repetitions within a pass, 1), +# ROUTE_SUBSET ("0 5 17": only these route ids), WORKERS, NUM_GPUS, STAGGER_S (15: seconds between worker launches), PORT_BASE +# (CARLA rpc port of worker 0; worker w uses PORT_BASE+50w, TM PORT_BASE+6000+50w), OMP_NUM_THREADS (policy server +# threads; 2 from 16 workers on, else 4), LOGGING (1, default: telemetry, world log and the HTML report per route; 0: scores +# only), CARLA_VIEW (0, default: no chase-cam; 1: chase-cam mp4 per route, needs LOGGING=1), OBS_HTML (1 = also +# the interactive obs replay per route, rendered into $OUT/obs_html at the end, large; needs LOGGING=1), REPORT +# (0 = skip the HTML report), MAX_ATTEMPTS, CARLA_ROOT, GARAGE_WORK_DIR, GARAGE_PY, SERVER_PY, PD, +# COSIM_MAX_SPEED_MPS (ego speed cap, default 30), COSIM_ZERO_PARTNER_STOPPED_TIME (default 1), +# COSIM_VRU_SIZE_PAD_M (default 0.8: added to the length and width of every walker and two-wheeler box, so they read +# strictly larger than the smallest training vehicle; 0 = true boxes), LIBTIFF5_COMPAT (1: when `import carla` fails for +# the egg's missing libtiff.so.5, try the garage env's lib dir (`conda install -n garage "libtiff<4.5"`), then link the +# system libtiff.so.6 under that name; 0: never link). +# +# Completing a partial run (routes lost to CARLA crashes or the job time limit): COMPLETE_RUN= +# reruns, in place, only its routes without a valid record (no result.json, or a crash status the leaderboard would +# resubmit; ROUTE_SUBSET narrows the candidates, deleting a route run's result.json forces it). SCENARIOS, SEEDS, REPETITIONS +# and COSIM_VRU_SIZE_PAD_M default to the run's own values from run_info.txt (a run from before the seeds layout has no +# seeds line and is completed as it was driven: single default seed, route_ folders), CKPT must be the run's checkpoint. The replaced +# route folders and worker logs move to $OUT/pre_completion_/, timing.csv grows, run_info.txt gets a completion +# line, and results.csv, the report and the gallery are rebuilt over all routes. +# scripts/kesai/17_carla_longest6_v1_complete_local.sh does this on a desktop for a run folder synced from the cluster. +set -u +source scripts/kesai/cluster_env.sh || exit 1 + +export GARAGE_PY=${GARAGE_PY:-$HOME/miniconda3/envs/garage/bin/python} # cp37: the CARLA 0.9.10 egg +export SERVER_PY=${SERVER_PY:-$HOME/miniconda3/envs/carl/bin/python} # any PufferDrive-capable python +export CARLA_ROOT=${CARLA_ROOT:-$HOME/CARLA_0.9.10} +export GARAGE_WORK_DIR=${GARAGE_WORK_DIR:-$HOME/carla_garage_1} +RUN_DIR=${RUN_DIR:-$PD/experiments/k_scaled_0045_1000} +# the agent finds config.yaml next to final_model.pt (or one level above a models/*.pt) +export CKPT=${CKPT:-$RUN_DIR/final_model.pt} +ROUTE_DIR=${ROUTE_DIR:-$GARAGE_WORK_DIR/leaderboard/data/longest6_split} +ROUTES_XML=${ROUTES_XML:-$GARAGE_WORK_DIR/leaderboard/data/longest6.xml} +COMPLETE_RUN=${COMPLETE_RUN:-} +if [ -n "$COMPLETE_RUN" ]; then + [ -d "$COMPLETE_RUN/routes" ] && [ -f "$COMPLETE_RUN/run_info.txt" ] || { echo "COMPLETE_RUN=$COMPLETE_RUN is not a run folder (needs routes/ and run_info.txt)"; exit 1; } + run_info_value() { awk -v key="$1" '$1 == key {print $2; exit}' "$COMPLETE_RUN/run_info.txt"; } + # the run's own settings, so the completed routes are scored like the rest + PREVIOUS_SCENARIO_FILE=$(run_info_value scenarios) + [ -n "$PREVIOUS_SCENARIO_FILE" ] && SCENARIOS=${SCENARIOS:-$([ "$(basename "$PREVIOUS_SCENARIO_FILE")" = "no_scenarios.json" ] && echo 0 || echo 1)} + REPETITIONS=${REPETITIONS:-$(run_info_value repetitions)} + # the rest of the line: several seeds; a run from before the seeds layout has no such line -> SEEDS stays empty + SEEDS=${SEEDS-$(awk '$1 == "seeds" {$1 = ""; print substr($0, 2); exit}' "$COMPLETE_RUN/run_info.txt")} + COSIM_VRU_SIZE_PAD_M=${COSIM_VRU_SIZE_PAD_M:-$(run_info_value vru_size_pad_m)} +fi +SCENARIOS=${SCENARIOS:-1} +if [ "$SCENARIOS" = "1" ]; then + SCENARIO_FILE=$GARAGE_WORK_DIR/leaderboard/data/scenarios/eval_scenarios.json +else + SCENARIO_FILE=$GARAGE_WORK_DIR/leaderboard/data/scenarios/no_scenarios.json +fi +REPETITIONS=${REPETITIONS:-1} +SEEDS=${SEEDS-0 1 2} # traffic-manager seeds, one pass over the routes each (the carla_garage / CaRL protocol); empty = legacy +for seed in $SEEDS; do [[ "$seed" =~ ^[0-9]+$ ]] || { echo "SEEDS=\"$SEEDS\" must be integers"; exit 1; }; done +if [ -n "$SEEDS" ]; then read -r -a SEED_LIST <<< "$SEEDS"; else SEED_LIST=(""); fi # "" = the evaluator's default seed, route_ folders +unit_dir() { printf "route_%02d%s" "$1" "${2:+_s$2}"; } # $1 route id, $2 seed +unit_json() { printf "%s%s/longest_weathers_%s.json" "$OUT/aggregate" "${2:+/s$2}" "$1"; } # the parser's copy of the record +ROUTE_SUBSET=${ROUTE_SUBSET:-} +PORT_BASE=${PORT_BASE:-2000} +MAX_ATTEMPTS=${MAX_ATTEMPTS:-3} +WORKERS=${WORKERS:-16} +STAGGER_S=${STAGGER_S:-15} +LOGGING=${LOGGING:-1} +CARLA_VIEW=${CARLA_VIEW:-0} +[ "$CARLA_VIEW" = "1" ] && [ "$LOGGING" != "1" ] && { echo "CARLA_VIEW=1 needs LOGGING=1"; exit 1; } +OBS_HTML=${OBS_HTML:-1} +REPORT=${REPORT:-$LOGGING} +export COSIM_VRU_SIZE_PAD_M=${COSIM_VRU_SIZE_PAD_M:-0.8} +EVALUATOR=$GARAGE_WORK_DIR/leaderboard/leaderboard/leaderboard_evaluator_local.py +CARLA_EGG=$CARLA_ROOT/PythonAPI/carla/dist/carla-0.9.10-py3.7-linux-x86_64.egg +AGENT=$PD/pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py +for path in "$CKPT" "$GARAGE_PY" "$SERVER_PY" "$CARLA_ROOT/CarlaUE4.sh" "$CARLA_EGG" "$EVALUATOR" "$ROUTE_DIR" "$ROUTES_XML" \ + "$SCENARIO_FILE" "$AGENT" "$GARAGE_WORK_DIR/tools/result_parser.py"; do + [ -e "$path" ] || { echo "missing $path"; exit 1; } +done +for path in "$GARAGE_PY" "$SERVER_PY" "$CARLA_ROOT/CarlaUE4.sh"; do + [ -x "$path" ] || { echo "not executable: $path (chmod +x, or the filesystem is mounted noexec)"; exit 1; } +done +echo "Evaluating checkpoint: $CKPT" + +# GPUs of this allocation; the workers share them round-robin +NUM_GPUS=${NUM_GPUS:-$(nvidia-smi -L 2>/dev/null | grep -c "^GPU")} +if [ "$NUM_GPUS" -lt 1 ] || [ "$NUM_GPUS" -gt 8 ]; then echo "NUM_GPUS=$NUM_GPUS must be 1..8"; exit 1; fi +if [ "$WORKERS" -lt 1 ] || [ "$WORKERS" -gt 64 ]; then echo "WORKERS=$WORKERS must be 1..64"; exit 1; fi + +TAG=longest6_v1$([ "$SCENARIOS" = "1" ] || echo "_noscen")$([ "$REPETITIONS" = "1" ] || echo "_rep$REPETITIONS")${SEEDS:+_seeds$(echo "$SEEDS" | tr ' ' '-')} +TAG=$TAG$([ "$COSIM_VRU_SIZE_PAD_M" = "0" ] || echo "_vrupad$COSIM_VRU_SIZE_PAD_M") +# results live in the model's own eval folder, next to the PufferDrive benchmark evals +OUT=${COMPLETE_RUN:-$RUN_DIR/eval/carla_${TAG}_$(date +%Y%m%d_%H%M%S)_${SLURM_JOB_ID:-local}} +mkdir -p "$OUT/routes" "$OUT/aggregate" +: > "$OUT/server_pids" +GIT_STATE=$(git -C "$PD" rev-parse --short HEAD 2>/dev/null)$(git -C "$PD" diff --quiet 2>/dev/null || echo " dirty") +if [ -z "$COMPLETE_RUN" ]; then + { + echo "checkpoint $CKPT"; echo "routes $ROUTE_DIR ($ROUTES_XML)"; echo "scenarios $SCENARIO_FILE" + echo "seeds $SEEDS"; echo "repetitions $REPETITIONS"; echo "gpus $NUM_GPUS"; echo "vru_size_pad_m $COSIM_VRU_SIZE_PAD_M" + echo "git $GIT_STATE" + } > "$OUT/run_info.txt" +else + PREVIOUS_CKPT=$(run_info_value checkpoint) + PREVIOUS_GIT=$(awk '$1 == "git" {$1 = ""; print substr($0, 2); exit}' "$COMPLETE_RUN/run_info.txt") + [ "$(basename "$(dirname "$PREVIOUS_CKPT")")" = "$(basename "$(dirname "$CKPT")")" ] \ + || echo "WARNING: the run was driven with $PREVIOUS_CKPT, this completion uses $CKPT" + [ "$PREVIOUS_GIT" = "$GIT_STATE" ] \ + || echo "WARNING: the run used git '$PREVIOUS_GIT', this checkout is '$GIT_STATE' (PD= keeps the routes comparable)" +fi +echo "Results -> $OUT" + +carla_imports() { PYTHONPATH="$CARLA_EGG" "$GARAGE_PY" -c "import carla" >/dev/null 2>&1; } +# the 0.9.10 egg links libtiff.so.5, which newer distros no longer ship +if ! carla_imports; then + GARAGE_LIB=$(cd "$(dirname "$GARAGE_PY")/../lib" 2>/dev/null && pwd) + if [ -n "$GARAGE_LIB" ] && [ -e "$GARAGE_LIB/libtiff.so.5" ]; then + export LD_LIBRARY_PATH="$GARAGE_LIB${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + echo "libtiff.so.5 from the garage env: $GARAGE_LIB" + fi +fi +if ! carla_imports && [ "${LIBTIFF5_COMPAT:-1}" = "1" ]; then + SYSTEM_LIBTIFF=$( (ldconfig -p 2>/dev/null || /sbin/ldconfig -p 2>/dev/null) | awk '/libtiff.so.6 /{print $NF; exit}') + if [ -n "$SYSTEM_LIBTIFF" ]; then + mkdir -p "$OUT/compat_lib" && ln -sf "$SYSTEM_LIBTIFF" "$OUT/compat_lib/libtiff.so.5" + export LD_LIBRARY_PATH="$OUT/compat_lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + echo "libtiff.so.5 missing: using $SYSTEM_LIBTIFF via $OUT/compat_lib" + fi +fi +carla_imports || { echo "ERROR: $GARAGE_PY cannot import the CARLA egg:"; PYTHONPATH="$CARLA_EGG" "$GARAGE_PY" -c "import carla"; exit 1; } + +# the garage evaluator and the CARLA egg only; the policy server gets PYTHONPATH=$PD from the agent +export PYTHONPATH="$CARLA_ROOT/PythonAPI/carla:$CARLA_EGG:$GARAGE_WORK_DIR/scenario_runner:$GARAGE_WORK_DIR/leaderboard" +export SCENARIO_RUNNER_ROOT=$GARAGE_WORK_DIR/scenario_runner +export LEADERBOARD_ROOT=$GARAGE_WORK_DIR/leaderboard +export BENCHMARK=longest6 DATAGEN=0 +cd "$PD" || exit 1 +OMP_THREADS=${OMP_NUM_THREADS:-$([ "$WORKERS" -ge 16 ] && echo 2 || echo 4)} +export OMP_NUM_THREADS=$OMP_THREADS MKL_NUM_THREADS=$OMP_THREADS OPENBLAS_NUM_THREADS=$OMP_THREADS NUMEXPR_NUM_THREADS=$OMP_THREADS +export COSIM_SERVER_PYTHON=$SERVER_PY +export COSIM_DEVICE=${COSIM_DEVICE:-cpu} +export COSIM_DYNAMICS_SOURCE=pufferdrive +export COSIM_MAX_SPEED_MPS=${COSIM_MAX_SPEED_MPS:-30} +export COSIM_ZERO_PARTNER_STOPPED_TIME=${COSIM_ZERO_PARTNER_STOPPED_TIME:-1} +export COSIM_OBS_HTML_MAX_STEPS=${COSIM_OBS_HTML_MAX_STEPS:-20000} +export COSIM_OBS_HTML_RENDER=0 # routes save the compact replay only; render_carla_obs_html.py renders all pages into $OUT/obs_html + +mkdir -p "$PD/experiments/logs" # the build helper keeps its source hash there +PATH="$(dirname "$SERVER_PY"):$PATH" bash "$PD/scripts/kesai/build_ext_if_changed.sh" "$PD" || exit 1 +IMPORTED_CORE=$(env -u PYTHONPATH PYTHONPATH="$PD" "$SERVER_PY" -c "import pufferlib.ocean.cosim.carla.shadow_ego as s; print(s.__file__)") \ + || { echo "ERROR: shadow env import failed in $SERVER_PY (see traceback above)"; exit 1; } +case "$IMPORTED_CORE" in + "$PD"/*) echo "shadow env: $IMPORTED_CORE" ;; + *) echo "ERROR: shadow env imported from '$IMPORTED_CORE', expected $PD"; exit 1 ;; +esac +"$GARAGE_PY" -c "import carla, leaderboard.leaderboard_evaluator_local" || { echo "ERROR: the garage python cannot import the evaluator (see traceback above)"; exit 1; } + +route_done() { # 0 when the route json holds a record that is not one of the crash states garage resubmits + "$SERVER_PY" - "$1" <<'EOF' +import json, sys +try: + records = json.load(open(sys.argv[1]))["_checkpoint"]["records"] +except Exception: + sys.exit(1) +crashed = ("Failed - Agent couldn't be set up", "Failed - Simulation crashed", "Failed - Agent crashed") +sys.exit(0 if records and all(r["status"] not in crashed for r in records) else 1) +EOF +} + +# route ids = the split files longest_weathers_.xml +mapfile -t ROUTE_IDS < <(ls "$ROUTE_DIR"/longest_weathers_*.xml | sed 's/.*longest_weathers_\([0-9]*\)\.xml/\1/' | sort -n) +TOTAL_ROUTES=${#ROUTE_IDS[@]} +[ -n "$ROUTE_SUBSET" ] && read -r -a ROUTE_IDS <<< "$ROUTE_SUBSET" +# work units = route x seed (":"), each its own server run and result folder +UNITS=() +for seed in "${SEED_LIST[@]}"; do for route_id in "${ROUTE_IDS[@]}"; do UNITS+=("$route_id:$seed"); done; done +TOTAL_UNITS=$((TOTAL_ROUTES * ${#SEED_LIST[@]})) +if [ -n "$COMPLETE_RUN" ]; then + PENDING=() + for unit in "${UNITS[@]}"; do + route_done "$OUT/routes/$(unit_dir "${unit%%:*}" "${unit##*:}")/result.json" || PENDING+=("$unit") + done + [ "${#PENDING[@]}" -gt 0 ] || { echo "nothing to complete: every route run of $OUT has a valid record"; exit 0; } + echo "completing $OUT: ${#PENDING[@]} of $TOTAL_UNITS route runs without a valid record: ${PENDING[*]}" + UNITS=("${PENDING[@]}") + # what the completion replaces stays readable next to the live results + COMPLETION_STAMP=$(date +%Y%m%d_%H%M%S) + PREVIOUS=$OUT/pre_completion_$COMPLETION_STAMP + mkdir -p "$PREVIOUS/routes" + for unit in "${UNITS[@]}"; do + route_name=$(unit_dir "${unit%%:*}" "${unit##*:}") + [ -d "$OUT/routes/$route_name" ] && mv "$OUT/routes/$route_name" "$PREVIOUS/routes/" + rm -f "$(unit_json "${unit%%:*}" "${unit##*:}")" + done + for f in "$OUT"/worker_w*.log; do [ -e "$f" ] && mv "$f" "$PREVIOUS/"; done + echo "completion $COMPLETION_STAMP on $(hostname): route runs ${UNITS[*]}; checkpoint $CKPT; scenarios $SCENARIO_FILE; seeds $SEEDS; repetitions $REPETITIONS; vru_size_pad_m $COSIM_VRU_SIZE_PAD_M; gpus $NUM_GPUS; git $GIT_STATE" >> "$OUT/run_info.txt" +fi +[ "$WORKERS" -gt "${#UNITS[@]}" ] && WORKERS=${#UNITS[@]} +# longest route first to the least loaded worker; length = the xml's waypoint polyline (an earlier run's duration would +# read a route that failed early as short) +mapfile -t ASSIGNMENT < <("$SERVER_PY" - "$ROUTE_DIR" "$WORKERS" "${UNITS[@]}" <<'EOF' +import math, sys, xml.etree.ElementTree as ET +route_dir, workers, units = sys.argv[1], int(sys.argv[2]), sys.argv[3:] # unit = ":" +SETUP_M = 300.0 # server start + world load per route, in route-length-equivalent metres +def length(rid): + pts = [(float(w.get("x")), float(w.get("y"))) for w in ET.parse(f"{route_dir}/longest_weathers_{rid}.xml").iter("waypoint")] + return sum(math.dist(a, b) for a, b in zip(pts, pts[1:])) +lengths = {u.split(":")[0]: length(int(u.split(":")[0])) for u in units} +loads = [0.0] * workers; assign = [[] for _ in range(workers)] +for unit in sorted(units, key=lambda u: -lengths[u.split(":")[0]]): + w = min(range(workers), key=lambda i: loads[i]); assign[w].append(unit); loads[w] += lengths[unit.split(":")[0]] + SETUP_M +for w in range(workers): + print(w, " ".join(assign[w]), f"# {loads[w] / 1000:.1f} km") +EOF +) +echo "${#ROUTE_IDS[@]} routes x ${#SEED_LIST[@]} seeds = ${#UNITS[@]} route runs over $WORKERS workers on $NUM_GPUS GPUs, policy server threads $OMP_THREADS" +printf " worker %s\n" "${ASSIGNMENT[@]}" +[ -f "$OUT/timing.csv" ] || echo "route,worker,gpu,attempt,start,wall_s,game_s,pace,status,driving_score,seed" > "$OUT/timing.csv" +TIMING_ROWS_BEFORE=$(($(wc -l < "$OUT/timing.csv") - 1)) + +log_route_timing() { # $1 worker, $2 gpu, $3 route id, $4 attempt, $5 start epoch, $6 result json, $7 seed + "$SERVER_PY" - "$@" "$(date +%s)" "$OUT/timing.csv" <<'EOF' +import json, sys, time +w, gpu, route_id, attempt, t0, path, seed, t1, csv_path = sys.argv[1:10] +rec = json.load(open(path))["_checkpoint"]["records"][0] +wall = int(t1) - int(t0); game = float(rec["meta"]["duration_game"]); ds = float(rec["scores"]["score_composed"]) +pace = game / wall if wall > 0 else 0.0 +with open(csv_path, "a") as f: + f.write(f"{route_id},{w},{gpu},{attempt},{time.strftime('%H:%M:%S', time.localtime(int(t0)))},{wall},{game:.1f},{pace:.3f},{rec['status']},{ds:.2f},{seed}\n") +print(f"[w{w} gpu{gpu} route{route_id}{' s' + seed if seed else ''}] done: {wall} s wall, {game:.0f} s game, {pace:.2f}x real time, DS {ds:.1f}, {rec['status']}") +EOF +} + +run_route() { # $1 worker, $2 gpu, $3 route id, $4 attempt, $5 seed: one CARLA server + one evaluator run, into $OUT/routes/route__s + local w=$1 gpu=$2 route_id=$3 attempt=$4 seed=$5 + local route_dir=$OUT/routes/$(unit_dir "$route_id" "$seed") + local seed_args=() + [ -n "$seed" ] && seed_args=(--trafficManagerSeed "$seed") # leaderboard 1.0 spelling; empty seed = the evaluator's default (runs before the seeds layout) + local routes_xml=$ROUTE_DIR/longest_weathers_$route_id.xml + local port=$((PORT_BASE + w * 50)) tm_port=$((PORT_BASE + 6000 + w * 50)) + local server_log=$route_dir/carla_server_attempt$attempt.log evaluator_log=$route_dir/evaluator_attempt$attempt.log + local t_start + t_start=$(date +%s) + mkdir -p "$route_dir" + rm -f "$route_dir/result.json" + local pin + pin=$(numa_pin_for_gpu "$gpu") + # offscreen OpenGL on the worker's GPU (UE 4.24 has no -RenderOffScreen/-nullrhi/-graphicsadapter) + $pin env DISPLAY= SDL_VIDEODRIVER=offscreen SDL_HINT_CUDA_DEVICE="$gpu" "$CARLA_ROOT/CarlaUE4.sh" -opengl -nosound \ + -carla-rpc-port="$port" -carla-streaming-port=$((port + 1)) > "$server_log" 2>&1 & + local server_pid=$! up=1 + echo "$server_pid" >> "$OUT/server_pids" + for _ in $(seq 1 60); do + kill -0 "$server_pid" 2>/dev/null || { echo "[w$w gpu$gpu route$route_id${seed:+ s$seed}] CARLA server died"; break; } + "$GARAGE_PY" -c " +import carla, sys +try: + c = carla.Client('localhost', $port); c.set_timeout(10.0); c.get_world().get_map() +except Exception: + sys.exit(1)" 2>/dev/null && up=0 && break + sleep 5 + done + if [ $up -ne 0 ]; then + pkill -9 -P "$server_pid" 2>/dev/null; kill -9 "$server_pid" 2>/dev/null + echo "[w$w gpu$gpu route$route_id${seed:+ s$seed}] last lines of $server_log:"; tail -n 20 "$server_log" + return 1 + fi + local log_env=(ROUTES="$routes_xml") + if [ "$LOGGING" = "1" ]; then + log_env+=(COSIM_TELEMETRY="$route_dir/telemetry" COSIM_WORLD_LOG="$route_dir/world_log") + [ "$CARLA_VIEW" = "1" ] && log_env+=(COSIM_DEBUG_CARLA_VIEW="$route_dir/carla_view") + [ "$OBS_HTML" = "1" ] && log_env+=(COSIM_OBS_HTML="$route_dir/obs_html") + fi + $pin env CUDA_VISIBLE_DEVICES="$gpu" "${log_env[@]}" \ + "$GARAGE_PY" -u "$EVALUATOR" --routes "$routes_xml" --scenarios "$SCENARIO_FILE" --repetitions "$REPETITIONS" \ + --agent "$AGENT" --agent-config "$CKPT" --checkpoint "$route_dir/result.json" --track MAP \ + --port "$port" --trafficManagerPort "$tm_port" --timeout 600 "${seed_args[@]}" \ + > "$evaluator_log" 2>&1 & + local evaluator_pid=$! + wait -n 2>/dev/null # the evaluator, or the server when it dies first (render-thread crash a minute in) + if kill -0 "$evaluator_pid" 2>/dev/null; then + echo "[w$w gpu$gpu route$route_id${seed:+ s$seed}] CARLA server died mid-route; stopping the evaluator (600 s client timeouts otherwise)" + pkill -9 -P "$evaluator_pid" 2>/dev/null; kill -9 "$evaluator_pid" 2>/dev/null; wait "$evaluator_pid" 2>/dev/null + fi + # children first: killing the wrapper shell first reparents the UE4 binary and leaves it running + pkill -9 -P "$server_pid" 2>/dev/null; kill -9 "$server_pid" 2>/dev/null; wait "$server_pid" 2>/dev/null + if ! route_done "$route_dir/result.json"; then + echo "[w$w gpu$gpu route$route_id${seed:+ s$seed}] no valid route record; last lines of $evaluator_log:"; tail -n 20 "$evaluator_log" + return 1 + fi + log_route_timing "$w" "$gpu" "$route_id" "$attempt" "$t_start" "$route_dir/result.json" "$seed" + return 0 +} + +run_worker() { # $1 worker: its assigned route runs, sequentially, each retried up to MAX_ATTEMPTS on the next GPU + local w=$1 units unit route_id seed attempt gpu + units=${ASSIGNMENT[w]#* } + units=${units%%#*} + for unit in $units; do + route_id=${unit%%:*}; seed=${unit##*:} + for ((attempt = 1; attempt <= MAX_ATTEMPTS; attempt++)); do + gpu=$(( (w + attempt - 1) % NUM_GPUS )) + echo "[w$w] route $route_id${seed:+ seed $seed} attempt $attempt gpu $gpu $(date +%H:%M:%S)" + run_route "$w" "$gpu" "$route_id" "$attempt" "$seed" && break + echo "[w$w] route $route_id${seed:+ seed $seed} attempt $attempt failed" + done + done +} + +# never attach to a server that is not ours (a peer's test server, a leftover of an earlier job) +if command -v ss >/dev/null 2>&1; then + for ((w = 0; w < WORKERS; w++)); do + for port in $((PORT_BASE + w * 50)) $((PORT_BASE + w * 50 + 1)) $((PORT_BASE + 6000 + w * 50)); do + [ -z "$(ss -Hltn "sport = :$port")" ] || { echo "port $port is already in use (ss -ltn); pick another PORT_BASE"; exit 1; } + done + done +fi +# only the servers this job started (children first: the UE4 binary outlives its wrapper shell otherwise) +trap 'for p in $(cat "$OUT/server_pids" 2>/dev/null); do pkill -9 -P "$p" 2>/dev/null; kill -9 "$p" 2>/dev/null; done' EXIT +WORKER_PIDS=() +ROUTES_T0=$(date +%s) +for ((w = 0; w < WORKERS; w++)); do + [ "$w" -gt 0 ] && sleep "$STAGGER_S" # simultaneous engine starts are the one thing the crashed runs shared + run_worker "$w" > "$OUT/worker_w$w.log" 2>&1 & + WORKER_PIDS+=($!) +done +wait "${WORKER_PIDS[@]}" +ROUTES_WALL_S=$(( $(date +%s) - ROUTES_T0 )) +cat "$OUT"/worker_w*.log +"$SERVER_PY" - "$OUT/timing.csv" "$ROUTES_WALL_S" "$WORKERS" "$NUM_GPUS" "${#UNITS[@]}" "$TIMING_ROWS_BEFORE" <<'EOF' | tee -a "$OUT/run_info.txt" +import csv, sys +rows = list(csv.DictReader(open(sys.argv[1]))); wall, workers, gpus, n_routes, rows_before = (int(x) for x in sys.argv[2:7]) +rows = rows[rows_before:] # a completion appends to the previous run's timing.csv +if not rows: + print(f"speed: no route finished in {wall / 3600:.2f} h wall ({workers} workers on {gpus} GPUs)"); sys.exit(0) +game = sum(float(r["game_s"]) for r in rows); route_wall = sum(int(r["wall_s"]) for r in rows) +per_worker = {} +for r in rows: + per_worker[r["worker"]] = per_worker.get(r["worker"], 0) + int(r["wall_s"]) +print(f"speed: {len(rows)}/{n_routes} route runs in {wall / 3600:.2f} h wall, {workers} workers on {gpus} GPUs; game time {game / 3600:.2f} h; " + f"per-route pace {game / route_wall:.2f}x real time; busiest worker {max(per_worker.values()) / 3600:.2f} h, " + f"idlest {min(per_worker.values()) / 3600:.2f} h; route runs/h {len(rows) / (wall / 3600):.1f}") +EOF + +# aggregate: garage's parser over the per-route jsons, one subfolder per seed (the parser walks $OUT/aggregate and +# reports mean/std over the seeds per route and town) -> $OUT/aggregate/results.csv, plus a per-route table; a crash +# record has no global_record and would abort the parser +for f in "$OUT"/routes/route_*/result.json; do + name=$(basename "$(dirname "$f")") + case $name in *_s*) seed=${name##*_s}; name=${name%_s*} ;; *) seed= ;; esac + route_id=$(echo "$name" | sed 's/route_0*//; s/^$/0/') + dest=$(unit_json "$route_id" "$seed"); mkdir -p "$(dirname "$dest")" + route_done "$f" && cp "$f" "$dest" && continue + echo "route $route_id${seed:+ seed $seed} has no valid record (all attempts crashed): left out of results.csv" +done +# garage's parser imports torch/pygame and carla: the garage python gets the egg; the server python (the fallback where +# the garage env's torch does not load) has its own carla and must not see the py3.7 egg +PARSER_ARGS=(--xml "$ROUTES_XML" --results "$OUT/aggregate" --log_dir "$OUT/aggregate") +env PYTHONPATH="$CARLA_EGG" "$GARAGE_PY" "$GARAGE_WORK_DIR/tools/result_parser.py" "${PARSER_ARGS[@]}" \ + || env -u PYTHONPATH "$SERVER_PY" "$GARAGE_WORK_DIR/tools/result_parser.py" "${PARSER_ARGS[@]}" \ + || echo "result_parser.py failed with $GARAGE_PY and $SERVER_PY (per-route jsons in $OUT/aggregate are unaffected)" +"$SERVER_PY" - "$OUT/aggregate" "$TOTAL_UNITS" "**/longest_weathers_*.json" <<'EOF' +import glob, json, os, re, statistics, sys +rows = [] # (seed, route_id, status, DS, RC, IP); seed "" = a run from before the seeds layout +for f in glob.glob(os.path.join(sys.argv[1], sys.argv[3]), recursive=True): + parent = os.path.basename(os.path.dirname(f)) + seed = parent[1:] if parent.startswith("s") and parent[1:].isdigit() else "" + for r in json.load(open(f))["_checkpoint"]["records"]: + rows.append((seed, r["route_id"], r["status"], r["scores"]["score_composed"], r["scores"]["score_route"], r["scores"]["score_penalty"])) +def route_index(route_id): + m = re.search(r"_(\d+)", route_id) + return int(m.group(1)) if m else 0 +rows.sort(key=lambda r: (int(r[0]) if r[0] else -1, route_index(r[1]))) +print(f"\n{len(rows)} route records of {sys.argv[2]} route runs") +for seed, route_id, status, ds, rc, ip in rows: + print(f" {('s' + seed) if seed else '-':>3s} {route_id:22s} {status:34s} DS {ds:6.2f} RC {rc:6.2f} IP {ip:5.3f}") +if rows: + n = len(rows) + print(f"mean DS {sum(r[3] for r in rows)/n:.2f} RC {sum(r[4] for r in rows)/n:.2f} IP {sum(r[5] for r in rows)/n:.3f} (all {n} route runs)") + by_seed = {} + for r in rows: + by_seed.setdefault(r[0], []).append(r) + if len(by_seed) > 1: # the paper-style number: the benchmark DS per seed, mean +- std over the seeds + ds = {s: sum(r[3] for r in v) / len(v) for s, v in by_seed.items()} + rc = {s: sum(r[4] for r in v) / len(v) for s, v in by_seed.items()} + for s in sorted(by_seed, key=lambda s: int(s) if s else -1): + print(f"seed {s}: DS {ds[s]:.2f} RC {rc[s]:.2f} ({len(by_seed[s])} routes)") + print(f"over seeds: DS {statistics.mean(ds.values()):.2f} +- {statistics.pstdev(ds.values()):.2f} " + f"RC {statistics.mean(rc.values()):.2f} +- {statistics.pstdev(rc.values()):.2f} (population std over {len(by_seed)} seeds, as in results.csv)") +EOF + +REPORT_ARGS=() +[ "$OBS_HTML" = "1" ] && REPORT_ARGS+=(--obs-html-dir "$OUT/obs_html") +if [ "$REPORT" = "1" ] && [ "$LOGGING" = "1" ]; then + env -u PYTHONPATH "$SERVER_PY" "$PD/scripts/eval/analyze_carla_cosim.py" "$OUT"/routes/route_* "$OUT/report" "${REPORT_ARGS[@]}" \ + || echo "HTML report failed (results above are unaffected)" +fi +# after the report: the gallery links to it when it exists +if [ "$OBS_HTML" = "1" ] && [ "$LOGGING" = "1" ]; then + env -u PYTHONPATH "$SERVER_PY" "$PD/scripts/eval/render_carla_obs_html.py" "$OUT" --routes "$ROUTES_XML" \ + || echo "obs replay gallery failed (results above are unaffected)" +fi +[ -f "$OUT/report/index.html" ] && echo "report -> $OUT/report/index.html" +[ -f "$OUT/obs_html/index.html" ] && echo "obs replays -> $OUT/obs_html/index.html" +echo "Done -> $OUT" diff --git a/scripts/kesai/15_tune_fused_slot_encoder_slurm.sh b/scripts/kesai/15_tune_fused_slot_encoder_slurm.sh new file mode 100755 index 0000000000..ce01940983 --- /dev/null +++ b/scripts/kesai/15_tune_fused_slot_encoder_slurm.sh @@ -0,0 +1,50 @@ +#!/bin/bash +#SBATCH --job-name tune_fused_encoder +#SBATCH --ntasks 1 +#SBATCH --nodes 1 +#SBATCH --time 3:00:00 +#SBATCH --gres gpu:1 +#SBATCH --gres-flags=enforce-binding # CPUs on the GPU's socket +#SBATCH --mem=125G +#SBATCH --cpus-per-task 18 +#SBATCH --output experiments/logs/tune_fused_%j.out +#SBATCH --error experiments/logs/tune_fused_%j.err +#SBATCH --partition dev + +source scripts/kesai/cluster_env.sh || exit 1 + +# Times the fused slot-encoder kernels on this node's GPU and prints the KernelConfig to paste into +# CONFIGS_BY_CAPABILITY in pufferlib/ocean/fused_slot_encoder.py. Submit once per GPU type (H100, B200, ...); +# pick the node type with --constraint or --gres gpu::1 if the partition mixes GPUs. + +echo "START TIME: $(date)" +start=$(date +%s) +cd $PD || exit 1 +nvidia-smi --query-gpu=name,compute_cap,memory.total --format=csv + +export OMP_NUM_THREADS=1 +source .venv/bin/activate +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 + +RESULT_FILE=$PD/experiments/logs/tune_fused_${SLURM_JOB_ID}_result.txt +CONFIG_JSON=$PD/experiments/logs/tune_fused_${SLURM_JOB_ID}_config.json +python scripts/tune_fused_slot_encoder.py --batch 131072 --shapes all --write-config ${CONFIG_JSON} 2>&1 | tee ${RESULT_FILE}.full +# The pasteable block is the tail of the full log. +sed -n '/=== fused_slot_encoder tuning result/,/=== end of tuning result ===/p' ${RESULT_FILE}.full > ${RESULT_FILE} + +# Real policy update + rollout step on this GPU: compiled reference vs fused kernels with the config just tuned. +echo "=== step timing (paste this block back too) ===" | tee -a ${RESULT_FILE} +python scripts/profile_policy_step.py --rows 131072 2>&1 | grep "^RESULT" | tee -a ${RESULT_FILE} +python scripts/profile_policy_step.py --rows 131072 --fused --kernel-config ${CONFIG_JSON} 2>&1 | grep "^RESULT" | tee -a ${RESULT_FILE} +awk '/^RESULT fused=0/ {for (i=1;i<=NF;i++) if ($i ~ /^update_ms=/) {split($i,a,"="); base=a[2]}} + /^RESULT fused=1/ {for (i=1;i<=NF;i++) if ($i ~ /^update_ms=/) {split($i,a,"="); fused=a[2]}} + END {if (base && fused) printf "update speedup fused vs reference: %.2fx\n", base / fused}' ${RESULT_FILE} | tee -a ${RESULT_FILE} +echo "=== end of step timing ===" | tee -a ${RESULT_FILE} + +echo +echo "Result block (also in ${RESULT_FILE}):" +cat ${RESULT_FILE} + +end=$(date +%s) +echo "END TIME: $(date)" +echo "Runtime: $((end-start)) s" diff --git a/scripts/kesai/16_mfu_bench_slurm.sh b/scripts/kesai/16_mfu_bench_slurm.sh new file mode 100755 index 0000000000..b4a8671ef7 --- /dev/null +++ b/scripts/kesai/16_mfu_bench_slurm.sh @@ -0,0 +1,70 @@ +#!/bin/bash +#SBATCH --job-name mfu_bench +#SBATCH --ntasks 1 +#SBATCH --nodes 1 +#SBATCH --time 1:00:00 +#SBATCH --gres gpu:1 +#SBATCH --gres-flags=enforce-binding # CPUs on the GPU's socket +#SBATCH --mem=125G +#SBATCH --cpus-per-task 18 +#SBATCH --output experiments/logs/mfu_bench_%j.out +#SBATCH --error experiments/logs/mfu_bench_%j.err +#SBATCH --partition dev + +source scripts/kesai/cluster_env.sh || exit 1 + +# Measures MFU of the real PPO loop (rollout, update, whole epoch) and SPS on this node's GPU: reference encoders, +# policy.fused_slot_encoder with the Triton kernels, and the same with the CUDA kernels (PUFFER_SLOT_ENCODER_BACKEND=cuda), +# and adds hardware-counter utilization of one step if Nsight Compute is available. +# 32 envs on 18 threads, like the multinode training script (blocking workers allow 2 per hardware thread). + +echo "START TIME: $(date)" +start=$(date +%s) +cd $PD || exit 1 +nvidia-smi --query-gpu=name,compute_cap,memory.total --format=csv +echo "node $(hostname)"; nvidia-smi --query-gpu=clocks.sm,clocks.max.sm,clocks.mem,power.draw,power.limit,temperature.gpu,clocks_throttle_reasons.active --format=csv + +export OMP_NUM_THREADS=1 +source .venv/bin/activate +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 + +RESULT_FILE=$PD/experiments/logs/mfu_bench_${SLURM_JOB_ID}_result.txt +NUM_ENVS=${NUM_ENVS:-32} # override: NUM_ENVS=48 sbatch scripts/kesai/16_mfu_bench_slurm.sh +MINIBATCH=${MINIBATCH:-131072} +EXTRA_ARGS=${EXTRA_ARGS:-} # e.g. EXTRA_ARGS="--kernel-table train.compile_mode=max-autotune" +echo "=== MFU result (paste this block back) ===" | tee ${RESULT_FILE} +python scripts/mfu_bench.py --num-envs ${NUM_ENVS} --minibatch ${MINIBATCH} ${EXTRA_ARGS} 2>&1 | grep -v "^MFU_RESULT" | grep "reference\|capability\|kernel table\|gemm shapes\|measured bf16\|reduced-precision\|^ *[0-9][0-9.]*%\|^ *share" | tee -a ${RESULT_FILE} +python scripts/mfu_bench.py --num-envs ${NUM_ENVS} --minibatch ${MINIBATCH} --fused ${EXTRA_ARGS} 2>&1 | grep -v "^MFU_RESULT" | grep "fused\|kernel table\|gemm shapes\|measured bf16\|reduced-precision\|^ *[0-9][0-9.]*%\|^ *share" | sed 's/^fused/fused+triton/' | tee -a ${RESULT_FILE} +env PUFFER_SLOT_ENCODER_BACKEND=cuda python scripts/mfu_bench.py --num-envs ${NUM_ENVS} --minibatch ${MINIBATCH} --fused ${EXTRA_ARGS} 2>&1 | grep -v "^MFU_RESULT" | grep "fused\|Error\|kernel table\|gemm shapes\|measured bf16\|reduced-precision\|^ *[0-9][0-9.]*%\|^ *share" | sed 's/^fused/fused+cuda/' | tee -a ${RESULT_FILE} +echo "--- rollout split, fused+triton (env = waiting for the simulation workers, the rest is the trainer process) ---" | tee -a ${RESULT_FILE} +python scripts/mfu_bench.py --num-envs ${NUM_ENVS} --minibatch ${MINIBATCH} --fused --rollout-only --profile --warmup 1 --epochs 2 ${EXTRA_ARGS} 2>&1 | grep -A5 "section split" | tee -a ${RESULT_FILE} +echo "--- rollout split, fused+cuda ---" | tee -a ${RESULT_FILE} +env PUFFER_SLOT_ENCODER_BACKEND=cuda python scripts/mfu_bench.py --num-envs ${NUM_ENVS} --minibatch ${MINIBATCH} --fused --rollout-only --profile --warmup 1 --epochs 2 ${EXTRA_ARGS} 2>&1 | grep -A5 "section split" | tee -a ${RESULT_FILE} + +NCU=$(command -v ncu || ls /opt/nvidia/nsight-compute/*/ncu 2>/dev/null | tail -1) +if [ -n "${NCU}" ]; then + echo "--- hardware counters (ncu range replay, one update minibatch of ${MINIBATCH} rows + one rollout forward) ---" | tee -a ${RESULT_FILE} + METRICS=gpu__time_duration.sum,sm__ops_path_tensor_op_hmma_src_bf16_dst_fp32.sum,dram__bytes.sum + for variant in reference fused fused_cuda; do + FLAG=""; [ ${variant} != reference ] && FLAG="--fused" + BACKEND=triton; [ ${variant} = fused_cuda ] && BACKEND=cuda + CSV=$PD/experiments/logs/mfu_bench_${SLURM_JOB_ID}_${variant}.csv + ( export PUFFER_SLOT_ENCODER_BACKEND=${BACKEND}; timeout 600 ${NCU} --replay-mode range --target-processes application-only --clock-control none --nvtx \ + --nvtx-include "update" --nvtx-include "rollout" --nvtx-include "calibrate" --print-nvtx-rename kernel \ + --metrics ${METRICS} --csv --page raw --print-units base \ + python scripts/profile_policy_step.py --rows ${MINIBATCH} ${FLAG} > ${CSV} 2>&1 ) \ + && { echo "${variant}:" | tee -a ${RESULT_FILE}; python scripts/profile_policy_step.py --parse ${CSV} --rows ${MINIBATCH} | tee -a ${RESULT_FILE}; } \ + || echo "${variant}: ncu failed (counters restricted or unsupported), see ${CSV}" | tee -a ${RESULT_FILE} + done +else + echo "Nsight Compute not found on this node; hardware-counter utilization skipped." | tee -a ${RESULT_FILE} +fi +echo "=== end of MFU result ===" | tee -a ${RESULT_FILE} + +echo +echo "Result block (also in ${RESULT_FILE}):" +cat ${RESULT_FILE} +end=$(date +%s) +echo "node $(hostname)"; nvidia-smi --query-gpu=clocks.sm,clocks.max.sm,clocks.mem,power.draw,power.limit,temperature.gpu,clocks_throttle_reasons.active --format=csv +echo "END TIME: $(date)" +echo "Runtime: $((end-start)) s" diff --git a/scripts/kesai/17_carla_longest6_v1_complete_local.sh b/scripts/kesai/17_carla_longest6_v1_complete_local.sh new file mode 100755 index 0000000000..1fd936a55d --- /dev/null +++ b/scripts/kesai/17_carla_longest6_v1_complete_local.sh @@ -0,0 +1,44 @@ +#!/bin/bash +# Complete a longest6 v1 run that lost routes on the cluster (CARLA 0.9.10 render-thread crashes, the job time limit) on +# this desktop: scripts/kesai/14_carla_longest6_v1.sh in its COMPLETE_RUN mode with the local CARLA 0.9.10, carla_garage, +# the `garage` conda env and the repo .venv as policy server. Sync the run folder from the cluster as is (it lives in the +# model's eval/ folder); only its routes without a valid record are driven, in place, WORKERS servers in parallel on the +# one GPU, then results.csv, the HTML report and the obs replay gallery are rebuilt over all routes. Nothing from the +# cluster is overwritten: the replaced route folders and worker logs move to /pre_completion_/. +# +# Usage: bash scripts/kesai/17_carla_longest6_v1_complete_local.sh experiments/k_scaled_0050_1000/eval/carla_longest6_v1_vrupad0.8__ +# +# Overridable env: CKPT (default: the model folder the run sits in, else the model named in its run_info.txt under +# $PD/experiments), WORKERS (8: ~2 cores and 2.5-5 GB each, the GPU stays idle), PORT_BASE (2200; never 2000, the +# interactive server), STAGGER_S, ROUTE_SUBSET, and every other variable of script 14 (SCENARIOS, SEEDS, REPETITIONS and +# COSIM_VRU_SIZE_PAD_M default to the run's own values; a run from before the seeds layout keeps its single seed). Expect ~1 h of wall time per hour of game time per worker. +set -u + +RUN=${1:?usage: $0 } +RUN=$(cd "$RUN" 2>/dev/null && pwd) || { echo "no such folder: $1"; exit 1; } +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +export PD=${PD:-$(cd "$SCRIPT_DIR/../.." && pwd)} +export GARAGE_PY=${GARAGE_PY:-$HOME/miniconda3/envs/garage/bin/python} +export SERVER_PY=${SERVER_PY:-$PD/.venv/bin/python} +export CARLA_ROOT=${CARLA_ROOT:-$HOME/ordnung/internal/CARLA_0.9.10} +export GARAGE_WORK_DIR=${GARAGE_WORK_DIR:-$HOME/ordnung/internal/carla_garage_1} +export WORKERS=${WORKERS:-8} PORT_BASE=${PORT_BASE:-2200} STAGGER_S=${STAGGER_S:-20} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-2} +# garage's torch 1.12 needs an executable stack that glibc 2.41 refuses by default (result_parser.py) +export GLIBC_TUNABLES=${GLIBC_TUNABLES:-glibc.rtld.execstack=2} +unset WAYLAND_DISPLAY # DISPLAY= alone still opens one window per server on this desktop + +if [ -z "${CKPT:-}" ]; then + MODEL_DIR=$(cd "$RUN/../.." && pwd) + if [ ! -f "$MODEL_DIR/final_model.pt" ]; then + PREVIOUS_CKPT=$(awk '$1 == "checkpoint" {print $2; exit}' "$RUN/run_info.txt" 2>/dev/null) + MODEL_DIR=$PD/experiments/$(basename "$(dirname "$PREVIOUS_CKPT")") + fi + CKPT=$MODEL_DIR/final_model.pt +fi +export CKPT +[ -f "$CKPT" ] || { echo "checkpoint not found: $CKPT (set CKPT=/path/to/final_model.pt)"; exit 1; } +[ -f "$(dirname "$CKPT")/config.yaml" ] || { echo "no config.yaml next to $CKPT"; exit 1; } + +export COMPLETE_RUN=$RUN +exec bash "$SCRIPT_DIR/14_carla_longest6_v1.sh" diff --git a/scripts/kesai/18_train_16GPU_gigaflow_slurm.sh b/scripts/kesai/18_train_16GPU_gigaflow_slurm.sh new file mode 100755 index 0000000000..aae73a75fa --- /dev/null +++ b/scripts/kesai/18_train_16GPU_gigaflow_slurm.sh @@ -0,0 +1,152 @@ +#!/bin/bash +#SBATCH --job-name train_puffer +#SBATCH --nodes 1 # Number of nodes requested +#SBATCH --ntasks-per-node 1 # Run 1 srun task per node (which fires up torchrun) +#SBATCH --gres gpu:8 # GPUs per node +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --mem=0 +#SBATCH --time 6-00:00 +#SBATCH --output experiments/logs/train_gigaflow_%j.out +#SBATCH --error experiments/logs/train_gigaflow_%j.err +#SBATCH --partition dev + +source scripts/kesai/cluster_env.sh || exit 1 + +# Set up PyTorch Distributed Rendezvous parameters from Slurm variables +export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) +export MASTER_PORT=29500 +export WORLD_SIZE=$SLURM_JOB_NUM_NODES + +echo "START TIME: $(date)" +echo "Master node: ${MASTER_ADDR}:${MASTER_PORT}" +echo "Total nodes: ${WORLD_SIZE}" +start=$(date +%s) + +export SEED=1000 + +export RUN_NAME=${RUN_NAME:-k_scaled_0060_${SEED}} +echo ${RUN_NAME} + +export DATA_DIR=$PD/experiments/${RUN_NAME} +echo ${DATA_DIR} + +# Name the training run writes its final model under, so the eval steps below can +# find it without knowing the last epoch number. Passed to train as +# train.final_model_name so the two can never disagree. +export FINAL_MODEL_NAME=final_model.pt +export MODEL_PATH=${DATA_DIR}/${FINAL_MODEL_NAME} +echo ${MODEL_PATH} + +# Thread limit limits CPU thrashing across worker environments +export NUMEXPR_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export VECLIB_MAXIMUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export OMP_NUM_THREADS=1 + +# Hand-written CUDA slot-encoder kernels instead of the Triton ones (same math); triton to compare. Passed per command, +# never exported: sbatch copies this shell's environment into the eval jobs below, whose interpreters lack cuda-bindings. +SLOT_ENCODER_BACKEND=${SLOT_ENCODER_BACKEND:-cuda} +echo "slot encoder backend: ${SLOT_ENCODER_BACKEND}" + +source .venv/bin/activate +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 +# GigaFlow (arXiv 2502.03349, Table A3) match on one node: 8 x 32768 = 262k rows per gradient update (paper: 8 x 32000), +# 3 PPO epochs, 1200-step episodes, cosine LR 5e-4. One node because the gradient all-reduce over the 10 GbE link costs +# ~35 ms per update whatever the node count (run 0055, 2 nodes: 66 ms per 32k update vs ~32 ms of compute); inside a node it +# goes over PCIe. vec.num_envs=60 is one blocking worker per kesai thread; env.num_agents=4096 per worker gives 31.5 M +# transitions per GPU per rollout (paper: ~41.6 M) and a 52 GB float16 obs buffer (60 GB total per GPU). +# Minibatch ablation 2026-10-27: k_scaled_0057 (MINIBATCH=32768) vs k_scaled_0058 (131072, old kernels throughout); 0058 +# led on every per-step metric, so 131072 is the default. k_scaled_0060 = the 0058 config on the 2026-10-27 fused-kernel +# and NUMA-bind code, the speed/quality A/B against 0058: +# sbatch scripts/kesai/18_train_16GPU_gigaflow_slurm.sh +# Same job on the scaleway H100 nodes (18 threads per rank, the 60 blocking workers oversubscribe the CPUs 3.3x): +# RUN_NAME=k_scaled_0061_1000 sbatch scripts/kesai/18_train_16GPU_gigaflow_slurm.sh # NUM_ENVS=32 fits the CPUs but shrinks the rollout batch +MINIBATCH=${MINIBATCH:-131072} +echo "minibatch per GPU: ${MINIBATCH}" +# Execute torchrun across all nodes using srun +PUFFER_SLOT_ENCODER_BACKEND=${SLOT_ENCODER_BACKEND} srun --cpus-per-task="$SLURM_CPUS_PER_TASK" torchrun \ + --nnodes=${SLURM_JOB_NUM_NODES} \ + --nproc-per-node=8 \ + --rdzv_id=${SLURM_JOB_ID} \ + --rdzv_backend=c10d \ + --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \ + --max_restarts=0 \ + --start-method spawn \ + -m pufferlib.pufferl train puffer_drive \ + wandb=True \ + wandb_project=nightly-multi-long \ + wandb_group=emerge_ \ + train.data_dir=${DATA_DIR} \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla_128_affine_zones_stop_ndshoulders_v19 \ + env.num_maps=128 \ + env.obs_stride_merge_segments=true \ + env.partner_blindness_corridor=false \ + env.partner_blindness_prob=0.05 \ + env.partner_blindness_prob_min=0.05 \ + env.partner_blindness_trigger_prob=1.0 \ + train.name=${RUN_NAME} \ + run_name=${RUN_NAME} \ + train.total_timesteps=1000000000000 \ + vec.num_envs=${NUM_ENVS:-60} \ + env.num_agents=4096 \ + env.scenario_length=1200 \ + train.rollout_dtype=float16 \ + env.preload_map_cache=True \ + train.max_minibatch_size=${MINIBATCH} \ + train.minibatch_size=${MINIBATCH} \ + train.update_epochs=3 \ + policy.fused_slot_encoder=true \ + train.final_model_name=${FINAL_MODEL_NAME} \ + train.seed=${SEED} \ + tb=True + +if [ ! -f ${MODEL_PATH} ]; then + echo "Training did not produce ${MODEL_PATH}; skipping evaluation." + exit 1 +fi + +# nuPlan (reactive and non-reactive), longest6 (v2 and v1) and AlpaSim evals need one 8-GPU node each: submit them as their own jobs before the long self-play eval below so they queue while it runs. +echo "Training done, submitting nuPlan reactive eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh \ + || echo "nuPlan eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh" +echo "Submitting nuPlan non-reactive eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh \ + || echo "nuPlan non-reactive eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh" +echo "Submitting longest6 v2 eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6_v2.sh \ + || echo "longest6 v2 eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6_v2.sh" +echo "Submitting longest6 v1 eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/14_carla_longest6_v1.sh \ + || echo "longest6 v1 eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/14_carla_longest6_v1.sh" +echo "Submitting AlpaSim eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/12_alpasim.sh \ + || echo "AlpaSim eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/12_alpasim.sh" + +# parallel_eval places one shard per allocated node via srun, so each shard's 128 +# env workers get a full node's cores instead of sharing the batch host. +echo "Follow-up evals submitted, evaluating ${MODEL_PATH}" +PUFFER_SLOT_ENCODER_BACKEND=${SLOT_ENCODER_BACKEND} .venv/bin/python scripts/parallel_eval.py carla \ + --total-scenarios 40000 \ + --num-nodes 1 \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ + vec.num_envs=128 \ + eval.reward_comfort=0.0 \ + eval.reward_lane_center=0.0075 \ + env.eval_perceived_size_margin_m=0.2 \ + eval.min_goal_spacing=20 \ + eval.max_goal_spacing=200 \ + env.disable_red_light_infractions=1 \ + env.disable_stop_sign_infractions=1 \ + env.traffic_light_junction_phases=0 \ + env.eval_standstill_jerk_deadband_mps3=1.5 \ + eval.render_filter=all_infractions \ + eval.capture_observations=true \ + eval.output_name=${RUN_NAME} \ + load_model_path=${MODEL_PATH} \ + wandb=True + +end=$(date +%s) +runtime=$((end-start)) +echo "END TIME: $(date)" +echo "Runtime: ${runtime}" diff --git a/scripts/kesai/19_rollout_batch_sweep_slurm.sh b/scripts/kesai/19_rollout_batch_sweep_slurm.sh new file mode 100755 index 0000000000..ac4caa2816 --- /dev/null +++ b/scripts/kesai/19_rollout_batch_sweep_slurm.sh @@ -0,0 +1,48 @@ +#!/bin/bash +#SBATCH --job-name rollout_sweep +#SBATCH --ntasks 1 +#SBATCH --nodes 1 +#SBATCH --time 2:00:00 +#SBATCH --gres gpu:1 +#SBATCH --mem=250G +#SBATCH --cpus-per-task 60 +#SBATCH --output experiments/logs/rollout_sweep_%j.out +#SBATCH --error experiments/logs/rollout_sweep_%j.err +#SBATCH --partition low + +source scripts/kesai/cluster_env.sh || exit 1 + +# Rollout-only timing of one training rank (60 blocking workers on 60 threads, env.num_agents per worker, float16 obs buffer) +# as a function of vec.batch_size, the number of workers one recv waits for: 30 (the "auto" default) means two buffers +# alternate between simulation and GPU forward; smaller values keep more buffers in flight at the cost of more recv calls. +# Values must divide NUM_ENVS (zero_copy). Prints the trainer's section split (env wait, copy, forward) per batch size. +# Usage: sbatch scripts/kesai/19_rollout_batch_sweep_slurm.sh # override with NUM_ENVS=.. NUM_AGENTS=.. BATCH_SIZES=".." PUFFER_NUMA_BIND=0 +# NUMA A/B, 61 processes on 64 threads in both arms: sbatch -m block:fcyclic -c 128 (binds to the GPU socket half) vs PUFFER_NUMA_BIND=0 with -m block:fcyclic -c 64 (32 per socket) +echo "START TIME: $(date)" +start=$(date +%s) +cd $PD || exit 1 +nvidia-smi --query-gpu=name,compute_cap,memory.total --format=csv +echo "node $(hostname)"; nvidia-smi --query-gpu=clocks.sm,clocks.max.sm,clocks.mem,power.draw,power.limit,temperature.gpu,clocks_throttle_reasons.active --format=csv +export OMP_NUM_THREADS=1 +source .venv/bin/activate +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 + +NUM_ENVS=${NUM_ENVS:-60} +NUM_AGENTS=${NUM_AGENTS:-4096} +BATCH_SIZES=${BATCH_SIZES:-"60 30 20 15 12 10"} +RESULT_FILE=$PD/experiments/logs/rollout_sweep_${SLURM_JOB_ID}_result.txt +echo "=== rollout sweep: ${NUM_ENVS} workers x ${NUM_AGENTS} agents, float16 obs, fused+cuda encoder ===" | tee ${RESULT_FILE} +for B in ${BATCH_SIZES}; do + echo "--- vec.batch_size=${B} ---" | tee -a ${RESULT_FILE} + env PUFFER_SLOT_ENCODER_BACKEND=cuda python scripts/mfu_bench.py --num-envs ${NUM_ENVS} --minibatch 32768 --fused \ + --rollout-only --profile --warmup 1 --epochs 3 \ + env.num_agents=${NUM_AGENTS} train.rollout_dtype=float16 env.scenario_length=1200 vec.batch_size=${B} 2>&1 \ + | grep -E "ROLLOUT_RESULT|NUMA|^ (env|eval_copy|eval_forward|eval_misc|eval) |^epoch|Error|error" | tee -a ${RESULT_FILE} +done +echo "=== end of rollout sweep ===" | tee -a ${RESULT_FILE} +echo +cat ${RESULT_FILE} +end=$(date +%s) +echo "node $(hostname)"; nvidia-smi --query-gpu=clocks.sm,clocks.max.sm,clocks.mem,power.draw,power.limit,temperature.gpu,clocks_throttle_reasons.active --format=csv +echo "END TIME: $(date)" +echo "Runtime: $((end-start)) s" diff --git a/scripts/kesai/1_eval_model_slurm.sh b/scripts/kesai/1_eval_model_slurm.sh index 0a6590e5a0..47a8047036 100644 --- a/scripts/kesai/1_eval_model_slurm.sh +++ b/scripts/kesai/1_eval_model_slurm.sh @@ -4,21 +4,23 @@ #SBATCH --nodes=1 #SBATCH --time=1-00:00 #SBATCH --gres=gpu:8 -#SBATCH --mem=1007G -#SBATCH --cpus-per-task=144 -#SBATCH --output=/home/bjaeger/PufferDrive/experiments/logs/eval_%a_%A.out -#SBATCH --error=/home/bjaeger/PufferDrive/experiments/logs/eval_%a_%A.err +#SBATCH --mem=0 +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --output=experiments/logs/eval_%a_%A.out +#SBATCH --error=experiments/logs/eval_%a_%A.err #SBATCH --partition=dev +source scripts/kesai/cluster_env.sh || exit 1 + # print info about current job echo "START TIME: $(date)" start=`date +%s` export RUN_NAME=k_scaled_0008_1000 -MODEL_PATH=/home/bjaeger/PufferDrive/experiments/${RUN_NAME}/final_model.pt +MODEL_PATH=$PD/experiments/${RUN_NAME}/final_model.pt source .venv/bin/activate -bash scripts/kesai/build_ext_if_changed.sh /home/bjaeger/PufferDrive || exit 1 +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 #.venv/bin/python scripts/parallel_eval.py carla \ # --total-scenarios 40000 \ @@ -29,7 +31,7 @@ bash scripts/kesai/build_ext_if_changed.sh /home/bjaeger/PufferDrive || exit 1 # wandb=True python -m pufferlib.pufferl eval puffer_drive nuplan_single \ - eval.map_dir=/home/shared/data/nuplan/PufferDrive \ + eval.map_dir=$SHARED_DATA/nuplan/PufferDrive \ vec.num_envs=64 \ eval.max_sdc_replay_workers=64 \ env.goal_source=gt \ diff --git a/scripts/kesai/20_rollout_kernel_table_slurm.sh b/scripts/kesai/20_rollout_kernel_table_slurm.sh new file mode 100755 index 0000000000..0982bd4883 --- /dev/null +++ b/scripts/kesai/20_rollout_kernel_table_slurm.sh @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name rollout_kernels +#SBATCH --ntasks 1 +#SBATCH --nodes 1 +#SBATCH --time 0:30:00 +#SBATCH --gres gpu:1 +#SBATCH --gres-flags=enforce-binding # CPUs on the GPU's socket +#SBATCH --mem=120G +#SBATCH --cpus-per-task 18 +#SBATCH --output experiments/logs/rollout_kernels_%j.out +#SBATCH --error experiments/logs/rollout_kernels_%j.err +#SBATCH --partition low + +source scripts/kesai/cluster_env.sh || exit 1 + +# Rollout kernel table of one trainer rank in the script-16 configuration (32 workers x 1024 agents), ~3 minutes. +# Compare the fused_slot_encoder_fwd_cuda<256, 16, false> row across code versions; the cutlass 256x128 GEMM, the +# simt sgemm and the Memcpy rows are clock anchors that must not move. +# Usage: sbatch -p low -w rtx-00 scripts/kesai/20_rollout_kernel_table_slurm.sh # BACKEND=triton NUM_ENVS=.. EXTRA_ARGS=.. +echo "START TIME: $(date)" +start=$(date +%s) +cd $PD || exit 1 +nvidia-smi --query-gpu=name,compute_cap,memory.total --format=csv +echo "node $(hostname)"; nvidia-smi --query-gpu=clocks.sm,clocks.max.sm,clocks.mem,power.draw,power.limit,temperature.gpu,clocks_throttle_reasons.active --format=csv +export OMP_NUM_THREADS=1 +source .venv/bin/activate +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 + +NUM_ENVS=${NUM_ENVS:-32} +MINIBATCH=${MINIBATCH:-131072} +BACKEND=${BACKEND:-cuda} +EXTRA_ARGS=${EXTRA_ARGS:-} +RESULT_FILE=$PD/experiments/logs/rollout_kernels_${SLURM_JOB_ID}_result.txt +echo "=== rollout kernel table, fused+${BACKEND}, ${NUM_ENVS} envs, git $(git rev-parse --short HEAD) ===" | tee ${RESULT_FILE} +python -c "from pufferlib.ocean.cuda_slot_encoder import _compile; print('fwd no-store kernel on this GPU: regs/local', _compile(256, 16, False, 0)[3:5])" 2>&1 | tail -1 | tee -a ${RESULT_FILE} +env PUFFER_SLOT_ENCODER_BACKEND=${BACKEND} python scripts/mfu_bench.py --num-envs ${NUM_ENVS} --minibatch ${MINIBATCH} --fused \ + --rollout-only --kernel-table --warmup 1 --epochs 1 ${EXTRA_ARGS} 2>&1 \ + | grep -v "^MFU_RESULT" | grep "capability\|kernel table, rollout\|^ *share\|^ *[0-9][0-9.]*%\|ROLLOUT_RESULT\|^epoch\|Error" | tee -a ${RESULT_FILE} +echo "node $(hostname)"; nvidia-smi --query-gpu=clocks.sm,clocks.max.sm,clocks.mem,power.draw,power.limit,temperature.gpu,clocks_throttle_reasons.active --format=csv +end=$(date +%s) +echo "END TIME: $(date)" +echo "Runtime: $((end-start)) s" diff --git a/scripts/kesai/2_train_8GPU_nightly_SLURM.sh b/scripts/kesai/2_train_8GPU_nightly_SLURM.sh index b8811e6c8f..fcf2c704ae 100644 --- a/scripts/kesai/2_train_8GPU_nightly_SLURM.sh +++ b/scripts/kesai/2_train_8GPU_nightly_SLURM.sh @@ -4,13 +4,15 @@ #SBATCH --nodes 1 #SBATCH --time 2-00:00 #SBATCH --gres gpu:8 -#SBATCH --mem=1007G -#SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/nightly_PufferDrive/experiments/logs/log_%a_%A.out -#SBATCH --error /home/bjaeger/nightly_PufferDrive/experiments/logs/log_%a_%A.err +#SBATCH --mem=0 +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --output experiments/logs/train_nightly_%a_%A.out +#SBATCH --error experiments/logs/train_nightly_%a_%A.err #SBATCH --partition dev #SBATCH --array=0-2 +source scripts/kesai/cluster_env.sh || exit 1 + # print info about current job echo "START TIME: $(date)" start=$(date +%s) @@ -22,7 +24,7 @@ echo "SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID} -> train.seed=${SEED}" export RUN_NAME=k_nightly_"$(date +%Y-%m-%d)"_${SEED} echo ${RUN_NAME} -export DATA_DIR=/home/bjaeger/nightly_PufferDrive/experiments/${RUN_NAME} +export DATA_DIR=$PD/experiments/${RUN_NAME} echo ${DATA_DIR} export FINAL_MODEL_NAME=final_model.pt @@ -37,14 +39,14 @@ export VECLIB_MAXIMUM_THREADS=1 export OPENBLAS_NUM_THREADS=1 export OMP_NUM_THREADS=1 -REPO_DIR=/home/bjaeger/nightly_PufferDrive +REPO_DIR=$PD cd ${REPO_DIR} || exit 1 source .venv/bin/activate # Only task 0 builds; concurrent in-place builds race on shared NFS build files. -BUILD_STATUS_FILE=/home/bjaeger/nightly_PufferDrive/experiments/logs/build_status_${SLURM_ARRAY_JOB_ID} +BUILD_STATUS_FILE=$PD/experiments/logs/build_status_${SLURM_ARRAY_JOB_ID} if [ "${SLURM_ARRAY_TASK_ID}" -eq 0 ]; then - bash scripts/kesai/build_ext_if_changed.sh /home/bjaeger/nightly_PufferDrive + bash scripts/kesai/build_ext_if_changed.sh $PD BUILD_STATUS=$? echo ${BUILD_STATUS} > ${BUILD_STATUS_FILE} else @@ -71,7 +73,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met wandb_project=nightly-multi-long \ wandb_group="$(date +%Y-%m-%d)" \ train.data_dir=${DATA_DIR} \ - env.map_dir=/home/bjaeger/nightly_PufferDrive/pufferlib/resources/drive/binaries/carla \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ train.name=${RUN_NAME} \ run_name=${RUN_NAME} \ train.total_timesteps=100000000000 \ @@ -97,7 +99,7 @@ python -m pufferlib.pufferl eval puffer_drive carla \ wandb=True python -m pufferlib.pufferl eval puffer_drive nuplan_single \ - eval.map_dir=/home/shared/data/nuplan/PufferDrive \ + eval.map_dir=$SHARED_DATA/nuplan/PufferDrive \ eval.output_name=${RUN_NAME} \ load_model_path=${MODEL_PATH} \ wandb=True diff --git a/scripts/kesai/3_train_64GPU_multinode_slurm.sh b/scripts/kesai/3_train_64GPU_multinode_slurm.sh index 39334d9503..1e99d19a95 100644 --- a/scripts/kesai/3_train_64GPU_multinode_slurm.sh +++ b/scripts/kesai/3_train_64GPU_multinode_slurm.sh @@ -3,13 +3,15 @@ #SBATCH --nodes 4 # Number of nodes requested #SBATCH --ntasks-per-node 1 # Run 1 srun task per node (which fires up torchrun) #SBATCH --gres gpu:8 # GPUs per node -#SBATCH --cpus-per-task 144 -#SBATCH --mem=1007G +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --mem=0 #SBATCH --time 6-00:00 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.err +#SBATCH --output experiments/logs/train_multinode_%j.out +#SBATCH --error experiments/logs/train_multinode_%j.err #SBATCH --partition dev +source scripts/kesai/cluster_env.sh || exit 1 + # Set up PyTorch Distributed Rendezvous parameters from Slurm variables export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) export MASTER_PORT=29500 @@ -22,10 +24,10 @@ start=$(date +%s) export SEED=1000 -export RUN_NAME=k_scaled_0044_${SEED} +export RUN_NAME=${RUN_NAME:-k_scaled_0059_${SEED}} echo ${RUN_NAME} -export DATA_DIR=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} +export DATA_DIR=$PD/experiments/${RUN_NAME} echo ${DATA_DIR} # Name the training run writes its final model under, so the eval steps below can @@ -42,10 +44,20 @@ export VECLIB_MAXIMUM_THREADS=1 export OPENBLAS_NUM_THREADS=1 export OMP_NUM_THREADS=1 +# Hand-written CUDA slot-encoder kernels instead of the Triton ones (same math); triton to compare. Passed per command, +# never exported: sbatch copies this shell's environment into the eval jobs below, whose interpreters lack cuda-bindings. +SLOT_ENCODER_BACKEND=cuda +echo "slot encoder backend: ${SLOT_ENCODER_BACKEND}" + source .venv/bin/activate -bash scripts/kesai/build_ext_if_changed.sh /home/bjaeger/PufferDrive || exit 1 +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 +# Blocking env workers: kesai nodes give 60 threads per rank (--exclusive), vec.num_envs=60 is one worker per thread (run 0053, +# ~25 M SPS vs ~21.5 M with 32); scaleway nodes have 18 threads per rank, 32 there. More workers enlarge the rollout batch +# (= workers x env.num_agents x horizon); forked workers share the preloaded map cache. ROLLOUT_DTYPE=float16 halves the obs buffer on the GPU. +# A second instance needs its own RUN_NAME: RUN_NAME=k_scaled_0055_1000 sbatch -p high scripts/kesai/3_train_64GPU_multinode_slurm.sh +# 0059 baseline: its env/train/policy settings are the puffer_drive.yaml defaults; overrides below are run identity, logging, maps and hardware. # Execute torchrun across all nodes using srun -srun torchrun \ +PUFFER_SLOT_ENCODER_BACKEND=${SLOT_ENCODER_BACKEND} srun --cpus-per-task="$SLURM_CPUS_PER_TASK" torchrun \ --nnodes=${SLURM_JOB_NUM_NODES} \ --nproc-per-node=8 \ --rdzv_id=${SLURM_JOB_ID} \ @@ -58,32 +70,15 @@ srun torchrun \ wandb_project=nightly-multi-long \ wandb_group=emerge_ \ train.data_dir=${DATA_DIR} \ - env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla_128_affine_zones_stop \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla_128_affine_zones_stop_ndshoulders_v19 \ env.num_maps=128 \ train.name=${RUN_NAME} \ run_name=${RUN_NAME} \ train.total_timesteps=1000000000000 \ - vec.num_envs=16 \ - train.compile=True \ + vec.num_envs=60 \ + ${ROLLOUT_DTYPE:+train.rollout_dtype=$ROLLOUT_DTYPE} \ train.max_minibatch_size=131072 \ train.minibatch_size=131072 \ - train.precision=bfloat16 \ - policy.fp32_heads=true \ - train.tf32=false \ - env.goal_speed_randomization=false \ - env.goal_reach_requires_speed=true \ - env.obs_partner_relative_velocity=true \ - env.obs_lane_heading_signed=true \ - env.obs_lane_speed_limit=true \ - env.stop_signs_enabled=true \ - env.pose_noise_xy_m=0.025 \ - env.pose_noise_yaw_deg=0.25 \ - env.speed_limit_random_prob=1.0 \ - env.stagger_first_episode=true \ - env.resample_frequency=0 \ - env.conditioning_speed_scale=2.0 \ - policy.mask_padded_features=true \ - train.evaluation_benchmarks=carla_fast \ train.final_model_name=${FINAL_MODEL_NAME} \ train.seed=${SEED} \ tb=True @@ -93,14 +88,31 @@ if [ ! -f ${MODEL_PATH} ]; then exit 1 fi -# parallel_eval places one shard per allocated node via srun, so each shard's 64 +# nuPlan (reactive and non-reactive), longest6 (v2 and v1) and AlpaSim evals need one 8-GPU node each: submit them as their own jobs before the long self-play eval below so they queue while it runs. +echo "Training done, submitting nuPlan reactive eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh \ + || echo "nuPlan eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh" +echo "Submitting nuPlan non-reactive eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh \ + || echo "nuPlan non-reactive eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh" +echo "Submitting longest6 v2 eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6_v2.sh \ + || echo "longest6 v2 eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6_v2.sh" +echo "Submitting longest6 v1 eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/14_carla_longest6_v1.sh \ + || echo "longest6 v1 eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/14_carla_longest6_v1.sh" +echo "Submitting AlpaSim eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/12_alpasim.sh \ + || echo "AlpaSim eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/12_alpasim.sh" + +# parallel_eval places one shard per allocated node via srun, so each shard's 128 # env workers get a full node's cores instead of sharing the batch host. -echo "Training done, evaluating ${MODEL_PATH}" -.venv/bin/python scripts/parallel_eval.py carla \ +echo "Follow-up evals submitted, evaluating ${MODEL_PATH}" +PUFFER_SLOT_ENCODER_BACKEND=${SLOT_ENCODER_BACKEND} .venv/bin/python scripts/parallel_eval.py carla \ --total-scenarios 40000 \ --num-nodes 4 \ - env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ - vec.num_envs=64 \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ + vec.num_envs=128 \ eval.reward_comfort=0.0 \ eval.reward_lane_center=0.0075 \ env.eval_perceived_size_margin_m=0.2 \ @@ -116,20 +128,6 @@ echo "Training done, evaluating ${MODEL_PATH}" load_model_path=${MODEL_PATH} \ wandb=True -# nuPlan (reactive and non-reactive), longest6 and AlpaSim evals need one 8-GPU node each: submit them as their own jobs so this 4-node allocation ends now. -echo "CARLA eval done, submitting nuPlan reactive eval for ${MODEL_PATH}" -RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh \ - || echo "nuPlan eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh" -echo "Submitting nuPlan non-reactive eval for ${MODEL_PATH}" -RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh \ - || echo "nuPlan non-reactive eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh" -echo "Submitting longest6 eval for ${MODEL_PATH}" -RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6.sh \ - || echo "longest6 eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6.sh" -echo "Submitting AlpaSim eval for ${MODEL_PATH}" -RUN_DIR=${DATA_DIR} sbatch scripts/kesai/12_alpasim.sh \ - || echo "AlpaSim eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/12_alpasim.sh" - end=$(date +%s) runtime=$((end-start)) echo "END TIME: $(date)" diff --git a/scripts/kesai/4_train_1GPU_slurm.sh b/scripts/kesai/4_train_1GPU_slurm.sh index 84e7226abd..ac14d6e283 100644 --- a/scripts/kesai/4_train_1GPU_slurm.sh +++ b/scripts/kesai/4_train_1GPU_slurm.sh @@ -4,13 +4,16 @@ #SBATCH --nodes 1 #SBATCH --time 2-00:00 #SBATCH --gres gpu:1 +#SBATCH --gres-flags=enforce-binding # CPUs on the GPU's socket #SBATCH --mem=125G #SBATCH --cpus-per-task 18 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.err +#SBATCH --output experiments/logs/train_1gpu_%a_%A.out +#SBATCH --error experiments/logs/train_1gpu_%a_%A.err #SBATCH --partition dev #SBATCH --array=0 +source scripts/kesai/cluster_env.sh || exit 1 + # print info about current job echo "START TIME: $(date)" start=$(date +%s) @@ -22,7 +25,7 @@ echo "SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID} -> train.seed=${SEED}" export RUN_NAME=k_fast_0001_${SEED} echo ${RUN_NAME} -export DATA_DIR=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} +export DATA_DIR=$PD/experiments/${RUN_NAME} echo ${DATA_DIR} export FINAL_MODEL_NAME=final_model.pt @@ -37,14 +40,14 @@ export OPENBLAS_NUM_THREADS=1 export OMP_NUM_THREADS=1 source .venv/bin/activate -bash scripts/kesai/build_ext_if_changed.sh /home/bjaeger/PufferDrive || exit 1 +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 torchrun --standalone --nnodes=1 --nproc-per-node=1 --max_restarts=0 --start-method spawn \ -m pufferlib.pufferl train puffer_drive \ wandb=True \ wandb_project=nightly-multi-long \ wandb_group=emerge_ \ train.data_dir=${DATA_DIR} \ - env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ train.name=${RUN_NAME} \ run_name=${RUN_NAME} \ train.total_timesteps=1000000000 \ @@ -71,7 +74,7 @@ python -m pufferlib.pufferl eval puffer_drive carla \ wandb=True python -m pufferlib.pufferl eval puffer_drive nuplan_single \ - eval.map_dir=/home/shared/data/nuplan/PufferDrive \ + eval.map_dir=$SHARED_DATA/nuplan/PufferDrive \ eval.output_name=${RUN_NAME} \ load_model_path=${MODEL_PATH} \ wandb=True diff --git a/scripts/kesai/5_train_8GPU_experimental_slurm.sh b/scripts/kesai/5_train_8GPU_experimental_slurm.sh index 3494083e4e..c879161329 100644 --- a/scripts/kesai/5_train_8GPU_experimental_slurm.sh +++ b/scripts/kesai/5_train_8GPU_experimental_slurm.sh @@ -4,13 +4,15 @@ #SBATCH --nodes 1 #SBATCH --time 2-00:00 #SBATCH --gres gpu:8 -#SBATCH --mem=1007G -#SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.err +#SBATCH --mem=0 +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --output experiments/logs/train_experimental_%a_%A.out +#SBATCH --error experiments/logs/train_experimental_%a_%A.err #SBATCH --partition dev #SBATCH --array=0-2 +source scripts/kesai/cluster_env.sh || exit 1 + # print info about current job echo "START TIME: $(date)" start=$(date +%s) @@ -22,7 +24,7 @@ echo "SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID} -> train.seed=${SEED}" export RUN_NAME=k_exp_0022_${SEED} echo ${RUN_NAME} -export DATA_DIR=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} +export DATA_DIR=$PD/experiments/${RUN_NAME} echo ${DATA_DIR} export FINAL_MODEL_NAME=final_model.pt @@ -40,9 +42,9 @@ export OMP_NUM_THREADS=1 source .venv/bin/activate # Only task 0 builds; concurrent in-place builds race on shared NFS build files. -BUILD_STATUS_FILE=/home/bjaeger/PufferDrive/experiments/logs/build_status_${SLURM_ARRAY_JOB_ID} +BUILD_STATUS_FILE=$PD/experiments/logs/build_status_${SLURM_ARRAY_JOB_ID} if [ "${SLURM_ARRAY_TASK_ID}" -eq 0 ]; then - bash scripts/kesai/build_ext_if_changed.sh /home/bjaeger/PufferDrive + bash scripts/kesai/build_ext_if_changed.sh $PD BUILD_STATUS=$? echo ${BUILD_STATUS} > ${BUILD_STATUS_FILE} else @@ -69,7 +71,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met wandb_project=nightly-multi-long \ wandb_group=kesai \ train.data_dir=${DATA_DIR} \ - env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ train.name=${RUN_NAME} \ run_name=${RUN_NAME} \ train.total_timesteps=100000000000 \ @@ -112,7 +114,7 @@ python -m pufferlib.pufferl eval puffer_drive carla \ python -m pufferlib.pufferl eval puffer_drive nuplan_single \ eval.render_filter=all_infractions \ - eval.map_dir=/home/shared/data/nuplan/PufferDrive \ + eval.map_dir=$SHARED_DATA/nuplan/PufferDrive \ eval.capture_observations=true \ eval.output_name=${RUN_NAME} \ load_model_path=${MODEL_PATH} \ diff --git a/scripts/kesai/7_eval_8node_slurm.sh b/scripts/kesai/7_eval_8node_slurm.sh index 58a71cf93d..d04d767a4f 100644 --- a/scripts/kesai/7_eval_8node_slurm.sh +++ b/scripts/kesai/7_eval_8node_slurm.sh @@ -1,22 +1,32 @@ #!/bin/bash -#SBATCH --job-name eval_puffer_8node -#SBATCH --nodes 8 +#SBATCH --job-name eval_puffer_2node +#SBATCH --nodes 2 #SBATCH --ntasks-per-node 1 #SBATCH --gres gpu:8 -#SBATCH --cpus-per-task 144 -#SBATCH --mem=1007G +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --mem=0 #SBATCH --time 1-00:00 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/eval_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/eval_%a_%A.err +#SBATCH --output experiments/logs/eval_%a_%A.out +#SBATCH --error experiments/logs/eval_%a_%A.err #SBATCH --partition dev +source scripts/kesai/cluster_env.sh || exit 1 + +# Eval-only twin of the post-training eval in 3_train_64GPU_multinode_slurm.sh: same CARLA +# self-play benchmark and follow-up co-sim jobs, on 2 nodes (16 GPUs) instead of 4. echo "START TIME: $(date)" start=$(date +%s) -export RUN_NAME=k_scaled_0038_1000 +export SEED=1000 + +export RUN_NAME=k_scaled_0045_${SEED} echo ${RUN_NAME} -export MODEL_PATH=/home/bjaeger/PufferDrive/experiments/${RUN_NAME}/final_model.pt +export DATA_DIR=$PD/experiments/${RUN_NAME} +echo ${DATA_DIR} + +export FINAL_MODEL_NAME=final_model.pt +export MODEL_PATH=${DATA_DIR}/${FINAL_MODEL_NAME} echo ${MODEL_PATH} # Thread limit limits CPU thrashing across worker environments @@ -27,22 +37,21 @@ export OPENBLAS_NUM_THREADS=1 export OMP_NUM_THREADS=1 source .venv/bin/activate -bash scripts/kesai/build_ext_if_changed.sh /home/bjaeger/PufferDrive || exit 1 +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 if [ ! -f ${MODEL_PATH} ]; then echo "${MODEL_PATH} is missing; nothing to evaluate." exit 1 fi -# parallel_eval places one shard per allocated node via srun, so each shard's 64 +# parallel_eval places one shard per allocated node via srun, so each shard's 128 # env workers get a full node's cores instead of sharing the batch host. +echo "Evaluating ${MODEL_PATH}" .venv/bin/python scripts/parallel_eval.py carla \ - --total-scenarios 5000 \ - --num-nodes 8 \ - env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ - vec.num_envs=64 \ - eval.render_filter=all_infractions \ - eval.capture_observations=true \ + --total-scenarios 40000 \ + --num-nodes 2 \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ + vec.num_envs=128 \ eval.reward_comfort=0.0 \ eval.reward_lane_center=0.0075 \ env.eval_perceived_size_margin_m=0.2 \ @@ -52,33 +61,25 @@ fi env.disable_stop_sign_infractions=1 \ env.traffic_light_junction_phases=0 \ env.eval_standstill_jerk_deadband_mps3=1.5 \ - eval.output_name=${RUN_NAME}_random_tl2 \ - load_model_path=${MODEL_PATH} \ - wandb=True - -# eval.obs_slots_partners_n=40 \ -# eval.goal_regen_mode=rolling \ -# env.max_speed_mps=13.33 \ - -# No srun: single-node eval, run here on the batch host. -python -m pufferlib.pufferl eval puffer_drive nuplan_multi \ - eval.map_dir=/home/shared/data/nuplan/PufferDrive_broken_noconn \ - vec.num_envs=64 \ - eval.max_sdc_replay_workers=64 \ - eval.num_agents=300 \ - eval.reward_comfort=0.0 \ - eval.reward_lane_center=0.0075 \ - env.eval_perceived_size_margin_m=0.0 \ - eval.disable_red_light_infractions=1 \ - eval.disable_stop_sign_infractions=1 \ - env.eval_standstill_jerk_deadband_mps3=1.5 \ eval.render_filter=all_infractions \ eval.capture_observations=true \ - eval.output_name=${RUN_NAME}_random_tl2 \ + eval.output_name=${RUN_NAME} \ load_model_path=${MODEL_PATH} \ wandb=True -# env.max_speed_mps=13.33 \ +# nuPlan (reactive and non-reactive), longest6 and AlpaSim evals need one 8-GPU node each: submit them as their own jobs so this 2-node allocation ends now. +echo "CARLA eval done, submitting nuPlan reactive eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh \ + || echo "nuPlan eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh" +echo "Submitting nuPlan non-reactive eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh \ + || echo "nuPlan non-reactive eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh" +echo "Submitting longest6 v2 eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6_v2.sh \ + || echo "longest6 v2 eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6_v2.sh" +echo "Submitting AlpaSim eval for ${MODEL_PATH}" +RUN_DIR=${DATA_DIR} sbatch scripts/kesai/12_alpasim.sh \ + || echo "AlpaSim eval submission failed; run by hand: RUN_DIR=${DATA_DIR} sbatch scripts/kesai/12_alpasim.sh" end=$(date +%s) runtime=$((end-start)) diff --git a/scripts/kesai/8_train_64GPU_multinode_slurm.sh b/scripts/kesai/8_train_64GPU_multinode_slurm.sh index f7ea094770..b9ad229177 100644 --- a/scripts/kesai/8_train_64GPU_multinode_slurm.sh +++ b/scripts/kesai/8_train_64GPU_multinode_slurm.sh @@ -3,12 +3,15 @@ #SBATCH --nodes 8 # Number of nodes requested #SBATCH --ntasks-per-node 1 # Run 1 srun task per node (which fires up torchrun) #SBATCH --gres gpu:8 # GPUs per node -#SBATCH --cpus-per-task 144 -#SBATCH --mem=1007G +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --mem=0 #SBATCH --time 3-00:00 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/log_%a_%A.err +#SBATCH --output experiments/logs/train_64gpu_%j.out +#SBATCH --error experiments/logs/train_64gpu_%j.err #SBATCH --partition dev +# kesai: 64 GPUs exceed the 32-GPU per-user cap of dev -> submit with -p high (uncapped, internal approval) or -p low (preemptible) + +source scripts/kesai/cluster_env.sh || exit 1 # Set up PyTorch Distributed Rendezvous parameters from Slurm variables export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) @@ -25,7 +28,7 @@ export SEED=1000 export RUN_NAME=k_scaled_0034_${SEED} echo ${RUN_NAME} -export DATA_DIR=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} +export DATA_DIR=$PD/experiments/${RUN_NAME} echo ${DATA_DIR} # Name the training run writes its final model under, so the eval steps below can @@ -43,9 +46,9 @@ export OPENBLAS_NUM_THREADS=1 export OMP_NUM_THREADS=1 source .venv/bin/activate -bash scripts/kesai/build_ext_if_changed.sh /home/bjaeger/PufferDrive || exit 1 +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 # Execute torchrun across all nodes using srun -srun torchrun \ +srun --cpus-per-task="$SLURM_CPUS_PER_TASK" torchrun \ --nnodes=${SLURM_JOB_NUM_NODES} \ --nproc-per-node=8 \ --rdzv_id=${SLURM_JOB_ID} \ @@ -58,7 +61,7 @@ srun torchrun \ wandb_project=nightly-multi-long \ wandb_group=emerge_ \ train.data_dir=${DATA_DIR} \ - env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla_128_affine \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla_128_affine \ env.num_maps=128 \ env.goal_speed_randomization=false \ env.goal_reach_requires_speed=true \ @@ -89,7 +92,7 @@ echo "Training done, evaluating ${MODEL_PATH}" .venv/bin/python scripts/parallel_eval.py carla \ --total-scenarios 40000 \ --num-nodes 8 \ - env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ vec.num_envs=64 \ eval.reward_comfort=0.0 \ eval.reward_lane_center=0.0075 \ @@ -108,7 +111,7 @@ echo "Training done, evaluating ${MODEL_PATH}" wandb=True python -m pufferlib.pufferl eval puffer_drive nuplan_multi \ - eval.map_dir=/home/shared/data/nuplan/PufferDrive \ + eval.map_dir=$SHARED_DATA/nuplan/PufferDrive \ vec.num_envs=64 \ eval.num_agents=300 \ eval.reward_comfort=0.0 \ diff --git a/scripts/kesai/9_nuPlan.sh b/scripts/kesai/9_nuPlan.sh index 53cc904afd..a39aa8428c 100644 --- a/scripts/kesai/9_nuPlan.sh +++ b/scripts/kesai/9_nuPlan.sh @@ -4,17 +4,17 @@ #SBATCH --nodes 1 #SBATCH --time 1-00:00 #SBATCH --gres gpu:8 -#SBATCH --mem=1007G -#SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/logs/log_%j.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/logs/log_%j.err +#SBATCH --mem=0 +#SBATCH --exclusive # whole node on either cluster, --mem=0 = all of its memory +#SBATCH --output experiments/logs/nuPlan_reactive_%j.out +#SBATCH --error experiments/logs/nuPlan_reactive_%j.err #SBATCH --partition dev set -u +source scripts/kesai/cluster_env.sh || exit 1 -export PD=/home/bjaeger/PufferDrive export PY=$(conda info --base)/envs/carl_nuplan/bin/python # Overridable env: RUN_DIR (3_train_64GPU_multinode_slurm.sh submits this job with RUN_DIR=), CHALLENGES -RUN_DIR=${RUN_DIR:-/home/bjaeger/PufferDrive/experiments/k_scaled_0041_1000} +RUN_DIR=${RUN_DIR:-$PD/experiments/k_scaled_0045_1000} # the planner finds config.yaml next to final_model.pt (or one level above a models/*.pt) export CKPT=$RUN_DIR/final_model.pt [ -f "$CKPT" ] || { echo "missing $CKPT"; exit 1; } @@ -26,12 +26,14 @@ case "$CHALLENGES" in *) echo "CHALLENGES must be closed_loop_reactive_agents_pufferdrive or closed_loop_nonreactive_agents_pufferdrive"; exit 1 ;; esac echo "Evaluating checkpoint: $CKPT" -export CITY_BIN_DIR=/home/shared/data/nuplan/PufferDrive +export CITY_BIN_DIR=$SHARED_DATA/nuplan/PufferDrive # shell env still exports the pre-rename nuPlan casing; pin the renamed paths here -export NUPLAN_DATA_ROOT=/home/shared/data/nuplan -export NUPLAN_MAPS_ROOT=/home/shared/data/nuplan/maps +export NUPLAN_DATA_ROOT=$SHARED_DATA/nuplan +export NUPLAN_MAPS_ROOT=$SHARED_DATA/nuplan/maps +# the runner's default would mkdir a /scratch path of the old cluster (harmless, but a Permission denied line in every kesai log) +export NUPLAN_EXP_ROOT=${NUPLAN_EXP_ROOT:-$HOME/nuplan_exp} # PD first, and drop inherited entries from other PufferDrive checkouts (the login shell still exports -# /home/bjaeger/cosim_Puffer): Python would otherwise import that checkout's pufferlib package. +# $HOME/cosim_Puffer): Python would otherwise import that checkout's pufferlib package. INHERITED_PYTHONPATH=$(echo "${PYTHONPATH:-}" | tr ':' '\n' | grep -v "cosim_Puffer\|/PufferDrive" | paste -sd: -) export PYTHONPATH=$PD:$CARL_DEVKIT_ROOT:$NUPLAN_DEVKIT_ROOT${INHERITED_PYTHONPATH:+:$INHERITED_PYTHONPATH} cd "$PD" || exit 1 @@ -41,7 +43,9 @@ export COSIM_OBS_HTML=infractions # obs replays (what the policy saw) only for export CALLBACKS="[simulation_log_callback, carl_visualization_callback]" export WORKER=ray_distributed # Per-worker memory is small now (1 policy agent + partner slots, light buffers sized to the scenario), so this is CPU-bound. -export THREADS_PER_NODE=128 +# 128 fits the 144-CPU scaleway nodes. A kesai RTX node gives the job 480 CPUs and a worker needs ~1.9 GB: 256 doubles the +# throughput (128 workers used 29 % of the CPUs there, 26 Oct) and leaves half the memory to the page cache of the logs. +export THREADS_PER_NODE=${THREADS_PER_NODE:-$([ "$CLUSTER" = kesai ] && echo 256 || echo 128)} # route: lane-graph route goals; roadblock: route roadblock centroids; roadblock_lane: roadblock goals on the lane continuing the ego lane export GOAL_SOURCE=roadblock # Ablations (see run_nuplan_planner.sh): refill the goal window after every consumed goal; cap the ego's @@ -49,24 +53,29 @@ export GOAL_SOURCE=roadblock export SLIDING_GOAL_WINDOW=true export STARTUP_ACCEL_JERK_CAP=2 export STARTUP_BRAKE_JERK_CAP=2 -# pedestrian/bicycle partner boxes grown to the training minimum agent size in the policy's obs (0 = off) -export PEDESTRIAN_MIN_SIZE_M=0.8 +# pedestrian and bicycle partner boxes grow by this much in length and width in the policy's obs (0 = off): 0.8 puts them strictly above the smallest training vehicle +export VRU_SIZE_PAD_M=0.8 # Shadow-env obs overrides: partner slots (training 20, GIGAFLOW eval 40; 40 real partners max-pooled is OOD in # sidewalk crowds) and the ego perceived-size margin [m per side] (0.2 = native eval; nuPlan's 2.3 x 5.2 m ego # reads as 2.7 x 5.6 m). Unset = planner yaml / arch.py defaults (40 / 0.2). export OBS_SLOTS_PARTNERS_N=40 export EVAL_PERCEIVED_SIZE_MARGIN_M=0.2 +# overspeed conditioning alpha (unset = arch.py 0.1 = max of the U(0, 0.1) training range; 0.05 kept the ego at the limit, 0.005 let it speed) +export REWARD_OVERSPEED=0.1 ABLATION_TAG="" [ "$SLIDING_GOAL_WINDOW" = "true" ] && ABLATION_TAG="${ABLATION_TAG}_slide" [ "$STARTUP_ACCEL_JERK_CAP" != "0" ] && ABLATION_TAG="${ABLATION_TAG}_jacc${STARTUP_ACCEL_JERK_CAP}" [ "$STARTUP_BRAKE_JERK_CAP" != "0" ] && ABLATION_TAG="${ABLATION_TAG}_jbrk${STARTUP_BRAKE_JERK_CAP}" -[ "$PEDESTRIAN_MIN_SIZE_M" != "0" ] && ABLATION_TAG="${ABLATION_TAG}_ped${PEDESTRIAN_MIN_SIZE_M}m" +[ "$VRU_SIZE_PAD_M" != "0" ] && ABLATION_TAG="${ABLATION_TAG}_vrupad${VRU_SIZE_PAD_M}m" [ -n "${OBS_SLOTS_PARTNERS_N:-}" ] && ABLATION_TAG="${ABLATION_TAG}_slots${OBS_SLOTS_PARTNERS_N}" [ -n "${EVAL_PERCEIVED_SIZE_MARGIN_M:-}" ] && ABLATION_TAG="${ABLATION_TAG}_margin${EVAL_PERCEIVED_SIZE_MARGIN_M}" +[ -n "${REWARD_OVERSPEED:-}" ] && ABLATION_TAG="${ABLATION_TAG}_overspeed${REWARD_OVERSPEED}" # results live in the model's own eval folder, next to the PufferDrive benchmark evals export GROUP=$RUN_DIR/eval/nuplan_val14${CHALLENGE_TAG}_${GOAL_SOURCE}${ABLATION_TAG}_$(date +%Y%m%d_%H%M%S)_${SLURM_JOB_ID:-local} -PATH="$(dirname "$PY"):$PATH" bash "$PD/scripts/kesai/build_ext_if_changed.sh" "$PD" || exit 1 +# NO_TRAIN skips the training-only CUDA extension (_C, never imported by the planner): the cp39 nuPlan env builds only +# the C binding, so nvcc 13.4 on kesai never meets this env's cu128 torch (script 12 builds the same way) +NO_TRAIN=1 PATH="$(dirname "$PY"):$PATH" bash "$PD/scripts/kesai/build_ext_if_changed.sh" "$PD" || exit 1 # The planner must come from this checkout (needs the freshly built C extension above). IMPORTED_PLANNER=$("$PY" -c "import pufferlib.ocean.cosim.nuplan.planner as p; print(p.__file__)") || { echo "ERROR: planner import failed (see traceback above)"; exit 1; } diff --git a/scripts/kesai/cluster_env.sh b/scripts/kesai/cluster_env.sh new file mode 100644 index 0000000000..096c458f4d --- /dev/null +++ b/scripts/kesai/cluster_env.sh @@ -0,0 +1,41 @@ +# Cluster detection for the scripts in this folder, sourced from the checkout root after the #SBATCH block: +# source scripts/kesai/cluster_env.sh || exit 1 +# Always submit from the checkout root (sbatch keeps it as the job's working directory) and create experiments/logs once. +# Sets CLUSTER (kesai | scaleway | local), HOME_BASE (where user homes live; other users' paths are $HOME_BASE//...), +# SHARED_DATA (the shared datasets: nuplan, nurec, ... live directly below it), PD (this checkout) and SLURM_CPUS_PER_TASK +# (whole-node jobs carry #SBATCH --exclusive instead of a per-cluster CPU count; srun steps and scripts/parallel_eval.py +# read the count from here). Your own paths stay $HOME/..., which is /home/ on scaleway and /cephfs/home/ on kesai. +if [ -d /cephfs/data ]; then + CLUSTER=kesai + HOME_BASE=/cephfs/home + SHARED_DATA=/cephfs/data + # the nodes also carry the BMC's USB NIC (169.254.x.x): pin NCCL and Gloo to the cluster network + CLUSTER_IFACE=$(ip -o -4 addr show to 192.168.20.0/24 2>/dev/null | head -n 1 | cut -d' ' -f2) + [ -n "$CLUSTER_IFACE" ] && export NCCL_SOCKET_IFNAME=$CLUSTER_IFACE GLOO_SOCKET_IFNAME=$CLUSTER_IFACE + # the 400G NICs have no working fabric yet: NCCL must not try RoCE over them (remove with the fabric) + export NCCL_IB_DISABLE=1 +elif [ -d /home/shared/data ]; then + CLUSTER=scaleway + HOME_BASE=/home + SHARED_DATA=/home/shared/data +else + CLUSTER=local + HOME_BASE=/home + SHARED_DATA= +fi +PD=${PD:-${SLURM_SUBMIT_DIR:-$PWD}} +if [ ! -f "$PD/setup.py" ]; then + echo "cluster_env.sh: $PD is not the PufferDrive checkout; submit from the checkout root or set PD" >&2 + return 1 2>/dev/null || exit 1 +fi +export SLURM_CPUS_PER_TASK=${SLURM_CPUS_ON_NODE:-${SLURM_CPUS_PER_TASK:-$(nproc)}} +export CLUSTER HOME_BASE SHARED_DATA PD +echo "cluster: $CLUSTER checkout: $PD shared data: ${SHARED_DATA:-none} cpus per task: $SLURM_CPUS_PER_TASK" + +# "taskset -c " for the NUMA node of GPU index $1 (nvidia-smi order), empty when unknown; prefix per-GPU workers with it +numa_pin_for_gpu() { + local bus node + bus=$(nvidia-smi --query-gpu=pci.bus_id --format=csv,noheader -i "$1" 2>/dev/null | tr 'A-F' 'a-f') + node=$(cat "/sys/bus/pci/devices/${bus:4}/numa_node" 2>/dev/null) || return 0 + [ "$node" -ge 0 ] 2>/dev/null && echo "taskset -c $(cat "/sys/devices/system/node/node${node}/cpulist")" +} diff --git a/scripts/kesai/docker_archive_to_sqsh.py b/scripts/kesai/docker_archive_to_sqsh.py new file mode 100644 index 0000000000..a5df2ed826 --- /dev/null +++ b/scripts/kesai/docker_archive_to_sqsh.py @@ -0,0 +1,232 @@ +"""Convert a `docker save` archive into an enroot squashfs image (.sqsh) on a machine without Docker. + +usage: python scripts/kesai/docker_archive_to_sqsh.py alpasim-base-0.134.0.tar.zst /cephfs/data/nurec/sqsh/alpasim_base_0.134.0.sqsh \ + --work_dir /cephfs/scratch/$USER/sqsh_build [--keep_work] [--processors 32] + +Flattens the image layers (OCI whiteouts applied), writes enroot's /etc/environment, /etc/fstab and /etc/rc from the +image config (ENV, VOLUME, WORKDIR, ENTRYPOINT, CMD), the way `enroot import` does, and packs the rootfs with mksquashfs. +work_dir needs about twice the unpacked image size. The sqsh name must be what the AlpaSim wizard looks for: +image basename with ':' and '-' replaced by '_', e.g. alpasim-base:0.134.0 -> alpasim_base_0.134.0.sqsh. +Needs tar, zstd (for .zst archives) and mksquashfs (squashfs-tools). +""" + +import argparse +import json +import os +import shlex +import shutil +import stat +import subprocess +import tarfile +from pathlib import Path + +COPY_BUFFER_BYTES = 1 << 20 +MAX_SYMLINK_HOPS = 40 +OPAQUE_WHITEOUT = ".wh..wh..opq" +WHITEOUT_PREFIX = ".wh." +SQUASH_OPTIONS = [ + "-noappend", + "-comp", + "lz4", + "-Xhc", + "-b", + "262144", + "-no-xattrs", + "-force-uid", + "0", + "-force-gid", + "0", +] +EXTRACTABLE_TYPES = (tarfile.REGTYPE, tarfile.AREGTYPE, tarfile.DIRTYPE, tarfile.SYMTYPE, tarfile.LNKTYPE) + + +def extract_archive(archive, archive_dir): + archive_dir.mkdir(parents=True) + if archive.suffix == ".zst": + decompress = subprocess.Popen(["zstd", "-dc", "-T0", str(archive)], stdout=subprocess.PIPE) + untar = subprocess.run(["tar", "-x", "-C", str(archive_dir)], stdin=decompress.stdout) + decompress.stdout.close() + if decompress.wait() != 0 or untar.returncode != 0: + raise SystemExit(f"extracting {archive} failed (zstd {decompress.returncode}, tar {untar.returncode})") + return + subprocess.run(["tar", "-xf", str(archive), "-C", str(archive_dir)], check=True) + + +def rootfs_path(rootfs, relative_name): + """Path of an image entry inside rootfs, following parent symlinks within the image like a chroot would.""" + parts = [part for part in relative_name.split("/") if part not in ("", ".")] + if ".." in parts: + raise SystemExit(f"refusing image entry with '..': {relative_name}") + if not parts: + return rootfs + current = rootfs + pending = parts[:-1] + hops = 0 + while pending: + candidate = current / pending.pop(0) + if not candidate.is_symlink(): + current = candidate + continue + hops += 1 + if hops > MAX_SYMLINK_HOPS: + raise SystemExit(f"symlink loop while resolving {relative_name}") + target = os.readlink(candidate) + target_parts = [part for part in target.split("/") if part not in ("", ".")] + if target.startswith("/"): + current = rootfs + if ".." in target_parts: + raise SystemExit(f"refusing symlink with '..' inside the image: {candidate} -> {target}") + pending = target_parts + pending + return current / parts[-1] + + +def remove_existing(path): + if path.is_dir() and not path.is_symlink(): + shutil.rmtree(path) + elif path.is_symlink() or path.exists(): + path.unlink() + + +def apply_whiteouts(layer_members, rootfs): + removed_count = 0 + for member in layer_members: + basename = os.path.basename(member.name.rstrip("/")) + if not basename.startswith(WHITEOUT_PREFIX): + continue + parent = rootfs_path(rootfs, member.name).parent + if basename == OPAQUE_WHITEOUT: + targets = list(parent.iterdir()) if parent.is_dir() else [] + else: + targets = [parent / basename[len(WHITEOUT_PREFIX) :]] + for target in targets: + if target.is_symlink() or target.exists(): + remove_existing(target) + removed_count += 1 + return removed_count + + +def extract_layer(layer_tar, rootfs, dir_modes): + """Extract one layer over rootfs. -> (entries written, entries skipped: devices and fifos).""" + written_count = skipped_count = 0 + for member in layer_tar.getmembers(): + basename = os.path.basename(member.name.rstrip("/")) + if basename.startswith(WHITEOUT_PREFIX) or member.name.strip("/") in ("", "."): + continue + if member.type not in EXTRACTABLE_TYPES: + skipped_count += 1 + continue + destination = rootfs_path(rootfs, member.name) + destination.parent.mkdir(parents=True, exist_ok=True) + if member.isdir(): + if destination.is_symlink() or (destination.exists() and not destination.is_dir()): + destination.unlink() + destination.mkdir(exist_ok=True) + destination.chmod((member.mode & 0o7777) | stat.S_IRWXU) # owner access while later layers write into it + dir_modes[destination] = member.mode & 0o7777 + elif member.issym(): + remove_existing(destination) + os.symlink(member.linkname, destination) + elif member.islnk(): + remove_existing(destination) + os.link(rootfs_path(rootfs, member.linkname), destination) + else: + remove_existing(destination) + with layer_tar.extractfile(member) as source, open(destination, "wb") as target: + shutil.copyfileobj(source, target, COPY_BUFFER_BYTES) + destination.chmod(member.mode & 0o7777) + written_count += 1 + return written_count, skipped_count + + +def write_enroot_config(rootfs, image_config): + """/etc/environment, /etc/fstab and /etc/rc as `enroot import` generates them from the image config.""" + etc = rootfs / "etc" + etc.mkdir(exist_ok=True) + (etc / "environment").write_text("".join(line + "\n" for line in image_config.get("Env") or [])) + volumes = image_config.get("Volumes") or {} + (etc / "fstab").write_text( + "".join(f"tmpfs {volume} tmpfs x-create=dir,rw,nosuid,nodev\n" for volume in sorted(volumes)) + ) + workdir = image_config.get("WorkingDir") or "/" + rootfs_path(rootfs, workdir).mkdir(parents=True, exist_ok=True) + entrypoint = " ".join(shlex.quote(part) for part in image_config.get("Entrypoint") or []) + cmd = " ".join(shlex.quote(part) for part in image_config.get("Cmd") or []) + labels = image_config.get("Labels") or {} + rc_lines = [f"# {key} {value}" for key, value in sorted(labels.items())] + [ + f"cd {shlex.quote(workdir)} && unset OLDPWD || exit 1", + "", + "if [ -s /etc/rc.local ]; then", + " . /etc/rc.local", + "fi", + "", + "if [ $# -gt 0 ]; then", + f' exec {entrypoint} "$@"'.replace("exec ", "exec "), + "else", + f" exec {entrypoint} {cmd}".replace("exec ", "exec "), + "fi", + ] + (etc / "rc").write_text("".join(line + "\n" for line in rc_lines)) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("archive", type=Path, help="docker save archive: .tar, .tar.gz or .tar.zst") + ap.add_argument("output", type=Path, help="squashfs image to write") + ap.add_argument("--work_dir", type=Path, default=None, help="scratch space, default .work") + ap.add_argument("--keep_work", action="store_true", help="leave the unpacked archive and rootfs in work_dir") + ap.add_argument("--processors", type=int, default=os.cpu_count(), help="mksquashfs threads") + args = ap.parse_args() + for tool in ("tar", "mksquashfs") + (("zstd",) if args.archive.suffix == ".zst" else ()): + if shutil.which(tool) is None: + raise SystemExit(f"missing tool: {tool}") + if not args.archive.is_file(): + raise SystemExit(f"no archive at {args.archive}") + if args.output.exists(): + raise SystemExit(f"{args.output} already exists") + work_dir = args.work_dir or args.output.with_name(args.output.name + ".work") + if work_dir.exists(): + raise SystemExit(f"work dir {work_dir} already exists; remove it or pass another --work_dir") + archive_dir, rootfs = work_dir / "archive", work_dir / "rootfs" + + print(f"[docker_archive_to_sqsh] extracting {args.archive} -> {archive_dir}", flush=True) + extract_archive(args.archive, archive_dir) + manifest = json.loads((archive_dir / "manifest.json").read_text()) + if len(manifest) != 1: + raise SystemExit(f"expected one image in the archive, got {[entry.get('RepoTags') for entry in manifest]}") + image_config = json.loads((archive_dir / manifest[0]["Config"]).read_text())["config"] + layers = [archive_dir / layer for layer in manifest[0]["Layers"]] + print(f"[docker_archive_to_sqsh] image {manifest[0].get('RepoTags')}: {len(layers)} layers", flush=True) + + rootfs.mkdir() + dir_modes = {} + for layer_idx, layer in enumerate(layers): + with tarfile.open(layer, "r:*") as layer_tar: + removed_count = apply_whiteouts(layer_tar.getmembers(), rootfs) + written_count, skipped_count = extract_layer(layer_tar, rootfs, dir_modes) + if layer not in layers[layer_idx + 1 :]: # identical layers share one blob + layer.unlink() + print( + f"[docker_archive_to_sqsh] layer {layer_idx + 1}/{len(layers)}: {written_count} entries, " + f"{removed_count} whiteout removals, {skipped_count} device/fifo entries skipped", + flush=True, + ) + write_enroot_config(rootfs, image_config) + for directory in sorted(dir_modes, key=lambda path: len(path.parts), reverse=True): + if directory.is_dir() and not directory.is_symlink(): # a later layer may have replaced or removed it + directory.chmod(dir_modes[directory] | stat.S_IRUSR | stat.S_IXUSR) # mksquashfs must be able to read it + + partial_output = args.output.with_name(args.output.name + ".partial") + partial_output.unlink(missing_ok=True) + print(f"[docker_archive_to_sqsh] mksquashfs {rootfs} -> {args.output}", flush=True) + subprocess.run( + ["mksquashfs", str(rootfs), str(partial_output), "-processors", str(args.processors), *SQUASH_OPTIONS], + check=True, + ) + partial_output.replace(args.output) + if not args.keep_work: + shutil.rmtree(work_dir) + print(f"[docker_archive_to_sqsh] done: {args.output} ({args.output.stat().st_size / 1e9:.2f} GB)") + + +if __name__ == "__main__": + main() diff --git a/scripts/kesai/nccl_check.py b/scripts/kesai/nccl_check.py new file mode 100644 index 0000000000..f84f7f55d6 --- /dev/null +++ b/scripts/kesai/nccl_check.py @@ -0,0 +1,39 @@ +"""NCCL smoke test for a torchrun launch: one tiny all_reduce (hangs = transport broken), then a 256 MiB all_reduce +bandwidth figure. 2 min timeout instead of NCCL's 10 min default. + one node: srun -p dev -N1 --gres=gpu:8 --exclusive -t 10 torchrun --standalone --nproc-per-node=8 scripts/kesai/nccl_check.py + two nodes: see 3_train_64GPU_multinode_slurm.sh for the rdzv flags; NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=INIT,NET shows the transport. +""" + +import datetime +import os +import time + +import torch +import torch.distributed as dist + +dist.init_process_group("nccl", timeout=datetime.timedelta(seconds=120)) +rank = dist.get_rank() +local_rank = int(os.environ["LOCAL_RANK"]) +torch.cuda.set_device(local_rank) +node = os.uname().nodename + +x = torch.ones(1, device="cuda") +t = time.time() +dist.all_reduce(x) +torch.cuda.synchronize() +print( + f"rank {rank} on {node}: all_reduce ok, sum={x.item():.0f} of {dist.get_world_size()}, {time.time() - t:.2f}s", + flush=True, +) + +y = torch.empty(64 * 1024 * 1024, device="cuda") # 256 MiB +dist.barrier() +torch.cuda.synchronize() +t = time.time() +for _ in range(5): + dist.all_reduce(y) +torch.cuda.synchronize() +dt = (time.time() - t) / 5 +if rank == 0: + print(f"all_reduce 256 MiB: {dt * 1000:.0f} ms per call, {y.numel() * 4 / dt / 1e9:.2f} GB/s", flush=True) +dist.destroy_process_group() diff --git a/scripts/kesai/nightly_scrontab.txt b/scripts/kesai/nightly_scrontab.txt index 17c07754f0..2f6defdf70 100644 --- a/scripts/kesai/nightly_scrontab.txt +++ b/scripts/kesai/nightly_scrontab.txt @@ -1,7 +1,7 @@ # Nightly training submitter. Install by pasting into: scrontab -e -# The scron job only runs sbatch; the training job itself carries its own #SBATCH resources. +# Paths are relative to $HOME, the working directory of scron jobs. The scron job only runs sbatch; the training job itself carries its own #SBATCH resources. #SCRON --partition=dev #SCRON --time=00:05:00 #SCRON --job-name=nightly_submit -#SCRON --output=/home/bjaeger/nightly_PufferDrive/experiments/logs/scron_submit_%j.out -0 21 * * * cd /home/bjaeger/nightly_PufferDrive && sbatch scripts/kesai/2_train_8GPU_nightly_SLURM.sh +#SCRON --output=nightly_PufferDrive/experiments/logs/scron_submit_%j.out +0 21 * * * cd $HOME/nightly_PufferDrive && sbatch scripts/kesai/2_train_8GPU_nightly_SLURM.sh diff --git a/scripts/kesai/numa_check.sh b/scripts/kesai/numa_check.sh new file mode 100755 index 0000000000..a9b5f70d3c --- /dev/null +++ b/scripts/kesai/numa_check.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Where a running training job's env workers execute relative to their rank's GPU socket on a kesai node +# (GPUs 0-3 on NUMA node 0, 4-7 on node 1). Run on the compute node: bash scripts/kesai/numa_check.sh +# A trainer rank is a process with LOCAL_RANK in its environment that has children (the env workers); +# torchrun has children but no LOCAL_RANK, workers have LOCAL_RANK but no children. +found=0 +# torchrun's children are the trainer ranks; without torchrun fall back to scanning every pufferl process (slow: ~500 workers) +candidates=$(for launcher in $(pgrep -f "bin/torchrun"); do pgrep -P "$launcher"; done) +[ -n "$candidates" ] || candidates=$(pgrep -f "pufferlib.pufferl") +for pid in $candidates; do + rank=$(tr '\0' '\n' < /proc/"$pid"/environ 2>/dev/null | awk -F= '$1 == "LOCAL_RANK" {print $2}') + [ -n "$rank" ] || continue + kids=$(pgrep -P "$pid") + [ -n "$kids" ] || continue + found=1 + node0=0; node1=0 + for kid in $kids; do + psr=$(ps -o psr= -p "$kid" | tr -d ' ') + [ -n "$psr" ] || continue + if [ "$psr" -lt 128 ] || { [ "$psr" -ge 256 ] && [ "$psr" -lt 384 ]; }; then node0=$((node0 + 1)); else node1=$((node1 + 1)); fi + done + echo "rank $rank (GPU on node $((rank / 4))): $node0 workers on node 0, $node1 on node 1; trainer mask:$(taskset -cp "$pid" | cut -d: -f2)" +done +[ "$found" = 1 ] || echo "no pufferlib.pufferl trainer processes with LOCAL_RANK found on $(hostname); is the training job on this node?" diff --git a/scripts/mfu_bench.py b/scripts/mfu_bench.py new file mode 100644 index 0000000000..d41560d9d7 --- /dev/null +++ b/scripts/mfu_bench.py @@ -0,0 +1,307 @@ +"""Model FLOP utilization of the real PPO loop: rollout, update and whole epoch, plus agent steps per second. + +Builds the trainer exactly as `puffer train puffer_drive` does, runs warm-up epochs (torch.compile), then measures +steady-state epochs. Model FLOPs are counted analytically from the policy's linear layers, so no profiler touches +the compiled policy (a TorchDispatchMode around it makes dynamo skip the frame for the rest of the process). + +Usage: + python scripts/mfu_bench.py --num-envs 16 --minibatch 131072 [--fused] [--profile] [-- train.seed=3 ...] +""" + +import argparse +import json +import os +import sys +import tempfile +import time +from collections import defaultdict + + +os.environ.setdefault("OMP_NUM_THREADS", "1") +import torch +import torch._dynamo +from profile_policy_step import peak_bf16_tflops +from torch import nn + + +ENV_NAME = "puffer_drive" +RECOMPILE_EPOCH_FACTOR = 1.25 +PROFILE_SECTIONS = ( + "env", + "eval_copy", + "eval_forward", + "eval_misc", + "train_copy", + "train_forward", + "learn", + "train_misc", +) +SLOT_ENCODER_COUNTS = { + "lane_encoder": "obs_slots_lane_kept", + "boundary_encoder": "obs_slots_boundary_kept", + "partner_encoder": "obs_slots_partners_n", + "traffic_control_encoder": "obs_slots_traffic_controls_n", +} + + +def model_flops_per_transition(policy): + """Forward and forward+backward FLOPs of one transition through the policy (matmuls only, the MFU convention).""" + forward = 0.0 + backward = 0.0 + for name, module in policy.named_modules(): + if not isinstance(module, nn.Linear): + continue + rows = 1 + backbone_name = name.split(".")[0] + backbone = getattr(policy, backbone_name, None) + for encoder_name, count_attribute in SLOT_ENCODER_COUNTS.items(): + if encoder_name in name and backbone is not None: + rows = getattr(backbone, count_attribute) + layer_flops = 2.0 * rows * module.in_features * module.out_features + forward += layer_flops + # the first layer of every encoder reads observations, which need no input gradient + first_encoder_layer = "encoder" in name and name.endswith(".0") + backward += layer_flops if first_encoder_layer else 2.0 * layer_flops + return forward, forward + backward + + +def print_kernel_table(prof, label, wall_seconds, top=18): + """Per-kernel self CUDA time of one profiled phase, the share of all kernel time, and the GPU busy fraction.""" + events = [ + e for e in prof.key_averages() if getattr(e, "device_type", None) is not None and e.self_device_time_total > 0 + ] + events = [e for e in events if str(e.device_type).endswith("CUDA")] + total_us = sum(e.self_device_time_total for e in events) + print( + f"kernel table, {label}: {total_us / 1e6:.3f} s of CUDA kernel time in {wall_seconds:.3f} s wall " + f"(GPU busy {100 * total_us / 1e6 / max(wall_seconds, 1e-9):.0f}%), {len(events)} distinct kernels" + ) + print(f" {'share':>6s} {'ms':>9s} {'calls':>7s} kernel") + for e in sorted(events, key=lambda e: -e.self_device_time_total)[:top]: + print( + f" {100 * e.self_device_time_total / total_us:5.1f}% {e.self_device_time_total / 1e3:9.1f} {e.count:7d} {e.key[:110]}" + ) + + +def measure_bf16_peak(seconds=3.0, size=8192): + """Achieved dense bf16 matmul TFLOP/s with fp32 accumulation on this device (cuBLAS, large square shapes).""" + a = torch.randn(size, size, device="cuda", dtype=torch.bfloat16) + b = torch.randn(size, size, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + a @ b + torch.cuda.synchronize() + start = time.perf_counter() + iterations = 0 + while time.perf_counter() - start < seconds: + a @ b + iterations += 1 + torch.cuda.synchronize() + elapsed = time.perf_counter() - start + return 2.0 * size**3 * iterations / elapsed / 1e12 + + +GEMM_OPS = ("aten::mm", "aten::addmm", "aten::bmm", "aten::baddbmm") # leaf GEMM ops; linear/matmul wrap these + + +def print_gemm_shape_table(prof, label, top=12): + """GEMM calls grouped by operand shapes with their CUDA time, so each slow GEMM kernel can be tied to a layer.""" + rows = [e for e in prof.key_averages(group_by_input_shape=True) if e.key in GEMM_OPS and e.device_time_total > 0] + total_us = sum(e.device_time_total for e in rows) + if not rows: + return + print(f"gemm shapes, {label}: {total_us / 1e6:.3f} s of CUDA time in GEMM ops") + print(f" {'share':>6s} {'ms':>9s} {'calls':>7s} {'ms/call':>8s} op input shapes") + for e in sorted(rows, key=lambda e: -e.device_time_total)[:top]: + print( + f" {100 * e.device_time_total / total_us:5.1f}% {e.device_time_total / 1e3:9.1f} {e.count:7d} {e.device_time_total / 1e3 / e.count:8.3f} {e.key} {e.input_shapes}" + ) + + +def timed(fn): + torch.cuda.synchronize() + start = time.perf_counter() + result = fn() + torch.cuda.synchronize() + return time.perf_counter() - start, result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--num-envs", type=int, default=16) + parser.add_argument("--minibatch", type=int, default=131072) + parser.add_argument("--warmup", type=int, default=2, help="epochs before measuring (compile, recompiles)") + parser.add_argument("--epochs", type=int, default=3, help="measured epochs") + parser.add_argument("--fused", action="store_true", help="policy.fused_slot_encoder=true") + parser.add_argument("--rollout-only", action="store_true", help="time only the rollout (data collection)") + parser.add_argument( + "--profile", action="store_true", help="print the trainer's section split (env wait, copy, forward, learn)" + ) + parser.add_argument("--load-model", help="policy checkpoint (.pt) so the rollout shows trained behavior") + parser.add_argument( + "--measure-peak", + action="store_true", + help="time large bf16 matmuls (fp32 accumulate, as the model uses) and print the achieved TFLOP/s next to the " + "table peak the MFU figures divide by", + ) + parser.add_argument( + "--kernel-table", + action="store_true", + help="after warmup, profile one update and one rollout with torch.profiler and print CUDA kernel time by kernel " + "(timeline only, needs no GPU performance-counter permission)", + ) + parser.add_argument("overrides", nargs="*", help="extra hydra overrides, e.g. train.seed=3") + args = parser.parse_args() + + data_dir = tempfile.mkdtemp(prefix="mfu_bench_") + sys.argv = [ + sys.argv[0], + f"vec.num_envs={args.num_envs}", + "env.preload_map_cache=True", + f"train.minibatch_size={args.minibatch}", + f"train.max_minibatch_size={args.minibatch}", + f"train.data_dir={data_dir}", + f"policy.fused_slot_encoder={str(args.fused).lower()}", + "wandb=False", + "tb=False", + "neptune=False", + *args.overrides, + ] + import pufferlib.pufferl as P + + config = P.load_config(ENV_NAME) + config = P.normalize_puffer_drive_config(config, "training") + P.validate_puffer_drive_config(config, "training") + P.validate_puffer_drive_resources(config, "training") + torch_seed, env_seed = P.derive_rank_seeds(config["vec"]["seed"], config["train"]["seed"], 1, 0) + torch.manual_seed(torch_seed) + P.apply_matmul_env_flags() + P.bind_process_to_gpu_numa_node(0) + vecenv = P.load_env(ENV_NAME, config, seed=env_seed) + policy = P.load_policy(config, vecenv, ENV_NAME) + train_config = dict(**config["train"], env=ENV_NAME, eval=config.get("eval", {}), run_name=config["run_name"]) + pufferl = P.PuffeRL(train_config, vecenv, policy) + pufferl.profile.frequency = 1 if args.profile else 10**9 + if args.load_model: + state = torch.load(args.load_model, map_location=config["train"]["device"], weights_only=False) + state = state.get("policy_state_dict", state) + pufferl.uncompiled_policy.load_state_dict(P.clean_policy_state_dict(state)) + + forward_flops, train_flops = model_flops_per_transition(pufferl.uncompiled_policy) + counters = defaultdict(int) + original_ppo_loss = pufferl._ppo_loss + + def counting_ppo_loss(mb_obs, *rest, **kwargs): + counters["samples"] += mb_obs.shape[0] + counters["minibatches"] += 1 + return original_ppo_loss(mb_obs, *rest, **kwargs) + + pufferl._ppo_loss = counting_ppo_loss + properties = torch.cuda.get_device_properties(0) + capability = torch.cuda.get_device_capability(0) + peak_tflops = peak_bf16_tflops(properties.name) + if args.measure_peak: + measured = measure_bf16_peak() + print( + f"measured bf16 matmul peak (fp32 accumulate, {8192}^3 cuBLAS): {measured:.0f} TFLOP/s " + f"(table value {peak_tflops}); MFU below is relative to the measured peak", + flush=True, + ) + peak_tflops = round(measured) + agent_steps = train_config["batch_size"] + print( + f"{properties.name} capability {capability}, {args.num_envs} envs, {vecenv.num_agents} agents, " + f"minibatch {pufferl.minibatch_size}, fused_slot_encoder={args.fused}, " + f"model FLOP per transition: forward {forward_flops / 1e6:.2f} MFLOP, update {train_flops / 1e6:.2f} MFLOP, dense bf16 peak {peak_tflops} TFLOP/s", + flush=True, + ) + + try: + + def update_or_skip(): + if args.rollout_only: + pufferl.epoch += 1 # keep the trainer's epoch bookkeeping moving without an update + return 0.0, None + return timed(pufferl.train) + + for epoch in range(args.warmup): + rollout_seconds, _ = timed(pufferl.evaluate) + update_seconds, _ = update_or_skip() + print(f"warmup epoch {epoch}: rollout {rollout_seconds:.2f} s, update {update_seconds:.2f} s", flush=True) + frames = dict(torch._dynamo.utils.counters["frames"]) + if args.kernel_table: + from torch.profiler import ProfilerActivity, profile + + for label, fn in (("rollout", pufferl.evaluate), ("update", update_or_skip)): + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof: + wall_seconds, _ = timed(fn) + print_kernel_table(prof, label, wall_seconds) + print_gemm_shape_table(prof, label) + profile_baseline = {name: pufferl.profile.profiles[name]["elapsed"] for name in PROFILE_SECTIONS} + epochs = [] + for epoch in range(args.epochs): + counters.clear() + rollout_seconds, _ = timed(pufferl.evaluate) + update_seconds, _ = update_or_skip() + epochs.append((rollout_seconds, update_seconds, counters["samples"])) + print( + f"epoch {epoch}: rollout {rollout_seconds:.2f} s, update {update_seconds:.2f} s, " + f"kept transitions {counters['samples']}", + flush=True, + ) + finally: + pufferl.utilization.stop() + vecenv.close() + + if args.profile: + print("section split per epoch (trainer profiler, CUDA-synchronized at every boundary):") + for name in PROFILE_SECTIONS: + seconds = (pufferl.profile.profiles[name]["elapsed"] - profile_baseline[name]) / len(epochs) + if seconds > 0: + print(f" {name:14s} {seconds:6.2f} s") + if args.rollout_only: + rollout = sum(e[0] for e in epochs) / len(epochs) + print( + f"ROLLOUT_RESULT envs={args.num_envs} agents={vecenv.num_agents} rollout_s={rollout:.3f} agent_steps_per_s={agent_steps / rollout:.0f}" + ) + return + fastest_update = min(update for _, update, _ in epochs) + steady = [e for e in epochs if e[1] < RECOMPILE_EPOCH_FACTOR * fastest_update] + rollout = sum(e[0] for e in steady) / len(steady) + update = sum(e[1] for e in steady) / len(steady) + kept = sum(e[2] for e in steady) / len(steady) + phases = { + "update": (kept * train_flops, update), + "rollout": (agent_steps * forward_flops, rollout), + "full_loop": (kept * train_flops + agent_steps * forward_flops, update + rollout), + } + result = { + "gpu": properties.name, + "capability": list(capability), + "fused": args.fused, + "num_envs": args.num_envs, + "minibatch": pufferl.minibatch_size, + "steady_epochs": len(steady), + "kept_transitions_per_epoch": kept, + "agent_steps_per_epoch": agent_steps, + "sps": agent_steps / (update + rollout), + "peak_tflops": peak_tflops, + "dynamo_frames": frames, + "max_memory_gb": torch.cuda.max_memory_allocated() / 2**30, + } + for phase, (flops, seconds) in phases.items(): + result[f"{phase}_seconds"] = seconds + result[f"{phase}_tflops"] = flops / seconds / 1e12 + if peak_tflops: + result[f"{phase}_mfu"] = 100.0 * flops / seconds / (peak_tflops * 1e12) + print("MFU_RESULT " + json.dumps(result)) + label = "fused" if args.fused else "reference" + for phase in phases: + mfu = f", MFU {result[f'{phase}_mfu']:.1f}%" if peak_tflops else "" + print( + f"{label} {phase:9s}: {result[f'{phase}_seconds']:6.2f} s/epoch, {result[f'{phase}_tflops']:6.1f} TFLOP/s{mfu}" + ) + print(f"{label} SPS {result['sps'] / 1e3:.0f}k agent-steps/s, peak GPU memory {result['max_memory_gb']:.1f} GB") + + +if __name__ == "__main__": + main() diff --git a/scripts/profile_policy_step.py b/scripts/profile_policy_step.py new file mode 100644 index 0000000000..9312b906f7 --- /dev/null +++ b/scripts/profile_policy_step.py @@ -0,0 +1,265 @@ +"""Run one PPO update minibatch and one rollout forward of the real Drive policy inside NVTX ranges for Nsight Compute. + +Usage: + ncu --nvtx --nvtx-include "update/" --nvtx-include "rollout/" --nvtx-include "calibrate/" \\ + --metrics gpu__time_duration.sum,sm__ops_path_tensor_src_bf16_dst_fp32.sum,sm__ops_path_tensor_src_fp16_dst_fp32.sum,smsp__sass_thread_inst_executed_op_ffma_pred_on.sum,dram__bytes.sum \\ + --csv --page raw python scripts/profile_policy_step.py --rows 65536 [--fused] > ncu_out.csv + python scripts/profile_policy_step.py --parse ncu_out.csv --rows 65536 +The calibrate range runs a matmul of known size so the tensor-op metric is converted to FLOPs without guessing. +""" + +import argparse +import copy +import csv +import json +import os +import sys + + +os.environ.setdefault("OMP_NUM_THREADS", "1") +import numpy as np +import torch + + +TRAIN_FLOP_PER_TRANSITION = 112145920.0 +ROLLOUT_FLOP_PER_TRANSITION = 37654113.0 +CALIBRATE_N = 8192 +CALIBRATE_FLOPS = 2.0 * CALIBRATE_N**3 +# dense bf16 tensor peak with fp32 accumulation, by device-name prefix; the compute capability does not +# identify the GPU: (12, 0) is the RTX 5090 (GeForce, half rate with fp32 accumulation) and the RTX PRO 6000 +PEAK_BF16_TFLOPS_BY_DEVICE_NAME = { + "NVIDIA GeForce RTX 5090": 209.5, + "NVIDIA RTX PRO 6000 Blackwell Server Edition": 471.0, + "NVIDIA H100 80GB HBM3": 989.0, + "NVIDIA B200": 2250.0, + "NVIDIA B300": 2250.0, +} +ROLLOUT_ROWS = 16384 +WARMUP_STEPS = 4 +TIMED_STEPS = 5 +ENV_AGENTS = 512 +ROLLOUT_STEPS = 40 +TRAIN_OVERRIDES = [ + "vec.backend=Serial", + "vec.num_envs=1", + "vec.num_workers=1", + "vec.batch_size=1", + f"env.num_agents={ENV_AGENTS}", + "train.minibatch_size=8192", + "train.max_minibatch_size=8192", + "wandb=False", + "tb=False", + "neptune=False", +] + + +def peak_bf16_tflops(device_name): + for name_prefix, peak in PEAK_BF16_TFLOPS_BY_DEVICE_NAME.items(): + if device_name.startswith(name_prefix): + return peak + print(f"no bf16 peak known for {device_name}, MFU not reported; add it to PEAK_BF16_TFLOPS_BY_DEVICE_NAME") + return None + + +def build_policy_and_obs(rows, fused): + sys.argv = ["profile_policy_step"] + TRAIN_OVERRIDES + import pufferlib.pufferl as P + + args = P.load_config("puffer_drive") + args = P.normalize_puffer_drive_config(args, "training") + P.validate_puffer_drive_config(args, "training") + args = copy.deepcopy(args) + args["policy"]["fused_slot_encoder"] = fused + vecenv = P.load_env("puffer_drive", args, seed=3) + torch.manual_seed(0) + policy = P.load_policy(args, vecenv, "puffer_drive") + vecenv.async_reset(3) + obs_rows = [] + for _ in range(ROLLOUT_STEPS): + o, *_ = vecenv.recv() + obs_rows.append(np.array(o, copy=True)) + vecenv.send(np.random.uniform(-1, 1, size=vecenv.action_space.shape).astype(vecenv.action_space.dtype)) + vecenv.close() + obs = torch.as_tensor(np.concatenate(obs_rows)).cuda() + repeats = -(-max(rows, ROLLOUT_ROWS) // obs.shape[0]) + return policy, obs.repeat(repeats, 1) + + +def fse_config_name(): + import pufferlib.ocean.fused_slot_encoder as fse + + if fse.CONFIG_OVERRIDE is not None: + return "override" + return "table" + + +def apply_kernel_config(path): + import pufferlib.ocean.fused_slot_encoder as fse + + with open(path) as handle: + fse.CONFIG_OVERRIDE = fse.KernelConfig(**json.load(handle)) + print(f"kernel config override: {fse.CONFIG_OVERRIDE}") + + +def run_phases(rows, fused): + import pufferlib.pytorch + + torch.set_float32_matmul_precision("highest") + torch.backends.cuda.matmul.allow_tf32 = False + policy, obs = build_policy_and_obs(rows, fused) + compiled = torch.compile(policy) + compiled_eval = torch.compile(policy.forward_eval) + sample_logits = torch.compile(pufferlib.pytorch.sample_logits) + optimizer = torch.optim.AdamW(policy.parameters(), lr=1e-4, weight_decay=0.01) + update_obs = obs[:rows] + actions = torch.randint(0, policy.action_dim, (rows,), device="cuda", dtype=torch.int32) + rollout_obs = obs[:ROLLOUT_ROWS] + amp = torch.autocast("cuda", dtype=torch.bfloat16) + + def update_step(): + optimizer.zero_grad(set_to_none=True) + with amp: + logits, value = compiled(update_obs, {"action": actions, "lstm_h": None, "lstm_c": None}) + logits = logits.float() + _, logprob, entropy, _ = sample_logits(logits, action=actions) + loss = -logprob.float().mean() + 0.5 * value.float().square().mean() - 0.01 * entropy.mean() + loss.backward() + torch.nn.utils.clip_grad_norm_(policy.parameters(), 0.5) + optimizer.step() + + def rollout_step(): + with torch.no_grad(), amp: + logits, _ = compiled_eval(rollout_obs, {}) + sample_logits(logits.float(), env_continuous=True, policy=policy) + + for _ in range(WARMUP_STEPS): + update_step() + rollout_step() + wall_ms = {} + for name, step in (("update", update_step), ("rollout", rollout_step)): + samples = [] + for _ in range(TIMED_STEPS): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + step() + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end)) + wall_ms[name] = sorted(samples)[len(samples) // 2] + print( + f"unprofiled wall time (median of {TIMED_STEPS}): update {wall_ms['update']:.2f} ms, rollout {wall_ms['rollout']:.2f} ms" + ) + peak = peak_bf16_tflops(torch.cuda.get_device_name(0)) + update_rate = rows * TRAIN_FLOP_PER_TRANSITION / (wall_ms["update"] * 1e-3) / 1e12 + rollout_rate = ROLLOUT_ROWS * ROLLOUT_FLOP_PER_TRANSITION / (wall_ms["rollout"] * 1e-3) / 1e12 + mfu = f" update_mfu={100 * update_rate / peak:.1f}% rollout_mfu={100 * rollout_rate / peak:.1f}%" if peak else "" + print( + f"RESULT fused={int(fused)} rows={rows} update_ms={wall_ms['update']:.2f} rollout_ms={wall_ms['rollout']:.2f}" + f" update_tflops={update_rate:.1f} rollout_tflops={rollout_rate:.1f}{mfu} config={fse_config_name()}" + ) + a = torch.randn(CALIBRATE_N, CALIBRATE_N, device="cuda", dtype=torch.bfloat16) + b = torch.randn_like(a) + a @ b + torch.cuda.synchronize() + # Start/end ranges are process-wide and also cover the autograd thread's kernels; push/pop ranges are per thread. + for name, step in (("calibrate", lambda: a @ b), ("update", update_step), ("rollout", rollout_step)): + range_id = torch.cuda.nvtx.range_start(name) + step() + torch.cuda.nvtx.range_end(range_id) + torch.cuda.synchronize() + print( + f"analytic model FLOPs: update {rows * TRAIN_FLOP_PER_TRANSITION:.4e}, rollout {ROLLOUT_ROWS * ROLLOUT_FLOP_PER_TRANSITION:.4e}, calibrate {CALIBRATE_FLOPS:.4e}" + ) + + +METRIC_COLUMNS = { + "time": "gpu__time_duration.sum", + "bf16": "sm__ops_path_tensor_op_hmma_src_bf16_dst_fp32.sum", + "fp16": "sm__ops_path_tensor_op_hmma_src_fp16_dst_fp32.sum", + "tf32": "sm__ops_path_tensor_op_hmma_src_tf32_dst_fp32.sum", + "ffma": "sm__sass_thread_inst_executed_op_ffma_pred_on.sum", + "fadd": "sm__sass_thread_inst_executed_op_fadd_pred_on.sum", + "fmul": "sm__sass_thread_inst_executed_op_fmul_pred_on.sum", + "dram": "dram__bytes.sum", + "local_ld": "sm__sass_inst_executed_op_local_ld.sum", + "local_st": "sm__sass_inst_executed_op_local_st.sum", +} +HEAD_FLOP_PER_TRANSITION = 3166208.0 +ROLLOUT_HEAD_FLOP_PER_TRANSITION = 1055232.0 + + +RANGE_NAMES = ("calibrate", "update", "rollout") + + +def _phase_of(record, kernel_col): + kernel = record[kernel_col] + if "/" in kernel: + return kernel.split("/")[0] + text = " ".join(record) + for name in RANGE_NAMES: + if f":{name}:" in text: + return name + return "unlabeled" + + +def parse(path, rows): + with open(path) as handle: + table = list(csv.reader(line for line in handle if line.startswith('"'))) + header, data_rows = table[0], table[2:] + kernel_col = header.index("Kernel Name") + columns = {key: header.index(name) for key, name in METRIC_COLUMNS.items() if name in header} + sums = {} + for record in data_rows: + phase = _phase_of(record, kernel_col) + phase_sums = sums.setdefault(phase, {key: 0.0 for key in columns}) + phase_sums["kernels"] = phase_sums.get("kernels", 0) + 1 + for key, col in columns.items(): + phase_sums[key] += float(record[col].replace(",", "") or 0.0) + calibrate = sums.get("calibrate") + if calibrate: + print(f"calibration: counter {calibrate['bf16']:.6e} FLOP vs exact {CALIBRATE_FLOPS:.6e} for the 8192^3 matmul") + model_tensor = { + "update": rows * (TRAIN_FLOP_PER_TRANSITION - HEAD_FLOP_PER_TRANSITION), + "rollout": ROLLOUT_ROWS * (ROLLOUT_FLOP_PER_TRANSITION - ROLLOUT_HEAD_FLOP_PER_TRANSITION), + "calibrate": CALIBRATE_FLOPS, + } + model_fp32 = {"update": rows * HEAD_FLOP_PER_TRANSITION, "rollout": ROLLOUT_ROWS * ROLLOUT_HEAD_FLOP_PER_TRANSITION} + peak = peak_bf16_tflops(torch.cuda.get_device_name(0)) if torch.cuda.is_available() else None + for phase, m in sums.items(): + seconds = m["time"] * 1e-9 + tensor = m["bf16"] + m.get("fp16", 0.0) + m.get("tf32", 0.0) + fp32 = 2.0 * m.get("ffma", 0.0) + m.get("fadd", 0.0) + m.get("fmul", 0.0) + print(f"== {phase}: {m['kernels']} kernels, summed kernel time {seconds * 1e3:.2f} ms") + print( + f" tensor FLOP counted {tensor:.4e} (bf16 {m['bf16']:.3e}, tf32 {m.get('tf32', 0.0):.2e}) vs model GEMM FLOP {model_tensor.get(phase, 0.0):.4e}" + ) + print(f" fp32 CUDA-core FLOP counted {fp32:.4e} vs model fp32 head FLOP {model_fp32.get(phase, 0.0):.3e}") + tensor_rate = tensor / seconds / 1e12 + model_rate = model_tensor.get(phase, 0.0) / seconds / 1e12 + tensor_share = f" = {100 * tensor_rate / peak:.1f}% of bf16 peak" if peak else "" + model_share = f" = {100 * model_rate / peak:.1f}%" if peak else "" + print( + f" tensor-FLOP rate over kernel time {tensor_rate:.1f} TFLOP/s{tensor_share}; model-FLOP rate {model_rate:.1f} TFLOP/s{model_share}" + ) + print( + f" DRAM {m['dram'] / 1e9:.2f} GB -> {m['dram'] / seconds / 1e9:.0f} GB/s; local (spill) ld/st instructions {m.get('local_ld', 0.0):.3e} / {m.get('local_st', 0.0):.3e}" + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--rows", type=int, default=65536, help="update minibatch rows") + parser.add_argument("--fused", action="store_true", help="enable policy.fused_slot_encoder") + parser.add_argument("--parse", help="ncu raw-page CSV to summarize instead of running") + parser.add_argument("--kernel-config", help="JSON KernelConfig (from tune_fused_slot_encoder.py --write-config)") + args = parser.parse_args() + if args.parse: + parse(args.parse, args.rows) + return + if args.kernel_config: + apply_kernel_config(args.kernel_config) + run_phases(args.rows, args.fused) + + +if __name__ == "__main__": + main() diff --git a/scripts/tune_fused_slot_encoder.py b/scripts/tune_fused_slot_encoder.py new file mode 100644 index 0000000000..2002bcb520 --- /dev/null +++ b/scripts/tune_fused_slot_encoder.py @@ -0,0 +1,204 @@ +"""Time fused_slot_encoder kernel configurations on the current GPU and print one KernelConfig for its GPU class. + +Usage (needs CUDA + Triton; one process per GPU type): + python scripts/tune_fused_slot_encoder.py # all encoder shapes, 131072 transitions + python scripts/tune_fused_slot_encoder.py --batch 65536 --quick +Paste the printed KernelConfig into CONFIGS_BY_CAPABILITY in pufferlib/ocean/fused_slot_encoder.py. +""" + +import argparse +import dataclasses +import itertools +import json +import time + +import torch +import triton +from torch import nn + +import pufferlib.ocean.fused_slot_encoder as fse +import pufferlib.pytorch + + +ENCODER_SHAPES = {"lane": (35, 10, 256), "boundary": (30, 6, 256), "partner": (20, 11, 256), "traffic": (4, 14, 128)} +OBS_PREFIX = 100 +TIMING_ITERS = 15 +WARMUP_ITERS = 3 +POOLED_TOLERANCE = 2.0**-6 +GRAD_TOLERANCE = 1e-3 +FWD_GRID = {"block_batch": (16, 32, 64), "block_k": (32, 64), "num_warps": (4, 8), "num_stages": (2, 3)} +BWD_GRID = {"block_rows": (32, 64, 128), "block_k": (32, 64), "num_warps": (4, 8), "num_stages": (2, 3)} +POOL_GRID = {"block_rows": (16, 32, 64), "num_warps": (4, 8)} +QUICK_FWD_GRID = {"block_batch": (32,), "block_k": (32, 64), "num_warps": (4,), "num_stages": (2,)} +QUICK_BWD_GRID = {"block_rows": (64,), "block_k": (64,), "num_warps": (8,), "num_stages": (2,)} +QUICK_POOL_GRID = {"block_rows": (32,), "num_warps": (4,)} +LAUNCH_ERRORS = (triton.runtime.errors.OutOfResources, triton.compiler.errors.CompilationError, RuntimeError) +PARTS = ("fwd", "bwd", "pool") + + +def make_case(slots, features, width, batch, seed): + generator = torch.Generator(device="cuda").manual_seed(seed) + encoder = nn.Sequential( + pufferlib.pytorch.layer_init(nn.Linear(features, width)), + nn.LayerNorm(width), + nn.ReLU(), + pufferlib.pytorch.layer_init(nn.Linear(width, width)), + ).cuda() + obs = torch.zeros(batch, OBS_PREFIX + slots * features, device="cuda") + counts = torch.randint(0, slots + 1, (batch,), device="cuda", generator=generator) + objects = obs[:, OBS_PREFIX:].view(batch, slots, features) + valid = torch.arange(slots, device="cuda")[None, :] < counts[:, None] + objects.copy_( + (torch.rand(batch, slots, features, device="cuda", generator=generator) * 1.6 - 0.8) * valid[:, :, None] + ) + grad_pooled = (torch.randn(batch, width, device="cuda", generator=generator) * 0.1).to(torch.bfloat16) + return encoder, objects, counts.long(), grad_pooled + + +def time_call(fn): + for _ in range(WARMUP_ITERS): + fn() + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(TIMING_ITERS): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - start) / TIMING_ITERS * 1e3 + + +def with_part(base, part, params): + fields = {name: getattr(base, name) for name in base.__dataclass_fields__} + fields.update({f"{part}_{key}": value for key, value in params.items()}) + return fse.KernelConfig(**fields) + + +class ShapeCase: + def __init__(self, name, batch): + slots, features, width = ENCODER_SHAPES[name] + self.name = name + self.encoder, self.objects, self.counts, self.grad_pooled = make_case(slots, features, width, batch, seed=slots) + fse.CONFIG_OVERRIDE = fse.GB202_CONFIG + self.reference_fwd = self.forward() + self.reference_bwd = self.backward(self.reference_fwd) + fse.CONFIG_OVERRIDE = None + + def forward(self): + linear_in, norm, _, linear_out = self.encoder + return fse.fused_slot_encoder( + self.objects, + self.counts, + linear_in.weight, + linear_in.bias, + norm.weight, + norm.bias, + linear_out.weight, + linear_out.bias, + True, + ) + + def backward(self, fwd): + pooled, argmax, normed_relu, w1_padded, b1_values, w2_bf16 = fwd + _, norm, _, _ = self.encoder + return fse.fused_slot_encoder_bwd( + self.objects, + self.counts, + w1_padded, + b1_values, + norm.weight, + norm.bias, + w2_bf16, + normed_relu, + argmax, + self.grad_pooled, + ) + + def measure(self, part, config): + """Milliseconds for one part under `config`, or None when its numerics differ from the default config.""" + fse.CONFIG_OVERRIDE = config + try: + if part == "fwd": + pooled = self.forward()[0] + if (pooled.float() - self.reference_fwd[0].float()).abs().max().item() > POOLED_TOLERANCE: + return None + return time_call(self.forward) + grads = self.backward(self.reference_fwd) + for candidate, reference in zip(grads, self.reference_bwd): + if ((candidate - reference).norm() / reference.norm().clamp_min(1e-12)).item() > GRAD_TOLERANCE: + return None + return time_call(lambda: self.backward(self.reference_fwd)) + finally: + fse.CONFIG_OVERRIDE = None + + +def sweep(cases, part, grid, base): + totals = {} + per_shape_best = {case.name: (float("inf"), None) for case in cases} + for values in itertools.product(*grid.values()): + params = dict(zip(grid.keys(), values)) + config = with_part(base, part, params) + total = 0.0 + for case in cases: + try: + ms = case.measure(part, config) + except LAUNCH_ERRORS as exc: + ms = None + print(f" {part} {params} on {case.name}: skipped ({str(exc).splitlines()[0][:70]})", flush=True) + if ms is None: + total = None + break + total += ms + if ms < per_shape_best[case.name][0]: + per_shape_best[case.name] = (ms, params) + if total is not None: + totals[tuple(params.items())] = total + print(f" {part} {params}: {total:7.3f} ms over all shapes", flush=True) + best_key = min(totals, key=totals.get) + return dict(best_key), totals[best_key], per_shape_best + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--shapes", default="all", help="comma-separated encoder names or 'all'") + parser.add_argument("--batch", type=int, default=131072, help="transitions per minibatch") + parser.add_argument("--quick", action="store_true", help="tiny grids, for checking that the script runs") + parser.add_argument("--write-config", help="write the recommended KernelConfig as JSON to this path") + args = parser.parse_args() + names = list(ENCODER_SHAPES) if args.shapes == "all" else args.shapes.split(",") + grids = {"fwd": FWD_GRID, "bwd": BWD_GRID, "pool": POOL_GRID} + if args.quick: + grids = {"fwd": QUICK_FWD_GRID, "bwd": QUICK_BWD_GRID, "pool": QUICK_POOL_GRID} + properties = torch.cuda.get_device_properties(0) + capability = torch.cuda.get_device_capability(0) + device_line = ( + f"{properties.name}, capability {capability}, {properties.multi_processor_count} SMs, " + f"{properties.shared_memory_per_block_optin // 1024} KB shared memory per block, " + f"torch {torch.__version__}, triton {triton.__version__}" + ) + print(device_line, flush=True) + cases = [ShapeCase(name, args.batch) for name in names] + base = fse.GB202_CONFIG + recommended = base + report_lines = [] + for part in PARTS: + best_params, best_total, per_shape = sweep(cases, part, grids[part], base) + baseline_total = sum(case.measure(part, base) for case in cases) + recommended = with_part(recommended, part, best_params) + report_lines.append(f"{part}: GB202_CONFIG {baseline_total:.3f} ms -> best {best_total:.3f} ms {best_params}") + for name, (ms, params) in per_shape.items(): + report_lines.append(f" {part} best for {name}: {ms:.3f} ms {params}") + combined_total = sum(case.measure(part, recommended) for case in cases for part in PARTS) + baseline_combined = sum(case.measure(part, base) for case in cases for part in PARTS) + print("\n=== fused_slot_encoder tuning result (paste this block back) ===") + print(device_line) + print(f"batch {args.batch}, shapes {','.join(names)}") + print("\n".join(report_lines)) + print(f"all parts, all shapes: GB202_CONFIG {baseline_combined:.3f} ms -> recommended {combined_total:.3f} ms") + print(f"RECOMMENDED for capability {capability}: {recommended}") + print("=== end of tuning result ===") + if args.write_config: + with open(args.write_config, "w") as handle: + json.dump(dataclasses.asdict(recommended), handle) + + +if __name__ == "__main__": + main() diff --git a/scripts/visualize_map_bins.py b/scripts/visualize_map_bins.py new file mode 100644 index 0000000000..35186aa23d --- /dev/null +++ b/scripts/visualize_map_bins.py @@ -0,0 +1,613 @@ +#!/usr/bin/env python +"""Render PufferDrive map .bin files into one self-contained HTML map inspector (lanes, road lines, road edges), +with click-to-mark: every click drops a numbered mark and appends a line with the bin coordinates and the nearest +road edge / lane to a report you can copy. Optionally overlays the road edges of a second bin directory for comparison. +The view is a 3D orbit camera: right/middle drag tilts and turns the map, a slider exaggerates heights, edges can be +coloured by height, and road edges whose height steps between vertical layers (bridge seams) are listed and ringed. + + python scripts/visualize_map_bins.py --bin-dir ~/ordnung/data/CARLA/puffer_bins_nd_shoulders \\ + --compare-dir ~/ordnung/data/CARLA/puffer_bins_stops_tol01 --output /tmp/map_inspector.html + +Keys: 1/2/3 mark category, u undo, c copy report, f fit, t top view, Esc clear. Left drag pans, right/middle drag +rotates, wheel zooms. +""" + +import argparse +import json +import math +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / "data_utils")) +from mirror_map_bin import read_bin # noqa: E402 + +LANE_TYPE_MAX = 9 +ROAD_LINE_TYPES = range(10, 20) +ROAD_EDGE_TYPES = (20, 21, 22) +SEAM_MIN_DZ_M = 1.5 # a road-edge segment climbing more than this ... +SEAM_MIN_GRADE = 0.5 # ... at a grade steeper than any road steps between vertical layers + + +def _polyline(road): + return [ + [round(float(x), 2), round(float(y), 2), round(float(z), 2)] for x, y, z in zip(road["x"], road["y"], road["z"]) + ] + + +def _seams(edges, kind): + seams = [] + for edge in edges: + for (x0, y0, z0), (x1, y1, z1) in zip(edge["pts"], edge["pts"][1:]): + dz = abs(z1 - z0) + if dz > SEAM_MIN_DZ_M and dz > SEAM_MIN_GRADE * max(math.hypot(x1 - x0, y1 - y0), 0.05): + seams.append( + { + "kind": kind, + "id": edge["id"], + "x": round((x0 + x1) / 2, 2), + "y": round((y0 + y1) / 2, 2), + "z0": z0, + "z1": z1, + } + ) + return seams + + +def _z_range(polylines): + zs = [p[2] for pts in polylines for p in pts] + return [min(zs), max(zs)] if zs else [0.0, 0.0] + + +def _edges(data): + return [ + {"id": int(r["id"]), "type": int(r["type"]), "pts": _polyline(r)} + for r in data["roads"] + if r["type"] in ROAD_EDGE_TYPES + ] + + +def collect_map(bin_file: pathlib.Path, compare_file: pathlib.Path | None) -> dict: + data = read_bin(bin_file) + lanes, lines = [], [] + for road in data["roads"]: + if 0 <= road["type"] <= LANE_TYPE_MAX: + lanes.append( + { + "id": int(road["id"]), + "junction": bool(data.get("has_zone_section")) and road.get("speed_zone_idx", 0) < 0, + "length": round(float(road["length"]), 1), + "pts": _polyline(road), + } + ) + elif road["type"] in ROAD_LINE_TYPES: + lines.append(_polyline(road)) + compare = _edges(read_bin(compare_file)) if compare_file is not None else [] + edges = _edges(data) + return { + "name": bin_file.stem.split("__")[-1], + "file": bin_file.name, + "lanes": lanes, + "lines": lines, + "edges": edges, + "compare": compare, + "seams": _seams(edges, "edge"), + "compareSeams": _seams(compare, "old edge"), + "zrange": _z_range([lane["pts"] for lane in lanes] + [edge["pts"] for edge in edges]), + } + + +STYLE = """ +:root { + --ground: #F2F3F1; --panel: #FAFAF8; --line: #D9DCD8; --text: #1C2126; --muted: #6E7882; + --accent: #C9552B; --accent-soft: rgba(201, 85, 43, 0.12); --lane: #B9BEB9; --junction: #D3D7D2; --roadline: #E3E6E2; + --edge: #C9552B; --median: #8C3D1E; --compare: #E6A23C; --row-hover: #ECEEEA; --tooltip: #1C2126; --tooltip-text: #F2F3F1; + --cat1: #C9552B; --cat2: #2F6FB3; --cat3: #B07C1A; + color-scheme: light; +} +@media (prefers-color-scheme: dark) { + :root:not([data-theme="light"]) { + --ground: #141719; --panel: #1B1F22; --line: #2C3237; --text: #E4E8EA; --muted: #8B95A0; + --accent: #E2764B; --accent-soft: rgba(226, 118, 75, 0.18); --lane: #4A5159; --junction: #363C42; --roadline: #2A3035; + --edge: #E2764B; --median: #C45A33; --compare: #E6B45C; --row-hover: #23282C; --tooltip: #E4E8EA; --tooltip-text: #141719; + --cat1: #E2764B; --cat2: #6FA5E6; --cat3: #D9A43A; + color-scheme: dark; + } +} +:root[data-theme="dark"] { + --ground: #141719; --panel: #1B1F22; --line: #2C3237; --text: #E4E8EA; --muted: #8B95A0; + --accent: #E2764B; --accent-soft: rgba(226, 118, 75, 0.18); --lane: #4A5159; --junction: #363C42; --roadline: #2A3035; + --edge: #E2764B; --median: #C45A33; --compare: #E6B45C; --row-hover: #23282C; --tooltip: #E4E8EA; --tooltip-text: #141719; + --cat1: #E2764B; --cat2: #6FA5E6; --cat3: #D9A43A; + color-scheme: dark; +} +* { box-sizing: border-box; } +html, body { height: 100%; } +body { margin: 0; background: var(--ground); color: var(--text); font-family: "IBM Plex Sans", "Segoe UI", system-ui, sans-serif; font-size: 13px; line-height: 1.4; } +.app { display: grid; grid-template-rows: auto 1fr; grid-template-columns: 1fr 360px; height: 100vh; min-height: 560px; } +header { grid-column: 1 / -1; display: flex; align-items: center; gap: 16px; padding: 10px 16px; border-bottom: 1px solid var(--line); background: var(--panel); flex-wrap: wrap; } +header h1 { font-size: 15px; font-weight: 600; margin: 0; white-space: nowrap; } +.tabs { display: flex; gap: 4px; flex-wrap: wrap; } +.tabs button, .seg button, .controls button { font: inherit; color: var(--text); background: transparent; border: 1px solid var(--line); border-radius: 4px; padding: 4px 10px; cursor: pointer; } +.tabs button[aria-pressed="true"], .seg button[aria-pressed="true"] { background: var(--accent-soft); border-color: var(--accent); } +.seg { display: flex; } +.seg button:first-child { border-radius: 4px 0 0 4px; } +.seg button:last-child { border-radius: 0 4px 4px 0; } +.seg button:not(:first-child) { margin-left: -1px; } +.controls { display: flex; align-items: center; gap: 12px; margin-left: auto; color: var(--muted); flex-wrap: wrap; } +.controls label { display: flex; align-items: center; gap: 5px; cursor: pointer; } +.stage { position: relative; overflow: hidden; } +canvas { display: block; width: 100%; height: 100%; cursor: crosshair; } +.hint { position: absolute; left: 12px; bottom: 10px; color: var(--muted); font-size: 12px; pointer-events: none; } +.legend { position: absolute; right: 12px; top: 10px; display: flex; flex-direction: column; gap: 3px; font-family: "IBM Plex Mono", ui-monospace, monospace; font-size: 11px; pointer-events: none; } +.legend span { display: inline-flex; align-items: center; gap: 6px; } +.legend i { width: 18px; height: 3px; border-radius: 2px; display: inline-block; } +.tooltip { position: absolute; pointer-events: none; background: var(--tooltip); color: var(--tooltip-text); padding: 6px 9px; border-radius: 4px; font-family: "IBM Plex Mono", ui-monospace, monospace; font-size: 11.5px; white-space: nowrap; transform: translate(12px, 12px); line-height: 1.5; } +.tooltip[hidden] { display: none; } +aside { border-left: 1px solid var(--line); background: var(--panel); display: flex; flex-direction: column; min-height: 0; } +.stats { display: grid; grid-template-columns: 1fr 1fr; gap: 8px 12px; padding: 12px 14px; border-bottom: 1px solid var(--line); } +.stat .k { color: var(--muted); font-size: 11px; text-transform: uppercase; letter-spacing: 0.06em; } +.stat .v { font-family: "IBM Plex Mono", ui-monospace, monospace; font-size: 15px; font-variant-numeric: tabular-nums; } +.seams { max-height: 170px; overflow: auto; border-bottom: 1px solid var(--line); padding: 6px 0; } +.seam-row { padding: 4px 14px; font-family: "IBM Plex Mono", ui-monospace, monospace; font-size: 11.5px; cursor: pointer; } +.seam-row:hover { background: var(--row-hover); } +.seam-row.old { color: var(--muted); } +.seam-row.none { color: var(--muted); cursor: default; } +.controls input[type="range"] { width: 90px; accent-color: var(--accent); } +.marks { overflow: auto; flex: 1; min-height: 0; padding: 6px 0; } +.mark-row { display: grid; grid-template-columns: 26px 1fr 22px; gap: 6px; align-items: start; padding: 5px 14px; font-family: "IBM Plex Mono", ui-monospace, monospace; font-size: 11.5px; cursor: pointer; } +.mark-row:hover { background: var(--row-hover); } +.mark-row .n { display: inline-flex; width: 20px; height: 20px; border-radius: 50%; color: #fff; align-items: center; justify-content: center; font-size: 11px; } +.mark-row .d { color: var(--muted); } +.mark-row button { font: inherit; border: 0; background: transparent; color: var(--muted); cursor: pointer; } +.report { border-top: 1px solid var(--line); padding: 10px 14px; display: flex; flex-direction: column; gap: 6px; } +.report textarea { width: 100%; height: 120px; font-family: "IBM Plex Mono", ui-monospace, monospace; font-size: 11px; background: var(--ground); color: var(--text); border: 1px solid var(--line); border-radius: 4px; padding: 6px; resize: vertical; } +.report .row { display: flex; gap: 8px; align-items: center; color: var(--muted); font-size: 11.5px; } +.note { padding: 8px 14px; color: var(--muted); font-size: 11.5px; border-top: 1px solid var(--line); } +@media (max-width: 860px) { .app { grid-template-columns: 1fr; grid-template-rows: auto 1fr 360px; } aside { border-left: 0; border-top: 1px solid var(--line); } } +""" + +SCRIPT = r""" +const MAPS = window.__MAP_INSPECTOR__.maps; +const META = window.__MAP_INSPECTOR__.meta; +const CATEGORIES = { "1": ["not_drivable", "--cat1", "NOT drivable in CARLA but drivable in the map"], + "2": ["missing_drivable", "--cat2", "drivable in CARLA but blocked in the map"], + "3": ["other", "--cat3", "other issue"] }; +const MAX_PITCH = 80 * Math.PI / 180; +const ORBIT_RAD_PER_PX = 0.008; +const FOV_DEG = 50; +const NEAR_FRACTION = 0.05; +const PICK_PX = 8; +const VERTEX_PICK_PX = 12; +const state = { map: 0, category: "1", hover: null, view: null, compare: true, lines: true, lanes: true, heightColor: false, zscale: 1, marks: MAPS.map(() => []) }; +const canvas = document.getElementById("map"); +const ctx = canvas.getContext("2d"); +const tooltip = document.getElementById("tooltip"); +const css = () => getComputedStyle(document.documentElement); +const color = v => css().getPropertyValue(v).trim(); +function current() { return MAPS[state.map]; } + +// Perspective orbit camera: the target (cx, cy, cz) sits at the screen centre, scale is px per metre at the target, +// yaw turns the map about the vertical axis and pitch tilts the camera (0 = top-down). +function cameraDistance(v) { return (canvas.clientHeight / v.scale) / (2 * Math.tan(FOV_DEG * Math.PI / 360)); } +function project(p, v) { + const ca = Math.cos(v.yaw), sa = Math.sin(v.yaw), cb = Math.cos(v.pitch), sb = Math.sin(v.pitch); + const rx = p[0] - v.cx, ry = p[1] - v.cy, rz = ((p[2] || 0) - v.cz) * state.zscale; + const x1 = rx * ca - ry * sa, y1 = rx * sa + ry * ca; + const dist = cameraDistance(v), depth = dist + y1 * sb - rz * cb; + if (depth < NEAR_FRACTION * dist) return null; + const f = v.scale * dist / depth; + return [canvas.clientWidth / 2 + f * x1, canvas.clientHeight / 2 - f * (y1 * cb + rz * sb)]; +} +function toScreen(p) { return project(p, state.view); } +function toWorld(sx, sy, z = 0) { + const v = state.view, ca = Math.cos(v.yaw), sa = Math.sin(v.yaw), cb = Math.cos(v.pitch), sb = Math.sin(v.pitch); + const dist = cameraDistance(v), rz = (z - v.cz) * state.zscale; + const u = (sx - canvas.clientWidth / 2) / (v.scale * dist), w = (canvas.clientHeight / 2 - sy) / (v.scale * dist); + const denom = w * sb - cb; + if (Math.abs(denom) < 1e-9) return null; + const y1 = (rz * (sb + w * cb) - w * dist) / denom; + const depth = dist + y1 * sb - rz * cb; + if (depth <= 0) return null; + const x1 = u * depth; + return [v.cx + x1 * ca + y1 * sa, v.cy - x1 * sa + y1 * ca, z]; +} +function groundOffset(sx, sy) { + const w = toWorld(sx, sy, state.view.cz); + return w ? [w[0] - state.view.cx, w[1] - state.view.cy] : null; +} +function allPoints(m) { const out = []; for (const l of m.lanes) out.push(...l.pts); for (const e of m.edges) out.push(...e.pts); return out; } +function fitPoints(points, pad) { + const w = canvas.clientWidth, h = canvas.clientHeight; + const yaw = state.view ? state.view.yaw : 0, pitch = state.view ? state.view.pitch : 0; + let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity, sumZ = 0; + for (const p of points) { + if (p[0] < minX) minX = p[0]; if (p[0] > maxX) maxX = p[0]; if (p[1] < minY) minY = p[1]; if (p[1] > maxY) maxY = p[1]; sumZ += p[2] || 0; + } + const cx = (minX + maxX) / 2, cy = (minY + maxY) / 2, cz = points.length ? sumZ / points.length : 0; + const ca = Math.cos(yaw), sa = Math.sin(yaw), cb = Math.cos(pitch), sb = Math.sin(pitch); + let ex = 1, ey = 1; + for (const p of points) { + const rx = p[0] - cx, ry = p[1] - cy, rz = ((p[2] || 0) - cz) * state.zscale; + ex = Math.max(ex, Math.abs(rx * ca - ry * sa)); ey = Math.max(ey, Math.abs((rx * sa + ry * ca) * cb + rz * sb)); + } + const scale = Math.min((w - 2 * pad) / (2 * ex), (h - 2 * pad) / (2 * ey)) * (1 - 0.35 * sb); + state.view = { cx, cy, cz, scale, yaw, pitch }; +} +function fitView() { fitPoints(allPoints(current()), 30); } +function fitBox(box, pad, z = 0) { fitPoints([[box[0], box[1], z], [box[2], box[1], z], [box[2], box[3], z], [box[0], box[3], z]], pad); } +function orbit(dyaw, dpitch) { const v = state.view; v.yaw += dyaw; v.pitch = Math.max(0, Math.min(MAX_PITCH, v.pitch + dpitch)); } +function topView() { state.view.yaw = 0; state.view.pitch = 0; } +function setZScale(z) { state.zscale = z; document.getElementById("zscale-val").textContent = `${z}×`; } +function viewLabel() { const v = state.view, deg = r => Math.round(r * 180 / Math.PI); return v.pitch < 1e-3 ? "top view" : `yaw ${deg(v.yaw)}° · tilt ${deg(v.pitch)}°`; } + +function resize() { + const dpr = window.devicePixelRatio || 1; + canvas.width = Math.round(canvas.clientWidth * dpr); + canvas.height = Math.round(canvas.clientHeight * dpr); + ctx.setTransform(dpr, 0, 0, dpr, 0, 0); + if (!state.view) fitView(); + draw(); +} + +function strokePolyline(pts) { + ctx.beginPath(); + let pen = false; + for (const p of pts) { + const s = toScreen(p); + if (!s) { pen = false; continue; } + if (pen) ctx.lineTo(s[0], s[1]); else ctx.moveTo(s[0], s[1]); + pen = true; + } + ctx.stroke(); +} +function heightColor(t) { return `hsl(${Math.round(230 - 230 * Math.max(0, Math.min(1, t)))}, 85%, 45%)`; } +function strokeByHeight(pts) { + const [z0, z1] = current().zrange, span = Math.max(1e-6, z1 - z0); + let prev = toScreen(pts[0]); + for (let i = 1; i < pts.length; i++) { + const s = toScreen(pts[i]); + if (prev && s) { + ctx.strokeStyle = heightColor(((pts[i - 1][2] + pts[i][2]) / 2 - z0) / span); + ctx.beginPath(); ctx.moveTo(prev[0], prev[1]); ctx.lineTo(s[0], s[1]); ctx.stroke(); + } + prev = s; + } +} +function drawSeam(seam, dashed) { + const s = toScreen([seam.x, seam.y, (seam.z0 + seam.z1) / 2]); + if (!s) return; + ctx.setLineDash(dashed ? [3, 3] : []); ctx.lineWidth = 2; ctx.strokeStyle = color("--accent"); + ctx.beginPath(); ctx.arc(s[0], s[1], 7, 0, 2 * Math.PI); ctx.stroke(); ctx.setLineDash([]); +} + +function draw() { + const m = current(); + ctx.clearRect(0, 0, canvas.clientWidth, canvas.clientHeight); + ctx.lineCap = "round"; ctx.lineJoin = "round"; ctx.globalAlpha = 1; + const base = Math.max(1.2, Math.min(4, 0.8 * state.view.scale)); + if (state.lines) { ctx.strokeStyle = color("--roadline"); ctx.lineWidth = Math.max(0.8, base * 0.5); ctx.setLineDash([]); m.lines.forEach(strokePolyline); } + if (state.lanes) { + for (const lane of m.lanes) { + ctx.setLineDash(lane.junction ? [4, 4] : []); + ctx.strokeStyle = lane.junction ? color("--junction") : color("--lane"); ctx.lineWidth = base * 0.8; + strokePolyline(lane.pts); + } + } + ctx.setLineDash([]); + if (state.compare && m.compare.length) { + ctx.strokeStyle = color("--compare"); ctx.lineWidth = base * 2.2; ctx.globalAlpha = 0.55; + m.compare.forEach(e => strokePolyline(e.pts)); ctx.globalAlpha = 1; + } + ctx.lineWidth = base * 1.1; + for (const e of m.edges) { + if (state.heightColor) { strokeByHeight(e.pts); continue; } + ctx.strokeStyle = e.type === 22 ? color("--median") : color("--edge"); strokePolyline(e.pts); + } + if (state.hover) { ctx.lineWidth = base * 2.6; ctx.strokeStyle = color("--accent"); strokePolyline(state.hover.pts); } + m.seams.forEach(s => drawSeam(s, false)); + if (state.compare) m.compareSeams.forEach(s => drawSeam(s, true)); + for (const mk of state.marks[state.map]) { + const s = toScreen([mk.x, mk.y, mk.z]); + if (!s) continue; + const [sx, sy] = s; + ctx.beginPath(); ctx.arc(sx, sy, 8, 0, 2 * Math.PI); ctx.fillStyle = color(CATEGORIES[mk.cat][1]); ctx.fill(); + ctx.fillStyle = "#fff"; ctx.font = "bold 10px IBM Plex Mono, ui-monospace, monospace"; ctx.textAlign = "center"; ctx.textBaseline = "middle"; + ctx.fillText(String(mk.n), sx, sy); + } +} + +function distToSegment(p, a, b) { + const dx = b[0] - a[0], dy = b[1] - a[1]; + const len2 = dx * dx + dy * dy; + const t = len2 === 0 ? 0 : Math.max(0, Math.min(1, ((p[0] - a[0]) * dx + (p[1] - a[1]) * dy) / len2)); + return Math.hypot(p[0] - (a[0] + t * dx), p[1] - (a[1] + t * dy)); +} +function polylineDistance(p, pts) { + let best = Infinity; + for (let i = 1; i < pts.length; i++) best = Math.min(best, distToSegment(p, pts[i - 1], pts[i])); + return pts.length === 1 ? Math.hypot(p[0] - pts[0][0], p[1] - pts[0][1]) : best; +} +function screenDistance(sx, sy, pts) { + let prev = toScreen(pts[0]), best = prev ? Math.hypot(sx - prev[0], sy - prev[1]) : Infinity; + for (let i = 1; i < pts.length; i++) { + const s = toScreen(pts[i]); + if (prev && s) best = Math.min(best, distToSegment([sx, sy], prev, s)); + prev = s; + } + return best; +} +function candidates() { + const m = current(); + const items = m.edges.map(e => ({ kind: "edge", id: e.id, type: e.type, pts: e.pts })); + if (state.compare) items.push(...m.compare.map(e => ({ kind: "old edge", id: e.id, type: e.type, pts: e.pts }))); + if (state.lanes) items.push(...m.lanes.map(l => ({ kind: l.junction ? "junction lane" : "lane", id: l.id, type: null, pts: l.pts, length: l.length }))); + return items; +} +function pick(sx, sy) { + let best = null, bestD = PICK_PX; + for (const it of candidates()) { const d = screenDistance(sx, sy, it.pts); if (d < bestD) { bestD = d; best = it; } } + return best; +} +function pickVertex(sx, sy) { + let best = null, bestD = VERTEX_PICK_PX; + for (const it of candidates()) for (const p of it.pts) { const s = toScreen(p); if (!s) continue; const d = Math.hypot(sx - s[0], sy - s[1]); if (d < bestD) { bestD = d; best = p; } } + return best; +} +function nearest(p, kind) { + let best = null, bestD = Infinity; + for (const it of candidates()) { + if (!it.kind.includes(kind)) continue; + const d = polylineDistance(p, it.pts); if (d < bestD) { bestD = d; best = it; } + } + return best ? { id: best.id, type: best.type, kind: best.kind, d: bestD } : null; +} +function showTooltip(item, x, y) { + if (!item) { tooltip.hidden = true; return; } + const len = item.pts.reduce((acc, p, i) => i ? acc + Math.hypot(p[0] - item.pts[i - 1][0], p[1] - item.pts[i - 1][1]) : 0, 0); + const zs = item.pts.map(p => p[2]), zText = `z ${Math.min(...zs).toFixed(1)}..${Math.max(...zs).toFixed(1)} m`; + const text = item.kind.includes("edge") ? `${item.kind} ${item.id} · type ${item.type} · ${item.pts.length} pts · ${len.toFixed(0)} m · ${zText}` + : `${item.kind} ${item.id} · ${len.toFixed(0)} m · ${zText}`; + tooltip.textContent = text; tooltip.style.left = `${x}px`; tooltip.style.top = `${y}px`; tooltip.hidden = false; +} + +function markLine(mk) { + const parts = []; + if (mk.edge) parts.push(`${mk.edge.kind} ${mk.edge.id} type ${mk.edge.type} at ${mk.edge.d.toFixed(2)} m`); + if (mk.oldEdge) parts.push(`old edge ${mk.oldEdge.id} at ${mk.oldEdge.d.toFixed(2)} m`); + if (mk.lane) parts.push(`${mk.lane.kind} ${mk.lane.id} at ${mk.lane.d.toFixed(2)} m`); + return `#${mk.n} ${CATEGORIES[mk.cat][0]} at (${mk.x.toFixed(1)}, ${mk.y.toFixed(1)}, z ${mk.z.toFixed(1)}) [${parts.join("; ")}]`; +} +function report() { + const m = current(), marks = state.marks[state.map]; + const head = `Map marks for ${m.name} (PufferDrive bin coordinates, ${META.bin_dir}${META.compare_dir ? ", old edges from " + META.compare_dir : ""}):`; + return [head, ...marks.map(mk => " " + markLine(mk))].join("\n"); +} +function seamLine(s) { return `${s.kind} ${s.id} at (${s.x.toFixed(1)}, ${s.y.toFixed(1)}) z ${s.z0.toFixed(1)} → ${s.z1.toFixed(1)}`; } +function focusSeam(s) { + state.view.yaw = 0.6; state.view.pitch = 0.95; + fitBox([s.x - 35, s.y - 35, s.x + 35, s.y + 35], 20, (s.z0 + s.z1) / 2); draw(); +} +function renderPanel() { + const m = current(), marks = state.marks[state.map]; + const stats = [["lanes", m.lanes.length], ["road edges", m.edges.length], ["old edges", m.compare.length], ["marks", marks.length], + ["seams", m.seams.length], ["old seams", m.compareSeams.length]]; + document.getElementById("stats").innerHTML = stats.map(([k, v]) => `
${k}
${v}
`).join(""); + const seams = [...m.seams, ...(state.compare ? m.compareSeams : [])]; + document.getElementById("seams").innerHTML = seams.length + ? seams.map((s, i) => `
${seamLine(s)}
`).join("") + : `
no height seams in the bin road edges
`; + document.querySelectorAll(".seam-row[data-i]").forEach(row => row.addEventListener("click", () => focusSeam(seams[Number(row.dataset.i)]))); + document.getElementById("marks").innerHTML = marks.map((mk, i) => + `
${mk.n}${markLine(mk).replace(/^#\d+ /, "")}
${CATEGORIES[mk.cat][2]}
`).join(""); + document.querySelectorAll(".mark-row").forEach(row => row.addEventListener("click", e => { + const i = Number(row.dataset.i); + if (e.target.dataset.rm !== undefined) { removeMark(i); return; } + const mk = marks[i]; fitBox([mk.x - 40, mk.y - 40, mk.x + 40, mk.y + 40], 20, mk.z); draw(); + })); + document.getElementById("report").value = report(); + document.querySelectorAll(".seg button").forEach(b => b.setAttribute("aria-pressed", String(b.dataset.cat === state.category))); +} +function addMark(x, y, z) { + const marks = state.marks[state.map], p = [x, y]; + const lane = nearest(p, "lane"); + const own = current().edges.map(e => ({ id: e.id, type: e.type, d: polylineDistance(p, e.pts) })).sort((a, b) => a.d - b.d)[0]; + const old = state.compare ? current().compare.map(e => ({ id: e.id, type: e.type, d: polylineDistance(p, e.pts) })).sort((a, b) => a.d - b.d)[0] : null; + marks.push({ n: marks.length + 1, cat: state.category, x, y, z, + edge: own && own.d < 5 ? { ...own, kind: "edge" } : null, + oldEdge: old && old.d < 5 ? old : null, + lane: lane && lane.d < 8 ? lane : null }); + renderPanel(); draw(); +} +function removeMark(i) { + const marks = state.marks[state.map]; marks.splice(i, 1); marks.forEach((mk, k) => mk.n = k + 1); renderPanel(); draw(); +} + +let drag = null, moved = false; +canvas.addEventListener("contextmenu", e => e.preventDefault()); +canvas.addEventListener("pointerdown", e => { + const orbiting = e.button === 1 || e.button === 2 || (e.button === 0 && e.shiftKey); + const rect = canvas.getBoundingClientRect(); + drag = { mode: orbiting ? "orbit" : "pan", x: e.clientX, y: e.clientY, t0: [state.view.cx, state.view.cy], + g0: groundOffset(e.clientX - rect.left, e.clientY - rect.top) }; + moved = false; canvas.setPointerCapture(e.pointerId); +}); +canvas.addEventListener("pointermove", e => { + const rect = canvas.getBoundingClientRect(); + const sx = e.clientX - rect.left, sy = e.clientY - rect.top; + if (drag) { + const dx = e.clientX - drag.x, dy = e.clientY - drag.y; + if (Math.abs(dx) + Math.abs(dy) > 2) moved = true; + if (drag.mode === "orbit") { orbit(dx * ORBIT_RAD_PER_PX, dy * ORBIT_RAD_PER_PX); drag.x = e.clientX; drag.y = e.clientY; } + else { const g = groundOffset(sx, sy); if (g && drag.g0) { state.view.cx = drag.t0[0] + drag.g0[0] - g[0]; state.view.cy = drag.t0[1] + drag.g0[1] - g[1]; } } + tooltip.hidden = true; document.getElementById("coords").textContent = viewLabel(); requestAnimationFrame(draw); return; + } + const item = pick(sx, sy); + if (item !== state.hover) { state.hover = item; requestAnimationFrame(draw); } + showTooltip(item, sx, sy); + const w = toWorld(sx, sy); document.getElementById("coords").textContent = `${w ? `bin (${w[0].toFixed(1)}, ${w[1].toFixed(1)})` : "above horizon"} · ${viewLabel()}`; +}); +canvas.addEventListener("pointerup", e => { + canvas.releasePointerCapture(e.pointerId); + if (!moved && e.button === 0) { + const rect = canvas.getBoundingClientRect(), sx = e.clientX - rect.left, sy = e.clientY - rect.top; + const w = pickVertex(sx, sy) || toWorld(sx, sy); + if (w) addMark(w[0], w[1], w[2] || 0); + } + drag = null; +}); +canvas.addEventListener("pointerleave", () => { state.hover = null; tooltip.hidden = true; draw(); }); +canvas.addEventListener("wheel", e => { + e.preventDefault(); + const rect = canvas.getBoundingClientRect(); + const sx = e.clientX - rect.left, sy = e.clientY - rect.top; + const factor = Math.exp(-e.deltaY * 0.0015); + const v = state.view, before = groundOffset(sx, sy); + v.scale *= factor; + const after = groundOffset(sx, sy); + if (before && after) { v.cx += before[0] - after[0]; v.cy += before[1] - after[1]; } + requestAnimationFrame(draw); +}, { passive: false }); +window.addEventListener("keydown", e => { + if (e.target.tagName === "TEXTAREA" || e.target.tagName === "INPUT") return; + if (CATEGORIES[e.key]) { state.category = e.key; renderPanel(); } + else if (e.key === "u") { const marks = state.marks[state.map]; if (marks.length) removeMark(marks.length - 1); } + else if (e.key === "f") { fitView(); draw(); } + else if (e.key === "t") { topView(); draw(); } + else if (e.key === "c") { copyReport(); } + else if (e.key === "Escape") { state.hover = null; tooltip.hidden = true; draw(); } +}); +function copyReport() { + const ta = document.getElementById("report"); ta.select(); + const done = () => { document.getElementById("copied").textContent = "copied"; setTimeout(() => { document.getElementById("copied").textContent = ""; }, 1500); }; + if (navigator.clipboard) navigator.clipboard.writeText(ta.value).then(done, () => { document.execCommand("copy"); done(); }); else { document.execCommand("copy"); done(); } +} +function selectMap(i) { + state.map = i; state.hover = null; + document.querySelectorAll(".tabs button").forEach((b, j) => b.setAttribute("aria-pressed", String(i === j))); + fitView(); renderPanel(); draw(); +} +const tabs = document.querySelector(".tabs"); +MAPS.forEach((m, i) => { + const b = document.createElement("button"); b.textContent = m.name; b.type = "button"; + b.setAttribute("aria-pressed", String(i === 0)); b.addEventListener("click", () => selectMap(i)); tabs.appendChild(b); +}); +document.querySelectorAll(".seg button").forEach(b => b.addEventListener("click", () => { state.category = b.dataset.cat; renderPanel(); })); +document.getElementById("compare").addEventListener("change", e => { state.compare = e.target.checked; renderPanel(); draw(); }); +document.getElementById("lines").addEventListener("change", e => { state.lines = e.target.checked; draw(); }); +document.getElementById("lanes").addEventListener("change", e => { state.lanes = e.target.checked; draw(); }); +document.getElementById("height").addEventListener("change", e => { state.heightColor = e.target.checked; draw(); }); +document.getElementById("zscale").addEventListener("input", e => { setZScale(Number(e.target.value)); draw(); }); +document.getElementById("fit").addEventListener("click", () => { fitView(); draw(); }); +document.getElementById("top").addEventListener("click", () => { topView(); draw(); }); +document.getElementById("copy").addEventListener("click", copyReport); +document.getElementById("clear").addEventListener("click", () => { state.marks[state.map] = []; renderPanel(); draw(); }); +document.getElementById("legend").innerHTML = `road edge (bin)median edgeold edge (compare)lane centerlinejunction laneheight (low → high)height seam in an edge`; +window.addEventListener("resize", resize); +window.matchMedia("(prefers-color-scheme: dark)").addEventListener("change", () => { renderPanel(); draw(); }); +new MutationObserver(() => { renderPanel(); draw(); }).observe(document.documentElement, { attributes: true, attributeFilter: ["data-theme"] }); +function applyHashView() { + const q = new URLSearchParams(location.hash.replace(/^#/, "")); + if (!q.size) return; + const mapIdx = MAPS.findIndex(m => m.name === q.get("map")); + if (mapIdx >= 0) { state.map = mapIdx; document.querySelectorAll(".tabs button").forEach((b, j) => b.setAttribute("aria-pressed", String(mapIdx === j))); } + if (q.has("z")) { state.zscale = Number(q.get("z")); document.getElementById("zscale").value = state.zscale; document.getElementById("zscale-val").textContent = `${state.zscale}×`; } + if (q.has("height")) { state.heightColor = q.get("height") === "1"; document.getElementById("height").checked = state.heightColor; } + const deg = Math.PI / 180; + state.view = { cx: 0, cy: 0, cz: 0, scale: 1, yaw: Number(q.get("yaw") || 0) * deg, pitch: Math.min(MAX_PITCH, Number(q.get("pitch") || 0) * deg) }; + fitView(); + if (q.has("cx") && q.has("cy")) { + const r = Number(q.get("r") || 60), cx = Number(q.get("cx")), cy = Number(q.get("cy")); + fitBox([cx - r, cy - r, cx + r, cy + r], 20, Number(q.get("cz") || 0)); + } + renderPanel(); +} +applyHashView(); +renderPanel(); +resize(); +""" + + +def build_html(maps: list[dict], meta: dict) -> str: + payload = json.dumps({"maps": maps, "meta": meta}, separators=(",", ":")) + head = ( + "PufferDrive Map Inspector\n" + '\n' + '\n' + f"\n" + ) + compare_note = f" · old edges: {meta['compare_dir']}" if meta.get("compare_dir") else "" + body = f""" +
+
+

PufferDrive Map Inspector

+
+
+ + + +
+
+ + + + + + + + +
+
+
+ +
+ +
left-drag pan · right/middle-drag (or shift+drag) rotate · wheel zoom · t top view · click drops a mark (1/2/3) · u undo · c copy report · view preset via #map=Town05&yaw=30&pitch=55&z=3&height=1&cx=-190&cy=-2&r=60 · bins: {meta["bin_dir"]}{compare_note}
+
+ +
+ + +""" + return f'\n\n\n\n\n{head}\n{body}\n\n' + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--bin-dir", required=True, help="Directory of PufferDrive *.bin map files to inspect") + parser.add_argument( + "--compare-dir", default=None, help="Second bin directory whose road edges are overlaid in orange" + ) + parser.add_argument( + "--towns", nargs="*", default=None, help="Subset of towns, e.g. Town04 Town06 (default: all bins)" + ) + parser.add_argument("--output", default=None, help="Output .html (default: /map_inspector.html)") + args = parser.parse_args() + + bin_dir = pathlib.Path(args.bin_dir).expanduser() + compare_dir = pathlib.Path(args.compare_dir).expanduser() if args.compare_dir else None + bin_files = sorted(bin_dir.glob("*.bin")) + if args.towns: + bin_files = [f for f in bin_files if f.stem.split("__")[-1] in set(args.towns)] + if not bin_files: + sys.exit(f"no matching .bin files in {bin_dir}") + maps = [] + for bin_file in bin_files: + compare_file = compare_dir / bin_file.name if compare_dir and (compare_dir / bin_file.name).exists() else None + maps.append(collect_map(bin_file, compare_file)) + meta = {"bin_dir": str(bin_dir), "compare_dir": str(compare_dir) if compare_dir else None} + output = pathlib.Path(args.output) if args.output else bin_dir / "map_inspector.html" + output.write_text(build_html(maps, meta)) + print(f"wrote {output} ({output.stat().st_size / 1e6:.1f} MB, {len(maps)} maps)") + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 134826fdfc..1b8b04ed8d 100644 --- a/setup.py +++ b/setup.py @@ -62,10 +62,11 @@ else: extra_compile_args += [ "-O2", - "-flto=auto" if PROFILE else "-flto", + "-flto=auto", ] extra_link_args += [ "-O2", + "-flto=auto", ] cxx_args += [ "-O3", @@ -84,7 +85,6 @@ ] extra_link_args += [ "-g", - "-flto=auto", "-fno-inline", "-fno-builtin", ] @@ -218,6 +218,7 @@ def run(self): "hydra-core", "omegaconf", "pandas", + "pyarrow", "tqdm", "matplotlib==3.11.0", "imageio", @@ -244,7 +245,7 @@ def run(self): ext_modules=c_extensions + torch_extensions, cmdclass={ "build_ext": BuildExt, - "build_torch": TorchBuildExt, + "build_torch": TorchBuildExt.with_options(use_ninja=cpp_extension.is_ninja_available()), "build_c": CBuildExt, }, include_dirs=[numpy.get_include()], diff --git a/tests/drive/include/drive_fixture.h b/tests/drive/include/drive_fixture.h index 71fed52d33..7aa6cc6393 100644 --- a/tests/drive/include/drive_fixture.h +++ b/tests/drive/include/drive_fixture.h @@ -130,7 +130,11 @@ static inline Drive drive_test_env_config( env.obs_range_road_behind_m = 20.0f; env.obs_range_road_side_m = 30.0f; env.partner_blindness_prob = 0.0f; + env.partner_blindness_prob_min = 0.0f; env.partner_blindness_trigger_prob = 0.1f; + env.partner_blindness_duration = 10; + env.partner_blindness_corridor = 1; + env.partner_blindness_time_fraction_max = 0.0f; env.phantom_braking_prob = 0.0f; env.phantom_braking_trigger_prob = 0.0f; env.phantom_braking_duration = 10; diff --git a/tests/drive/test_drive_cosim_setters.c b/tests/drive/test_drive_cosim_setters.c index 16986abce0..f8b97415f6 100644 --- a/tests/drive/test_drive_cosim_setters.c +++ b/tests/drive/test_drive_cosim_setters.c @@ -65,9 +65,10 @@ static TrafficControlElement make_traffic_element(int type, int state_size, int return element; } -static int test_set_traffic_light_states_writes_current_timestep_for_lights_only(void) { - // Only TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT elements are written, only at env->timestep, and every - // other timestep slot must be left alone. +static int test_set_traffic_light_states_writes_current_and_next_timestep_for_lights_only(void) { + // Only TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT elements are written, at env->timestep and the step after + // it (c_step increments the timestep before its metrics and obs read the state, and the next sync + // only arrives after that step); every other timestep slot must be left alone. int light_states[5] = {99, 99, 99, 99, 99}; int stop_sign_states[5] = {99, 99, 99, 99, 99}; TrafficControlElement elements[2]; @@ -83,9 +84,9 @@ static int test_set_traffic_light_states_writes_current_timestep_for_lights_only EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states), 0); EXPECT_EQ_INT(light_states[2], 3); // written at the current timestep + EXPECT_EQ_INT(light_states[3], 3); // and at the one c_step reads next EXPECT_EQ_INT(light_states[0], 99); EXPECT_EQ_INT(light_states[1], 99); - EXPECT_EQ_INT(light_states[3], 99); EXPECT_EQ_INT(light_states[4], 99); for (int t = 0; t < 5; t++) { EXPECT_EQ_INT(stop_sign_states[t], 99); // non-light element untouched @@ -118,6 +119,12 @@ static int test_set_traffic_light_states_rejects_out_of_range_timestep_and_state EXPECT_EQ_INT(light_states[1], 99); EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states), 0); EXPECT_EQ_INT(light_states[1], 1); + EXPECT_EQ_INT(light_states[2], 1); + + env.timestep = 2; // last slot of the schedule: no next timestep to fill, still a success + int last_states[2] = {3, 3}; + EXPECT_EQ_INT(c_set_traffic_light_states(&env, last_states), 0); + EXPECT_EQ_INT(light_states[2], 3); return 0; } @@ -141,7 +148,7 @@ static int test_set_agent_states_teleport_resets_prev_pose(void) { int idx[1] = {0}; float x[1] = {10.0f}, y[1] = {20.0f}, z[1] = {0.0f}, h[1] = {0.5f}; float vx[1] = {1.0f}, vy[1] = {0.0f}, yr[1] = {0.0f}, al[1] = {0.0f}; - EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, NULL), 0); + EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, NULL, NULL), 0); EXPECT_NEAR(agent.prev_x, agent.sim_x, 1e-6f); EXPECT_NEAR(agent.prev_y, agent.sim_y, 1e-6f); @@ -163,19 +170,51 @@ static int test_set_agent_states_seconds_stopped_injects_or_preserves(void) { float x[1] = {1.0f}, y[1] = {2.0f}, z[1] = {0.0f}, h[1] = {0.0f}; float vx[1] = {0.0f}, vy[1] = {0.0f}, yr[1] = {0.0f}, al[1] = {0.0f}; - EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, NULL), 0); + EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, NULL, NULL), 0); EXPECT_NEAR(agent.seconds_stopped, 12.5f, 1e-6f); float seconds_stopped[1] = {3.0f}; - EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, seconds_stopped), 0); + EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, seconds_stopped, NULL), 0); EXPECT_NEAR(agent.seconds_stopped, 3.0f, 1e-6f); float negative_seconds_stopped[1] = {-1.0f}; - EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, negative_seconds_stopped), -1); + EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, vx, vy, yr, al, negative_seconds_stopped, NULL), -1); EXPECT_NEAR(agent.seconds_stopped, 3.0f, 1e-6f); return 0; } +static int test_set_agent_states_steering_is_state_below_1mps(void) { + // A co-sim that tracks steering as state can hand it over exactly: the angle is kept and the + // lateral acceleration follows from it as in c_step (v^2 tan(steering) / wheelbase). Without it, + // steering is derived from yaw_rate with the speed floored at 1 m/s, so a stopped car that still + // holds its wheels turned reads 0 -- the case the steering input exists for. + Drive env = {0}; + Agent agent = drive_test_agent(0.0f, 0.0f, 0.0f); + env.agents = &agent; + env.num_total_agents = 1; + + int idx[1] = {0}; + float x[1] = {1.0f}, y[1] = {2.0f}, z[1] = {0.0f}, h[1] = {0.0f}; + float vy[1] = {0.0f}, yr[1] = {0.0f}, al[1] = {0.0f}; + + float stopped[1] = {0.0f}, steering[1] = {0.14f}; + EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, stopped, vy, yr, al, NULL, steering), 0); + EXPECT_NEAR(agent.steering_angle, 0.14f, 1e-6f); + EXPECT_NEAR(agent.accel_lat, 0.0f, 1e-6f); + + float crawling[1] = {0.5f}; + EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, crawling, vy, yr, al, NULL, steering), 0); + EXPECT_NEAR(agent.steering_angle, 0.14f, 1e-6f); + EXPECT_NEAR(agent.accel_lat, 0.25f * tanf(0.14f) / agent.wheelbase, 1e-6f); + + EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, stopped, vy, yr, al, NULL, NULL), 0); + EXPECT_NEAR(agent.steering_angle, 0.0f, 1e-6f); + + float not_finite[1] = {NAN}; + EXPECT_EQ_INT(c_set_agent_states(&env, 1, idx, x, y, z, h, stopped, vy, yr, al, NULL, not_finite), -1); + return 0; +} + // --------------------------------------------------------------------------- // c_set_agent_goals // --------------------------------------------------------------------------- @@ -327,16 +366,60 @@ static int test_set_stop_signs_rejects_bad_input_untouched(void) { return 0; } +// --------------------------------------------------------------------------- +// c_set_traffic_light_lines +// --------------------------------------------------------------------------- + +static int test_set_traffic_light_lines_moves_lights_only_and_rejects_bad_input(void) { + // Only the addressed light elements move (world_mean subtracted like every setter input), their state + // schedule stays; a stop-sign index, an out-of-range index, a non-finite coordinate or a zero-length + // line fails the whole call (-1) with every line untouched. + int light_states[2] = {1, 1}; + TrafficControlElement elements[2]; + elements[0] = make_traffic_element(TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT, 2, light_states); + elements[1] = make_traffic_element(TRAFFIC_CONTROL_TYPE_STOP_SIGN, 0, NULL); + Drive env = {0}; + env.traffic_elements = elements; + env.num_traffic_elements = 2; + env.world_mean_x = 100.0f; + env.world_mean_y = -50.0f; + + int idx[1] = {0}; + float line[6] = {110.0f, -52.0f, 0.5f, 110.0f, -48.0f, 0.5f}; + EXPECT_EQ_INT(c_set_traffic_light_lines(&env, 1, idx, line), 0); + EXPECT_NEAR(elements[0].stop_line[0], 10.0f, 1e-6f); + EXPECT_NEAR(elements[0].stop_line[1], -2.0f, 1e-6f); + EXPECT_NEAR(elements[0].stop_line[2], 0.5f, 1e-6f); + EXPECT_NEAR(elements[0].stop_line[3], 10.0f, 1e-6f); + EXPECT_NEAR(elements[0].stop_line[4], 2.0f, 1e-6f); + EXPECT_EQ_INT(elements[0].states[0], 1); + + int stop_sign_idx[1] = {1}; + EXPECT_EQ_INT(c_set_traffic_light_lines(&env, 1, stop_sign_idx, line), -1); + EXPECT_NEAR(elements[1].stop_line[0], 0.0f, 1e-6f); + int bad_idx[1] = {7}; + EXPECT_EQ_INT(c_set_traffic_light_lines(&env, 1, bad_idx, line), -1); + float degenerate[6] = {1.0f, 1.0f, 0.0f, 1.0f, 1.0f, 0.0f}; + EXPECT_EQ_INT(c_set_traffic_light_lines(&env, 1, idx, degenerate), -1); + float nan_line[6] = {0.0f, NAN, 0.0f, 0.0f, 4.0f, 0.0f}; + EXPECT_EQ_INT(c_set_traffic_light_lines(&env, 1, idx, nan_line), -1); + EXPECT_NEAR(elements[0].stop_line[0], 10.0f, 1e-6f); + EXPECT_EQ_INT(c_set_traffic_light_lines(&env, 0, NULL, NULL), 0); + return 0; +} + int main(void) { int failures = 0; RUN_TEST(test_set_agent_sizes_updates_dimensions_radius_and_wheelbase); RUN_TEST(test_set_agent_sizes_rejects_out_of_range_index_and_bad_size); - RUN_TEST(test_set_traffic_light_states_writes_current_timestep_for_lights_only); + RUN_TEST(test_set_traffic_light_states_writes_current_and_next_timestep_for_lights_only); RUN_TEST(test_set_traffic_light_states_rejects_out_of_range_timestep_and_state); RUN_TEST(test_set_stop_signs_retires_map_signs_and_appends_external_ones); RUN_TEST(test_set_stop_signs_rejects_bad_input_untouched); + RUN_TEST(test_set_traffic_light_lines_moves_lights_only_and_rejects_bad_input); RUN_TEST(test_set_agent_states_teleport_resets_prev_pose); RUN_TEST(test_set_agent_states_seconds_stopped_injects_or_preserves); + RUN_TEST(test_set_agent_states_steering_is_state_below_1mps); RUN_TEST(test_set_agent_goals_sets_positions_lane_and_count); RUN_TEST(test_set_agent_goals_rejects_more_than_max_goals); RUN_TEST(test_set_agent_goals_rejects_out_of_range_agent_idx); diff --git a/tests/drive/test_drive_dynamics.c b/tests/drive/test_drive_dynamics.c index 487f862188..f7b74e4476 100644 --- a/tests/drive/test_drive_dynamics.c +++ b/tests/drive/test_drive_dynamics.c @@ -98,6 +98,51 @@ static int test_neutral_actions_zero_out(void) { return 0; } +// Braking to a standstill with reset_accel_on_stop must report the driver's jerk, not the zeroed acceleration +static int test_jerk_stop_reset_keeps_driver_jerk(void) { + srand(29); + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); + env.action_type = 1; + env.dynamics_model = DYNAMICS_MODEL_JERK; + env.reset_accel_on_stop = 1; + env.dt = 0.3f; + Agent *agent = &env.agents[env.active_agent_indices[0]]; + agent->reward_coefs[REWARD_COEF_THROTTLE] = 1.0f; + agent->reward_coefs[REWARD_COEF_STEER] = 1.0f; + agent->reward_coefs[REWARD_COEF_ACC] = 1.0f; + agent->reward_coefs[REWARD_COEF_SPEED] = 1.0f; + float (*actions)[2] = (float (*)[2]) env.actions; + const float braking_jerk_action = JERK_LONG[1] / -JERK_LONG[0]; // decodes to -4 m/s^3 + const float stop_speed_mps = 0.3f; + const float stop_accel_mps2 = -3.0f; + { + agent->sim_vx = stop_speed_mps * agent->cos_heading; + agent->sim_vy = stop_speed_mps * agent->sin_heading; + update_agent_speed(agent); + agent->accel_long = stop_accel_mps2; + actions[0][0] = 0.0f; + actions[0][1] = 0.0f; + move_dynamics(&env, 0, env.active_agent_indices[0]); + EXPECT_NEAR(agent->sim_speed, 0.0f, 1e-6f); + EXPECT_NEAR(agent->accel_long, 0.0f, 1e-6f); + EXPECT_NEAR(agent->jerk_long, 0.0f, 1e-6f); + } + { + agent->sim_vx = stop_speed_mps * agent->cos_heading; + agent->sim_vy = stop_speed_mps * agent->sin_heading; + update_agent_speed(agent); + agent->accel_long = stop_accel_mps2; + actions[0][0] = braking_jerk_action; + actions[0][1] = 0.0f; + move_dynamics(&env, 0, env.active_agent_indices[0]); + EXPECT_NEAR(agent->sim_speed, 0.0f, 1e-6f); + EXPECT_NEAR(agent->accel_long, 0.0f, 1e-6f); + EXPECT_NEAR(agent->jerk_long, JERK_LONG[1], 1e-3f); + } + free_allocated(&env); + return 0; +} + int main(void) { int failures = 0; RUN_TEST(test_classic_action_clipping); @@ -105,5 +150,6 @@ int main(void) { RUN_TEST(test_dynamics_stopped_agent_clears_motion); RUN_TEST(test_dynamics_removed_agent_invalidated); RUN_TEST(test_neutral_actions_zero_out); + RUN_TEST(test_jerk_stop_reset_keeps_driver_jerk); return test_summary(failures); } diff --git a/tests/drive/test_drive_map_cache.c b/tests/drive/test_drive_map_cache.c index a1dd9ff6d8..a4ff0a936c 100644 --- a/tests/drive/test_drive_map_cache.c +++ b/tests/drive/test_drive_map_cache.c @@ -330,8 +330,8 @@ static int test_road_observation_cache_handles_widths_and_zero_length_segments(v Drive env = {.road_elements = roads, .num_road_elements = 3}; cache_road_observation_geometry(&env); const RoadObservationSegment expected_lane[] = { - {3.0f, 4.0f, 2.0f, 5.0f, 0.6f, 0.8f, 3.0f}, - {6.0f, 8.0f, 3.0f, 0.0f, 0.0f, 0.0f, 5.0f}, + {3.0f, 4.0f, 2.0f, 5.0f, 0.6f, 0.8f, 3.0f, 1}, + {6.0f, 8.0f, 3.0f, 0.0f, 0.0f, 0.0f, 5.0f, 1}, }; RoadObservationSegment expected_boundary[2]; memcpy(expected_boundary, expected_lane, sizeof(expected_boundary)); @@ -345,10 +345,47 @@ static int test_road_observation_cache_handles_widths_and_zero_length_segments(v return 0; } +// 5 straight 10 m segments then a 90 degree turn of two 10 m segments; stride 2. +static int test_stride_merge_chords_stay_connected_and_split_at_bends(void) { + float x[] = {0.0f, 10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 50.0f, 50.0f}; + float y[] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 10.0f, 20.0f}; + float z[8] = {0.0f}; + float headings[] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.5708f, 1.5708f, 1.5708f}; + RoadMapElement roads[] = { + {.type = ROAD_EDGE_BOUNDARY, .segment_size = 8, .x = x, .y = y, .z = z, .headings = headings}, + }; + Drive legacy = {.road_elements = roads, .num_road_elements = 1, .obs_boundary_stride = 2}; + cache_road_observation_geometry(&legacy); + const int expected_legacy_observed[] = {1, 0, 1, 0, 1, 1, 0}; + for (int segment_idx = 0; segment_idx < 7; segment_idx++) { + EXPECT_EQ_INT(roads[0].observation_segments[segment_idx].observed, expected_legacy_observed[segment_idx]); + EXPECT_NEAR(roads[0].observation_segments[segment_idx].half_length_m, 5.0f, 1e-5f); + } + free(roads[0].observation_segments); + + Drive merged + = {.road_elements = roads, .num_road_elements = 1, .obs_boundary_stride = 2, .obs_stride_merge_segments = 1}; + cache_road_observation_geometry(&merged); + // runs [0,1] [2,3] [4] [5,6] -> carriers 0, 2, 4, 5 spanning x 0-20, 20-40, 40-50, then y 0-20 + const RoadObservationSegment expected[] = { + {10.0f, 0.0f, 0.0f, 10.0f, 1.0f, 0.0f, 0.0f, 1}, + {15.0f, 0.0f, 0.0f, 5.0f, 1.0f, 0.0f, 0.0f, 0}, + {30.0f, 0.0f, 0.0f, 10.0f, 1.0f, 0.0f, 0.0f, 1}, + {35.0f, 0.0f, 0.0f, 5.0f, 1.0f, 0.0f, 0.0f, 0}, + {45.0f, 0.0f, 0.0f, 5.0f, 1.0f, 0.0f, 0.0f, 1}, + {50.0f, 10.0f, 0.0f, 10.0f, 0.0f, 1.0f, 0.0f, 1}, + {50.0f, 15.0f, 0.0f, 5.0f, 0.0f, 1.0f, 0.0f, 0}, + }; + EXPECT_EQ_INT(memcmp(roads[0].observation_segments, expected, sizeof(expected)), 0); + free(roads[0].observation_segments); + return 0; +} + int main(void) { int failures = 0; RUN_TEST(test_road_observation_cache_follows_map_ownership); RUN_TEST(test_road_observation_cache_handles_widths_and_zero_length_segments); + RUN_TEST(test_stride_merge_chords_stay_connected_and_split_at_bends); RUN_TEST(test_all_cache_modes_produce_identical_step_outputs); RUN_TEST(test_cache_mode_matrix_has_expected_map_and_neighbor_allocations); RUN_TEST(test_mixed_neighbor_modes_share_map_and_populate_neighbor_cache); diff --git a/tests/drive/test_drive_observations_rewards.c b/tests/drive/test_drive_observations_rewards.c index 2ca5f6e099..6e631a6ac9 100644 --- a/tests/drive/test_drive_observations_rewards.c +++ b/tests/drive/test_drive_observations_rewards.c @@ -276,19 +276,23 @@ static int test_partner_obs_relative_velocity(void) { update_agent_speed(&agents[0]); update_agent_speed(&agents[1]); env.agents = agents; + PartnerScanEntry partner_scan[2] = { + {agents[0].sim_x, agents[0].sim_y, agents[0].sim_z, 0}, + {agents[1].sim_x, agents[1].sim_y, agents[1].sim_z, 1}, + }; float obs[2 * (PARTNER_FEATURES + PARTNER_RELATIVE_VELOCITY_FEATURES)]; int partner_count = 0; env.obs_partner_relative_velocity = 0; memset(obs, 0, sizeof(obs)); - int end_idx = write_partner_obs(&env, &agents[0], 0, obs, 0, &partner_count); + int end_idx = write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); EXPECT_EQ_INT(partner_count, 1); EXPECT_EQ_INT(end_idx, 2 * PARTNER_FEATURES); EXPECT_NEAR(obs[7], 0.5f, 1e-5f); env.obs_partner_relative_velocity = 1; memset(obs, 0, sizeof(obs)); - end_idx = write_partner_obs(&env, &agents[0], 0, obs, 0, &partner_count); + end_idx = write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); EXPECT_EQ_INT(end_idx, 2 * (PARTNER_FEATURES + PARTNER_RELATIVE_VELOCITY_FEATURES)); EXPECT_NEAR(obs[7], 0.5f, 1e-5f); EXPECT_NEAR(obs[PARTNER_FEATURES], -1.0f, 1e-5f); @@ -296,6 +300,230 @@ static int test_partner_obs_relative_velocity(void) { return 0; } +static int test_partner_obs_split_long_partner(void) { + // Ego heading +x; 4 m car 10 m behind; 13 m truck 20 m ahead heading +y -> split at 7 m into 6.5 m pieces at y + // -/+3.25. + Drive env = {0}; + env.obs_norm_speed_mps = 10.0f; + env.obs_norm_xy_offset_m = 100.0f; + env.obs_norm_z_m = 10.0f; + env.obs_norm_veh_length_m = 10.0f; + env.obs_norm_veh_width_m = 5.0f; + env.obs_range_partner_m = 100.0f; + env.obs_partner_relative_velocity = 1; + env.num_agents = 3; + env.active_agent_count = 3; + int active[3] = {0, 1, 2}; + env.active_agent_indices = active; + Agent agents[3]; + agents[0] = drive_test_agent(0.0f, 0.0f, 0.0f); + agents[1] = drive_test_agent(20.0f, 0.0f, (float) M_PI / 2.0f); + agents[1].sim_length = 13.0f; + agents[1].sim_vy = 5.0f; + agents[2] = drive_test_agent(-10.0f, 0.0f, 0.0f); + for (int agent_idx = 0; agent_idx < 3; agent_idx++) { + update_agent_speed(&agents[agent_idx]); + } + env.agents = agents; + PartnerScanEntry partner_scan[3] = { + {agents[0].sim_x, agents[0].sim_y, agents[0].sim_z, 0}, + {agents[1].sim_x, agents[1].sim_y, agents[1].sim_z, 1}, + {agents[2].sim_x, agents[2].sim_y, agents[2].sim_z, 2}, + }; + const int row_features = PARTNER_FEATURES + PARTNER_RELATIVE_VELOCITY_FEATURES; + float obs[3 * (PARTNER_FEATURES + PARTNER_RELATIVE_VELOCITY_FEATURES)]; + int partner_count = 0; + + env.obs_slots_partners_n = 3; + env.obs_partner_split_length_m = 0.0f; + memset(obs, 0, sizeof(obs)); + int end_idx = write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); + EXPECT_EQ_INT(partner_count, 2); + EXPECT_EQ_INT(end_idx, 3 * row_features); + EXPECT_NEAR(obs[0], -0.1f, 1e-6f); + EXPECT_NEAR(obs[row_features], 0.2f, 1e-6f); + EXPECT_NEAR(obs[row_features + 1], 0.0f, 1e-6f); + EXPECT_NEAR(obs[row_features + 3], 1.3f, 1e-6f); + + env.obs_partner_split_length_m = 7.0f; + memset(obs, 0, sizeof(obs)); + end_idx = write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); + EXPECT_EQ_INT(partner_count, 3); + EXPECT_EQ_INT(end_idx, 3 * row_features); + EXPECT_NEAR(obs[0], -0.1f, 1e-6f); + EXPECT_NEAR(obs[3], 0.4f, 1e-6f); + EXPECT_NEAR(obs[row_features], 0.2f, 1e-6f); + EXPECT_NEAR(obs[row_features + 1], -0.0325f, 1e-6f); + EXPECT_NEAR(obs[row_features + 3], 0.65f, 1e-6f); + EXPECT_NEAR(obs[2 * row_features], 0.2f, 1e-6f); + EXPECT_NEAR(obs[2 * row_features + 1], 0.0325f, 1e-6f); + EXPECT_EQ_INT(memcmp(&obs[row_features + 2], &obs[2 * row_features + 2], (row_features - 2) * sizeof(float)), 0); + + env.obs_slots_partners_n = 2; + memset(obs, 0, sizeof(obs)); + end_idx = write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); + EXPECT_EQ_INT(partner_count, 2); + EXPECT_EQ_INT(end_idx, 2 * row_features); + EXPECT_NEAR(obs[0], -0.1f, 1e-6f); + EXPECT_NEAR(obs[row_features + 1], -0.0325f, 1e-6f); + return 0; +} + +static int test_partner_obs_blindness_corridor_toggle(void) { + // Ego heading +x, blind for the frame; car 20 m ahead in the corridor, car 20 m to the left outside it + Drive env = {0}; + env.obs_norm_speed_mps = 10.0f; + env.obs_norm_xy_offset_m = 100.0f; + env.obs_norm_z_m = 10.0f; + env.obs_norm_veh_length_m = 10.0f; + env.obs_norm_veh_width_m = 5.0f; + env.obs_range_partner_m = 100.0f; + env.obs_slots_partners_n = 3; + env.num_agents = 3; + env.active_agent_count = 3; + int active[3] = {0, 1, 2}; + env.active_agent_indices = active; + Agent agents[3]; + agents[0] = drive_test_agent(0.0f, 0.0f, 0.0f); + agents[1] = drive_test_agent(20.0f, 0.0f, 0.0f); + agents[2] = drive_test_agent(0.0f, 20.0f, 0.0f); + for (int agent_idx = 0; agent_idx < 3; agent_idx++) { + update_agent_speed(&agents[agent_idx]); + } + env.agents = agents; + PartnerScanEntry partner_scan[3] = { + {agents[0].sim_x, agents[0].sim_y, agents[0].sim_z, 0}, + {agents[1].sim_x, agents[1].sim_y, agents[1].sim_z, 1}, + {agents[2].sim_x, agents[2].sim_y, agents[2].sim_z, 2}, + }; + float obs[3 * PARTNER_FEATURES]; + int partner_count = 0; + + env.partner_blindness_corridor = 1; + agents[0].partner_blindness_counter = 2; + memset(obs, 0, sizeof(obs)); + int end_idx = write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); + EXPECT_EQ_INT(partner_count, 1); + EXPECT_EQ_INT(end_idx, 3 * PARTNER_FEATURES); + EXPECT_EQ_INT(agents[0].partner_blindness_counter, 1); + EXPECT_NEAR(obs[0], 0.2f, 1e-6f); + EXPECT_NEAR(obs[1], 0.0f, 1e-6f); + + env.partner_blindness_corridor = 0; + agents[0].partner_blindness_counter = 2; + memset(obs, 0, sizeof(obs)); + end_idx = write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); + EXPECT_EQ_INT(partner_count, 0); + EXPECT_EQ_INT(end_idx, 3 * PARTNER_FEATURES); + EXPECT_EQ_INT(agents[0].partner_blindness_counter, 1); + for (int i = 0; i < 3 * PARTNER_FEATURES; i++) { + EXPECT_NEAR(obs[i], 0.0f, 0.0f); + } + + agents[0].partner_blindness_counter = 0; + memset(obs, 0, sizeof(obs)); + end_idx = write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); + EXPECT_EQ_INT(partner_count, 2); + EXPECT_EQ_INT(end_idx, 3 * PARTNER_FEATURES); + return 0; +} + +static int test_partner_blindness_episode_rate_floor(void) { + Drive env = {0}; + env.init_seed = 7; + env.partner_blindness_prob_min = 0.2f; + env.partner_blindness_prob = 0.4f; + for (int episode = 0; episode < 16; episode++) { + begin_episode_rng(&env); + EXPECT_TRUE(env.episode_partner_blindness_prob >= 0.2f); + EXPECT_TRUE(env.episode_partner_blindness_prob <= 0.4f); + } + env.partner_blindness_prob_min = 0.4f; + begin_episode_rng(&env); + EXPECT_NEAR(env.episode_partner_blindness_prob, 0.4f, 1e-6f); + return 0; +} + +static int test_partner_blindness_time_fraction_scales_trigger_prob(void) { + // 10-step bursts at trigger prob 0.03 leave a blind partner blind 10 / (10 + 0.97 / 0.03) = 23.6% of the time + Drive env = {0}; + rng_seed(&env.rng_state, 7); + env.episode_partner_blindness_prob = 1.0f; + env.partner_blindness_trigger_prob = 0.03f; + env.partner_blindness_duration = 10; + Agent agent = drive_test_agent(0.0f, 0.0f, 0.0f); + + sample_erratic_flags(&env, &agent); + EXPECT_EQ_INT(agent.is_blind_partner, 1); + EXPECT_NEAR(agent.partner_blindness_trigger_prob, 0.03f, 1e-6f); + + float base_fraction = 10.0f * 0.03f / (10.0f * 0.03f + 1.0f - 0.03f); + env.partner_blindness_time_fraction_max = base_fraction; + sample_erratic_flags(&env, &agent); + EXPECT_NEAR(agent.partner_blindness_trigger_prob, 0.03f, 1e-5f); + + env.partner_blindness_time_fraction_max = 1.0f; + for (int draw = 0; draw < 32; draw++) { + sample_erratic_flags(&env, &agent); + EXPECT_EQ_INT(agent.is_blind_partner, 1); + EXPECT_TRUE(agent.partner_blindness_trigger_prob >= 0.03f - 1e-6f); + EXPECT_TRUE(agent.partner_blindness_trigger_prob <= 1.0f); + } + + env.partner_blindness_trigger_prob = 1.0f; + sample_erratic_flags(&env, &agent); + EXPECT_NEAR(agent.partner_blindness_trigger_prob, 1.0f, 1e-6f); + return 0; +} + +static int test_partner_blindness_full_time_fraction_blinds_every_frame(void) { + // Ego heading +x with trigger prob 1: the corridor hides the car on the left in every frame, the first included + Drive env = {0}; + rng_seed(&env.rng_state, 7); + env.obs_norm_speed_mps = 10.0f; + env.obs_norm_xy_offset_m = 100.0f; + env.obs_norm_z_m = 10.0f; + env.obs_norm_veh_length_m = 10.0f; + env.obs_norm_veh_width_m = 5.0f; + env.obs_range_partner_m = 100.0f; + env.obs_slots_partners_n = 3; + env.num_agents = 3; + env.active_agent_count = 3; + env.partner_blindness_duration = 1; + env.partner_blindness_corridor = 1; + int active[3] = {0, 1, 2}; + env.active_agent_indices = active; + Agent agents[3]; + agents[0] = drive_test_agent(0.0f, 0.0f, 0.0f); + agents[1] = drive_test_agent(20.0f, 0.0f, 0.0f); + agents[2] = drive_test_agent(0.0f, 20.0f, 0.0f); + for (int agent_idx = 0; agent_idx < 3; agent_idx++) { + update_agent_speed(&agents[agent_idx]); + } + env.agents = agents; + agents[0].is_blind_partner = 1; + agents[0].partner_blindness_trigger_prob = 1.0f; + PartnerScanEntry partner_scan[3] = { + {agents[0].sim_x, agents[0].sim_y, agents[0].sim_z, 0}, + {agents[1].sim_x, agents[1].sim_y, agents[1].sim_z, 1}, + {agents[2].sim_x, agents[2].sim_y, agents[2].sim_z, 2}, + }; + float obs[3 * PARTNER_FEATURES]; + int partner_count = 0; + for (int frame = 0; frame < 20; frame++) { + memset(obs, 0, sizeof(obs)); + write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); + EXPECT_EQ_INT(partner_count, 1); + } + + agents[0].partner_blindness_trigger_prob = 0.0f; + agents[0].partner_blindness_counter = 0; + memset(obs, 0, sizeof(obs)); + write_partner_obs(&env, &agents[0], 0, partner_scan, obs, 0, &partner_count); + EXPECT_EQ_INT(partner_count, 2); + return 0; +} + static int test_reward_lane_align_wrong_way(void) { Drive env; Agent agent; @@ -350,7 +578,7 @@ static int test_lane_graph_distance_to_goal_sentinels(void) { EXPECT_TRUE(env.lane_graph.n_lanes >= 2); int lane_idx = env.lane_graph.lane_ids[0]; int goal_graph_idx = 1; - float *stored_distance_m = &env.lane_graph.distances[0 * env.lane_graph.n_lanes + goal_graph_idx]; + float *stored_distance_m = &env.lane_graph.distances[goal_graph_idx * env.lane_graph.n_lanes + 0]; *stored_distance_m = 123.0f; EXPECT_NEAR(lane_graph_distance_to_goal_m(&env, lane_idx, goal_graph_idx), 123.0f, 1e-6f); @@ -376,6 +604,11 @@ int main(void) { RUN_TEST(test_reward_goal_speed_gating); RUN_TEST(test_reward_coef_goal_speed_pinned); RUN_TEST(test_partner_obs_relative_velocity); + RUN_TEST(test_partner_obs_split_long_partner); + RUN_TEST(test_partner_obs_blindness_corridor_toggle); + RUN_TEST(test_partner_blindness_episode_rate_floor); + RUN_TEST(test_partner_blindness_time_fraction_scales_trigger_prob); + RUN_TEST(test_partner_blindness_full_time_fraction_blinds_every_frame); RUN_TEST(test_reward_lane_align_wrong_way); return test_summary(failures); } diff --git a/tests/drive/test_drive_stop_sign.c b/tests/drive/test_drive_stop_sign.c index 472ccff4a2..4baff302e2 100644 --- a/tests/drive/test_drive_stop_sign.c +++ b/tests/drive/test_drive_stop_sign.c @@ -294,6 +294,34 @@ static int test_standstill_clears_beyond_proximity(void) { return 0; } +// Two lanes' signs 4 m apart: 19 m out on a curved approach the agent sits in the neighbour line's frame +// (|across| 0.1 m) and latches it; once it is in its own lane the own sign must take over and complete. +static int test_neighbour_lane_target_yields_to_own_lane_sign(void) { + TrafficControlElement tcs[2]; + tcs[0] = drive_test_stop_sign(); + tcs[0].stop_line[1] = 2.0f; + tcs[0].stop_line[4] = 6.0f; + tcs[1] = drive_test_stop_sign(); + Agent agent = drive_test_stop_agent(-20.0f, 3.9f, 0.0f); + Drive env = drive_test_stop_env(tcs, &agent); + env.num_traffic_elements = 2; + drive_test_step_to(&agent, -19.0f, 3.9f, 8.0f); + EXPECT_FALSE(update_stop_sign_state(&env, 0)); + EXPECT_EQ_INT(agent.stop_sign_target_idx, 0); + drive_test_step_to(&agent, -10.0f, 0.5f, 8.0f); + EXPECT_FALSE(update_stop_sign_state(&env, 0)); + EXPECT_EQ_INT(agent.stop_sign_target_idx, 1); + EXPECT_EQ_INT(agent.stop_sign_stop_completed, 0); + drive_test_step_to(&agent, -2.5f, 0.0f, 0.0f); + EXPECT_FALSE(update_stop_sign_state(&env, 0)); + EXPECT_EQ_INT(agent.stop_sign_target_idx, 1); + EXPECT_EQ_INT(agent.stop_sign_stop_completed, 1); + drive_test_step_to(&agent, 0.5f, 0.0f, 3.0f); + EXPECT_FALSE(update_stop_sign_state(&env, 0)); + EXPECT_EQ_INT(agent.stop_sign_last_failed_idx, -1); + return 0; +} + static int test_traffic_light_ignored(void) { TrafficControlElement tc = drive_test_stop_sign(); tc.type = TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT; @@ -331,6 +359,7 @@ int main(void) { RUN_TEST(test_state_clears_beyond_proximity); RUN_TEST(test_spawn_inside_box_counts_as_stopped); RUN_TEST(test_standstill_before_targeting_counts); + RUN_TEST(test_neighbour_lane_target_yields_to_own_lane_sign); RUN_TEST(test_standstill_outside_box_before_targeting_does_not_count); RUN_TEST(test_standstill_clears_beyond_proximity); RUN_TEST(test_traffic_light_ignored); diff --git a/tests/drive/test_drive_traffic_lights.c b/tests/drive/test_drive_traffic_lights.c index 525fbf3662..9c7b3c088c 100644 --- a/tests/drive/test_drive_traffic_lights.c +++ b/tests/drive/test_drive_traffic_lights.c @@ -160,6 +160,41 @@ static int test_eval_never_removes_lights(void) { return 0; } +// Gigaflow schedules include the terminal step, so the bootstrap observation reads a real light state +static int test_schedule_covers_terminal_observation(void) { + Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); + env.eval_mode = 1; // default CARLA cycles: no light is removed, so the terminal state is never OFF + allocate(&env); + c_reset(&env); + int light_count = 0; + TrafficControlElement *probe_light = NULL; + for (int i = 0; i < env.num_traffic_elements; i++) { + TrafficControlElement *tc = &env.traffic_elements[i]; + if (tc->type != TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT) { + continue; + } + EXPECT_EQ_INT(tc->state_size, env.scenario_length + 1); + EXPECT_TRUE(tc->states[env.scenario_length] != TRAFFIC_CONTROL_STATE_OFF); + probe_light = tc; + light_count++; + } + EXPECT_TRUE(light_count > 0); + Agent *agent = &env.agents[env.active_agent_indices[0]]; + agent->sim_x = (probe_light->stop_line[0] + probe_light->stop_line[3]) * 0.5f; + agent->sim_y = (probe_light->stop_line[1] + probe_light->stop_line[4]) * 0.5f; + agent->sim_z = (probe_light->stop_line[2] + probe_light->stop_line[5]) * 0.5f; + copy_pose_to_prev(agent); + env.timestep = env.scenario_length; + compute_observations(&env); + int valid_base = compute_observation_size(&env) - OBS_VALID_COUNT_FEATURES; + int traffic_base = valid_base - env.obs_slots_traffic_controls_n * TRAFFIC_CONTROL_FEATURES; + EXPECT_TRUE(env.observations[valid_base + 3] >= 1.0f); + int observed_state = (int) env.observations[traffic_base + TRAFFIC_CONTROL_FEATURES - 1]; + EXPECT_EQ_INT(observed_state, probe_light->states[env.scenario_length]); + free_allocated(&env); + return 0; +} + int main(void) { int failures = 0; RUN_TEST(test_junction_lights_never_green_together); @@ -169,5 +204,6 @@ int main(void) { RUN_TEST(test_training_lights_stay_exclusive_unless_removed); RUN_TEST(test_training_group_removal_turns_off_whole_junction); RUN_TEST(test_eval_never_removes_lights); + RUN_TEST(test_schedule_covers_terminal_observation); return test_summary(failures); } diff --git a/tests/eval/test_eval.py b/tests/eval/test_eval.py index bb19d2a3ff..ae5c261ead 100644 --- a/tests/eval/test_eval.py +++ b/tests/eval/test_eval.py @@ -102,6 +102,7 @@ def _write_benchmark_config( "termination_mode": False, "obs_dropout_lane": 0.0, "obs_dropout_boundary": 0.0, + "speed_limit_random_prob": 0.0, }, "benchmarks": [ { @@ -381,6 +382,7 @@ def _replay_render_args(): "goal_source": "gt", "num_goals": 3, "compute_eval_metrics": True, + "speed_limit_random_prob": 0.0, } ) _set_small_observation_config(args) @@ -556,6 +558,7 @@ def _training_args(tmp_path, benchmark_config_path, evaluation_enabled): "actor_num_layers": 0, "critic_hidden_size": 16, "critic_num_layers": 0, + "fused_slot_encoder": False, # tiny CPU encoders fail the fused width check } ) args["train"].update( diff --git a/tests/smoke_tests/data/drive_rollout_golden.json b/tests/smoke_tests/data/drive_rollout_golden.json index 4a56a6eee1..5b983ca767 100644 --- a/tests/smoke_tests/data/drive_rollout_golden.json +++ b/tests/smoke_tests/data/drive_rollout_golden.json @@ -1,43 +1,43 @@ { "env": { - "avg_distance_per_infraction": 13.006657809019089, - "avg_speed_per_agent": 1.3086299777030945, - "collision_rate": 0.022455357760190964, - "comfort_violation_count": 0.6155481282621622, - "dnf_rate": 0.5376322198659181, + "avg_distance_per_infraction": 13.21597754804394, + "avg_speed_per_agent": 1.3028922790213475, + "collision_rate": 0.029705106080332888, + "comfort_violation_count": 0.5997052290771581, + "dnf_rate": 0.5287974786154831, "early_reset_short": 0.0, - "episode_length": 14.7, - "episode_return": -1.1682074233889579, - "lane_center_rate": 0.2795072810724378, - "n": 15.0125, - "num_goals_reached": 0.008072916883975267, - "offroad_rate": 0.43584993463009597, + "episode_length": 15.0, + "episode_return": -1.2338795367675492, + "lane_center_rate": 0.27735796957453596, + "n": 14.531645569620252, + "num_goals_reached": 0.013422706459141984, + "offroad_rate": 0.43427748921551285, "red_light_violation_rate": 0.0, "reward_components/ade": 0.0, - "reward_components/collision": -0.03335988249164075, - "reward_components/comfort": -0.4453425757586956, - "reward_components/goal": 0.004010416753590107, - "reward_components/lane_align": -0.027996644657105208, - "reward_components/lane_center": -0.011789687484269961, - "reward_components/offroad": -0.64305120408535, - "reward_components/overspeed": 0.0, + "reward_components/collision": -0.04564126991207087, + "reward_components/comfort": -0.45430679939970187, + "reward_components/goal": 0.0034433394859108746, + "reward_components/lane_align": -0.025490657886183713, + "reward_components/lane_center": -0.012532398257923277, + "reward_components/offroad": -0.6707258812988861, + "reward_components/overspeed": -0.01763227161040671, "reward_components/red_light": 0.0, - "reward_components/reverse": -0.01078360435203649, + "reward_components/reverse": -0.011113312932415098, "reward_components/stop_sign": 0.0, - "reward_components/timestep": -8.974586080512381e-05, - "reward_components/velocity": 0.0001955053230744852, + "reward_components/timestep": -9.126150063268511e-05, + "reward_components/velocity": 0.00021098649279027292, "score": 0.0, "spawn_failed": 0.0, "spawn_failed_goal": 0.0, - "spawn_reject_collision": 0.0664726434275508, + "spawn_reject_collision": 0.06356702546907377, "spawn_reject_empty_cell": 0.0, - "spawn_reject_offroad": 0.43846612256020306, - "spawn_reject_stop_line": 0.01114011025056243, + "spawn_reject_offroad": 0.40240530132115643, + "spawn_reject_stop_line": 0.020339523783967466, "stop_sign_violation_rate": 0.0, "stopped_at_reset": 0.0, - "total_distance_travelled_sum": 86.93375225067139, - "total_infraction_count": 6.7375, - "velocity_progress_sum": 0.017799060756806283 + "total_distance_travelled_sum": 85.30317799049088, + "total_infraction_count": 6.544303797468355, + "velocity_progress_sum": 0.019357881364353662 }, "meta": { "bptt_horizon": 64, diff --git a/tests/smoke_tests/data/drive_smoke_golden.json b/tests/smoke_tests/data/drive_smoke_golden.json index c78d781b00..df6979618d 100644 --- a/tests/smoke_tests/data/drive_smoke_golden.json +++ b/tests/smoke_tests/data/drive_smoke_golden.json @@ -1,65 +1,65 @@ { "env": { - "avg_distance_per_infraction": 12.77927256565468, - "avg_speed_per_agent": 1.298731041161013, - "collision_rate": 0.012987013438563327, - "comfort_violation_count": 0.5917845969592338, - "distance_per_agent": 5.642795678341027, - "dnf_rate": 0.545454554743581, + "avg_distance_per_infraction": 14.347808436343545, + "avg_speed_per_agent": 1.308473795368558, + "collision_rate": 0.0, + "comfort_violation_count": 0.6011366736321222, + "distance_per_agent": 6.4906752450125555, + "dnf_rate": 0.5428571552038193, "early_reset_short": 0.0, - "episode_length": 14.380952380952381, - "episode_return": -1.17685332004126, - "infractions_per_agent": 0.44155844155844154, - "lane_center_rate": 0.2932935758328541, - "n": 231.0, - "num_goals_reached": 0.012987013438563327, - "obs/max": 1.0, - "obs/mean": 0.10966602416010574, - "obs/min": -1.0000000018626451, - "offroad_rate": 0.43290043560973496, + "episode_length": 16.538095238095238, + "episode_return": -1.3155656732264018, + "infractions_per_agent": 0.4523809523809524, + "lane_center_rate": 0.2616634910305341, + "n": 210.0, + "num_goals_reached": 0.01904761968624024, + "obs/max": 1.0000000009313226, + "obs/mean": 0.11018761165905744, + "obs/min": -1.0000000009313226, + "offroad_rate": 0.45238095649651117, "red_light_violation_rate": 0.0, "reward_components/ade": 0.0, - "reward_components/collision": -0.022441596373335107, - "reward_components/comfort": -0.4126336298463665, - "reward_components/goal": 0.004329004554779498, - "reward_components/lane_align": -0.027322572619690524, - "reward_components/lane_center": -0.012265666613750385, - "reward_components/offroad": -0.6959504439975276, - "reward_components/overspeed": 0.0, + "reward_components/collision": 0.0, + "reward_components/comfort": -0.4954122905220304, + "reward_components/goal": 0.014285714746940703, + "reward_components/lane_align": -0.02797369836341767, + "reward_components/lane_center": -0.014327265708042042, + "reward_components/offroad": -0.7670932460398902, + "reward_components/overspeed": -0.012818281075340651, "reward_components/red_light": 0.0, - "reward_components/reverse": -0.010638060892863578, + "reward_components/reverse": -0.012323847098187322, "reward_components/stop_sign": 0.0, - "reward_components/timestep": -8.928571238451601e-05, - "reward_components/velocity": 0.00015886408038481715, + "reward_components/timestep": -9.978571310750253e-05, + "reward_components/velocity": 0.000197039804022227, "score": 0.0, "spawn_failed": 0.0, "spawn_failed_goal": 0.0, - "spawn_reject_collision": 0.056277057051142575, + "spawn_reject_collision": 0.07142857298964546, "spawn_reject_empty_cell": 0.0, - "spawn_reject_offroad": 0.3636363656360866, - "spawn_reject_stop_line": 0.004329004329004329, + "spawn_reject_offroad": 0.40476190976443743, + "spawn_reject_stop_line": 0.014285714285714285, "stop_sign_violation_rate": 0.0, "stopped_at_reset": 0.0, - "total_distance_travelled": 1303.4858016967773, - "total_infractions": 102.0, - "velocity_progress_sum": 0.014079634136484042 + "total_distance_travelled": 1363.0418014526367, + "total_infractions": 95.0, + "velocity_progress_sum": 0.01484889875476559 }, "losses": { - "actor_grad_norm": 2.6260669231414795, - "approx_kl": 0.015240681357681751, - "clipfrac": 0.21147167682647705, - "critic_grad_norm": 0.75099778175354, - "ema_max": 3.856906831264496, + "actor_grad_norm": 1.2954603433609009, + "approx_kl": 0.011499513871967793, + "clipfrac": 0.22194667160511017, + "critic_grad_norm": 0.14509320259094238, + "ema_max": 3.0478071570396423, "ent_coef": 0.01, - "entropy": 2.447709321975708, - "explained_variance": 0.04323387145996094, - "filter_threshold": 0.03856906831264496, - "filtered_fraction": 0.11793611793611791, - "kept_fraction": 0.8820638820638821, - "masked_fraction": 0.205078125, - "old_approx_kl": 0.01860375516116619, - "policy_loss": -0.0062951878644526005, - "value_loss": 0.25460731983184814 + "entropy": 2.447284460067749, + "explained_variance": 0.002093076705932617, + "filter_threshold": 0.030478071570396423, + "filtered_fraction": 0.11620400258231112, + "kept_fraction": 0.8837959974176889, + "masked_fraction": 0.24365234375, + "old_approx_kl": 0.009929797612130642, + "policy_loss": -0.00016886356752365828, + "value_loss": 0.19429592788219452 }, "meta": { "bptt_horizon": 64, diff --git a/tests/smoke_tests/test_drive_train.py b/tests/smoke_tests/test_drive_train.py index b677a5f1a9..e99a666889 100644 --- a/tests/smoke_tests/test_drive_train.py +++ b/tests/smoke_tests/test_drive_train.py @@ -161,6 +161,7 @@ def _build_config(): "backbone_hidden_size": 256, "actor_hidden_size": 128, "critic_hidden_size": 64, + "fused_slot_encoder": False, # the CPU-only image has no triton }, ) _set_existing(args["rnn"], {"input_size": 256, "hidden_size": 256}) diff --git a/tests/unit_tests/test_carla_light_stop_lines.py b/tests/unit_tests/test_carla_light_stop_lines.py new file mode 100644 index 0000000000..ff78146a5c --- /dev/null +++ b/tests/unit_tests/test_carla_light_stop_lines.py @@ -0,0 +1,117 @@ +"""light_stop_line_overrides moves a bin light stop line along its lane onto the light's stop waypoint (the +junction entry RunningRedLightTest scores) and leaves lines within tolerance as exported.""" + +import math +import os +import sys +from pathlib import Path + +import numpy as np + + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +import data_utils.mirror_map_bin as mbin +from pufferlib.ocean.cosim import carla_bridge as cb + + +TOWN_BIN = Path(__file__).resolve().parents[2] / "pufferlib/resources/drive/binaries/carla/opendrive__Town01.bin" + + +def _line_mid(element): + line = element["stop_line"] + return 0.5 * (line[0] + line[3]), 0.5 * (line[1] + line[4]) + + +def _light_with_waypoints(transform, elements, along_offsets_m): + """One light whose stop waypoints sit along_offsets_m along the lane from each element's bin stop line.""" + waypoints = [] + for element, along_m in zip(elements, along_offsets_m): + heading = element["heading"] + mx, my = _line_mid(element) + cx, cy = transform.bin_to_loc(mx + along_m * math.cos(heading), my + along_m * math.sin(heading)) + waypoints.append( + { + "x": cx, + "y": cy, + "yaw_deg": transform.bin_heading_to_yaw(heading), + "lane_width": 3.5, + "backward": [], + "forward": [], + } + ) + return { + "id": 1, + "x": waypoints[0]["x"], + "y": waypoints[0]["y"], + "trigger_x": 0.0, + "trigger_y": 0.0, + "stop_waypoints": waypoints, + } + + +def _setup(): + traffic = mbin.read_bin(TOWN_BIN)["traffic"] + light_elements = [j for j, t in enumerate(traffic) if t["type"] == 1] + transform = cb.CarlaTransform("Town01", offset=cb.town_offset(TOWN_BIN)) + return traffic, light_elements, transform + + +def test_line_moves_onto_waypoint_eight_metres_upstream(): + traffic, light_elements, transform = _setup() + j = light_elements[0] + light = _light_with_waypoints(transform, [traffic[j]], [-8.0]) + + indices, lines = cb.light_stop_line_overrides([light], [[j]], transform, TOWN_BIN) + + assert indices.tolist() == [j] + heading = traffic[j]["heading"] + mx, my = _line_mid(traffic[j]) + moved_mid = 0.5 * (lines[0, 0:2] + lines[0, 3:5]) + assert np.allclose(moved_mid, [mx - 8.0 * math.cos(heading), my - 8.0 * math.sin(heading)], atol=1e-3) + x1, y1, z1, x2, y2, z2 = traffic[j]["stop_line"] + moved_length = float(np.hypot(lines[0, 3] - lines[0, 0], lines[0, 4] - lines[0, 1])) + assert math.isclose(moved_length, math.hypot(x2 - x1, y2 - y1), abs_tol=1e-3) + assert lines[0, 2] == np.float32(z1) and lines[0, 5] == np.float32(z2) + + +def test_line_within_tolerance_is_left_as_exported(): + traffic, light_elements, transform = _setup() + j = light_elements[1] + light = _light_with_waypoints(transform, [traffic[j]], [0.5]) + + indices, lines = cb.light_stop_line_overrides([light], [[j]], transform, TOWN_BIN) + + assert len(indices) == 0 and lines.shape == (0, 6) + + +def test_waypoints_go_to_the_nearest_of_the_lights_elements(): + traffic, light_elements, transform = _setup() + a, b = light_elements[2], light_elements[3] + light = _light_with_waypoints(transform, [traffic[a], traffic[b]], [-8.0, 0.3]) + + indices, lines = cb.light_stop_line_overrides([light], [[a, b]], transform, TOWN_BIN) + + assert indices.tolist() == [a] + heading = traffic[a]["heading"] + mx, my = _line_mid(traffic[a]) + moved_mid = 0.5 * (lines[0, 0:2] + lines[0, 3:5]) + assert np.allclose(moved_mid, [mx - 8.0 * math.cos(heading), my - 8.0 * math.sin(heading)], atol=1e-3) + + +def test_only_lines_the_route_crosses_in_travel_direction_move(): + traffic, light_elements, transform = _setup() + a, b = light_elements[2], light_elements[3] + light = _light_with_waypoints(transform, [traffic[a], traffic[b]], [-8.0, -8.0]) + heading = traffic[a]["heading"] + mx, my = _line_mid(traffic[a]) + route = np.array([[mx + s * math.cos(heading), my + s * math.sin(heading)] for s in np.arange(-20.0, 20.0, 1.0)]) + + indices, lines = cb.light_stop_line_overrides([light], [[a, b]], transform, TOWN_BIN, route_xy=route) + assert indices.tolist() == [a] and lines.shape == (1, 6) + + indices, _ = cb.light_stop_line_overrides([light], [[a, b]], transform, TOWN_BIN, route_xy=route[::-1]) + assert len(indices) == 0 # opposing traffic crosses the same line against its travel direction + + indices, _ = cb.light_stop_line_overrides([light], [[a, b]], transform, TOWN_BIN) + assert indices.tolist() == [a, b] # no route: every off line moves diff --git a/tests/unit_tests/test_carla_offset_calibration.py b/tests/unit_tests/test_carla_offset_calibration.py index 3102986b52..e80c718446 100644 --- a/tests/unit_tests/test_carla_offset_calibration.py +++ b/tests/unit_tests/test_carla_offset_calibration.py @@ -65,7 +65,7 @@ def generate_waypoints(self, spacing_m): def test_calibration_recovers_translation_and_height(): transform = cb.CarlaTransform("Town01", offset=cb.town_offset(str(TOWN_BIN))) offset, z_offset, residual_before, residual_after = cb.calibrate_town_offset( - _Map(transform), transform, str(TOWN_BIN) + cb.driving_waypoint_samples(_Map(transform)), transform, str(TOWN_BIN) ) assert residual_before > 0.2 assert residual_after < 0.02 diff --git a/tests/unit_tests/test_config_schema.py b/tests/unit_tests/test_config_schema.py index c6656586dc..9e26a2c8b8 100644 --- a/tests/unit_tests/test_config_schema.py +++ b/tests/unit_tests/test_config_schema.py @@ -304,6 +304,15 @@ def test_preload_map_cache_requires_use_map_cache(self): with self.assertRaisesRegex(pufferlib.APIUsageError, "preload_map_cache"): validate_puffer_drive_config(args, "test") + @patch("sys.argv", ["pufferl.py"]) + def test_partner_blindness_prob_min_must_not_exceed_max(self): + args = load_config("puffer_drive") + args["env"]["partner_blindness_prob_min"] = 0.5 + args["env"]["partner_blindness_prob"] = 0.1 + + with self.assertRaisesRegex(pufferlib.APIUsageError, "partner_blindness_prob_min"): + validate_puffer_drive_config(args, "test") + if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/test_drive_config.py b/tests/unit_tests/test_drive_config.py index 3e9928ae12..364da1755a 100644 --- a/tests/unit_tests/test_drive_config.py +++ b/tests/unit_tests/test_drive_config.py @@ -96,7 +96,7 @@ def test_cli_override(self): @patch("sys.argv", ["pufferl.py"]) def test_training_performance_defaults(self): args = load_config("puffer_drive") - self.assertFalse(args["train"]["compile"]) + self.assertTrue(args["train"]["compile"]) self.assertEqual(args["train"]["precision"], "bfloat16") @patch("sys.argv", ["pufferl.py", "--train.learning-rate=0.5"]) diff --git a/tests/unit_tests/test_fused_slot_encoder.py b/tests/unit_tests/test_fused_slot_encoder.py new file mode 100644 index 0000000000..e7e478e71b --- /dev/null +++ b/tests/unit_tests/test_fused_slot_encoder.py @@ -0,0 +1,153 @@ +"""The fused Triton slot encoder must reproduce DriveBackbone._encode_and_pool under bf16 autocast. + +Agreement is checked at bf16 rounding level against the eager autocast path (gradient to the first maximal slot, as +torch.max does) and against an fp32 ground truth: +the fused path may not be further from the truth than the eager bf16 path it replaces. +""" + +import pytest +import torch +from torch import nn + +import pufferlib.pytorch + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="fused slot encoder needs CUDA + Triton") + +ENCODER_SHAPES = [(35, 10, 256), (30, 6, 256), (20, 11, 256), (4, 14, 128)] +# None = the table config of this GPU; the second config exercises 64-row tiles with 8 warps +KERNEL_CONFIGS = {"table": None, "wide": (64, 64, 8, 2, 128, 64, 8, 2, 64, 4)} +BATCH = 1024 +BF16_ULP_AT_FOUR = 2.0**-6 + + +def make_encoder(in_features, width): + layers = [ + pufferlib.pytorch.layer_init(nn.Linear(in_features, width)), + nn.LayerNorm(width), + nn.ReLU(), + pufferlib.pytorch.layer_init(nn.Linear(width, width)), + ] + return nn.Sequential(*layers).cuda() + + +def reference_encode_and_pool(objects, valid_counts, encoder): + valid_mask = torch.arange(objects.shape[1], device=objects.device) < valid_counts.unsqueeze(1) + encoded = encoder(objects) + masked = encoded.masked_fill(~valid_mask.unsqueeze(2), torch.finfo(encoded.dtype).min) + pooled = masked.max(dim=1).values + return torch.where(valid_counts.unsqueeze(1) == 0, pooled.new_zeros(()), pooled) + + +def make_inputs(slots, features, width, seed): + generator = torch.Generator(device="cuda").manual_seed(seed) + obs_width = 100 + slots * features + obs = torch.zeros(BATCH, obs_width, device="cuda") + counts = torch.randint(0, slots + 1, (BATCH,), device="cuda", generator=generator) + counts[: BATCH // 8] = 0 + counts[BATCH // 8 : BATCH // 4] = slots + objects = obs[:, 100:].view(BATCH, slots, features) + data = torch.rand(BATCH, slots, features, device="cuda", generator=generator) * 1.6 - 0.8 + valid = torch.arange(slots, device="cuda")[None, :] < counts[:, None] + objects.copy_(data * valid[:, :, None]) + duplicated = counts >= 2 + objects[duplicated, 1] = objects[duplicated, 0] + grad_pooled = (torch.randn(BATCH, width, device="cuda", generator=generator) * 0.1).to(torch.bfloat16) + return objects, counts.long(), grad_pooled + + +def run(encoder, objects, counts, grad_pooled, path): + from pufferlib.ocean.fused_slot_encoder import fused_encode_and_pool + + for parameter in encoder.parameters(): + parameter.grad = None + if path == "fp32": + pooled = reference_encode_and_pool(objects, counts, encoder) + pooled.backward(grad_pooled.float()) + else: + with torch.autocast("cuda", dtype=torch.bfloat16): + if path == "eager": + pooled = reference_encode_and_pool(objects, counts, encoder) + else: + pooled = fused_encode_and_pool(objects, counts, encoder) + pooled.backward(grad_pooled) + grads = {name: parameter.grad.detach().clone() for name, parameter in encoder.named_parameters()} + return pooled.detach().float(), grads + + +def relative_error(a, b): + return ((a - b).norm() / b.norm().clamp_min(1e-12)).item() + + +@pytest.fixture(params=list(KERNEL_CONFIGS), ids=list(KERNEL_CONFIGS)) +def kernel_config(request): + import pufferlib.ocean.fused_slot_encoder as fse + + fields = KERNEL_CONFIGS[request.param] + fse.CONFIG_OVERRIDE = fse.KernelConfig(*fields) if fields else None + yield request.param + fse.CONFIG_OVERRIDE = None + + +@pytest.mark.parametrize("slots,features,width", ENCODER_SHAPES) +def test_fused_matches_autocast_reference(slots, features, width, kernel_config): + torch.manual_seed(0) + encoder = make_encoder(features, width) + objects, counts, grad_pooled = make_inputs(slots, features, width, seed=slots) + truth_pooled, truth_grads = run(encoder, objects, counts, grad_pooled, "fp32") + eager_pooled, eager_grads = run(encoder, objects, counts, grad_pooled, "eager") + fused_pooled, fused_grads = run(encoder, objects, counts, grad_pooled, "fused") + + assert fused_pooled.dtype == torch.float32 and fused_pooled.shape == eager_pooled.shape + pooled_diff = (fused_pooled - eager_pooled).abs() + assert pooled_diff.max().item() <= BF16_ULP_AT_FOUR + assert (pooled_diff > 0).float().mean().item() < 5e-3 + assert (fused_pooled - truth_pooled).abs().max().item() <= 1.05 * ( + eager_pooled - truth_pooled + ).abs().max().item() + 1e-3 + for name in eager_grads: + assert relative_error(fused_grads[name], eager_grads[name]) < 2e-2, name + fused_truth_error = relative_error(fused_grads[name], truth_grads[name]) + eager_truth_error = relative_error(eager_grads[name], truth_grads[name]) + assert fused_truth_error <= 1.05 * eager_truth_error + 2e-3, name + + +@pytest.mark.parametrize("slots,features,width", ENCODER_SHAPES[:1]) +def test_fused_inference_matches_training_forward(slots, features, width): + from pufferlib.ocean.fused_slot_encoder import fused_encode_and_pool + + torch.manual_seed(0) + encoder = make_encoder(features, width) + objects, counts, _ = make_inputs(slots, features, width, seed=7) + with torch.autocast("cuda", dtype=torch.bfloat16): + pooled_train = fused_encode_and_pool(objects, counts, encoder) + with torch.no_grad(): + pooled_eval = fused_encode_and_pool(objects, counts, encoder) + assert torch.equal(pooled_train.detach(), pooled_eval) + assert pooled_eval.dtype == torch.bfloat16 + + +def test_fused_handles_strided_upstream_gradient(): + """The cat backward hands the encoder a strided slice of the feature gradient.""" + from pufferlib.ocean.fused_slot_encoder import fused_encode_and_pool + + torch.manual_seed(0) + slots, features, width = ENCODER_SHAPES[0] + encoder = make_encoder(features, width) + head = pufferlib.pytorch.layer_init(nn.Linear(2 * width, 8)).cuda() + objects, counts, _ = make_inputs(slots, features, width, seed=11) + other_features = torch.randn(BATCH, width, device="cuda") + grads = {} + for path in ("eager", "fused"): + for parameter in encoder.parameters(): + parameter.grad = None + with torch.autocast("cuda", dtype=torch.bfloat16): + if path == "eager": + pooled = reference_encode_and_pool(objects, counts, encoder) + else: + pooled = fused_encode_and_pool(objects, counts, encoder) + output = head(torch.cat([other_features.to(pooled.dtype), pooled], dim=1)) + output.float().square().mean().backward() + grads[path] = {name: parameter.grad.detach().clone() for name, parameter in encoder.named_parameters()} + for name in grads["eager"]: + assert relative_error(grads["fused"][name], grads["eager"][name]) < 2e-2, name diff --git a/tests/unit_tests/test_lb1_ground_truth.py b/tests/unit_tests/test_lb1_ground_truth.py new file mode 100644 index 0000000000..1a5e7d67ae --- /dev/null +++ b/tests/unit_tests/test_lb1_ground_truth.py @@ -0,0 +1,90 @@ +"""lb1/ground_truth.py: the Leaderboard 1.0 criteria recomputed from plain CARLA-frame geometry.""" + +import os +import sys + + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from pufferlib.ocean.cosim.carla.lb1.ground_truth import LeaderboardOneGroundTruth, segments_intersect + + +EGO_EXTENT_X = 2.4 +LANE_WIDTH = 3.5 +# one light whose stop line crosses the +x lane at x = 50 (ego drives along +x, y = 0) +LIGHT = { + "id": 11, + "x": 55.0, + "y": 6.0, + "trigger_x": 50.0, + "trigger_y": 0.0, + "stop_waypoints": [{"x": 50.0, "y": 0.0, "yaw_deg": 0.0, "lane_width": LANE_WIDTH, "backward": [], "forward": []}], +} +STOP_SIGN = { + "id": 21, + "center": [120.0, 0.0, 0.0], + "extent": [2.0, 1.5, 1.0], + "yaw_deg": 0.0, + "lane": [120.0, 0.0, 0.0, 0.0], +} +IN_LANE = [0.0, LANE_WIDTH, -1.0, 0.0, 0.0, 0.0] # on the driving lane centre, no parking lane, lane yaw 0, no junction + + +def ego(x, speed=5.0, yaw_deg=0.0): + return [x, 0.0, 0.0, yaw_deg, speed, 0.0, 0.0, 0.0, 0.0, 0.0] + + +def snapshot(x, light_state="Red", speed=5.0, yaw_deg=0.0): + return {"ego": ego(x, speed, yaw_deg), "partners": [], "light_states": [[11, light_state]], "ego_light_id": -1} + + +def test_segments_intersect(): + assert segments_intersect((0, 0), (2, 0), (1, -1), (1, 1)) + assert not segments_intersect((0, 0), (2, 0), (3, -1), (3, 1)) + assert segments_intersect((0, 0), (2, 0), (2, 0), (2, 1)) # touching end point + + +def test_red_light_fires_once_when_the_tail_crosses_the_stop_line(): + truth = LeaderboardOneGroundTruth([LIGHT], [], EGO_EXTENT_X) + observation = {"collision": [], "lane": IN_LANE} + assert truth.update(snapshot(45.0), observation)["red_light"] == 0.0 # nose over, tail not yet + assert truth.update(snapshot(52.5), observation)["red_light"] == 1.0 # tail line 50.1..49.6 straddles x = 50 + assert truth.update(snapshot(52.6), observation)["red_light"] == 0.0 # same light never fires twice + + +def test_green_light_and_oncoming_lane_do_not_fire(): + truth = LeaderboardOneGroundTruth([LIGHT], [], EGO_EXTENT_X) + observation = {"collision": [], "lane": IN_LANE} + assert truth.update(snapshot(52.5, light_state="Green"), observation)["red_light"] == 0.0 + oncoming = snapshot(47.5, yaw_deg=180.0) # driving -x: the stop line faces the other lane + oncoming["ego"][4] = -5.0 + assert truth.update(oncoming, {"collision": [], "lane": IN_LANE})["red_light"] == 0.0 + + +def test_stop_sign_run_vs_stopped(): + observation = {"collision": [], "lane": IN_LANE} + run = LeaderboardOneGroundTruth([], [STOP_SIGN], EGO_EXTENT_X) + flags = [run.update(snapshot(x, speed=5.0), observation)["stop_sign"] for x in (100.0, 110.0, 119.5, 123.0, 150.0)] + assert flags == [0.0, 0.0, 0.0, 1.0, 0.0] # targeted at 100 m (20 m lookahead), inside at 119.5, fires once past it + + stopped = LeaderboardOneGroundTruth([], [STOP_SIGN], EGO_EXTENT_X) + for x, speed in ((100.0, 5.0), (119.5, 0.05), (123.0, 3.0), (150.0, 5.0)): + assert stopped.update(snapshot(x, speed=speed), observation)["stop_sign"] == 0.0 + + +def test_outside_lane_and_wrong_lane(): + truth = LeaderboardOneGroundTruth([], [], EGO_EXTENT_X) + far = [LANE_WIDTH / 2.0 + 1.4, LANE_WIDTH, -1.0, 0.0, 0.0, 0.0] + assert truth.update(snapshot(10.0), {"collision": [], "lane": far})["offroad"] == 1.0 + on_parking = [LANE_WIDTH / 2.0 + 1.4, LANE_WIDTH, 0.5, 2.5, 0.0, 0.0] # closer to a parking lane centre + assert truth.update(snapshot(10.0), {"collision": [], "lane": on_parking})["offroad"] == 0.0 + wrong_way = snapshot(10.0, yaw_deg=170.0) + assert truth.update(wrong_way, {"collision": [], "lane": IN_LANE})["offroad"] == 1.0 + junction = [0.0, LANE_WIDTH, -1.0, 0.0, 0.0, 1.0] + assert truth.update(wrong_way, {"collision": [], "lane": junction})["offroad"] == 0.0 + + +def test_collision_flag_follows_the_sensor_events(): + truth = LeaderboardOneGroundTruth([], [], EGO_EXTENT_X) + assert truth.update(snapshot(10.0), {"collision": ["vehicle.audi.a2"], "lane": IN_LANE})["collision"] == 1.0 + assert truth.update(snapshot(11.0), {"collision": [], "lane": IN_LANE})["collision"] == 0.0 diff --git a/tests/unit_tests/test_lb1_protocol.py b/tests/unit_tests/test_lb1_protocol.py new file mode 100644 index 0000000000..5ce03a306e --- /dev/null +++ b/tests/unit_tests/test_lb1_protocol.py @@ -0,0 +1,79 @@ +"""Leaderboard 1.0 co-sim client: wire protocol round trips, and the client stays Python 3.7 / pufferlib-free +(it runs inside carla_garage's `garage` conda env, the only Python the CARLA 0.9.10 egg supports).""" + +import ast +import os +import socket +import subprocess +import sys +from pathlib import Path + +import pytest + + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from pufferlib.ocean.cosim.carla.lb1 import lb1_protocol + + +LB1_DIR = Path(__file__).resolve().parents[2] / "pufferlib" / "ocean" / "cosim" / "carla" / "lb1" +CLIENT_SIDE_FILES = ("leaderboard_agent.py", "lb1_protocol.py", "run_evaluator.py") +CLIENT_PYTHON_FEATURE_VERSION = (3, 7) +GARAGE_PYTHON = Path.home() / "miniconda3" / "envs" / "garage" / "bin" / "python" +LARGE_PAYLOAD_FLOATS = 300_000 # a 5-route-km dense plan plus every Town04 driving waypoint + + +def test_round_trip_preserves_nested_message(): + client, server = socket.socketpair() + message = { + "type": "tick", + "snapshot": {"ego": [1.5, -2.0, 0.0], "partners": [], "light_states": [[7, "Red"]]}, + "x": None, + } + lb1_protocol.send_message(client, message) + assert lb1_protocol.recv_message(server) == message + client.close() + server.close() + + +def test_large_message_arrives_whole_over_chunked_reads(): + client, server = socket.socketpair() + payload = { + "type": "init", + "route": {"dense_plan": [[float(k), 0.5 * k, 0.0] for k in range(LARGE_PAYLOAD_FLOATS // 3)]}, + } + import threading + + sender = threading.Thread(target=lb1_protocol.send_message, args=(client, payload)) + sender.start() + received = lb1_protocol.recv_message(server) + sender.join() + assert received == payload + client.close() + server.close() + + +def test_closed_peer_reads_as_none(): + client, server = socket.socketpair() + client.close() + assert lb1_protocol.recv_message(server) is None + server.close() + + +@pytest.mark.parametrize("filename", CLIENT_SIDE_FILES) +def test_client_side_files_parse_as_python_37_without_pufferlib(filename): + source = (LB1_DIR / filename).read_text() + tree = ast.parse(source, filename=filename, feature_version=CLIENT_PYTHON_FEATURE_VERSION) + imported = [] + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported += [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom) and node.module: + imported.append(node.module) + assert not [name for name in imported if name.split(".")[0] in ("pufferlib", "data_utils", "torch")], imported + + +@pytest.mark.skipif(not GARAGE_PYTHON.is_file(), reason="garage conda env not installed on this machine") +@pytest.mark.parametrize("filename", CLIENT_SIDE_FILES) +def test_client_side_files_compile_with_garage_python(filename): + subprocess.run([str(GARAGE_PYTHON), "-m", "py_compile", str(LB1_DIR / filename)], check=True) diff --git a/tests/unit_tests/test_nuplan_bridge_partner_sizes.py b/tests/unit_tests/test_nuplan_bridge_partner_sizes.py index 704edc5bea..19a31d22d0 100644 --- a/tests/unit_tests/test_nuplan_bridge_partner_sizes.py +++ b/tests/unit_tests/test_nuplan_bridge_partner_sizes.py @@ -1,4 +1,5 @@ -"""floor_vru_partner_sizes: pedestrian/bicycle boxes grow to the floor on each axis, vehicles keep their size.""" +"""pad_vru_partner_sizes: pedestrian and bicycle boxes grow by the pad on both axes, vehicles keep their size, and +the inputs are never modified.""" import os import sys @@ -8,24 +9,22 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) -nuplan = pytest.importorskip("nuplan") - from pufferlib.ocean.cosim import nuplan_bridge as nb -def test_floor_grows_only_vru_boxes_below_the_floor(): +def test_pad_grows_pedestrian_and_bicycle_boxes_only(): types = np.array([1, 2, 3, 2], np.int32) # vehicle, pedestrian, bicycle, pedestrian lengths = np.array([4.5, 0.5, 1.7, 0.9], np.float32) widths = np.array([1.9, 0.6, 0.5, 0.7], np.float32) - ln, wd = nb.floor_vru_partner_sizes(types, lengths, widths, 0.8) - np.testing.assert_allclose(ln, [4.5, 0.8, 1.7, 0.9]) - np.testing.assert_allclose(wd, [1.9, 0.8, 0.8, 0.8]) + ln, wd = nb.pad_vru_partner_sizes(types, lengths, widths, 0.8) + np.testing.assert_allclose(ln, [4.5, 1.3, 2.5, 1.7], rtol=1e-6) + np.testing.assert_allclose(wd, [1.9, 1.4, 1.3, 1.5], rtol=1e-6) assert ln.dtype == np.float32 and wd.dtype == np.float32 -def test_floor_leaves_inputs_untouched(): +def test_pad_leaves_inputs_untouched(): types = np.array([2], np.int32) lengths = np.array([0.4], np.float32) widths = np.array([0.4], np.float32) - nb.floor_vru_partner_sizes(types, lengths, widths, 0.8) + nb.pad_vru_partner_sizes(types, lengths, widths, 0.8) assert lengths[0] == pytest.approx(0.4) and widths[0] == pytest.approx(0.4) diff --git a/tests/unit_tests/test_render_infraction_steps.py b/tests/unit_tests/test_render_infraction_steps.py new file mode 100644 index 0000000000..9a56b7b4e7 --- /dev/null +++ b/tests/unit_tests/test_render_infraction_steps.py @@ -0,0 +1,181 @@ +"""Infraction -> replay step links in the obs replay panels of the CARLA and nuPlan renderers.""" + +import os +import sys + +import numpy as np +import pandas as pd + +from pufferlib import viz + + +sys.path.insert( + 0, os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "scripts", "eval") +) + +import render_carla_obs_html as carla_render +import render_obs_html as nuplan_render + + +FRAMES, AGENT_CAP = 3000, 3 +TICK_DT = 0.05 +OFFSET = (10.0, 20.0) +COLLISION_FRAME = 600 + + +def _replay(ego_xy, collision_frames=()): + replay = { + "env": { + "num_goals": 2, + "goal_radius": 2.0, + "reward_conditioning": False, + "obs_slots_partners_n": 4, + "obs_slots_lane_n": 4, + "obs_slots_boundary_n": 4, + "obs_slots_traffic_controls_n": 1, + }, + "agent_f32": np.zeros((FRAMES, AGENT_CAP, 13), np.float32), + "agent_i32": np.zeros((FRAMES, AGENT_CAP, 10), np.int32), + "metrics_f32": np.zeros((FRAMES, AGENT_CAP, len(viz.METRIC_LABELS)), np.float32), + "puffer_f32": np.zeros((FRAMES, AGENT_CAP, 4), np.float32), + "traffic_i16": np.zeros((FRAMES, 1, 3), np.int16), + "raw_action": np.ones((FRAMES, 1, 2), np.float32), + "clipped_action": np.ones((FRAMES, 1, 2), np.float32), + "value": np.zeros((FRAMES, 1), np.float32), + "entropy": np.zeros((FRAMES, 1), np.float32), + "obs": None, + } + replay["agent_f32"][:, 0, :2] = ego_xy + replay["metrics_f32"][list(collision_frames), 0, nuplan_render.COLLISION_METRIC_IDX] = 1.0 + return replay + + +def _straight_ego_xy(): + return np.column_stack([0.5 * np.arange(FRAMES), np.full(FRAMES, 100.0)]) + + +def _carla_record(): + # bin (300, 100) = ego at frame 600; CARLA x = bin_x - tx, CARLA y = -(bin_y - ty) + return { + "route_id": "RouteScenario_3_rep0", + "status": "Completed", + "scores": {"score_composed": 50.0, "score_route": 100.0, "score_penalty": 0.5}, + "infractions": { + "collisions_vehicle": [ + "Agent collided against object with type=vehicle and id=1 at (x=290.0, y=-80.0, z=0.0)" + ], + "route_dev": ["Agent deviated from the route at (x=5000.0, y=0.0, z=0.0)"], + "outside_route_lanes": [ + "Agent went outside its route lanes for about 18.0 meters (1.5% of the completed route)" + ], + }, + } + + +def test_read_replay_zlib_chunk_subset_matches_full_read(tmp_path): + path = tmp_path / "route.replay.zlib" + viz.save_interactive_replay_zlib({}, _replay(_straight_ego_xy(), collision_frames=(7,)), str(path)) + header_full, chunks_full = viz.read_replay_zlib(path) + header_part, chunks_part = viz.read_replay_zlib(path, chunk_names=("agent_f32", "metrics_f32")) + assert header_part == header_full + assert set(chunks_part) == {"agent_f32", "metrics_f32"} + for name, arr in chunks_part.items(): + assert np.array_equal(arr, chunks_full[name]) + assert chunks_full["agent_f32"].nbytes > viz.REPLAY_HEADER_PROBE_BYTES + + +def test_carla_infraction_steps_link_nearest_ego_frame(tmp_path): + path = tmp_path / "route.replay.zlib" + viz.save_interactive_replay_zlib({}, _replay(_straight_ego_xy()), str(path)) + record = _carla_record() + log_meta = {"offset": list(OFFSET), "tick_dt": TICK_DT} + + steps = carla_render.infraction_steps(record, str(path), "Town01", log_meta) + collision, deviation, outside = ( + record["infractions"][key][0] for key in ("collisions_vehicle", "route_dev", "outside_route_lanes") + ) + assert steps[collision] == viz.replay_step_link(COLLISION_FRAME, COLLISION_FRAME * TICK_DT) + assert steps[deviation] == "beyond the replay" + assert outside not in steps + assert carla_render.infraction_steps(record, str(path), "", {}) == {} + + panel = carla_render.score_panel(record, "Town01", steps) + assert f'onclick="step={COLLISION_FRAME};' in panel + assert "beyond the replay" in panel + + +def test_carla_infraction_steps_fall_back_to_town_offset(tmp_path): + path = tmp_path / "route.replay.zlib" + viz.save_interactive_replay_zlib({}, _replay(_straight_ego_xy()), str(path)) + tx, ty = carla_render.TOWN_OFFSETS["Town01"] + message = f"Agent ran a red light 12 at (x={300.0 - tx:.3f}, y={-(100.0 - ty):.3f}, z=0.0)" + record = {"infractions": {"red_light": [message]}} + steps = carla_render.infraction_steps(record, str(path), "Town01", {}) + assert steps[message] == viz.replay_step_link(COLLISION_FRAME) + + +def _write_series(metrics_dir, name, values, extra=None): + row = {"scenario_name": "tok", "time_series_values": None if values is None else np.asarray(values, np.float64)} + row.update(extra or {}) + pd.DataFrame([row]).to_parquet(metrics_dir / f"{name}.parquet") + + +def test_nuplan_violation_steps_from_metric_series(tmp_path): + metrics_dir = tmp_path / "simulation" / "closed_loop" / "2026.01.01.00.00.00" / "metrics" + metrics_dir.mkdir(parents=True) + _write_series(metrics_dir, "corners_in_drivable_area", [1.0, 1.0, 0.0, 0.0]) + _write_series(metrics_dir, "driving_direction_compliance", [0.0, -1.0, -3.0]) + _write_series(metrics_dir, "time_to_collision_within_bound", [np.inf, 2.0, 0.9]) + _write_series(metrics_dir, "speed_limit_compliance", [0.0, 0.0, 0.0, 0.4]) + assert nuplan_render.violation_steps(str(tmp_path)) == { + "tok": { + "drivable_area_compliance": 2, + "driving_direction_compliance": 2, + "time_to_collision_within_bound": 2, + "speed_limit_compliance": 3, + } + } + + +def test_nuplan_violation_steps_skip_compliant_and_missing_series(tmp_path): + metrics_dir = tmp_path / "simulation" / "closed_loop" / "2026.01.01.00.00.00" / "metrics" + metrics_dir.mkdir(parents=True) + _write_series(metrics_dir, "speed_limit_compliance", [0.0, 0.0]) + _write_series(metrics_dir, "corners_in_drivable_area", None) + assert nuplan_render.violation_steps(str(tmp_path)) == {} + + +def test_nuplan_comfort_failures_carry_first_out_of_bound_step(tmp_path): + metrics_dir = tmp_path / "simulation" / "closed_loop" / "2026.01.01.00.00.00" / "metrics" + metrics_dir.mkdir(parents=True) + stats = { + "abs_ego_lon_jerk_within_bounds_stat_value": False, + "min_ego_lon_jerk_stat_value": 1.0, + "max_ego_lon_jerk_stat_value": 5.0, + } + _write_series(metrics_dir, "ego_lon_jerk", [1.0, 5.0, 1.0], stats) + comfort = nuplan_render.comfort_failures(str(tmp_path)) + assert comfort == { + "tok": [f"ego_lon_jerk 1.00..5.00 (bound |j| <= 4.13 m/s^3) · {nuplan_render.step_link(1)}"] + } + + +def test_nuplan_replay_collision_step(tmp_path): + path = tmp_path / "tok.replay.zlib" + viz.save_interactive_replay_zlib({}, _replay(np.zeros((FRAMES, 2)), collision_frames=(40, 41)), str(path)) + assert nuplan_render.replay_collision_step(str(path)) == 40 + viz.save_interactive_replay_zlib({}, _replay(np.zeros((FRAMES, 2))), str(path)) + assert nuplan_render.replay_collision_step(str(path)) is None + + +def test_nuplan_score_panel_links_violated_metrics_only(): + row = pd.Series( + {metric: 1.0 for metric, _, _ in nuplan_render.SCORED_METRICS} | {"score": 0.4, "scenario_type": "t"} + ) + row["speed_limit_compliance"] = 0.8 + row[nuplan_render.COLLISION_METRIC] = 0.0 + steps = {"speed_limit_compliance": 12, nuplan_render.COLLISION_METRIC: 40, "drivable_area_compliance": 3} + panel = nuplan_render.score_panel("tok", row, [], steps) + assert 'onclick="step=12;' in panel + assert 'onclick="step=40;' in panel and "(shadow env flag)" in panel + assert 'onclick="step=3;' not in panel diff --git a/tests/unit_tests/test_replay_ghost.py b/tests/unit_tests/test_replay_ghost.py index c7fc571960..7e589bd0f1 100644 --- a/tests/unit_tests/test_replay_ghost.py +++ b/tests/unit_tests/test_replay_ghost.py @@ -117,3 +117,13 @@ def ego_state(i): assert np.allclose(boxes[:, 0], [0.0, 1.0, 2.0, 3.0]) assert np.allclose(boxes[:, 1], 0.0) assert np.allclose(boxes[:, 2:], [0.25, 5.1, 2.3]) + + +def test_read_replay_zlib_chunk_subset(tmp_path): + path = tmp_path / "e.replay.zlib" + viz.save_interactive_replay_zlib({}, _replay(_ghost_values()), str(path)) + header_full, chunks_full = viz.read_replay_zlib(path) + header_part, chunks_part = viz.read_replay_zlib(path, chunk_names=("ghost_f32",)) + assert header_part == header_full + assert list(chunks_part) == ["ghost_f32"] + assert np.array_equal(chunks_part["ghost_f32"], chunks_full["ghost_f32"]) diff --git a/tests/unit_tests/test_vector_driver.py b/tests/unit_tests/test_vector_driver.py index 8a2f808961..8b82eedf78 100644 --- a/tests/unit_tests/test_vector_driver.py +++ b/tests/unit_tests/test_vector_driver.py @@ -5,9 +5,11 @@ """ import multiprocessing +import time import gymnasium import numpy as np +import pytest import pufferlib.vector from pufferlib import PufferEnv @@ -115,3 +117,68 @@ def test_preloaded_resources_are_inherited_when_default_is_forkserver(): assert process.exitcode == 0 assert queue.get() == ("forkserver", 1, True) + + +def _make_multiprocessing(worker_count, batch_size): + return pufferlib.vector.make( + [_make_minimal_env] * worker_count, + env_args=[[]] * worker_count, + env_kwargs=[{}] * worker_count, + backend="Multiprocessing", + num_envs=worker_count, + num_workers=worker_count, + batch_size=batch_size, + ) + + +def test_reset_drains_pending_workers_and_steps_again(): + """A reset issued right after recv() finds half the workers idle and half + mid-step; both halves must be reset without hanging and step afterwards.""" + worker_count = 4 + vecenv = _make_multiprocessing(worker_count, batch_size=2) + try: + obs, _ = vecenv.reset(seed=1) + assert obs.shape == (2, OBSERVATION_SIZE) + actions = np.zeros(vecenv.action_space.shape, dtype=np.int32) + vecenv.send(actions) + vecenv.recv() + obs, _ = vecenv.reset(seed=2) + assert obs.shape == (2, OBSERVATION_SIZE) + assert vecenv.pending_result.sum() == 2 + for _ in range(worker_count): + obs, *_ = vecenv.step(actions) + assert obs.shape == (2, OBSERVATION_SIZE) + finally: + vecenv.close() + + +class SlowStepEnv(MinimalEnv): + def step(self, actions): + time.sleep(30) + return super().step(actions) + + +def _make_slow_step_env(**kwargs): + return SlowStepEnv(**kwargs) + + +def test_dead_worker_raises_instead_of_hanging(monkeypatch): + monkeypatch.setattr(pufferlib.vector, "WORKER_POLL_TIMEOUT_SECONDS", 0.2) + vecenv = pufferlib.vector.make( + [_make_minimal_env, _make_slow_step_env], + env_args=[[]] * 2, + env_kwargs=[{}] * 2, + backend="Multiprocessing", + num_envs=2, + num_workers=2, + batch_size=2, + ) + try: + vecenv.reset(seed=1) + vecenv.send(np.zeros(vecenv.action_space.shape, dtype=np.int32)) + vecenv.processes[1].kill() + vecenv.processes[1].join() + with pytest.raises((RuntimeError, EOFError)): + vecenv.recv() + finally: + vecenv.close()