10 bernhard dev Cosim PR

README.md

@@ -247,6 +247,10 @@ To add a metric to the report: add its key to `TREND_METRICS` /
 | `collision_behavior` | `1` | `0` ignore, `1` stop, `2` remove |
 | `offroad_behavior` | `1` | Same options |
 | `traffic_light_behavior` | `1` | Same options |
+| `stop_sign_behavior` | `1` | Same options |
+| `traffic_lights_enabled` | `True` | Observe and enforce traffic lights |
+| `stop_signs_enabled` | `False` | Observe and enforce stop signs (CaRL RunStopSign2 semantics, penalised with `reward_stop_line`) |
+| `yield_signs_enabled` | `False` | Observe yield signs |
 | `control_mode` | `"control_vehicles"` | `"control_vehicles"`, `"control_agents"`, `"control_sdc_only"` |
 | `reward_conditioning` | `False` | Condition policy on reward weights |
 | `reward_randomization` | `False` | Randomize reward weights each episode |

data_utils/mirror_map_bin.py

@@ -29,6 +29,7 @@
 
 TRAFFIC_PHASE_SECTION_TAG = b"TLPHASE1"
 LANE_WIDTH_SECTION_TAG = b"LANEWID1"
+SPEED_ZONE_SECTION_TAG = b"SPDZONE1"
 
 
 def _is_lane(road_type: int) -> bool:
@@ -155,6 +156,13 @@ def read_bin(path: Path) -> dict:
                 if _is_lane(r["type"]):
                     r["widths"] = _read_f_array(f, r["S"])
 
+        zone_tag = f.read(len(SPEED_ZONE_SECTION_TAG))
+        has_zone_section = zone_tag == SPEED_ZONE_SECTION_TAG
+        assert has_zone_section or zone_tag == b"", f"unexpected bytes after width section in {path}"
+        for r in roads:
+            if _is_lane(r["type"]):
+                (r["speed_zone_idx"],) = _read("<i", f) if has_zone_section else (-1,)
+
         trailing = f.read()
         assert not trailing, f"{len(trailing)} unparsed trailing bytes in {path}"
 
@@ -172,6 +180,7 @@ def read_bin(path: Path) -> dict:
         "tracks_to_predict": tracks_to_predict,
         "has_phase_section": has_phase_section,
         "has_width_section": has_width_section,
+        "has_zone_section": has_zone_section,
     }
 
 
@@ -292,6 +301,11 @@ def write_bin(data: dict, path: Path):
             for r in roads:
                 if _is_lane(r["type"]) and r["S"]:
                     f.write(struct.pack(f"<{r['S']}f", *r["widths"]))
+        if data["has_zone_section"]:
+            f.write(SPEED_ZONE_SECTION_TAG)
+            for r in roads:
+                if _is_lane(r["type"]):
+                    f.write(struct.pack("<i", r["speed_zone_idx"]))
 
 
 def main():

docs/evaluation.md

@@ -133,7 +133,7 @@ puffer eval puffer_drive carla_fast \
   eval.capture_observations=false
 ```
 
-The default `eval.render_filter: null` disables filtered rendering. Multiple comma-separated columns use OR: `collision_rate,offroad_rate` selects scenarios where either metric is greater than zero. Use `eval.render_filter=all_infractions` to select collision, at-fault collision, offroad, and red-light failures.
+The default `eval.render_filter: null` disables filtered rendering. Multiple comma-separated columns use OR: `collision_rate,offroad_rate` selects scenarios where either metric is greater than zero. Use `eval.render_filter=all_infractions` to select collision, at-fault collision, offroad, red-light, and stop-sign failures.
 
 The filtered pass replays the selected map/seed pairs, captures standard
 interactive `.replay.zlib` files, renders one HTML page per replay, and builds a

pufferlib/config/evaluation/benchmark.yaml

@@ -16,7 +16,7 @@ env:
   reward_comfort: 0.05
   reward_velocity: 0.0025
   reward_lane_align: 0.025
-  reward_lane_center: 0.0038
+  reward_lane_center: 0.0075  # training-range max (REWARD_BOUNDS), not the paper's 0.0038
   reward_timestep: 0.000025
   reward_reverse: 0.005
   goal_speed: 3.0
@@ -37,6 +37,8 @@ env:
   spawn_heading_max_deg: 0.0
   pose_noise_xy_m: 0.0
   pose_noise_yaw_deg: 0.0
+  speed_limit_random_prob: 0.0
+  stagger_first_episode: false
 
 benchmarks:
   - name: carla
@@ -113,6 +115,7 @@ benchmarks:
       map_dir: pufferlib/resources/drive/binaries/nuplan
       use_neighbor_cache: false
       disable_red_light_infractions: true
+      disable_stop_sign_infractions: true
       goal_reach_requires_speed: false
 
   - name: nuplan_single
@@ -128,4 +131,14 @@ benchmarks:
       map_dir: pufferlib/resources/drive/binaries/nuplan
       use_neighbor_cache: false
       disable_red_light_infractions: true
+      disable_stop_sign_infractions: true
       goal_reach_requires_speed: false
+
+  # Closed-loop co-sim against the real external simulators (see
+  # pufferlib/ocean/evaluation_utils/cosim_evaluator.py), unlike the
+  # benchmarks above which replay pre-baked .bin maps in-process.
+  - name: carla_cosim
+    simulation_mode: carla_cosim
+    route_ids: [0, 1, 2]
+    base_carla_port: 2000
+    compute_config: scripts/cluster_configs/nyu_greene.yaml

pufferlib/config/puffer_drive.yaml

@@ -54,8 +54,10 @@ env:
   dt: 0.3
   # Base forward speed cap
   base_max_speed_mps: 20.0
-  # 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 [0.667, 1.5] -> [13.3, 30] m/s)
+  # 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
   # 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
@@ -64,6 +66,13 @@ env:
   # 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
+  # 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_delta_mps: 9.72
+  speed_limit_random_min_mps: 1.39
+  speed_limit_random_max_mps: 36.11
   # Collision behavior - options: "ignore", "stop", "remove"
   collision_behavior: stop
   # Offroad behavior - options: "ignore", "stop", "remove"
@@ -72,11 +81,14 @@ env:
   traffic_light_behavior: stop
   # Stop sign behavior - options: "ignore", "stop", "remove"
   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
   yield_signs_enabled: false
   # Skip red light violation detection (metrics, reward, behavior)
   disable_red_light_infractions: false
+  # Skip stop sign violation detection (metrics, reward, behavior); the sign obs still turn red/green
+  disable_stop_sign_infractions: false
   # Cycle the lights of a junction phase-by-phase (TLPHASE1 bins); false keeps independent lights (legacy)
   traffic_light_junction_phases: false
   # Share static map geometry (roads/grid/lane-graph) across envs using the same map
@@ -102,7 +114,11 @@ env:
   init_step_spread: false
   # sets an upper bound on the random init step
   # the sampled init step can be at most this close to the end of the episode;
+  # also the shortest first episode when stagger_first_episode is on
   init_step_min_horizon: 20
+  # training only: the first episode after every env (re)creation ends at a random step in
+  # [init_step_min_horizon, scenario_length], so sub-envs do not run their episodes in lockstep
+  stagger_first_episode: true
   # options: "control_vehicles", "control_agents", "control_tracks_to_predict", "control_sdc_only"
   control_mode: control_vehicles
   # Controller used by agent 0, the canonical SDC/target.
@@ -126,7 +142,8 @@ env:
   # --- Goal / Target ---
   # Goal regeneration policy - options: "finite" (regen full set when all reached), "rolling" (sliding window)
   goal_regen_mode: "finite"
-  # How the goal set is seeded - options: "route" (agent forward route), "map" (uniform free-roam)
+  # How the goal set is seeded - options: "route" (agent forward route), "map" (uniform free-roam),
+  # replay only: "gt" (raw logged trajectory points), "gt_map" (logged points snapped to the nearest lane center)
   goal_source: map
   # Append normalized lane-graph distance to the current goal lane on each lane obs row
   obs_goal_lane_distance: True
@@ -163,6 +180,8 @@ env:
   reward_reverse: 0.005
   reward_timestep: 2.5e-05
   reward_overspeed: 0.05
+  # m/s above the lane limit before the overspeed penalty fires
+  overspeed_tolerance_mps: 0.0
   reward_ade: 0.0
   # --- Map ---
   # Path to map used for training
@@ -177,6 +196,10 @@ env:
   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
+  # Ego lane-angle slot: true = signed heading error theta_f / pi (GIGAFLOW B.1), false = cos(theta_f)
+  obs_lane_heading_signed: false
+  # 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_slots_traffic_controls_n: 4
   # Fraction of segment observation slots to drop (reduces obs size)
   obs_dropout_lane: 0.5
@@ -265,6 +288,8 @@ train:
   evaluation_interval_epochs: null
   # Select multiple benchmarks with a comma-separated value. eg: carla_fast,womd_single
   evaluation_benchmarks: carla_fast
+  # Submit the carla_cosim SLURM debug hook at start/end of training.
+  cosim_debug_evals: False
   torch_deterministic: false
   cpu_offload: false
   device: cuda
@@ -354,7 +379,7 @@ eval:
   base_max_speed_mps: null
   # Overrides env.goal_regen_mode ("finite" or "rolling") of every selected benchmark; null keeps the benchmark's value.
   goal_regen_mode: null
-  # Overrides env.goal_source ("route", "map" or "gt") of every selected benchmark; null keeps the benchmark's value.
+  # Overrides env.goal_source ("route", "map", "gt" or "gt_map") of every selected benchmark; null keeps the benchmark's value.
   goal_source: null
   # Overrides env.goal_speed of every selected benchmark; null keeps the benchmark's value.
   goal_speed: null
@@ -367,6 +392,8 @@ eval:
   obs_slots_partners_n: null
   # Overrides env.disable_red_light_infractions of every selected benchmark; null keeps the benchmark's value.
   disable_red_light_infractions: null
+  # Overrides env.disable_stop_sign_infractions of every selected benchmark; null keeps the benchmark's value.
+  disable_stop_sign_infractions: null
   output_name: null
   # Results folder inside the run dir; also names the wandb key prefix final_<output_dir_name>_.
   output_dir_name: eval

pufferlib/config_schema.py

@@ -161,6 +161,8 @@ class GoalSource(Enum):
     route = 0
     map = 1
     gt = 2
+    external = 3
+    gt_map = 4
 
 
 class PackageName(Enum):
@@ -251,16 +253,22 @@ class DriveEnvConfig:
     dt: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT)
     base_max_speed_mps: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT)
     max_speed_mps: float | None = _constrained_field(POSITIVE_NUMBER_CONSTRAINT, default=None)
+    conditioning_speed_scale: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT)
     spawn_initial_speed: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT)
     spawn_lateral_offset_max_frac: float = _constrained_field(PROBABILITY_CONSTRAINT)
     spawn_heading_max_deg: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT)
     pose_noise_xy_m: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT)
     pose_noise_yaw_deg: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT)
+    speed_limit_random_prob: float = _constrained_field(PROBABILITY_CONSTRAINT)
+    speed_limit_random_delta_mps: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT)
+    speed_limit_random_min_mps: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT)
+    speed_limit_random_max_mps: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT)
     collision_behavior: InfractionBehavior = MISSING
     offroad_behavior: InfractionBehavior = MISSING
     traffic_light_behavior: InfractionBehavior = MISSING
     stop_sign_behavior: InfractionBehavior = MISSING
     disable_red_light_infractions: bool = MISSING
+    disable_stop_sign_infractions: bool = MISSING
     traffic_light_junction_phases: bool = MISSING
     traffic_lights_enabled: bool = MISSING
     stop_signs_enabled: bool = MISSING
@@ -276,6 +284,7 @@ class DriveEnvConfig:
     init_step: int = _constrained_field(NONNEGATIVE_INT_CONSTRAINT)
     init_step_spread: bool = MISSING
     init_step_min_horizon: int = _constrained_field(POSITIVE_INT_CONSTRAINT)
+    stagger_first_episode: bool = MISSING
     control_mode: ControlMode = MISSING
     sdc_controller: Controller = MISSING
     non_sdc_controller: Controller = MISSING
@@ -311,6 +320,7 @@ class DriveEnvConfig:
     reward_reverse: float = _constrained_field(FINITE_NUMBER_CONSTRAINT)
     reward_timestep: float = _constrained_field(FINITE_NUMBER_CONSTRAINT)
     reward_overspeed: float = _constrained_field(FINITE_NUMBER_CONSTRAINT)
+    overspeed_tolerance_mps: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT)
     reward_ade: float = _constrained_field(FINITE_NUMBER_CONSTRAINT)
     map_dir: str = MISSING
     num_maps: int = _constrained_field(POSITIVE_INT_CONSTRAINT)
@@ -319,6 +329,8 @@ class DriveEnvConfig:
     obs_slots_boundary_n: int = _constrained_field(NONNEGATIVE_INT_CONSTRAINT)
     obs_slots_partners_n: int = _constrained_field(NONNEGATIVE_INT_CONSTRAINT)
     obs_partner_relative_velocity: bool = MISSING
+    obs_lane_heading_signed: bool = MISSING
+    obs_lane_speed_limit: bool = MISSING
     obs_slots_traffic_controls_n: int = _constrained_field(NONNEGATIVE_INT_CONSTRAINT)
     obs_dropout_lane: float = _constrained_field(PROBABILITY_CONSTRAINT)
     obs_dropout_boundary: float = _constrained_field(PROBABILITY_CONSTRAINT)
@@ -392,6 +404,7 @@ class TrainingConfig:
     final_model_name: str = _constrained_field(NONEMPTY_STRING_CONSTRAINT)
     evaluation_interval_epochs: int | None = _constrained_field(POSITIVE_INT_CONSTRAINT)
     evaluation_benchmarks: str | None = MISSING
+    cosim_debug_evals: bool = MISSING
     torch_deterministic: bool = MISSING
     cpu_offload: bool = MISSING
     device: str | int = MISSING
@@ -466,6 +479,7 @@ class EvaluationConfig:
     max_goal_spacing: float | None = _constrained_field(POSITIVE_NUMBER_CONSTRAINT)
     obs_slots_partners_n: int | None = _constrained_field(POSITIVE_INT_CONSTRAINT)
     disable_red_light_infractions: bool | None = MISSING
+    disable_stop_sign_infractions: bool | None = MISSING
     output_name: str | None = MISSING
     output_dir_name: str = _constrained_field(NONEMPTY_STRING_CONSTRAINT)
     scenario_offset: int = _constrained_field(NONNEGATIVE_INT_CONSTRAINT)
@@ -605,8 +619,14 @@ def _validate_cross_field_constraints(config, context):
         _raise_config_error(context, "env.init_step_spread", "is only supported in replay mode")
     if env["init_step_spread"] and env["init_step_min_horizon"] >= env["scenario_length"]:
         _raise_config_error(context, "env.init_step_min_horizon", "must be smaller than env.scenario_length")
-    if env["goal_source"] == "gt" and env["simulation_mode"] != "replay":
-        _raise_config_error(context, "env.goal_source", "'gt' is only supported in replay mode")
+    if env["goal_source"] in ("gt", "gt_map") and env["simulation_mode"] != "replay":
+        _raise_config_error(context, "env.goal_source", f"'{env['goal_source']}' is only supported in replay mode")
+    if env["stagger_first_episode"] and env["init_step"] + env["init_step_min_horizon"] > env["scenario_length"]:
+        _raise_config_error(
+            context,
+            "env.stagger_first_episode",
+            "requires env.init_step + env.init_step_min_horizon <= env.scenario_length",
+        )
     if env["terminate_on_goal"] and (env["simulation_mode"] != "replay" or env["control_mode"] != "control_sdc_only"):
         _raise_config_error(context, "env.terminate_on_goal", "requires replay mode with control_sdc_only")
     if env.get("eval_mode") is not None and (