diff --git a/Makefile b/Makefile index d45918d330..2fa90dde08 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ PYTHON := .venv/bin/python # sources, so `make test` rebuilds it only when a .c/.h file actually changed. EXT_SUFFIX := $(shell $(PYTHON) -c "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX'))" 2>/dev/null) BINDING := pufferlib/ocean/drive/binding$(EXT_SUFFIX) -DRIVE_SOURCES := $(wildcard pufferlib/ocean/drive/*.c) $(wildcard pufferlib/ocean/drive/*.h) setup.py +DRIVE_SOURCES := $(wildcard pufferlib/ocean/drive/*.c) $(wildcard pufferlib/ocean/drive/*.h) pufferlib/ocean/env_binding.h setup.py .PHONY: help test rebuild ensure-test-deps test-unit test-eval test-c test-notebooks test-docker-smoke diff --git a/data_utils/datasets.yaml b/data_utils/datasets.yaml index c839a13496..85452c6dbe 100644 --- a/data_utils/datasets.yaml +++ b/data_utils/datasets.yaml @@ -74,3 +74,15 @@ carla_128_affine_zones_stop: behind the k_scaled_004x runs. license: >- CARLA assets are CC BY 4.0 — attribution required. + +carla_128_affine_zones_stop_ndshoulders_v20: + s3_uri: s3://pufferdrive-bins/carla/0.3.2/128_affine_zones_stop_ndshoulders_v20/ + size: ~107 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 + shoulders outside junctions made off-road (map set v20). The training + set behind the runs from k_scaled_0067 on. + license: >- + CARLA assets are CC BY 4.0 — attribution required. diff --git a/pufferlib/config/evaluation/benchmark.yaml b/pufferlib/config/evaluation/benchmark.yaml index fc28008bd1..d842d108ae 100644 --- a/pufferlib/config/evaluation/benchmark.yaml +++ b/pufferlib/config/evaluation/benchmark.yaml @@ -41,6 +41,7 @@ env: pose_noise_yaw_deg: 0.0 speed_limit_random_prob: 0.0 stagger_first_episode: false + scripted_pedestrian_ratio: 0.0 benchmarks: - name: carla @@ -51,7 +52,7 @@ benchmarks: num_maps: 8 max_agents_per_env: 50 max_scenarios_per_batch: null - scenario_length: 6000 + scenario_length: 9000 # 600 s at dt 0.0667, GIGAFLOW Table A4 control_mode: control_vehicles map_dir: pufferlib/resources/drive/binaries/carla diff --git a/pufferlib/config/puffer_drive.yaml b/pufferlib/config/puffer_drive.yaml index a435b22647..9510ea6c4e 100644 --- a/pufferlib/config/puffer_drive.yaml +++ b/pufferlib/config/puffer_drive.yaml @@ -87,10 +87,23 @@ env: yield_signs_enabled: false # Skip red light violation detection (metrics, reward, behavior) disable_red_light_infractions: false + # Red light violation = the stop line lies between the ego's tail clearance point (rear edge - 1 m, what CARLA's + # RunningRedLightTest scores) and its front edge while red, moving or standing: a yellow entry must clear the line + # by 1 m. Training-only margins (m): the front one makes the front edge stop that much earlier, the rear one would + # extend the tail clearance (0 = CARLA's clearance as is); eval and co-sim use the real edges + red_light_train_front_margin_m: 0.3 + red_light_train_rear_margin_m: 0.0 # 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 + # Training light cycles (s): green ~ U(green_min, 10), yellow ~ U(yellow_min, yellow_max), red ~ U(0.3, 10); eval runs a + # fixed 10 / 3 / 2 cycle. GIGAFLOW used U(1, 10) / U(1.5, 2.25) with a centre-crossing test; with CARLA's tail test the + # yellow must cover CARLA's 2 s (Leaderboard 1 scenarios 7-9) to 3.1 s (natural) or committed slow crossings become + # unavoidable violations, and CARLA has no green below 5 s + traffic_light_green_min_s: 3.0 + traffic_light_yellow_min_s: 2.0 + traffic_light_yellow_max_s: 3.0 # 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 @@ -201,6 +214,9 @@ env: # 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: true obs_slots_traffic_controls_n: 4 + # Append to every traffic-control row the seconds the light stays yellow / traffic_light_yellow_max_s (clipped at 1, + # 0 unless yellow); co-sims feed the simulator's own timer, clipped to the same training max + obs_traffic_light_yellow_remaining: false # Fraction of segment observation slots to drop (reduces obs size) obs_dropout_lane: 0.5 obs_dropout_boundary: 0.4 @@ -223,6 +239,9 @@ env: 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: observed lane widths (m) are clipped to [min, max]; the CARLA training maps span 2.85-4.4 m. max 0 disables. + obs_lane_width_clip_min_m: 0.0 + obs_lane_width_clip_max_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 @@ -254,6 +273,17 @@ env: phantom_braking_duration_seconds: 3.0 # Zero the steering command while a phantom brake is active (false keeps the policy's steering) phantom_braking_freeze_steering: false + # Scripted pedestrians: each env gets ceil(ratio * agents) walker slots (1.0-1.5 m boxes = CARLA walker plus the 0.8 m + # eval VRU pad). A slot appears beside a random moving car's lane, plays one behaviour and vanishes for 5-20 s: + # 40% pop-in crosser (CARLA Leaderboard 1.0 Scenario3): standing at the right kerb 10-17 m ahead, half a lane outside + # the road edge, waits 0-0.6 s, ramps over 0.3-0.7 s to 3.4 m/s + 0.4 m/s per lane (+-1) and crosses; + # 15% slow crosser: either kerb, 25-70 m ahead, waits 0-5 s, crosses at 0.8-2 m/s up to 34 deg off the perpendicular; + # 20% stander: either kerb, 15-60 m ahead, stands 8-25 s with a stopped time already running (0-60 s); + # 15% kerb walker: either kerb, walks 8-25 m along it at 0.7-1.7 m/s, with or against traffic, never on the road; + # 10% kerb stepper: pops in like a crosser, steps 0.5-1.5 m toward the road edge or away from it, stands 8-25 s. + # Shares and ranges are the PEDESTRIAN_* constants in constants.h. Scripted, never policy-controlled, no observations + # of their own; cars see them as partners and collide with them normally. 0 disables. + scripted_pedestrian_ratio: 0.0 policy: # Encoder layer (per-encoder embedding width) ego_input_size: 128 @@ -411,7 +441,7 @@ eval: # 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__. + # Results folder inside the run dir. output_dir_name: eval # First global scenario index this process evaluates; used to shard one benchmark across GPUs. scenario_offset: 0 @@ -429,6 +459,9 @@ eval: # Existing episode metrics CSV to replay without running the benchmark first. failure_replay_csv: null capture_observations: false + # Agents whose observations the replay pages keep: all, or infracting (agents with a collision, offroad, red light + # or stop sign infraction in the episode; the others' observation rows are zeroed, a 50-agent page shrinks ~4x). + capture_observations_agents: all observation_replay_wave_size: 16 observation_replay_writer_count: 4 # Options: sample | mode | mean diff --git a/pufferlib/config_schema.py b/pufferlib/config_schema.py index a9802c2ecf..81b3eee52c 100644 --- a/pufferlib/config_schema.py +++ b/pufferlib/config_schema.py @@ -268,8 +268,13 @@ class DriveEnvConfig: traffic_light_behavior: InfractionBehavior = MISSING stop_sign_behavior: InfractionBehavior = MISSING disable_red_light_infractions: bool = MISSING + red_light_train_front_margin_m: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT) + red_light_train_rear_margin_m: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT) disable_stop_sign_infractions: bool = MISSING traffic_light_junction_phases: bool = MISSING + traffic_light_green_min_s: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) + traffic_light_yellow_min_s: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) + traffic_light_yellow_max_s: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) traffic_lights_enabled: bool = MISSING stop_signs_enabled: bool = MISSING yield_signs_enabled: bool = MISSING @@ -332,6 +337,7 @@ class DriveEnvConfig: obs_lane_heading_signed: bool = MISSING obs_lane_speed_limit: bool = MISSING obs_slots_traffic_controls_n: int = _constrained_field(NONNEGATIVE_INT_CONSTRAINT) + obs_traffic_light_yellow_remaining: bool = MISSING obs_dropout_lane: float = _constrained_field(PROBABILITY_CONSTRAINT) obs_dropout_boundary: float = _constrained_field(PROBABILITY_CONSTRAINT) obs_lane_stride: int = _constrained_field(POSITIVE_INT_CONSTRAINT) @@ -347,6 +353,8 @@ class DriveEnvConfig: 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) + obs_lane_width_clip_min_m: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT) + obs_lane_width_clip_max_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) @@ -363,6 +371,7 @@ class DriveEnvConfig: phantom_braking_trigger_prob: float = _constrained_field(PROBABILITY_CONSTRAINT) phantom_braking_duration_seconds: float = _constrained_field(POSITIVE_NUMBER_CONSTRAINT) phantom_braking_freeze_steering: bool = MISSING + scripted_pedestrian_ratio: float = _constrained_field(NONNEGATIVE_NUMBER_CONSTRAINT) # Added by benchmark overlays after the base Hydra config is composed. eval_mode: bool | int = 0 @@ -498,6 +507,8 @@ class EvaluationConfig: max_rendered_failures: int | None = _constrained_field(POSITIVE_INT_CONSTRAINT) failure_replay_csv: str | None = MISSING capture_observations: bool = MISSING + # all | infracting: whose observations the replay pages keep (infracting zeroes the other agents' rows) + capture_observations_agents: str = "all" observation_replay_wave_size: int = _constrained_field(POSITIVE_INT_CONSTRAINT) observation_replay_writer_count: int = _constrained_field(POSITIVE_INT_CONSTRAINT) action_selection: ActionSelection = MISSING diff --git a/pufferlib/ensemble.py b/pufferlib/ensemble.py new file mode 100644 index 0000000000..fd4d5bf253 --- /dev/null +++ b/pufferlib/ensemble.py @@ -0,0 +1,222 @@ +"""Evaluation-time policy ensembles. + +An ensemble checkpoint carries the weights of several trained policies of one architecture. load_policy_weights turns a +freshly built policy plus a loaded checkpoint into either that policy or an EnsemblePolicy: its mean action (what the +evals execute) is the mean of the members' mean actions, while its action distribution and value, used for mode and +sampled actions, logged probabilities and distillation labels, are the equal-weight averages over the members. +write_ensemble_checkpoint packs member checkpoints into a run-dir-shaped folder (final_model.pt + config.yaml) that +every eval harness takes like a training run. +""" + +import copy +import math +import os +import threading + +import torch +import yaml +from torch import nn + +from pufferlib.pytorch import clean_policy_state_dict + + +ENSEMBLE_CHECKPOINT_KEY = "pufferdrive_ensemble" +ENSEMBLE_FORMAT_VERSION = 1 +ENSEMBLE_CONFIG_KEY = "ensemble" +MIN_MEMBER_COUNT = 2 +MIN_MIXTURE_VARIANCE = 1e-8 +CHECKPOINT_FILE_NAME = "final_model.pt" +RUN_CONFIG_FILE_NAME = "config.yaml" +ARCHITECTURE_CONFIG_KEYS = ("policy_name", "rnn_name", "rnn") +POLICY_CONFIG_KEYS_IGNORED = ("fused_slot_encoder",) # kernel choice, not architecture + + +class EnsemblePolicy(nn.Module): + def __init__(self, members): + super().__init__() + if len(members) < MIN_MEMBER_COUNT: + raise ValueError(f"an ensemble needs at least {MIN_MEMBER_COUNT} members, got {len(members)}") + for member in members: + if hasattr(member, "lstm"): + raise ValueError("recurrent policies cannot be ensembled: every member would need its own state") + self.members = nn.ModuleList(members) + self.is_continuous = members[0].is_continuous + self.fp32_heads = getattr(members[0], "fp32_heads", False) + # sample_logits' mean-action hook only sees the averaged probabilities; it runs right after forward_eval + self._member_probs_store = threading.local() + + @property + def member_count(self): + return len(self.members) + + def forward(self, observations, state=None): + return self.forward_eval(observations, state) + + def forward_eval(self, observations, state=None): + member_logits = [] + member_values = [] + for member in self.members: + logits, value = member.forward_eval(observations, state) + member_logits.append(logits) + member_values.append(value.float()) + if isinstance(member_logits[0], torch.Tensor): + self._member_probs_store.probs = [logits.float().softmax(dim=-1) for logits in member_logits] + return average_member_predictions(member_logits), torch.stack(member_values).mean(dim=0) + + def discrete_actions_to_continuous(self, actions): + return self.members[0].discrete_actions_to_continuous(actions) + + def discrete_probs_to_continuous_mean(self, probs): + """Mean of the members' mean actions for the batch of the preceding forward_eval call.""" + member_probs = getattr(self._member_probs_store, "probs", None) + if member_probs is None: + raise RuntimeError("forward_eval must run before discrete_probs_to_continuous_mean") + class_count = probs.shape[-1] + if member_probs[0].shape != (probs.numel() // class_count, class_count): + raise RuntimeError("discrete_probs_to_continuous_mean needs the member probabilities of the same batch") + # averaged in physical units: the env's long scaling is piecewise, so normalised member means do not average + mean_member_probs = torch.stack(member_probs).mean(dim=0) + return self.members[0].discrete_probs_to_continuous_mean(mean_member_probs).reshape(*probs.shape[:-1], -1) + + +def _average_discrete_logits(member_logits): + log_probs = torch.stack([logits.float() for logits in member_logits]).log_softmax(dim=-1) + return torch.logsumexp(log_probs, dim=0) - math.log(len(member_logits)) + + +def _average_normal_distributions(member_distributions): + locs = torch.stack([distribution.loc.float() for distribution in member_distributions]) + scales = torch.stack([distribution.scale.float() for distribution in member_distributions]) + mean = locs.mean(dim=0) + # moment-matched Gaussian: the mixture itself is not Gaussian + variance = (scales.square() + locs.square()).mean(dim=0) - mean.square() + return torch.distributions.Normal(mean, variance.clamp_min(MIN_MIXTURE_VARIANCE).sqrt()) + + +def average_member_predictions(member_logits): + """Logits of the equal-weight mixture of the members' action distributions, in the form sample_logits takes.""" + first = member_logits[0] + if isinstance(first, torch.Tensor): + return _average_discrete_logits(member_logits) + if isinstance(first, torch.distributions.Normal): + return _average_normal_distributions(member_logits) + if isinstance(first, (tuple, list)): + return type(first)(_average_discrete_logits(component) for component in zip(*member_logits)) + raise TypeError(f"cannot average policy outputs of type {type(first).__name__}") + + +def describe_policy(policy): + if isinstance(policy, EnsemblePolicy): + return f"ensemble of {policy.member_count} policies" + return "policy" + + +def is_ensemble_checkpoint(checkpoint): + return isinstance(checkpoint, dict) and ENSEMBLE_CHECKPOINT_KEY in checkpoint + + +def member_state_dicts(checkpoint): + """Cleaned policy state dicts of a loaded checkpoint: one for a trained policy, every member of an ensemble.""" + if not is_ensemble_checkpoint(checkpoint): + return [clean_policy_state_dict(checkpoint)] + ensemble = checkpoint[ENSEMBLE_CHECKPOINT_KEY] + format_version = ensemble.get("format_version") + if format_version != ENSEMBLE_FORMAT_VERSION: + raise ValueError(f"ensemble checkpoint format {format_version} is not the supported {ENSEMBLE_FORMAT_VERSION}") + state_dicts = ensemble["member_state_dicts"] + if len(state_dicts) < MIN_MEMBER_COUNT: + raise ValueError(f"ensemble checkpoint holds {len(state_dicts)} members; at least {MIN_MEMBER_COUNT} needed") + return [clean_policy_state_dict(state_dict) for state_dict in state_dicts] + + +def load_policy_weights(policy, checkpoint): + """Load a torch.load()ed checkpoint into a freshly built policy: the policy itself, or an EnsemblePolicy of copies.""" + state_dicts = member_state_dicts(checkpoint) + if len(state_dicts) == 1: + policy.load_state_dict(state_dicts[0]) + return policy + members = [] + for state_dict in state_dicts: + member = copy.deepcopy(policy) + member.load_state_dict(state_dict) + members.append(member) + return EnsemblePolicy(members) + + +def _policy_architecture(config): + policy_config = config.get("policy") or {} + return {key: value for key, value in policy_config.items() if key not in POLICY_CONFIG_KEYS_IGNORED} + + +def _check_member_matches_first(member_paths, member_configs, member_weights, member_idx, allow_env_mismatch): + first_path, path = member_paths[0], member_paths[member_idx] + first_config, config = member_configs[0], member_configs[member_idx] + for key in ARCHITECTURE_CONFIG_KEYS: + if first_config.get(key) != config.get(key): + raise ValueError(f"{path}: {key} differs from {first_path}; an ensemble needs one architecture") + if _policy_architecture(first_config) != _policy_architecture(config): + raise ValueError(f"{path}: policy config differs from {first_path}; an ensemble needs one architecture") + first_shapes = {name: tuple(tensor.shape) for name, tensor in member_weights[0].items()} + shapes = {name: tuple(tensor.shape) for name, tensor in member_weights[member_idx].items()} + if shapes != first_shapes: + differing = sorted( + name for name in shapes.keys() | first_shapes.keys() if shapes.get(name) != first_shapes.get(name) + ) + raise ValueError(f"{path}: parameter shapes differ from {first_path}: {differing}") + first_env, env = first_config.get("env") or {}, config.get("env") or {} + differing_env_keys = sorted(key for key in env.keys() | first_env.keys() if env.get(key) != first_env.get(key)) + if differing_env_keys and not allow_env_mismatch: + raise ValueError( + f"{path}: env config differs from {first_path} in {differing_env_keys}; allow_env_mismatch evaluates the " + "ensemble with the first member's env config" + ) + + +def write_ensemble_checkpoint(member_checkpoint_paths, output_dir, run_name, allow_env_mismatch=False): + """Pack member checkpoints into /final_model.pt + config.yaml (the first member's training config).""" + from pufferlib.ocean.cosim.arch import checkpoint_config_path + + member_paths = [os.path.abspath(path) for path in member_checkpoint_paths] + if len(member_paths) < MIN_MEMBER_COUNT: + raise ValueError(f"an ensemble needs at least {MIN_MEMBER_COUNT} member checkpoints, got {len(member_paths)}") + if len(set(member_paths)) != len(member_paths): + raise ValueError(f"member checkpoints repeat: {member_paths}") + member_configs = [] + member_weights = [] + for path in member_paths: + if not os.path.isfile(path): + raise FileNotFoundError(f"member checkpoint does not exist: {path}") + checkpoint = torch.load(path, map_location="cpu", weights_only=False) + if is_ensemble_checkpoint(checkpoint): + raise ValueError(f"{path} is itself an ensemble checkpoint; members must be trained policies") + with open(checkpoint_config_path(path)) as config_file: + member_configs.append(yaml.safe_load(config_file)) + member_weights.append(clean_policy_state_dict(checkpoint)) + if member_configs[0].get("rnn_name") is not None: + raise ValueError("recurrent policies cannot be ensembled: every member would need its own state") + for member_idx in range(1, len(member_paths)): + _check_member_matches_first(member_paths, member_configs, member_weights, member_idx, allow_env_mismatch) + + checkpoint_path = os.path.join(output_dir, CHECKPOINT_FILE_NAME) + config_path = os.path.join(output_dir, RUN_CONFIG_FILE_NAME) + if os.path.exists(checkpoint_path): + existing = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + if not is_ensemble_checkpoint(existing): + raise FileExistsError(f"{checkpoint_path} is a trained checkpoint, not an ensemble; refusing to overwrite") + os.makedirs(output_dir, exist_ok=True) + config = copy.deepcopy(member_configs[0]) + config["run_name"] = run_name + config["wandb"] = False # no training run to attach eval results to + config[ENSEMBLE_CONFIG_KEY] = { + "member_checkpoints": member_paths, + "member_run_names": [member_config.get("run_name") for member_config in member_configs], + } + ensemble = { + "format_version": ENSEMBLE_FORMAT_VERSION, + "member_checkpoints": member_paths, + "member_state_dicts": member_weights, + } + torch.save({ENSEMBLE_CHECKPOINT_KEY: ensemble}, checkpoint_path) + with open(config_path, "w") as config_file: + yaml.dump(config, config_file) + return checkpoint_path, config_path diff --git a/pufferlib/ocean/cosim/arch.py b/pufferlib/ocean/cosim/arch.py index 1ab69f335b..5fda96301a 100644 --- a/pufferlib/ocean/cosim/arch.py +++ b/pufferlib/ocean/cosim/arch.py @@ -50,6 +50,7 @@ "pose_noise_yaw_deg": 0.0, "speed_limit_random_prob": 0.0, "stagger_first_episode": False, + "scripted_pedestrian_ratio": 0.0, } diff --git a/pufferlib/ocean/cosim/carla/README.md b/pufferlib/ocean/cosim/carla/README.md index 6e72985eb8..86b9a03cc8 100644 --- a/pufferlib/ocean/cosim/carla/README.md +++ b/pufferlib/ocean/cosim/carla/README.md @@ -42,10 +42,15 @@ srun --jobid=$JOBID --overlap \ 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_VRU_SIZE_PAD_M` +`COSIM_ZERO_PARTNER_STOPPED_TIME` (1 = vehicle partners' `seconds_stopped` obs feature held at 0, ego +and walkers 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`, +longest6 scripts set 0.8 m), `COSIM_OBS_SLOTS_PARTNERS_N` (partner slots in the policy's obs; +default = the checkpoint's training value, the slots are max-pooled so any count transfers; the +longest6 v1 script sets 40, GIGAFLOW's eval value), `COSIM_WALKER_SPEED_FLOOR_MPS` (diagnostic, +default 0: for 1.5 s after a walker first appears the policy reads it as moving at least this fast +along its facing direction, so CARLA's standing-then-ramping Scenario3 walker looks like a training +walker; world logs and collision checks stay raw), `COSIM_OBS_HTML`, `COSIM_DEBUG_CARLA_VIEW`, `COSIM_RECORD_INFRACTIONS`. ### Ego dynamics source @@ -78,6 +83,41 @@ of the ego no longer being a physically-simulated CARLA vehicle; the leaderboard's route/collision/infraction criteria still see it correctly since they only read its pose. +`COSIM_DYNAMICS_SOURCE=carla_pedals` drives the CARLA vehicle the way the leaderboard +prescribes: the agent only returns a `carla.VehicleControl` (steer, throttle, brake, reverse) from +`run_step`, nothing else touches the ego, and the gearbox stays automatic: the leaderboard resets +`manual_gear_shift` on every control, so the agent cannot select gears. Like nuPlan's +`ego_dynamics=nuplan_bicycle`, CARLA is the plant: pose, speed, yaw rate and the current gear are +re-read every tick, while acceleration and steering angle stay the jerk integrator's own state. The +jerk model moves the box centre along its heading, whereas the CARLA car turns about its rear axle, +so its centre travels up to 13 degrees inside its heading in a tight turn. The policy is therefore +given the direction of travel as its heading (body heading + the slip the commanded steering gives +at `rear_axle_to_center_m` ahead of the rear axle); with the body heading it cut every corner +(0.6 m at a 7 m radius). `controller.py`'s `PedalMap` turns acceleration and steering angle +into pedals and steer with constant gains and no feedback (the policy closes the loop). The gains +are per ego blueprint and per gear (`PEDAL_GAINS_BY_BLUEPRINT`: `mkz_2020` for longest6 v2, +`mkz2017` for v1, measured open loop on a straight road with each gear held); an unmeasured car is +refused. For the gear CARLA reports, throttle = the throttle that holds the current speed + the +target acceleration divided by that gear's acceleration per unit throttle, clipped at zero: CARLA's +engine braking makes the throttle the main brake too. The brake pedal supplies what exceeds the +gear's engine braking, capped at `max_engine_braking_mps2` so that hard braking acts at once, and +the whole deceleration while the gearbox shifts (gear 0: the engine is disconnected). Steer = path +curvature divided by the curvature one unit of steer gives, which falls with speed. A car that +has braked to below `STANDSTILL_SPEED_MPS` is held on the full brake, reads exactly stopped and, +like the jerk model, starts its acceleration again from zero. From there a positive acceleration +launches it forward; on a negative one it stays on the brake and never backs up +(`COSIM_PEDAL_REVERSE=1`: it launches backward with `VehicleControl.reverse`, the throttle from +the reverse gains, and the brake for a positive acceleration while it rolls backward; on longest6 +v1 that scored 88.9 against 91.3 driving score). The route +summary line prints how far the gains are off (`accel_bias_*`, per gear +`accel_bias_by_gear_mps2` = (ticks, bias) with reverse as gear -1, `yaw_rate_gain`). An automatic +shift leaves the car without drive for `gear_shift_seconds` (0.5 s on the v1 car), which no pedal +value can make up for and the policy has never seen: it would sprint through a yellow and arrive +late. So in this mode the policy sees that much less remaining yellow than CARLA reports and stops +at those lights instead (`COSIM_PEDAL_YELLOW_MARGIN=0`: it sees CARLA's value and sprints). Scripts +11 and 14 pass both switches through and tag the result folder `_reverse` / `_noyellowmargin`. +Not handled: road grade (gravity is not compensated). + ### Bin frame calibration The exported bins sit a constant fraction of a metre off the CARLA frame (2026-09-23, stored diff --git a/pufferlib/ocean/cosim/carla/controller.py b/pufferlib/ocean/cosim/carla/controller.py index 2bc8a20da2..c8416cabb5 100644 --- a/pufferlib/ocean/cosim/carla/controller.py +++ b/pufferlib/ocean/cosim/carla/controller.py @@ -16,6 +16,7 @@ """ import math +from typing import NamedTuple import numpy as np @@ -134,3 +135,156 @@ def stats(self): "yaw_err_max": float(ye.max()), "ticks": len(se), } + + +YAW_TRACKING_MIN_SPEED_MPS = 2.0 +REVERSE_GEAR = -1 # what CARLA reports as the gear while VehicleControl.reverse is set + + +class GearGains(NamedTuple): + speeds_mps: tuple # span where the automatic gearbox uses the gear; the values below are at these speeds + throttle_holding_speed: tuple + accel_per_throttle_mps2: tuple + + +class PedalGains(NamedTuple): + gears: tuple # GearGains of forward gear 1, 2, ... + reverse: GearGains + gear_shift_seconds: float # the automatic gearbox gives no drive for this long while it shifts + max_engine_braking_mps2: float # the brake starts beyond this even in a gear that brakes harder: it acts at once + brake_per_mps2: float + curvature_per_steer_at_rest: float + curvature_per_steer_per_mps: float + rear_axle_to_center_m: float # the car turns about its rear axle, so its centre travels at an angle to its heading + + +# measured open loop on a straight road with each gear held (pedal steps from rest and from speed, steer steps): +# CARLA 0.9.15 and 0.9.10.1. First gear and reverse hold at rest the throttle at which the standing car starts to move. +PEDAL_GAINS_BY_BLUEPRINT = { + "vehicle.lincoln.mkz_2020": PedalGains( + gears=( + GearGains((0.0, 2.6, 5.3), (0.30, 0.35, 0.47), (4.4, 8.0, 10.4)), + GearGains((5.7, 10.2), (0.40, 0.45), (8.4, 13.2)), + GearGains((10.0, 16.8), (0.37, 0.47), (7.8, 9.7)), + GearGains((12.7, 22.1), (0.38, 0.47), (5.8, 8.0)), + GearGains((18.3, 26.5), (0.38, 0.48), (4.3, 5.3)), + GearGains((24.5, 28.0), (0.44, 0.52), (3.2, 3.6)), + ), + reverse=GearGains((0.0, 2.0), (0.21, 0.21), (4.8, 4.8)), + gear_shift_seconds=0.1, + max_engine_braking_mps2=2.0, + brake_per_mps2=0.22, + curvature_per_steer_at_rest=0.376, + curvature_per_steer_per_mps=-0.0023, + rear_axle_to_center_m=1.25, + ), + "vehicle.lincoln.mkz2017": PedalGains( + gears=( + GearGains((0.0, 2.0, 4.6), (0.27, 0.30, 0.46), (4.7, 5.6, 6.9)), + GearGains((1.3, 8.4), (0.18, 0.49), (4.1, 5.9)), + GearGains((2.6, 12.3), (0.21, 0.49), (3.0, 4.3)), + GearGains((10.2, 15.3), (0.38, 0.49), (2.8, 3.2)), + GearGains((13.9, 28.0), (0.44, 0.71), (2.7, 2.7)), + ), + reverse=GearGains((0.0, 2.0), (0.20, 0.20), (4.4, 4.4)), + gear_shift_seconds=0.5, + max_engine_braking_mps2=1.0, + brake_per_mps2=0.16, + curvature_per_steer_at_rest=0.387, + curvature_per_steer_per_mps=-0.0031, + rear_axle_to_center_m=1.3, + ), +} + + +def pedal_gains_for(blueprint): + if blueprint not in PEDAL_GAINS_BY_BLUEPRINT: + raise ValueError( + f"no measured pedal gains for ego blueprint {blueprint!r}; known: {sorted(PEDAL_GAINS_BY_BLUEPRINT)}" + ) + return PEDAL_GAINS_BY_BLUEPRINT[blueprint] + + +def pedals(gear_gains, max_engine_braking_mps2, brake_per_mps2, shifting, target_accel_mps2, speed_mps): + """(throttle, brake) for a target acceleration in one gear: one straight line each. Acceleration and speed are + along the direction the gear drives in.""" + accel_per_throttle = float(np.interp(speed_mps, gear_gains.speeds_mps, gear_gains.accel_per_throttle_mps2)) + throttle_holding_speed = float(np.interp(speed_mps, gear_gains.speeds_mps, gear_gains.throttle_holding_speed)) + # a released throttle brakes by the holding throttle's worth; while the gearbox shifts the engine is disconnected + engine_braking_mps2 = 0.0 if shifting else min(throttle_holding_speed * accel_per_throttle, max_engine_braking_mps2) + throttle = float(np.clip(throttle_holding_speed + target_accel_mps2 / accel_per_throttle, 0.0, 1.0)) + brake = float(np.clip(brake_per_mps2 * (-target_accel_mps2 - engine_braking_mps2), 0.0, 1.0)) + return throttle, brake + + +class PedalMap: + """Jerk-integrator targets -> (steer, throttle, brake, reverse) through the gains of the gear the automatic gearbox + is in, or of reverse; the policy is the only feedback loop.""" + + def __init__(self, gains): + self.gains = gains + self.reset() + + def reset(self): + self._engaged_gear = 1 + self._command = None + self._accel_rows = [] + self._yaw_rate_rows = [] + + def observe(self, speed_mps, yaw_rate_radps, tick_dt): + """CARLA's response to the previous tick's command, kept for stats().""" + if self._command is None: + return + target_accel_mps2, target_curvature_per_m, commanded_speed_mps, brake, gear = self._command + self._accel_rows.append((target_accel_mps2, (speed_mps - commanded_speed_mps) / tick_dt, brake, gear)) + if speed_mps > YAW_TRACKING_MIN_SPEED_MPS: + self._yaw_rate_rows.append((target_curvature_per_m * speed_mps, yaw_rate_radps)) + + def step(self, target_accel_mps2, target_curvature_per_m, speed_mps, gear, hold_standstill, reverse): + """One tick -> (steer, throttle, brake, reverse) of the carla.VehicleControl. Speed and acceleration are along + the car's forward axis, curvature and yaw rate positive to the right, like CARLA's steer. gear is the one CARLA + reports for the ego: 1.. forward, 0 while the gearbox shifts, REVERSE_GEAR in reverse.""" + if gear > 0: + self._engaged_gear = gear + forward_speed_mps = max(speed_mps, 0.0) + curvature_per_steer = ( + self.gains.curvature_per_steer_at_rest + self.gains.curvature_per_steer_per_mps * forward_speed_mps + ) + steer = float(np.clip(target_curvature_per_m / curvature_per_steer, -1.0, 1.0)) + if hold_standstill: + self._command = None + return steer, 0.0, 1.0, False + drive_direction = -1.0 if reverse else 1.0 + throttle, brake = pedals( + self.gains.reverse if reverse else self.gains.gears[self._engaged_gear - 1], + self.gains.max_engine_braking_mps2, + self.gains.brake_per_mps2, + gear == 0, + drive_direction * target_accel_mps2, + max(drive_direction * speed_mps, 0.0), + ) + self._command = (target_accel_mps2, target_curvature_per_m, speed_mps, brake, REVERSE_GEAR if reverse else gear) + return steer, throttle, brake, reverse + + def stats(self): + """How far the gains are off: realised minus target acceleration per pedal and per gear (ticks, bias), + realised over target yaw rate.""" + if not self._accel_rows: + return {} + target_accel, realised_accel, brake, gear = np.array(self._accel_rows).T + accel_error = realised_accel - target_accel + on_brake = brake > 0.0 + out = {"accel_abs_err_mps2": float(np.abs(accel_error).mean()), "ticks": len(accel_error)} + out["accel_bias_by_gear_mps2"] = { + int(g): (int((gear == g).sum()), round(float(accel_error[gear == g].mean()), 2)) for g in np.unique(gear) + } + if not on_brake.all(): + out["accel_bias_throttle_mps2"] = float(accel_error[~on_brake].mean()) + if on_brake.any(): + out["accel_bias_brake_mps2"] = float(accel_error[on_brake].mean()) + if self._yaw_rate_rows: + target_yaw_rate, realised_yaw_rate = np.array(self._yaw_rate_rows).T + target_power = float(target_yaw_rate @ target_yaw_rate) + if target_power > 0.0: + out["yaw_rate_gain"] = float(target_yaw_rate @ realised_yaw_rate) / target_power + return out diff --git a/pufferlib/ocean/cosim/carla/lb1/README.md b/pufferlib/ocean/cosim/carla/lb1/README.md index 03e28e2059..d4ef2c896f 100644 --- a/pufferlib/ocean/cosim/carla/lb1/README.md +++ b/pufferlib/ocean/cosim/carla/lb1/README.md @@ -68,8 +68,9 @@ launch config "Cosim: CARLA leaderboard 1.0 agent (garage)" runs `run_evaluator. `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, +All 36 longest6 routes on the cluster: `scripts/kesai/14_carla_longest6_v1.sh` (WORKERS CARLA 0.9.10 server ++ evaluator pairs sharing the GPUs, each pulling the next `longest6_split` route run from one longest-first +queue, 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 @@ -97,3 +98,7 @@ As for the Leaderboard 2.0 agent: `result.json` (the leaderboard's), `telemetry/ `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/`. +The cluster scripts (`scripts/kesai/11_carla_longest6_v2.sh`, `14_carla_longest6_v1.sh`) render the pages of the +failed route runs only and then flatten every route folder with `scripts/eval/tidy_carla_run.py`: `results.csv` +at the top, `routes/route__s.result.json` for every route run plus `.telemetry.csv` and +`.world_log.npz` of the failed ones, `details/` for the timing and logs; `TIDY=0` keeps the route folders. diff --git a/pufferlib/ocean/cosim/carla/lb1/lb1_protocol.py b/pufferlib/ocean/cosim/carla/lb1/lb1_protocol.py index 184d9379ff..b5019bc947 100644 --- a/pufferlib/ocean/cosim/carla/lb1/lb1_protocol.py +++ b/pufferlib/ocean/cosim/carla/lb1/lb1_protocol.py @@ -8,7 +8,7 @@ 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, + "motion": [...] | null, "control": [steer, throttle, brake, reverse] | null, "pd_flags": {...} | null, "carla_flags": {...} | null} client -> server {"type": "finish"} -> {"type": "done"} """ diff --git a/pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py b/pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py index aab4292083..8f08c04ac5 100644 --- a/pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py +++ b/pufferlib/ocean/cosim/carla/lb1/leaderboard_agent.py @@ -454,6 +454,7 @@ def _init_on_first_step(self): "extent": [body.extent.x, body.extent.y, body.extent.z], "box_z": body.location.z, "wheels": wheels, + "blueprint": self.vehicle.type_id, "state": actor_state_row(self.vehicle), }, "sparse_plan": plan_xyz(self._global_plan_world_coord), @@ -512,8 +513,10 @@ def _snapshot(self): ego_light = self.vehicle.get_traffic_light() return { "ego": actor_state_row(self.vehicle), + "ego_gear": self.vehicle.get_control().gear, "partners": [partner_row(a) for a in self._nearby_actors()], "light_states": [[lt.id, lt.get_state().name] for lt in self.lights], + "light_timers": [[lt.id, lt.get_elapsed_time(), lt.get_yellow_time()] for lt in self.lights], "ego_light_id": NO_LIGHT_ID if ego_light is None else ego_light.id, } @@ -632,8 +635,8 @@ def run_step(self, input_data, timestamp, sensors=None): 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) + steer, throttle, brake, reverse = reply["control"] + return carla.VehicleControl(steer=steer, throttle=throttle, brake=brake, reverse=bool(reverse)) def destroy(self, results=None): if not hasattr(self, "server"): diff --git a/pufferlib/ocean/cosim/carla/leaderboard_agent.py b/pufferlib/ocean/cosim/carla/leaderboard_agent.py index 75905a15ea..86bf076466 100644 --- a/pufferlib/ocean/cosim/carla/leaderboard_agent.py +++ b/pufferlib/ocean/cosim/carla/leaderboard_agent.py @@ -184,6 +184,7 @@ def _init_on_first_step(self): "extent": [body.extent.x, body.extent.y, body.extent.z], "box_z": body.location.z, "wheels": wheels, + "blueprint": self.vehicle.type_id, "state": actor_state_row(self.vehicle), }, "sparse_plan": plan_xyz(self._global_plan_world_coord), @@ -235,8 +236,10 @@ def _snapshot(self): ego_light = self.vehicle.get_traffic_light() return { "ego": actor_state_row(self.vehicle), + "ego_gear": self.vehicle.get_control().gear, "partners": [partner_row(a) for a in self._nearby_actors()], "light_states": [[lt.id, lt.get_state().name] for lt in self.lights], + "light_timers": [[lt.id, lt.get_elapsed_time(), lt.get_yellow_time()] for lt in self.lights], "ego_light_id": -1 if ego_light is None else ego_light.id, } @@ -445,8 +448,8 @@ def run_step(self, input_data, timestamp, sensors=None): if result.control is None: return carla.VehicleControl() - steer, throttle, brake = result.control - return carla.VehicleControl(steer=steer, throttle=throttle, brake=brake) + steer, throttle, brake, reverse = result.control + return carla.VehicleControl(steer=steer, throttle=throttle, brake=brake, reverse=reverse) def destroy(self, results=None): if not self.initialized: diff --git a/pufferlib/ocean/cosim/carla/shadow_ego.py b/pufferlib/ocean/cosim/carla/shadow_ego.py index ea4ae7f950..fc6954a585 100644 --- a/pufferlib/ocean/cosim/carla/shadow_ego.py +++ b/pufferlib/ocean/cosim/carla/shadow_ego.py @@ -16,14 +16,17 @@ 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}, + ego: {extent: [x, y, z], box_z, wheels: controller.wheel_rows, blueprint: actor type_id, 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 } + ego: ACTOR_STATE, ego_gear: the gear CARLA reports for the ego (1.. forward, 0 while the gearbox shifts, -1 reverse), + 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, + light_timers: [[light_id, get_elapsed_time(), get_yellow_time()], ...] (the light's own phase timer; required when + the checkpoint observes the remaining yellow, obs_traffic_light_yellow_remaining) } Environment variables: COSIM_DEVICE=cpu torch device for the policy @@ -34,18 +37,34 @@ "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). + "carla_pedals": CARLA's physics moves the ego, driven only by the (steer, throttle, + brake, reverse) the leaderboard applies. controller.PedalMap maps the jerk integrator's + acceleration and steering to them with the gains of the gear the automatic gearbox + is in; pose, speed, yaw rate and gear are re-read every tick. + COSIM_PEDAL_REVERSE=0 carla_pedals: 0 (default) = a standing car that is asked to go backwards stays on + the brake; 1 = it reverses (scored lower on longest6 v1). + COSIM_PEDAL_YELLOW_MARGIN=1 carla_pedals: 1 (default) = the policy sees the remaining yellow shortened by the + car's gear shift time, so it stops where a shift would make it late; 0 = it sees + CARLA's remaining yellow. 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. + 1: every vehicle 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. Walkers keep their real counter. 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_WALKER_SPEED_FLOOR_MPS=0 diagnostic: for WALKER_SPEED_FLOOR_WINDOW_S after a walker is first seen, the + policy's obs carry at least this speed along the walker's facing direction (CARLA's + Scenario3 walker stands ~0.3 s and ramps up over ~0.5 s; training walkers set off + at once). Shadow-env state only; world logs and collision checks stay raw. + COSIM_OBS_SLOTS_PARTNERS_N=40 + partner slots in the policy's obs (default: the checkpoint's training value; + the slots are max-pooled, so any count transfers; 40 = GIGAFLOW's eval value) 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 + @@ -69,8 +88,14 @@ 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.carla.controller import ( + PedalMap, + TrackingController, + pedal_gains_for, + vehicle_geometry_from_wheels, +) from pufferlib.ocean.cosim.goals import RouteGoalWindow, route_goals_from_target_points +from pufferlib.ocean.cosim.light_timing import YellowCountdown from pufferlib.ocean.drive.drive import Drive @@ -100,6 +125,7 @@ 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 +WALKER_SPEED_FLOOR_WINDOW_S = 1.5 # covers the standstill + ramp of a CARLA Scenario3 walker 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 @@ -107,7 +133,10 @@ 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") +DYNAMICS_SOURCES = ("carla", "pufferdrive", "carla_pedals") +EGO_OBS_STEERING_IDX = 3 # write_ego_obs column order: speed, width, length, steering, accel_long, ... +DRIVE_WHEELBASE_PER_LENGTH = 0.6 # must equal the jerk model's wheelbase / sim_length +STANDSTILL_SPEED_MPS = 0.1 # carla_pedals: below this, in either direction, the ego counts as standing TELEMETRY_HEADER = ( "step,current_speed,target_speed,ego_action,goal_cursor,goal_dist_m," "near_light_dist_m,near_light_state," @@ -119,18 +148,6 @@ ) -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 @@ -166,6 +183,17 @@ def cosim_max_speed_mps(env_cfg): return max_speed_mps +def cosim_obs_slots_partners_n(): + """COSIM_OBS_SLOTS_PARTNERS_N as a positive int, or None when unset/empty (the checkpoint's training value).""" + raw = os.environ.get("COSIM_OBS_SLOTS_PARTNERS_N") + if not raw: + return None + obs_slots_partners_n = int(raw) + if obs_slots_partners_n < 1: + raise ValueError(f"COSIM_OBS_SLOTS_PARTNERS_N must be >= 1, got {raw!r}") + return obs_slots_partners_n + + def env_flag(name, default="0"): raw = os.environ.get(name, default) if raw not in ("0", "1"): @@ -182,7 +210,7 @@ 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.""" + throttle, brake, reverse) in the carla modes; `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 @@ -202,10 +230,16 @@ def __init__(self, path_to_conf_file, route_tag): 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.obs_slots_partners_n = cosim_obs_slots_partners_n() self.zero_partner_stopped_time = env_flag("COSIM_ZERO_PARTNER_STOPPED_TIME") + self.pedal_reverse = env_flag("COSIM_PEDAL_REVERSE", "0") + self.pedal_yellow_margin = env_flag("COSIM_PEDAL_YELLOW_MARGIN", "1") 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}") + self.walker_speed_floor_mps = float(os.environ.get("COSIM_WALKER_SPEED_FLOOR_MPS", "0.0")) + if not (self.walker_speed_floor_mps >= 0.0): + raise ValueError(f"COSIM_WALKER_SPEED_FLOOR_MPS must be >= 0, got {self.walker_speed_floor_mps}") # Shadow agent pool == the training per-env cap self.num_agents = int(env_cfg["max_agents_per_env"]) @@ -215,12 +249,23 @@ def __init__(self, path_to_conf_file, route_tag): 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._world_log = {"ego": [], "partners": [], "lights": [], "light_yellow_s": []} self._last_partners = None + self._partner_raw_velocity = None self._partner_stopped_s = {} # CARLA actor id -> seconds at standstill + self._walker_first_seen_s = {} # CARLA actor id -> sim time the walker first appeared from pufferlib.ocean.drive import binding self._partner_stopped_speed_threshold = float(binding.AGENT_STOPPED_SPEED_THRESHOLD) + self._steering_limit_rad = float(binding.STEERING_LIMIT) + self._intent_accel_mps2 = 0.0 # carla_pedals: the jerk integrator's own state, never re-read from CARLA + self._intent_steering_rad = 0.0 + self._plant_speed_mps = 0.0 # along the car's forward axis, negative in reverse + self._plant_gear = 0 + self._standing = False # carla_pedals: too slow to count as moving + self._held = False # carla_pedals: braked to a stop and kept on the brake until the policy asks to move + self._moving_direction = 1.0 # carla_pedals: +1 forward, -1 reverse, of the last motion + self._yellow_margin_s = 0.0 # carla_pedals: the policy sees this much less yellow than CARLA has left 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 @@ -264,6 +309,7 @@ def _load_policy_and_env(self, town_bin): offroad_behavior="ignore", traffic_light_behavior="ignore", **({} if self.max_speed_mps is None else {"max_speed_mps": self.max_speed_mps}), + **({} if self.obs_slots_partners_n is None else {"obs_slots_partners_n": self.obs_slots_partners_n}), ), ) self.dynamics_model = arch.get("dynamics_model", "classic") @@ -275,11 +321,11 @@ def _load_policy_and_env(self, town_bin): import torch + import pufferlib.ensemble 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) - ) + checkpoint = torch.load(self.checkpoint, map_location=self.device, weights_only=False) + policy_state_dict = pufferlib.ensemble.member_state_dicts(checkpoint)[0] env_for_policy = self.env self._boundary_feat_ckpt = None bw = policy_state_dict.get("actor_backbone.boundary_encoder.0.weight") @@ -299,10 +345,13 @@ def _load_policy_and_env(self, town_bin): 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) + policy = policy_cls(env_for_policy, **self.cfg["policy"]).to(self.device) + self.policy = pufferlib.ensemble.load_policy_weights(policy, checkpoint) self.policy.eval() - print(f"[puffer_agent] loaded policy from {self.checkpoint}") + print( + f"[puffer_agent] loaded {pufferlib.ensemble.describe_policy(self.policy)} from {self.checkpoint} " + f"(obs partner slots {self.env.obs_slots_partners_n})" + ) def init_route(self, route): """Build the shadow env for this route's town and seed it from the simulator's ground truth (see @@ -395,9 +444,17 @@ def init_route(self, route): # 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) + self.last_yellow_remaining_s = np.zeros(self.num_traffic, np.float32) + self._yellow_countdown = YellowCountdown(self.env.traffic_light_yellow_min_s, self.dt) + + if self.dynamics_source == "carla_pedals": + self.controller = PedalMap(pedal_gains_for(ego.get("blueprint"))) + if self.pedal_yellow_margin: + # a sprint through a yellow can lose one gear shift without drive, which the policy does not know about + self._yellow_margin_s = self.controller.gains.gear_shift_seconds + else: + 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 @@ -416,15 +473,23 @@ def init_route(self, route): 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}" + f"vru_size_pad_m={self.vru_size_pad_m:g} walker_speed_floor_mps={self.walker_speed_floor_mps:g}" ) self.initialized = True # --- shadow-env sync (read-only w.r.t. the simulator) ------------------- + def _floored_walker_velocity(self, actor_id, now_s, heading, vx, vy): + first_seen_s = self._walker_first_seen_s.setdefault(actor_id, now_s) + if now_s - first_seen_s > WALKER_SPEED_FLOOR_WINDOW_S or math.hypot(vx, vy) >= self.walker_speed_floor_mps: + return vx, vy + return self.walker_speed_floor_mps * math.cos(heading), self.walker_speed_floor_mps * math.sin(heading) + def _read_states(self, partners): idx, x, y, z, h, vx, vy, yaw_rate, accel_long, stopped_s = [], [], [], [], [], [], [], [], [], [] + raw_vx, raw_vy = [], [] stopped_by_id = {} + now_s = self.step_count * self.tick_dt 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( @@ -440,6 +505,10 @@ def _read_states(self, partners): row[PARTNER_BOX_Z], row[PARTNER_EXTENT_Z], ) + raw_vx.append(bvx) + raw_vy.append(bvy) + if self.walker_speed_floor_mps > 0.0 and row[PARTNER_IS_WALKER] > 0.0: + bvx, bvy = self._floored_walker_velocity(int(row[PARTNER_ACTOR_ID]), now_s, bh, bvx, bvy) x.append(bx) y.append(by) z.append(bz) @@ -454,8 +523,10 @@ def _read_states(self, partners): 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 + self._partner_raw_velocity = (np.array(raw_vx, np.float32), np.array(raw_vy, np.float32)) if self.zero_partner_stopped_time: - stopped_s = [0.0] * len(partners) + # walkers keep their counter: held at 0, a walker that stays standing reads as about to cross forever + stopped_s = [s if row[PARTNER_IS_WALKER] > 0.0 else 0.0 for s, row in zip(stopped_s, partners)] return ( np.array(idx, np.int32), np.array(x, np.float32), @@ -479,27 +550,46 @@ def _read_sizes(self, partners): 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 + def _read_light_states(self, snapshot): + """Ground truth for every mapped bin traffic element (state, seconds it stays yellow), 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) + yellow_remaining_s = np.zeros(self.num_traffic, np.float32) + state_by_id = {int(light_id): cb.light_state_from_name(name) for light_id, name in snapshot["light_states"]} + remaining_by_id = self._yellow_remaining_by_id(snapshot, state_by_id) + ego_light_id = int(snapshot["ego_light_id"]) + # the ego's own light is written last so it wins over any other light mapped onto the same element + ordered_ids = [light_id for light_id in state_by_id if light_id != 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]: + ordered_ids.append(ego_light_id) + for light_id in ordered_ids: + for j in self.light_bin_indices_by_id[light_id]: if 0 <= j < self.num_traffic: - states[j] = state - return states - - def _sync_ego(self, ego_state, zero_velocity=False): + states[j] = state_by_id[light_id] + yellow_remaining_s[j] = max(remaining_by_id.get(light_id, 0.0) - self._yellow_margin_s, 0.0) + return states, yellow_remaining_s + + def _yellow_remaining_by_id(self, snapshot, state_by_id): + """{light_id: seconds it stays yellow} for the yellow lights: CARLA's own phase timer, or the + shortest-trained-yellow countdown from the onset where a scenario froze that timer.""" + timers = snapshot.get("light_timers") + if timers is None: + if self.env.obs_traffic_light_yellow_remaining: + raise KeyError("snapshot has no light_timers, which obs_traffic_light_yellow_remaining needs") + return {} + yellow_ids = {light_id for light_id, state in state_by_id.items() if state == cb.TRAFFIC_LIGHT_STATE_YELLOW} + countdown_s = self._yellow_countdown.tick(yellow_ids, self.step_count * self.tick_dt) + remaining_by_id = {} + for light_id, elapsed_s, yellow_time_s in timers: + light_id = int(light_id) + if light_id not in yellow_ids: + continue + timed_s = cb.carla_yellow_remaining_s(elapsed_s, yellow_time_s) + remaining_by_id[light_id] = countdown_s[light_id] if timed_s is None else timed_s + return remaining_by_id + + def _sync_ego(self, ego_state, zero_velocity=False, keep_integrator_state=False): """Overwrite agent 0 (ego) from the simulator's ground-truth pose/size.""" bin_state = self.transform.state_to_bin( ego_state[STATE_X], @@ -517,13 +607,44 @@ def _sync_ego(self, ego_state, zero_velocity=False): 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]) + steering = None + if keep_integrator_state: + bin_state = self._plant_state(bin_state) + steering = np.array([self._intent_steering_rad], np.float32) + self.env.set_agent_states( + np.array([0], np.int32), *[np.array([v], np.float32) for v in bin_state], steering=steering + ) 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 _plant_state(self, bin_state): + """carla_pedals: pose, velocity and yaw rate from CARLA's physics, acceleration from the jerk integrator.""" + x, y, z, heading, vx, vy, yaw_rate, _ = bin_state + self._plant_speed_mps = vx * math.cos(heading) + vy * math.sin(heading) + self.controller.observe(self._plant_speed_mps, -yaw_rate, self.tick_dt) + # without reverse a car that rolls backwards is stopped like a standing one + backward_standstill_mps = STANDSTILL_SPEED_MPS if self.pedal_reverse else math.inf + self._standing = -backward_standstill_mps < self._plant_speed_mps < STANDSTILL_SPEED_MPS + if not self._standing: + self._held = False + self._moving_direction = math.copysign(1.0, self._plant_speed_mps) + else: + braked_to_a_stop = not self._held and self._intent_accel_mps2 * self._moving_direction <= 0.0 + if braked_to_a_stop or (self._intent_accel_mps2 < 0.0 and not self.pedal_reverse): + # like the jerk model's, the acceleration starts again from zero; without reverse it cannot go below + self._held, self._intent_accel_mps2 = True, 0.0 + if self._plant_speed_mps * self._intent_accel_mps2 <= 0.0: + # held, or creeping against the launch: the jerk model would read that as a stop and cancel the launch + vx, vy = 0.0, 0.0 + # the jerk model moves the box centre along its heading; CARLA's centre swings about the rear axle instead + wheelbase_m = DRIVE_WHEELBASE_PER_LENGTH * 2.0 * self.ego_extent[0] + curvature_per_m = math.tan(self._intent_steering_rad) / wheelbase_m + travel_heading = heading + math.atan(self.controller.gains.rear_axle_to_center_m * curvature_per_m) + return (x, y, z, travel_heading, vx, vy, yaw_rate, self._intent_accel_mps2) + 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 @@ -532,6 +653,9 @@ def _sync(self, snapshot): self._ego_speed_carla = ego_speed_from_state(ego_state) if self.dynamics_source == "carla": self._sync_ego(ego_state) + if self.dynamics_source == "carla_pedals": + self._plant_gear = int(snapshot["ego_gear"]) + self._sync_ego(ego_state, keep_integrator_state=True) ego = self.env.get_global_agent_state() ebx, eby = float(ego["x"][0]), float(ego["y"][0]) @@ -542,7 +666,8 @@ def _sync(self, snapshot): 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) + raw_vx, raw_vy = self._partner_raw_velocity + self._last_partners = (states[1], states[2], states[4], raw_vx, raw_vy, sizes[1], sizes[2], is_walker) else: self._last_partners = None n_used = 1 + len(partners) @@ -552,8 +677,11 @@ def _sync(self, snapshot): 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.last_light_states, self.last_yellow_remaining_s = self._read_light_states(snapshot) + self.env.set_traffic_light_states( + self.last_light_states, + self.last_yellow_remaining_s if self.env.obs_traffic_light_yellow_remaining else None, + ) self.goal_window.sync(ebx, eby, float(ego["heading"][0])) if self._goal_lane_debug_file is not None: @@ -625,13 +753,33 @@ def _integrate(self, actions, ego_state): actions, self.target, motion=[x0, y0, yaw0_deg, speed_before, x1, y1, target_yaw_deg, speed_after, z1] ) + if self.dynamics_source == "carla_pedals": + self._intent_accel_mps2 = accel_after + steering_obs = float(np.asarray(self.env.observations)[0][EGO_OBS_STEERING_IDX]) + self._intent_steering_rad = steering_obs * self._steering_limit_rad + wheelbase_m = DRIVE_WHEELBASE_PER_LENGTH * 2.0 * self.ego_extent[0] + # the bin frame is y-mirrored: its left-positive curvature is a negative CARLA steer + curvature_per_m = -math.tan(self._intent_steering_rad) / wheelbase_m + self.target = (speed_after, target_yaw_deg) + launch_backwards = self._standing and accel_after < 0.0 + reverse = self.pedal_reverse and (launch_backwards or not self._standing and self._plant_speed_mps < 0.0) + control = self.controller.step( + accel_after, + curvature_per_m, + self._plant_speed_mps, + self._plant_gear, + self._held and (accel_after == 0.0 or (launch_backwards and not self.pedal_reverse)), + reverse, + ) + return StepResult(actions, self.target, control=control) + 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( + steer, throttle, brake = 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) + return StepResult(actions, self.target, control=(steer, throttle, brake, False)) def step(self, snapshot): """Shadow env <- simulator ground truth, policy, one dt of dynamics -> StepResult.""" @@ -752,6 +900,7 @@ def _record_world_state(self, obs): 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)) + self._world_log["light_yellow_s"].append(self.last_yellow_remaining_s.astype(np.float32)) def _write_world_log(self): Path(self.world_log_dir).mkdir(parents=True, exist_ok=True) @@ -762,6 +911,9 @@ def _write_world_log(self): 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), + light_yellow_s=np.stack(self._world_log["light_yellow_s"]) + if self._world_log["light_yellow_s"] + else np.zeros((0, 0), np.float32), route_goals=self.route_goals, stop_signs=self._stop_sign_lines, dense_route=dense, diff --git a/pufferlib/ocean/cosim/carla_bridge.py b/pufferlib/ocean/cosim/carla_bridge.py index 7ab3b37fc2..f8d986a8cf 100644 --- a/pufferlib/ocean/cosim/carla_bridge.py +++ b/pufferlib/ocean/cosim/carla_bridge.py @@ -262,8 +262,18 @@ def bin_heading_to_yaw(self, heading_rad): # CARLA traffic-light state name -> PufferDrive enum (datatypes.h:61-67) # UNKNOWN=0 RED=1 YELLOW=2 GREEN=3 OFF=4 +TRAFFIC_LIGHT_STATE_YELLOW = 2 TRAFFIC_LIGHT_STATE_OFF = 4 -LIGHT_STATE_BY_NAME = {"Red": 1, "Yellow": 2, "Green": 3, "Off": 4, "Unknown": 0} +LIGHT_STATE_BY_NAME = { + "Red": 1, + "Yellow": TRAFFIC_LIGHT_STATE_YELLOW, + "Green": 3, + "Off": TRAFFIC_LIGHT_STATE_OFF, + "Unknown": 0, +} +# scenario_runner's TrafficLightManipulator (Leaderboard 1 scenarios 7-9) sets a light's phase times to 1e9 s while it +# drives the states itself; the light's own timer then says nothing about the remaining yellow +CARLA_LIGHT_TIMER_FROZEN_S = 1.0e6 def light_state_from_name(state_name) -> int: @@ -275,6 +285,14 @@ def carla_light_to_puffer(state) -> int: return light_state_from_name(state.name) +def carla_yellow_remaining_s(elapsed_s, yellow_time_s): + """Seconds a CARLA light stays yellow by its own phase timer (get_yellow_time - get_elapsed_time), or None + when a scenario froze that timer.""" + if float(yellow_time_s) >= CARLA_LIGHT_TIMER_FROZEN_S: + return None + return max(float(yellow_time_s) - float(elapsed_s), 0.0) + + 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 = ( diff --git a/pufferlib/ocean/cosim/carla_cosim.py b/pufferlib/ocean/cosim/carla_cosim.py index 0fbb07fd4c..9e07b93d6a 100644 --- a/pufferlib/ocean/cosim/carla_cosim.py +++ b/pufferlib/ocean/cosim/carla_cosim.py @@ -40,6 +40,7 @@ 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.goals import RouteGoalWindow +from pufferlib.ocean.cosim.light_timing import YellowCountdown # carla_scenarios imports carla at module scope (it's pure CARLA-actor # manipulation, unlike this module's carla_bridge deps) and is only used # inside main() below, so it's imported there too. @@ -52,19 +53,6 @@ EGO_BLUEPRINT = "vehicle.lincoln.mkz_2017" # the longest6/CARLA-leaderboard hero vehicle -def clean_policy_state_dict(state_dict): - """Strip torch.compile / DDP prefixes. Inlined from pufferlib.pufferl to - avoid importing the training stack (wandb, neptune, the _C kernel, ...) in - the co-sim environment — same as cosim/carla/leaderboard_agent.py.""" - - 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 load_checkpoint_config(checkpoint): """The checkpoint's sibling config.yaml — the single source of truth for the obs/action layout the policy expects (and the policy arch). None if no ckpt.""" @@ -396,15 +384,16 @@ def main(): policy = None if cfg is not None: import torch + import pufferlib.ensemble import pufferlib.ocean.torch as drive_torch # pre-3.0 checkpoints keep action_type only in the env section cfg["policy"].setdefault("action_type", cfg["env"]["action_type"]) policy = getattr(drive_torch, cfg.get("policy_name", "Drive"))(env, **cfg["policy"]).to(args.device) - sd = torch.load(args.checkpoint, map_location=args.device, weights_only=False) - policy.load_state_dict(clean_policy_state_dict(sd)) + checkpoint = torch.load(args.checkpoint, map_location=args.device, weights_only=False) + policy = pufferlib.ensemble.load_policy_weights(policy, checkpoint) policy.eval() - print(f"[cosim] loaded policy from {args.checkpoint}") + print(f"[cosim] loaded {pufferlib.ensemble.describe_policy(policy)} from {args.checkpoint}") # jerk action: idx = j_long_idx*3 + j_lat_idx; (3,1) = +long jerk, straight (accelerate) dummy_action = np.full((1, 1), 3 * 3 + 1, dtype=np.int32) @@ -421,6 +410,7 @@ def main(): 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(cb.light_geometry_from_carla(lights, cmap), transform, town_bin) + yellow_countdown = YellowCountdown(env.traffic_light_yellow_min_s, dt) print( f"[cosim] carla: ego + {len(bg)} background + {len(lights)} lights; " f"offset={transform.tx:.1f},{transform.ty:.1f} z={transform.tz:+.2f}" @@ -510,12 +500,20 @@ def main(): zz = np.zeros_like(zf) env.set_agent_states(sp, zf, zf, zf, zz, zz, zz, zz, zz) states = np.zeros(num_traffic, np.int32) + yellow_remaining_s = np.zeros(num_traffic, np.float32) + light_states = [cb.carla_light_to_puffer(lt.get_state()) for lt in lights] + yellow_lights = [li for li, state in enumerate(light_states) if state == cb.TRAFFIC_LIGHT_STATE_YELLOW] + countdown_s = yellow_countdown.tick(yellow_lights, step * dt) for li, lt in enumerate(lights): - state = cb.carla_light_to_puffer(lt.get_state()) + remaining_s = 0.0 + if li in countdown_s: + timed_s = cb.carla_yellow_remaining_s(lt.get_elapsed_time(), lt.get_yellow_time()) + remaining_s = countdown_s[li] if timed_s is None else timed_s for j in light_map[li]: if 0 <= j < num_traffic: - states[j] = state - env.set_traffic_light_states(states) + states[j] = light_states[li] + yellow_remaining_s[j] = remaining_s + env.set_traffic_light_states(states, yellow_remaining_s if env.obs_traffic_light_yellow_remaining else None) ebx, eby = float(ego_bin["x"][0]), float(ego_bin["y"][0]) goal_window.sync(ebx, eby, float(ego_bin["heading"][0])) gx, gy = goal_window.window[:, 0], goal_window.window[:, 1] diff --git a/pufferlib/ocean/cosim/goals.py b/pufferlib/ocean/cosim/goals.py index d486481ca8..f48b45e38f 100644 --- a/pufferlib/ocean/cosim/goals.py +++ b/pufferlib/ocean/cosim/goals.py @@ -79,10 +79,11 @@ def sync(self, ego_x, ego_y, ego_heading): """Push the next window when the current one is exhausted (sliding: when any goal was consumed), or when the ego has driven past its current goal along the route (e.g. in the adjacent lane) without consuming it. An ego far off the route keeps its window: skipping - goals there cascades through the whole route.""" + goals there cascades through the whole route. Returns True when a window was pushed (the + env's observations are stale until recomputed).""" if not self.pushed: self._push(0) - return + return True current_goal_idx, goal_count = self.env.get_agent_goal_progress(self.agent_idx) if current_goal_idx >= goal_count: next_start = self.window_start + goal_count @@ -98,8 +99,23 @@ def sync(self, ego_x, ego_y, ego_heading): if self.sliding and not passed: next_start = self.current_index if next_start == self.window_start or next_start >= len(self.goals): - return # window still live, or route exhausted: the last window stays saturated + return False # window still live, or route exhausted: the last window stays saturated self._push(next_start) + return True + + def snapshot(self): + """State to hand back to restore(): the window and how many of its goals the env has consumed.""" + if not self.pushed: + raise RuntimeError("RouteGoalWindow.snapshot needs a pushed window") + current_goal_idx, _ = self.env.get_agent_goal_progress(self.agent_idx) + return self.window_start, self.current_index, current_goal_idx + + def restore(self, snapshot): + """Undo everything since snapshot(), e.g. the goals a forecast rollout consumed.""" + window_start, current_index, current_goal_idx = snapshot + self._push(window_start) + self.current_index = current_index + self.env.set_agent_goal_progress(self.agent_idx, current_goal_idx) def _push(self, start): sel = self.goals[start : start + self.num_goals] diff --git a/pufferlib/ocean/cosim/light_timing.py b/pufferlib/ocean/cosim/light_timing.py new file mode 100644 index 0000000000..5fa0a2d968 --- /dev/null +++ b/pufferlib/ocean/cosim/light_timing.py @@ -0,0 +1,24 @@ +"""Remaining-yellow timing for co-sim shadow envs (obs_traffic_light_yellow_remaining).""" + + +class YellowCountdown: + """Remaining yellow seconds for lights whose simulator exposes no usable phase timer: the shortest + trained yellow (traffic_light_yellow_min_s) counted down from the first tick the light was seen + yellow, never below one env step (a yellow light always has the current step left).""" + + def __init__(self, assumed_total_s, min_remaining_s): + if not assumed_total_s > 0.0 or not min_remaining_s > 0.0: + raise ValueError(f"YellowCountdown needs positive durations. Got: {assumed_total_s}, {min_remaining_s}") + self.assumed_total_s = float(assumed_total_s) + self.min_remaining_s = float(min_remaining_s) + self._onset_s = {} + + def tick(self, yellow_keys, now_s): + """-> {key: remaining_s} for the lights yellow at now_s; a light absent from a tick restarts its + countdown when it is next seen yellow.""" + now_s = float(now_s) + self._onset_s = {key: self._onset_s.get(key, now_s) for key in yellow_keys} + return { + key: max(self.assumed_total_s - (now_s - onset_s), self.min_remaining_s) + for key, onset_s in self._onset_s.items() + } diff --git a/pufferlib/ocean/cosim/nuplan/README.md b/pufferlib/ocean/cosim/nuplan/README.md index 4384784678..e8a6eba9b3 100644 --- a/pufferlib/ocean/cosim/nuplan/README.md +++ b/pufferlib/ocean/cosim/nuplan/README.md @@ -20,6 +20,90 @@ Zero changes to nuplan-devkit or CaRL's `carl_nuplan`. Hydra loads any planner v iteration; afterwards the shadow env owns it and the planner raises if nuPlan's ego pose ever diverges from the integrated one (the configs pin `perfect_tracking_controller`). +- Ego plant `ego_dynamics: nuplan_bicycle` (env `EGO_DYNAMICS=nuplan_bicycle`, script 9's default since + 2026-10-30 together with floor 0, the model's own actuator lag, the center cap at margin 0.95 and the floor + clamp: run 431 = 93.01 vs 93.12 for perfect tracking, driven exactly as CaRL drives the model; the + lag-compensated variant scores 93.48 (run 428) but inverts the plant's actuator model and is kept off; + `EGO_DYNAMICS=pufferdrive` is the old eval; script 9 tags the run + `_bicycle_vmin[_nolag][_centercap][_clamp]`): nuPlan's + `KinematicBicycleModel` is the plant, driven through + CaRL's `one_stage_controller` exactly like CaRL. Every step the shadow ego takes nuPlan's pose, speed + and yaw rate but keeps its own jerk-integrated accel/steering (the state the policy was trained to + observe; restamping the lagged plant's values made Drive's zero-crossing snaps fire every step and + swallow the policy's steering), and its post-step accel/steering become the model's acceleration / + steering-rate commands (`cosim/nuplan_bicycle.py`). The model low-passes both (gain dt/(dt+tau), tau + 0.2 s / 0.05 s) and has neither a speed floor nor a stop snap: `bicycle_min_speed_mps` + (`BICYCLE_MIN_SPEED_MPS`; 0 = CaRL's no-reverse driving, -2 = Drive's floor) is enforced on the command, + a shadow stop snap commands the model to exactly zero speed, and `bicycle_undo_low_pass` + (`BICYCLE_UNDO_LOW_PASS`) pre-compensates the lag so one model step lands on the targets (off = CaRL's + lag, the model realises a third of each accel step). The planner checks every step that nuPlan's accel and speed are the + model's response to its last command, so a perfect-tracking controller or mismatched time constants + abort the scenario. Comfort under the bicycle model: nuPlan scores the acceleration at the vehicle + center, which the devkit computes as rear-axle accel + d*yaw_rate^2 + d*yaw_accel (d = 1.46 m), so + pulling away at the 2.30 m/s^2 cap into a left turn reads above the 2.40 bound (Val14: ~250 scenarios; + perfect tracking drops the angular velocity and never sees it). `bicycle_center_accel_cap` + (`BICYCLE_CENTER_ACCEL_CAP`, tag `_centercap`) is CaRL's clip_angular_adjustment: the plant's accel is + reduced while turning so that reading stays under the cap (`bicycle_center_accel_cap_margin`, + `BICYCLE_CENTER_ACCEL_CAP_MARGIN`, leaves headroom for nuPlan's Savitzky-Golay overshoot); + `bicycle_max_yaw_jerk_radps3` (`BICYCLE_MAX_YAW_JERK`, tag `_yawjerk`) slows the steering command so the + yaw acceleration, whose derivative enters nuPlan's longitudinal jerk through the same formula, cannot jump + within a step; `accel_cap_mps2` (`EGO_ACCEL_CAP_MPS2`, tag `_acap`, any ego dynamics) lowers the ego's + accel cap itself (Val14: no gain over the center cap, costs progress). Run-426 anatomy knobs (2026-10-30, + default off): `partner_split_max_length_m` (`PARTNER_SPLIT_MAX_LENGTH_M`, tag + `_split`): partners longer than this are shown as equal boxes along their heading (training spawns at + most 7 m, nuPlan buses reach 12.7 m). `STANDSTILL_JERK_DEADBAND` (tag `_deadband`) overrides the shadow + env's standstill jerk deadband (arch.py 1.5 m/s^3; it was an anti-reverse-creep hack for perfect tracking + and froze five run-426 stalls whose start intent was below it; the bicycle floor does that job now). + `bicycle_floor_accel_clamp` (`BICYCLE_FLOOR_ACCEL_CLAMP`, default true, tag `_clamp`): the policy's + integrated acceleration is clamped to zero while the plant sits on the speed floor, mirroring Drive's stop + reset (nine run-426 stalls had it wound down to -5 m/s^2 at standstill); false reproduces run 426. Known + and left to training: Drive's jerk model keeps its kinematic reference at the box center while nuPlan's + bicycle rotates about the rear axle 1.46 m behind it, so the same steering cuts the nuPlan center to the + inside of a turn (four inside-rear corner collisions in run 426); no eval-side conversion. +- Ego plant `ego_dynamics: nuplan_tracker` (env `EGO_DYNAMICS=nuplan_tracker`, script 9 tag + `_tracker__step_veh[_reverse]`): GIGAFLOW's nuPlan evaluation ("we roll out an entire scenario + using GIGAFLOW policies for all agents ... and use the subsequent observed trajectory from the driver as its + trajectory"). The planner returns a real trajectory and nuPlan's own `two_stage_controller` drives the + `KinematicBicycleModel` along it, the setup every trajectory planner on the benchmark is scored with. + - Rollout: every planning step the shadow env takes the whole scene from nuPlan and is rolled out for + `horizon_seconds` (8 s) in `rollout_step_seconds` steps (`ROLLOUT_STEP_SECONDS`, 0.3 s = the training step, 27 + steps). The ego starts from the plant's full state (pose, speed, yaw rate, acceleration, steering angle) and is + driven by the policy; its rolled-out states, one every 0.3 s, are the trajectory. The rollout is thrown away + afterwards: a goal the ego consumes in its first step stays consumed (the other modes consume on their one + step too), later ones are undone (`RouteGoalWindow.snapshot` / `restore`). + - Other agents (`rollout_policy_vehicles`, `ROLLOUT_POLICY_VEHICLES`, default 0): every other agent keeps the + velocity observed at this step (a forecast, nothing is read from the log ahead); static objects stay where they + are; lights hold their state, a yellow runs out into red. With a count above 0 the nearest that many vehicles + and bicycles are driven by the policy instead. nuPlan gives no route for them, so their goals are invented: one + every `goal_spacing` along the lane they are on and, at a fork, the exit that turns least + (`Drive.set_agent_lane_goals`, renewed by `c_step` whenever one is reached); they start with zero acceleration + and steering, which nuPlan's tracks do not carry. Pedestrians always keep their velocity: nuPlan replays them, + so they do not yield to the ego, which a policy-driven one would. The reactive challenge simulates no bicycles + at all (IDM vehicles, replayed pedestrians and static objects). Why 0: with 32 (Val14 runs 2566 / 2567, + k_scaled_0072) a policy-driven queue leader pulls up ~6 m to the bin's stop line while nuPlan's vehicle stays + where it stopped, the ego's plan follows it a second later and both trackers act on that plan: LQR creeps into + the standing lead at 0.05 m/s, iLQR drives into it at 0.5 m/s (about 10 of the 31 and 10 of the 22 at-fault + collisions those runs added over the bicycle plant; 5 of 10 and 5 of 5 rerun locally with 0 are gone). + - What the trackers cost (k_scaled_0072, Val14 reactive, policy-driven forecast): LQR 0.897, iLQR 0.870, bicycle + plant 0.923. nuPlan's LQR picks one acceleration and one steering rate that reach the trajectory's state 1.0 s + ahead (accel = 0.91 x (speed in 1 s - speed now)). Once the plan stops within that second it brakes at 0.91 x + speed per second (0.9 m/s^2 at 1 m/s) and below 0.2 m/s with a 2 s time constant: a stop the policy plans + within 0.6 s takes seconds, never reaches rest and ends in the pedestrian or lead it was for (nuPlan counts a + moving ego as at fault). It also turns in before the plan does and drops a sub-second swerve. iLQR follows the + poses, but reaches 2.7 m/s^2 for a planned 2.3 on a straight launch (251 scenarios above nuPlan's 2.40 comfort + bound; `accel_cap_mps2: 2.0` still gives 2.45). + - No reversing (`rollout_allow_reverse: false`, `ROLLOUT_ALLOW_REVERSE`): an agent the policy would back up stays + where it stood, at rest, like `bicycle_min_speed_mps: 0`. Without it 3 of 40 sampled Val14 scenarios backed up + by up to 4.5 m and one left the drivable area. + - Tracker: `ego_controller=two_stage_controller` is nuPlan's default LQR tracker, `two_stage_ilqr_controller` + (this package) its iLQR tracker (`EGO_TRACKER=lqr|ilqr`). The iLQR preset drops nuPlan's `max_solve_time: 0.05`: + the full 20 iterations take ~45 ms on an idle workstation, so on a loaded node the wall-clock cap cuts them + short and the command depends on the hardware (CaRL's `LogActionPlanner` keeps the cap). The LQR tracker reads + the trajectory's speed 1.0 s ahead and iLQR tracks its poses for the whole horizon, which is why the rollout + covers the horizon instead of one step: a trajectory that ends after 0.3 s makes the LQR tracker accelerate + and brake at a fifth of what the policy asked for and makes the iLQR warm start raise. + - Cost: 27 policy calls per planning step instead of 1 (batched over the policy-driven vehicles, if any); 20 s to + 1.5 min per scenario on one core. - Route goals (`goal_source: external`): PDM-style lane-graph route -> goals every `goal_spacing` m. Planner `goal_source: roadblock` (env `GOAL_SOURCE=roadblock`) uses the centroids of the CaRL-corrected route @@ -35,7 +119,11 @@ Zero changes to nuplan-devkit or CaRL's `carl_nuplan`. Hydra loads any planner v 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 + 0.8 m. Eval hack `obs_lane_width_clip_min_m` / `obs_lane_width_clip_max_m` (env + `LANE_WIDTH_CLIP_MIN_M` / `LANE_WIDTH_CLIP_MAX_M`, Drive env keys via `env_overrides`): the + lane widths the policy observes are clipped to this range (the CARLA training maps span + 2.85-4.4 m, the nuPlan city bins reach 50 m at junction connectors and multi-lane + roadblocks). Off by default, script 9 sets 2.85-4.4 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 @@ -57,24 +145,39 @@ Zero changes to nuplan-devkit or CaRL's `carl_nuplan`. Hydra loads any planner v - `config/` — the Hydra configs that plug the planner in without touching nuplan-devkit or `carl_nuplan`: `planner/pufferdrive_planner.yaml` and two challenge presets, `closed_loop_{nonreactive,reactive}_agents_pufferdrive` - (differing only in the `observation` override: log-replay vs IDM agents). + (differing only in the `observation` override: log-replay vs IDM agents), + plus `ego_controller/two_stage_ilqr_controller.yaml` and its tracker + `ego_controller/tracker/ilqr_tracker_deterministic.yaml` for `nuplan_tracker`. Found via `hydra.searchpath=[pkg://pufferlib.ocean.cosim.nuplan.config, ...]`. Use city bins -## Visualization +## Visualization and run layout -- `COSIM_OBS_HTML=all|failures|0` (default `failures`): one interactive pufferlib.viz +- `COSIM_OBS_HTML=all|failures|infractions|0` (script 9: `infractions`): one interactive pufferlib.viz replay per scenario -> `$GROUP/obs_html/.html`, showing the exact observation vector the policy received each step (partner/lane/boundary/light slots, goals), its action/value/entropy/probabilities and the encoder max-pool winners. This is what the agent saw; road geometry is cropped to 250 m around the driven path. `failures` renders - (and keeps) only scenarios scoring below `COSIM_OBS_HTML_MAX_SCORE` (default 0.9); - `all` renders every scenario, `infractions` only collision / drivable-area / direction / - no-progress failures. `obs_html/index.html` is the same navigator as the self-play eval - replays (pages named `s___.html`, worst first). Pages can be - rendered later from the saved `.replay.zlib` files with `scripts/eval/render_obs_html.py`. -- nuPlan ground truth: `carl_visualization_callback` in `CALLBACKS`, or nuBoard on the - simulation logs (`scripts/eval/nuboard_failures.py` builds a failures-only nuBoard folder). + only scenarios scoring below `COSIM_OBS_HTML_MAX_SCORE` (default 0.9), `all` every + scenario, `infractions` only collision / drivable-area / direction / no-progress failures. + `obs_html/index.html` is the same navigator as the self-play eval replays (pages named + `s___.html`, worst first). Every `.replay.zlib` is deleted after + rendering: a page embeds its replay, `scripts/eval/render_obs_html.py` re-renders pages from + themselves and `pufferlib.viz.read_replay_zlib` reads a page like a `.replay.zlib`. +- nuPlan ground-truth videos are opt-in: `CALLBACKS="[simulation_log_callback, carl_visualization_callback]"` + renders one `.avi` per scenario (every step of all 1118 scenarios through matplotlib, 14 GB, plus a frame-shape + print per step in the job log) for the few the tidy keeps. Off since 2026-11-09: the kept simulation logs + render the same view offline (`analyze_nuplan_cosim.py`, CaRL's `render_simulation_log`). +- After the run `scripts/eval/tidy_nuplan_run.py` (runner env `TIDY`, default on) flattens `$GROUP` to + `overview.csv` (CLS, sub-metrics, loss importance), `scenarios.csv` (per scenario, worst first), `obs_html/`, + `simulation_log/.msgpack.xz` of the scenarios `COSIM_OBS_HTML` selects, and `details/` + (`aggregated_metrics.csv`, `metrics/` parquets, `config.yaml`, `overrides.yaml`, `runner_report.parquet`, the + SLURM log, and `videos/.mp4` of the kept scenarios when the video callback ran). A Val14 run shrinks + from 2300 files in 2000 folders (45 GB with videos) to ~220 files in 4 folders and ~0.2 GB for ~100 flagged + scenarios. The planner no longer pickles its policy into the simulation logs (28 MB -> 0.2 MB each); + `--strip-planner` rewrites the kept logs of older runs. `scripts/eval/nuplan_run_layout.py` resolves both + layouts (and the `details/simulation_log/` of runs tidied before 2026-11-09) for `compare_nuplan_evals.py`, + `nuplan_comfort_breakdown.py`, `analyze_nuplan_cosim.py` and the renderer. ## How to run @@ -106,6 +209,14 @@ uv run convert --preset nuplan --map_only \ --output /scratch/yw4142/datasets/ad/nuplan/maps ``` +Lane widths: the converter writes left + right boundary distance, and nuPlan lanes that include a +parking strip have their baseline 1.7 m from one boundary and 4.3 m from the other, so the policy +believed it had 3 m of room towards the oncoming lane. Eval hack `symmetric_lane_widths` (env +`SYMMETRIC_LANE_WIDTHS`, script 9 sets it): the planner rewrites every lane width as twice the +centerline's distance to the nearer GPKG boundary (`nuplan_bridge.symmetric_lane_widths`) and the +shadow env loads that copy from `bin_cache_dir` (`..symw.bin`, one rewrite per city). No bin +regeneration needed; 123Drive's `compute_lane_widths` writes the same definition since 2026-10-29. + Do NOT use `nuplan_bridge.write_drive_bin` for city bins — it writes an empty lane graph (`n: 0`), which silently zeroes the GPS goal columns. @@ -115,8 +226,11 @@ lane graph (`n: 0`), which silently zeroes the GPS goal columns. perception tracks, not telemetry. Zeros are exact for partner observations (never read for non-ego agents); ego telemetry comes from `EgoState.dynamic_car_state`. -- Partner `seconds_stopped` is approximated at the planner's 10 Hz cadence from - the synced speeds (same caveat as the CARLA side). +- Partner `seconds_stopped` is counted per nuPlan track token at the planner's + 10 Hz cadence from the synced speeds and injected with the partner states + (shadow slots change owner whenever a track appears or ends). It starts from + the planner's history buffer, about 2 s before the scenario; a track that + leaves the detections starts over. ## Remaining TODOs diff --git a/pufferlib/ocean/cosim/nuplan/config/ego_controller/tracker/ilqr_tracker_deterministic.yaml b/pufferlib/ocean/cosim/nuplan/config/ego_controller/tracker/ilqr_tracker_deterministic.yaml new file mode 100644 index 0000000000..9f39ab75a1 --- /dev/null +++ b/pufferlib/ocean/cosim/nuplan/config/ego_controller/tracker/ilqr_tracker_deterministic.yaml @@ -0,0 +1,10 @@ +# nuPlan's iLQR tracker without its wall-clock cap (stock max_solve_time 0.05 s: the full 20 iterations take ~45 ms on +# an idle workstation, so a loaded node cuts them short and the command depends on the hardware). Every solve now runs +# until convergence or max_ilqr_iterations. +defaults: + - ilqr_tracker + - _self_ + +ilqr_solver: + solver_params: + max_solve_time: null diff --git a/pufferlib/ocean/cosim/nuplan/config/ego_controller/two_stage_ilqr_controller.yaml b/pufferlib/ocean/cosim/nuplan/config/ego_controller/two_stage_ilqr_controller.yaml new file mode 100644 index 0000000000..733fd95e9b --- /dev/null +++ b/pufferlib/ocean/cosim/nuplan/config/ego_controller/two_stage_ilqr_controller.yaml @@ -0,0 +1,7 @@ +# nuPlan's two_stage_controller with the iLQR tracker instead of its default LQR tracker (ego_dynamics: nuplan_tracker). +_target_: nuplan.planning.simulation.controller.two_stage_controller.TwoStageController +_convert_: 'all' + +defaults: + - tracker: ilqr_tracker_deterministic + - motion_model: kinematic_bicycle_model diff --git a/pufferlib/ocean/cosim/nuplan/config/planner/pufferdrive_planner.yaml b/pufferlib/ocean/cosim/nuplan/config/planner/pufferdrive_planner.yaml index 5e9dc697ab..57a5c40421 100644 --- a/pufferlib/ocean/cosim/nuplan/config/planner/pufferdrive_planner.yaml +++ b/pufferlib/ocean/cosim/nuplan/config/planner/pufferdrive_planner.yaml @@ -23,6 +23,44 @@ pufferdrive_planner: startup_brake_jerk_cap_mps3: 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 + # eval hack: shadow env loads a cached copy of the city bin with lane width = 2 x distance to the nearer GPKG boundary + symmetric_lane_widths: false + # Ego plant. pufferdrive: the shadow env's jerk model integrates the ego, nuPlan's perfect_tracking_controller + # replays the pose. nuplan_bicycle: nuPlan's KinematicBicycleModel is the plant (run with + # ego_controller=one_stage_controller); the shadow ego takes nuPlan's pose/speed/yaw rate every step, keeps its own + # jerk-integrated accel/steering, and its post-step accel/steering are the model's commands. nuplan_tracker: the + # planner returns the ego trajectory of a forecast rollout and nuPlan's two-stage controller tracks it (see rollout_*). + # bicycle_min_speed_mps: speed floor enforced on the command (0 = CaRL's + # no-reverse driving, -2 = Drive's MAX_BACKWARD_SPEED). bicycle_undo_low_pass: pre-compensate the model's command + # lag (CaRL default off). The time constants must match the motion model (devkit 0.2 / 0.05). + ego_dynamics: pufferdrive + bicycle_min_speed_mps: 0.0 + bicycle_undo_low_pass: false + bicycle_accel_time_constant_s: 0.2 + bicycle_steering_time_constant_s: 0.05 + # bicycle_center_accel_cap (nuplan_bicycle only, CaRL's clip_angular_adjustment): nuPlan scores the center accel, + # which the devkit computes as rear accel + d*yaw_rate^2 + d*yaw_accel; the plant's accel is reduced while turning + # to keep it under accel_cap_mps2. accel_cap_mps2: ego accel cap [m/s^2] (null = 0.96 x nuPlan's 2.40 comfort bound); + # the policy's c_acc conditioning = cap / 2.5 must stay inside its training range. + bicycle_center_accel_cap: false + bicycle_center_accel_cap_margin: 1.0 # factor on accel_cap_mps2 for the center cap (savgol overshoot headroom) + bicycle_max_yaw_jerk_radps3: 0.0 # bound on the per-step change of yaw accel via the steering command; 0 = off + # clamp the policy's integrated accel to 0 while the plant sits on the speed floor (Drive's stop reset); false = run 426 + bicycle_floor_accel_clamp: true + accel_cap_mps2: null + # eval hack: partners longer than this [m] are shown as equal boxes along their heading (training max 7 m); 0 = off + partner_split_max_length_m: 0.0 + # ego_dynamics: nuplan_tracker, GIGAFLOW's nuPlan evaluation (run with ego_controller=two_stage_controller for the + # LQR tracker or two_stage_ilqr_controller for the iLQR tracker). Every step the shadow env takes the whole scene from + # nuPlan and is rolled out for horizon_seconds in rollout_step_seconds steps (0.3 = the training step); the policy + # drives the ego (from the plant's accel and steering) and the nearest rollout_policy_vehicles vehicles and bicycles, + # which follow their lane; pedestrians and the rest keep their observed velocity. The ego's rolled-out states are + # the trajectory nuPlan's tracker follows. rollout_allow_reverse: false = an agent the policy would back up stays at + # rest (as bicycle_min_speed_mps 0). rollout_policy_vehicles 0 = every other agent keeps its observed velocity: a + # policy-driven vehicle pulls up to the stop line where nuPlan's stays put, and the trackers drive the ego after it. + rollout_step_seconds: 0.3 + rollout_policy_vehicles: 0 + rollout_allow_reverse: false 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 dd23c1a4d3..d1a75ad458 100644 --- a/pufferlib/ocean/cosim/nuplan/planner.py +++ b/pufferlib/ocean/cosim/nuplan/planner.py @@ -1,7 +1,7 @@ import json import math from pathlib import Path -from typing import Dict, Optional, Type +from typing import Dict, NamedTuple, Optional, Type import numpy as np @@ -16,16 +16,37 @@ ) from nuplan.planning.simulation.trajectory.abstract_trajectory import AbstractTrajectory from nuplan.planning.simulation.trajectory.interpolated_trajectory import InterpolatedTrajectory +from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling + +from carl_nuplan.planning.simulation.trajectory.action_trajectory import ActionTrajectory from pufferlib.ocean.cosim import nuplan_bridge as nb from pufferlib.ocean.cosim.arch import checkpoint_config_path from pufferlib.ocean.cosim.goals import RouteGoalWindow, route_goals_from_xy +from pufferlib.ocean.cosim.light_timing import YellowCountdown +from pufferlib.ocean.cosim.nuplan_bicycle import ( + DEVKIT_ACCEL_TIME_CONSTANT_S, + DEVKIT_STEERING_TIME_CONSTANT_S, + BicycleCommandConfig, + bicycle_commands, + bicycle_response, + floor_clamped_intent_accel, +) +from pufferlib.ocean.cosim.nuplan_rollout import forecast_light_states, rollout_step_count, split_rollout_slots from pufferlib.ocean.drive.drive import Drive from pufferlib.ocean.cosim.nuplan_bridge import DEFAULT_ARCH, FAR_AWAY from pufferlib.ocean.drive import binding EGO_OBS_ACCEL_LONG_IDX = 4 # write_ego_obs column order: speed, width, length, steering, accel_long, ... +EGO_OBS_STEERING_IDX = 3 +EGO_DYNAMICS_PUFFERDRIVE = "pufferdrive" # the shadow env integrates the ego, perfect_tracking_controller replays it +EGO_DYNAMICS_NUPLAN_BICYCLE = "nuplan_bicycle" # nuPlan's KinematicBicycleModel is the plant (one_stage_controller) +EGO_DYNAMICS_NUPLAN_TRACKER = "nuplan_tracker" # nuPlan's two_stage_controller tracks the ego of a forecast rollout +EGO_DYNAMICS_MODES = (EGO_DYNAMICS_PUFFERDRIVE, EGO_DYNAMICS_NUPLAN_BICYCLE, EGO_DYNAMICS_NUPLAN_TRACKER) +DUMMY_ROLLOUT_EGO_ACTION = (1.0, 0.0) # no-checkpoint wiring test: full accelerating jerk, no steering +ROLLOUT_POLICY_TYPES = ("VEHICLE", "BICYCLE") # nuPlan types the policy drives in a rollout; pedestrians never yield +BICYCLE_RESPONSE_TOL = 1e-4 # m/s^2 and m/s: nuPlan's ego state must be the bicycle model's response to our command EGO_TRACKING_TOL_M = 0.05 # perfect_tracking_controller hands back our integrated pose up to log-timestamp jitter EGO_TRACKING_TOL_RAD = 0.005 SCENARIO_LENGTH_MARGIN_STEPS = 2 # shadow env must never hit its own truncation while nuPlan still plans @@ -33,6 +54,8 @@ TRAJECTORY_TAIL_SPACING_S = 0.5 # constant-velocity tail past the integrated step; the controller only reads t+dt COMFORT_ACCEL_MARGIN = 0.96 # nuPlan Savitzky-Golay-smooths the reported accel and overshoots a hard cap by ~0.5% ROADBLOCK_GOAL_MARGIN_M = 5.0 # first roadblock goal sits at least this far beyond the goal radius ahead of the ego +ROLLOUT_SLOT_PROBE_SPACING_M = 10.0 # distinct x per policy slot in the build-time slot round trip +LOG_EXCLUDED_STATE = ("_env", "_policy", "_goal_window", "_obs_replay") # per-scenario runtime, never pickled def startup_jerk_action_bounds(accel_cap_mps3: float, brake_cap_mps3: float): @@ -43,21 +66,8 @@ def startup_jerk_action_bounds(accel_cap_mps3: float, brake_cap_mps3: float): return lo, hi -def clean_policy_state_dict(state_dict): - """Strip torch.compile / DDP prefixes. Inlined from pufferlib.pufferl to - avoid importing the training stack (wandb, neptune, the _C kernel, ...) - inside the 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()} - - # per-process caches: one planner per scenario, same checkpoint/bins each time -_state_dict_cache: Dict[str, dict] = {} +_checkpoint_cache: Dict[str, dict] = {} _bin_geometry_cache: Dict[str, dict] = {} @@ -65,6 +75,19 @@ def _angle_diff(a: float, b: float) -> float: return (a - b + math.pi) % (2.0 * math.pi) - math.pi +class BicycleStep(NamedTuple): + """One nuplan_bicycle planning call: the plant state it saw, the command it sent and the jerk integrator's + post-step state (the accel/steering the policy observes and the next call restamps).""" + + time_us: int + plant_accel_mps2: float + plant_speed_mps: float + accel_cmd_mps2: float + intent_accel_mps2: float + intent_steering_rad: float + plant_yaw_rate_radps: float + + def _steering_from(yaw_rate: float, speed: float, wheel_base: float) -> float: if abs(speed) < 0.1: return 0.0 @@ -72,7 +95,8 @@ def _steering_from(yaw_rate: float, speed: float, wheel_base: float) -> float: class PufferDrivePlanner(AbstractPlanner): - """nuPlan planner that delegates observation + ego dynamics to PufferDrive.""" + """nuPlan planner that delegates the observation and, with ego_dynamics=pufferdrive, the ego dynamics to + PufferDrive.""" requires_scenario: bool = True @@ -94,6 +118,21 @@ def __init__( startup_accel_jerk_cap_mps3: float = 0.0, startup_brake_jerk_cap_mps3: float = 0.0, vru_size_pad_m: float = 0.0, + symmetric_lane_widths: bool = False, + ego_dynamics: str = EGO_DYNAMICS_PUFFERDRIVE, + bicycle_min_speed_mps: float = 0.0, + bicycle_undo_low_pass: bool = False, + bicycle_accel_time_constant_s: float = DEVKIT_ACCEL_TIME_CONSTANT_S, + bicycle_steering_time_constant_s: float = DEVKIT_STEERING_TIME_CONSTANT_S, + bicycle_center_accel_cap: bool = False, + bicycle_center_accel_cap_margin: float = 1.0, + bicycle_max_yaw_jerk_radps3: float = 0.0, + bicycle_floor_accel_clamp: bool = True, + accel_cap_mps2: Optional[float] = None, + partner_split_max_length_m: float = 0.0, + rollout_step_seconds: float = 0.3, + rollout_policy_vehicles: int = 0, + rollout_allow_reverse: bool = False, env_overrides: Optional[Dict] = None, obs_html_dir: Optional[str] = None, obs_html_max_steps: int = 800, @@ -129,6 +168,58 @@ def __init__( 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 symmetric_lane_widths: eval hack: the shadow env loads a copy of the city bin whose lane widths + are twice the centerline's distance to the nearer GPKG boundary instead of the converter's + left + right (nuPlan lanes with a parking strip have their baseline 1.7 m from one boundary and + 4.3 m from the other; the sum made the policy believe it had 3 m towards the oncoming lane). + The copy is cached in bin_cache_dir, one rewrite per city. + :param ego_dynamics: "pufferdrive": the shadow env integrates the ego with its jerk model and nuPlan's + perfect_tracking_controller replays the pose. "nuplan_bicycle": nuPlan's KinematicBicycleModel is + the plant, driven through CaRL's one_stage_controller like CaRL: every step the shadow ego takes + nuPlan's pose, speed and yaw rate, keeps its own jerk-integrated accel/steering (the state the + policy was trained to observe; the plant lags behind it), and its post-step accel/steering + become the model's commands. "nuplan_tracker": the planner returns a trajectory and nuPlan's + two_stage_controller (LQR or iLQR tracker + KinematicBicycleModel) drives the ego along it, as in + GIGAFLOW's nuPlan evaluation. Every step the shadow env takes the whole scene from nuPlan (the ego + with the plant's acceleration and steering angle) and is rolled out for horizon_seconds in + rollout_step_seconds steps, the policy driving the ego (and, if asked for, the nearest + rollout_policy_vehicles vehicles); the ego's rolled-out states are the trajectory. + :param rollout_step_seconds: nuplan_tracker: step of the forecast rollout [s] (0.3 = the training step). + :param rollout_policy_vehicles: nuplan_tracker: at most this many vehicles and bicycles, nearest first, are + driven by the policy in the rollout. nuPlan gives no route for them, so their goals follow their lane and + its straightest exits. Farther ones, those off the lane graph and pedestrians keep the velocity observed + at this step (nuPlan replays pedestrians, they do not yield); static objects stay where they are. + 0 (default) = everyone keeps their observed velocity: policy-driven vehicles pull up to the bin's stop + line, nuPlan's stay where they stopped, and both trackers drive the ego into that forecast gap. + :param rollout_allow_reverse: nuplan_tracker: false = an agent the policy would back up stays where it + stood, at rest (what bicycle_min_speed_mps 0 does for nuplan_bicycle; nuPlan's IDM vehicles never + reverse). true = Drive's own floor (MAX_BACKWARD_SPEED). + :param bicycle_min_speed_mps: nuplan_bicycle speed floor [m/s], enforced on the command (the model has + none): 0 = CaRL's no-reverse driving, -2 = Drive's MAX_BACKWARD_SPEED. + :param bicycle_undo_low_pass: nuplan_bicycle: pre-compensate the model's first-order command lag so one + model step lands on the shadow env's targets (CaRL default: off, the model realises dt/(dt+tau)). + :param bicycle_accel_time_constant_s: the motion model's acceleration low-pass time constant [s] + (devkit default 0.2); a mismatch fails the per-step response check. + :param bicycle_steering_time_constant_s: same for steering [s] (devkit default 0.05). + :param bicycle_center_accel_cap: nuplan_bicycle: CaRL's clip_angular_adjustment. nuPlan scores the + acceleration at the vehicle center, which the devkit computes as rear-axle accel + d*yaw_rate^2 + + d*yaw_accel (d = rear axle to center), so the plant's acceleration is reduced while turning to keep + that reading under accel_cap_mps2 x bicycle_center_accel_cap_margin. The policy keeps its own + intent and observations. + :param bicycle_center_accel_cap_margin: factor on accel_cap_mps2 for the center cap (nuPlan smooths the + reported acceleration with a Savitzky-Golay window and overshoots a peaky capped profile). + :param bicycle_max_yaw_jerk_radps3: nuplan_bicycle: bound on the per-step change of the yaw + acceleration [rad/s^3], enforced by slowing the steering command (the derivative sibling of CaRL's + clip_max_abs_yaw_accel): nuPlan's longitudinal jerk is the derivative of the center acceleration, + which carries d*yaw_accel. 0 = off. + :param bicycle_floor_accel_clamp: nuplan_bicycle: clamp the policy's integrated acceleration to zero while + the plant sits on the speed floor (Drive zeroes it when the car stops; on the floor a negative intent + would wind down to the clip at zero speed, a state the policy never saw). Off = run 426 behaviour. + :param partner_split_max_length_m: partners longer than this [m] are shown to the policy as equal boxes + along their heading (training spawns at most 7 m; nuPlan buses reach 12.7 m). 0 = off. + :param accel_cap_mps2: cap on the ego's longitudinal acceleration (the policy's c_acc conditioning = + cap / ACCEL_LONG_MAX, which must stay inside its training range); None = 0.96 x nuPlan's 2.40 + comfort bound. :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). @@ -161,6 +252,58 @@ def __init__( 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._symmetric_lane_widths = bool(symmetric_lane_widths) + if ego_dynamics not in EGO_DYNAMICS_MODES: + raise ValueError(f"ego_dynamics must be one of {EGO_DYNAMICS_MODES}, got {ego_dynamics!r}") + self._ego_dynamics = ego_dynamics + self._rollout_step_seconds = float(rollout_step_seconds) + if not self._rollout_step_seconds > 0: + raise ValueError(f"rollout_step_seconds must be > 0, got {rollout_step_seconds}") + self._rollout_policy_vehicles = int(rollout_policy_vehicles) + if self._rollout_policy_vehicles < 0: + raise ValueError(f"rollout_policy_vehicles must be >= 0, got {rollout_policy_vehicles}") + if ego_dynamics == EGO_DYNAMICS_NUPLAN_TRACKER and self._rollout_policy_vehicles > self._num_agents - 1: + raise ValueError( + f"rollout_policy_vehicles={rollout_policy_vehicles} exceeds the {self._num_agents - 1} partners " + "the shadow env takes from nuPlan (num_agents - 1)" + ) + self._rollout_allow_reverse = bool(rollout_allow_reverse) + self._bicycle_min_speed_mps = float(bicycle_min_speed_mps) + if not binding.MAX_BACKWARD_SPEED <= self._bicycle_min_speed_mps <= 0.0: + raise ValueError( + f"bicycle_min_speed_mps must be within [{binding.MAX_BACKWARD_SPEED}, 0], got {bicycle_min_speed_mps}" + ) + self._bicycle_undo_low_pass = bool(bicycle_undo_low_pass) + self._bicycle_accel_time_constant_s = float(bicycle_accel_time_constant_s) + self._bicycle_steering_time_constant_s = float(bicycle_steering_time_constant_s) + if not (self._bicycle_accel_time_constant_s >= 0 and self._bicycle_steering_time_constant_s >= 0): + raise ValueError("the bicycle model time constants must be >= 0") + self._bicycle_center_accel_cap = bool(bicycle_center_accel_cap) + self._bicycle_center_accel_cap_margin = float(bicycle_center_accel_cap_margin) + if not 0 < self._bicycle_center_accel_cap_margin <= 1: + raise ValueError( + f"bicycle_center_accel_cap_margin must be in (0, 1], got {bicycle_center_accel_cap_margin}" + ) + self._bicycle_max_yaw_jerk_radps3 = float(bicycle_max_yaw_jerk_radps3) + if not self._bicycle_max_yaw_jerk_radps3 >= 0: + raise ValueError(f"bicycle_max_yaw_jerk_radps3 must be >= 0, got {bicycle_max_yaw_jerk_radps3}") + self._bicycle_floor_accel_clamp = bool(bicycle_floor_accel_clamp) + self._partner_split_max_length_m = float(partner_split_max_length_m) + if not self._partner_split_max_length_m >= 0: + raise ValueError(f"partner_split_max_length_m must be >= 0, got {partner_split_max_length_m}") + if accel_cap_mps2 is None: + self._accel_cap_mps2 = COMFORT_ACCEL_MARGIN * NUPLAN_COMFORT_MAX_LON_ACCEL_MPS2 + else: + self._accel_cap_mps2 = float(accel_cap_mps2) + conditioning = self._accel_cap_mps2 / binding.ACCEL_LONG_MAX + if not binding.CONDITIONING_ACC_MIN <= conditioning <= binding.CONDITIONING_ACC_MAX: + raise ValueError( + f"accel_cap_mps2={accel_cap_mps2} is outside the policy's c_acc training range " + f"[{binding.CONDITIONING_ACC_MIN * binding.ACCEL_LONG_MAX:.3f}, " + f"{binding.CONDITIONING_ACC_MAX * binding.ACCEL_LONG_MAX:.3f}] m/s^2" + ) + self._bicycle_cfg: Optional[BicycleCommandConfig] = None # resolved in _build from the scenario dt + self._last_bicycle_step: Optional[BicycleStep] = None 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 {} @@ -181,7 +324,16 @@ def __init__( self._goal_window: Optional[RouteGoalWindow] = None self._last_integrated = None # (bin_x, bin_y, heading) the shadow ego was left at by the last step self._last_light_states = None # traffic-light state array from the most recent _sync - self._static_on_lane: Dict[str, bool] = {} # static object track_token -> stands inside a lane polygon + self._yellow_countdown = None # YellowCountdown, built with the env + self._static_on_lane: Dict[str, bool] = {} # static object track_token -> its box overlaps a lane polygon + self._partner_stopped_s: Dict[str, float] = {} # track_token -> seconds below the stopped speed threshold + self._planning_step_s: Optional[float] = None # nuPlan's planning interval, resolved in _build + # nuplan_tracker rollout state, refreshed by every _sync + self._rollout_steps = 0 # resolved in _build from the horizon and the rollout step + self._rollout_policy_rows = np.zeros(1, np.int64) # env rows the policy drives: the ego and vehicles with goals + self._forecast_objects = None # (idx, x, y, z, heading, vx, vy) of the static-slot objects that move + self._last_yellow_remaining_s = None # per traffic element, 0 unless yellow + self._ego_stopped_s = 0.0 # seconds the plant's speed has been below the stopped threshold # --- AbstractPlanner interface ------------------------------------------- def initialize(self, initialization: PlannerInitialization) -> None: @@ -203,6 +355,13 @@ def generate_planner_report(self, clear_stats: bool = True): self._obs_replay = None return report + def __getstate__(self): + # nuPlan pickles the planner into every simulation log: 30 MB of policy weights per scenario without this + state = dict(self.__dict__) + for name in LOG_EXCLUDED_STATE: + state[name] = None + return state + # --- whole-city map-only bins ------------------------------------------- def _find_city_bin(self) -> Path: """City bin for this scenario's map (the map id the devkit read from @@ -219,16 +378,45 @@ def _gpkg_layer(self, layer_name: str): except AttributeError: return map_api._maps_db.load_vector_layer(map_api.map_name, layer_name) + @staticmethod + def _bin_fingerprint(bin_path: Path) -> str: + """Size and mtime of a bin; files derived from it carry this, so a regenerated bin never reuses them.""" + stat = bin_path.stat() + return f"{stat.st_size:x}-{stat.st_mtime_ns:x}" + + def _symmetric_width_bin(self, city_bin: Path, tf: nb.NuPlanTransform) -> Path: + """Cached copy of the city bin with every lane width rewritten from the GPKG boundaries.""" + cached = self._bin_cache_dir / f"{city_bin.stem}.{self._bin_fingerprint(city_bin)}.symw.bin" + if cached.exists(): + return cached + from nuplan.common.actor_state.state_representation import Point2D + from nuplan.common.maps.maps_datatypes import SemanticMapLayer + + pts = nb.read_bin_lane_points(city_bin) + lo, hi = pts.min(axis=0), pts.max(axis=0) + cx, cy = tf.bin_to_loc(*((lo + hi) / 2.0)) + radius = float(np.hypot(*(hi - lo)) / 2.0) + self._map_radius + layers = [SemanticMapLayer.LANE, SemanticMapLayer.LANE_CONNECTOR] + objs = self._initialization.map_api.get_proximal_map_objects(Point2D(cx, cy), radius, layers) + patched, unmatched = nb.patch_bin_lane_widths(city_bin, cached, tf, objs[layers[0]] + objs[layers[1]]) + print( + f"[pufferdrive_planner] {city_bin.name}: symmetric lane widths for {patched} lanes " + f"({unmatched} unmatched keep theirs) -> {cached}" + ) + return cached + def _city_bin_origin(self, bin_path: Path, stop_line_centers: np.ndarray) -> nb.NuPlanTransform: """Recover a city bin's frame origin (it stores no centroid): global translation vote on stop-line landmarks, then translation-only ICP of the bin's lane centerlines onto the GPKG baseline paths. Cached in a sidecar json — one registration per city, ever.""" sidecar = bin_path.with_suffix(".origin.json") + fingerprint = self._bin_fingerprint(bin_path) cache = sidecar if sidecar.exists() else self._bin_cache_dir / sidecar.name - if cache.exists(): - origin = json.loads(cache.read_text())["origin"] - return nb.NuPlanTransform(origin[0], origin[1]) + cached = json.loads(cache.read_text()) if cache.exists() else None + # sidecars written before the fingerprint existed carry no "source" and stay trusted + if cached is not None and cached.get("source", fingerprint) == fingerprint: + return nb.NuPlanTransform(cached["origin"][0], cached["origin"][1]) if not len(stop_line_centers): raise RuntimeError(f"{bin_path}: no traffic elements to register the origin with") @@ -243,7 +431,7 @@ def _city_bin_origin(self, bin_path: Path, stop_line_centers: np.ndarray) -> nb. f"{bin_path}: lane-geometry registration residual {resid:.2f} m — " "bin does not match this city's GPKG map, refusing to guess the origin." ) - payload = json.dumps({"origin": [float(t[0]), float(t[1])], "residual_m": resid}) + payload = json.dumps({"origin": [float(t[0]), float(t[1])], "residual_m": resid, "source": fingerprint}) for target in (sidecar, self._bin_cache_dir / sidecar.name): try: target.parent.mkdir(parents=True, exist_ok=True) @@ -269,7 +457,8 @@ def _resolve_map_bin(self): if geo["origin"] is not None else self._city_bin_origin(city_bin, geo["stop_line_centers"]) ) - return city_bin, tf, geo + env_bin = self._symmetric_width_bin(city_bin, tf) if self._symmetric_lane_widths else city_bin + return env_bin, tf, geo def _match_traffic_lights(self, geo: dict, ex: float, ey: float) -> Dict[str, int]: """Signalized lane connectors within map_radius -> bin traffic elements, @@ -299,27 +488,36 @@ def _resolve_arch(self, cfg: Dict, bin_path: Path) -> Dict: if self._scenario is None: raise RuntimeError("PufferDrivePlanner needs the scenario injected (requires_scenario)") + policy_agent_count = 1 + step_seconds = float(self._scenario.database_interval) + scenario_steps = int(self._scenario.get_number_of_iterations()) + SCENARIO_LENGTH_MARGIN_STEPS + rollout_overrides = {} + if self._ego_dynamics == EGO_DYNAMICS_NUPLAN_TRACKER: + # the rollout also drives the nearest vehicles, at its own step, for the whole horizon every iteration + policy_agent_count += self._rollout_policy_vehicles + step_seconds = self._rollout_step_seconds + scenario_steps *= rollout_step_count(self._horizon_seconds, step_seconds) + rollout_overrides = {"max_agents_per_env": policy_agent_count} return shadow_env_kwargs( cfg, defaults=DEFAULT_ARCH, overrides={ **self._env_overrides, + **rollout_overrides, "map_dir": str(bin_path), "num_maps": 1, # One policy agent (the ego); every other slot is a static partner streamed from nuPlan. - "num_agents": 1, - "min_agents_per_env": 1, + "num_agents": policy_agent_count, + "min_agents_per_env": policy_agent_count, "cosim_partner_slots": self._num_agents - 1, "goal_source": "external", # Training semantics: a window's final goal clears only below goal_speed (3 m/s) "goal_reach_requires_speed": True, - # C_acc conditioning: cap the jerk model's positive accel at nuPlan's comfort bound instead of 2.5 - "conditioning_accel_scale": COMFORT_ACCEL_MARGIN - * NUPLAN_COMFORT_MAX_LON_ACCEL_MPS2 - / binding.ACCEL_LONG_MAX, + # C_acc conditioning: cap the jerk model's positive accel at accel_cap_mps2 instead of 2.5 + "conditioning_accel_scale": self._accel_cap_mps2 / binding.ACCEL_LONG_MAX, # lockstep with nuPlan's planning iterations, not the training dt (mimolette: 0.3) - "dt": float(self._scenario.database_interval), - "scenario_length": int(self._scenario.get_number_of_iterations()) + SCENARIO_LENGTH_MARGIN_STEPS, + "dt": step_seconds, + "scenario_length": scenario_steps, "resample_frequency": 0, # External sim owns the episode: a training-config # termination_mode=1 would c_reset() the pool (ego included) @@ -345,6 +543,7 @@ def _build(self, ego_state: EgoState) -> None: import torch import yaml + import pufferlib.ensemble import pufferlib.ocean.torch as drive_torch init = self._initialization @@ -359,6 +558,15 @@ def _build(self, ego_state: EgoState) -> None: self._arch = self._resolve_arch(cfg, bin_path) self._env = Drive(**self._arch) self._env.reset() + self._planning_step_s = float(self._scenario.database_interval) + if self._ego_dynamics == EGO_DYNAMICS_NUPLAN_TRACKER: + import pufferlib.spaces + + if not isinstance(self._env.single_action_space, pufferlib.spaces.Box): + raise RuntimeError(f"ego_dynamics={EGO_DYNAMICS_NUPLAN_TRACKER} needs a continuous (jerk) action space") + self._rollout_steps = rollout_step_count(self._horizon_seconds, self._rollout_step_seconds) + # nuPlan reports light status only: a yellow counts down from the shortest trained yellow + self._yellow_countdown = YellowCountdown(self._env.traffic_light_yellow_min_s, self._env.dt) torch.set_num_threads(1) # one scenario per worker process; BLAS threads only fight the other workers if max(self._startup_accel_jerk_cap_mps3, self._startup_brake_jerk_cap_mps3) > 0: import pufferlib.spaces @@ -370,6 +578,22 @@ def _build(self, ego_state: EgoState) -> None: self._startup_accel_jerk_cap_mps3, self._startup_brake_jerk_cap_mps3 ) + if self._ego_dynamics == EGO_DYNAMICS_NUPLAN_BICYCLE: + vehicle = ego_state.car_footprint.vehicle_parameters + self._bicycle_cfg = BicycleCommandConfig( + dt_s=float(self._arch["dt"]), + min_speed_mps=self._bicycle_min_speed_mps, + undo_low_pass=self._bicycle_undo_low_pass, + accel_time_constant_s=self._bicycle_accel_time_constant_s, + steering_time_constant_s=self._bicycle_steering_time_constant_s, + center_accel_cap_mps2=self._accel_cap_mps2 * self._bicycle_center_accel_cap_margin + if self._bicycle_center_accel_cap + else 0.0, + max_yaw_jerk_radps3=self._bicycle_max_yaw_jerk_radps3, + wheel_base_m=float(vehicle.wheel_base), + rear_axle_to_center_m=float(vehicle.rear_axle_to_center), + ) + # Round-trip probe: slot 0 must be the ego we set. It is silently wrong when # gigaflow spawning failed (e.g. a city bin without lane connectivity). probe_x, probe_y = self._transform.loc_to_bin(ex, ey) @@ -393,21 +617,23 @@ def _build(self, ego_state: EgoState) -> None: "Gigaflow spawning likely failed on this bin -- check for '[GIGAFLOW WARNING]'/'[ERROR]' in the " "worker stdout; a map-only city bin without lane connectivity (exit_lanes) is the known cause." ) + if self._ego_dynamics == EGO_DYNAMICS_NUPLAN_TRACKER: + self._check_rollout_policy_slots(probe_x, probe_y) if self._dummy: self._policy = None else: # pre-3.0 checkpoints keep action_type only in the env section cfg["policy"].setdefault("action_type", cfg["env"]["action_type"]) - policy = getattr(drive_torch, cfg.get("policy_name", "Drive"))(self._env, **cfg["policy"]) - sd = _state_dict_cache.get(str(self._checkpoint_path)) - if sd is None: - sd = clean_policy_state_dict( - torch.load(self._checkpoint_path, map_location=self._device, weights_only=False) - ) - _state_dict_cache[str(self._checkpoint_path)] = sd - policy.load_state_dict(sd) - self._policy = policy.to(self._device).eval() + policy_kwargs = dict(cfg["policy"]) + if torch.device(self._device).type == "cpu": + policy_kwargs["fused_slot_encoder"] = False # CUDA-only triton kernel; CPU uses the reference path + policy = getattr(drive_torch, cfg.get("policy_name", "Drive"))(self._env, **policy_kwargs) + checkpoint = _checkpoint_cache.get(str(self._checkpoint_path)) + if checkpoint is None: + checkpoint = torch.load(self._checkpoint_path, map_location=self._device, weights_only=False) + _checkpoint_cache[str(self._checkpoint_path)] = checkpoint + self._policy = pufferlib.ensemble.load_policy_weights(policy, checkpoint).to(self._device).eval() # The lane-graph route through the CaRL-corrected route roadblocks always decides the shared # stop-line elements and seeds the route goals. @@ -466,55 +692,136 @@ def _build(self, ego_state: EgoState) -> None: f"{self._num_traffic} goals={len(goals)} goal_source={self._goal_source}" ) + def _check_rollout_policy_slots(self, probe_x: float, probe_y: float) -> None: + """nuplan_tracker: observation row k must be env agent k for every policy slot. A failed gigaflow spawn + drops its agent from the active list and shifts every later row onto another vehicle.""" + slots = np.arange(1 + self._rollout_policy_vehicles, dtype=np.int32) + slot_x = (probe_x + ROLLOUT_SLOT_PROBE_SPACING_M * slots).astype(np.float32) + slot_y = np.full(len(slots), probe_y, np.float32) + zeros = np.zeros(len(slots), np.float32) + self._env.set_agent_states(slots, slot_x, slot_y, zeros, zeros, zeros, zeros, zeros, zeros) + readback = self._env.get_global_agent_state() + if not np.allclose(readback["x"], slot_x, atol=0.5): + raise RuntimeError( + f"shadow env policy slots are not agents 0..{len(slots) - 1} (set x {slot_x.tolist()}, read " + f"{readback['x'].tolist()}): a gigaflow spawn failed, see '[GIGAFLOW WARNING]' in the worker stdout" + ) + # --- world sync ----------------------------------------------------------- + def _advance_partner_stopped_s(self, tracked_objects) -> None: + """One observation step of the stopped time per nuPlan track; a track missing from the step starts over.""" + step_s = self._planning_step_s + stopped_s = {} + for obj in tracked_objects: + velocity = getattr(obj, "velocity", None) + speed_mps = math.hypot(velocity.x, velocity.y) if velocity is not None else 0.0 + is_stopped = speed_mps < binding.AGENT_STOPPED_SPEED_THRESHOLD + previous_s = self._partner_stopped_s.get(obj.track_token, 0.0) + stopped_s[obj.track_token] = previous_s + step_s if is_stopped else 0.0 + self._partner_stopped_s = stopped_s + def _sync(self, ego_state: EgoState, detections: DetectionsTracks, traffic_light_data) -> None: tf = self._transform env = self._env - # ego (slot 0): synced from nuPlan telemetry once, at the scenario's initial state. Afterwards the - # shadow env owns the ego: nuPlan's perfect_tracking_controller hands back exactly the pose we - # integrated, and re-deriving accel_long/steering from telemetry every tick re-injects - # finite-difference accel after every stop and drives the ego into reverse. + # ego (slot 0). pufferdrive dynamics: synced from nuPlan telemetry once, at the scenario's initial state; + # afterwards the shadow env owns the ego (nuPlan's perfect_tracking_controller hands back exactly the pose + # we integrated, and re-deriving accel_long/steering from telemetry every tick re-injects finite-difference + # accel after every stop and drives the ego into reverse). nuplan_bicycle dynamics: nuPlan's bicycle model + # is the plant: pose, speed and yaw rate are restamped every step, accel/steering stay the jerk integrator's + # own state (the lagged plant never lands exactly on zero, which would make the zero-crossing snaps stick). + # nuplan_tracker dynamics: every rollout starts from the plant's full state, accel and steering included. h = float(ego_state.center.heading) dcs = ego_state.dynamic_car_state - v_local = dcs.center_velocity_2d - vx_g = v_local.x * math.cos(h) - v_local.y * math.sin(h) - vy_g = v_local.x * math.sin(h) + v_local.y * math.cos(h) ex, ey = tf.loc_to_bin(ego_state.center.x, ego_state.center.y) - if self._last_integrated is not None: + if self._ego_dynamics == EGO_DYNAMICS_NUPLAN_TRACKER: + speed = float(dcs.rear_axle_velocity_2d.x) + is_stopped = abs(speed) < binding.AGENT_STOPPED_SPEED_THRESHOLD + self._ego_stopped_s = self._ego_stopped_s + self._planning_step_s if is_stopped else 0.0 + self._stamp_ego( + ex, + ey, + h, + speed * math.cos(h), + speed * math.sin(h), + float(dcs.angular_velocity), + float(dcs.rear_axle_acceleration_2d.x), + float(ego_state.tire_steering_angle), + self._ego_stopped_s, + ) + elif self._ego_dynamics == EGO_DYNAMICS_NUPLAN_BICYCLE: + self._check_bicycle_response(ego_state) + speed = float(dcs.rear_axle_velocity_2d.x) + last = self._last_bicycle_step + intent_accel = float(dcs.rear_axle_acceleration_2d.x) if last is None else last.intent_accel_mps2 + if self._bicycle_floor_accel_clamp: + intent_accel = floor_clamped_intent_accel(self._bicycle_cfg, speed, intent_accel) + self._stamp_ego( + ex, + ey, + h, + speed * math.cos(h), + speed * math.sin(h), + float(dcs.angular_velocity), + intent_accel, + float(ego_state.tire_steering_angle) if last is None else last.intent_steering_rad, + ) + elif self._last_integrated is not None: self._check_ego_tracking(ex, ey, h) else: - env.set_agent_states( - np.array([0], np.int32), - np.array([ex], np.float32), - np.array([ey], np.float32), - np.array([0.0], np.float32), - np.array([h], np.float32), - np.array([vx_g], np.float32), - np.array([vy_g], np.float32), - np.array([float(dcs.angular_velocity)], np.float32), - np.array([float(dcs.center_acceleration_2d.x)], np.float32), + v_local = dcs.center_velocity_2d + vx_g = v_local.x * math.cos(h) - v_local.y * math.sin(h) + vy_g = v_local.x * math.sin(h) + v_local.y * math.cos(h) + self._stamp_ego( + ex, ey, h, vx_g, vy_g, float(dcs.angular_velocity), float(dcs.center_acceleration_2d.x), None ) # background (slots 1..): streamed from nuPlan, never simulated here. nuPlan's TrackedObject # carries no acceleration/angular-velocity (perception detections, not ego telemetry), but # write_partner_obs never reads those fields for non-ego agents, so zero is exact, not a stopgap. + # stopped time follows the nuPlan track, not the shadow slot (slots shift whenever a track appears or ends) + self._advance_partner_stopped_s(detections.tracked_objects) objs = nb.partner_tracked_objects( detections.tracked_objects, self._initialization.map_api, self._static_on_lane ) - if len(objs) > self._num_agents - 1: - # keep the nearest: devkit ordering is by type/token, so blind - # truncation could drop a close vehicle while keeping far ones - ox, oy = ego_state.center.x, ego_state.center.y - objs.sort(key=lambda o: (o.center.x - ox) ** 2 + (o.center.y - oy) ** 2) - objs = objs[: self._num_agents - 1] + rollout = self._ego_dynamics == EGO_DYNAMICS_NUPLAN_TRACKER + partner_count = 0 + idx = np.zeros(0, np.int32) + if rollout: + self._rollout_policy_rows = np.zeros(1, np.int64) + self._forecast_objects = None if objs: - idx, x, y, z, hh, vx, vy, tp, ln, wd = nb.tracked_objects_to_arrays(objs, tf) + _, x, y, z, hh, vx, vy, tp, ln, wd = nb.tracked_objects_to_arrays(objs, tf) + stopped_s = np.array([self._partner_stopped_s[obj.track_token] for obj in objs], np.float32) + # the nuPlan type decides who the policy can drive in a rollout: the bridge types static clutter as vehicle + policy_eligible = np.array([obj.tracked_object_type.name in ROLLOUT_POLICY_TYPES for obj in objs], bool) 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) + if self._partner_split_max_length_m > 0: + max_length_m = self._partner_split_max_length_m + piece_counts = np.maximum(1, np.ceil(ln.astype(np.float64) / max_length_m)).astype(np.int64) + stopped_s = np.repeat(stopped_s, piece_counts) + policy_eligible = np.repeat(policy_eligible & (piece_counts == 1), piece_counts) + x, y, z, hh, vx, vy, tp, ln, wd = nb.split_long_partners(x, y, z, hh, vx, vy, tp, ln, wd, max_length_m) + assert len(stopped_s) == len(x) + # nearest first only when the slots run out: devkit ordering is by type/token, so blind truncation + # could drop a close vehicle while keeping far ones + keep = nb.nearest_partner_slots(x, y, ex, ey, self._num_agents - 1) + partner_count = len(keep) + idx = np.arange(1, 1 + partner_count, dtype=np.int32) + if rollout: + idx = self._rollout_partner_slots(x[keep], y[keep], policy_eligible[keep], ex, ey) + zeros = np.zeros(partner_count, np.float32) + env.set_agent_states( + idx, x[keep], y[keep], z[keep], hh[keep], vx[keep], vy[keep], zeros, zeros, stopped_s[keep] + ) + env.set_agent_sizes(idx, ln[keep], wd[keep]) + if rollout: + self._arm_rollout_partners(idx, x[keep], y[keep], z[keep], hh[keep], vx[keep], vy[keep]) + surplus = np.arange(1 + partner_count, self._num_agents, dtype=np.int32) + if rollout: + partner_slots = np.arange(1, self._rollout_policy_vehicles + self._num_agents, dtype=np.int32) + surplus = np.setdiff1d(partner_slots, idx) if len(surplus): far = np.full(len(surplus), FAR_AWAY, np.float32) zero = np.zeros_like(far) @@ -525,10 +832,115 @@ def _sync(self, ego_state: EgoState, detections: DetectionsTracks, traffic_light self._last_light_states = nb.traffic_light_states( traffic_light_data, self._connector_map, self._num_traffic, self._route_connector_ids ) - env.set_traffic_light_states(self._last_light_states) + yellow_remaining_s = None + if env.obs_traffic_light_yellow_remaining or rollout: + yellow_elements = [int(j) for j in np.flatnonzero(self._last_light_states == nb.LIGHT_STATE_YELLOW)] + countdown_s = self._yellow_countdown.tick(yellow_elements, ego_state.time_point.time_s) + yellow_remaining_s = np.zeros(self._num_traffic, np.float32) + for j, remaining_s in countdown_s.items(): + yellow_remaining_s[j] = remaining_s + self._last_yellow_remaining_s = yellow_remaining_s + env.set_traffic_light_states( + self._last_light_states, yellow_remaining_s if env.obs_traffic_light_yellow_remaining else None + ) self._goal_window.sync(ex, ey, h) + def _rollout_partner_slots(self, x, y, policy_eligible, ego_x: float, ego_y: float) -> np.ndarray: + """nuplan_tracker: env slot per partner. The nearest whole vehicles and bicycles take the policy slots 1..K, + every other partner a static slot after them.""" + distance_sq = (x.astype(np.float64) - ego_x) ** 2 + (y.astype(np.float64) - ego_y) ** 2 + policy_rows, static_rows = split_rollout_slots(distance_sq, policy_eligible, self._rollout_policy_vehicles) + idx = np.empty(len(x), np.int32) + idx[policy_rows] = 1 + np.arange(len(policy_rows)) + idx[static_rows] = 1 + self._rollout_policy_vehicles + np.arange(len(static_rows)) + return idx + + def _arm_rollout_partners(self, idx, x, y, z, heading, vx, vy) -> None: + """nuplan_tracker: vehicles in policy slots get goals along their lane and are driven by the policy in the + rollout; one without a lane to follow keeps coasting. The moving partners in static slots are remembered + so the rollout can carry them on at their observed velocity.""" + in_policy_slot = idx <= self._rollout_policy_vehicles + policy_idx = idx[in_policy_slot] + goal_counts = self._env.set_agent_lane_goals(policy_idx, self._goal_spacing) + self._rollout_policy_rows = np.concatenate([[0], policy_idx[goal_counts > 0]]).astype(np.int64) + moving = ~in_policy_slot & (np.hypot(vx, vy) > 0.0) + if moving.any(): + self._forecast_objects = ( + idx[moving], + x[moving], + y[moving], + z[moving], + heading[moving], + vx[moving], + vy[moving], + ) + + def _hold_reversing_agents(self, policy_rows, obs, poses_before, poses_after) -> bool: + """nuplan_tracker without reversing: a policy-driven agent whose step ended at a negative speed is put back + where it stood, at rest (Drive stops exactly at zero before it reverses, so the pose before the step is + the standstill). poses_after is patched in place; returns True when an agent was held.""" + held = policy_rows[obs[policy_rows, 0] < 0.0] + if len(held) == 0: + return False + for key in ("x", "y", "z", "heading"): + poses_after[key][held] = poses_before[key][held] + zeros = np.zeros(len(held), np.float32) + self._env.set_agent_states( + held.astype(np.int32), + poses_after["x"][held], + poses_after["y"][held], + poses_after["z"][held], + poses_after["heading"][held], + zeros, + zeros, + zeros, + zeros, + steering=obs[held, EGO_OBS_STEERING_IDX] * binding.STEERING_LIMIT, + ) + return True + + def _stamp_forecast(self, elapsed_s: float) -> None: + """nuplan_tracker: the scene elapsed_s after the last sync, as far as the policy does not drive it: static-slot + objects at their observed velocity, lights holding their state (a yellow runs out into red).""" + if self._forecast_objects is not None: + idx, x, y, z, heading, vx, vy = self._forecast_objects + zeros = np.zeros(len(idx), np.float32) + self._env.set_agent_states(idx, x + vx * elapsed_s, y + vy * elapsed_s, z, heading, vx, vy, zeros, zeros) + if self._num_traffic: + states, yellow_remaining_s = forecast_light_states( + self._last_light_states, self._last_yellow_remaining_s, elapsed_s + ) + self._env.set_traffic_light_states( + states, yellow_remaining_s if self._env.obs_traffic_light_yellow_remaining else None + ) + + def _stamp_ego( + self, + bin_x: float, + bin_y: float, + heading: float, + vx: float, + vy: float, + yaw_rate: float, + accel_long: float, + steering: Optional[float], + seconds_stopped: Optional[float] = None, + ) -> None: + self._env.set_agent_states( + np.array([0], np.int32), + np.array([bin_x], np.float32), + np.array([bin_y], np.float32), + np.array([0.0], np.float32), + np.array([heading], np.float32), + np.array([vx], np.float32), + np.array([vy], np.float32), + np.array([yaw_rate], np.float32), + np.array([accel_long], np.float32), + seconds_stopped=None if seconds_stopped is None else np.array([seconds_stopped], np.float32), + steering=None if steering is None else np.array([steering], np.float32), + ) + def _check_ego_tracking(self, bin_x: float, bin_y: float, heading: float) -> None: lx, ly, lh = self._last_integrated if ( @@ -541,71 +953,112 @@ def _check_ego_tracking(self, bin_x: float, bin_y: float, heading: float) -> Non f"shadow pose ({lx:.3f}, {ly:.3f}, {lh:.4f}); the planner assumes perfect_tracking_controller" ) + def _check_bicycle_response(self, ego_state: EgoState) -> None: + last = self._last_bicycle_step + if last is None: + return + elapsed_s = (ego_state.time_point.time_us - last.time_us) * 1e-6 + if elapsed_s <= 0: + raise RuntimeError(f"nuPlan ego time did not advance since the last command ({elapsed_s:.6f} s)") + expected_accel, expected_speed = bicycle_response( + self._bicycle_cfg, last.plant_accel_mps2, last.accel_cmd_mps2, last.plant_speed_mps, elapsed_s + ) + dcs = ego_state.dynamic_car_state + accel = float(dcs.rear_axle_acceleration_2d.x) + speed = float(dcs.rear_axle_velocity_2d.x) + if abs(accel - expected_accel) > BICYCLE_RESPONSE_TOL or abs(speed - expected_speed) > BICYCLE_RESPONSE_TOL: + raise RuntimeError( + f"nuPlan ego accel/speed ({accel:.4f}, {speed:.4f}) are not the kinematic bicycle model's response " + f"({expected_accel:.4f}, {expected_speed:.4f}) to the last command; ego_dynamics=nuplan_bicycle needs " + "ego_controller=one_stage_controller and the planner's bicycle time constants must match its model" + ) + # --- planning ------------------------------------------------------------- - def compute_planner_trajectory(self, current_input: PlannerInput) -> AbstractTrajectory: + def _policy_forward(self, obs: np.ndarray): + """One policy call on observation rows -> (env actions [row, action dim], obs tensor, logits, value, + entropy, discrete action classes).""" import torch import pufferlib.pytorch + import pufferlib.spaces + + # 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) + if self._deterministic: + action_selection = ( + pufferlib.pytorch.ACTION_SELECT_MEAN + if env_continuous and not self._policy.is_continuous + else pufferlib.pytorch.ACTION_SELECT_MODE + ) + else: + action_selection = pufferlib.pytorch.ACTION_SELECT_SAMPLE + obs_tensor = torch.as_tensor(obs).to(self._device) + with torch.no_grad(): + logits, value = self._policy.forward_eval(obs_tensor) + 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 + act = env_action.cpu().numpy().reshape(len(obs), -1) + if not env_continuous: + act = act.astype(np.int32) + return act, obs_tensor, logits, value, entropy, action + + def _capture_obs_replay(self, obs, act, obs_tensor, logits, value, entropy, action) -> None: + # sample_logits' discrete `action` is the class taken (mode/sample) or the argmax under mean selection + action_index = action.cpu().numpy().reshape(-1) if not self._policy.is_continuous else None + self._obs_replay.capture( + obs, act, self._obs_replay.policy_outputs(obs_tensor, logits, value, entropy), action_index + ) + def _write_obs_replay_at_scenario_end(self, iteration_index: int) -> None: + if self._obs_replay is None or iteration_index < self._scenario.get_number_of_iterations() - 2: + return + out = self._obs_replay.write(render_html=self._obs_html_render) + print(f"[pufferdrive_planner] wrote obs replay ({len(self._obs_replay)} steps) -> {out}") + self._obs_replay = None + + def compute_planner_trajectory(self, current_input: PlannerInput) -> AbstractTrajectory: ego_state, detections = current_input.history.current_state if self._env is None: self._build(ego_state) + # nuPlan's history buffer: its last entry is the current observation, which _sync counts + for past_observation in current_input.history.observations[:-1]: + self._advance_partner_stopped_s(past_observation.tracked_objects) self._sync(ego_state, detections, current_input.traffic_light_data) - if self._env.tick >= self._arch["scenario_length"] - 1: + rollout = self._ego_dynamics == EGO_DYNAMICS_NUPLAN_TRACKER + steps_this_call = self._rollout_steps if rollout else 1 + if self._env.tick + steps_this_call > self._arch["scenario_length"] - 1: raise RuntimeError( f"shadow env tick {self._env.tick} reached scenario_length {self._arch['scenario_length']}; " "nuPlan planned more iterations than the scenario declared" ) + if rollout: + trajectory = self._rollout_trajectory(ego_state, current_input.iteration.index) + self._write_obs_replay_at_scenario_end(current_input.iteration.index) + return trajectory obs = np.asarray(self._env.recompute_observations()) if self._policy is None: # dummy wiring test: constant forward jerk act = np.full((self._env.num_agents, 1), 3 * 3 + 1, dtype=np.int32) else: - import pufferlib.spaces - - # 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) - if self._deterministic: - action_selection = ( - pufferlib.pytorch.ACTION_SELECT_MEAN - if env_continuous and not self._policy.is_continuous - else pufferlib.pytorch.ACTION_SELECT_MODE - ) - else: - action_selection = pufferlib.pytorch.ACTION_SELECT_SAMPLE - obs_tensor = torch.as_tensor(obs).to(self._device) - with torch.no_grad(): - logits, value = self._policy.forward_eval(obs_tensor) - 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 - act = env_action.cpu().numpy().reshape(self._env.num_agents, -1) - if not env_continuous: - act = act.astype(np.int32) + act, obs_tensor, logits, value, entropy, action = self._policy_forward(obs) if self._startup_action_bounds is not None and self._env.tick < self._startup_jerk_cap_steps: np.clip(act[:, 0], *self._startup_action_bounds, out=act[:, 0]) if self._obs_replay is not None: - # sample_logits' discrete `action` is the class taken (mode/sample) or the argmax under mean selection - action_index = action.cpu().numpy().reshape(-1) if not self._policy.is_continuous else None - self._obs_replay.capture( - obs, act, self._obs_replay.policy_outputs(obs_tensor, logits, value, entropy), action_index - ) + self._capture_obs_replay(obs, act, obs_tensor, logits, value, entropy, action) self._env.step(act) # integrates the ego one dt; background is re-synced next call - if ( - self._obs_replay is not None - and current_input.iteration.index >= self._scenario.get_number_of_iterations() - 2 - ): - out = self._obs_replay.write(render_html=self._obs_html_render) - print(f"[pufferdrive_planner] wrote obs replay ({len(self._obs_replay)} steps) -> {out}") - self._obs_replay = None + self._write_obs_replay_at_scenario_end(current_input.iteration.index) + + if self._ego_dynamics == EGO_DYNAMICS_NUPLAN_BICYCLE: + return self._bicycle_action_trajectory(ego_state, act[0]) # integrated ego -> nuPlan EgoState at t + dt dt = float(self._arch["dt"]) @@ -650,3 +1103,110 @@ def _state_at(i: int) -> EgoState: if pose_indices[-1] != n_poses: pose_indices.append(n_poses) return InterpolatedTrajectory([_state_at(i) for i in pose_indices]) + + def _rollout_trajectory(self, ego_state: EgoState, iteration_index: int) -> InterpolatedTrajectory: + """nuplan_tracker: roll the synced scene out for the horizon; the ego's states are the trajectory. The policy + drives the rows of _rollout_policy_rows, every other policy slot coasts on a zero jerk action. A goal the + ego consumes in the first step stays consumed (the other modes consume on their one step too); what the + later steps consume is undone.""" + env = self._env + step_s = self._rollout_step_seconds + policy_rows = self._rollout_policy_rows + scenario_elapsed_s = iteration_index * self._planning_step_s + actions = np.zeros((env.num_agents, env.single_action_space.shape[0]), np.float32) + obs = np.asarray(env.recompute_observations()) + vehicle = ego_state.car_footprint.vehicle_parameters + start_time_us = ego_state.time_point.time_us + previous_heading = float(ego_state.center.heading) + trajectory = [ego_state] + goal_snapshot = None + poses = env.get_global_agent_state() + for step_idx in range(self._rollout_steps): + actions[:] = 0.0 + if self._policy is None: + actions[0] = DUMMY_ROLLOUT_EGO_ACTION + else: + act, obs_tensor, logits, value, entropy, action = self._policy_forward(obs[policy_rows]) + actions[policy_rows] = act + in_startup_window = scenario_elapsed_s + step_idx * step_s < self._startup_jerk_cap_seconds + if self._startup_action_bounds is not None and in_startup_window: + actions[0, 0] = np.clip(actions[0, 0], *self._startup_action_bounds) + if step_idx == 0 and self._obs_replay is not None: + # the ego is row 0 of the batch; sample_logits squeezes entropy to 0-dim for a one-row batch + row_count = len(policy_rows) + ego_logits = logits[:1] if hasattr(logits, "shape") else type(logits)(head[:1] for head in logits) + self._capture_obs_replay( + obs[:1], + actions[:1].copy(), + obs_tensor[:1], + ego_logits, + value.reshape(row_count, -1)[:1], + entropy.reshape(row_count)[:1], + action.reshape(row_count, -1)[:1], + ) + self._stamp_forecast((step_idx + 1) * step_s) + env.step(actions) + obs = np.asarray(env.observations) + poses_before, poses = poses, env.get_global_agent_state() + observations_stale = False + if not self._rollout_allow_reverse: + observations_stale = self._hold_reversing_agents(policy_rows, obs, poses_before, poses) + bin_x, bin_y, heading = float(poses["x"][0]), float(poses["y"][0]), float(poses["heading"][0]) + if self._goal_window.sync(bin_x, bin_y, heading) or observations_stale: + env.recompute_observations() + if step_idx == 0: + goal_snapshot = self._goal_window.snapshot() + ego_obs = obs[0] + x, y = self._transform.bin_to_loc(bin_x, bin_y) + trajectory.append( + EgoState.build_from_center( + center=StateSE2(x, y, heading), + center_velocity_2d=StateVector2D(float(ego_obs[0]) * env.obs_norm_speed_mps, 0.0), + center_acceleration_2d=StateVector2D( + float(ego_obs[EGO_OBS_ACCEL_LONG_IDX]) * binding.ACCEL_LONG_NORM, 0.0 + ), + tire_steering_angle=float(ego_obs[EGO_OBS_STEERING_IDX]) * binding.STEERING_LIMIT, + time_point=TimePoint(start_time_us + int(round((step_idx + 1) * step_s * 1e6))), + vehicle_parameters=vehicle, + angular_vel=_angle_diff(heading, previous_heading) / step_s, + ) + ) + previous_heading = heading + self._goal_window.restore(goal_snapshot) + return InterpolatedTrajectory(trajectory) + + def _bicycle_action_trajectory(self, ego_state: EgoState, action: np.ndarray) -> ActionTrajectory: + """Shadow-env post-step accel/steering targets -> commands for nuPlan's bicycle model over the next dt.""" + ego_obs = np.asarray(self._env.observations)[0] + intent_accel = float(ego_obs[EGO_OBS_ACCEL_LONG_IDX]) * binding.ACCEL_LONG_NORM + intent_steering = float(ego_obs[EGO_OBS_STEERING_IDX]) * binding.STEERING_LIMIT + shadow_speed = float(ego_obs[0]) * self._env.obs_norm_speed_mps + dcs = ego_state.dynamic_car_state + accel = float(dcs.rear_axle_acceleration_2d.x) + speed = float(dcs.rear_axle_velocity_2d.x) + yaw_rate = float(dcs.angular_velocity) + last = self._last_bicycle_step + accel_cmd, steering_rate_cmd = bicycle_commands( + self._bicycle_cfg, + intent_accel, + intent_steering, + accel, + float(ego_state.tire_steering_angle), + speed, + yaw_rate_radps=yaw_rate, + yaw_rate_prev_radps=None if last is None else last.plant_yaw_rate_radps, + stop_this_step=speed != 0.0 and shadow_speed == 0.0, + ) + self._last_bicycle_step = BicycleStep( + ego_state.time_point.time_us, accel, speed, accel_cmd, intent_accel, intent_steering, yaw_rate + ) + sampling = TrajectorySampling( + num_poses=int(round(self._horizon_seconds / self._bicycle_cfg.dt_s)), interval_length=self._bicycle_cfg.dt_s + ) + return ActionTrajectory( + accel_cmd, + steering_rate_cmd, + ego_state, + raw_action=[float(a) for a in action], + trajectory_sampling=sampling, + ) diff --git a/pufferlib/ocean/cosim/nuplan/run_nuplan_planner.sh b/pufferlib/ocean/cosim/nuplan/run_nuplan_planner.sh index 6656e6bce2..e265be0d80 100755 --- a/pufferlib/ocean/cosim/nuplan/run_nuplan_planner.sh +++ b/pufferlib/ocean/cosim/nuplan/run_nuplan_planner.sh @@ -1,8 +1,8 @@ #!/bin/bash # Run pufferlib/ocean/cosim/nuplan/planner.py under nuPlan's unmodified # closed-loop run_simulation.py: BOTH agent-reactivity challenges (nonreactive -# log-replay + reactive IDM), with both debug videos enabled. See README.md -# "How to run" and planner.py's module docstring. +# log-replay + reactive IDM). See README.md "How to run" and planner.py's +# module docstring. # # Usage (from a compute-node allocation's job id -- squeue -u $USER to find it): # CKPT=/path/to/model_puffer_drive_xxxxxx.pt \ @@ -29,6 +29,10 @@ # session -- only worker=sequential has been exercised end-to-end # (one scenario). If ray_distributed misbehaves here, retry with # WORKER=sequential (much slower, but proven) while investigating. +# CALLBACKS nuPlan simulation callbacks (default "[simulation_log_callback]": one +# msgpack.xz log per scenario). Add carl_visualization_callback for nuPlan's +# ground-truth .avi of every scenario; the tidy keeps the selected ones as +# details/videos/.mp4. # LIMIT_TOTAL_SCENARIOS caps scenario_filter.limit_total_scenarios (unset = # no cap, SPLIT's own scenario count applies). Handy for a quick # sanity check before committing to the full SPLIT. @@ -40,18 +44,48 @@ # 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. +# SYMMETRIC_LANE_WIDTHS=true|false the shadow env loads a cached copy of the city bin whose lane widths are +# twice the centerline's distance to the nearer GPKG boundary (the converter writes left + right, +# which overstates nuPlan's parking-strip lanes). Unset = planner yaml, false. +# EGO_DYNAMICS=pufferdrive|nuplan_bicycle|nuplan_tracker ego plant (unset = planner yaml, pufferdrive). nuplan_bicycle drives +# nuPlan's KinematicBicycleModel through CaRL's one_stage_controller (the ego_controller override is +# added here). BICYCLE_MIN_SPEED_MPS [m/s, <= 0: 0 = no reverse driving, -2 = Drive's floor] and +# BICYCLE_UNDO_LOW_PASS=true|false (pre-compensate the model's command lag) tune it; unset = planner yaml. +# BICYCLE_CENTER_ACCEL_CAP=true|false: reduce the plant's accel while turning so nuPlan's center accel +# (rear accel + d*yaw_rate^2 + d*yaw_accel) stays under the accel cap (CaRL clip_angular_adjustment); +# BICYCLE_CENTER_ACCEL_CAP_MARGIN (factor on the cap, savgol headroom) and BICYCLE_MAX_YAW_JERK +# [rad/s^3, bound on the per-step change of yaw accel via the steering command; 0 = off] tune comfort. +# BICYCLE_FLOOR_ACCEL_CLAMP=true|false: clamp the policy's integrated accel to 0 while the plant sits on +# the speed floor (unset = planner yaml, true; false reproduces run 426's standstill windup). +# EGO_DYNAMICS=nuplan_tracker is GIGAFLOW's nuPlan evaluation: every step the planner rolls the scene out +# for the trajectory horizon (the policy drives the ego, everyone else keeps its velocity) and nuPlan's +# two_stage_controller tracks the ego's rolled-out trajectory. EGO_TRACKER=lqr|ilqr picks the tracker +# (default lqr = nuPlan's default; ilqr = nuPlan's iLQR without its wall-clock solve cap, so the result +# does not depend on the node's load). ROLLOUT_STEP_SECONDS [s, rollout step, 0.3 = the training step], +# ROLLOUT_POLICY_VEHICLES [vehicles and bicycles the policy also drives in the rollout, nearest first] and +# ROLLOUT_ALLOW_REVERSE=true|false [false: an agent the policy would back up stays at rest] tune it; +# unset = planner yaml. +# EGO_ACCEL_CAP_MPS2 [m/s^2] ego accel cap, any ego dynamics (unset = planner yaml, 0.96 x 2.40). +# PARTNER_SPLIT_MAX_LENGTH_M [m] partners longer than this are shown as equal boxes along their heading +# (training spawns at most 7 m, nuPlan buses reach 12.7 m); unset/0 = off. +# STANDSTILL_JERK_DEADBAND [m/s^3] shadow-env standstill jerk deadband (unset = arch.py, 1.5; 0 = off: with the +# bicycle plant's speed floor the deadband's anti-reverse-creep job is done by the floor). # 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). +# LANE_WIDTH_CLIP_MIN_M / LANE_WIDTH_CLIP_MAX_M [m] clip the lane widths the policy observes to this +# range (training maps span 2.85-4.4 m; nuPlan city bins reach 50 m); unset = off. # 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 # COSIM_OBS_HTML_MAX_SCORE (default 0.9). infractions: only scenarios with an # at-fault collision, drivable-area or driving-direction violation, or no -# progress. Non-selected replays are deleted after the run. 0: off. -# (nuPlan's own ground-truth video is always requested via -# carl_visualization_callback in the callback list below -- one .avi per -# scenario under $GROUP/simulation///visualization/) +# progress. Every .replay.zlib is deleted after rendering (the pages embed them). 0: off. +# TIDY=1 scripts/eval/tidy_nuplan_run.py then flattens $GROUP to overview.csv, scenarios.csv, obs_html/, +# simulation_log/ (the logs of the scenarios COSIM_OBS_HTML selects, infractions when it is 0) and +# details/ (metrics, the aggregator table, hydra config, runner report; the selected scenarios' +# videos as .mp4 when carl_visualization_callback ran). TIDY=0 leaves nuPlan's +# simulation// tree with every log (2300 files in 2000 folders per Val14 run). if [ -z "$INSIDE_PUFFER_CONTAINER" ] && [ -f /share/apps/images/slurm-ib-bind-apptainer.sh ]; then source /share/apps/images/slurm-ib-bind-apptainer.sh export APPTAINERENV_PREPEND_PATH=/share/apps/apptainer/bin:$APPTAINERENV_PREPEND_PATH @@ -88,9 +122,36 @@ EXTRA_ARGS+=("planner.pufferdrive_planner.city_bin_dir=$CITY_BIN_DIR") [ -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 "${VRU_SIZE_PAD_M:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.vru_size_pad_m=$VRU_SIZE_PAD_M") +[ -n "${SYMMETRIC_LANE_WIDTHS:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.symmetric_lane_widths=$SYMMETRIC_LANE_WIDTHS") +if [ "${EGO_DYNAMICS:-}" = "nuplan_bicycle" ]; then + EXTRA_ARGS+=("ego_controller=one_stage_controller" "planner.pufferdrive_planner.ego_dynamics=nuplan_bicycle") +elif [ "${EGO_DYNAMICS:-}" = "nuplan_tracker" ]; then + case "${EGO_TRACKER:-lqr}" in + lqr) TRACKER_CONTROLLER=two_stage_controller ;; + ilqr) TRACKER_CONTROLLER=two_stage_ilqr_controller ;; + *) echo "EGO_TRACKER must be lqr or ilqr, got '${EGO_TRACKER}'"; exit 1 ;; + esac + EXTRA_ARGS+=("ego_controller=$TRACKER_CONTROLLER" "planner.pufferdrive_planner.ego_dynamics=nuplan_tracker") +elif [ -n "${EGO_DYNAMICS:-}" ]; then + EXTRA_ARGS+=("planner.pufferdrive_planner.ego_dynamics=$EGO_DYNAMICS") +fi +[ -n "${ROLLOUT_STEP_SECONDS:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.rollout_step_seconds=$ROLLOUT_STEP_SECONDS") +[ -n "${ROLLOUT_POLICY_VEHICLES:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.rollout_policy_vehicles=$ROLLOUT_POLICY_VEHICLES") +[ -n "${ROLLOUT_ALLOW_REVERSE:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.rollout_allow_reverse=$ROLLOUT_ALLOW_REVERSE") +[ -n "${BICYCLE_MIN_SPEED_MPS:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.bicycle_min_speed_mps=$BICYCLE_MIN_SPEED_MPS") +[ -n "${BICYCLE_UNDO_LOW_PASS:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.bicycle_undo_low_pass=$BICYCLE_UNDO_LOW_PASS") +[ -n "${BICYCLE_CENTER_ACCEL_CAP:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.bicycle_center_accel_cap=$BICYCLE_CENTER_ACCEL_CAP") +[ -n "${BICYCLE_CENTER_ACCEL_CAP_MARGIN:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.bicycle_center_accel_cap_margin=$BICYCLE_CENTER_ACCEL_CAP_MARGIN") +[ -n "${BICYCLE_MAX_YAW_JERK:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.bicycle_max_yaw_jerk_radps3=$BICYCLE_MAX_YAW_JERK") +[ -n "${EGO_ACCEL_CAP_MPS2:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.accel_cap_mps2=$EGO_ACCEL_CAP_MPS2") +[ -n "${BICYCLE_FLOOR_ACCEL_CLAMP:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.bicycle_floor_accel_clamp=$BICYCLE_FLOOR_ACCEL_CLAMP") +[ -n "${PARTNER_SPLIT_MAX_LENGTH_M:-}" ] && EXTRA_ARGS+=("planner.pufferdrive_planner.partner_split_max_length_m=$PARTNER_SPLIT_MAX_LENGTH_M") +[ -n "${STANDSTILL_JERK_DEADBAND:-}" ] && EXTRA_ARGS+=("++planner.pufferdrive_planner.env_overrides.eval_standstill_jerk_deadband_mps3=$STANDSTILL_JERK_DEADBAND") [ -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 "${LANE_WIDTH_CLIP_MIN_M:-}" ] && EXTRA_ARGS+=("++planner.pufferdrive_planner.env_overrides.obs_lane_width_clip_min_m=$LANE_WIDTH_CLIP_MIN_M") +[ -n "${LANE_WIDTH_CLIP_MAX_M:-}" ] && EXTRA_ARGS+=("++planner.pufferdrive_planner.env_overrides.obs_lane_width_clip_max_m=$LANE_WIDTH_CLIP_MAX_M") [ -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") @@ -114,7 +175,7 @@ for CHALLENGE in $CHALLENGES; do scenario_filter=$SPLIT \ planner=pufferdrive_planner \ planner.pufferdrive_planner.checkpoint_path="$CKPT" \ - callback="${CALLBACKS:-[simulation_log_callback, carl_visualization_callback]}" \ + callback="${CALLBACKS:-[simulation_log_callback]}" \ main_callback="[time_callback, metric_file_callback, metric_aggregator_callback, csv_main_callback]" \ hydra.searchpath="[pkg://pufferlib.ocean.cosim.nuplan.config, pkg://carl_nuplan.planning.script.config.common, pkg://carl_nuplan.planning.script.config.simulation, pkg://nuplan.planning.script.config.common, pkg://nuplan.planning.script.experiments]" \ group="$GROUP" \ @@ -132,6 +193,13 @@ elif [ "$OBS_HTML" = "infractions" ]; then --metric no_ego_at_fault_collisions --metric drivable_area_compliance \ --metric driving_direction_compliance --metric ego_is_making_progress fi +TIDY_KEEP=$OBS_HTML +[ "$OBS_HTML" = "0" ] && TIDY_KEEP=infractions +if [ "${TIDY:-1}" = "1" ]; then + "$PY" "$PD/scripts/eval/tidy_nuplan_run.py" "$GROUP" \ + --keep "$TIDY_KEEP" --max-score "${COSIM_OBS_HTML_MAX_SCORE:-0.9}" \ + || echo "tidy failed (results above are unaffected)" +fi echo "=========== SUMMARY ===========" rc=0 @@ -141,4 +209,5 @@ for c in "${!STATUS[@]}"; do done echo "group -> $GROUP" [ -f "$GROUP/obs_html/index.html" ] && echo "obs replays -> $GROUP/obs_html/index.html" +[ -f "$GROUP/overview.csv" ] && echo "overview -> $GROUP/overview.csv" exit $rc diff --git a/pufferlib/ocean/cosim/nuplan_bicycle.py b/pufferlib/ocean/cosim/nuplan_bicycle.py new file mode 100644 index 0000000000..6639ae6d1f --- /dev/null +++ b/pufferlib/ocean/cosim/nuplan_bicycle.py @@ -0,0 +1,144 @@ +"""Commands for nuPlan's KinematicBicycleModel from the shadow Drive env's post-step ego targets. + +The devkit model low-passes acceleration and steering with gain dt/(dt+tau) before a forward Euler step and +has neither a speed floor nor a stop snap: both are enforced here on the realised (low-passed) acceleration. +""" + +import math +from dataclasses import dataclass +from typing import Optional + +DEVKIT_ACCEL_TIME_CONSTANT_S = 0.2 +DEVKIT_STEERING_TIME_CONSTANT_S = 0.05 +YAW_JERK_MIN_SPEED_MPS = 0.1 # below this the yaw rate is ~0 whatever the wheel does; the atan bounds degenerate +FLOOR_SPEED_TOL_MPS = 1e-6 # the floor lands the model on min_speed up to float rounding + + +@dataclass(frozen=True) +class BicycleCommandConfig: + dt_s: float + min_speed_mps: float # floor on the model's next speed; 0 = CaRL's no-reverse driving + undo_low_pass: bool # pre-compensate the command lag so one model step lands on the targets + accel_time_constant_s: float = DEVKIT_ACCEL_TIME_CONSTANT_S + steering_time_constant_s: float = DEVKIT_STEERING_TIME_CONSTANT_S + center_accel_cap_mps2: float = 0.0 # cap on the devkit's center accel (rear + d*yaw_rate^2 + d*yaw_accel); 0 = off + max_yaw_jerk_radps3: float = 0.0 # per-step change of yaw accel bounded via the steering command; 0 = off + wheel_base_m: float = 0.0 + rear_axle_to_center_m: float = 0.0 + + def __post_init__(self): + if not self.dt_s > 0: + raise ValueError(f"dt_s must be > 0, got {self.dt_s}") + if not self.min_speed_mps <= 0: + raise ValueError(f"min_speed_mps must be <= 0, got {self.min_speed_mps}") + if not (self.accel_time_constant_s >= 0 and self.steering_time_constant_s >= 0): + raise ValueError("the low-pass time constants must be >= 0") + if not self.center_accel_cap_mps2 >= 0: + raise ValueError(f"center_accel_cap_mps2 must be >= 0, got {self.center_accel_cap_mps2}") + if self.center_accel_cap_mps2 > 0 and not (self.wheel_base_m > 0 and self.rear_axle_to_center_m >= 0): + raise ValueError("center_accel_cap_mps2 needs wheel_base_m > 0 and rear_axle_to_center_m >= 0") + if not self.max_yaw_jerk_radps3 >= 0: + raise ValueError(f"max_yaw_jerk_radps3 must be >= 0, got {self.max_yaw_jerk_radps3}") + if self.max_yaw_jerk_radps3 > 0 and not self.wheel_base_m > 0: + raise ValueError("max_yaw_jerk_radps3 needs wheel_base_m > 0") + + +def low_pass_gain(dt_s: float, time_constant_s: float) -> float: + return dt_s / (dt_s + time_constant_s) + + +def low_pass_response(current: float, command: float, gain: float) -> float: + return current + gain * (command - current) + + +def low_pass_command(current: float, target: float, gain: float) -> float: + return current + (target - current) / gain + + +def floor_clamped_intent_accel(cfg: BicycleCommandConfig, speed_mps: float, intent_accel_mps2: float) -> float: + """Drive zeroes the acceleration when the car stops; with the model held on the floor a negative intent would + wind down to the clip while the speed never moves, a state the policy never saw.""" + if speed_mps <= cfg.min_speed_mps + FLOOR_SPEED_TOL_MPS and intent_accel_mps2 < 0: + return 0.0 + return intent_accel_mps2 + + +def center_accel_shift( + cfg: BicycleCommandConfig, accel_mps2: float, speed_mps: float, steering_next_rad: float, yaw_rate_radps: float +) -> float: + """What the devkit adds to the rear-axle acceleration at the vehicle center after one model step.""" + speed_next = speed_mps + accel_mps2 * cfg.dt_s + yaw_rate_next = speed_next * math.tan(steering_next_rad) / cfg.wheel_base_m + yaw_accel_next = (yaw_rate_next - yaw_rate_radps) / cfg.dt_s + return cfg.rear_axle_to_center_m * (yaw_rate_next * yaw_rate_next + yaw_accel_next) + + +def clip_steering_for_yaw_jerk( + cfg: BicycleCommandConfig, + steering_next_rad: float, + speed_next_mps: float, + yaw_rate_radps: float, + yaw_rate_prev_radps: float, +) -> float: + """Next steering angle such that the yaw acceleration changes by at most max_yaw_jerk * dt within the step.""" + if speed_next_mps < YAW_JERK_MIN_SPEED_MPS: + return steering_next_rad + yaw_rate_window = cfg.max_yaw_jerk_radps3 * cfg.dt_s * cfg.dt_s + yaw_rate_center = 2.0 * yaw_rate_radps - yaw_rate_prev_radps + steering_lo = math.atan((yaw_rate_center - yaw_rate_window) * cfg.wheel_base_m / speed_next_mps) + steering_hi = math.atan((yaw_rate_center + yaw_rate_window) * cfg.wheel_base_m / speed_next_mps) + return min(max(steering_next_rad, steering_lo), steering_hi) + + +def bicycle_commands( + cfg: BicycleCommandConfig, + target_accel_mps2: float, + target_steering_rad: float, + current_accel_mps2: float, + current_steering_rad: float, + speed_mps: float, + yaw_rate_radps: float = 0.0, + yaw_rate_prev_radps: Optional[float] = None, + stop_this_step: bool = False, +): + """-> (acceleration [m/s^2], steering rate [rad/s]) for a CaRL ActionTrajectory over the next dt. + stop_this_step: the shadow env snapped its speed to zero this step, so the realised acceleration brings + the model to exactly zero speed too. With center_accel_cap_mps2 the realised acceleration is reduced + while turning so the center acceleration nuPlan scores stays under the cap (CaRL clip_angular_adjustment); + with max_yaw_jerk_radps3 the steering command is slowed so the yaw acceleration (and with it the center + acceleration's derivative, nuPlan's longitudinal jerk) cannot jump within a step.""" + accel_gain = low_pass_gain(cfg.dt_s, cfg.accel_time_constant_s) + steering_gain = low_pass_gain(cfg.dt_s, cfg.steering_time_constant_s) + accel_cmd = target_accel_mps2 + steering_cmd = target_steering_rad + if cfg.undo_low_pass: + accel_cmd = low_pass_command(current_accel_mps2, target_accel_mps2, accel_gain) + steering_cmd = low_pass_command(current_steering_rad, target_steering_rad, steering_gain) + if stop_this_step: + accel_cmd = low_pass_command(current_accel_mps2, -speed_mps / cfg.dt_s, accel_gain) + accel_realised = low_pass_response(current_accel_mps2, accel_cmd, accel_gain) + if cfg.max_yaw_jerk_radps3 > 0 and yaw_rate_prev_radps is not None: + steering_next = clip_steering_for_yaw_jerk( + cfg, + low_pass_response(current_steering_rad, steering_cmd, steering_gain), + speed_mps + accel_realised * cfg.dt_s, + yaw_rate_radps, + yaw_rate_prev_radps, + ) + steering_cmd = low_pass_command(current_steering_rad, steering_next, steering_gain) + if cfg.center_accel_cap_mps2 > 0: + steering_next = low_pass_response(current_steering_rad, steering_cmd, steering_gain) + shift = center_accel_shift(cfg, accel_realised, speed_mps, steering_next, yaw_rate_radps) + if accel_realised > cfg.center_accel_cap_mps2 - shift: + accel_cmd = low_pass_command(current_accel_mps2, cfg.center_accel_cap_mps2 - shift, accel_gain) + floor_accel_mps2 = (cfg.min_speed_mps - speed_mps) / cfg.dt_s + accel_cmd = max(accel_cmd, low_pass_command(current_accel_mps2, floor_accel_mps2, accel_gain)) + return accel_cmd, (steering_cmd - current_steering_rad) / cfg.dt_s + + +def bicycle_response( + cfg: BicycleCommandConfig, current_accel_mps2: float, accel_cmd_mps2: float, speed_mps: float, elapsed_s: float +): + """(acceleration, speed) the devkit model reaches after elapsed_s under the command: the tracking guard.""" + accel = low_pass_response(current_accel_mps2, accel_cmd_mps2, low_pass_gain(elapsed_s, cfg.accel_time_constant_s)) + return accel, speed_mps + accel * elapsed_s diff --git a/pufferlib/ocean/cosim/nuplan_bridge.py b/pufferlib/ocean/cosim/nuplan_bridge.py index b07288adb9..f2032a51a5 100644 --- a/pufferlib/ocean/cosim/nuplan_bridge.py +++ b/pufferlib/ocean/cosim/nuplan_bridge.py @@ -22,8 +22,12 @@ writer stays testable from an env without the nuplan-devkit installed. """ +import math +import os from pathlib import Path +from typing import Tuple + import numpy as np import data_utils.mirror_map_bin as _mbin @@ -40,8 +44,15 @@ # nuPlan TrafficLightStatusType -> PufferDrive light enum (datatypes.h: RED=1 YELLOW=2 GREEN=3). # Training never produces UNKNOWN (0), so an unknown or unreported light is treated as GREEN. +LIGHT_STATE_RED = 1 +LIGHT_STATE_YELLOW = 2 LIGHT_STATE_GREEN = 3 -_NUPLAN_LIGHT_STATE = {"RED": 1, "YELLOW": 2, "GREEN": LIGHT_STATE_GREEN, "UNKNOWN": LIGHT_STATE_GREEN} +_NUPLAN_LIGHT_STATE = { + "RED": LIGHT_STATE_RED, + "YELLOW": LIGHT_STATE_YELLOW, + "GREEN": LIGHT_STATE_GREEN, + "UNKNOWN": LIGHT_STATE_GREEN, +} FAR_AWAY = 1.0e6 # park surplus PufferDrive agents out of observation range ROUTE_SEARCH_DEPTH_BLOCKS = 30 # PDM's Dijkstra window over route roadblocks @@ -575,17 +586,30 @@ def match_connectors_to_stop_lines( MOVING_PARTNER_TYPES = ("VEHICLE", "PEDESTRIAN", "BICYCLE") # the only types the training bins turn into agents STATIC_PARTNER_TYPES = ("TRAFFIC_CONE", "BARRIER", "CZONE_SIGN", "GENERIC_OBJECT") +# get_proximal_map_objects matches lane connectors by centerline (lanes by polygon): the patch must reach the +# centerline of every connector whose polygon can overlap the object's box +STATIC_PARTNER_LANE_QUERY_RADIUS_M = 10.0 + + +def lane_polygon_overlaps(lane, box) -> bool: + try: + polygon = lane.polygon + except AssertionError: # devkit: lane connector without a scaled-width polygon row, nothing to hit + return False + return bool(polygon.intersects(box)) def partner_tracked_objects(tracked_objects, map_api, static_on_lane: dict): """Objects the shadow env should see: moving agents always; static clutter (cones, barriers, signs, - generic objects) only when it stands inside a lane or lane-connector polygon, where the ego can hit it. - The wider drivable area also covers intersection corners and crosswalks, where barrier rows line the - curb: those filled the nearest-N partner slots as 0.6 x 3 m stopped vehicles the policy never saw in - training. static_on_lane: track_token -> bool, filled here (statics never move).""" + generic objects) only when its box overlaps a lane or lane-connector polygon, where the ego can hit it + (a center-point test dropped a curbside object protruding 0.7 m into the lane). The wider drivable area + also covers intersection corners and crosswalks, where barrier rows line the curb: those filled the + nearest-N partner slots as 0.6 x 3 m stopped vehicles the policy never saw in training. + static_on_lane: track_token -> bool, filled here (statics never move).""" from nuplan.common.actor_state.state_representation import Point2D from nuplan.common.maps.maps_datatypes import SemanticMapLayer + lane_layers = [SemanticMapLayer.LANE, SemanticMapLayer.LANE_CONNECTOR] kept = [] for obj in tracked_objects: type_name = obj.tracked_object_type.name @@ -596,12 +620,11 @@ def partner_tracked_objects(tracked_objects, map_api, static_on_lane: dict): continue on_lane = static_on_lane.get(obj.track_token) if on_lane is None: - point = Point2D(obj.center.x, obj.center.y) - # is_in_layer(LANE_CONNECTOR) tests the centerline layer; the polygon lookup is get_all_map_objects - on_lane = bool( - map_api.is_in_layer(point, SemanticMapLayer.LANE) - or map_api.get_all_map_objects(point, SemanticMapLayer.LANE_CONNECTOR) + box = obj.box.geometry + proximal = map_api.get_proximal_map_objects( + Point2D(obj.center.x, obj.center.y), STATIC_PARTNER_LANE_QUERY_RADIUS_M, lane_layers ) + on_lane = any(lane_polygon_overlaps(lane, box) for layer in lane_layers for lane in proximal[layer]) static_on_lane[obj.track_token] = on_lane if on_lane: kept.append(obj) @@ -620,6 +643,40 @@ def pad_vru_partner_sizes(types, lengths, widths, pad_m: float): return lengths, widths +def split_long_partners(x, y, z, h, vx, vy, types, lengths, widths, max_length_m: float): + """Partners longer than max_length_m become ceil(length / max_length_m) equal boxes along their heading, in + place of the original (training spawns nothing longer than 7 m; a 12 m bus is one box the policy never saw).""" + out = [[] for _ in range(9)] + for j in range(len(x)): + pieces = max(1, int(math.ceil(float(lengths[j]) / max_length_m))) + piece_length = float(lengths[j]) / pieces + for k in range(pieces): + along = (k + 0.5) * piece_length - 0.5 * float(lengths[j]) + values = ( + x[j] + along * math.cos(h[j]), + y[j] + along * math.sin(h[j]), + z[j], + h[j], + vx[j], + vy[j], + types[j], + piece_length, + widths[j], + ) + for column, value in zip(out, values): + column.append(value) + dtypes = (np.float32,) * 6 + (np.int32, np.float32, np.float32) + return tuple(np.array(column, dtype) for column, dtype in zip(out, dtypes)) + + +def nearest_partner_slots(x, y, ego_x: float, ego_y: float, slots: int): + """Row order for at most `slots` partners: the devkit order when they fit, else nearest to the ego first.""" + if len(x) <= slots: + return np.arange(len(x)) + distance_sq = (np.asarray(x, np.float64) - ego_x) ** 2 + (np.asarray(y, np.float64) - ego_y) ** 2 + return np.argsort(distance_sq, kind="stable")[:slots] + + def tracked_objects_to_arrays(tracked_objects, transform: NuPlanTransform, first_slot: int = 1): """nuPlan DetectionsTracks agents -> (idx, x, y, z, h, vx, vy, types, lengths, widths) arrays in the bin frame, filling PufferDrive slots first_slot..N.""" @@ -676,3 +733,61 @@ def traffic_light_states( elif j not in route_elements and _LIGHT_RESTRICTIVENESS[state] < _LIGHT_RESTRICTIVENESS[int(states[j])]: states[j] = state return states + + +SYMMETRIC_WIDTH_MIN_M = 1.0 # floor where a baseline touches its polygon edge (lane tips, fan-outs) +SYMMETRIC_WIDTH_MAX_M = 50.0 # PufferDrive's loader limit +SYMMETRIC_WIDTH_MATCH_TOL_M = 0.5 # mean vertex distance to the matched baseline; farther lanes keep their widths +SYMMETRIC_WIDTH_SEARCH_M = 3.0 + + +def symmetric_lane_widths(roads, transform: NuPlanTransform, map_objects) -> Tuple[int, int]: + """Rewrite every bin lane's per-vertex widths as twice the distance to the nearer boundary of the nuPlan + lane / lane connector whose baseline the bin lane lies on: the widest symmetric corridor around the + centerline. `map_objects` expose baseline_path / left_boundary / right_boundary with `.linestring`. + Returns (lanes patched, lanes left unchanged because no baseline matched).""" + import shapely + from shapely.strtree import STRtree + + baselines = [obj.baseline_path.linestring for obj in map_objects] + tree = STRtree(baselines) + patched = unmatched = 0 + for road in roads: + if "widths" not in road or not road["S"]: + continue + ux, uy = transform.bin_to_loc(np.asarray(road["x"], np.float64), np.asarray(road["y"], np.float64)) + points = shapely.points(ux, uy) + search = shapely.box(ux.min(), uy.min(), ux.max(), uy.max()).buffer(SYMMETRIC_WIDTH_SEARCH_M) + candidates = tree.query(search) + if not len(candidates): + unmatched += 1 + continue + mean_dist = [float(np.mean(shapely.distance(baselines[c], points))) for c in candidates] + best = int(candidates[int(np.argmin(mean_dist))]) + if min(mean_dist) > SYMMETRIC_WIDTH_MATCH_TOL_M: + unmatched += 1 + continue + try: + left = map_objects[best].left_boundary.linestring + right = map_objects[best].right_boundary.linestring + except Exception: # noqa: BLE001 -- a connector without a scaled-width polygon row keeps its widths + unmatched += 1 + continue + nearer = np.minimum(shapely.distance(left, points), shapely.distance(right, points)) + road["widths"] = tuple(np.clip(2.0 * nearer, SYMMETRIC_WIDTH_MIN_M, SYMMETRIC_WIDTH_MAX_M).tolist()) + patched += 1 + return patched, unmatched + + +def patch_bin_lane_widths(bin_path: Path, out_path: Path, transform: NuPlanTransform, map_objects) -> Tuple[int, int]: + """Copy of `bin_path` with symmetric_lane_widths applied, written atomically to `out_path`.""" + data = _mbin.read_bin(Path(bin_path)) + if not data["has_width_section"]: + raise ValueError(f"{bin_path}: no LANEWID1 section to rewrite") + counts = symmetric_lane_widths(data["roads"], transform, map_objects) + out_path = Path(out_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + tmp = out_path.with_name(f"{out_path.name}.{os.getpid()}.tmp") + _mbin.write_bin(data, tmp) + os.replace(tmp, out_path) + return counts diff --git a/pufferlib/ocean/cosim/nuplan_rollout.py b/pufferlib/ocean/cosim/nuplan_rollout.py new file mode 100644 index 0000000000..f86e5fa154 --- /dev/null +++ b/pufferlib/ocean/cosim/nuplan_rollout.py @@ -0,0 +1,45 @@ +"""Forecast rollout helpers for the nuPlan planner's ego_dynamics=nuplan_tracker. + +Every planning step the shadow env is rolled out for the trajectory horizon: the policy drives the ego and the nearest +on-lane vehicles, every other object moves at its observed velocity and the lights hold their state. The ego's +rolled-out states are the trajectory nuPlan's two-stage controller tracks. +""" + +import math + +import numpy as np + +from pufferlib.ocean.cosim.nuplan_bridge import LIGHT_STATE_RED, LIGHT_STATE_YELLOW + + +STEP_COUNT_TOLERANCE = 1e-9 # 8.1 s / 0.3 s is 27 steps, not 28, despite float rounding + + +def rollout_step_count(horizon_seconds: float, step_seconds: float) -> int: + """Rollout steps whose end reaches the trajectory horizon.""" + if not (horizon_seconds > 0 and step_seconds > 0): + raise ValueError(f"horizon and step must be > 0, got {horizon_seconds} and {step_seconds}") + return max(1, math.ceil(horizon_seconds / step_seconds - STEP_COUNT_TOLERANCE)) + + +def split_rollout_slots(distance_sq, policy_eligible, policy_slot_count: int): + """-> (policy_rows, static_rows) of the partner arrays: the nearest policy-eligible partners take the policy slots, + every other partner a static slot.""" + distance_sq = np.asarray(distance_sq, np.float64) + eligible_rows = np.flatnonzero(np.asarray(policy_eligible, bool)) + nearest_first = eligible_rows[np.argsort(distance_sq[eligible_rows], kind="stable")] + policy_rows = np.sort(nearest_first[:policy_slot_count]) + is_static = np.ones(len(distance_sq), bool) + is_static[policy_rows] = False + return policy_rows, np.flatnonzero(is_static) + + +def forecast_light_states(states, yellow_remaining_seconds, elapsed_seconds: float): + """Light states elapsed_seconds into a forecast: a yellow whose remaining time has run out is red, every other + light holds its state (nuPlan reports no future phases). -> (states, yellow_remaining_seconds).""" + states = np.asarray(states, np.int32) + remaining = np.maximum(np.asarray(yellow_remaining_seconds, np.float32) - np.float32(elapsed_seconds), 0.0) + expired = (states == LIGHT_STATE_YELLOW) & (remaining <= 0.0) + forecast_states = np.where(expired, LIGHT_STATE_RED, states).astype(np.int32) + forecast_remaining = np.where(forecast_states == LIGHT_STATE_YELLOW, remaining, 0.0).astype(np.float32) + return forecast_states, forecast_remaining diff --git a/pufferlib/ocean/cosim/obs_replay.py b/pufferlib/ocean/cosim/obs_replay.py index 444b442a33..f392495bf3 100644 --- a/pufferlib/ocean/cosim/obs_replay.py +++ b/pufferlib/ocean/cosim/obs_replay.py @@ -139,6 +139,7 @@ def _env_cfg(self): "obs_norm_speed_mps": float(env.obs_norm_speed_mps), "obs_partner_relative_velocity": bool(env.obs_partner_relative_velocity), "obs_lane_speed_limit": bool(env.obs_lane_speed_limit), + "obs_traffic_light_yellow_remaining": bool(env.obs_traffic_light_yellow_remaining), "obs_goal_lane_distance": bool(env.obs_goal_lane_distance), "ego_features": int(env.ego_features), "num_reward_coefs": int(env.num_reward_coefs), @@ -158,7 +159,7 @@ def _check_layout(self, obs_dim): + env.obs_slots_partners_n * env.partner_features + env.obs_slots_lane_kept * env.lane_features + env.obs_slots_boundary_kept * env.boundary_features - + env.obs_slots_traffic_controls_n * binding.TRAFFIC_CONTROL_FEATURES + + env.obs_slots_traffic_controls_n * env.traffic_control_features ) if obs_dim < viewer_dim or obs_dim - viewer_dim > OBS_TRAILING_COUNT_FEATURES: raise RuntimeError( diff --git a/pufferlib/ocean/drive/binding.c b/pufferlib/ocean/drive/binding.c index d21d15200b..34767f9ca9 100644 --- a/pufferlib/ocean/drive/binding.c +++ b/pufferlib/ocean/drive/binding.c @@ -1830,7 +1830,10 @@ static PyObject *my_shared(PyObject *self, PyObject *args, PyObject *kwargs) { env->num_max_agents = max_agents_per_env; env->eval_mode = eval_mode; env->goal_radius = goal_radius; - load_map_binary(map_file, env); + if (load_map_binary(map_file, env) != 0) { + PyErr_Format(PyExc_RuntimeError, "Failed to load map binary: %s", map_file); + return NULL; + } if (init_grid_map(env) != 0) { PyErr_Format(PyExc_RuntimeError, "Failed to build grid map for map: %s", map_file); return NULL; @@ -2023,6 +2026,18 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) { env->disable_red_light_infractions = (int) unpack(kwargs, "disable_red_light_infractions"); env->disable_stop_sign_infractions = (int) unpack(kwargs, "disable_stop_sign_infractions"); env->traffic_light_junction_phases = (int) unpack(kwargs, "traffic_light_junction_phases"); + env->traffic_light_green_min_s = (float) unpack(kwargs, "traffic_light_green_min_s"); + env->traffic_light_yellow_min_s = (float) unpack(kwargs, "traffic_light_yellow_min_s"); + env->traffic_light_yellow_max_s = (float) unpack(kwargs, "traffic_light_yellow_max_s"); + if (!(env->traffic_light_green_min_s > 0.0f) || env->traffic_light_green_min_s > TL_DEFAULT_GREEN_DURATION + || !(env->traffic_light_yellow_min_s > 0.0f) + || !(env->traffic_light_yellow_max_s >= env->traffic_light_yellow_min_s) + || !isfinite(env->traffic_light_yellow_max_s)) { + PyErr_SetString( + PyExc_ValueError, + "traffic light cycle: need 0 < green_min <= 10 and 0 < yellow_min <= yellow_max"); + return -1; + } env->traffic_light_behavior = (int) unpack(kwargs, "traffic_light_behavior"); env->stop_sign_behavior = (int) unpack(kwargs, "stop_sign_behavior"); env->use_map_cache = (int) unpack(kwargs, "use_map_cache"); @@ -2060,6 +2075,7 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) { env->obs_lane_heading_signed = (int) unpack(kwargs, "obs_lane_heading_signed"); env->obs_lane_speed_limit = (int) unpack(kwargs, "obs_lane_speed_limit"); env->obs_slots_traffic_controls_n = (int) unpack(kwargs, "obs_slots_traffic_controls_n"); + env->obs_traffic_light_yellow_remaining = (int) unpack(kwargs, "obs_traffic_light_yellow_remaining"); env->traffic_lights_enabled = (bool) unpack(kwargs, "traffic_lights_enabled"); env->stop_signs_enabled = (bool) unpack(kwargs, "stop_signs_enabled"); env->yield_signs_enabled = (bool) unpack(kwargs, "yield_signs_enabled"); @@ -2114,6 +2130,11 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) { env->num_max_agents = (int) unpack(kwargs, "max_agents_per_env"); env->cosim_partner_slots = (int) unpack(kwargs, "cosim_partner_slots"); env->cosim_eval_semantics = (int) unpack(kwargs, "cosim_eval_semantics"); + env->scripted_pedestrian_slots = (int) unpack(kwargs, "scripted_pedestrian_slots"); + if (env->scripted_pedestrian_slots < 0) { + PyErr_SetString(PyExc_ValueError, "scripted_pedestrian_slots must be >= 0"); + return -1; + } int init_step = (int) unpack(kwargs, "init_step"); env->init_step = init_step; env->timestep = init_step; @@ -2133,6 +2154,13 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) { env->non_vehicle_controller = (int) unpack(kwargs, "non_vehicle_controller"); env->replay_expert_agents = (int) unpack(kwargs, "replay_expert_agents"); env->simulation_mode = (int) unpack(kwargs, "simulation_mode"); + if (env->scripted_pedestrian_slots > 0 + && (env->simulation_mode != SIMULATION_MODE_GIGAFLOW || env->cosim_partner_slots > 0)) { + PyErr_SetString( + PyExc_ValueError, + "scripted_pedestrian_slots needs gigaflow simulation_mode without cosim_partner_slots"); + return -1; + } env->reward_conditioning = (bool) unpack(kwargs, "reward_conditioning"); env->reward_randomization = (bool) unpack(kwargs, "reward_randomization"); env->reward_log_sampling = (bool) unpack(kwargs, "reward_log_sampling"); @@ -2148,11 +2176,29 @@ 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->red_light_train_front_margin_m = (float) unpack(kwargs, "red_light_train_front_margin_m"); + if (!(env->red_light_train_front_margin_m >= 0.0f) || !isfinite(env->red_light_train_front_margin_m)) { + PyErr_SetString(PyExc_ValueError, "red_light_train_front_margin_m must be finite and >= 0"); + return -1; + } + env->red_light_train_rear_margin_m = (float) unpack(kwargs, "red_light_train_rear_margin_m"); + if (!(env->red_light_train_rear_margin_m >= 0.0f) || !isfinite(env->red_light_train_rear_margin_m)) { + PyErr_SetString(PyExc_ValueError, "red_light_train_rear_margin_m must be finite and >= 0"); + return -1; + } 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->obs_lane_width_clip_min_m = (float) unpack(kwargs, "obs_lane_width_clip_min_m"); + env->obs_lane_width_clip_max_m = (float) unpack(kwargs, "obs_lane_width_clip_max_m"); + if (!(env->obs_lane_width_clip_min_m >= 0.0f) || !isfinite(env->obs_lane_width_clip_min_m) + || !(env->obs_lane_width_clip_max_m >= 0.0f) || !isfinite(env->obs_lane_width_clip_max_m) + || (env->obs_lane_width_clip_max_m > 0.0f && env->obs_lane_width_clip_min_m > env->obs_lane_width_clip_max_m)) { + PyErr_SetString(PyExc_ValueError, "obs_lane_width_clip_min_m/max_m must be finite, >= 0 and min <= max"); + 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"); @@ -2181,7 +2227,10 @@ static int my_init(Env *env, PyObject *args, PyObject *kwargs) { env->phantom_braking_duration = (int) ceilf((float) unpack(kwargs, "phantom_braking_duration_seconds") / env->dt); env->phantom_braking_freeze_steering = (int) unpack(kwargs, "phantom_braking_freeze_steering"); - init(env); + if (init(env) != 0) { + PyErr_Format(PyExc_RuntimeError, "Failed to initialise the env for map %s (reason on stderr)", env->map_name); + return -1; + } // Episodes must always be generated by c_reset's full path so a logged // episode_seed replays identically regardless of which reset created it. env->timestep = -1; @@ -2253,6 +2302,9 @@ static int my_log(PyObject *dict, Env *env, Log *log, float n) { assign_to_dict(dict, "spawn_failed_goal", log->spawn_failed_goal); assign_to_dict(dict, "stopped_at_reset", log->stopped_at_reset); assign_to_dict(dict, "early_reset_short", log->early_reset_short); + assign_to_dict(dict, "pedestrian_collision_rate", log->pedestrian_collision_rate); + assign_to_dict(dict, "pedestrian_crossings", log->pedestrian_crossings); + assign_to_dict(dict, "pedestrian_bystanders", log->pedestrian_bystanders); assign_to_dict(dict, "offroad_rate", log->offroad_rate); assign_to_dict(dict, "episode_length", log->episode_length); assign_to_dict(dict, "collision_rate", log->collision_rate); diff --git a/pufferlib/ocean/drive/constants.h b/pufferlib/ocean/drive/constants.h index 02f1c8f908..a8836a7a85 100644 --- a/pufferlib/ocean/drive/constants.h +++ b/pufferlib/ocean/drive/constants.h @@ -184,6 +184,67 @@ static const float STEERING_VALUES[NUM_STEERING_ACTIONS] #define REPLAY_MAX_CONTROLLED_LENGTH_M 6.0f // Longer vehicles (trucks, buses) are made static #define REPLAY_SPAWN_MAX_LANE_DISTANCE_M 2.5f // Car must be this close to a lane center line +// Scripted pedestrians: one behaviour per appearance, drawn with the PEDESTRIAN_*_PROB shares +#define PEDESTRIAN_STATE_IDLE 0 // hidden, cooldown running +#define PEDESTRIAN_STATE_WALKING 1 // stands where it appeared, then walks its leg +#define PEDESTRIAN_STATE_STANDING 2 // leg done: stands until it vanishes +#define PEDESTRIAN_MODE_POP_IN_CROSS 0 // CARLA Scenario3: pops in at the right kerb close ahead, crosses fast +#define PEDESTRIAN_MODE_SLOW_CROSS 1 // crosses at walking speed from either kerb, visible from afar +#define PEDESTRIAN_MODE_STAND 2 // stands beside the road, never enters it +#define PEDESTRIAN_MODE_KERB_WALK 3 // walks along the kerb, never enters the road +#define PEDESTRIAN_MODE_KERB_STEP 4 // pops in like a crosser, steps toward or away from the road, then stands +#define PEDESTRIAN_POP_IN_CROSS_PROB 0.40f +#define PEDESTRIAN_SLOW_CROSS_PROB 0.15f +#define PEDESTRIAN_STAND_PROB 0.20f +#define PEDESTRIAN_KERB_WALK_PROB 0.15f // the rest: PEDESTRIAN_MODE_KERB_STEP +// Leaderboard 1.0 Scenario3 is the centre of every LB1 range; lanes = driving lanes from the car to the kerb, >= 1 +#define PEDESTRIAN_LB1_LANE_WIDTH_M 3.5f +#define PEDESTRIAN_LB1_APPEAR_BASE_M 12.0f // walker pops in 12 m + 1.5 m per extra lane ahead of the car +#define PEDESTRIAN_LB1_APPEAR_PER_LANE_M 1.5f +#define PEDESTRIAN_APPEAR_JITTER_M 2.0f +#define PEDESTRIAN_KERB_OFFSET_MIN_M 0.75f // LB1 stands half a lane (1.75 m) outside the road edge +#define PEDESTRIAN_KERB_OFFSET_MAX_M 2.75f +#define PEDESTRIAN_KERB_SEARCH_MAX_M 10.0f // lateral reach per side; wider roads get no crossing +#define PEDESTRIAN_FAR_KERB_OVERSHOOT_M 1.0f // walk this far past the far road edge +#define PEDESTRIAN_LB1_SPEED_BASE_MPS 3.4f // LB1 walks at 3 + 0.4 m/s x (lanes + 1) +#define PEDESTRIAN_LB1_SPEED_PER_LANE_MPS 0.4f +#define PEDESTRIAN_SPEED_JITTER_MPS 1.0f +#define PEDESTRIAN_WAIT_MIN_S 0.0f // LB1 walker stands ~0.25 s after popping in +#define PEDESTRIAN_WAIT_MAX_S 0.6f +#define PEDESTRIAN_RAMP_MIN_S 0.3f // LB1 walker reaches its speed ~0.5 s after setting off +#define PEDESTRIAN_RAMP_MAX_S 0.7f +#define PEDESTRIAN_SIZE_MIN_M 1.0f // CARLA walker box (0.4-0.5 m) plus the 0.8 m eval VRU pad +#define PEDESTRIAN_SIZE_MAX_M 1.5f +#define PEDESTRIAN_HEIGHT_M 1.8f +#define PEDESTRIAN_LINGER_MIN_S 1.0f // crossers and kerb walkers at the end of their leg +#define PEDESTRIAN_LINGER_MAX_S 3.0f +#define PEDESTRIAN_FAR_APPEAR_MIN_M 15.0f // standers and kerb walkers appear this far ahead of the car +#define PEDESTRIAN_FAR_APPEAR_MAX_M 60.0f +#define PEDESTRIAN_SLOW_APPEAR_MIN_M 25.0f // slow crossers: far enough to yield without an emergency stop +#define PEDESTRIAN_SLOW_APPEAR_MAX_M 70.0f +#define PEDESTRIAN_SLOW_SPEED_MIN_MPS 0.8f // nuPlan pedestrians cross at 0.6-1.7 m/s +#define PEDESTRIAN_SLOW_SPEED_MAX_MPS 2.0f +#define PEDESTRIAN_SLOW_WAIT_MAX_S 5.0f // standstill at the kerb before setting off +#define PEDESTRIAN_SLOW_SKEW_MAX_RAD 0.6f // crossing direction off the perpendicular +#define PEDESTRIAN_SLOW_AGAINST_TRAFFIC_PROB 0.5f // which way a skewed crossing leans +#define PEDESTRIAN_KERB_WALK_SPEED_MIN_MPS 0.7f +#define PEDESTRIAN_KERB_WALK_SPEED_MAX_MPS 1.7f +#define PEDESTRIAN_KERB_WALK_MIN_M 8.0f +#define PEDESTRIAN_KERB_WALK_MAX_M 25.0f +#define PEDESTRIAN_KERB_WALK_BEND_TOLERANCE_M 0.5f // the straight walk may leave the kerb line by this much mid-way +#define PEDESTRIAN_KERB_WALK_AGAINST_TRAFFIC_PROB 0.5f +#define PEDESTRIAN_KERB_STEP_MIN_M 0.5f // CARLA's stuck Scenario3 walkers move 0.7-1.4 m and stand +#define PEDESTRIAN_KERB_STEP_MAX_M 1.5f +#define PEDESTRIAN_KERB_STEP_EDGE_CLEARANCE_M 0.3f // a step toward the road ends this far outside the road edge +#define PEDESTRIAN_KERB_STEP_AWAY_PROB 0.5f +#define PEDESTRIAN_BYSTANDER_STAND_MIN_S 8.0f // standers and kerb steppers before they vanish +#define PEDESTRIAN_BYSTANDER_STAND_MAX_S 25.0f +#define PEDESTRIAN_LEFT_KERB_PROB 0.5f // every mode but the two Scenario3 ones +#define PEDESTRIAN_COOLDOWN_MIN_S 5.0f +#define PEDESTRIAN_COOLDOWN_MAX_S 20.0f +#define PEDESTRIAN_TARGET_MIN_SPEED_MPS 1.0f // only moving cars get a crossing +#define PEDESTRIAN_ARM_ATTEMPTS 4 // target cars tried per step once the cooldown ran out + // Stop line geometry (spawn overlap check and IDM braking) #define STOP_LINE_DIST_SQ (10.0f * 10.0f) #define STOP_LINE_EXTENSION_FACTOR 1.5f @@ -193,6 +254,7 @@ static const float STEERING_VALUES[NUM_STEERING_ACTIONS] #define RED_LIGHT_TRIGGER_DIST_SQ (30.0f * 30.0f) #define RED_LIGHT_LATERAL_EXTENSION_M 15.0f // Beyond each stop line endpoint, so the line cannot be driven around #define RED_LIGHT_ENTER_HEADING_THRESHOLD (M_PI / 2.0f) +#define RED_LIGHT_TAIL_CLEARANCE_M 1.0f // CARLA's RunningRedLightTest scores the tail down to 1 m behind the rear edge // Stop sign violation detection (CaRL RunStopSign2 semantics, trigger box centred on the stop line) #define STOP_SIGN_PROXIMITY_DIST_SQ (20.0f * 20.0f) @@ -215,6 +277,7 @@ static const float STEERING_VALUES[NUM_STEERING_ACTIONS] // => For each entity type in gridmap, diagonal poly-lines -> sqrt(2), include diagonal ends -> 2 #define MAX_ENTITIES_PER_CELL 64 #define ROAD_QUERY_CELL_COUNT 25 // 5x5 cell neighborhood +#define ROAD_QUERY_HALF_CELLS 2 // cells on each side of the window's centre cell #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 @@ -271,6 +334,7 @@ static const int ROAD_OFFSETS[25][2] #define PARTNER_FEATURES 9 #define PARTNER_RELATIVE_VELOCITY_FEATURES 2 // appended per partner row when obs_partner_relative_velocity is set #define TRAFFIC_CONTROL_FEATURES 7 +#define TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES 1 // before type/state when obs_traffic_light_yellow_remaining is set #define GOAL_FEATURES 3 #define OBS_VALID_COUNT_FEATURES 4 // Heading deviation since last kept point that forces a keep when obs stride > 1 (~15 degrees). diff --git a/pufferlib/ocean/drive/datatypes.h b/pufferlib/ocean/drive/datatypes.h index 16e7a35a90..3ad43d6082 100644 --- a/pufferlib/ocean/drive/datatypes.h +++ b/pufferlib/ocean/drive/datatypes.h @@ -110,6 +110,7 @@ struct Agent { float current_goal_z; // alias = list_goal_z[current_goal_idx] int current_goal_idx; // index of next goal to reach (0..N-1) int goal_count; // number of active goals (<= num_goals) + float lane_goal_spacing_m; // co-sim forecast agent: goals follow the lane ahead at this spacing; 0 = off float gt_goal_x; // Last valid ground-truth goal position x float gt_goal_y; // Last valid ground-truth goal position y float gt_goal_z; // Last valid ground-truth goal position z @@ -121,6 +122,7 @@ struct Agent { int stop_sign_stop_completed; // 0/1 -> stood still inside the target's trigger box int stop_sign_last_failed_idx; // element run without stopping, not re-targeted until 20 m away int stop_sign_standstill_idx; // element whose trigger box the agent last stood still in, -1 if none + int red_light_latched_idx; // light whose line the footprint has sat on since the violation, -1 if none // Jerk dynamics float accel_long; @@ -140,6 +142,36 @@ struct Agent { unsigned char is_phantom_braker; // episode-level flag: agent may phantom-brake }; +struct ScriptedPedestrian { + int agent_idx; // slot in Drive.agents + int target_agent_idx; // car it appeared for, -1 while idle + int state; // PEDESTRIAN_STATE_* + int mode; // PEDESTRIAN_MODE_*, drawn when it goes idle + float timer_s; // cooldown left while idle, standstill left before the leg, standing time left after it + float dir_x; // unit walking (and facing) direction + float dir_y; + float path_m; // length of the leg, 0 for a stander + float remaining_m; + float speed_mps; // walking speed once the ramp is over + float ramp_s; // seconds from setting off to speed_mps + float current_speed_mps; + float stand_s; // standing time after the leg + float climb_per_m; // z change per metre walked: legs that run along a sloped road +}; + +// Lane point with the road edge distance on both sides +typedef struct { + float x; + float y; + float z; + float right_x; // unit vector to the right of the lane heading + float right_y; + float right_kerb_m; + float left_kerb_m; + int lane_idx; + float lane_s_m; +} KerbSite; + typedef struct { float mid_x_m; float mid_y_m; diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index 7599963e0c..ae68dd050a 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -22,6 +22,7 @@ typedef struct { typedef struct Drive Drive; typedef struct Log Log; typedef struct Agent Agent; +typedef struct ScriptedPedestrian ScriptedPedestrian; typedef struct RoadMapElement RoadMapElement; typedef struct TrafficControlElement TrafficControlElement; typedef struct GridMapEntity GridMapEntity; @@ -95,6 +96,9 @@ struct Log { float spawn_failed_goal; float stopped_at_reset; float early_reset_short; + float pedestrian_collision_rate; + float pedestrian_crossings; + float pedestrian_bystanders; }; struct GridMapEntity { @@ -172,7 +176,9 @@ struct Drive { // Agents Agent *agents; int num_controllable_agents; - int cosim_partner_slots; // gigaflow co-sim: static slots owned by the external sim (never spawned/stepped) + int cosim_partner_slots; // gigaflow co-sim: static slots owned by the external sim (never spawned/stepped) + int scripted_pedestrian_slots; // gigaflow: scripted pedestrians appended after the policy agents + ScriptedPedestrian *pedestrians; int active_agent_count; int *active_agent_indices; int num_total_agents; @@ -243,6 +249,9 @@ struct Drive { int disable_red_light_infractions; int disable_stop_sign_infractions; int traffic_light_junction_phases; + float traffic_light_green_min_s; // training cycles: green ~ U(min, TL_DEFAULT_GREEN_DURATION) + float traffic_light_yellow_min_s; // training cycles: yellow ~ U(min, max) + float traffic_light_yellow_max_s; int stop_sign_behavior; int sdc_controller; int non_sdc_controller; @@ -296,6 +305,7 @@ struct Drive { int obs_lane_heading_signed; int obs_lane_speed_limit; int obs_slots_traffic_controls_n; + int obs_traffic_light_yellow_remaining; // row gains remaining yellow / traffic_light_yellow_max_s, clipped at 1 int traffic_lights_enabled; int stop_signs_enabled; int yield_signs_enabled; @@ -314,7 +324,11 @@ struct Drive { float obs_norm_road_seg_width_m; float obs_norm_z_m; float eval_perceived_size_margin_m; + float red_light_train_front_margin_m; // training-only gap the front edge must keep before a red stop line + float red_light_train_rear_margin_m; // training-only extension of the tail clearance behind the rear edge float obs_partner_split_length_m; // 0 disables; longer partners are observed as ceil(length / this) equal pieces + float obs_lane_width_clip_min_m; // observed lane width floor; 0 = no floor + float obs_lane_width_clip_max_m; // observed lane width ceiling; 0 disables both clips 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; @@ -535,6 +549,7 @@ static void reset_agent_state(Agent *agent) { agent->stop_sign_target_idx = -1; agent->stop_sign_stop_completed = 0; agent->stop_sign_last_failed_idx = -1; + agent->red_light_latched_idx = -1; agent->stop_sign_standstill_idx = -1; agent->current_lane_idx = -1; agent->previous_lane_idx = -1; @@ -1085,6 +1100,75 @@ static void commit_goals( agent->current_goal_z = agent->list_goal_z[start_slot_idx]; } +static int pick_straightest_exit_lane(const RoadMapElement *road_elements, int lane_idx) { + const RoadMapElement *lane = &road_elements[lane_idx]; + int straightest_exit_idx = -1; + float smallest_turn_rad = 2.0f * (float) M_PI; + for (int exit_slot = 0; exit_slot < lane->num_exits; exit_slot++) { + int exit_idx = lane->exit_lanes[exit_slot]; + if (exit_idx == -1 || road_elements[exit_idx].segment_size < 2) { + continue; + } + const RoadMapElement *exit_lane = &road_elements[exit_idx]; + float turn_rad + = fabsf(normalize_heading(exit_lane->headings[exit_lane->segment_size - 2] - exit_lane->headings[0])); + if (turn_rad < smallest_turn_rad) { + smallest_turn_rad = turn_rad; + straightest_exit_idx = exit_idx; + } + } + return straightest_exit_idx; +} + +// Up to num_goals goals every spacing_m along the agent's lane and its straightest exits; returns the count placed. +static int place_lane_goals_ahead(Drive *env, Agent *agent, float spacing_m) { + if (agent->current_lane_idx == -1) { + return 0; + } + RoadMapElement *start_lane = &env->road_elements[agent->current_lane_idx]; + float start_s_on_lane = compute_lane_progress( + start_lane, + agent->sim_x, + agent->sim_y, + agent->cos_heading, + agent->sin_heading, + true, + NULL); + float needed_route_m = start_s_on_lane + spacing_m * (float) env->num_goals; + int route[MAX_ROUTE_LENGTH]; + int route_length = 0; + float route_m = 0.0f; + int lane_idx = agent->current_lane_idx; + for (int hop = 0; hop < MAX_ROUTE_LENGTH && lane_idx != -1 && route_m < needed_route_m; hop++) { + route[route_length++] = lane_idx; + route_m += env->road_elements[lane_idx].length; + lane_idx = pick_straightest_exit_lane(env->road_elements, lane_idx); + } + float goal_spacings_m[MAX_GOALS]; + for (int goal_idx = 0; goal_idx < env->num_goals; goal_idx++) { + goal_spacings_m[goal_idx] = spacing_m; + } + float goal_x[MAX_GOALS], goal_y[MAX_GOALS], goal_z[MAX_GOALS]; + int goal_lane[MAX_GOALS]; + int placed_goal_count = chain_goals( + env, + route, + route_length, + 0, + start_s_on_lane, + agent->sim_heading, + goal_spacings_m, + env->num_goals, + goal_x, + goal_y, + goal_z, + goal_lane); + if (placed_goal_count > 0) { + commit_goals(env, agent, goal_x, goal_y, goal_z, goal_lane, placed_goal_count, 0); + } + return placed_goal_count; +} + static bool compute_new_route(Drive *env, Agent *agent, int current_lane_idx) { RoadMapElement *road_elements = env->road_elements; int candidate_route[MAX_ROUTE_LENGTH]; @@ -1655,6 +1739,13 @@ static bool check_segment_crosses_moving_box(float ax, float ay, float bx, float agent->sin_heading, &b_cur[0], &b_cur[1]); + // all four endpoint positions beyond one face of the box: the swept quad cannot reach it + if (fminf(fminf(a_prev[0], b_prev[0]), fminf(a_cur[0], b_cur[0])) > half_length + || fmaxf(fmaxf(a_prev[0], b_prev[0]), fmaxf(a_cur[0], b_cur[0])) < -half_length + || fminf(fminf(a_prev[1], b_prev[1]), fminf(a_cur[1], b_cur[1])) > half_width + || fmaxf(fmaxf(a_prev[1], b_prev[1]), fmaxf(a_cur[1], b_cur[1])) < -half_width) { + return false; + } if (check_segment_intersects_aabb(a_prev, b_prev, half_length, half_width) || check_segment_intersects_aabb(a_cur, b_cur, half_length, half_width) || check_segment_intersects_aabb(a_prev, a_cur, half_length, half_width) @@ -1694,67 +1785,6 @@ static bool check_segment_crosses_moving_box(float ax, float ay, float bx, float return has_positive_cross != has_negative_cross; } -static bool check_agent_on_stop_line(Drive *env, Agent *agent, bool include_yellow_violation) { - 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; - } - if (tc->num_controlled_lanes == 0) { - continue; - } - - int controls_lane = 0; - for (int j = 0; j < tc->num_controlled_lanes; j++) { - if (tc->controlled_lanes[j] == agent->current_lane_idx) { - controls_lane = 1; - break; - } - } - if (!controls_lane) { - continue; - } - if (env->timestep >= tc->state_size) { - continue; - } - int light_state = tc->states[env->timestep]; - if (light_state != TRAFFIC_CONTROL_STATE_RED - && !(include_yellow_violation && light_state == TRAFFIC_CONTROL_STATE_YELLOW)) { - continue; - } - - // Pre-filter: distance to stop line midpoint - float mid_x = (tc->stop_line[0] + tc->stop_line[3]) * 0.5f; - float mid_y = (tc->stop_line[1] + tc->stop_line[4]) * 0.5f; - float dx = agent->sim_x - mid_x; - float dy = agent->sim_y - mid_y; - if (dx * dx + dy * dy > STOP_LINE_DIST_SQ) { - continue; - } - - // Heading check: agent must be heading towards the stop line - float heading_diff = compute_heading_diff(agent->sim_heading, tc->heading); - if (fabsf(heading_diff) > STOP_LINE_HEADING_THRESHOLD) { - continue; - } - - // Stop line segment vector (endpoint 0 -> endpoint 1) - float sl_dx = tc->stop_line[3] - tc->stop_line[0]; - float sl_dy = tc->stop_line[4] - tc->stop_line[1]; - // Lengthen the segment by STOP_LINE_EXTENSION_FACTOR, growing equally - // from both endpoints so agents crossing near the edges are still caught. - float ext = (STOP_LINE_EXTENSION_FACTOR - 1.0f) * 0.5f; - float ext_p1[2] = {tc->stop_line[0] - ext * sl_dx, tc->stop_line[1] - ext * sl_dy}; - float ext_p2[2] = {tc->stop_line[3] + ext * sl_dx, tc->stop_line[4] + ext * sl_dy}; - - if (check_segment_crosses_moving_box(ext_p1[0], ext_p1[1], ext_p2[0], ext_p2[1], agent)) { - return true; - } - } - return false; -} - static bool check_corner_boxes_overlap(float corners_a[4][2], float corners_b[4][2], float axes[4][2]) { for (int i = 0; i < 4; i++) { float min_a = INFINITY, max_a = -INFINITY; @@ -1774,10 +1804,60 @@ static bool check_corner_boxes_overlap(float corners_a[4][2], float corners_b[4] return true; } -static bool check_red_light_violation(Drive *env, int agent_idx) { - Agent *agent = &env->agents[agent_idx]; - // Vehicle center crosses on red; the laterally extended line cannot be driven - // around, and the heading gate exempts exiting/opposing traffic. +static float red_light_train_margin_m(Drive *env, float train_margin_m) { + bool eval_semantics = (env->eval_mode && !env->eval_training_render) || env->cosim_eval_semantics; + return eval_semantics ? 0.0f : train_margin_m; +} + +// directed segment a->b crosses the laterally extended stop line in its travel direction: a before it, b at or past +static bool segment_crosses_stop_line(TrafficControlElement *tc, float a_x, float a_y, float b_x, float b_y) { + float mid_x = (tc->stop_line[0] + tc->stop_line[3]) * 0.5f; + float mid_y = (tc->stop_line[1] + tc->stop_line[4]) * 0.5f; + float line_dx = tc->stop_line[3] - tc->stop_line[0]; + float line_dy = tc->stop_line[4] - tc->stop_line[1]; + float line_len = sqrtf(line_dx * line_dx + line_dy * line_dy); + if (line_len <= 0.0f) { + return false; + } + float line_ux = line_dx / line_len, line_uy = line_dy / line_len; + // normal oriented along travel direction, so s < 0 is before the line + float normal_x = -line_uy, normal_y = line_ux; + if (normal_x * cosf(tc->heading) + normal_y * sinf(tc->heading) < 0.0f) { + normal_x = -normal_x; + normal_y = -normal_y; + } + float s_a = (a_x - mid_x) * normal_x + (a_y - mid_y) * normal_y; + float s_b = (b_x - mid_x) * normal_x + (b_y - mid_y) * normal_y; + if (!(s_a < 0.0f && s_b >= 0.0f)) { + return false; + } + float crossing_frac = s_a / (s_a - s_b); + float cross_x = a_x + crossing_frac * (b_x - a_x); + float cross_y = a_y + crossing_frac * (b_y - a_y); + float lateral = (cross_x - mid_x) * line_ux + (cross_y - mid_y) * line_uy; + return fabsf(lateral) <= 0.5f * line_len + RED_LIGHT_LATERAL_EXTENSION_M; +} + +// include_yellow: spawn probe, a yellow line has to be cleared before red just like a red one +static bool check_red_light_footprint(Drive *env, Agent *agent, bool include_yellow) { + // On red the line may not lie between the tail clearance point and the front edge, moving or standing + float front_offset_m + = 0.5f * agent->sim_length + red_light_train_margin_m(env, env->red_light_train_front_margin_m); + float rear_offset_m = 0.5f * agent->sim_length + RED_LIGHT_TAIL_CLEARANCE_M + + red_light_train_margin_m(env, env->red_light_train_rear_margin_m); + float front_x = agent->sim_x + front_offset_m * agent->cos_heading; + float front_y = agent->sim_y + front_offset_m * agent->sin_heading; + float rear_x = agent->sim_x - rear_offset_m * agent->cos_heading; + float rear_y = agent->sim_y - rear_offset_m * agent->sin_heading; + // one violation per stay: nothing fires while the footprint is on the latched line, re-armed once it has left + bool was_latched = agent->red_light_latched_idx != -1; + if (was_latched) { + TrafficControlElement *latched_light = &env->traffic_elements[agent->red_light_latched_idx]; + if (segment_crosses_stop_line(latched_light, rear_x, rear_y, front_x, front_y)) { + return false; + } + agent->red_light_latched_idx = -1; + } for (int i = 0; i < env->num_traffic_elements; i++) { TrafficControlElement *tc = &env->traffic_elements[i]; if (tc->type != TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT) { @@ -1786,7 +1866,9 @@ static bool check_red_light_violation(Drive *env, int agent_idx) { if (env->timestep >= tc->state_size) { continue; } - if (tc->states[env->timestep] != TRAFFIC_CONTROL_STATE_RED) { + int light_state = tc->states[env->timestep]; + if (light_state != TRAFFIC_CONTROL_STATE_RED + && !(include_yellow && light_state == TRAFFIC_CONTROL_STATE_YELLOW)) { continue; } float mid_x = (tc->stop_line[0] + tc->stop_line[3]) * 0.5f; @@ -1803,36 +1885,19 @@ static bool check_red_light_violation(Drive *env, int agent_idx) { if (fabsf(compute_heading_diff(agent->sim_heading, tc->heading)) >= RED_LIGHT_ENTER_HEADING_THRESHOLD) { continue; } - - float line_dx = tc->stop_line[3] - tc->stop_line[0]; - float line_dy = tc->stop_line[4] - tc->stop_line[1]; - float line_len = sqrtf(line_dx * line_dx + line_dy * line_dy); - if (line_len <= 0.0f) { - continue; - } - float line_ux = line_dx / line_len, line_uy = line_dy / line_len; - // normal oriented along travel direction, so s < 0 is before the line - float normal_x = -line_uy, normal_y = line_ux; - if (normal_x * cosf(tc->heading) + normal_y * sinf(tc->heading) < 0.0f) { - normal_x = -normal_x; - normal_y = -normal_y; - } - float s_prev = (agent->prev_x - mid_x) * normal_x + (agent->prev_y - mid_y) * normal_y; - float s_cur = (agent->sim_x - mid_x) * normal_x + (agent->sim_y - mid_y) * normal_y; - if (!(s_prev < 0.0f && s_cur >= 0.0f)) { + if (!segment_crosses_stop_line(tc, rear_x, rear_y, front_x, front_y)) { continue; } - float crossing_frac = s_prev / (s_prev - s_cur); - float cross_x = agent->prev_x + crossing_frac * (agent->sim_x - agent->prev_x); - float cross_y = agent->prev_y + crossing_frac * (agent->sim_y - agent->prev_y); - float lateral = (cross_x - mid_x) * line_ux + (cross_y - mid_y) * line_uy; - if (fabsf(lateral) <= 0.5f * line_len + RED_LIGHT_LATERAL_EXTENSION_M) { - return true; - } + agent->red_light_latched_idx = i; + return !was_latched; // an overlapping red line takes over the latch without a second violation } return false; } +static bool check_red_light_violation(Drive *env, int agent_idx) { + return check_red_light_footprint(env, &env->agents[agent_idx], false); +} + typedef struct { float mid_x; float mid_y; @@ -2491,6 +2556,9 @@ static void add_log(Drive *env) { episode_log.spawn_failed_goal += env->logs[i].spawn_failed_goal; episode_log.stopped_at_reset += env->logs[i].stopped_at_reset; episode_log.early_reset_short += env->logs[i].early_reset_short; + episode_log.pedestrian_collision_rate += env->logs[i].pedestrian_collision_rate; + episode_log.pedestrian_crossings += env->logs[i].pedestrian_crossings; + episode_log.pedestrian_bystanders += env->logs[i].pedestrian_bystanders; // Per-component reward sums (mirrors compute_rewards' env->rewards[i]+= sites). episode_log.reward_collision += env->logs[i].reward_collision; episode_log.reward_offroad += env->logs[i].reward_offroad; @@ -2674,8 +2742,9 @@ static int traffic_light_phase_count(Drive *env, int junction_id) { return max_phase_idx + 1; } +// Rounded up: a phase never lasts less than its configured duration static int duration_to_steps(float duration_seconds, float dt) { - int steps = (int) (duration_seconds / dt); + int steps = (int) ceilf(duration_seconds / dt); return steps < 1 ? 1 : steps; } @@ -2686,9 +2755,8 @@ static TrafficLightCycle sample_traffic_light_cycle(Drive *env, int junction_id, dur_yellow = TL_DEFAULT_YELLOW_DURATION; dur_red = TL_DEFAULT_RED_DURATION; } else { - dur_green = sample_uniform(&env->rng_state, 0.1 * TL_DEFAULT_GREEN_DURATION, TL_DEFAULT_GREEN_DURATION); - dur_yellow - = sample_uniform(&env->rng_state, 0.5f * TL_DEFAULT_YELLOW_DURATION, 0.75f * TL_DEFAULT_YELLOW_DURATION); + dur_green = sample_uniform(&env->rng_state, env->traffic_light_green_min_s, TL_DEFAULT_GREEN_DURATION); + dur_yellow = sample_uniform(&env->rng_state, env->traffic_light_yellow_min_s, env->traffic_light_yellow_max_s); dur_red = sample_uniform(&env->rng_state, 0.15f * TL_DEFAULT_RED_DURATION, 5.0f * TL_DEFAULT_RED_DURATION); } TrafficLightCycle cycle; @@ -2924,6 +2992,448 @@ static bool check_spawn_offroad(Drive *env, Agent *tmp_agent, float edge_clearan return false; } +// ======================================== +// Scripted Pedestrians +// ======================================== + +// Ray from (qx, qy) along the unit direction against segment AB; t_out = ray distance of the crossing. +static bool ray_hits_segment( + float qx, + float qy, + float dir_x, + float dir_y, + float ax, + float ay, + float bx, + float by, + float *t_out) { + float seg_dx = bx - ax; + float seg_dy = by - ay; + float denom = dir_x * seg_dy - dir_y * seg_dx; + if (fabsf(denom) < 1e-6f) { + return false; + } + float to_a_x = ax - qx; + float to_a_y = ay - qy; + float t = (to_a_x * seg_dy - to_a_y * seg_dx) / denom; + float u = (to_a_x * dir_y - to_a_y * dir_x) / denom; + if (t < 0.0f || u < 0.0f || u > 1.0f) { + return false; + } + *t_out = t; + return true; +} + +// Nearest road-edge crossings of the lateral ray through a lane point: near kerb along +right, far kerb along -right. +static bool find_kerb_distances( + Drive *env, + float qx, + float qy, + float qz, + float right_x, + float right_y, + float *near_m, + float *far_m) { + GridMapEntity entity_list[ROAD_QUERY_ENTITY_COUNT]; + int list_size = get_neighbors_entities(env, qx, qy, entity_list, ROAD_QUERY_ENTITY_COUNT, ROAD_OFFSETS, 25); + *near_m = PEDESTRIAN_KERB_SEARCH_MAX_M; + *far_m = PEDESTRIAN_KERB_SEARCH_MAX_M; + for (int i = 0; i < list_size; i++) { + RoadMapElement *element = &env->road_elements[entity_list[i].entity_idx]; + int geometry_idx = entity_list[i].geometry_idx; + if (!is_road_edge(element->type) || fabsf(element->z[geometry_idx] - qz) > Z_BUFFER) { + continue; + } + float ax = element->x[geometry_idx]; + float ay = element->y[geometry_idx]; + float bx = element->x[geometry_idx + 1]; + float by = element->y[geometry_idx + 1]; + float t; + if (ray_hits_segment(qx, qy, right_x, right_y, ax, ay, bx, by, &t) && t < *near_m) { + *near_m = t; + } + if (ray_hits_segment(qx, qy, -right_x, -right_y, ax, ay, bx, by, &t) && t < *far_m) { + *far_m = t; + } + } + return *near_m < PEDESTRIAN_KERB_SEARCH_MAX_M && *far_m < PEDESTRIAN_KERB_SEARCH_MAX_M; +} + +// Lane point `ahead_m` beyond (lane_idx, lane_s_m) and its road edges; false at dead ends or with an edge out of reach. +static bool find_kerb_site(Drive *env, int lane_idx, float lane_s_m, float ahead_m, KerbSite *site) { + int cursor_idx; + float lane_heading; + if (!route_point_at_distance( + env, + NULL, + 0, + lane_idx, + lane_s_m, + ahead_m, + &site->x, + &site->y, + &site->z, + &site->lane_idx, + &cursor_idx, + &site->lane_s_m, + &lane_heading)) { + return false; + } + site->right_x = sinf(lane_heading); + site->right_y = -cosf(lane_heading); + return find_kerb_distances( + env, + site->x, + site->y, + site->z, + site->right_x, + site->right_y, + &site->right_kerb_m, + &site->left_kerb_m); +} + +// Point `offset_m` outside the road edge of a site; side = +1 right kerb, -1 left kerb. +static void kerb_point(const KerbSite *site, float side, float offset_m, float *x, float *y) { + float kerb_m = side > 0.0f ? site->right_kerb_m : site->left_kerb_m; + *x = site->x + side * site->right_x * (kerb_m + offset_m); + *y = site->y + side * site->right_y * (kerb_m + offset_m); +} + +// Straight walk along the kerb: the chord between two kerb points of one lane, refused where the road bends. +static bool find_kerb_walk( + Drive *env, + ScriptedPedestrian *pedestrian, + const KerbSite *site, + float side, + float kerb_offset_m, + float *start_x, + float *start_y, + float *start_z) { + float walk_m = sample_uniform(&env->rng_state, PEDESTRIAN_KERB_WALK_MIN_M, PEDESTRIAN_KERB_WALK_MAX_M); + KerbSite mid_site, end_site; + if (!find_kerb_site(env, site->lane_idx, site->lane_s_m, 0.5f * walk_m, &mid_site) + || !find_kerb_site(env, site->lane_idx, site->lane_s_m, walk_m, &end_site) + || end_site.lane_idx != site->lane_idx) { + return false; + } + float near_x, near_y, mid_x, mid_y, far_x, far_y; + kerb_point(site, side, kerb_offset_m, &near_x, &near_y); + kerb_point(&mid_site, side, kerb_offset_m, &mid_x, &mid_y); + kerb_point(&end_site, side, kerb_offset_m, &far_x, &far_y); + float bend_m = hypotf(0.5f * (near_x + far_x) - mid_x, 0.5f * (near_y + far_y) - mid_y); + if (bend_m > PEDESTRIAN_KERB_WALK_BEND_TOLERANCE_M) { + return false; + } + float chord_m = hypotf(far_x - near_x, far_y - near_y); + bool against_traffic = sample_uniform(&env->rng_state, 0.0f, 1.0f) < PEDESTRIAN_KERB_WALK_AGAINST_TRAFFIC_PROB; + float travel_sign = against_traffic ? -1.0f : 1.0f; + *start_x = against_traffic ? far_x : near_x; + *start_y = against_traffic ? far_y : near_y; + *start_z = against_traffic ? end_site.z : site->z; + pedestrian->dir_x = travel_sign * (far_x - near_x) / chord_m; + pedestrian->dir_y = travel_sign * (far_y - near_y) / chord_m; + pedestrian->path_m = chord_m; + pedestrian->climb_per_m = travel_sign * (end_site.z - site->z) / chord_m; + return true; +} + +// Crossing from one kerb to past the other, up to PEDESTRIAN_SLOW_SKEW_MAX_RAD off the perpendicular along one lane. +static bool find_slow_crossing( + Drive *env, + ScriptedPedestrian *pedestrian, + const KerbSite *site, + float side, + float kerb_offset_m, + float *start_x, + float *start_y, + float *start_z) { + float skew_rad = sample_uniform(&env->rng_state, 0.0f, PEDESTRIAN_SLOW_SKEW_MAX_RAD); + float along_m = (site->right_kerb_m + site->left_kerb_m) * tanf(skew_rad); + KerbSite ahead_site; + if (!find_kerb_site(env, site->lane_idx, site->lane_s_m, along_m, &ahead_site) + || ahead_site.lane_idx != site->lane_idx) { + return false; + } + bool against_traffic = sample_uniform(&env->rng_state, 0.0f, 1.0f) < PEDESTRIAN_SLOW_AGAINST_TRAFFIC_PROB; + const KerbSite *start_site = against_traffic ? &ahead_site : site; + const KerbSite *end_site = against_traffic ? site : &ahead_site; + float end_x, end_y; + kerb_point(start_site, side, kerb_offset_m, start_x, start_y); + kerb_point(end_site, -side, PEDESTRIAN_FAR_KERB_OVERSHOOT_M, &end_x, &end_y); + *start_z = start_site->z; + pedestrian->path_m = hypotf(end_x - *start_x, end_y - *start_y); + pedestrian->dir_x = (end_x - *start_x) / pedestrian->path_m; + pedestrian->dir_y = (end_y - *start_y) / pedestrian->path_m; + pedestrian->climb_per_m = (end_site->z - start_site->z) / pedestrian->path_m; + return true; +} + +static int sample_pedestrian_mode(Drive *env) { + float draw = sample_uniform(&env->rng_state, 0.0f, 1.0f); + if (draw < PEDESTRIAN_POP_IN_CROSS_PROB) { + return PEDESTRIAN_MODE_POP_IN_CROSS; + } + draw -= PEDESTRIAN_POP_IN_CROSS_PROB; + if (draw < PEDESTRIAN_SLOW_CROSS_PROB) { + return PEDESTRIAN_MODE_SLOW_CROSS; + } + draw -= PEDESTRIAN_SLOW_CROSS_PROB; + if (draw < PEDESTRIAN_STAND_PROB) { + return PEDESTRIAN_MODE_STAND; + } + draw -= PEDESTRIAN_STAND_PROB; + return draw < PEDESTRIAN_KERB_WALK_PROB ? PEDESTRIAN_MODE_KERB_WALK : PEDESTRIAN_MODE_KERB_STEP; +} + +// Start point and leg of the pedestrian's mode at `site`; false where a kerb walk or skewed crossing leaves the lane. +static bool plan_pedestrian_leg( + Drive *env, + ScriptedPedestrian *pedestrian, + const KerbSite *site, + int lanes_to_kerb, + float *start_x, + float *start_y, + float *start_z) { + int mode = pedestrian->mode; + bool is_scenario3 = mode == PEDESTRIAN_MODE_POP_IN_CROSS || mode == PEDESTRIAN_MODE_KERB_STEP; + float side = 1.0f; // +1 right kerb, -1 left kerb + float scenario3_speed_mps = 0.0f; + if (is_scenario3) { + scenario3_speed_mps = PEDESTRIAN_LB1_SPEED_BASE_MPS + PEDESTRIAN_LB1_SPEED_PER_LANE_MPS * (float) lanes_to_kerb + + sample_uniform(&env->rng_state, -PEDESTRIAN_SPEED_JITTER_MPS, PEDESTRIAN_SPEED_JITTER_MPS); + } else if (sample_uniform(&env->rng_state, 0.0f, 1.0f) < PEDESTRIAN_LEFT_KERB_PROB) { + side = -1.0f; + } + float kerb_offset_m = sample_uniform(&env->rng_state, PEDESTRIAN_KERB_OFFSET_MIN_M, PEDESTRIAN_KERB_OFFSET_MAX_M); + kerb_point(site, side, kerb_offset_m, start_x, start_y); + *start_z = site->z; + float into_road_x = -side * site->right_x; + float into_road_y = -side * site->right_y; + float road_crossing_m = kerb_offset_m + site->right_kerb_m + site->left_kerb_m + PEDESTRIAN_FAR_KERB_OVERSHOOT_M; + pedestrian->dir_x = into_road_x; + pedestrian->dir_y = into_road_y; + pedestrian->timer_s = 0.0f; + pedestrian->climb_per_m = 0.0f; + pedestrian->ramp_s = sample_uniform(&env->rng_state, PEDESTRIAN_RAMP_MIN_S, PEDESTRIAN_RAMP_MAX_S); + pedestrian->stand_s = sample_uniform(&env->rng_state, PEDESTRIAN_LINGER_MIN_S, PEDESTRIAN_LINGER_MAX_S); + + if (mode == PEDESTRIAN_MODE_POP_IN_CROSS) { + pedestrian->path_m = road_crossing_m; + pedestrian->speed_mps = scenario3_speed_mps; + pedestrian->timer_s = sample_uniform(&env->rng_state, PEDESTRIAN_WAIT_MIN_S, PEDESTRIAN_WAIT_MAX_S); + return true; + } + if (mode == PEDESTRIAN_MODE_SLOW_CROSS) { + pedestrian->speed_mps + = sample_uniform(&env->rng_state, PEDESTRIAN_SLOW_SPEED_MIN_MPS, PEDESTRIAN_SLOW_SPEED_MAX_MPS); + pedestrian->timer_s = sample_uniform(&env->rng_state, 0.0f, PEDESTRIAN_SLOW_WAIT_MAX_S); + return find_slow_crossing(env, pedestrian, site, side, kerb_offset_m, start_x, start_y, start_z); + } + if (mode == PEDESTRIAN_MODE_STAND) { + float facing_rad = sample_uniform(&env->rng_state, -(float) M_PI, (float) M_PI); + pedestrian->dir_x = cosf(facing_rad); + pedestrian->dir_y = sinf(facing_rad); + pedestrian->path_m = 0.0f; + pedestrian->speed_mps = 0.0f; + pedestrian->stand_s + = sample_uniform(&env->rng_state, PEDESTRIAN_BYSTANDER_STAND_MIN_S, PEDESTRIAN_BYSTANDER_STAND_MAX_S); + return true; + } + if (mode == PEDESTRIAN_MODE_KERB_STEP) { + float step_m = sample_uniform(&env->rng_state, PEDESTRIAN_KERB_STEP_MIN_M, PEDESTRIAN_KERB_STEP_MAX_M); + if (sample_uniform(&env->rng_state, 0.0f, 1.0f) < PEDESTRIAN_KERB_STEP_AWAY_PROB) { + pedestrian->dir_x = -into_road_x; + pedestrian->dir_y = -into_road_y; + } else { + step_m = fminf(step_m, fmaxf(0.0f, kerb_offset_m - PEDESTRIAN_KERB_STEP_EDGE_CLEARANCE_M)); + } + pedestrian->path_m = step_m; + pedestrian->speed_mps = scenario3_speed_mps; + pedestrian->timer_s = sample_uniform(&env->rng_state, PEDESTRIAN_WAIT_MIN_S, PEDESTRIAN_WAIT_MAX_S); + pedestrian->stand_s + = sample_uniform(&env->rng_state, PEDESTRIAN_BYSTANDER_STAND_MIN_S, PEDESTRIAN_BYSTANDER_STAND_MAX_S); + return true; + } + pedestrian->speed_mps + = sample_uniform(&env->rng_state, PEDESTRIAN_KERB_WALK_SPEED_MIN_MPS, PEDESTRIAN_KERB_WALK_SPEED_MAX_MPS); + return find_kerb_walk(env, pedestrian, site, side, kerb_offset_m, start_x, start_y, start_z); +} + +static bool place_scripted_pedestrian(Drive *env, ScriptedPedestrian *pedestrian, const Agent *car) { + RoadMapElement *car_lane = &env->road_elements[car->current_lane_idx]; + float car_s_m + = compute_lane_progress(car_lane, car->sim_x, car->sim_y, car->cos_heading, car->sin_heading, true, NULL); + KerbSite car_site, site; + if (!find_kerb_site(env, car->current_lane_idx, car_s_m, 0.0f, &car_site)) { + return false; + } + int mode = pedestrian->mode; + int lanes_to_kerb = (int) (car_site.right_kerb_m / PEDESTRIAN_LB1_LANE_WIDTH_M + 1.0f); + float ahead_m; + if (mode == PEDESTRIAN_MODE_POP_IN_CROSS || mode == PEDESTRIAN_MODE_KERB_STEP) { + ahead_m = PEDESTRIAN_LB1_APPEAR_BASE_M + PEDESTRIAN_LB1_APPEAR_PER_LANE_M * (float) (lanes_to_kerb - 1) + + sample_uniform(&env->rng_state, -PEDESTRIAN_APPEAR_JITTER_M, PEDESTRIAN_APPEAR_JITTER_M); + } else if (mode == PEDESTRIAN_MODE_SLOW_CROSS) { + ahead_m = sample_uniform(&env->rng_state, PEDESTRIAN_SLOW_APPEAR_MIN_M, PEDESTRIAN_SLOW_APPEAR_MAX_M); + } else { + ahead_m = sample_uniform(&env->rng_state, PEDESTRIAN_FAR_APPEAR_MIN_M, PEDESTRIAN_FAR_APPEAR_MAX_M); + } + if (!find_kerb_site(env, car->current_lane_idx, car_s_m, ahead_m, &site)) { + return false; + } + if (env->num_speed_zones > 0 && env->road_elements[site.lane_idx].speed_zone_idx < 0) { + return false; // junction lane + } + float start_x, start_y, start_z; + if (!plan_pedestrian_leg(env, pedestrian, &site, lanes_to_kerb, &start_x, &start_y, &start_z)) { + return false; + } + + Agent *agent = &env->agents[pedestrian->agent_idx]; + agent->sim_length = sample_uniform(&env->rng_state, PEDESTRIAN_SIZE_MIN_M, PEDESTRIAN_SIZE_MAX_M); + agent->sim_width = sample_uniform(&env->rng_state, PEDESTRIAN_SIZE_MIN_M, PEDESTRIAN_SIZE_MAX_M); + agent->sim_height = PEDESTRIAN_HEIGHT_M; + update_agent_radius(agent); + + // The whole leg as one box: nothing may stand on it when the pedestrian appears + Agent path_probe = *agent; + path_probe.sim_x = start_x + pedestrian->dir_x * 0.5f * pedestrian->path_m; + path_probe.sim_y = start_y + pedestrian->dir_y * 0.5f * pedestrian->path_m; + path_probe.sim_z = start_z; + path_probe.sim_heading = atan2f(pedestrian->dir_y, pedestrian->dir_x); + path_probe.cos_heading = pedestrian->dir_x; + path_probe.sin_heading = pedestrian->dir_y; + path_probe.sim_length = fmaxf(pedestrian->path_m, agent->sim_length); + update_agent_radius(&path_probe); + copy_pose_to_prev(&path_probe); + if (check_spawn_collision(env, &path_probe)) { + return false; + } + + agent->sim_x = start_x; + agent->sim_y = start_y; + agent->sim_z = start_z; + agent->sim_heading = path_probe.sim_heading; + agent->cos_heading = pedestrian->dir_x; + agent->sin_heading = pedestrian->dir_y; + copy_pose_to_prev(agent); + clear_agent_motion(agent); + agent->seconds_stopped = 0.0f; + agent->sim_valid = 1; + pedestrian->remaining_m = pedestrian->path_m; + pedestrian->current_speed_mps = 0.0f; + pedestrian->state = PEDESTRIAN_STATE_WALKING; + if (mode == PEDESTRIAN_MODE_STAND) { + // a stander has been there for a while + agent->seconds_stopped = sample_uniform(&env->rng_state, 0.0f, MAX_STOPPED_SECONDS); + pedestrian->state = PEDESTRIAN_STATE_STANDING; + pedestrian->timer_s = pedestrian->stand_s; + } + if (mode == PEDESTRIAN_MODE_KERB_WALK) { + // already walking when it appears + pedestrian->current_speed_mps = pedestrian->speed_mps; + agent->sim_vx = pedestrian->dir_x * pedestrian->speed_mps; + agent->sim_vy = pedestrian->dir_y * pedestrian->speed_mps; + update_agent_speed(agent); + } + return true; +} + +static bool arm_scripted_pedestrian(Drive *env, ScriptedPedestrian *pedestrian) { + if (env->active_agent_count == 0) { + return false; + } + bool enters_road + = pedestrian->mode == PEDESTRIAN_MODE_POP_IN_CROSS || pedestrian->mode == PEDESTRIAN_MODE_SLOW_CROSS; + for (int attempt = 0; attempt < PEDESTRIAN_ARM_ATTEMPTS; attempt++) { + int target_slot = rng_below(&env->rng_state, env->active_agent_count); + int target_agent_idx = env->active_agent_indices[target_slot]; + Agent *car = &env->agents[target_agent_idx]; + bool car_eligible = car->sim_valid == 1 && !car->stopped && !car->removed && car->current_lane_idx != -1 + && car->sim_speed >= PEDESTRIAN_TARGET_MIN_SPEED_MPS + && car->metrics_array[LANE_ANGLE_IDX] >= LANE_ALIGN_COS_THRESHOLD; + if (!car_eligible || !place_scripted_pedestrian(env, pedestrian, car)) { + continue; + } + pedestrian->target_agent_idx = target_agent_idx; + env->logs[target_slot].pedestrian_crossings += enters_road ? 1.0f : 0.0f; + env->logs[target_slot].pedestrian_bystanders += enters_road ? 0.0f : 1.0f; + return true; + } + return false; +} + +static void idle_scripted_pedestrian(Drive *env, ScriptedPedestrian *pedestrian) { + Agent *agent = &env->agents[pedestrian->agent_idx]; + invalidate_agent(agent); + copy_pose_to_prev(agent); + pedestrian->target_agent_idx = -1; + pedestrian->state = PEDESTRIAN_STATE_IDLE; + pedestrian->mode = sample_pedestrian_mode(env); + pedestrian->timer_s = sample_uniform(&env->rng_state, PEDESTRIAN_COOLDOWN_MIN_S, PEDESTRIAN_COOLDOWN_MAX_S); + pedestrian->remaining_m = 0.0f; +} + +static void reset_scripted_pedestrians(Drive *env) { + for (int k = 0; k < env->scripted_pedestrian_slots; k++) { + idle_scripted_pedestrian(env, &env->pedestrians[k]); + } +} + +static void step_pedestrian_idle(Drive *env, ScriptedPedestrian *pedestrian) { + pedestrian->timer_s -= env->dt; + if (pedestrian->timer_s > 0.0f) { + return; + } + arm_scripted_pedestrian(env, pedestrian); // a failed attempt retries next step +} + +static void step_pedestrian_walking(Drive *env, ScriptedPedestrian *pedestrian, Agent *agent) { + if (pedestrian->timer_s > 0.0f) { + pedestrian->timer_s -= env->dt; // standing where it appeared + return; + } + float ramp_step_mps = pedestrian->speed_mps * env->dt / pedestrian->ramp_s; + pedestrian->current_speed_mps = fminf(pedestrian->speed_mps, pedestrian->current_speed_mps + ramp_step_mps); + agent->sim_vx = pedestrian->dir_x * pedestrian->current_speed_mps; + agent->sim_vy = pedestrian->dir_y * pedestrian->current_speed_mps; + update_agent_speed(agent); + float step_m = fminf(pedestrian->current_speed_mps * env->dt, pedestrian->remaining_m); + agent->sim_x += pedestrian->dir_x * step_m; + agent->sim_y += pedestrian->dir_y * step_m; + agent->sim_z += pedestrian->climb_per_m * step_m; + pedestrian->remaining_m -= step_m; + if (pedestrian->remaining_m > 0.0f) { + return; + } + clear_agent_motion(agent); + pedestrian->state = PEDESTRIAN_STATE_STANDING; + pedestrian->timer_s = pedestrian->stand_s; +} + +static void step_pedestrian_standing(Drive *env, ScriptedPedestrian *pedestrian) { + pedestrian->timer_s -= env->dt; + if (pedestrian->timer_s > 0.0f) { + return; + } + idle_scripted_pedestrian(env, pedestrian); +} + +static void move_scripted_pedestrians(Drive *env) { + for (int k = 0; k < env->scripted_pedestrian_slots; k++) { + ScriptedPedestrian *pedestrian = &env->pedestrians[k]; + Agent *agent = &env->agents[pedestrian->agent_idx]; + copy_pose_to_prev(agent); + if (pedestrian->state == PEDESTRIAN_STATE_WALKING) { + step_pedestrian_walking(env, pedestrian, agent); + } else if (pedestrian->state == PEDESTRIAN_STATE_STANDING) { + step_pedestrian_standing(env, pedestrian); + } else { + step_pedestrian_idle(env, pedestrian); + } + } +} + static bool spawn_agent(Drive *env, int agent_idx, int num_agents) { Agent *agent = &env->agents[agent_idx]; for (int reason = 0; reason < SPAWN_REJECT_REASON_COUNT; reason++) { @@ -3039,7 +3549,10 @@ static bool spawn_agent(Drive *env, int agent_idx, int num_agents) { continue; } - if (check_agent_on_stop_line(env, &tmp_agent, true)) { + // the detector's own footprint test, on yellow too, so no spawn starts on a line it would have to clear + Agent stop_line_probe = tmp_agent; + stop_line_probe.red_light_latched_idx = -1; + if (check_red_light_footprint(env, &stop_line_probe, true)) { env->spawn_reject_counts[SPAWN_REJECT_STOP_LINE]++; continue; } @@ -3219,6 +3732,9 @@ static void set_start_position(Drive *env) { int i = is_active ? env->active_agent_indices[list_idx] : env->static_agent_indices[list_idx - env->active_agent_count]; Agent *agent = &env->agents[i]; + if (env->simulation_mode == SIMULATION_MODE_GIGAFLOW && agent->removed) { + continue; // failed init spawn: reset_agent_state would clear the removed latch + } // Initialize simulation trajectory from logged trajectory at init_step if (env->simulation_mode == SIMULATION_MODE_REPLAY) { @@ -3350,26 +3866,30 @@ void set_active_agents(Drive *env) { if (env->simulation_mode == SIMULATION_MODE_GIGAFLOW) { int num_agents_to_create = env->num_controllable_agents; int partner_slots = env->cosim_partner_slots; + int pedestrian_slots = env->scripted_pedestrian_slots; + int static_slots = partner_slots + pedestrian_slots; - // Initialize agents for GIGAFLOW mode - env->agents = (Agent *) calloc(num_agents_to_create + partner_slots, sizeof(Agent)); - env->num_total_agents = num_agents_to_create + partner_slots; - int *active_agent_indices = (int *) malloc(num_agents_to_create * sizeof(int)); - + // Initialize agents for GIGAFLOW mode; the bin's logged agents are not used + for (int i = 0; i < env->num_total_agents; i++) { + free_agent(&env->agents[i]); + } + free(env->agents); + env->agents = (Agent *) calloc(num_agents_to_create + static_slots, sizeof(Agent)); + env->num_total_agents = num_agents_to_create + static_slots; int successfully_created = 0; for (int i = 0; i < num_agents_to_create; i++) { if (spawn_agent(env, i, i)) { - active_agent_indices[successfully_created] = i; successfully_created++; } else { - // Failed spawn: ensure agent is properly invalidated + // A failed spawn keeps its policy slot, removed and masked until the next c_reset respawns it invalidate_agent(&env->agents[i]); + reset_agent_state(&env->agents[i]); env->agents[i].removed = 1; } } - env->active_agent_indices = (int *) malloc(successfully_created * sizeof(int)); - env->static_agent_indices = partner_slots > 0 ? (int *) malloc(partner_slots * sizeof(int)) : NULL; + env->active_agent_indices = (int *) malloc(num_agents_to_create * sizeof(int)); + env->static_agent_indices = static_slots > 0 ? (int *) malloc(static_slots * sizeof(int)) : NULL; env->expert_static_agent_indices = NULL; for (int k = 0; k < partner_slots; k++) { int partner_idx = num_agents_to_create + k; @@ -3385,14 +3905,29 @@ void set_active_agents(Drive *env) { invalidate_agent(partner); env->static_agent_indices[k] = partner_idx; } - env->static_agent_count = partner_slots; + env->pedestrians + = pedestrian_slots > 0 ? (ScriptedPedestrian *) calloc(pedestrian_slots, sizeof(ScriptedPedestrian)) : NULL; + for (int k = 0; k < pedestrian_slots; k++) { + int pedestrian_idx = num_agents_to_create + partner_slots + k; + Agent *pedestrian = &env->agents[pedestrian_idx]; + pedestrian->id = pedestrian_idx; + pedestrian->type = PEDESTRIAN; + pedestrian->controller = CONTROLLER_STATIC; + pedestrian->sim_length = PEDESTRIAN_SIZE_MIN_M; + pedestrian->sim_width = PEDESTRIAN_SIZE_MIN_M; + pedestrian->sim_height = PEDESTRIAN_HEIGHT_M; + update_agent_radius(pedestrian); + invalidate_agent(pedestrian); + env->static_agent_indices[partner_slots + k] = pedestrian_idx; + env->pedestrians[k].agent_idx = pedestrian_idx; + } + reset_scripted_pedestrians(env); + env->static_agent_count = static_slots; - for (int i = 0; i < successfully_created; i++) { - env->active_agent_indices[i] = active_agent_indices[i]; - env->agents[active_agent_indices[i]].controller - = resolve_agent_controller(env, active_agent_indices[i], 1, 0); + for (int i = 0; i < num_agents_to_create; i++) { + env->active_agent_indices[i] = i; + env->agents[i].controller = resolve_agent_controller(env, i, 1, 0); } - free(active_agent_indices); if (successfully_created == 0) { fprintf( @@ -3402,8 +3937,8 @@ void set_active_agents(Drive *env) { num_agents_to_create, env->map_name); } - env->active_agent_count = successfully_created; - env->num_agents = successfully_created + partner_slots; + env->active_agent_count = num_agents_to_create; + env->num_agents = num_agents_to_create + static_slots; return; } @@ -3618,7 +4153,8 @@ void remove_bad_trajectories(Drive *env) { env->timestep = 0; } -void init(Drive *env) { +// 0 on success, -1 when the map cannot be loaded or is unusable (reason on stderr); the env must not be used then. +int init(Drive *env) { env->timestep = 0; env->first_reset_pending = 1; struct SharedMapData *shared = env->use_map_cache ? map_cache_lookup(env) : NULL; @@ -3627,7 +4163,7 @@ void init(Drive *env) { // then discard the freshly-loaded geometry and borrow the shared copy. if (load_map_binary(env->map_name, env) != 0) { fprintf(stderr, "[ERROR] -> Failed to load map binary: %s\n", env->map_name); - return; + return -1; } for (int i = 0; i < env->num_road_elements; i++) { free_road_element(&env->road_elements[i]); @@ -3645,11 +4181,11 @@ void init(Drive *env) { // Cache miss (or caching off): load and build the geometry as usual. if (load_map_binary(env->map_name, env) != 0) { fprintf(stderr, "[ERROR] -> Failed to load map binary: %s\n", env->map_name); - return; + return -1; } if (init_grid_map(env) != 0) { fprintf(stderr, "[ERROR] -> Failed to build grid map for map: %s\n", env->map_name); - return; + return -1; } env->grid_map->vision_range = compute_vision_range(env); init_neighbor_offsets(env); @@ -3661,7 +4197,7 @@ void init(Drive *env) { } if (env->simulation_mode == SIMULATION_MODE_REPLAY && (!isfinite(env->log_dt) || env->log_dt <= 0.0f)) { fprintf(stderr, "[ERROR] -> Replay map %s has invalid log_dt %f\n", env->map_name, (double) env->log_dt); - return; + return -1; } if (env->use_neighbor_cache && env->grid_map->neighbor_cache_pool_idx == NULL) { cache_neighbor_offsets(env); @@ -3733,6 +4269,7 @@ void init(Drive *env) { generate_new_goals_from_route(env, agent); } } + return 0; } void c_close(Drive *env) { @@ -3744,6 +4281,7 @@ void c_close(Drive *env) { free(env->obs_neighbor_scratch); free(env->static_agent_indices); free(env->expert_static_agent_indices); + free(env->pedestrians); free_loaded_map_data(env); } @@ -3755,15 +4293,22 @@ static inline int lane_feature_count(const Drive *env) { return LANE_FEATURES + (env->obs_lane_speed_limit ? LANE_SPEED_LIMIT_FEATURES : 0); } +static inline int traffic_control_feature_count(const Drive *env) { + return TRAFFIC_CONTROL_FEATURES + + (env->obs_traffic_light_yellow_remaining ? TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES : 0); +} + static int compute_observation_size(Drive *env) { return EGO_FEATURES + partner_feature_count(env) * env->obs_slots_partners_n + lane_feature_count(env) * env->obs_slots_lane_kept + BOUNDARY_FEATURES * env->obs_slots_boundary_kept - + TRAFFIC_CONTROL_FEATURES * env->obs_slots_traffic_controls_n + OBS_VALID_COUNT_FEATURES + + traffic_control_feature_count(env) * env->obs_slots_traffic_controls_n + OBS_VALID_COUNT_FEATURES + env->reward_conditioning * NUM_REWARD_COEFS + env->num_goals * GOAL_FEATURES; } void allocate(Drive *env) { - init(env); + if (init(env) != 0) { + raise_error_with_message(ERROR_INITIALIZATION_FAILED, "init failed for map %s", env->map_name); + } int max_obs = compute_observation_size(env); env->observations = (float *) calloc(env->active_agent_count * max_obs, sizeof(float)); env->actions = (float *) calloc(env->active_agent_count * 2, sizeof(float)); @@ -3931,6 +4476,29 @@ static bool find_closest_lane_segment( return true; } +// Road edges listed by one cell's 5x5 window against the agent's swept box. +static bool crosses_road_edge_in_window(Drive *env, Agent *agent, int cell_idx) { + GridMap *grid = env->grid_map; + const LaneQueryEntry *entries = &grid->lane_query_pool[grid->lane_query_offset[cell_idx]]; + int entry_count = grid->lane_query_offset[cell_idx + 1] - grid->lane_query_offset[cell_idx]; + for (int i = 0; i < entry_count; i++) { + RoadMapElement *element = &env->road_elements[entries[i].entity_idx]; + int geometry_idx = entries[i].seg_first; + if (!is_road_edge(element->type) || fabsf(element->z[geometry_idx] - agent->sim_z) > Z_BUFFER) { + continue; + } + if (check_segment_crosses_moving_box( + element->x[geometry_idx], + element->y[geometry_idx], + element->x[geometry_idx + 1], + element->y[geometry_idx + 1], + agent)) { + return true; + } + } + return false; +} + static void find_lane_and_offroad( Drive *env, Agent *agent, @@ -4053,6 +4621,14 @@ static void find_lane_and_offroad( } } + // A step longer than the window margin sweeps road edges that only the previous cell's window lists + float step_m = compute_euclidean_distance(agent->prev_x, agent->prev_y, agent->sim_x, agent->sim_y); + if (!is_offroad && step_m + agent->radius > window_margin_m) { + int prev_cell_idx = get_grid_index(env, agent->prev_x, agent->prev_y); + is_offroad = prev_cell_idx != -1 && prev_cell_idx != cell_idx + && crosses_road_edge_in_window(env, agent, prev_cell_idx); + } + *is_offroad_out = is_offroad; *lane_idx_out = lane_idx; *lane_seg_idx_out = lane_seg_idx; @@ -4197,8 +4773,13 @@ int c_set_agent_sizes(Drive *env, int count, const int *idx, const float *length return 0; } -int c_set_traffic_light_states(Drive *env, const int *states) { +// yellow_remaining_s (optional, seconds per element) keeps a yellow schedule yellow through ts + ceil(s / dt) - 1, then +// red +int c_set_traffic_light_states(Drive *env, const int *states, const float *yellow_remaining_s) { int ts = env->timestep; + if (env->obs_traffic_light_yellow_remaining && yellow_remaining_s == NULL) { + return -1; + } for (int i = 0; i < env->num_traffic_elements; i++) { TrafficControlElement *traffic = &env->traffic_elements[i]; if (traffic->type != TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT) { @@ -4213,11 +4794,30 @@ int c_set_traffic_light_states(Drive *env, const int *states) { if (states[i] < TRAFFIC_CONTROL_STATE_UNKNOWN || states[i] > TRAFFIC_CONTROL_STATE_OFF) { return -1; } + if (yellow_remaining_s != NULL + && (!isfinite(yellow_remaining_s[i]) || yellow_remaining_s[i] < 0.0f + || yellow_remaining_s[i] > env->traffic_light_yellow_max_s)) { + 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]; } + if (yellow_remaining_s == NULL || states[i] != TRAFFIC_CONTROL_STATE_YELLOW) { + continue; + } + int yellow_steps = (int) ceilf(yellow_remaining_s[i] / env->dt); + if (yellow_steps < 2) { + yellow_steps = 2; // the held step stays yellow + } + int yellow_end = ts + yellow_steps < traffic->state_size ? ts + yellow_steps : traffic->state_size; + for (int t = ts + 2; t < yellow_end; t++) { + traffic->states[t] = TRAFFIC_CONTROL_STATE_YELLOW; + } + if (ts + yellow_steps < traffic->state_size) { + traffic->states[ts + yellow_steps] = TRAFFIC_CONTROL_STATE_RED; + } } return 0; } @@ -4275,6 +4875,7 @@ int c_set_stop_signs(Drive *env, int count, const float *lines, const float *hea env->agents[i].stop_sign_stop_completed = 0; env->agents[i].stop_sign_last_failed_idx = -1; env->agents[i].stop_sign_standstill_idx = -1; + env->agents[i].red_light_latched_idx = -1; } return 0; } @@ -4320,6 +4921,24 @@ void c_get_agent_goal_progress(Drive *env, int agent_idx, int *current_goal_idx_ *goal_count_out = agent->goal_count; } +// Co-sim: rewind or advance the consumed-goal count of the current window (undoes a forecast rollout's progress). +int c_set_agent_goal_progress(Drive *env, int agent_idx, int current_goal_idx) { + if (agent_idx < 0 || agent_idx >= env->num_total_agents) { + return -1; + } + Agent *agent = &env->agents[agent_idx]; + if (current_goal_idx < 0 || current_goal_idx > agent->goal_count) { + return -1; + } + agent->current_goal_idx = current_goal_idx; + if (current_goal_idx < agent->goal_count) { + agent->current_goal_x = agent->list_goal_x[current_goal_idx]; + agent->current_goal_y = agent->list_goal_y[current_goal_idx]; + agent->current_goal_z = agent->list_goal_z[current_goal_idx]; + } + return 0; +} + int c_set_agent_goals( Drive *env, int agent_idx, @@ -4351,12 +4970,35 @@ int c_set_agent_goals( agent->current_goal_x = agent->list_goal_x[0]; agent->current_goal_y = agent->list_goal_y[0]; agent->current_goal_z = agent->list_goal_z[0]; + agent->lane_goal_spacing_m = 0.0f; + return 0; +} + +// Co-sim forecast agents have no known route: their goals follow the lane ahead and c_step renews them when reached. +int c_set_agent_lane_goals(Drive *env, int count, const int *idx, float spacing_m, int *placed_goal_counts_out) { + if (!isfinite(spacing_m) || spacing_m <= 0.0f) { + return -1; + } + for (int k = 0; k < count; k++) { + if (idx[k] < 0 || idx[k] >= env->num_total_agents || !env->agents[idx[k]].sim_valid) { + return -1; + } + } + for (int k = 0; k < count; k++) { + Agent *agent = &env->agents[idx[k]]; + agent->lane_goal_spacing_m = spacing_m; + agent->goal_count = 0; + agent->current_goal_idx = 0; + placed_goal_counts_out[k] = place_lane_goals_ahead(env, agent, spacing_m); + } return 0; } // ─────────────────────────────────────── +// Rows are num_timesteps wide starting at init_step; steps beyond the log read invalid. void c_get_global_ground_truth_trajectories( Drive *env, + int num_timesteps, float *x_out, float *y_out, float *z_out, @@ -4370,8 +5012,13 @@ void c_get_global_ground_truth_trajectories( id_out[i] = get_track_id_or_placeholder(env, agent_idx); scenario_id_out[i] = 0; // TODO: FIXME - for (int t = env->init_step; t < agent->trajectory_size; t++) { - int out_idx = i * (agent->trajectory_size - env->init_step) + (t - env->init_step); + for (int step_idx = 0; step_idx < num_timesteps; step_idx++) { + int t = env->init_step + step_idx; + int out_idx = i * num_timesteps + step_idx; + if (t >= agent->trajectory_size) { + valid_out[out_idx] = 0; + continue; + } // Add world means back to get original world coordinates x_out[out_idx] = agent->log_trajectory_x[t] + env->world_mean_x; y_out[out_idx] = agent->log_trajectory_y[t] + env->world_mean_y; @@ -4635,6 +5282,8 @@ static void compute_rewards(Drive *env, int i) { float reward_collision = -(agent->reward_coefs[REWARD_COEF_COLLISION] + 0.1f * agent->sim_speed); env->rewards[i] += reward_collision; agent_log->collision_rate = 1.0f; + agent_log->pedestrian_collision_rate + = env->agents[agent->first_collision_partner_idx].type == PEDESTRIAN ? 1.0f : 0.0f; agent_log->reward_collision += reward_collision; } @@ -5045,6 +5694,8 @@ static int write_road_obs_rows( int is_lane = neighbor_cursor->obs_kind == OBS_ENTITY_LANE; int segment_features = is_lane ? lane_feature_count(env) : BOUNDARY_FEATURES; int rows_written = 0; + float width_clip_min_m = env->obs_lane_width_clip_max_m > 0.0f ? env->obs_lane_width_clip_min_m : 0.0f; + float width_clip_max_m = env->obs_lane_width_clip_max_m > 0.0f ? env->obs_lane_width_clip_max_m : INFINITY; for (const GridMapEntity *entity = neighbor_cursor_next(env, neighbor_cursor); entity != NULL && rows_written < slot_cap; @@ -5073,7 +5724,7 @@ static int write_road_obs_rows( row[5] = rel_seg_dir_y; // Goal-distance features: absolute and relative to ego's lane->goal distance. if (is_lane) { - row[6] = segment->width_m / env->obs_norm_road_seg_width_m; + row[6] = clip(segment->width_m, width_clip_min_m, width_clip_max_m) / env->obs_norm_road_seg_width_m; float lane_dist_m = goal_graph_idx >= 0 ? lane_graph_distance_to_goal_m(env, entity_idx, goal_graph_idx) : -1.0f; write_lane_goal_distance_obs(lane_dist_m, ego_dist_to_goal_m, &row[7]); @@ -5177,6 +5828,19 @@ static int write_road_obs(Drive *env, Agent *ego, float *obs, int obs_idx, int * return obs_idx; } +// Consecutive yellow steps from timestep on, at most max_steps (the schedule's end truncates) +static int yellow_remaining_steps(const TrafficControlElement *tc, int timestep, int max_steps) { + int scan_end = timestep + max_steps < tc->state_size ? timestep + max_steps : tc->state_size; + int remaining_steps = 0; + for (int t = timestep; t < scan_end; t++) { + if (tc->states[t] != TRAFFIC_CONTROL_STATE_YELLOW) { + break; + } + remaining_steps++; + } + return remaining_steps; +} + static int write_traffic_control_obs(Drive *env, Agent *ego, float *obs, int obs_idx, int *traffic_control_count) { typedef struct { int idx; @@ -5226,6 +5890,7 @@ static int write_traffic_control_obs(Drive *env, Agent *ego, float *obs, int obs } } + int yellow_scan_steps = (int) ceilf(env->traffic_light_yellow_max_s / env->dt); int controls_written = 0; for (int j = 0; j < controls_to_observe && controls_written < env->obs_slots_traffic_controls_n; j++) { TrafficControlElement *tc = &env->traffic_elements[visible_controls[j].idx]; @@ -5248,13 +5913,20 @@ static int write_traffic_control_obs(Drive *env, Agent *ego, float *obs, int obs obs[obs_idx++] = rel_x2 / env->obs_norm_xy_offset_m; obs[obs_idx++] = rel_y2 / env->obs_norm_xy_offset_m; obs[obs_idx++] = rel_z / env->obs_norm_z_m; + if (env->obs_traffic_light_yellow_remaining) { + float yellow_remaining_s = light_state == TRAFFIC_CONTROL_STATE_YELLOW + ? yellow_remaining_steps(tc, env->timestep, yellow_scan_steps) * env->dt + : 0.0f; + obs[obs_idx++] + = fminf(yellow_remaining_s, env->traffic_light_yellow_max_s) / env->traffic_light_yellow_max_s; + } obs[obs_idx++] = tc->type; obs[obs_idx++] = light_state; controls_written++; } *traffic_control_count = controls_written; - return obs_idx + (env->obs_slots_traffic_controls_n - controls_written) * TRAFFIC_CONTROL_FEATURES; + return obs_idx + (env->obs_slots_traffic_controls_n - controls_written) * traffic_control_feature_count(env); } static void compute_observations(Drive *env) { @@ -5569,6 +6241,19 @@ static void update_rollout_masks(Drive *env) { } } +// c_reset scores no goals: one inside the radius at spawn stays for the first step to consume, reward and renew +static void restore_goal_reached_at_reset(Agent *agent, Log *agent_log) { + if (agent->metrics_array[REACHED_GOAL_IDX] == 0.0f) { + return; + } + agent->metrics_array[REACHED_GOAL_IDX] = 0.0f; + agent_log->num_goals_reached -= 1; + agent->current_goal_idx--; + agent->current_goal_x = agent->list_goal_x[agent->current_goal_idx]; + agent->current_goal_y = agent->list_goal_y[agent->current_goal_idx]; + agent->current_goal_z = agent->list_goal_z[agent->current_goal_idx]; +} + void c_reset(Drive *env) { env->autoreset_pending = 0; env->episode_end_timestep = env->scenario_length; @@ -5589,6 +6274,7 @@ void c_reset(Drive *env) { int agent_idx = env->active_agent_indices[x]; sample_erratic_flags(env, &env->agents[agent_idx]); compute_metrics(env, agent_idx, x); + restore_goal_reached_at_reset(&env->agents[agent_idx], &env->logs[x]); } update_rollout_masks(env); compute_observations(env); @@ -5605,6 +6291,7 @@ void c_reset(Drive *env) { for (int x = 0; x < env->active_agent_count; x++) { invalidate_agent(&env->agents[env->active_agent_indices[x]]); } + reset_scripted_pedestrians(env); int num_reset = 0; for (int x = 0; x < env->active_agent_count; x++) { int agent_idx = env->active_agent_indices[x]; @@ -5640,6 +6327,7 @@ void c_reset(Drive *env) { sample_erratic_flags(env, agent); generate_reward_coefs(env, agent); compute_metrics(env, agent_idx, x); + restore_goal_reached_at_reset(agent, &env->logs[x]); env->logs[x].stopped_at_reset = agent->stopped ? 1.0f : 0.0f; } update_rollout_masks(env); @@ -5668,6 +6356,7 @@ void c_reset(Drive *env) { generate_new_goals_from_route(env, agent); } compute_metrics(env, agent_idx, x); + restore_goal_reached_at_reset(agent, &env->logs[x]); } update_rollout_masks(env); compute_observations(env); @@ -5750,6 +6439,7 @@ void c_step(Drive *env) { move_expert(env, background_idx); } } + move_scripted_pedestrians(env); // Move active agents with policy actions for (int i = 0; i < env->active_agent_count; i++) { env->logs[i].score = 0.0f; @@ -5854,6 +6544,10 @@ void c_step(Drive *env) { if (agent->metrics_array[REACHED_GOAL_IDX] == 0.0f) { continue; } + if (env->goal_source == GOAL_SOURCE_EXTERNAL && agent->lane_goal_spacing_m > 0.0f) { + place_lane_goals_ahead(env, agent, agent->lane_goal_spacing_m); // a dead end keeps the goals still ahead + continue; + } if (env->goal_source == GOAL_SOURCE_GT || env->goal_source == GOAL_SOURCE_GT_MAP || env->goal_source == GOAL_SOURCE_EXTERNAL) { // Replay/co-sim: leave current_goal_idx saturated so the reached-goal condition won't fire diff --git a/pufferlib/ocean/drive/drive.py b/pufferlib/ocean/drive/drive.py index 0992f5980f..cb060dd66f 100644 --- a/pufferlib/ocean/drive/drive.py +++ b/pufferlib/ocean/drive/drive.py @@ -1,3 +1,4 @@ +import math import pickle import zlib import numpy as np @@ -47,6 +48,9 @@ def __init__( disable_red_light_infractions=False, disable_stop_sign_infractions=False, traffic_light_junction_phases=False, + traffic_light_green_min_s=3.0, + traffic_light_yellow_min_s=2.0, + traffic_light_yellow_max_s=3.0, stop_sign_behavior="ignore", use_map_cache=False, preload_map_cache=False, @@ -79,6 +83,7 @@ def __init__( max_agents_per_env=64, cosim_partner_slots=0, cosim_eval_semantics=False, + scripted_pedestrian_ratio=0.0, action_type="discrete", dynamics_model="classic", reset_accel_on_stop=False, @@ -124,6 +129,7 @@ def __init__( obs_lane_heading_signed=False, obs_lane_speed_limit=False, obs_slots_traffic_controls_n=4, + obs_traffic_light_yellow_remaining=False, traffic_lights_enabled=True, stop_signs_enabled=False, yield_signs_enabled=False, @@ -137,7 +143,11 @@ def __init__( obs_norm_road_seg_width_m=5.0, obs_norm_z_m=10.0, eval_perceived_size_margin_m=0.1, + red_light_train_front_margin_m=0.3, + red_light_train_rear_margin_m=0.0, obs_partner_split_length_m=0.0, + obs_lane_width_clip_min_m=0.0, + obs_lane_width_clip_max_m=0.0, eval_standstill_jerk_deadband_mps3=0.0, obs_range_traffic_control_m=100.0, obs_range_partner_m=100.0, @@ -252,6 +262,17 @@ def __init__( self.disable_red_light_infractions = bool(disable_red_light_infractions) self.disable_stop_sign_infractions = bool(disable_stop_sign_infractions) self.traffic_light_junction_phases = bool(traffic_light_junction_phases) + self.traffic_light_green_min_s = float(traffic_light_green_min_s) + self.traffic_light_yellow_min_s = float(traffic_light_yellow_min_s) + self.traffic_light_yellow_max_s = float(traffic_light_yellow_max_s) + if not (0.0 < self.traffic_light_green_min_s <= 10.0) or not ( + 0.0 < self.traffic_light_yellow_min_s <= self.traffic_light_yellow_max_s < float("inf") + ): + raise ValueError( + "traffic light cycle: need 0 < traffic_light_green_min_s <= 10 and " + f"0 < yellow_min <= yellow_max. Got: {traffic_light_green_min_s}, {traffic_light_yellow_min_s}, " + f"{traffic_light_yellow_max_s}" + ) if replay_expert_agents not in (0, 1): raise ValueError(f"replay_expert_agents must be 0 or 1. Got: {replay_expert_agents}") self.replay_expert_agents = replay_expert_agents @@ -299,6 +320,12 @@ def __init__( # Co-sim envs stay out of eval_mode (its scenario batching needs a full agent pool) but still want the # eval-only observation/action semantics: perceived-size margin and standstill jerk deadband. self.cosim_eval_semantics = int(bool(cosim_eval_semantics)) + # Scripted crossing walkers per env = ceil(ratio * agents in the env); static partners, never policy-controlled. + self.scripted_pedestrian_ratio = float(scripted_pedestrian_ratio) + if not np.isfinite(self.scripted_pedestrian_ratio) or self.scripted_pedestrian_ratio < 0.0: + raise ValueError(f"scripted_pedestrian_ratio must be finite and >= 0. Got: {scripted_pedestrian_ratio}") + if self.scripted_pedestrian_ratio > 0.0 and (simulation_mode != "gigaflow" or self.cosim_partner_slots > 0): + raise ValueError("scripted_pedestrian_ratio > 0 needs gigaflow simulation_mode without cosim_partner_slots") if goal_source == "external" and simulation_mode != "gigaflow": raise ValueError("goal_source 'external' is only supported in gigaflow simulation_mode (co-sim)") @@ -318,6 +345,7 @@ def __init__( self.stop_signs_enabled = stop_signs_enabled self.yield_signs_enabled = yield_signs_enabled self.obs_slots_traffic_controls_n = obs_slots_traffic_controls_n + self.obs_traffic_light_yellow_remaining = int(bool(obs_traffic_light_yellow_remaining)) self.obs_norm_speed_mps = float(obs_norm_speed_mps) if not np.isfinite(self.obs_norm_speed_mps) or self.obs_norm_speed_mps <= 0.0: raise ValueError(f"obs_norm_speed_mps must be finite and > 0. Got: {obs_norm_speed_mps}") @@ -329,9 +357,31 @@ 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.red_light_train_front_margin_m = float(red_light_train_front_margin_m) + if not np.isfinite(self.red_light_train_front_margin_m) or self.red_light_train_front_margin_m < 0.0: + raise ValueError( + f"red_light_train_front_margin_m must be finite and >= 0. Got: {red_light_train_front_margin_m}" + ) + self.red_light_train_rear_margin_m = float(red_light_train_rear_margin_m) + if not np.isfinite(self.red_light_train_rear_margin_m) or self.red_light_train_rear_margin_m < 0.0: + raise ValueError( + f"red_light_train_rear_margin_m must be finite and >= 0. Got: {red_light_train_rear_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.obs_lane_width_clip_min_m = float(obs_lane_width_clip_min_m) + self.obs_lane_width_clip_max_m = float(obs_lane_width_clip_max_m) + lane_width_clip_enabled = self.obs_lane_width_clip_max_m > 0.0 + if ( + not (0.0 <= self.obs_lane_width_clip_min_m < float("inf")) + or not (0.0 <= self.obs_lane_width_clip_max_m < float("inf")) + or (lane_width_clip_enabled and self.obs_lane_width_clip_min_m > self.obs_lane_width_clip_max_m) + ): + raise ValueError( + "obs_lane_width_clip_min_m/max_m must be finite, >= 0 and min <= max (max 0 disables), " + f"got {obs_lane_width_clip_min_m} / {obs_lane_width_clip_max_m}" + ) self.eval_standstill_jerk_deadband_mps3 = float(eval_standstill_jerk_deadband_mps3) if self.eval_standstill_jerk_deadband_mps3 < 0: raise ValueError( @@ -369,7 +419,9 @@ def __init__( binding.LANE_SPEED_LIMIT_FEATURES if self.obs_lane_speed_limit else 0 ) self.boundary_features = binding.BOUNDARY_FEATURES - self.traffic_control_features = binding.TRAFFIC_CONTROL_FEATURES + self.traffic_control_features = binding.TRAFFIC_CONTROL_FEATURES + ( + binding.TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES if self.obs_traffic_light_yellow_remaining else 0 + ) self.obs_valid_count_features = binding.OBS_VALID_COUNT_FEATURES self.num_reward_coefs = binding.NUM_REWARD_COEFS if reward_conditioning else 0 @@ -572,6 +624,9 @@ def _env_init_kwargs(self, map_file, max_agents): "disable_red_light_infractions": self.disable_red_light_infractions, "disable_stop_sign_infractions": self.disable_stop_sign_infractions, "traffic_light_junction_phases": self.traffic_light_junction_phases, + "traffic_light_green_min_s": self.traffic_light_green_min_s, + "traffic_light_yellow_min_s": self.traffic_light_yellow_min_s, + "traffic_light_yellow_max_s": self.traffic_light_yellow_max_s, "stop_sign_behavior": self.stop_sign_behavior, "use_map_cache": self.use_map_cache, "use_neighbor_cache": self.use_neighbor_cache, @@ -593,6 +648,7 @@ def _env_init_kwargs(self, map_file, max_agents): "obs_lane_heading_signed": self.obs_lane_heading_signed, "obs_lane_speed_limit": self.obs_lane_speed_limit, "obs_slots_traffic_controls_n": self.obs_slots_traffic_controls_n, + "obs_traffic_light_yellow_remaining": self.obs_traffic_light_yellow_remaining, "traffic_lights_enabled": self.traffic_lights_enabled, "stop_signs_enabled": self.stop_signs_enabled, "yield_signs_enabled": self.yield_signs_enabled, @@ -622,6 +678,7 @@ def _env_init_kwargs(self, map_file, max_agents): "max_agents_per_env": self.max_agents_per_env, "cosim_partner_slots": self.cosim_partner_slots, "cosim_eval_semantics": self.cosim_eval_semantics, + "scripted_pedestrian_slots": self._scripted_pedestrian_slots(max_agents), "init_step": self._sample_init_step(), "init_step_min_horizon": self.init_step_min_horizon, "stagger_first_episode": self.stagger_first_episode, @@ -648,7 +705,11 @@ 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, + "red_light_train_front_margin_m": self.red_light_train_front_margin_m, + "red_light_train_rear_margin_m": self.red_light_train_rear_margin_m, "obs_partner_split_length_m": self.obs_partner_split_length_m, + "obs_lane_width_clip_min_m": self.obs_lane_width_clip_min_m, + "obs_lane_width_clip_max_m": self.obs_lane_width_clip_max_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, @@ -669,6 +730,11 @@ def _env_init_kwargs(self, map_file, max_agents): "phantom_braking_freeze_steering": self.phantom_braking_freeze_steering, } + def _scripted_pedestrian_slots(self, agents_in_env): + if self.scripted_pedestrian_ratio <= 0.0: + return 0 + return int(math.ceil(round(self.scripted_pedestrian_ratio * agents_in_env, 6))) + def _sample_init_step(self): # randomizer for the initialization of the C environment if not self.init_step_spread: @@ -817,14 +883,20 @@ def step(self, actions): def get_global_agent_state(self, include_static=False): """Get current global state of all active agents. - include_static: also return the co-sim partner slots after the active agents - (num_active_agents + cosim_partner_slots * num_envs rows). + include_static: also return the static slots (co-sim partners, scripted pedestrians) after the + active agents, in env order. Returns: dict with keys 'x', 'y', 'z', 'heading', 'id', 'length', 'width' containing numpy arrays of shape (num_active_agents,) """ - num_agents = self.num_agents + (self.cosim_partner_slots * self.num_envs if include_static else 0) + static_slots = 0 + if include_static: + env_agent_counts = np.diff(np.asarray(self.agent_offsets)) + static_slots = sum( + self.cosim_partner_slots + self._scripted_pedestrian_slots(int(count)) for count in env_agent_counts + ) + num_agents = self.num_agents + static_slots states = { "x": np.zeros(num_agents, dtype=np.float32), @@ -899,10 +971,22 @@ def recompute_observations(self): binding.vec_recompute_observations(self.c_envs) return self.observations - def set_traffic_light_states(self, states): + def set_traffic_light_states(self, states, yellow_remaining_s=None): """override each traffic-light element's state at the current - timestep (states length == num_traffic_elements).""" - binding.vec_set_traffic_light_states(self.c_envs, np.ascontiguousarray(states, dtype=np.int32)) + timestep (states length == num_traffic_elements). yellow_remaining_s + (same length, seconds each light stays yellow from now, 0 unless + yellow) feeds the obs_traffic_light_yellow_remaining column and is + required when that flag is set; the co-sim's value is clipped to the + training max traffic_light_yellow_max_s.""" + states = np.ascontiguousarray(states, dtype=np.int32) + if yellow_remaining_s is None: + binding.vec_set_traffic_light_states(self.c_envs, states) + return + remaining = np.asarray(yellow_remaining_s, dtype=np.float32).reshape(-1) + if not np.all(np.isfinite(remaining)) or np.any(remaining < 0.0): + raise ValueError(f"yellow_remaining_s must be finite and >= 0. Got: {yellow_remaining_s}") + remaining = np.minimum(remaining, np.float32(self.traffic_light_yellow_max_s)) + binding.vec_set_traffic_light_states(self.c_envs, states, np.ascontiguousarray(remaining)) def set_stop_signs(self, lines, headings): """replace every stop-sign element with the external sim's own (co-sim): @@ -945,11 +1029,26 @@ def set_agent_goals(self, agent_idx, gx, gy, gz, gdir_x=None, gdir_y=None): np.ascontiguousarray(gdir_y, dtype=np.float32), ) + def set_agent_lane_goals(self, idx, spacing_m): + """Co-sim forecast: the agents at global indices `idx` have no known route, so their goals follow + the lane they are on and its straightest exits, one every `spacing_m` meters; c_step places new + ones whenever one is reached. Returns the number of goals placed per agent (0: no lane under the + agent, or the lane graph ends within spacing_m).""" + idx = np.ascontiguousarray(idx, dtype=np.int32) + placed_goal_counts = np.zeros(len(idx), dtype=np.int32) + binding.vec_set_agent_lane_goals(self.c_envs, idx, float(spacing_m), placed_goal_counts) + return placed_goal_counts + def get_agent_goal_progress(self, agent_idx): """Co-sim: (current_goal_idx, goal_count) of one agent. current_goal_idx == goal_count means every goal of the current window has been consumed.""" return binding.vec_get_agent_goal_progress(self.c_envs, int(agent_idx)) + def set_agent_goal_progress(self, agent_idx, current_goal_idx): + """Co-sim: set how many goals of one agent's current window count as consumed (0..goal_count), + e.g. to undo the goals a forecast rollout consumed.""" + binding.vec_set_agent_goal_progress(self.c_envs, int(agent_idx), int(current_goal_idx)) + # ─────────────────────────────────────── def get_ground_truth_trajectories(self): diff --git a/pufferlib/ocean/drive/map_data.h b/pufferlib/ocean/drive/map_data.h index 3b79c415ca..61be508884 100644 --- a/pufferlib/ocean/drive/map_data.h +++ b/pufferlib/ocean/drive/map_data.h @@ -123,6 +123,59 @@ static int append_cell_lane_query_entries(const Drive *env, int cell_idx, LaneQu return count; } +// Counts (and, with reach_entries, writes) one road-edge segment into every cell whose window it reaches from outside. +static void file_reaching_edge_segment( + Drive *env, + int entity_idx, + int seg_idx, + int *cell_cursor, + LaneQueryEntry *reach_entries) { + const GridMap *grid = env->grid_map; + const RoadMapElement *element = &env->road_elements[entity_idx]; + float x0 = element->x[seg_idx], y0 = element->y[seg_idx]; + float x1 = element->x[seg_idx + 1], y1 = element->y[seg_idx + 1]; + // the grid files a segment under its midpoint cell only, the window of every cell around that one lists it + int filed_cell_idx = get_grid_index(env, (x0 + x1) / 2, (y0 + y1) / 2); + int filed_x = filed_cell_idx % grid->grid_cols; + int filed_y = filed_cell_idx / grid->grid_cols; + int first_x = (int) ((fminf(x0, x1) - grid->top_left_x) / GRID_CELL_SIZE) - ROAD_QUERY_HALF_CELLS; + int last_x = (int) ((fmaxf(x0, x1) - grid->top_left_x) / GRID_CELL_SIZE) + ROAD_QUERY_HALF_CELLS; + int first_y = (int) ((fminf(y0, y1) - grid->bottom_right_y) / GRID_CELL_SIZE) - ROAD_QUERY_HALF_CELLS; + int last_y = (int) ((fmaxf(y0, y1) - grid->bottom_right_y) / GRID_CELL_SIZE) + ROAD_QUERY_HALF_CELLS; + first_x = first_x < 0 ? 0 : first_x; + first_y = first_y < 0 ? 0 : first_y; + last_x = last_x >= grid->grid_cols ? grid->grid_cols - 1 : last_x; + last_y = last_y >= grid->grid_rows ? grid->grid_rows - 1 : last_y; + int span_x = last_x - first_x + 1; + int span_y = last_y - first_y + 1; + if (span_x <= 0 || span_y <= 0) { + return; + } + for (int o = 0; o < span_x * span_y; o++) { + int cell_x = first_x + o % span_x; + int cell_y = first_y + o / span_x; + if (abs(cell_x - filed_x) <= ROAD_QUERY_HALF_CELLS && abs(cell_y - filed_y) <= ROAD_QUERY_HALF_CELLS) { + continue; + } + int cell_idx = cell_y * grid->grid_cols + cell_x; + if (reach_entries != NULL) { + reach_entries[cell_cursor[cell_idx]] = (LaneQueryEntry) {entity_idx, seg_idx, seg_idx}; + } + cell_cursor[cell_idx]++; + } +} + +static void file_reaching_edge_segments(Drive *env, int *cell_cursor, LaneQueryEntry *reach_entries) { + for (int i = 0; i < env->num_road_elements; i++) { + if (!is_road_edge(env->road_elements[i].type)) { + continue; + } + for (int j = 0; j < env->road_elements[i].segment_size - 1; j++) { + file_reaching_edge_segment(env, i, j, cell_cursor, reach_entries); + } + } +} + // 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; @@ -134,20 +187,43 @@ static void cache_lane_query_lists(Drive *env) { 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"); } + // Long road edges reach windows that do not list them: per cell, [reach_offset[c], reach_offset[c + 1]) adds those + int *reach_offset = (int *) calloc(cell_count + 1, sizeof(int)); + int *reach_cursor = (int *) calloc(cell_count, sizeof(int)); + if (reach_offset == NULL || reach_cursor == NULL) { + raise_error_with_message(ERROR_MEMORY_ALLOCATION, "Failed to allocate the road edge reach index"); + } + file_reaching_edge_segments(env, reach_cursor, NULL); + for (int cell_idx = 0; cell_idx < cell_count; cell_idx++) { + reach_offset[cell_idx + 1] = reach_offset[cell_idx] + reach_cursor[cell_idx]; + reach_cursor[cell_idx] = reach_offset[cell_idx]; + } + LaneQueryEntry *reach_entries = (LaneQueryEntry *) malloc((reach_offset[cell_count] + 1) * sizeof(LaneQueryEntry)); + if (reach_entries == NULL) { + raise_error_with_message(ERROR_MEMORY_ALLOCATION, "Failed to allocate the road edge reach entries"); + } + file_reaching_edge_segments(env, reach_cursor, reach_entries); + 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); + int reach_count = reach_offset[cell_idx + 1] - reach_offset[cell_idx]; + if (count + window_count + reach_count > capacity) { + capacity = 2 * (count + window_count + reach_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); + memcpy(&entries[count], &reach_entries[reach_offset[cell_idx]], reach_count * sizeof(LaneQueryEntry)); + count += reach_count; } + free(reach_offset); + free(reach_cursor); + free(reach_entries); grid->lane_query_offset[cell_count] = count; grid->lane_query_pool = entries; } diff --git a/pufferlib/ocean/env_binding.h b/pufferlib/ocean/env_binding.h index 0c6b91f9f4..1673caf360 100644 --- a/pufferlib/ocean/env_binding.h +++ b/pufferlib/ocean/env_binding.h @@ -1121,8 +1121,11 @@ static PyObject *vec_recompute_observations(PyObject *self, PyObject *args) { } static PyObject *vec_set_traffic_light_states(PyObject *self, PyObject *args) { - if (PyTuple_Size(args) != 2) { - PyErr_SetString(PyExc_TypeError, "vec_set_traffic_light_states requires 2 arguments"); + Py_ssize_t arg_count = PyTuple_Size(args); + if (arg_count != 2 && arg_count != 3) { + PyErr_SetString( + PyExc_TypeError, + "vec_set_traffic_light_states requires 2 or 3 arguments (vec_env, states[, yellow_remaining_s])"); return NULL; } VecEnv *vec = unpack_vecenv(args); @@ -1144,10 +1147,25 @@ static PyObject *vec_set_traffic_light_states(PyObject *self, PyObject *args) { (int) PyArray_SIZE((PyArrayObject *) states_arr)); return NULL; } - if (c_set_traffic_light_states(drive, states) != 0) { + float *yellow_remaining_s = NULL; + if (arg_count == 3 && PyTuple_GetItem(args, 2) != Py_None) { + PyObject *remaining_arr = PyTuple_GetItem(args, 2); + if (!PyArray_Check(remaining_arr) || PyArray_TYPE((PyArrayObject *) remaining_arr) != NPY_FLOAT32 + || (int) PyArray_SIZE((PyArrayObject *) remaining_arr) != drive->num_traffic_elements) { + PyErr_Format( + PyExc_ValueError, + "vec_set_traffic_light_states: yellow_remaining_s must be a float32 NumPy array of %d entries", + drive->num_traffic_elements); + return NULL; + } + yellow_remaining_s = (float *) PyArray_DATA((PyArrayObject *) remaining_arr); + } + if (c_set_traffic_light_states(drive, states, yellow_remaining_s) != 0) { PyErr_SetString( PyExc_ValueError, - "vec_set_traffic_light_states: timestep outside the light state buffer or invalid state"); + "vec_set_traffic_light_states: timestep outside the light state buffer, invalid state, or " + "yellow_remaining_s missing (obs_traffic_light_yellow_remaining needs it) or outside " + "[0, traffic_light_yellow_max_s]"); return NULL; } Py_RETURN_NONE; @@ -1235,6 +1253,31 @@ static PyObject *vec_get_agent_goal_progress(PyObject *self, PyObject *args) { return Py_BuildValue("(ii)", current_goal_idx, goal_count); } +static PyObject *vec_set_agent_goal_progress(PyObject *self, PyObject *args) { + if (PyTuple_Size(args) != 3) { + PyErr_SetString( + PyExc_TypeError, + "vec_set_agent_goal_progress requires 3 arguments (vec_env, agent_idx, current_goal_idx)"); + return NULL; + } + VecEnv *vec = unpack_vecenv(args); + if (!vec) { + return NULL; + } + int agent_idx = (int) PyLong_AsLong(PyTuple_GetItem(args, 1)); + int current_goal_idx = (int) PyLong_AsLong(PyTuple_GetItem(args, 2)); + if (PyErr_Occurred()) { + return NULL; + } + if (c_set_agent_goal_progress((Drive *) vec->envs[0], agent_idx, current_goal_idx) != 0) { + PyErr_SetString( + PyExc_ValueError, + "vec_set_agent_goal_progress: agent index out of range or current_goal_idx outside 0..goal_count"); + return NULL; + } + Py_RETURN_NONE; +} + static PyObject *vec_set_agent_goals(PyObject *self, PyObject *args) { if (PyTuple_Size(args) != 7) { PyErr_SetString(PyExc_TypeError, "vec_set_agent_goals requires 7 arguments"); @@ -1277,9 +1320,47 @@ static PyObject *vec_set_agent_goals(PyObject *self, PyObject *args) { Py_RETURN_NONE; } +static PyObject *vec_set_agent_lane_goals(PyObject *self, PyObject *args) { + if (PyTuple_Size(args) != 4) { + PyErr_SetString( + PyExc_TypeError, + "vec_set_agent_lane_goals requires 4 arguments (vec_env, idx, spacing_m, placed_goal_counts)"); + return NULL; + } + VecEnv *vec = unpack_vecenv(args); + if (!vec) { + return NULL; + } + PyObject *idx_arr = PyTuple_GetItem(args, 1); + PyObject *counts_arr = PyTuple_GetItem(args, 3); + double spacing_m = PyFloat_AsDouble(PyTuple_GetItem(args, 2)); + if (PyErr_Occurred()) { + return NULL; + } + if (!PyArray_Check(idx_arr) || !PyArray_Check(counts_arr) || PyArray_TYPE((PyArrayObject *) idx_arr) != NPY_INT32 + || PyArray_TYPE((PyArrayObject *) counts_arr) != NPY_INT32) { + PyErr_SetString(PyExc_TypeError, "vec_set_agent_lane_goals: idx and placed_goal_counts must be int32 arrays"); + return NULL; + } + int count = (int) PyArray_SIZE((PyArrayObject *) idx_arr); + if ((int) PyArray_SIZE((PyArrayObject *) counts_arr) != count) { + PyErr_SetString(PyExc_ValueError, "vec_set_agent_lane_goals: placed_goal_counts must have the length of idx"); + return NULL; + } + const int *idx = (const int *) PyArray_DATA((PyArrayObject *) idx_arr); + int *placed_goal_counts = (int *) PyArray_DATA((PyArrayObject *) counts_arr); + if (c_set_agent_lane_goals((Drive *) vec->envs[0], count, idx, (float) spacing_m, placed_goal_counts) != 0) { + PyErr_SetString( + PyExc_ValueError, + "vec_set_agent_lane_goals: agent index out of range, agent without a pose, or spacing_m not > 0"); + return NULL; + } + Py_RETURN_NONE; +} + static PyObject *get_ground_truth_trajectories(PyObject *self, PyObject *args) { - if (PyTuple_Size(args) != 7) { - PyErr_SetString(PyExc_TypeError, "get_ground_truth_trajectories requires 7 arguments"); + if (PyTuple_Size(args) != 8) { + PyErr_SetString(PyExc_TypeError, "get_ground_truth_trajectories requires 8 arguments"); return NULL; } @@ -1313,8 +1394,13 @@ static PyObject *get_ground_truth_trajectories(PyObject *self, PyObject *args) { int *id_data = (int *) PyArray_DATA((PyArrayObject *) id_arr); int *scenario_id_data = (int *) PyArray_DATA((PyArrayObject *) scenario_id_arr); + if (PyArray_NDIM((PyArrayObject *) x_arr) != 2) { + PyErr_SetString(PyExc_ValueError, "trajectory arrays must be 2D [num_agents, num_timesteps]"); + return NULL; + } c_get_global_ground_truth_trajectories( drive, + (int) PyArray_DIM((PyArrayObject *) x_arr, 1), x_data, y_data, z_data, @@ -1369,9 +1455,20 @@ static PyObject *vec_get_global_ground_truth_trajectories(PyObject *self, PyObje int *id_base = (int *) PyArray_DATA(id_array); int *scenario_id_base = (int *) PyArray_DATA(scenario_id_array); - // Get number of timesteps from array shape - npy_intp *x_shape = PyArray_DIMS(x_array); - int num_timesteps = x_shape[1]; // Second dimension for 2D arrays + PyArrayObject *trajectory_arrays[5] = {x_array, y_array, z_array, heading_array, valid_array}; + npy_intp total_agent_count = 0; + for (int i = 0; i < vec->num_envs; i++) { + total_agent_count += ((Drive *) vec->envs[i])->active_agent_count; + } + for (int k = 0; k < 5; k++) { + PyArrayObject *array = trajectory_arrays[k]; + if (PyArray_NDIM(array) != 2 || PyArray_NDIM(x_array) != 2 || PyArray_DIM(array, 0) < total_agent_count + || PyArray_DIM(array, 1) != PyArray_DIM(x_array, 1)) { + PyErr_SetString(PyExc_ValueError, "trajectory arrays must share one 2D shape [num_agents, num_timesteps]"); + return NULL; + } + } + int num_timesteps = (int) PyArray_DIM(x_array, 1); // Iterate through environments and write to correct offsets int agent_offset = 0; // Offset for 1D arrays (id, scenario_id) @@ -1382,6 +1479,7 @@ static PyObject *vec_get_global_ground_truth_trajectories(PyObject *self, PyObje c_get_global_ground_truth_trajectories( drive, + num_timesteps, &x_base[traj_offset], &y_base[traj_offset], &z_base[traj_offset], @@ -1567,6 +1665,14 @@ static PyMethodDef methods[] 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_set_agent_lane_goals", + vec_set_agent_lane_goals, + METH_VARARGS, + "Give agents goals along their lane ahead (co-sim forecast)"}, + {"vec_set_agent_goal_progress", + vec_set_agent_goal_progress, + METH_VARARGS, + "Set the consumed-goal count of one agent's window (co-sim)"}, {"vec_get_agent_goal_progress", vec_get_agent_goal_progress, METH_VARARGS, @@ -1629,6 +1735,7 @@ PyMODINIT_FUNC PyInit_binding(void) { PyModule_AddIntConstant(m, "PARTNER_FEATURES", PARTNER_FEATURES); PyModule_AddIntConstant(m, "PARTNER_RELATIVE_VELOCITY_FEATURES", PARTNER_RELATIVE_VELOCITY_FEATURES); PyModule_AddIntConstant(m, "TRAFFIC_CONTROL_FEATURES", TRAFFIC_CONTROL_FEATURES); + PyModule_AddIntConstant(m, "TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES", TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES); PyModule_AddIntConstant(m, "OBS_VALID_COUNT_FEATURES", OBS_VALID_COUNT_FEATURES); PyModule_AddIntConstant(m, "NUM_TRAFFIC_CONTROL_TYPES", NUM_TRAFFIC_CONTROL_TYPES); PyModule_AddIntConstant(m, "NUM_TRAFFIC_CONTROL_STATES", NUM_TRAFFIC_CONTROL_STATES); @@ -1674,6 +1781,8 @@ PyMODINIT_FUNC PyInit_binding(void) { // norm (obs[0]) is per-env config: env.obs_norm_speed_mps. PyModule_AddObject(m, "ACCEL_LONG_NORM", PyFloat_FromDouble(fabsf(ACCEL_LONG_LIMIT[0]))); PyModule_AddObject(m, "ACCEL_LONG_MAX", PyFloat_FromDouble(ACCEL_LONG_LIMIT[1])); + PyModule_AddObject(m, "STEERING_LIMIT", PyFloat_FromDouble(STEERING_LIMIT)); + PyModule_AddObject(m, "MAX_BACKWARD_SPEED", PyFloat_FromDouble(MAX_BACKWARD_SPEED)); PyModule_AddObject(m, "AGENT_STOPPED_SPEED_THRESHOLD", PyFloat_FromDouble(AGENT_STOPPED_SPEED_THRESHOLD)); PyModule_AddObject(m, "CONDITIONING_ACC_MIN", PyFloat_FromDouble(REWARD_BOUNDS[REWARD_COEF_ACC].min_val)); PyModule_AddObject(m, "CONDITIONING_ACC_MAX", PyFloat_FromDouble(REWARD_BOUNDS[REWARD_COEF_ACC].max_val)); diff --git a/pufferlib/ocean/evaluation_utils/eval_replay.py b/pufferlib/ocean/evaluation_utils/eval_replay.py index 301e1ffd49..a635841a42 100644 --- a/pufferlib/ocean/evaluation_utils/eval_replay.py +++ b/pufferlib/ocean/evaluation_utils/eval_replay.py @@ -15,6 +15,11 @@ ZLIB_REPLAY_DIR_NAME = "replays_zlib" ZLIB_REPLAY_SUFFIX = ".replay.zlib" +OBSERVATION_AGENT_MODES = ("all", "infracting") +INFRACTION_METRIC_INDICES = tuple( + pufferlib.viz.METRIC_LABELS.index(name) for name in ("collision", "offroad", "red_light", "stop_sign") +) +POOL_CHUNK_PREFIX = "pool_" def _eval_replay_stem(summary, episode_id): @@ -39,6 +44,11 @@ def __init__( if capture_batch_steps <= 0: raise RuntimeError("Replay capture requires a positive resample frequency") self.capture_observations = capture_observations + self.observation_agents = args["eval"]["capture_observations_agents"] + if self.observation_agents not in OBSERVATION_AGENT_MODES: + raise pufferlib.APIUsageError( + f"eval.capture_observations_agents must be one of {OBSERVATION_AGENT_MODES}, got {self.observation_agents!r}" + ) self.replay_writer_count = num_workers if self.capture_observations: observation_replay_writer_count = args["eval"]["observation_replay_writer_count"] @@ -144,6 +154,8 @@ def queue_replay(self, summary, episode_id): replay = {"env": self.env_config, **replay_environment["frames"]} for replay_key, history_values in self.policy_history.items(): replay[replay_key] = history_values[:episode_length, global_agent_start:global_agent_end] + if self.capture_observations and self.observation_agents == "infracting": + keep_observations_of_infracting_agents(replay) replay_stem = _eval_replay_stem(summary, self.episode_id_offset + episode_id) replay_path = os.path.abspath(os.path.join(self.replay_output_dir, f"{replay_stem}{ZLIB_REPLAY_SUFFIX}")) self.pending_replays.append((replay_environment["scenario"], replay, replay_path)) @@ -164,6 +176,17 @@ def write_pending(self): self.pending_replays = [] +def keep_observations_of_infracting_agents(replay): + """Zero the observations and encoder pool winners of every agent without an infraction in the episode (the + zeros deflate to nothing: a 50-agent page shrinks about 4x), listing the kept agents as replay["obs_agents"].""" + metrics = np.asarray(replay["metrics_f32"]) + infracting = np.any(metrics[:, :, list(INFRACTION_METRIC_INDICES)] > 0, axis=(0, 2)) + replay["obs_agents"] = np.flatnonzero(infracting).tolist() + for replay_key in ["obs"] + [key for key in replay if key.startswith(POOL_CHUNK_PREFIX)]: + if replay.get(replay_key) is not None: + replay[replay_key][:, ~infracting] = 0 + + def _render_eval_replays(episode_summaries, out_dir, keep_zlib_replays, max_renderer_count=None): """Render captured eval replays as navigable HTML pages plus an index.""" render_dir = os.path.join(out_dir, "rendered_replays") diff --git a/pufferlib/ocean/evaluation_utils/eval_wandb_report.py b/pufferlib/ocean/evaluation_utils/eval_wandb_report.py new file mode 100644 index 0000000000..8ce5c4cc75 --- /dev/null +++ b/pufferlib/ocean/evaluation_utils/eval_wandb_report.py @@ -0,0 +1,271 @@ +"""Write one finished evaluation's headline metric into the summary of the wandb run that trained the checkpoint.""" + +import argparse +import csv +import fcntl +import glob +import json +import math +import numbers +import os +import re +import time + +import yaml + + +EVAL_TYPE_LONGEST6_V1 = "longest6_v1" +EVAL_TYPE_LONGEST6_V2 = "longest6_v2" +EVAL_TYPE_NUPLAN_REACTIVE = "nuplan_reactive" +EVAL_TYPE_NUPLAN_NONREACTIVE = "nuplan_nonreactive" +EVAL_TYPE_SELFPLAY = "selfplay" +EVAL_TYPE_ALPASIM = "alpasim" + +HEADLINE_METRIC_KEYS = { + EVAL_TYPE_LONGEST6_V1: "final_eval/longest6_v1_DS", + EVAL_TYPE_LONGEST6_V2: "final_eval/longest6_v2_DS", + EVAL_TYPE_NUPLAN_REACTIVE: "final_eval/nuplan_reactive_CLS", + EVAL_TYPE_NUPLAN_NONREACTIVE: "final_eval/nuplan_nonreactive_CLS", + EVAL_TYPE_SELFPLAY: "final_eval/selfplay_km_per_infraction", + EVAL_TYPE_ALPASIM: "final_eval/alpasim_score", +} +SOURCE_KEY_SUFFIX = "_source" + +SELFPLAY_BENCHMARK_NAME = "carla" +SELFPLAY_SUMMARY_FILE_NAME = "evaluation_summary.json" +METERS_PER_KM = 1000.0 + +LONGEST6_DRIVING_SCORE_ROW = "Avg. driving score" +LONGEST6_DRIVING_SCORE_TOLERANCE = 0.01 +NUPLAN_SCORE_COLUMN = "CLS" +NUPLAN_SCENARIOS_FILE_NAME = "scenarios.csv" +LONGEST6_ROUTES_DIR_NAME = "routes" +LONGEST6_RECORD_SUFFIX = ".result.json" +ALPASIM_SCORE_LINE_PATTERN = re.compile(r"official_score (\S+)\s+scenes (\d+)\s+failed (\d+)") + +RUN_CONFIG_FILE_NAME = "config.yaml" +RUN_IDENTITY_KEYS = ("run_name", "wandb_project") +REPORT_LOCK_FILE_NAME = ".eval_wandb_report.lock" +WANDB_RUNNING_STATE = "running" +SUMMARY_FLOAT_REL_TOLERANCE = 1e-9 +SUMMARY_VERIFY_ATTEMPT_COUNT = 3 +SUMMARY_VERIFY_BACKOFF_SECONDS = 2.0 + + +class EvalReportError(Exception): + pass + + +def longest6_route_record_paths(results_csv_path): + """The route records results.csv summarizes: routes/*.result.json of a run tidy_carla_run.py flattened (results.csv + at its top), else every json next to a fresh run's aggregate/results.csv (the parser's copies of the records).""" + results_dir = os.path.dirname(results_csv_path) + flat_records = glob.glob(os.path.join(results_dir, LONGEST6_ROUTES_DIR_NAME, f"*{LONGEST6_RECORD_SUFFIX}")) + return flat_records or glob.glob(os.path.join(results_dir, "**", "*.json"), recursive=True) + + +def read_longest6_driving_score(results_csv_path, expected_route_run_count): + with open(results_csv_path, newline="") as results_file: + driving_scores = [ + row[1] for row in csv.reader(results_file) if len(row) >= 2 and row[0] == LONGEST6_DRIVING_SCORE_ROW + ] + if len(driving_scores) != 1: + raise EvalReportError(f"{results_csv_path} must hold exactly one '{LONGEST6_DRIVING_SCORE_ROW}' row") + driving_score = float(driving_scores[0]) + + route_scores = [] + for record_path in longest6_route_record_paths(results_csv_path): + with open(record_path) as record_file: + route_records = json.load(record_file)["_checkpoint"]["records"] + route_scores.extend(float(record["scores"]["score_composed"]) for record in route_records) + if len(route_scores) != expected_route_run_count: + raise EvalReportError( + f"incomplete longest6 eval: {len(route_scores)} of {expected_route_run_count} route runs have a record " + f"next to {results_csv_path}" + ) + mean_route_score = sum(route_scores) / len(route_scores) + if abs(driving_score - mean_route_score) > LONGEST6_DRIVING_SCORE_TOLERANCE: + raise EvalReportError( + f"{results_csv_path} is stale: its driving score {driving_score} is not the mean {mean_route_score} " + "of the route records next to it" + ) + return driving_score + + +def read_nuplan_closed_loop_score(overview_csv_path, expected_scenario_count): + with open(overview_csv_path, newline="") as overview_file: + overview_rows = list(csv.DictReader(overview_file)) + if len(overview_rows) != 1 or NUPLAN_SCORE_COLUMN not in overview_rows[0]: + raise EvalReportError(f"{overview_csv_path} must hold exactly one row with a {NUPLAN_SCORE_COLUMN} column") + + scenarios_csv_path = os.path.join(os.path.dirname(overview_csv_path), NUPLAN_SCENARIOS_FILE_NAME) + with open(scenarios_csv_path, newline="") as scenarios_file: + scenario_count = sum(1 for _ in csv.DictReader(scenarios_file)) + if scenario_count != expected_scenario_count: + raise EvalReportError( + f"incomplete nuPlan eval: {scenario_count} of {expected_scenario_count} scenarios scored in " + f"{scenarios_csv_path}" + ) + return float(overview_rows[0][NUPLAN_SCORE_COLUMN]) + + +def read_selfplay_km_per_infraction(summary_json_path): + with open(summary_json_path) as summary_file: + summary = json.load(summary_file) + if summary["num_episodes"] != summary["num_scenarios"]: + raise EvalReportError( + f"incomplete self-play eval: {summary['num_episodes']} of {summary['num_scenarios']} scenarios finished " + f"in {summary_json_path}" + ) + return float(summary["metrics_mean"]["avg_distance_per_infraction"]) / METERS_PER_KM + + +def read_alpasim_score(results_txt_path, expected_scene_count): + with open(results_txt_path) as results_file: + score_line = results_file.readline().strip() + score_match = ALPASIM_SCORE_LINE_PATTERN.fullmatch(score_line) + if score_match is None: + raise EvalReportError(f"{results_txt_path} holds no official score: '{score_line}'") + scored_scene_count = int(score_match.group(2)) + failed_scene_count = int(score_match.group(3)) + if scored_scene_count != expected_scene_count or failed_scene_count != 0: + raise EvalReportError( + f"incomplete AlpaSim eval: {scored_scene_count} of {expected_scene_count} scenes scored, " + f"{failed_scene_count} failed in {results_txt_path}" + ) + return float(score_match.group(1)) + + +def read_headline_value(eval_type, result_path, expected_episode_count): + """Read an eval type's headline number from its result file; refuses incomplete evaluations.""" + if not os.path.isfile(result_path): + raise EvalReportError(f"missing result file {result_path}") + if eval_type != EVAL_TYPE_SELFPLAY and expected_episode_count is None: + raise EvalReportError(f"{eval_type} needs the benchmark's episode count to tell a complete eval from a subset") + + if eval_type == EVAL_TYPE_SELFPLAY: + headline_value = read_selfplay_km_per_infraction(result_path) + elif eval_type == EVAL_TYPE_ALPASIM: + headline_value = read_alpasim_score(result_path, expected_episode_count) + elif eval_type in (EVAL_TYPE_NUPLAN_REACTIVE, EVAL_TYPE_NUPLAN_NONREACTIVE): + headline_value = read_nuplan_closed_loop_score(result_path, expected_episode_count) + elif eval_type in (EVAL_TYPE_LONGEST6_V1, EVAL_TYPE_LONGEST6_V2): + headline_value = read_longest6_driving_score(result_path, expected_episode_count) + else: + raise EvalReportError(f"unknown eval type '{eval_type}'; expected one of {sorted(HEADLINE_METRIC_KEYS)}") + + if not math.isfinite(headline_value): + raise EvalReportError(f"{result_path} holds a non-finite {eval_type} headline value: {headline_value}") + return headline_value + + +def load_run_identity(run_config_path): + """Resolve the wandb run a checkpoint belongs to from its training config.""" + if not os.path.isfile(run_config_path): + raise EvalReportError(f"missing run config {run_config_path}") + with open(run_config_path) as run_config_file: + run_config = yaml.safe_load(run_config_file) + if not isinstance(run_config, dict): + raise EvalReportError(f"run config {run_config_path} must be a mapping") + + run_identity = {} + for key in RUN_IDENTITY_KEYS: + value = run_config.get(key) + if not isinstance(value, str) or not value.strip(): + raise EvalReportError( + f"run config {run_config_path} is missing a usable {key}; cannot attach eval results to the " + "training run" + ) + run_identity[key] = value + return run_identity + + +def find_summary_mismatches(expected_summary, stored_summary): + """Keys whose stored value is not the expected one; numbers may differ by the server's float formatting.""" + mismatched_keys = [] + for key in sorted(expected_summary.keys() | stored_summary.keys()): + expected_value = expected_summary.get(key) + stored_value = stored_summary.get(key) + values_are_numbers = isinstance(expected_value, numbers.Real) and isinstance(stored_value, numbers.Real) + if values_are_numbers and math.isclose(expected_value, stored_value, rel_tol=SUMMARY_FLOAT_REL_TOLERANCE): + continue + if json.dumps(expected_value, sort_keys=True) != json.dumps(stored_value, sort_keys=True): + mismatched_keys.append(key) + return mismatched_keys + + +def write_run_summary(run_identity, reported_values): + """Set summary values of a finished wandb run and verify that nothing else in its summary changed.""" + import wandb + + run_path = f"{run_identity['wandb_project']}/{run_identity['run_name']}" + run = wandb.Api().run(run_path) + if run.state == WANDB_RUNNING_STATE: + raise EvalReportError( + f"wandb run {run_path} is still running and its own logger would overwrite the summary; report again " + "once training has finished" + ) + + expected_summary = {**run.summary_metrics, **reported_values} + run.summary.update(reported_values) + + mismatched_keys = [] + for attempt_idx in range(SUMMARY_VERIFY_ATTEMPT_COUNT): + time.sleep(attempt_idx * SUMMARY_VERIFY_BACKOFF_SECONDS) + mismatched_keys = find_summary_mismatches(expected_summary, wandb.Api().run(run_path).summary_metrics) + if not mismatched_keys: + return + raise EvalReportError( + f"wandb run {run_path}: summary keys differ from the intended update, check them: {mismatched_keys}" + ) + + +def report_headline_metric(eval_type, result_path, run_dir, expected_episode_count=None): + """Report one evaluation's headline metric; a repeated report replaces the run's value for that eval type.""" + headline_value = read_headline_value(eval_type, result_path, expected_episode_count) + run_identity = load_run_identity(os.path.join(run_dir, RUN_CONFIG_FILE_NAME)) + metric_key = HEADLINE_METRIC_KEYS[eval_type] + result_source = os.path.relpath(os.path.abspath(result_path), os.path.abspath(run_dir)) + reported_values = {metric_key: headline_value, metric_key + SOURCE_KEY_SUFFIX: result_source} + + # the wandb API replaces the whole summary: evals of one run finishing together must not interleave + with open(os.path.join(run_dir, REPORT_LOCK_FILE_NAME), "w") as lock_file: + fcntl.flock(lock_file, fcntl.LOCK_EX) + write_run_summary(run_identity, reported_values) + + print( + f"wandb run {run_identity['wandb_project']}/{run_identity['run_name']}: {metric_key} = {headline_value} " + f"({result_source})" + ) + return reported_values + + +def parse_arguments(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("eval_type", choices=sorted(HEADLINE_METRIC_KEYS)) + parser.add_argument( + "result_path", + help="longest6: results.csv (aggregate/results.csv before the tidy); nuPlan: overview.csv; selfplay: evaluation_summary.json; " + "alpasim: results.txt", + ) + parser.add_argument("--run-dir", required=True, help="Training run folder holding config.yaml") + parser.add_argument( + "--expected-episode-count", + type=int, + help="Route runs, scenarios or scenes of the complete benchmark; selfplay reads it from its summary", + ) + return parser.parse_args() + + +def main(): + arguments = parse_arguments() + try: + report_headline_metric( + arguments.eval_type, arguments.result_path, arguments.run_dir, arguments.expected_episode_count + ) + except EvalReportError as error: + raise SystemExit(f"eval_wandb_report: {error}") + + +if __name__ == "__main__": + main() diff --git a/pufferlib/ocean/evaluation_utils/evaluation_utils.py b/pufferlib/ocean/evaluation_utils/evaluation_utils.py index 878ed309c8..7ac45692df 100644 --- a/pufferlib/ocean/evaluation_utils/evaluation_utils.py +++ b/pufferlib/ocean/evaluation_utils/evaluation_utils.py @@ -180,24 +180,6 @@ def load_checkpoint_architecture(args): return merged, config_path -CHECKPOINT_RUN_IDENTITY_KEYS = ("run_name", "wandb_project", "wandb_group") - - -def load_checkpoint_run_identity(config_path): - """Resolve the tracker run a checkpoint belongs to from its training config.""" - checkpoint_config = _load_yaml_mapping(config_path, "checkpoint config") - identity = {} - for key in CHECKPOINT_RUN_IDENTITY_KEYS: - value = checkpoint_config.get(key) - if not isinstance(value, str) or not value.strip(): - raise pufferlib.APIUsageError( - f"Checkpoint config {config_path} is missing a usable {key}; cannot attach eval " - "results to the training run" - ) - identity[key] = value - return identity - - def summarize_benchmark_metrics(benchmark_results, key_prefix): """Flatten benchmark results into one scalar metric dict for the experiment logger.""" metrics = {} diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index 0e93120696..588a32e260 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -40,10 +40,13 @@ from pufferlib.ocean.evaluation_utils import evaluation_utils as drive_benchmark from pufferlib.ocean.evaluation_utils import eval_replay as drive_eval_replay from pufferlib.ocean.evaluation_utils import cosim_evaluator as drive_cosim_eval +from pufferlib.ocean.evaluation_utils import eval_wandb_report as drive_eval_report import pufferlib.sweep import pufferlib.utils import pufferlib.vector import pufferlib.pytorch +import pufferlib.ensemble +from pufferlib.pytorch import clean_policy_state_dict from pufferlib.config_schema import ( MAX_C_SEED, normalize_puffer_drive_config, @@ -119,17 +122,6 @@ def base_policy(policy): return policy.module if hasattr(policy, "module") else policy -def clean_state_key(key): - prefixes = ("module.", "_orig_mod.") - while key.startswith(prefixes): - key = key.split(".", 1)[1] - return key - - -def clean_policy_state_dict(state_dict): - return {clean_state_key(k): v for k, v in state_dict.items()} - - def logits_to_float(logits): if isinstance(logits, torch.distributions.Normal): return torch.distributions.Normal(logits.loc.float(), logits.scale.float()) @@ -785,6 +777,11 @@ def _train_ppo_transition(self, losses, profile, epoch): losses["filtered_fraction"] = 1.0 - kept_fraction else: keep_idx = valid_idx + if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1: + # Ranks mask different row counts: trim to the global minimum so minibatch counts stay in sync + kept_tensor = torch.tensor([keep_idx.numel()], device=device) + torch.distributed.all_reduce(kept_tensor, op=torch.distributed.ReduceOp.MIN) + keep_idx = keep_idx[torch.randperm(keep_idx.numel(), device=device)[: kept_tensor.item()]] if config["min_batch_size"] is not None: num_missing = config["min_batch_size"] - keep_idx.numel() @@ -1348,7 +1345,7 @@ def download(self): class WandbLogger: - def __init__(self, args, load_id=None, resume="allow", upload_config=True, disable_meta=False): + def __init__(self, args, load_id=None): import wandb # run_name is the run id: with resume="allow" wandb attaches to the @@ -1360,10 +1357,10 @@ def __init__(self, args, load_id=None, resume="allow", upload_config=True, disab group=args["wandb_group"], allow_val_change=True, save_code=False, - resume=resume, - config=args if upload_config else None, + resume="allow", + config=args, tags=[args["tag"]] if args["tag"] is not None else [], - settings=wandb.Settings(console="off", x_disable_meta=disable_meta), + settings=wandb.Settings(console="off"), ) self.wandb = wandb self.run_id = wandb.run.id @@ -1372,10 +1369,6 @@ def __init__(self, args, load_id=None, resume="allow", upload_config=True, disab def log(self, logs, step): self.wandb.log(logs, step=step) - def finish(self): - """End the wandb session without the model upload and early_stop that close() adds.""" - self.wandb.finish() - def upload_model(self, model_path): artifact = self.wandb.Artifact(self.run_id, type="model") artifact.add_file(model_path) @@ -1420,6 +1413,25 @@ def close(self, model_path, early_stop): self.local_writer.close() +class TensorBoardMirror: + """A remote run logger that also writes local TensorBoard event files.""" + + def __init__(self, remote_logger, tensorboard_logger): + self.remote_logger = remote_logger + self.tensorboard_logger = tensorboard_logger + + def __getattr__(self, name): + return getattr(self.remote_logger, name) + + def log(self, logs, step): + self.remote_logger.log(logs, step) + self.tensorboard_logger.log(logs, step) + + def close(self, model_path, early_stop): + self.tensorboard_logger.close(model_path, early_stop) + self.remote_logger.close(model_path, early_stop) + + def _get_git_metadata(): git_metadata = { "commit_hash": os.environ.get("GITHUB_SHA") or os.environ.get("COMMIT_SHA"), @@ -1610,6 +1622,8 @@ def train(env_name, args=None, vecenv=None, policy=None, logger=None, early_stop torch.manual_seed(torch_seed) vecenv = vecenv or load_env(env_name, args, seed=env_seed) policy = policy or load_policy(args, vecenv, env_name) + if isinstance(policy, pufferlib.ensemble.EnsemblePolicy): + raise pufferlib.APIUsageError("load_model_path is an ensemble checkpoint; ensembles are evaluation-only") if "LOCAL_RANK" in os.environ: args["train"]["device"] = "cuda" @@ -1635,11 +1649,13 @@ def train(env_name, args=None, vecenv=None, policy=None, logger=None, early_stop logger = NeptuneLogger(args) elif args["wandb"]: logger = WandbLogger(args) - elif args["tb"]: - logger = TensorBoardLogger( + if args["tb"]: + tensorboard_logger = TensorBoardLogger( run_id=args["run_name"], experiment_dir=args["train"]["data_dir"], ) + has_remote_logger = args["neptune"] or args["wandb"] + logger = TensorBoardMirror(logger, tensorboard_logger) if has_remote_logger else tensorboard_logger train_config = dict(**args["train"], env=env_name, eval=args.get("eval", {}), run_name=args["run_name"]) if torch.distributed.is_initialized(): @@ -1858,9 +1874,8 @@ def eval( base_args, checkpoint_config_path = drive_benchmark.load_checkpoint_architecture(args) base_args["env"]["eval_training_render"] = eval_training_render - wandb_run_identity = ( - drive_benchmark.load_checkpoint_run_identity(checkpoint_config_path) if report_to_wandb else None - ) + if report_to_wandb: + drive_eval_report.load_run_identity(checkpoint_config_path) if eval_output_dir is None: run_dir = drive_benchmark.resolve_run_dir(base_args["load_model_path"]) eval_output_dir = os.path.join(run_dir, eval_config["output_dir_name"]) @@ -1871,6 +1886,7 @@ def eval( failure_replay_csv = os.path.abspath(failure_replay_csv) failure_replay_output_dir = os.path.dirname(failure_replay_csv) benchmark_results = {} + selfplay_summary_path = None evaluation_policy_cache = {"policy": policy} for benchmark in benchmarks: if benchmark.get("simulation_mode") in drive_cosim_eval.COSIM_SIMULATION_MODES: @@ -1976,6 +1992,8 @@ def eval( "episodes": summaries, "summary": summary, } + if benchmark["name"] == drive_eval_report.SELFPLAY_BENCHMARK_NAME and summary is not None: + selfplay_summary_path = os.path.join(benchmark_output_dir, drive_eval_report.SELFPLAY_SUMMARY_FILE_NAME) if render_scenarios: max_renderer_count = ( @@ -1997,29 +2015,17 @@ def eval( evaluation_policy_cache=evaluation_policy_cache, ) - if wandb_run_identity is not None: - report_eval_to_wandb(args, benchmark_results, wandb_run_identity, eval_config["output_dir_name"]) + if report_to_wandb and selfplay_summary_path is not None: + drive_eval_report.report_headline_metric( + drive_eval_report.EVAL_TYPE_SELFPLAY, selfplay_summary_path, os.path.dirname(checkpoint_config_path) + ) + elif report_to_wandb: + print( + f"wandb: nothing reported; only the {drive_eval_report.SELFPLAY_BENCHMARK_NAME} benchmark has a headline metric." + ) return benchmark_results -def report_eval_to_wandb(args, benchmark_results, wandb_run_identity, output_dir_name): - """Attach standalone eval results to the wandb run that trained the checkpoint.""" - metrics = drive_benchmark.summarize_benchmark_metrics(benchmark_results, f"final_{output_dir_name}_") - if not metrics: - print("No evaluation metrics to report to wandb.") - return - - run_args = copy.copy(args) - run_args.update(wandb_run_identity) - - logger = WandbLogger(run_args, resume="must", upload_config=False, disable_meta=True) - try: - logger.log(metrics, step=None) - finally: - logger.finish() - print(f"Reported {len(metrics)} eval metrics to wandb run {wandb_run_identity['run_name']}") - - def sweep(args=None, env_name=None): args = args or load_config(env_name) if not args["wandb"] and not args["neptune"] and not args["tb"]: @@ -2805,16 +2811,18 @@ def load_policy(args, vecenv, env_name=""): elif args["wandb"]: path = WandbLogger(args, load_id).download() - state_dict = torch.load(path, map_location=device) - policy.load_state_dict(clean_policy_state_dict(state_dict)) + checkpoint = torch.load(path, map_location=device) + policy = pufferlib.ensemble.load_policy_weights(policy, checkpoint) load_path = args["load_model_path"] if load_path == "latest": load_path = max(glob.glob(f"experiments/{env_name}*.pt"), key=os.path.getctime) if load_path is not None: - state_dict = torch.load(load_path, map_location=device) - policy.load_state_dict(clean_policy_state_dict(state_dict)) + checkpoint = torch.load(load_path, map_location=device) + policy = pufferlib.ensemble.load_policy_weights(policy, checkpoint) + if isinstance(policy, pufferlib.ensemble.EnsemblePolicy): + print(f"Loaded an ensemble of {policy.member_count} policies from {load_path}") # state_path = os.path.join(*load_path.split('/')[:-1], 'state.pt') # optim_state = torch.load(state_path)['optimizer_state_dict'] # pufferl.optimizer.load_state_dict(optim_state) diff --git a/pufferlib/pytorch.py b/pufferlib/pytorch.py index 03ed92ed88..85aa749b57 100644 --- a/pufferlib/pytorch.py +++ b/pufferlib/pytorch.py @@ -126,6 +126,17 @@ def layer_init(layer, std=np.sqrt(2), bias_const=0.0): return layer +def clean_state_key(key): + prefixes = ("module.", "_orig_mod.") + while key.startswith(prefixes): + key = key.split(".", 1)[1] + return key + + +def clean_policy_state_dict(state_dict): + return {clean_state_key(k): v for k, v in state_dict.items()} + + # taken from torch.distributions.Categorical def log_prob(logits, value): value = value.long().unsqueeze(-1) diff --git a/pufferlib/resources/drive/binaries/carla/opendrive__Town01.bin b/pufferlib/resources/drive/binaries/carla/opendrive__Town01.bin index db51075617..fb3b23e777 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 3184cf7578..0b76f4e0a8 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 9ba4881637..158154d3fd 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 c50074700b..fe82e33b86 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 f413079330..067a53f942 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 04a1b9542f..07242d10ca 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 2a30a55a92..d149df9cca 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 c6587b07bc..4823de45a1 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 a40397adf3..886bc34e56 100644 --- a/pufferlib/vector.py +++ b/pufferlib/vector.py @@ -842,7 +842,7 @@ def make(env_creator_or_creators, env_args=None, env_kwargs=None, backend=Puffer # TODO: First step action space check - return backend(env_creators, env_args, env_kwargs, num_envs, **kwargs) + return backend(env_creators, env_args, env_kwargs, num_envs, seed=seed, **kwargs) def check_envs(envs, driver): diff --git a/pufferlib/viz.py b/pufferlib/viz.py index 7b8aa1a255..90d9ef3b5c 100644 --- a/pufferlib/viz.py +++ b/pufferlib/viz.py @@ -8,6 +8,7 @@ from matplotlib.patches import Circle import os import json +import re import zlib import base64 import struct @@ -80,6 +81,9 @@ OBS_QUANTIZE_FRAME_CHUNK = 64 REPLAY_COMPRESSION_LEVEL = 1 REPLAY_HEADER_PROBE_BYTES = 1 << 16 +PAYLOAD_CHUNK_PATTERN = re.compile( + rb'' +) REPLAY_INFLATE_MAX_ROUNDS = 1 << 12 SIMULATOR_FIGSIZE = (20.0, 20.0) SIMULATOR_DPI = 100 @@ -489,6 +493,7 @@ def unpack_obs( agent_idx: int = 0, obs_partner_relative_velocity: bool = False, obs_lane_speed_limit: bool = False, + obs_traffic_light_yellow_remaining: bool = False, ): """ Unpack the flattened observation into ego, map, partner, and traffic-control views. @@ -511,7 +516,9 @@ def unpack_obs( lane_feature_size = binding.LANE_FEATURES + (binding.LANE_SPEED_LIMIT_FEATURES if obs_lane_speed_limit else 0) boundary_feature_size = binding.BOUNDARY_FEATURES # Traffic control obs - traffic_control_feature_size = binding.TRAFFIC_CONTROL_FEATURES + traffic_control_feature_size = binding.TRAFFIC_CONTROL_FEATURES + ( + binding.TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES if obs_traffic_light_yellow_remaining else 0 + ) lane_segment_count = compute_effective_road_obs_count(obs_slots_lane_n, obs_dropout_lane) boundary_segment_count = compute_effective_road_obs_count(obs_slots_boundary_n, obs_dropout_boundary) @@ -593,6 +600,7 @@ def plot_observation( true_width_m=None, obs_partner_relative_velocity=False, obs_lane_speed_limit=False, + obs_traffic_light_yellow_remaining=False, ) -> np.ndarray: """Plot observation in ego-centric frame. @@ -615,6 +623,7 @@ def plot_observation( agent_idx=agent_idx, obs_partner_relative_velocity=obs_partner_relative_velocity, obs_lane_speed_limit=obs_lane_speed_limit, + obs_traffic_light_yellow_remaining=obs_traffic_light_yellow_remaining, ) scales = _obs_scales( obs_norm_goal_offset_m=obs_norm_goal_offset_m, @@ -786,7 +795,8 @@ def plot_observation( for i in range(traffic_controls_obs.shape[0]): if np.all(traffic_controls_obs[i] == 0): continue - rel_x1, rel_y1, rel_x2, rel_y2, _, control_type, state = traffic_controls_obs[i] + rel_x1, rel_y1, rel_x2, rel_y2 = traffic_controls_obs[i][:4] + control_type, state = traffic_controls_obs[i][-2:] control_type = int(control_type) if control_type == binding.TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT: ax.plot( @@ -1022,6 +1032,7 @@ def encode_interactive_replay(scenario, replay): "active_count": int(replay["raw_action"].shape[1]), "obs_dim": int(replay["obs"].shape[2]) if replay.get("obs") is not None else 0, "obs_layout": "agent_dim_frame_delta" if replay.get("obs") is not None else None, + "obs_agents": replay.get("obs_agents"), # None: every agent's observations are stored "action_type": env_cfg.get("action_type", "continuous"), "dynamics_model": env_cfg.get("dynamics_model", "classic"), "num_goals": int(env_cfg["num_goals"]), @@ -1035,7 +1046,12 @@ def encode_interactive_replay(scenario, replay): + (int(binding.LANE_SPEED_LIMIT_FEATURES) if env_cfg.get("obs_lane_speed_limit") else 0), "lane_speed_limit_idx": int(binding.LANE_FEATURES) if env_cfg.get("obs_lane_speed_limit") else -1, "boundary_features": int(binding.BOUNDARY_FEATURES), - "traffic_features": int(binding.TRAFFIC_CONTROL_FEATURES), + "traffic_features": int(binding.TRAFFIC_CONTROL_FEATURES) + + ( + int(binding.TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES) + if env_cfg.get("obs_traffic_light_yellow_remaining") + else 0 + ), "lane_count": int(lane_count), "boundary_count": int(boundary_count), "traffic_obs_count": int(env_cfg["obs_slots_traffic_controls_n"]), @@ -1232,6 +1248,8 @@ def _render_interactive_replay_payload(compressed_payload, filename): const PARTNER_BLINDNESS_OUTLINE_COLOR = "#6d28d9"; const PHANTOM_BRAKING_OUTLINE_COLOR = "#b45309"; const INFRACTION_METRIC_COUNT = 4; + // [agent idx, first frame, end frame, name]: infractions the embedding page reports besides the env's own metrics + const AGENT_MARKS = window.REPLAY_AGENT_MARKS || []; const DEFAULT_GOAL_RADIUS_METERS = 2; const SVG_PLAY = ''; const SVG_PAUSE = ''; @@ -1404,7 +1422,7 @@ def _render_interactive_replay_payload(compressed_payload, filename): for (let metricIdx=0;metricIdx 0) return true; } - return false; + return AGENT_MARKS.some(mark => mark[0] === idx && frame >= mark[1] && frame < mark[2]); } function rr(x, y, w, h, r) { ctx.beginPath(); if (ctx.roundRect) ctx.roundRect(x, y, w, h, r); else ctx.rect(x, y, w, h); } function drawAgentBody(a, outline) { @@ -1653,7 +1671,7 @@ def _render_interactive_replay_payload(compressed_payload, filename): 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 onset < f ? `${name} since step ${onset}` : name; ["COLLISION", "OFFROAD", "RED LIGHT", "STOP SIGN"].forEach((name, metricIdx) => { const onset = infractionOnset(f, agent.idx, metricIdx); - if (onset < 0) return; - warnings.push(onset < f ? `${name} since step ${onset}` : name); + if (onset >= 0) warnings.push(sinceOnset(name, onset)); + }); + AGENT_MARKS.forEach(mark => { + if (mark[0] === agent.idx && f >= mark[1] && f < mark[2]) warnings.push(sinceOnset(mark[3], mark[1])); }); 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(''); } @@ -2005,10 +2026,21 @@ def save_interactive_replay_zlib(scenario, replay, filename): return compressed_payload +def read_replay_source(path): + """-> the compressed replay bytes of a .replay.zlib, or of a rendered page (its base64 payload chunks).""" + with open(path, "rb") as source_file: + data = source_file.read() + if not str(path).endswith(".html"): + return data + chunks = PAYLOAD_CHUNK_PATTERN.findall(data) + if not chunks: + raise ValueError(f"{path}: no replay payload chunks in the page") + return b"".join(base64.b64decode(chunk) for chunk in chunks) + + def render_interactive_replay_zlib(replay_path, filename): - with open(replay_path, "rb") as replay_file: - compressed_payload = replay_file.read() - _render_interactive_replay_payload(compressed_payload, filename) + """replay_path: a .replay.zlib, or a rendered page whose embedded payload is rendered again.""" + _render_interactive_replay_payload(read_replay_source(replay_path), filename) def replay_step_link(step, seconds=None): @@ -2020,6 +2052,13 @@ def replay_step_link(step, seconds=None): ) +def replay_agent_marks(marks): + """" + + def _inflate_through(decompressor, raw, needed_bytes): rounds = 0 while len(raw) < needed_bytes and decompressor.unconsumed_tail and rounds < REPLAY_INFLATE_MAX_ROUNDS: @@ -2042,10 +2081,9 @@ def _inflate_replay_prefix(compressed, chunk_names): def read_replay_zlib(replay_path, chunk_names=None): - """-> (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: - compressed = replay_file.read() + """-> (header dict, chunk name -> ndarray view) of a saved .replay.zlib or of a rendered page; chunk_names + limits the chunks returned and the inflated prefix of the stream.""" + compressed = read_replay_source(replay_path) raw = zlib.decompress(compressed) if chunk_names is None else _inflate_replay_prefix(compressed, chunk_names) header_len = struct.unpack(".replay.zlib files of a -nuPlan run that was recorded before the planner wrote it, then re-render with render_obs_html.py. - -usage: python scripts/eval/add_nuplan_human_ghost.py [--data-root DIR] [--maps-root DIR] [--tokens tok ...] - -Needs the nuPlan devkit (carl_nuplan venv) and the split's log .db files under -/nuplan-v1.1/splits/val. Scenarios are rebuilt with the run's own hydra scenario mapping / -vehicle parameters; the bin-frame origin is recovered from the replay's frame-0 ego (synced from nuPlan). -""" - -import argparse -import glob -import os -import sys -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 -from pufferlib.ocean.cosim import nuplan_bridge as nb # noqa: E402 - -EGO_HEADING_IDX = 3 # agent_f32 columns: x, y, z, heading, ... -FRAME0_HEADING_TOL_RAD = 1e-3 # replay frame 0 is the nuPlan initial ego state, synced verbatim - - -def token_log_names(group_dir): - rows = {} - for csv in glob.glob(f"{group_dir}/simulation/*/20*/aggregator_metric/*_weighted_average_metrics_*.csv"): - agg = pd.read_csv(csv) - for _, row in agg[agg["scenario_type"] != "final_score"].iterrows(): - rows[str(row["scenario"])] = str(row["log_name"]) - return rows - - -def run_config(group_dir): - configs = glob.glob(f"{group_dir}/simulation/*/20*/code/hydra/config.yaml") - if not configs: - raise SystemExit(f"no code/hydra/config.yaml under {group_dir}") - return yaml.safe_load(open(configs[0])) - - -def build_scenarios(cfg, data_root, maps_root, wanted): - """token -> NuPlanScenario for every (token, log_name) in `wanted`, extracted like the run did.""" - from nuplan.common.actor_state.vehicle_parameters import VehicleParameters - from nuplan.planning.scenario_builder.nuplan_db.nuplan_scenario_builder import NuPlanScenarioBuilder - from nuplan.planning.scenario_builder.nuplan_db.nuplan_scenario_utils import ScenarioMapping - from nuplan.planning.scenario_builder.scenario_filter import ScenarioFilter - from nuplan.planning.utils.multithreading.worker_sequential import Sequential - - builder_cfg = cfg["scenario_builder"] - mapping_cfg = builder_cfg["scenario_mapping"] - vehicle_cfg = {k: v for k, v in builder_cfg["vehicle_parameters"].items() if not k.startswith("_")} - log_names = sorted(set(wanted.values())) - db_files = [str(Path(data_root) / "nuplan-v1.1" / "splits" / "val" / f"{log_name}.db") for log_name in log_names] - missing = [db for db in db_files if not os.path.exists(db)] - if missing: - raise SystemExit(f"{len(missing)} log db(s) missing locally, e.g. {missing[0]}") - builder = NuPlanScenarioBuilder( - data_root=str(Path(data_root) / "nuplan-v1.1" / "splits" / "val"), - map_root=str(maps_root), - sensor_root=None, - db_files=db_files, - map_version=builder_cfg["map_version"], - include_cameras=False, - verbose=False, - scenario_mapping=ScenarioMapping(mapping_cfg["scenario_map"], mapping_cfg.get("subsample_ratio_override")), - vehicle_parameters=VehicleParameters(**vehicle_cfg), - ) - scenario_filter = ScenarioFilter( - scenario_types=None, - scenario_tokens=sorted(wanted.keys()), - log_names=log_names, - map_names=None, - num_scenarios_per_type=None, - limit_total_scenarios=None, - timestamp_threshold_s=None, - ego_displacement_minimum_m=None, - expand_scenarios=False, - remove_invalid_goals=False, - shuffle=False, - ) - return {scenario.token: scenario for scenario in builder.get_scenarios(scenario_filter, Sequential())} - - -def replay_ghost(replay_path, scenario): - """(frames, 1, 5) ghost in the replay's bin frame, or None with a reason when the replay can't be matched.""" - header, chunks = viz.read_replay_zlib(replay_path) - if header["active_count"] != 1: - return None, f"active_count {header['active_count']} != 1" - ego0 = chunks["agent_f32"][0, 0] - initial = scenario.get_ego_state_at_iteration(0) - heading_delta = float(initial.center.heading) - float(ego0[EGO_HEADING_IDX]) - heading_error = abs((heading_delta + np.pi) % (2.0 * np.pi) - np.pi) - if heading_error > FRAME0_HEADING_TOL_RAD: - return None, f"frame-0 heading mismatch {heading_error:.4f} rad" - origin = nb.NuPlanTransform(float(initial.center.x) - float(ego0[0]), float(initial.center.y) - float(ego0[1])) - boxes = nb.logged_ego_boxes(scenario, origin) - frames = int(header["frames"]) - ghost = np.zeros((frames, 1, 5), np.float32) - count = min(frames, len(boxes)) - ghost[:count, 0] = boxes[:count] - return ghost, None - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("group_dir") - ap.add_argument("--data-root", default=os.environ.get("NUPLAN_DATA_ROOT")) - ap.add_argument("--maps-root", default=os.environ.get("NUPLAN_MAPS_ROOT")) - ap.add_argument("--tokens", nargs="*", default=[]) - args = ap.parse_args() - if not args.data_root or not args.maps_root: - raise SystemExit("--data-root / --maps-root (or NUPLAN_DATA_ROOT / NUPLAN_MAPS_ROOT) required") - obs_dir = Path(args.group_dir) / "obs_html" - replays = {p.name[: -len(".replay.zlib")]: p for p in obs_dir.glob("*.replay.zlib")} - if args.tokens: - replays = {token: path for token, path in replays.items() if token in set(args.tokens)} - if not replays: - raise SystemExit(f"no obs_html/*.replay.zlib under {args.group_dir}") - log_names = token_log_names(args.group_dir) - wanted = {token: log_names[token] for token in replays if token in log_names} - scenarios = build_scenarios(run_config(args.group_dir), args.data_root, args.maps_root, wanted) - written, skipped = 0, [] - for token, path in sorted(replays.items()): - scenario = scenarios.get(token) - if scenario is None: - skipped.append(f"{token}: no scenario ({'no log_name' if token not in wanted else 'not built'})") - continue - ghost, reason = replay_ghost(path, scenario) - if ghost is None: - skipped.append(f"{token}: {reason}") - continue - viz.set_replay_ghost(path, ghost) - written += 1 - for line in skipped: - print(f"[add_nuplan_human_ghost] skipped {line}") - print( - f"[add_nuplan_human_ghost] {written}/{len(replays)} replays updated -> now run scripts/eval/render_obs_html.py {args.group_dir}" - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/eval/analyze_carla_cosim.py b/scripts/eval/analyze_carla_cosim.py index fc69a275a7..fec1438c30 100644 --- a/scripts/eval/analyze_carla_cosim.py +++ b/scripts/eval/analyze_carla_cosim.py @@ -9,7 +9,8 @@ holds result.json (leaderboard_evaluator --checkpoint), world_log/*.npz, telemetry/*.csv and carla_view/*.mp4; the records of one run pair with the newest per-route files in timestamp order (retried -routes leave the files of crashed attempts behind). --obs-html-dir: the run's obs replay gallery +routes leave the files of crashed attempts behind); a run tidy_carla_run.py flattened has no such folders, so the +cluster scripts build the report (REPORT=1) before the tidy. --obs-html-dir: the run's obs replay gallery (scripts/eval/render_carla_obs_html.py), linked per route. """ diff --git a/scripts/eval/analyze_nuplan_cosim.py b/scripts/eval/analyze_nuplan_cosim.py index 21f811d969..8c86b65e6b 100644 --- a/scripts/eval/analyze_nuplan_cosim.py +++ b/scripts/eval/analyze_nuplan_cosim.py @@ -3,14 +3,14 @@ usage: python scripts/eval/analyze_nuplan_cosim.py [--workers N] [--max-inline N] [--no-video] - is the run_simulation.py `group` directory (holds simulation///...). + is the run_simulation.py `group` directory, fresh (simulation///...) or +tidied by tidy_nuplan_run.py (details/). Writes /index.html (score table sorted worst-first, failure-category counts, per-scenario diagnosis, the worst --max-inline scenarios inline with a six-frame strip + speed plot + video), plus strips/, speed/, videos/ and scenarios.csv. Per-scenario work runs in a multiprocessing pool. """ import argparse -import glob import math import sys from multiprocessing import Pool @@ -30,6 +30,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent)) import nuplan_log_remap # noqa: E402 logs written on another machine (cluster paths, planner classes) +import nuplan_run_layout # noqa: E402 REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT)) @@ -151,11 +152,11 @@ def trajectory_stats(hist, scenario): return ego, expert, stats -def comfort_failures(sim_dir, token): +def comfort_failures(metrics_dir, token): """nuPlan's own within-bound flags per comfort sub-metric -> ['ego_lon_jerk range -2.7..5.0', ...].""" out = [] for name in COMFORT_METRICS: - f = Path(sim_dir) / "metrics" / f"{name}.parquet" + f = Path(metrics_dir) / f"{name}.parquet" if not f.exists(): continue df = pd.read_parquet(f) @@ -218,12 +219,12 @@ def diagnose(row, stats, comfort): def process_scenario(job): - sim_dir, token, row, report_dir, make_video = job + metrics_dir, token, row, report_dir, make_video = job report_dir = Path(report_dir) log = nuplan_log_remap.load_log(row["log_path"]) hist, scenario = log.simulation_history, log.scenario ego, expert, stats = trajectory_stats(hist, scenario) - categories, reasons = diagnose(row, stats, comfort_failures(sim_dir, token)) + categories, reasons = diagnose(row, stats, comfort_failures(metrics_dir, token)) n = len(hist.data) picks = [0, n // 5, 2 * n // 5, 3 * n // 5, 4 * n // 5, n - 1] if stats["min_agent_dist_iter"] not in picks: @@ -275,21 +276,18 @@ def process_scenario(job): def collect_jobs(group_dir, report_dir, make_video): + agg = pd.read_csv(nuplan_run_layout.aggregated_metrics_csv(group_dir)) + per_scenario = agg[(agg["scenario_type"] != "final_score") & (agg["scenario"] != agg["scenario_type"])] + logs = nuplan_run_layout.simulation_logs(group_dir) + metrics_dir = str(nuplan_run_layout.metrics_dir(group_dir)) jobs = [] - for sim_dir in sorted(glob.glob(f"{group_dir}/simulation/*/20*")): - agg_files = glob.glob(f"{sim_dir}/aggregator_metric/*_weighted_average_metrics_*.csv") - if not agg_files: + for _, row in per_scenario.iterrows(): + token = str(row["scenario"]) + if token not in logs: continue - agg = pd.read_csv(agg_files[0]) - per_scenario = agg[(agg["scenario_type"] != "final_score") & (agg["scenario"] != agg["scenario_type"])] - logs = {Path(p).parent.name: p for p in glob.glob(f"{sim_dir}/simulation_log/**/*.msgpack.xz", recursive=True)} - for _, row in per_scenario.iterrows(): - token = str(row["scenario"]) - if token not in logs: - continue - row = row.to_dict() - row["log_path"] = logs[token] - jobs.append((sim_dir, token, row, str(report_dir), make_video)) + row = row.to_dict() + row["log_path"] = str(logs[token]) + jobs.append((metrics_dir, token, row, str(report_dir), make_video)) return jobs diff --git a/scripts/eval/build_ensemble_checkpoint.py b/scripts/eval/build_ensemble_checkpoint.py new file mode 100644 index 0000000000..60b4402812 --- /dev/null +++ b/scripts/eval/build_ensemble_checkpoint.py @@ -0,0 +1,45 @@ +"""Pack trained PufferDrive checkpoints into one ensemble checkpoint folder for the evals. + +Usage (from the repo root, venv active): + python scripts/eval/build_ensemble_checkpoint.py --output-dir experiments/k_scaled_0073_ensemble \ + experiments/k_scaled_0073_1000/final_model.pt experiments/k_scaled_0073_2000/final_model.pt + +The folder gets final_model.pt with every member's weights and config.yaml, the first member's training config, so the +evals build the env and the policy architecture from it; point RUN_DIR / load_model_path at the folder like at a +training run. At every policy step the members' action distributions and values are averaged and the action is +selected from the averaged distribution. +""" + +import argparse +import os + +from pufferlib.ensemble import write_ensemble_checkpoint + + +def parse_arguments(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("member_checkpoints", nargs="+", help="Trained .pt checkpoints, config.yaml beside each") + parser.add_argument("--output-dir", required=True, help="Folder to write final_model.pt and config.yaml into") + parser.add_argument("--run-name", help="run_name of the written config.yaml (default: the output folder name)") + parser.add_argument( + "--allow-env-mismatch", + action="store_true", + help="Accept members trained with different env configs; the evals use the first member's env config", + ) + return parser.parse_args() + + +def main(): + arguments = parse_arguments() + run_name = arguments.run_name or os.path.basename(os.path.normpath(arguments.output_dir)) + try: + checkpoint_path, config_path = write_ensemble_checkpoint( + arguments.member_checkpoints, arguments.output_dir, run_name, arguments.allow_env_mismatch + ) + except (ValueError, FileNotFoundError, FileExistsError) as error: + raise SystemExit(f"build_ensemble_checkpoint: {error}") + print(f"ensemble of {len(arguments.member_checkpoints)} checkpoints -> {checkpoint_path} (config {config_path})") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval/carla_run_layout.py b/scripts/eval/carla_run_layout.py new file mode 100644 index 0000000000..ea1a4a7a95 --- /dev/null +++ b/scripts/eval/carla_run_layout.py @@ -0,0 +1,119 @@ +"""Where the files of a CARLA longest6 run (scripts/kesai/11_carla_longest6_v2.sh, 14_carla_longest6_v1.sh) live. +While the run drives, every route run has its own folder, routes/route__s/: result.json (the leaderboard +record), evaluator_attempt.log, carla_server_attempt.log and, per attempt, telemetry/.csv, +world_log/.npz, obs_html/.replay.zlib and carla_view/.mp4 (retried routes leave the files of +crashed attempts behind; the newest pair with the records). scripts/eval/tidy_carla_run.py flattens the folders to +routes/[_rep].result.json plus .telemetry.csv / .world_log.npz / .carla_view.mp4 of the kept route +runs; the rendered pages in obs_html/ are the only replays left.""" + +import json +import re +from pathlib import Path + + +ROUTES_DIR = "routes" +OBS_HTML_DIR = "obs_html" +RAW_RESULT_NAME = "result.json" +RESULT_SUFFIX = ".result.json" +REPLAY_SUFFIX = ".replay.zlib" +RECORD_FILE_KINDS = (("telemetry", ".csv"), ("world_log", ".npz"), ("obs_html", REPLAY_SUFFIX), ("carla_view", ".mp4")) +ROUTE_ID_PATTERN = re.compile(r"RouteScenario_(\d+)(?:_rep(\d+))?") +SEED_SUFFIX_PATTERN = re.compile( + r"_s(\d+)$" +) # route_08_s1: seed 1 of route 8 (the eval scripts drive one pass per seed) +COMPLETED_STATUSES = ("Completed", "Perfect") +COLLISION_KEYS = ("collisions_pedestrian", "collisions_vehicle", "collisions_layout") +FLAG_INFRACTIONS = ( # page-name flag -> leaderboard infraction keys; min-speed and emergency-vehicle infractions carry none + ("COL", COLLISION_KEYS), + ("RED", ("red_light",)), + ("STOP", ("stop_infraction",)), + ("OFF", ("outside_route_lanes",)), + ("DEV", ("route_dev",)), + ("BLOCK", ("vehicle_blocked",)), + ("TIMEOUT", ("route_timeout", "scenario_timeouts")), +) + + +def infraction_messages(record, key): + return record.get("infractions", {}).get(key, []) or [] + + +def route_failed(record): + """A route run worth its replay page and logs: a flagged leaderboard infraction or a status other than + Completed / Perfect. Min-speed and emergency-vehicle infractions alone do not count (most v2 routes have some).""" + if record.get("status") not in COMPLETED_STATUSES: + return True + return any(infraction_messages(record, key) for _, keys in FLAG_INFRACTIONS for key in keys) + + +def route_index_rep(record): + """-> (route id, repetition) of a leaderboard record ('RouteScenario_12_rep0').""" + match = ROUTE_ID_PATTERN.fullmatch(str(record.get("route_id", ""))) + if match is None: + raise ValueError(f"unexpected route_id {record.get('route_id')!r}") + return int(match.group(1)), int(match.group(2) or 0) + + +def seed_of(route_run): + """-> the seed of a route run name ('route_08_s1' -> '1'), None for a run from before the seeds layout.""" + match = SEED_SUFFIX_PATTERN.search(route_run) + return match.group(1) if match else None + + +def flat_stem(route_run, record_index): + """-> the flat file stem of a route run's record: route_08_s1 for its first record, route_08_rep1_s1 for the next.""" + if record_index == 0: + return route_run + match = SEED_SUFFIX_PATTERN.search(route_run) + if match is None: + return f"{route_run}_rep{record_index}" + return f"{route_run[: match.start()]}_rep{record_index}{match.group(0)}" + + +def read_records(result_path): + with open(result_path) as result_file: + return json.load(result_file)["_checkpoint"]["records"] + + +def is_raw(result_path): + return Path(result_path).name == RAW_RESULT_NAME + + +def route_runs(run_dir): + """-> [(route run name, its record file)]: the route folders of a driving run and the flat records of a tidied + one, sorted by name.""" + routes_dir = Path(run_dir) / ROUTES_DIR + runs = {path.parent.name: path for path in routes_dir.glob(f"route_*/{RAW_RESULT_NAME}")} + runs.update({path.name[: -len(RESULT_SUFFIX)]: path for path in routes_dir.glob(f"route_*{RESULT_SUFFIX}")}) + return sorted(runs.items()) + + +def raw_attempt_stems(route_dir): + """-> the timestamp stems of every attempt that left a telemetry, world log, replay or chase-cam file, oldest first.""" + stems = set() + for kind, suffix in RECORD_FILE_KINDS: + stems.update(path.name[: -len(suffix)] for path in (Path(route_dir) / kind).glob(f"*{suffix}")) + return sorted(stems) + + +def record_stems(result_path, record_count): + """-> one file stem per record, oldest first: the newest attempt stems of a route folder (an attempt that crashed + before writing leaves the oldest records with None), the flat stems once tidied.""" + result_path = Path(result_path) + if not is_raw(result_path): + return [flat_stem(result_path.name[: -len(RESULT_SUFFIX)], k) for k in range(record_count)] + stems = raw_attempt_stems(result_path.parent)[-record_count:] if record_count else [] + return [None] * (record_count - len(stems)) + stems + + +def record_file(result_path, kind, suffix, stem): + """-> the `kind` file of a record (its stem from record_stems), None when it was not written or not kept.""" + result_path = Path(result_path) + if stem is None: + return None + path = ( + result_path.parent / kind / f"{stem}{suffix}" + if is_raw(result_path) + else result_path.parent / f"{stem}.{kind}{suffix}" + ) + return path if path.is_file() else None diff --git a/scripts/eval/compare_nuplan_evals.py b/scripts/eval/compare_nuplan_evals.py new file mode 100644 index 0000000000..19d5889a39 --- /dev/null +++ b/scripts/eval/compare_nuplan_evals.py @@ -0,0 +1,111 @@ +"""Paired comparison of two nuPlan closed-loop eval runs of the PufferDrive planner, overall and per city. + +usage: python scripts/eval/compare_nuplan_evals.py [--worst N] + +Each group dir is a run_simulation.py `group` directory (scripts/kesai/9_nuPlan.sh writes them under $RUN_DIR/eval, +tidy_nuplan_run.py flattens them); its scenarios.csv is compared on the scenarios both runs contain. The city of +a scenario is read from the nuPlan log database under $NUPLAN_DATA_ROOT/nuplan-v1.1/splits/val when that exists. +""" + +import argparse +import functools +import os +import sqlite3 + +import numpy as np +import pandas as pd + +import nuplan_run_layout + +METRICS = ( + "score", + "no_ego_at_fault_collisions", + "drivable_area_compliance", + "driving_direction_compliance", + "ego_is_making_progress", + "time_to_collision_within_bound", + "ego_progress_along_expert_route", + "speed_limit_compliance", + "ego_is_comfortable", +) +BOOTSTRAP_SAMPLES = 5000 +SCORE_EPSILON = 1e-9 + + +@functools.lru_cache(maxsize=None) +def city_of_log(log_name): + db_dir = os.path.join(os.environ.get("NUPLAN_DATA_ROOT", ""), "nuplan-v1.1", "splits", "val") + db_path = os.path.join(db_dir, f"{log_name}.db") + if not os.path.isfile(db_path): + return "unknown" + connection = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True) + location = connection.execute("select location from log").fetchone()[0] + connection.close() + return location + + +def load_scenarios(group_dir): + scenarios = pd.read_csv(nuplan_run_layout.scenarios_csv(group_dir)) + scenarios["city"] = scenarios["log_name"].map(city_of_log) + return scenarios.set_index("scenario") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("baseline_group_dir") + parser.add_argument("candidate_group_dir") + parser.add_argument("--worst", type=int, default=10, help="scenarios listed at both ends of the delta ranking") + args = parser.parse_args() + + baseline = load_scenarios(args.baseline_group_dir) + candidate = load_scenarios(args.candidate_group_dir) + common = baseline.index.intersection(candidate.index) + print(f"baseline {os.path.basename(os.path.normpath(args.baseline_group_dir))}: {len(baseline)} scenarios") + print(f"candidate {os.path.basename(os.path.normpath(args.candidate_group_dir))}: {len(candidate)} scenarios") + print(f"compared on {len(common)} common scenarios\n") + baseline, candidate = baseline.loc[common], candidate.loc[common] + + rows = [] + for metric in METRICS: + delta = candidate[metric] - baseline[metric] + rows.append( + { + "metric": metric, + "baseline": baseline[metric].mean(), + "candidate": candidate[metric].mean(), + "delta": delta.mean(), + "improved": int((delta > SCORE_EPSILON).sum()), + "worsened": int((delta < -SCORE_EPSILON).sum()), + } + ) + print(pd.DataFrame(rows).to_string(index=False, float_format=lambda value: f"{value:.4f}")) + + per_city = pd.DataFrame( + { + "n": baseline.groupby("city").size(), + "baseline": baseline.groupby("city")["score"].mean(), + "candidate": candidate.groupby("city")["score"].mean(), + } + ) + per_city["delta"] = per_city["candidate"] - per_city["baseline"] + print("\nscore per city:") + print(per_city.to_string(float_format=lambda value: f"{value:.4f}")) + + score_delta = candidate["score"] - baseline["score"] + rng = np.random.default_rng(0) + bootstrap_means = [ + rng.choice(score_delta.values, len(score_delta), replace=True).mean() for _ in range(BOOTSTRAP_SAMPLES) + ] + low, high = np.percentile(bootstrap_means, [2.5, 97.5]) + print(f"\nscore delta {score_delta.mean():+.4f} (paired bootstrap 95% CI [{low:+.4f}, {high:+.4f}])") + + changed = score_delta[score_delta.abs() > SCORE_EPSILON].sort_values() + listing = pd.concat([changed.head(args.worst), changed.tail(args.worst)]).to_frame("delta") + listing = listing.join(baseline[["city", "score"]].rename(columns={"score": "baseline"})) + listing = listing.join(candidate[["score"]].rename(columns={"score": "candidate"})) + print(f"\n{len(changed)} scenarios changed; largest regressions and gains:") + print(listing.to_string(float_format=lambda value: f"{value:.3f}")) + + +if __name__ == "__main__": + main() diff --git a/scripts/eval/nuboard_failures.py b/scripts/eval/nuboard_failures.py deleted file mode 100644 index 3aab43beb8..0000000000 --- a/scripts/eval/nuboard_failures.py +++ /dev/null @@ -1,124 +0,0 @@ -"""Build a nuBoard folder that contains only the failing scenarios of a nuPlan run, then launch nuBoard on it. - -usage: python scripts/eval/nuboard_failures.py [--max-score 0.5] [--metric no_ego_at_fault_collisions ...] - [--out ] [--launch] [--port 5006] - -nuBoard has no score filter of its own, so this mirrors /simulation/// into - with symlinks to the selected scenarios' simulation logs, metric/aggregator tables filtered to -those scenarios, and a .nuboard descriptor. --launch starts nuBoard on the mirror (Hydra -run_nuboard.py from $NUPLAN_DEVKIT_ROOT); logs written on another machine are loaded through -scripts/eval/nuplan_log_remap.py (set NUPLAN_DATA_ROOT / NUPLAN_MAPS_ROOT). - -Selection: scenarios with aggregated score < --max-score, OR any --metric below 1.0. -""" - -import argparse -import glob -import os -import shutil -import sys -from pathlib import Path - -import pandas as pd - - -def select_failures(sim_dir, max_score, metrics): - agg = pd.read_csv(glob.glob(f"{sim_dir}/aggregator_metric/*_weighted_average_metrics_*.csv")[0]) - per_scenario = agg[(agg["scenario_type"] != "final_score") & (agg["scenario"] != agg["scenario_type"])] - mask = per_scenario["score"] < max_score - for metric in metrics: - if metric not in per_scenario: - raise SystemExit( - f"unknown metric {metric}; available: {[c for c in per_scenario.columns if c not in ('scenario', 'scenario_type', 'log_name', 'planner_name', 'aggregator_type', 'num_scenarios')]}" - ) - mask |= per_scenario[metric] < 1.0 - return per_scenario[mask].sort_values("score") - - -def build_mirror(sim_dir, out_dir, selected): - from nuplan.planning.nuboard.base.data_class import NuBoardFile - - sim_dir, out_dir = ( - Path(sim_dir).resolve(), - Path(out_dir).resolve(), - ) # absolute: symlink targets must not be relative - tokens = set(selected["scenario"].astype(str)) - if out_dir.exists(): - shutil.rmtree(out_dir) - (out_dir / "simulation_log").mkdir(parents=True) - linked = 0 - for log_path in glob.glob(f"{sim_dir}/simulation_log/**/*.xz", recursive=True): - rel = Path(log_path).relative_to(sim_dir / "simulation_log") - if rel.parent.name not in tokens: - continue - target = out_dir / "simulation_log" / rel - target.parent.mkdir(parents=True, exist_ok=True) - os.symlink(log_path, target) - linked += 1 - (out_dir / "metrics").mkdir() - for f in glob.glob(f"{sim_dir}/metrics/*.parquet"): - df = pd.read_parquet(f) - df[df["scenario_name"].astype(str).isin(tokens)].to_parquet(out_dir / "metrics" / Path(f).name) - (out_dir / "aggregator_metric").mkdir() - for f in glob.glob(f"{sim_dir}/aggregator_metric/*"): - if not f.endswith((".parquet", ".csv")): - continue - df = pd.read_parquet(f) if f.endswith(".parquet") else pd.read_csv(f) - if "scenario" in df: # per-scenario rows -> keep failures plus the per-type / final summary rows - keep = df["scenario"].astype(str).isin(tokens) - if "scenario_type" in df: - keep |= (df["scenario"] == df["scenario_type"]) | (df["scenario_type"] == "final_score") - df = df[keep] - if f.endswith(".parquet"): - df.to_parquet(out_dir / "aggregator_metric" / Path(f).name) - else: - df.to_csv(out_dir / "aggregator_metric" / Path(f).name, index=False) - NuBoardFile( - simulation_main_path=str(out_dir), - metric_main_path=str(out_dir), - metric_folder="metrics", - aggregator_metric_folder="aggregator_metric", - simulation_folder="simulation_log", - ).save_nuboard_file(out_dir / "failures.nuboard") - return linked - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("group_dir") - ap.add_argument("--max-score", type=float, default=0.5) - ap.add_argument( - "--metric", action="append", default=[], help="also select scenarios where this metric is < 1 (repeatable)" - ) - ap.add_argument("--out", default=None, help="mirror directory (default /nuboard_failures)") - ap.add_argument("--launch", action="store_true") - ap.add_argument("--port", type=int, default=5006) - args = ap.parse_args() - sim_dirs = sorted(glob.glob(f"{args.group_dir}/simulation/*/20*")) - if len(sim_dirs) != 1: - raise SystemExit( - f"expected exactly one simulation// under {args.group_dir}, found {sim_dirs}" - ) - selected = select_failures(sim_dirs[0], args.max_score, args.metric) - out_dir = Path(args.out or f"{args.group_dir}/nuboard_failures").resolve() - linked = build_mirror(sim_dirs[0], out_dir, selected) - print(f"[nuboard_failures] {len(selected)} scenarios selected, {linked} logs linked -> {out_dir}") - print(selected[["scenario", "scenario_type", "score"]].head(15).to_string(index=False)) - devkit = os.environ.get("NUPLAN_DEVKIT_ROOT", "") - cmd = f"python {devkit}/nuplan/planning/script/run_nuboard.py simulation_path=[{out_dir}] port_number={args.port}" - if not args.launch: - print("launch with:\n " + cmd) - return - sys.path.insert(0, str(Path(__file__).resolve().parent)) - import nuplan_log_remap # noqa: F401 patches scenario/map paths for logs written elsewhere - from nuplan.planning.simulation import simulation_log as _sl - - _sl.SimulationLog.load_data = classmethod(lambda cls, file_path: nuplan_log_remap.load_log(file_path)) - sys.argv = [sys.argv[0], f"simulation_path=[{out_dir}]", f"port_number={args.port}"] - import runpy - - runpy.run_path(f"{devkit}/nuplan/planning/script/run_nuboard.py", run_name="__main__") - - -if __name__ == "__main__": - main() diff --git a/scripts/eval/nuplan_comfort_breakdown.py b/scripts/eval/nuplan_comfort_breakdown.py new file mode 100644 index 0000000000..4f9e50c741 --- /dev/null +++ b/scripts/eval/nuplan_comfort_breakdown.py @@ -0,0 +1,49 @@ +"""Which ego_is_comfortable sub-metric fails how often in nuPlan closed-loop eval runs of the PufferDrive planner. +usage: python scripts/eval/nuplan_comfort_breakdown.py [ ...] +Each group dir is a run_simulation.py `group` directory, fresh or tidied by tidy_nuplan_run.py; the per-scenario +metric parquets under details/metrics (simulation///metrics before the tidy) are read. nuPlan's +longitudinal acceleration is taken at the vehicle center, +which the devkit computes as rear-axle accel + d*yaw_rate^2 + d*yaw_accel (d = rear axle to center), so a yaw +acceleration while accelerating reads as longitudinal acceleration. +""" + +import sys + +import pandas as pd + +import nuplan_run_layout + +SUBMETRICS = ( + "ego_lon_acceleration", + "ego_lat_acceleration", + "ego_jerk", + "ego_lon_jerk", + "ego_yaw_rate", + "ego_yaw_acceleration", +) +LON_ACCEL_MAX_MPS2 = 2.40 # ego_lon_acceleration_statistics.yaml + + +def main() -> None: + for group in sys.argv[1:]: + metrics_dir = nuplan_run_layout.metrics_dir(group) + if not metrics_dir.is_dir(): + print(f"{group}: no metrics directory") + continue + comfortable = pd.read_parquet(f"{metrics_dir}/ego_is_comfortable.parquet") + failed = int((comfortable["metric_score"] < 1).sum()) + print(f"{group.rstrip('/').split('/')[-1]}") + print(f" ego_is_comfortable fails {failed}/{len(comfortable)}") + for name in SUBMETRICS: + df = pd.read_parquet(f"{metrics_dir}/{name}.parquet") + within = df[f"abs_{name}_within_bounds_stat_value"].astype(bool) + fails = df[~within] + detail = "" + if name == "ego_lon_acceleration": + positive = int((fails[f"max_{name}_stat_value"] > LON_ACCEL_MAX_MPS2).sum()) + detail = f" ({positive} above +{LON_ACCEL_MAX_MPS2}, {len(fails) - positive} below the braking bound)" + print(f" {name:22s} {len(fails):4d} fail{detail}") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval/nuplan_run_layout.py b/scripts/eval/nuplan_run_layout.py new file mode 100644 index 0000000000..d354ba034c --- /dev/null +++ b/scripts/eval/nuplan_run_layout.py @@ -0,0 +1,72 @@ +"""Where the files of a nuPlan closed-loop run of the PufferDrive planner live. A fresh run is nuPlan's own tree, +simulation///{aggregator_metric,metrics,simulation_log}; after scripts/eval/tidy_nuplan_run.py +the run holds overview.csv, scenarios.csv, obs_html/, simulation_log/ and details/{aggregated_metrics.csv,metrics} +(runs tidied before 2026-11-09 keep their logs under details/simulation_log/). +Every lookup raises FileNotFoundError when the run has neither layout.""" + +import glob +from pathlib import Path + +DETAILS_DIR = "details" +SIMULATION_LOG_DIR = "simulation_log" +AGGREGATED_METRICS_CSV = "aggregated_metrics.csv" +SIMULATION_LOG_SUFFIX = ".msgpack.xz" + + +def raw_simulation_dir(group_dir): + """-> nuPlan's simulation// directory of an untidied run, None when there is none.""" + matches = sorted(glob.glob(f"{group_dir}/simulation/*/20*")) + if len(matches) > 1: + raise FileNotFoundError(f"expected one simulation// under {group_dir}, found {matches}") + return Path(matches[0]) if matches else None + + +def details_dir(group_dir): + """-> the directory holding metrics/ and simulation_log/: details/ once tidied, else nuPlan's simulation dir.""" + tidied = Path(group_dir) / DETAILS_DIR + if tidied.is_dir(): + return tidied + sim_dir = raw_simulation_dir(group_dir) + if sim_dir is None: + raise FileNotFoundError(f"{group_dir}: neither {DETAILS_DIR}/ nor simulation//") + return sim_dir + + +def metrics_dir(group_dir): + return details_dir(group_dir) / "metrics" + + +def aggregated_metrics_csv(group_dir): + """-> the aggregator's full table: per-scenario, per-scenario-type and final_score rows.""" + base = details_dir(group_dir) + if (base / AGGREGATED_METRICS_CSV).is_file(): + return base / AGGREGATED_METRICS_CSV + matches = sorted(base.glob("aggregator_metric/*_weighted_average_metrics_*.csv")) + if len(matches) != 1: + raise FileNotFoundError(f"expected one aggregator csv under {base}, found {matches}") + return matches[0] + + +def scenarios_csv(group_dir): + """-> CaRL's per-scenario score table, worst first.""" + for path in (Path(group_dir) / "scenarios.csv", details_dir(group_dir) / "aggregator_metric" / "scenarios.csv"): + if path.is_file(): + return path + raise FileNotFoundError(f"no scenarios.csv under {group_dir}") + + +def simulation_log_dir(group_dir): + """-> the directory holding the .msgpack.xz logs: simulation_log/ once tidied, else nuPlan's nested tree.""" + group_dir = Path(group_dir) + for tidied in (group_dir / SIMULATION_LOG_DIR, group_dir / DETAILS_DIR / SIMULATION_LOG_DIR): + if tidied.is_dir(): + return tidied + return details_dir(group_dir) / SIMULATION_LOG_DIR + + +def simulation_logs(group_dir): + """-> {scenario token: its simulation log}; nuPlan names the file .msgpack.xz in both layouts.""" + log_dir = simulation_log_dir(group_dir) + return { + path.name[: -len(SIMULATION_LOG_SUFFIX)]: path for path in sorted(log_dir.rglob(f"*{SIMULATION_LOG_SUFFIX}")) + } diff --git a/scripts/eval/render_alpasim_obs_html.py b/scripts/eval/render_alpasim_obs_html.py index 377803d051..749d433b8b 100644 --- a/scripts/eval/render_alpasim_obs_html.py +++ b/scripts/eval/render_alpasim_obs_html.py @@ -1,21 +1,25 @@ """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 +/obs_html/.replay.zlib) into one page per selected 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. +usage: python scripts/eval/render_alpasim_obs_html.py [--select imperfect|failed|all] [--prune] [--workers N] + +--select imperfect (default) pages the scenes scoring below 1 (and scenes without a score row), failed only the +scenes AlpaSim marks as failed (at-fault collision, offroad, corridor exit), all every scene. 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 pages the run wrote (.html) and the pages of unselected scenes are deleted; --prune deletes +every .replay.zlib once the pages are written (a page embeds the same bytes and is the replay source of a later +re-render). Reads the tidy_alpasim_run.py layout (/details/*, obs_html/), the 2026-10-12 one (alpasim/*, +obs_replays/) and the untidied AlpaSim trees of older runs. """ import argparse import html import json import os +import re import sys from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -24,9 +28,11 @@ 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 +from pufferlib import viz + FLAG_METRICS = ( ("COL", "collision_at_fault"), @@ -34,6 +40,10 @@ ("CORR", "left_corridor_laterally"), ) PROGRESS_FLAG = "PROG" +SELECT_MODES = ("imperfect", "failed", "all") +REPLAY_SUFFIX = ".replay.zlib" +GALLERY_INDEX = "index.html" +CLIPGT_PATTERN = re.compile(r"clipgt-[0-9a-f-]{36}$") # the scene id ends every page and replay name SCORED_METRICS = ( ("collision_at_fault", "at-fault collision"), ("offroad", "offroad"), @@ -62,13 +72,18 @@ 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()] + layout (details/, then the 2026-10-12 alpasim/), else of an AlpaSim wizard tree (alpasim/wizard_0, then the + older alpagym/ nesting).""" + obs_dirs = [ + path + for path in (run_dir / "obs_html", run_dir / "obs_replays", run_dir / "alpagym" / "obs_replays") + if path.is_dir() + ] if not obs_dirs: - raise SystemExit(f"no obs_replays/ under {run_dir} or {run_dir / 'alpagym'}") - tidied = run_dir / "alpasim" + raise SystemExit(f"no obs_html/ or obs_replays/ under {run_dir}") candidates = [ - (tidied / "results-summary.json", tidied / "metrics_unprocessed.parquet", tidied / "eval-config.yaml") + (base / "results-summary.json", base / "metrics_unprocessed.parquet", base / "eval-config.yaml") + for base in (run_dir / "details", run_dir / "alpasim") ] for wizard_dir in (run_dir / "alpasim" / "wizard_0", run_dir / "alpagym" / "alpasim" / "wizard_0"): aggregate = wizard_dir / "aggregate" @@ -213,6 +228,33 @@ def score_panel(row, criteria, timeline): ) +def selected(row, select): + """A scene without a summary row never finished scoring and counts as failed.""" + if select == "all" or row is None: + return True + return row["status"] == "fail" if select == "failed" else float(row["score"]) < 1.0 + + +def replay_sources(obs_dir): + """clipgt -> its .replay.zlib, else a page of the scene (the run's own or a rendered one, embedding the replay).""" + sources = {path.name[: -len(REPLAY_SUFFIX)]: path for path in sorted(obs_dir.glob(f"*{REPLAY_SUFFIX}"))} + for page in sorted(obs_dir.glob("*.html")): + match = CLIPGT_PATTERN.search(page.stem) + if match and match.group(0) not in sources: + sources[match.group(0)] = page + return sources + + +def scene_pages(obs_dir): + """clipgt -> every page of the scene in obs_dir (the run's .html, earlier renders).""" + pages = {} + for page in sorted(obs_dir.glob("*.html")): + match = CLIPGT_PATTERN.search(page.stem) + if match: + pages.setdefault(match.group(0), []).append(page) + return pages + + def page_name(clipgt, row): if row is None: return f"{clipgt}.html" @@ -243,43 +285,60 @@ def render_page(zlib_path, html_path, panel): 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("run_dir", help="output dir of 12_alpasim.sh (holds obs_html/ and details/)") + ap.add_argument("--select", choices=SELECT_MODES, default="imperfect", help="scenes that get a page") + ap.add_argument("--prune", action="store_true", help="delete every .replay.zlib once the pages are 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") + sources = replay_sources(obs_dir) + if not sources: + raise SystemExit(f"no {obs_dir}/*{REPLAY_SUFFIX} or rendered pages") 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()] + old_pages = scene_pages(obs_dir) + jobs = [ + (clipgt, source, obs_dir / page_name(clipgt, rows.get(clipgt))) + for clipgt, source in sources.items() + if selected(rows.get(clipgt), args.select) + ] def render(job): - clipgt, replay, html_path = job + clipgt, source, 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"]) + frames = int(viz.read_replay_zlib(source, 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() + render_page(source, html_path, panel) + for page in old_pages.get(clipgt, []): + if page != html_path: + page.unlink() # the run's own page, or a render under a stale name with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool: list(pool.map(render, jobs)) + rendered = {clipgt for clipgt, _, _ in jobs} + dropped_count = 0 + for clipgt, pages in old_pages.items(): + for page in pages if clipgt not in rendered else []: + page.unlink() # an unselected scene keeps no page + dropped_count += 1 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 jobs: + viz.build_gallery_index(str(obs_dir), file_metrics=file_metrics) + pruned_count = 0 if args.prune: - for replay in replays.values(): - replay.unlink() - unscored = sum(clipgt not in rows for clipgt in replays) + for source in sources.values(): + if source.name.endswith(REPLAY_SUFFIX): + source.unlink() + pruned_count += 1 + unscored = sum(clipgt not in rows for clipgt in rendered) 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'}" + f"[render_alpasim_obs_html] {len(jobs)} of {len(sources)} scenes paged ({args.select})" + f"{f', {unscored} without a score (no panel)' if unscored else ''}" + f"{f', {dropped_count} unselected pages removed' if dropped_count else ''}" + f"{f', {pruned_count} replays pruned' if pruned_count else ''} -> {obs_dir / GALLERY_INDEX}" ) diff --git a/scripts/eval/render_carla_obs_html.py b/scripts/eval/render_carla_obs_html.py index 7911c902c7..047e8bf623 100644 --- a/scripts/eval/render_carla_obs_html.py +++ b/scripts/eval/render_carla_obs_html.py @@ -1,19 +1,26 @@ -"""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). +"""Render the interactive observation replays of a CARLA longest6 run (scripts/kesai/11_carla_longest6_v2.sh, +14_carla_longest6_v1.sh: /routes/route_XX[_s]/obs_html/*.replay.zlib) into /obs_html/, one +page per selected route run 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] +usage: python scripts/eval/render_carla_obs_html.py [--routes ] [--select failures|all] + [--workers N] -Pages are named ds___route.html so the gallery lists the worst first +--select failures (default) renders the route runs with a flagged leaderboard infraction or a failed status +(carla_run_layout.route_failed; min-speed infractions alone do not count), all renders every route run. +Pages are named ds___route_s.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, every positioned infraction linking to the replay step where the ego was nearest. +infractions as a panel, every positioned infraction linking to the replay step where the ego was nearest +(red lights: to the step the route's telemetry flagged, their message carries the light's position); from +that step on the viewer draws the ego as infringing for MARK_FRAMES, under a LEADERBOARD warning chip. --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. +Retried routes leave the files of crashed attempts behind: the newest files pair with the records. On a run +tidy_carla_run.py flattened, a route's rendered page is its replay source (the page embeds it); a page whose +route was driven again under another name is replaced. """ import argparse -import glob +import csv import html import json import os @@ -28,20 +35,15 @@ REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT)) +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import carla_run_layout as layout +from carla_run_layout import COLLISION_KEYS, FLAG_INFRACTIONS, REPLAY_SUFFIX, infraction_messages, route_index_rep + from pufferlib import viz from pufferlib.ocean.cosim.carla_bridge import TOWN_OFFSETS -COLLISION_KEYS = ("collisions_pedestrian", "collisions_vehicle", "collisions_layout") -FLAG_INFRACTIONS = ( - ("COL", COLLISION_KEYS), - ("RED", ("red_light",)), - ("STOP", ("stop_infraction",)), - ("OFF", ("outside_route_lanes",)), - ("DEV", ("route_dev",)), - ("BLOCK", ("vehicle_blocked",)), - ("TIMEOUT", ("route_timeout", "scenario_timeouts")), -) +SELECT_MODES = ("failures", "all") PANEL_INFRACTIONS = ( ("collisions_pedestrian", "pedestrian collisions"), ("collisions_vehicle", "vehicle collisions"), @@ -59,29 +61,34 @@ 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): - return record.get("infractions", {}).get(key, []) or [] - - -def route_index_rep(record): - """-> (route id, repetition) of a leaderboard record ('RouteScenario_12_rep0').""" - match = ROUTE_ID_PATTERN.fullmatch(str(record.get("route_id", ""))) - if match is None: - raise ValueError(f"unexpected route_id {record.get('route_id')!r}") - return int(match.group(1)), int(match.group(2) or 0) +RED_LIGHT_FLAG_COLUMN = "carla_infr_red" # telemetry: 1 on the step the simulator-side red light test fired +EGO_AGENT_IDX = 0 +MARK_FRAMES = 40 # 2 s at the leaderboards' 20 ticks per second +MARK_NAMES = { + "collisions_pedestrian": "LEADERBOARD COLLISION", + "collisions_vehicle": "LEADERBOARD COLLISION", + "collisions_layout": "LEADERBOARD COLLISION", + "red_light": "LEADERBOARD RED LIGHT", + "stop_infraction": "LEADERBOARD STOP SIGN", + "route_dev": "LEADERBOARD ROUTE DEVIATION", + "vehicle_blocked": "LEADERBOARD BLOCKED", +} +GALLERY_INDEX = "index.html" +PAGE_KEY_PATTERN = re.compile(r"route\d+(?:_rep\d+)?(?:_s\d+)?$") # the route-run part of a page name + + +def page_key(record, seed=None): + """-> 'route08_rep1_s1': the part of a page name that names the route run, whatever its score and flags.""" + route_idx, rep = route_index_rep(record) + return f"route{route_idx:02d}" + (f"_rep{rep}" if rep else "") + (f"_s{seed}" if seed is not None else "") 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 "") + (f"_s{seed}" if seed is not None else "")) + parts.append(page_key(record, seed)) return "_".join(parts) + ".html" @@ -126,29 +133,56 @@ def score_panel(record, town, steps, seed=None): ) -def infraction_steps(record, replay_path, town, log_meta): - """message -> 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") +def red_light_steps(telemetry_path): + """Steps on which the route's telemetry flagged a run red light (row k of the CSV = replay frame k).""" + with open(telemetry_path, newline="") as telemetry_file: + rows = csv.DictReader(telemetry_file) + return [step for step, row in enumerate(rows) if float(row[RED_LIGHT_FLAG_COLUMN]) > 0.0] + + +def infraction_frames(record, replay_path, town, log_meta, telemetry_path=None): + """-> (message -> replay frame of every infraction that can be placed, the replay's frame count): red lights at + the frame the telemetry flagged, the other positioned ones where the ego was nearest (CARLA -> bin offset: the + world log's, else the town's). A frame at or past the count: the infraction happened after the replay ended.""" _, 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: + frame_count = len(ego_xy) + frames = {} + red_messages = infraction_messages(record, "red_light") + flagged_steps = red_light_steps(telemetry_path) if telemetry_path and red_messages else [] + if len(flagged_steps) == len(red_messages): + frames.update(zip(red_messages, flagged_steps)) + offset = log_meta.get("offset") or TOWN_OFFSETS.get(town) + nearest_keys = [key for key, _ in PANEL_INFRACTIONS if key != "red_light" and offset is not None] + for message in [message for key in nearest_keys for message in infraction_messages(record, key)]: 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 + nearest_frame = int(np.argmin(distance_m)) + # a place the replay's ego never came near was reached after its last frame + frames[message] = nearest_frame if distance_m[nearest_frame] <= NEAREST_STEP_MAX_M else frame_count + return frames, frame_count + + +def step_links(frames, frame_count, tick_dt): + """message -> the panel's seek link to the infraction's frame, or the note that the replay ended before it.""" + links = {} + for message, frame in frames.items(): + seconds = None if tick_dt is None else frame * tick_dt + links[message] = viz.replay_step_link(frame, seconds) if frame < frame_count else "beyond the replay" + return links + + +def ego_marks(record, frames, frame_count): + """-> viz.replay_agent_marks rows: the ego counts as infringing for MARK_FRAMES from each placed infraction on.""" + keyed_messages = [(key, message) for key in MARK_NAMES for message in infraction_messages(record, key)] + return [ + (EGO_AGENT_IDX, frames[message], MARK_FRAMES, MARK_NAMES[key]) + for key, message in keyed_messages + if frames.get(message, frame_count) < frame_count + ] def render_page(zlib_path, html_path, panel): @@ -177,61 +211,68 @@ def world_log_meta(npz_path): return json.loads(str(np.load(npz_path)["meta"])) -def newest_files(route_dir, subdir, suffix, count): - """The `count` newest (timestamp-named) files of /.""" - files = sorted(glob.glob(str(route_dir / subdir / f"*{suffix}"))) - return files[max(0, len(files) - count) :] +def rendered_pages(obs_dir): + """-> page key -> the route run's rendered page: the replay source once the run is tidied (a page embeds it).""" + pages = {} + for page in sorted(obs_dir.glob("*.html")): + match = PAGE_KEY_PATTERN.search(page.stem) + if page.name != GALLERY_INDEX and match: + pages[match.group(0)] = page + return pages -def collect_jobs(run_dir, towns): - """-> [(record, replay zlib path, town, world log meta, seed or None)] over routes/route_*, the newest replay per record.""" +def collect_jobs(run_dir, towns, select): + """-> [(record, replay source, town, world log meta, seed, telemetry csv or None)] over the selected records of + every route run: the newest replay of a route folder, or the route run's rendered page once tidied.""" + pages = rendered_pages(Path(run_dir) / layout.OBS_HTML_DIR) 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") - continue - with open(result_path) as result_file: - records = json.load(result_file)["_checkpoint"]["records"] - replays = newest_files(route_dir, "obs_html", ".replay.zlib", len(records)) - if len(replays) != len(records): - print(f"[{route_dir.name}] {len(records)} route records but {len(replays)} replays; pairing the newest") - if not replays: - continue - records = records[-len(replays) :] - world_logs = newest_files(route_dir, "world_log", ".npz", len(replays)) - for k, (record, replay) in enumerate(zip(records, replays)): - log_meta = world_log_meta(world_logs[k]) if k < len(world_logs) else {} + for route_run, result_path in layout.route_runs(run_dir): + seed = layout.seed_of(route_run) + records = layout.read_records(result_path) + for record, stem in zip(records, layout.record_stems(result_path, len(records))): + if select == "failures" and not layout.route_failed(record): + continue + replay = layout.record_file(result_path, layout.OBS_HTML_DIR, REPLAY_SUFFIX, stem) + source = replay or pages.get(page_key(record, seed)) + if source is None: + print(f"[{route_run}] no replay for {record.get('route_id')}, skipped") + continue + world_log = layout.record_file(result_path, "world_log", ".npz", stem) + log_meta = world_log_meta(world_log) if world_log else {} town = towns.get(route_index_rep(record)[0]) or log_meta.get("town", "") - jobs.append((record, replay, town, log_meta, seed)) + telemetry = layout.record_file(result_path, "telemetry", ".csv", stem) + jobs.append((record, source, town, log_meta, seed, telemetry)) return jobs def main(): ap = argparse.ArgumentParser() - ap.add_argument("run_dir", help="output dir of 11_carla_longest6_v2.sh (holds routes/route_*/)") + ap.add_argument("run_dir", help="output dir of 11_carla_longest6_v2.sh / 14_carla_longest6_v1.sh (holds routes/)") ap.add_argument("--routes", help="evaluated route xml, names the town of every route") + ap.add_argument("--select", choices=SELECT_MODES, default="failures", help="route runs that get a page") ap.add_argument("--workers", type=int, default=min(16, os.cpu_count() or 1)) args = ap.parse_args() run_dir = Path(args.run_dir) towns = route_towns(args.routes) if args.routes else {} - jobs = collect_jobs(run_dir, towns) + jobs = collect_jobs(run_dir, towns, args.select) if not jobs: - raise SystemExit(f"no routes/route_*/obs_html/*.replay.zlib with a result.json under {run_dir}") - obs_dir = run_dir / "obs_html" + raise SystemExit(f"no selected route run with a replay or a rendered page under {run_dir}") + obs_dir = run_dir / layout.OBS_HTML_DIR obs_dir.mkdir(exist_ok=True) + stale_pages = rendered_pages(obs_dir) pages = [ - (obs_dir / page_name(record, town, seed), record, replay, town, log_meta, seed) - for record, replay, town, log_meta, seed in jobs + (obs_dir / page_name(record, town, seed), record, replay, town, log_meta, seed, telemetry) + for record, replay, town, log_meta, seed, telemetry 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) - ) + html_path, record, replay, town, log_meta, seed, telemetry = page + frames, frame_count = infraction_frames(record, replay, town, log_meta, telemetry) + panel = score_panel(record, town, step_links(frames, frame_count, log_meta.get("tick_dt")), seed) + render_page(replay, html_path, panel + viz.replay_agent_marks(ego_marks(record, frames, frame_count))) + stale = stale_pages.get(page_key(record, seed)) + if stale is not None and stale != html_path: + stale.unlink() # the route run was driven again: its old page would stay in the gallery with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool: list(pool.map(render, pages)) @@ -239,7 +280,7 @@ def render(page): 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) - print(f"[render_carla_obs_html] {len(pages)} route pages rendered -> {obs_dir / 'index.html'}") + print(f"[render_carla_obs_html] {len(pages)} route pages rendered ({args.select}) -> {obs_dir / GALLERY_INDEX}") if __name__ == "__main__": diff --git a/scripts/eval/render_obs_html.py b/scripts/eval/render_obs_html.py index 26130fb559..041ef058a1 100644 --- a/scripts/eval/render_obs_html.py +++ b/scripts/eval/render_obs_html.py @@ -9,11 +9,12 @@ 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). 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. +--prune deletes every .replay.zlib afterwards: a rendered page embeds the same bytes and is the source of a +later re-render (e.g. after a viewer change), the other replays belong to unselected scenarios. Works on fresh +runs and on runs tidied by tidy_nuplan_run.py. """ import argparse -import glob import os import sys from concurrent.futures import ThreadPoolExecutor @@ -26,6 +27,10 @@ sys.path.insert(0, str(REPO_ROOT)) from pufferlib import viz # noqa: E402 +import nuplan_run_layout # noqa: E402 + +REPLAY_SUFFIX = ".replay.zlib" +GALLERY_INDEX = "index.html" FLAG_METRICS = ( ("COL", "no_ego_at_fault_collisions"), ("OFF", "drivable_area_compliance"), @@ -74,14 +79,15 @@ def scenario_rows(group_dir): - """token -> aggregator row (score + per-metric values) over every challenge of the run.""" - rows = {} - for csv in glob.glob(f"{group_dir}/simulation/*/20*/aggregator_metric/*_weighted_average_metrics_*.csv"): - agg = pd.read_csv(csv) - per_scenario = agg[(agg["scenario_type"] != "final_score") & (agg["scenario"] != agg["scenario_type"])] - for _, row in per_scenario.iterrows(): - rows[str(row["scenario"])] = row - return rows + """token -> aggregator row (score + per-metric values); {} for a run that was never scored.""" + try: + csv = nuplan_run_layout.aggregated_metrics_csv(group_dir) + except FileNotFoundError as error: + print(f"[render_obs_html] no aggregator table, pages named by token only: {error}") + return {} + agg = pd.read_csv(csv) + per_scenario = agg[(agg["scenario_type"] != "final_score") & (agg["scenario"] != agg["scenario_type"])] + return {str(row["scenario"]): row for _, row in per_scenario.iterrows()} def first_out_of_bounds(values, lower=None, upper=None): @@ -102,13 +108,10 @@ def step_link(step): 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 - ] + metrics_dir = nuplan_run_layout.metrics_dir(group_dir) out = {} - for metric, parquet_path, lower, upper in series_files: + for metric, parquet_name, lower, upper in VIOLATION_SERIES: + parquet_path = metrics_dir / f"{parquet_name}.parquet" if not parquet_path.exists(): continue for _, r in pd.read_parquet(parquet_path).iterrows(): @@ -128,24 +131,24 @@ def replay_collision_step(replay_path): def comfort_failures(group_dir): """token -> ['ego_lon_jerk -2.7..5.0', ...]: nuPlan's own out-of-bound comfort sub-metrics.""" + metrics_dir = nuplan_run_layout.metrics_dir(group_dir) out = {} - for sim_dir in glob.glob(f"{group_dir}/simulation/*/20*"): - for name in COMFORT_METRICS: - f = Path(sim_dir) / "metrics" / f"{name}.parquet" - if not f.exists(): - continue - df = pd.read_parquet(f) - 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} " - f"(bound {bound_label}){step_html}" - ) + for name in COMFORT_METRICS: + f = metrics_dir / f"{name}.parquet" + if not f.exists(): + continue + df = pd.read_parquet(f) + 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} " + f"(bound {bound_label}){step_html}" + ) return out @@ -201,6 +204,20 @@ def page_name(token, row): return f"s{row['score']:.2f}_{flags}{row['scenario_type']}_{token}.html" +def page_token(page_path): + """-> scenario token of a rendered page: the last underscore-separated part of its name.""" + return Path(page_path).stem.rsplit("_", 1)[-1] + + +def replay_sources(obs_dir): + """token -> its .replay.zlib, or its rendered page once the replay was pruned (the page embeds the payload).""" + sources = {path.name[: -len(REPLAY_SUFFIX)]: path for path in obs_dir.glob(f"*{REPLAY_SUFFIX}")} + for page in obs_dir.glob("*.html"): + if page.name != GALLERY_INDEX and page_token(page) not in sources: + sources[page_token(page)] = page + return sources + + def gallery_metrics(row): """Failure flags in the names the self-play gallery filters on (rate > 0 = flagged).""" if row is None: @@ -223,43 +240,45 @@ def main(): ap.add_argument("--workers", type=int, default=min(16, os.cpu_count() or 1)) args = ap.parse_args() obs_dir = Path(args.group_dir) / "obs_html" - replays = {p.name[: -len(".replay.zlib")]: p for p in obs_dir.glob("*.replay.zlib")} - if not replays: - raise SystemExit(f"no obs_html/*.replay.zlib under {args.group_dir}") + sources = replay_sources(obs_dir) + if not sources: + raise SystemExit(f"no obs_html/*{REPLAY_SUFFIX} or rendered pages 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))) + (token, sources[token], obs_dir / page_name(token, rows.get(token))) for token in sorted(wanted) - if token in replays + if token in sources ] def render(job): - token, replay, html_path = job + token, source, 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 + collision_step = replay_collision_step(source) 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)) + render_page(source, html_path, score_panel(token, row, comfort.get(token, []), token_steps)) + if source.suffix == ".html" and source != html_path: + source.unlink() # re-rendered under a new name: the old page would stay in the gallery with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool: list(pool.map(render, jobs)) file_metrics = {html.name: gallery_metrics(rows.get(token)) for token, _, html in jobs} pruned = 0 if args.prune: - for token, path in replays.items(): - if token not in wanted: + for path in sources.values(): + if path.suffix == ".zlib": path.unlink() pruned += 1 if jobs: viz.build_gallery_index(str(obs_dir), file_metrics=file_metrics) missing = len(wanted) - len(jobs) print( - f"[render_obs_html] {len(replays)} replays, {len(wanted)} scenarios selected, {len(jobs)} pages rendered" + f"[render_obs_html] {len(sources)} replays, {len(wanted)} scenarios selected, {len(jobs)} pages rendered" f"{f', {missing} selected without a replay' if missing else ''}{f', {pruned} replays pruned' if pruned else ''}" f" -> {obs_dir / 'index.html'}" ) diff --git a/scripts/eval/tidy_alpasim_run.py b/scripts/eval/tidy_alpasim_run.py index 77884751c0..b638cb7f7b 100644 --- a/scripts/eval/tidy_alpasim_run.py +++ b/scripts/eval/tidy_alpasim_run.py @@ -1,57 +1,81 @@ """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 + results.txt official score, km between incidents and the per-metric means (written by 12_alpasim.sh) + run_info.txt checkpoint, scenes, topology, git states + obs_html/ the replay pages render_alpasim_obs_html.py renders for the selected scenes (+ index.html) + details/ results-summary.json (the scored per-scene rows), metrics_unprocessed.parquet (per-step metrics of + every rollout), eval-config.yaml, wizard-config.yaml, config.json, evaluation.yaml, evaluate.log, + evaluation.log, services.log (every AlpaSim service container log, concatenated) + videos/ the in-run mp4s, only with RENDER_VIDEO=1 + alpasim/wizard_0/rollouts///rollout.asl only with KEEP_ROLLOUT_LOGS all|imperfect|failed, where + AlpaSim's and alpasim-training's offline tools look 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. -""" +Everything else is deleted: the evaluated checkpoint copy (final_model.pt + config.yaml, once the source checkpoint +named in run_info.txt is confirmed to exist), nre_rc.local, per-rollout metrics.parquet and _complete markers +(concatenated in metrics_unprocessed.parquet), controller csvs, runtime worker logs (copies of the runtime lines of +evaluate.log), generated configs, docker-compose.yaml, prometheus configs, aggregate plots. evaluate.log and +services.log lose the containers' INFO lines unless --keep_container_info. A run without results-summary.json (never +aggregated: the wizard skips the aggregation when it takes a leaked SLURM_ARRAY_TASK_COUNT for an array job) keeps its +wizard tree minus the controller csvs and worker logs; `/.venv/bin/python -m eval.aggregation.main +--array_job_dir --config_path /eval-config.yaml --failed_rollouts /failed_rollouts.json` +scores it, then this script finishes the job. Runs tidied into the 2026-10-12 layout (obs_replays/, logs/, alpasim/ +files) are moved into this one.""" import argparse import json import re +import shutil 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") +DETAILS_DIR = "details" +OBS_HTML_DIR = "obs_html" +RUN_INFO = "run_info.txt" +EVALUATE_LOG = "evaluate.log" +SERVICES_LOG = "services.log" ALPAGYM_LOGS = ("config.json", "evaluation.yaml", "logs/evaluation.log") +LEGACY_LOGS = ("evaluate.log", "config.json", "evaluation.yaml", "evaluation.log") # logs/ of the 2026-10-12 layout WIZARD_KEPT = ( "aggregate/results-summary.json", "aggregate/metrics_unprocessed.parquet", "eval-config.yaml", "wizard-config.yaml", ) +RUN_SCRATCH = ("nre_rc.local",) +WORKER_LOG_GLOB = "runtime_worker_*.log" +CONTROLLER_CSV_GLOB = "controller/alpasim_controller_*.csv" 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 +CHECKPOINT_LINE = re.compile(r"^checkpoint (\S+)") -def move(src, dst_dir): +def move(src, dst): + """Rename src to dst; a src an interrupted earlier tidy already moved is skipped.""" 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) + dst.parent.mkdir(parents=True, exist_ok=True) src.rename(dst) +def remove(path): + """-> bytes freed by deleting the file, 0 when it does not exist.""" + if not path.is_file() and not path.is_symlink(): + return 0 + size = path.lstat().st_size + path.unlink() + return size + + 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 @@ -72,6 +96,76 @@ def remove_empty_dirs(root): root.rmdir() +def remove_all(paths): + """-> (files deleted, bytes freed).""" + freed_bytes = 0 + paths = list(paths) + for path in paths: + freed_bytes += remove(path) + return len(paths), freed_bytes + + +def append_log(src, target, header): + """Append src under a header line to target and delete it. -> bytes freed.""" + target.parent.mkdir(parents=True, exist_ok=True) + with open(target, "ab") as target_file, open(src, "rb") as src_file: + target_file.write(f"==== {header} ====\n".encode()) + shutil.copyfileobj(src_file, target_file) + target_file.write(b"\n") + return remove(src) + + +def tidy_service_logs(txt_logs_dir, details_dir, keep_container_info): + """txt-logs/: the runtime worker logs deleted (copies of evaluate.log's runtime lines), the service container logs + into details/services.log. -> (files deleted, bytes freed).""" + if not txt_logs_dir.is_dir(): + return 0, 0 + deleted_count, freed_bytes = remove_all(sorted(txt_logs_dir.glob(WORKER_LOG_GLOB))) + for service_log in sorted(path for path in txt_logs_dir.iterdir() if path.is_file()): + if not keep_container_info: + freed_bytes += drop_info_lines(service_log, SERVICE_INFO_LINE) + freed_bytes += append_log(service_log, details_dir / SERVICES_LOG, service_log.name) + remove_empty_dirs(txt_logs_dir) + return deleted_count, freed_bytes + + +def remove_checkpoint_copy(run_dir): + """The evaluated copy of the checkpoint and its config.yaml, once run_info.txt's source checkpoint still exists + with the same size (a downloaded run keeps them). -> bytes freed.""" + run_info = run_dir / RUN_INFO + if not run_info.is_file(): + return 0 + match = CHECKPOINT_LINE.match(run_info.read_text()) + if match is None: + return 0 + source, copy = Path(match.group(1)), run_dir / Path(match.group(1)).name + if not (source.is_file() and copy.is_file() and source.stat().st_size == copy.stat().st_size): + return 0 + return remove(copy) + remove(run_dir / "config.yaml") + + +def gather_run(run_dir, details_dir): + """A fresh run's alpagym/ tree and a 2026-10-12 layout into the current one; evaluate.log into details/.""" + alpagym_dir, legacy_logs, legacy_alpasim = run_dir / "alpagym", run_dir / "logs", run_dir / "alpasim" + move(alpagym_dir / "obs_replays", run_dir / OBS_HTML_DIR) + for name in ALPAGYM_LOGS: + move(alpagym_dir / name, details_dir / Path(name).name) + move(alpagym_dir / "alpasim" / "wizard_0", legacy_alpasim / "wizard_0") + remove_empty_dirs(alpagym_dir) + if alpagym_dir.exists(): + leftovers = sorted(str(path.relative_to(run_dir)) for path in alpagym_dir.rglob("*") if path.is_file()) + print(f"[tidy_alpasim_run] left in place, unknown: {leftovers}") + move(run_dir / "obs_replays", run_dir / OBS_HTML_DIR) + move(legacy_logs / RUN_INFO, run_dir / RUN_INFO) + for name in LEGACY_LOGS: + move(legacy_logs / name, details_dir / name) + move(legacy_logs / "txt-logs", legacy_alpasim / "wizard_0" / "txt-logs") + remove_empty_dirs(legacy_logs) + for name in WIZARD_KEPT: + move(legacy_alpasim / Path(name).name, details_dir / Path(name).name) + move(run_dir / EVALUATE_LOG, details_dir / EVALUATE_LOG) + + 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).""" @@ -82,70 +176,57 @@ def sweep_wizard(wizard_dir, rows, keep_rule): 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_bytes += remove(path) 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 = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) 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" - ) + ap.add_argument("--keep_container_info", action="store_true", help="keep the containers' INFO lines in the logs") 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" + run_dir = Path(args.run_dir).resolve() + details_dir, wizard_dir = run_dir / DETAILS_DIR, run_dir / "alpasim" / "wizard_0" + gather_run(run_dir, details_dir) + deleted_count, freed_bytes = remove_all(run_dir / name for name in RUN_SCRATCH) + freed_bytes += remove_checkpoint_copy(run_dir) + evaluate_log = details_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" + freed_bytes += drop_info_lines(evaluate_log, CONTAINER_INFO_LINE) + controller_count, controller_bytes = remove_all(sorted(wizard_dir.glob(CONTROLLER_CSV_GLOB))) + remove_empty_dirs(wizard_dir / "controller") + service_count, service_bytes = tidy_service_logs(wizard_dir / "txt-logs", details_dir, args.keep_container_info) + deleted_count += controller_count + service_count + freed_bytes += controller_bytes + service_bytes 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'}") + if wizard_dir.is_dir() and any(wizard_dir.rglob("metrics.parquet")): + print( + f"[tidy_alpasim_run] no {summary_path}: the run was never aggregated, its wizard tree stays for " + f"`python -m eval.aggregation.main --array_job_dir {wizard_dir}` ({deleted_count} files, " + f"{freed_bytes / 1e9:.2f} GB removed)" + ) + else: + remove_empty_dirs(wizard_dir.parent) + migrated = ( + f" ({deleted_count} files, {freed_bytes / 1e9:.2f} GB removed)" if deleted_count or freed_bytes else "" + ) + print(f"[tidy_alpasim_run] already tidied{migrated}: {run_dir}") return rows = json.loads(summary_path.read_text())["rollouts"] for mp4 in sorted(wizard_dir.glob("rollouts/*/*/*.mp4")): - move(mp4, run_dir / "videos") + move(mp4, run_dir / "videos" / mp4.name) 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) + move(wizard_dir / name, details_dir / Path(name).name) 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 + remove_empty_dirs(wizard_dir.parent) + deleted_count += wizard_count + freed_bytes += wizard_bytes print( - f"[tidy_alpasim_run] removed {removed_count} files ({removed_bytes / 1e9:.2f} GB), kept {kept_logs} rollout logs " + f"[tidy_alpasim_run] removed {deleted_count} files ({freed_bytes / 1e9:.2f} GB), kept {kept_logs} rollout logs " f"({args.keep_rollout_logs}) -> {run_dir}" ) diff --git a/scripts/eval/tidy_carla_run.py b/scripts/eval/tidy_carla_run.py new file mode 100644 index 0000000000..9709051617 --- /dev/null +++ b/scripts/eval/tidy_carla_run.py @@ -0,0 +1,183 @@ +"""Tidy a CARLA longest6 run (the $OUT folder of scripts/kesai/11_carla_longest6_v2.sh or 14_carla_longest6_v1.sh, +after its gallery was rendered) down to what is read after the run: + + results.csv the result parser's summary and per-route table (moved up from aggregate/) + run_info.txt the run's settings and speed line (COMPLETE_RUN reads it) + obs_html/ the replay pages render_carla_obs_html.py rendered before this script ran + routes/ flat: route__s.result.json of every route run; .telemetry.csv, .world_log.npz and + .carla_view.mp4 of the kept route runs (route__rep_s.* for further repetitions) + details/ timing.csv, workers.log (every worker log), evaluators.log (every evaluator log) + +usage: python scripts/eval/tidy_carla_run.py [--keep failures] + +--keep selects the route runs whose telemetry, world log and chase cam survive: failures (a flagged leaderboard +infraction or a failed status, the gallery's default selection), all, none. Everything else is deleted: the +aggregate/ copies of the records, every .replay.zlib (a rendered page embeds its replay), the files of crashed +earlier attempts, the CARLA server logs, server_pids, compat_lib/ and the queue. report/ and pre_completion_*/ +stay. A folder tidied before is left alone; a completed run (COMPLETE_RUN) is tidied again for its new routes.""" + +import argparse +import shutil +import sys +from pathlib import Path + + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import carla_run_layout as layout +from tidy_alpasim_run import remove_empty_dirs + + +KEEP_MODES = ("failures", "all", "none") +DETAILS_DIR = "details" +AGGREGATE_DIR = "aggregate" +RESULTS_CSV = "results.csv" +TIMING_CSV = "timing.csv" +WORKERS_LOG = "workers.log" +EVALUATORS_LOG = "evaluators.log" +WORKER_LOG_GLOB = "worker_w*.log" +EVALUATOR_LOG_GLOB = "evaluator_attempt*.log" +SERVER_LOG_GLOB = "carla_server_attempt*.log" +RUN_SCRATCH = ("server_pids", "queue", "queue.lock") +RUN_SCRATCH_DIRS = ("compat_lib",) + + +def move(src, dst): + """Rename src to dst; a src an interrupted earlier tidy already moved is skipped.""" + if not src.exists(): + return + if dst.exists(): + raise SystemExit(f"{dst} already exists, not overwriting it with {src}") + dst.parent.mkdir(parents=True, exist_ok=True) + src.rename(dst) + + +def remove(path): + """-> bytes freed by deleting the file, 0 when it does not exist.""" + if not path.is_file() and not path.is_symlink(): + return 0 + size = path.lstat().st_size + path.unlink() + return size + + +def append_log(src, target, header): + """Append src under a header line to target (details/*.log) and delete it. -> bytes freed.""" + target.parent.mkdir(parents=True, exist_ok=True) + with open(target, "ab") as target_file, open(src, "rb") as src_file: + target_file.write(f"==== {header} ====\n".encode()) + shutil.copyfileobj(src_file, target_file) + target_file.write(b"\n") + return remove(src) + + +def worker_index(path): + return int(path.stem[len("worker_w") :]) + + +def kept_record(record, keep): + if keep == "all": + return True + if keep == "none": + return False + return layout.route_failed(record) + + +def tidy_route_run(route_dir, routes_dir, details_dir, keep): + """Flatten one route folder: the record and the kept records' files up into routes/, the evaluator logs into + details/evaluators.log, the rest deleted. -> (files deleted, bytes freed, files kept).""" + result_path = route_dir / layout.RAW_RESULT_NAME + records = layout.read_records(result_path) if result_path.is_file() else [] + stems = layout.record_stems(result_path, len(records)) + kept_stems = {stem: k for k, (stem, record) in enumerate(zip(stems, records)) if stem and kept_record(record, keep)} + deleted_count = freed_bytes = kept_count = 0 + for kind, suffix in layout.RECORD_FILE_KINDS: + for path in sorted((route_dir / kind).glob(f"*{suffix}")): + stem = path.name[: -len(suffix)] + if kind != layout.OBS_HTML_DIR and stem in kept_stems: + move(path, routes_dir / f"{layout.flat_stem(route_dir.name, kept_stems[stem])}.{kind}{suffix}") + kept_count += 1 + continue + freed_bytes += remove(path) + deleted_count += 1 + for path in sorted(route_dir.glob(EVALUATOR_LOG_GLOB)): + freed_bytes += append_log(path, details_dir / EVALUATORS_LOG, f"{route_dir.name} {path.name}") + for path in sorted(route_dir.glob(SERVER_LOG_GLOB)): + freed_bytes += remove(path) + deleted_count += 1 + if records: + move(result_path, routes_dir / f"{route_dir.name}{layout.RESULT_SUFFIX}") + kept_count += 1 + else: + freed_bytes += remove(result_path) # the leaderboard's empty stub of a route whose every attempt crashed + remove_empty_dirs(route_dir) + return deleted_count, freed_bytes, kept_count + + +def tidy_run_files(run_dir, details_dir): + """results.csv up from aggregate/, the aggregate copies, worker logs, timing and run scratch into details/ or + away. -> (files deleted, bytes freed).""" + deleted_count = freed_bytes = 0 + aggregate_dir = run_dir / AGGREGATE_DIR + if (aggregate_dir / RESULTS_CSV).is_file(): + freed_bytes += remove(run_dir / RESULTS_CSV) # a completion re-parsed the run + move(aggregate_dir / RESULTS_CSV, run_dir / RESULTS_CSV) + if aggregate_dir.is_dir(): + for path in sorted(aggregate_dir.rglob("*")): + if path.is_file(): + freed_bytes += remove(path) + deleted_count += 1 + remove_empty_dirs(aggregate_dir) + for path in sorted(run_dir.glob(WORKER_LOG_GLOB), key=worker_index): + freed_bytes += append_log(path, details_dir / WORKERS_LOG, path.name) + timing = run_dir / TIMING_CSV + if timing.is_file() and (details_dir / TIMING_CSV).is_file(): + with open(timing) as timing_file, open(details_dir / TIMING_CSV, "a") as target_file: + target_file.writelines(timing_file.readlines()[1:]) # a completion's rows join the run's table + freed_bytes += remove(timing) + else: + move(timing, details_dir / TIMING_CSV) + for name in RUN_SCRATCH: + freed_bytes += remove(run_dir / name) + for name in RUN_SCRATCH_DIRS: + if (run_dir / name).is_dir(): + for path in sorted((run_dir / name).iterdir()): + freed_bytes += remove(path) + deleted_count += 1 + remove_empty_dirs(run_dir / name) + return deleted_count, freed_bytes + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("run_dir") + ap.add_argument("--keep", choices=KEEP_MODES, default="failures", help="route runs whose logs survive") + args = ap.parse_args() + run_dir = Path(args.run_dir).resolve() + routes_dir = run_dir / layout.ROUTES_DIR + details_dir = run_dir / DETAILS_DIR + if not routes_dir.is_dir(): + raise SystemExit(f"no {layout.ROUTES_DIR}/ under {run_dir}") + route_dirs = sorted(path for path in routes_dir.glob("route_*") if path.is_dir()) + if not route_dirs and not (run_dir / AGGREGATE_DIR).exists() and not list(run_dir.glob(WORKER_LOG_GLOB)): + print(f"[tidy_carla_run] already tidied: {run_dir}") + return + deleted_count = freed_bytes = kept_count = 0 + for route_dir in route_dirs: + route_deleted, route_freed, route_kept = tidy_route_run(route_dir, routes_dir, details_dir, args.keep) + deleted_count += route_deleted + freed_bytes += route_freed + kept_count += route_kept + run_deleted, run_freed = tidy_run_files(run_dir, details_dir) + deleted_count += run_deleted + freed_bytes += run_freed + leftovers = sorted(path.name for path in route_dirs if path.exists()) + if leftovers: + print(f"[tidy_carla_run] route folders left in place, unknown files: {leftovers}") + print( + f"[tidy_carla_run] removed {deleted_count} files ({freed_bytes / 1e9:.2f} GB freed), kept {kept_count} " + f"route files ({args.keep}) over {len(route_dirs)} route folders -> {run_dir}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/eval/tidy_nuplan_run.py b/scripts/eval/tidy_nuplan_run.py new file mode 100644 index 0000000000..5331fc2dfa --- /dev/null +++ b/scripts/eval/tidy_nuplan_run.py @@ -0,0 +1,233 @@ +"""Tidy a nuPlan closed-loop run of the PufferDrive planner (the run_simulation.py `group` directory that +scripts/kesai/9_nuPlan.sh writes under /eval) down to what is read after the run: + + overview.csv CLS, the eight sub-metrics and the loss_* importance of each (CaRL's csv_main_callback) + scenarios.csv per-scenario scores, worst first (compare_nuplan_evals.py) + obs_html/ the observation replay pages render_obs_html.py rendered before this script ran + simulation_log/.msgpack.xz nuPlan's simulation logs of the kept scenarios + details/ + aggregated_metrics.csv the aggregator's full table: scenario type, comfort sub-metrics, per-type and final rows + metrics/ nuPlan's per-metric parquets with the per-step series of every scenario + runner_report.parquet, config.yaml, overrides.yaml + videos/.mp4 carl_visualization_callback's ground-truth view of the kept scenarios, when it ran + +usage: python scripts/eval/tidy_nuplan_run.py [--keep infractions] [--max-score 0.9] [--strip-planner] + +--keep selects the scenarios whose simulation log and video survive, with render_obs_html.py's rules: infractions +(an at-fault collision, drivable-area or driving-direction violation, or no progress), failures (score below +--max-score), all, none. Everything else nuPlan wrote is deleted: the other logs and videos, the aggregator parquet, +hydra.yaml, log.txt (hydra's job log, repeated in the SLURM log), the nuBoard descriptor, and every +obs_html/*.replay.zlib whose page exists (the page embeds it). +Videos are re-encoded with ffmpeg (imageio-ffmpeg's binary or the system's; h264, ~1.3 MB each) or, without one, +with OpenCV (mp4v, ~4.4 MB). --strip-planner rewrites the kept logs of runs recorded before the planner stopped +pickling its policy into them (28 MB -> 0.2 MB each; needs the nuPlan devkit env). A run that was never scored +keeps nuPlan's tree, as does a group holding both challenges; an already tidied run is left alone, an interrupted +tidy resumes. +""" + +import argparse +import shutil +import subprocess +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import nuplan_run_layout # noqa: E402 +from render_obs_html import FLAG_METRICS, GALLERY_INDEX, REPLAY_SUFFIX, page_token, scenario_rows, selected_tokens # noqa: E402 +from tidy_alpasim_run import remove_empty_dirs # noqa: E402 + +KEEP_MODES = ("infractions", "failures", "all", "none") +SIMULATION_LOG_SUFFIX = nuplan_run_layout.SIMULATION_LOG_SUFFIX +RUN_FILES = ("runner_report.parquet",) +HYDRA_FILES = ("config.yaml", "overrides.yaml") +HYDRA_JOB_LOG = "log.txt" +VIDEO_CRF = 23 # h264 quality: an 11 MB MJPEG scenario video -> 1.3 MB with lines and labels still legible +MAX_VIDEO_FRAMES = 100000 + + +def kept_tokens(rows, keep, max_score): + if keep == "all": + return set(rows) + if keep == "none": + return set() + if keep == "failures": + return selected_tokens(rows, max_score, []) + return selected_tokens(rows, 0.0, [metric for _, metric in FLAG_METRICS]) + + +def move(src, dst): + """Rename src to dst; a src that an interrupted earlier tidy already moved is skipped.""" + if not src.exists(): + return + if dst.exists(): + raise SystemExit(f"{dst} already exists, not overwriting it with {src}") + dst.parent.mkdir(parents=True, exist_ok=True) + src.rename(dst) + + +def remove(path): + """-> bytes freed by deleting the file, 0 when it does not exist.""" + if not path.is_file(): + return 0 + size = path.stat().st_size + path.unlink() + return size + + +def ffmpeg_executable(): + """-> imageio-ffmpeg's bundled ffmpeg or the system's, None when there is neither.""" + try: + import imageio_ffmpeg + + return imageio_ffmpeg.get_ffmpeg_exe() + except (ImportError, RuntimeError): + return shutil.which("ffmpeg") + + +def transcode_video(avi_path, mp4_path, ffmpeg): + mp4_path.parent.mkdir(parents=True, exist_ok=True) + if ffmpeg: + command = [ffmpeg, "-v", "error", "-y", "-i", str(avi_path), "-c:v", "libx264", "-preset", "fast"] + command += ["-crf", str(VIDEO_CRF), "-pix_fmt", "yuv420p", str(mp4_path)] + subprocess.run(command, check=True) + return + import cv2 + + reader = cv2.VideoCapture(str(avi_path)) + size = (int(reader.get(cv2.CAP_PROP_FRAME_WIDTH)), int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT))) + writer = cv2.VideoWriter(str(mp4_path), cv2.VideoWriter_fourcc(*"mp4v"), reader.get(cv2.CAP_PROP_FPS), size) + if not writer.isOpened(): + raise SystemExit(f"{mp4_path}: OpenCV could not open the video writer") + frame_count = 0 + while frame_count < MAX_VIDEO_FRAMES: + ok, frame = reader.read() + if not ok: + break + writer.write(frame) + frame_count += 1 + writer.release() + reader.release() + + +def strip_planner_from_log(log_path, target): + """Re-save a log recorded before the planner dropped its policy from its pickled state. -> bytes freed.""" + import nuplan_log_remap # remaps the scenario paths of logs written on another machine + + log = nuplan_log_remap.load_log(log_path) + target.parent.mkdir(parents=True, exist_ok=True) + log.file_path = target + log.save_to_file() + freed = log_path.stat().st_size - target.stat().st_size + log_path.unlink() + return freed + + +def tidy_summaries(sim_dir, group_dir, details_dir): + """The aggregator table into details/, overview.csv and scenarios.csv to the top of the run. -> bytes freed.""" + aggregator_dir = sim_dir / "aggregator_metric" + for csv_path in aggregator_dir.glob("*_weighted_average_metrics_*.csv"): + move(csv_path, details_dir / nuplan_run_layout.AGGREGATED_METRICS_CSV) + move(aggregator_dir / "overview.csv", group_dir / "overview.csv") + move(aggregator_dir / "scenarios.csv", group_dir / "scenarios.csv") + return sum(remove(path) for path in aggregator_dir.glob("*.parquet")) + + +def tidy_run_files(sim_dir, details_dir): + """metrics/, the runner report and the resolved hydra config into details/. -> bytes freed.""" + move(sim_dir / "metrics", details_dir / "metrics") + for name in RUN_FILES: + move(sim_dir / name, details_dir / name) + for name in HYDRA_FILES: + move(sim_dir / "code" / "hydra" / name, details_dir / name) + freed_bytes = remove(sim_dir / "code" / "hydra" / "hydra.yaml") + remove(sim_dir / HYDRA_JOB_LOG) + return freed_bytes + sum(remove(path) for path in sim_dir.glob("*.nuboard")) + + +def tidy_simulation_logs(sim_dir, log_dir, kept, strip_planner): + """-> (logs deleted, bytes freed, logs kept as /.msgpack.xz).""" + deleted_count = freed_bytes = kept_count = 0 + for log_path in sorted((sim_dir / nuplan_run_layout.SIMULATION_LOG_DIR).rglob(f"*{SIMULATION_LOG_SUFFIX}")): + token = log_path.name[: -len(SIMULATION_LOG_SUFFIX)] + target = log_dir / log_path.name + if token not in kept: + freed_bytes += remove(log_path) + deleted_count += 1 + elif strip_planner: + freed_bytes += strip_planner_from_log(log_path, target) + kept_count += 1 + else: + move(log_path, target) + kept_count += 1 + return deleted_count, freed_bytes, kept_count + + +def tidy_videos(sim_dir, details_dir, kept, ffmpeg): + """-> (videos deleted, bytes freed, videos kept as details/videos/.mp4); nothing without the callback.""" + deleted_count = freed_bytes = kept_count = 0 + for avi_path in sorted((sim_dir / "visualization").glob("*.avi")): + if avi_path.stem not in kept: + freed_bytes += remove(avi_path) + deleted_count += 1 + continue + mp4_path = details_dir / "videos" / f"{avi_path.stem}.mp4" + transcode_video(avi_path, mp4_path, ffmpeg) + freed_bytes += remove(avi_path) - mp4_path.stat().st_size + kept_count += 1 + return deleted_count, freed_bytes, kept_count + + +def prune_rendered_replays(obs_dir): + """-> (replays deleted, bytes freed): obs_html/*.replay.zlib whose rendered page exists (the page embeds it).""" + pages = {page_token(page) for page in obs_dir.glob("*.html") if page.name != GALLERY_INDEX} + deleted_count = freed_bytes = 0 + for replay_path in sorted(obs_dir.glob(f"*{REPLAY_SUFFIX}")): + if replay_path.name[: -len(REPLAY_SUFFIX)] not in pages: + continue + freed_bytes += remove(replay_path) + deleted_count += 1 + return deleted_count, freed_bytes + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("group_dir") + ap.add_argument("--keep", choices=KEEP_MODES, default="infractions", help="scenarios whose log and video survive") + ap.add_argument("--max-score", type=float, default=0.9, help="--keep failures: scenarios scoring below this") + ap.add_argument("--strip-planner", action="store_true", help="re-save the kept logs without the pickled policy") + args = ap.parse_args() + group_dir = Path(args.group_dir).resolve() + details_dir = group_dir / nuplan_run_layout.DETAILS_DIR + log_dir = group_dir / nuplan_run_layout.SIMULATION_LOG_DIR + simulation_root = group_dir / "simulation" + if details_dir.is_dir() and not simulation_root.exists(): + print(f"[tidy_nuplan_run] already tidied: {group_dir}") + return + sim_dir = nuplan_run_layout.raw_simulation_dir(group_dir) + if sim_dir is None: + raise SystemExit(f"no simulation// under {group_dir}") + try: + nuplan_run_layout.aggregated_metrics_csv(group_dir) + except FileNotFoundError: + print(f"[tidy_nuplan_run] {group_dir} was never scored (no aggregator csv): nuPlan's tree kept as is") + return + kept = kept_tokens(scenario_rows(group_dir), args.keep, args.max_score) + ffmpeg = ffmpeg_executable() + freed_bytes = tidy_summaries(sim_dir, group_dir, details_dir) + freed_bytes += tidy_run_files(sim_dir, details_dir) + log_count, log_bytes, kept_logs = tidy_simulation_logs(sim_dir, log_dir, kept, args.strip_planner) + video_count, video_bytes, kept_videos = tidy_videos(sim_dir, details_dir, kept, ffmpeg) + replay_count, replay_bytes = prune_rendered_replays(group_dir / "obs_html") + remove_empty_dirs(simulation_root) + leftovers = sorted(str(path.relative_to(group_dir)) for path in simulation_root.rglob("*") if path.is_file()) + if leftovers: + print(f"[tidy_nuplan_run] left in place, unknown: {leftovers}") + freed_bytes += log_bytes + video_bytes + replay_bytes + print( + f"[tidy_nuplan_run] removed {log_count} logs, {video_count} videos and {replay_count} replays " + f"({freed_bytes / 1e9:.2f} GB freed), kept {kept_logs} logs and {kept_videos} videos " + f"({args.keep}, {'h264' if ffmpeg else 'mp4v'}) -> {group_dir}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/eval/tidy_selfplay_eval.py b/scripts/eval/tidy_selfplay_eval.py new file mode 100644 index 0000000000..77a4d7289f --- /dev/null +++ b/scripts/eval/tidy_selfplay_eval.py @@ -0,0 +1,226 @@ +"""Tidy a CARLA self-play eval (scripts/parallel_eval.py, or one `puffer eval` run) down to what is read after it: + + evaluation_summary.json the metric averages over every scenario (eval_wandb_report.py reads them) + obs_html/ the rendered replay pages of every shard's failures + one index.html + details/ episode_metrics.csv (one row per scenario), failure_episodes.csv (the re-simulated + failure episodes behind the pages, with their shard), selected_failures.csv (the rows + the render filter picked, with their shard), resolved_benchmark.yaml, shards.log + +usage: python scripts/eval/tidy_selfplay_eval.py + + is /eval/carla_/ as parallel_eval.py leaves it, its shards next to it as +_shard/ (resolved_benchmark.yaml, episode_metrics.csv, evaluation_summary.json, failures/{..., +rendered_replays/}); a single-process `puffer eval` dir holds the same files itself. The shard folders are deleted: +their episode_metrics.csv rows are the aggregate's, their summaries are per-shard means, the per-shard gallery +indexes make way for one. resolved_benchmark.yaml is the first shard's (the others differ in their scenario window +only). A .replay.zlib kept by eval.keep_zlib_replays is deleted once its page exists. An already tidied dir is +left alone.""" + +import argparse +import shutil +import sys +from pathlib import Path + +import numpy as np +import pandas as pd + + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from tidy_alpasim_run import remove_empty_dirs + +from pufferlib import viz + + +OBS_HTML_DIR = "obs_html" +DETAILS_DIR = "details" +SUMMARY_JSON = "evaluation_summary.json" +EPISODE_METRICS_CSV = "episode_metrics.csv" +FAILURE_EPISODES_CSV = "failure_episodes.csv" +SELECTED_FAILURES_CSV = "selected_failures.csv" +RESOLVED_BENCHMARK_YAML = "resolved_benchmark.yaml" +SHARDS_LOG = "shards.log" +SHARD_LOG_GLOB = "shard*.log" +FAILURES_DIR = "failures" +RENDERED_REPLAYS_DIR = "rendered_replays" +ZLIB_REPLAYS_DIR = "replays_zlib" +REPLAY_SUFFIX = ".replay.zlib" +GALLERY_INDEX = "index.html" +SHARD_COLUMN = "shard" + + +def move(src, dst): + """Rename src to dst; a src an interrupted earlier tidy already moved is skipped.""" + if not src.exists(): + return + if dst.exists(): + raise SystemExit(f"{dst} already exists, not overwriting it with {src}") + dst.parent.mkdir(parents=True, exist_ok=True) + src.rename(dst) + + +def remove(path): + """-> bytes freed by deleting the file, 0 when it does not exist.""" + if not path.is_file(): + return 0 + size = path.stat().st_size + path.unlink() + return size + + +def append_log(src, target, header): + """Append src under a header line to target and delete it. -> bytes freed.""" + target.parent.mkdir(parents=True, exist_ok=True) + with open(target, "ab") as target_file, open(src, "rb") as src_file: + target_file.write(f"==== {header} ====\n".encode()) + shutil.copyfileobj(src_file, target_file) + target_file.write(b"\n") + return remove(src) + + +def shard_dirs(run_dir): + """-> the _shard folders next to a parallel_eval.py aggregate dir, then the dir itself when it holds + a run's own failures/ or resolved_benchmark.yaml (a single-process `puffer eval`).""" + shards = sorted(path for path in run_dir.parent.glob(f"{run_dir.name}_shard*") if path.is_dir()) + if (run_dir / FAILURES_DIR).is_dir() or (run_dir / RESOLVED_BENCHMARK_YAML).is_file(): + shards.append(run_dir) + return shards + + +def page_name(replay_path): + """-> the page rendered from a replay: .html for .replay.zlib.""" + name = Path(str(replay_path)).name + return name[: -len(REPLAY_SUFFIX)] + ".html" if name.endswith(REPLAY_SUFFIX) else name + + +def read_csv_with_shard(path, shard): + frame = pd.read_csv(path) + frame.insert(0, SHARD_COLUMN, shard) + return frame + + +def tidy_shard(shard_dir, run_dir, details_dir, obs_dir, keep_aggregate_metrics): + """Pages into obs_html/, the failure tables and the benchmark config into details/, the rest deleted. + -> (failure episode frame or None, selected failures frame or None, files deleted, bytes freed).""" + shard = shard_dir.name[len(run_dir.name) + 1 :] if shard_dir != run_dir else "single" + failures_dir = shard_dir / FAILURES_DIR + deleted_count = freed_bytes = 0 + for page in sorted((failures_dir / RENDERED_REPLAYS_DIR).glob("*.html")): + if page.name == GALLERY_INDEX: + freed_bytes += remove(page) + deleted_count += 1 + continue + move(page, obs_dir / page.name) + for replay in sorted((failures_dir / ZLIB_REPLAYS_DIR).glob(f"*{REPLAY_SUFFIX}")): + if (obs_dir / page_name(replay)).is_file(): + freed_bytes += remove(replay) + deleted_count += 1 + else: + move(replay, obs_dir / replay.name) + failure_frame = selected_frame = None + if (failures_dir / EPISODE_METRICS_CSV).is_file(): + failure_frame = read_csv_with_shard(failures_dir / EPISODE_METRICS_CSV, shard) + freed_bytes += remove(failures_dir / EPISODE_METRICS_CSV) + deleted_count += 1 + if (failures_dir / SELECTED_FAILURES_CSV).is_file(): + selected_frame = read_csv_with_shard(failures_dir / SELECTED_FAILURES_CSV, shard) + freed_bytes += remove(failures_dir / SELECTED_FAILURES_CSV) + deleted_count += 1 + freed_bytes += remove(failures_dir / SUMMARY_JSON) + if (details_dir / RESOLVED_BENCHMARK_YAML).is_file(): + freed_bytes += remove(shard_dir / RESOLVED_BENCHMARK_YAML) + else: + move(shard_dir / RESOLVED_BENCHMARK_YAML, details_dir / RESOLVED_BENCHMARK_YAML) + if shard_dir != run_dir or not keep_aggregate_metrics: + freed_bytes += remove(shard_dir / EPISODE_METRICS_CSV) + deleted_count += 1 + if shard_dir != run_dir: + freed_bytes += remove(shard_dir / SUMMARY_JSON) + deleted_count += 1 + remove_empty_dirs(failures_dir) + if shard_dir != run_dir: + remove_empty_dirs(shard_dir) + return failure_frame, selected_frame, deleted_count, freed_bytes + + +def gallery_metrics(failure_frame): + """page name -> its episode's finite numeric metrics, the gallery's filter inputs (as the shards' renderer).""" + if failure_frame is None or "replay_path" not in failure_frame.columns: + return {} + metrics = {} + for _, row in failure_frame.iterrows(): + values = {} + for column, value in row.items(): + if column in (SHARD_COLUMN, "replay_path") or not isinstance(value, (int, float, np.integer, np.floating)): + continue + if np.isfinite(value): + values[column] = float(value) + metrics[page_name(row["replay_path"])] = values + return metrics + + +def append_frames(frames, target): + """Concatenate the shards' tables under target (appending to a table an earlier tidy wrote).""" + frames = [frame for frame in frames if frame is not None] + if not frames: + return + merged = pd.concat(frames, ignore_index=True) + if target.is_file(): + merged = pd.concat([pd.read_csv(target), merged], ignore_index=True) + target.parent.mkdir(parents=True, exist_ok=True) + merged.to_csv(target, index=False) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("run_dir", help="parallel_eval.py aggregate dir (its shards next to it), or a `puffer eval` dir") + args = ap.parse_args() + run_dir = Path(args.run_dir).resolve() + if not (run_dir / SUMMARY_JSON).is_file(): + raise SystemExit(f"no {SUMMARY_JSON} under {run_dir}: not an evaluated run") + details_dir, obs_dir = run_dir / DETAILS_DIR, run_dir / OBS_HTML_DIR + shards = shard_dirs(run_dir) + shard_logs = sorted(run_dir.glob(SHARD_LOG_GLOB)) + if not shards and not shard_logs and not (run_dir / EPISODE_METRICS_CSV).is_file(): + print(f"[tidy_selfplay_eval] already tidied: {run_dir}") + return + move(run_dir / EPISODE_METRICS_CSV, details_dir / EPISODE_METRICS_CSV) + deleted_count = freed_bytes = 0 + for log in shard_logs: + freed_bytes += append_log(log, details_dir / SHARDS_LOG, log.name) + failure_frames, selected_frames = [], [] + for shard_dir in shards: + failure_frame, selected_frame, shard_deleted, shard_freed = tidy_shard( + shard_dir, run_dir, details_dir, obs_dir, keep_aggregate_metrics=True + ) + failure_frames.append(failure_frame) + selected_frames.append(selected_frame) + deleted_count += shard_deleted + freed_bytes += shard_freed + append_frames(failure_frames, details_dir / FAILURE_EPISODES_CSV) + append_frames(selected_frames, details_dir / SELECTED_FAILURES_CSV) + pages = sorted(path for path in obs_dir.glob("*.html") if path.name != GALLERY_INDEX) if obs_dir.is_dir() else [] + if pages: + file_metrics = ( + gallery_metrics(pd.read_csv(details_dir / FAILURE_EPISODES_CSV)) + if (details_dir / FAILURE_EPISODES_CSV).is_file() + else {} + ) + viz.build_gallery_index(str(obs_dir), file_metrics=file_metrics) + leftovers = sorted( + str(path.relative_to(run_dir.parent)) + for shard in shards + for path in shard.rglob("*") + if path.is_file() and shard != run_dir + ) + if leftovers: + print(f"[tidy_selfplay_eval] left in place, unknown: {leftovers}") + print( + f"[tidy_selfplay_eval] {len(shards)} shard folders folded into {run_dir}: {len(pages)} replay pages, " + f"{deleted_count} files removed ({freed_bytes / 1e6:.1f} MB)" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/kesai/11_carla_longest6_v2.sh b/scripts/kesai/11_carla_longest6_v2.sh index b0d9865e2f..0d843ae090 100644 --- a/scripts/kesai/11_carla_longest6_v2.sh +++ b/scripts/kesai/11_carla_longest6_v2.sh @@ -17,9 +17,11 @@ # 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. +# aggregated at the end with CaRL's tools/result_parser.py (results.csv) plus the obs replay gallery of the failed +# routes ($OUT/obs_html/index.html, as in 9_nuPlan.sh). scripts/eval/tidy_carla_run.py then leaves $OUT/results.csv, +# run_info.txt, obs_html/, routes/ (flat: route__s.result.json of every route run, .telemetry.csv + +# .world_log.npz of the failed ones) and details/ (timing.csv, workers.log, evaluators.log). Speed: wall and game +# time of every route go to 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 @@ -27,23 +29,34 @@ # 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 +# PORT_BASE (12000: CARLA rpc port of worker 0; worker w uses PORT_BASE+50w, TM PORT_BASE+6000+50w; a range with a port +# already listening is skipped for 12000, then 22000, see run_info.txt port_base), OMP_NUM_THREADS (policy +# threads; 2 from 16 workers on, else 4), LOGGING (1, default: telemetry, world log and obs replay 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), +# needs LOGGING=1), OBS_HTML (1 = failures, default: the interactive obs replay pages of the route runs with a +# flagged infraction or a failed status, rendered into $OUT/obs_html at the end; all = every route run; 0 = none; +# needs LOGGING=1), REPORT (0, default; 1 = the HTML report of scripts/eval/analyze_carla_cosim.py in $OUT/report, +# needs LOGGING=1), TIDY (1, default: scripts/eval/tidy_carla_run.py flattens $OUT at the end; 0 = keep every file), # 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_ZERO_PARTNER_STOPPED_TIME (default 1: vehicles' stopped-time obs held at 0, walkers keep theirs; 0 = all real), # 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). +# read strictly larger than the smallest training vehicle; 0 = true CARLA boxes), COSIM_WALKER_SPEED_FLOOR_MPS (default 0; +# the diagnostic walker speed floor of shadow_ego.py, tags the folder _walkerfloor), WANDB_REPORT (1, default: the +# DS of results.csv becomes the training run's final_eval/longest6_v2_DS in wandb, replacing an earlier report; 0 for +# ablations, so the run keeps the number of its standard eval; never reported: a run with missing route records, the +# scenario-free variant and a walker speed floor run). +# COSIM_DYNAMICS_SOURCE (carla_pedals, default: CARLA's physics moves the ego, driven only by the VehicleControl the agent +# returns; tags the folder _carla_pedals; pufferdrive: the ego is teleported along PufferDrive's own dynamics). +# Ablations of carla_pedals: COSIM_PEDAL_REVERSE (0, default: the car never backs up; 1 = it reverses, tags the folder _reverse) and +# COSIM_PEDAL_YELLOW_MARGIN (1, default; 0 = the policy sees CARLA's remaining yellow, tags the folder _noyellowmargin). 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} +RUN_DIR=${RUN_DIR:-$PD/experiments/k_scaled_0071_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} @@ -58,15 +71,18 @@ for seed in $SEEDS; do [[ "$seed" =~ ^[0-9]+$ ]] || { echo "SEEDS=\"$SEEDS\" mus 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} +PORT_BASE=${PORT_BASE:-12000} 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} +OBS_HTML=${OBS_HTML:-1} # 1|failures: replay pages of the failed route runs; all: every route run; 0: none +case "$OBS_HTML" in 1) OBS_HTML=failures ;; 0|failures|all) ;; *) echo "OBS_HTML=\"$OBS_HTML\" must be 0, 1, failures or all"; exit 1 ;; esac +REPORT=${REPORT:-0} +TIDY=${TIDY:-1} +WANDB_REPORT=${WANDB_REPORT:-1} 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; } @@ -81,6 +97,26 @@ 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 +# Every port of the worker layout (rpc, streaming = rpc+1, TM) must be free of listeners; the first clear base wins. +# Bases stay below the ephemeral range (32768+); the old 2000 put node_exporter (9100) and dcgm-exporter (9400) inside it. +pick_port_base() { + local listening base w port busy + listening=" $(ss -ltun 2>/dev/null | awk 'NR > 1 {sub(/.*:/, "", $5); print $5}' | sort -u | tr '\n' ' ') " + for base in "$@"; do + busy=0 + for ((w = 0; w < WORKERS; w++)); do + for port in $((base + 50 * w)) $((base + 50 * w + 1)) $((base + 6000 + 50 * w)); do + [[ "$listening" == *" $port "* ]] && { busy=$port; break 2; } + done + done + [ "$busy" = 0 ] && { echo "$base"; return 0; } + echo "port base $base: port $busy is in use, trying the next range" >&2 + done + return 1 +} +PORT_BASE=$(pick_port_base "$PORT_BASE" 12000 22000) || { echo "no free port range for $WORKERS workers (ss -ltun)"; exit 1; } +echo "port base $PORT_BASE: rpc $PORT_BASE+50w, TM $((PORT_BASE + 6000))+50w" + # 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}" @@ -89,22 +125,36 @@ 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_DYNAMICS_SOURCE=${COSIM_DYNAMICS_SOURCE:-carla_pedals} # carla_pedals: CARLA's physics moves the ego, see shadow_ego.py +[[ "$COSIM_DYNAMICS_SOURCE" =~ ^(pufferdrive|carla_pedals)$ ]] || { echo "COSIM_DYNAMICS_SOURCE=\"$COSIM_DYNAMICS_SOURCE\" must be pufferdrive or carla_pedals"; exit 1; } +export COSIM_PEDAL_REVERSE=${COSIM_PEDAL_REVERSE:-0} COSIM_PEDAL_YELLOW_MARGIN=${COSIM_PEDAL_YELLOW_MARGIN:-1} # carla_pedals ablations, see shadow_ego.py +[[ "$COSIM_PEDAL_REVERSE" =~ ^[01]$ && "$COSIM_PEDAL_YELLOW_MARGIN" =~ ^[01]$ ]] || { echo "COSIM_PEDAL_REVERSE=\"$COSIM_PEDAL_REVERSE\" and COSIM_PEDAL_YELLOW_MARGIN=\"$COSIM_PEDAL_YELLOW_MARGIN\" must be 0 or 1"; exit 1; } 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_WALKER_SPEED_FLOOR_MPS=${COSIM_WALKER_SPEED_FLOOR_MPS:-0} # diagnostic only, see shadow_ego.py +[[ "$COSIM_WALKER_SPEED_FLOOR_MPS" =~ ^[0-9]+([.][0-9]+)?$ ]] || { echo "COSIM_WALKER_SPEED_FLOOR_MPS=\"$COSIM_WALKER_SPEED_FLOOR_MPS\" must be a non-negative number"; exit 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 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") +TAG=$TAG$([ "$COSIM_WALKER_SPEED_FLOOR_MPS" = "0" ] || echo "_walkerfloor$COSIM_WALKER_SPEED_FLOOR_MPS") +TAG=$TAG$([ "$COSIM_DYNAMICS_SOURCE" = "pufferdrive" ] || echo "_$COSIM_DYNAMICS_SOURCE") +if [ "$COSIM_DYNAMICS_SOURCE" = "carla_pedals" ]; then + TAG=$TAG$([ "$COSIM_PEDAL_REVERSE" = "0" ] || echo "_reverse")$([ "$COSIM_PEDAL_YELLOW_MARGIN" = "1" ] || echo "_noyellowmargin") +fi # 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 "walker_speed_floor_mps $COSIM_WALKER_SPEED_FLOOR_MPS" + echo "dynamics_source $COSIM_DYNAMICS_SOURCE" + echo "pedal_reverse $COSIM_PEDAL_REVERSE"; echo "pedal_yellow_margin $COSIM_PEDAL_YELLOW_MARGIN" echo "git $(git -C "$PD" rev-parse --short HEAD 2>/dev/null) $(git -C "$PD" diff --quiet 2>/dev/null || echo dirty)" + echo "port_base $PORT_BASE" } > "$OUT/run_info.txt" echo "Results -> $OUT" @@ -120,10 +170,12 @@ esac 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") +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[@]})) [ "$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) @@ -209,7 +261,7 @@ except Exception: 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") + [ "$OBS_HTML" != "0" ] && 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" \ @@ -305,16 +357,29 @@ if rows: 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 +# the repo .venv holds wandb; the script reads results.csv and the route records and never imports pufferlib +if [ "$WANDB_REPORT" = "1" ] && [ "$SCENARIOS" = "1" ] && [ "$COSIM_WALKER_SPEED_FLOOR_MPS" = "0" ]; then + env -u PYTHONPATH "$PD/.venv/bin/python" "$PD/pufferlib/ocean/evaluation_utils/eval_wandb_report.py" \ + longest6_v2 "$OUT/aggregate/results.csv" --run-dir "$RUN_DIR" --expected-episode-count $((TOTAL_UNITS * REPETITIONS)) \ + || echo "wandb report failed (results above are unaffected)" +fi REPORT_ARGS=() -[ "$OBS_HTML" = "1" ] && REPORT_ARGS+=(--obs-html-dir "$OUT/obs_html") +[ "$OBS_HTML" != "0" ] && 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)" +if [ "$OBS_HTML" != "0" ] && [ "$LOGGING" = "1" ]; then + "$PY" "$PD/scripts/eval/render_carla_obs_html.py" "$OUT" --routes "$ROUTES" --select "$OBS_HTML" \ + || echo "obs replay gallery failed (results above are unaffected)" +fi +# the pages are rendered: flatten the route folders, keep the logs of the failed (OBS_HTML=all: every) route runs +if [ "$TIDY" = "1" ]; then + "$PY" "$PD/scripts/eval/tidy_carla_run.py" "$OUT" --keep "$([ "$OBS_HTML" = "all" ] && echo all || echo failures)" \ + || echo "tidy 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" +for results in "$OUT/results.csv" "$OUT/aggregate/results.csv"; do [ -f "$results" ] && { echo "results -> $results"; break; }; done echo "Done -> $OUT" diff --git a/scripts/kesai/12_alpasim.sh b/scripts/kesai/12_alpasim.sh index fadf133703..90f534b98f 100644 --- a/scripts/kesai/12_alpasim.sh +++ b/scripts/kesai/12_alpasim.sh @@ -36,10 +36,13 @@ # 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 +# 3_train / 7_eval submit this eval from inside an array job and sbatch hands its SLURM_ARRAY_* down: the wizard then +# splits the scenes as if this were array task N and skips the aggregation (k_scaled_0074_2000: no official score) +unset SLURM_ARRAY_TASK_COUNT SLURM_ARRAY_TASK_ID SLURM_ARRAY_TASK_MIN SLURM_ARRAY_TASK_MAX SLURM_ARRAY_JOB_ID 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} +RUN_DIR=${RUN_DIR:-$PD/experiments/k_scaled_0058_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 @@ -47,13 +50,14 @@ GOAL_RADIUS_M=2 # goal arrival radius Damiano's reference evals use; the te 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 +OBS_REPLAY=imperfect # replay pages in $OUT/obs_html: imperfect (score < 1) | failed (AlpaSim's failed scenes) | all | 0 (no 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) +WANDB_REPORT=1 # the official score becomes the training run's final_eval/alpasim_score in wandb, replacing an earlier report; 0 for ablations 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 @@ -80,6 +84,7 @@ export PATH="$HOME/.local/bin:$PATH" [ -n "$CKPT" ] || { echo "no checkpoint in $RUN_DIR (final_model.pt or models/model_*.pt)"; 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 +case "$OBS_REPLAY" in 1) OBS_REPLAY=imperfect ;; 0|imperfect|failed|all) ;; *) echo "OBS_REPLAY must be 0, 1, imperfect, failed or all, got $OBS_REPLAY"; 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 @@ -277,8 +282,8 @@ if [ "$CONTAINER_RUNTIME" != docker ]; then 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)" +[ "$OBS_REPLAY" != "0" ] && OBS_REPLAY_FLAG=true || OBS_REPLAY_FLAG=false +echo "${#SCENES[@]} scenes, $TOPOLOGY with $PARALLEL in flight on $CONTAINER_RUNTIME; full log: $OUT/evaluate.log (details/ 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 @@ -339,6 +344,12 @@ Path(sys.argv[3]).write_text(text) sys.stdout.write("\n" + text) EOF fi +# the repo .venv holds wandb; the script reads results.txt and never imports pufferlib. A subset (N_SCENES) or a run with failed scenes is not reported. +if [ "$WANDB_REPORT" = "1" ] && [ $EVAL_STATUS -eq 0 ] && [ "$N_SCENES" = "0" ]; then + env -u PYTHONPATH "$PD/.venv/bin/python" "$PD/pufferlib/ocean/evaluation_utils/eval_wandb_report.py" \ + alpasim "$OUT/results.txt" --run-dir "$RUN_DIR" --expected-episode-count ${#SCENES[@]} \ + || echo "wandb report failed (results above are unaffected)" +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 @@ -350,14 +361,14 @@ if [ $EVAL_STATUS -ne 0 ]; then 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 +# results.txt, run_info.txt, obs_html/, details/ (+ videos/); container INFO lines stay in the logs 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" +# score-named pages of the OBS_REPLAY scenes with the AlpaSim score panel + index.html gallery, as 9_nuPlan.sh and 11_carla_longest6_v2.sh +if [ -d "$OUT/obs_html" ] && [ -f "$OUT/details/results-summary.json" ]; then + "$ALPAGYM_PY" "$PD/scripts/eval/render_alpasim_obs_html.py" "$OUT" --select "$OBS_REPLAY" --prune || echo "obs replay gallery failed" fi -echo "results -> $OUT/results.txt; obs replays -> $OUT/obs_replays/index.html; logs -> $OUT/logs; per-scene scores + per-step metrics -> $OUT/alpasim" +echo "results -> $OUT/results.txt; obs replays -> $OUT/obs_html/index.html; logs, per-scene scores + per-step metrics -> $OUT/details" echo "Done -> $OUT" exit $EVAL_STATUS diff --git a/scripts/kesai/14_carla_longest6_v1.sh b/scripts/kesai/14_carla_longest6_v1.sh index ef36e977fc..84f245e96a 100755 --- a/scripts/kesai/14_carla_longest6_v1.sh +++ b/scripts/kesai/14_carla_longest6_v1.sh @@ -8,22 +8,25 @@ #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 --open-mode=append # a requeued job keeps the log of the instance it resumes #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 +# WORKERS parallel workers (default 32), 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. +# runs, $OUT/routes/route__s/) wait in one queue, longest first (route length from the xml), and a worker pulls +# the next one whenever it is free (a static split left one worker alone for the last hour whenever its routes ran their +# clock out). 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 +# (results.csv) plus the obs replay gallery of the failed routes ($OUT/obs_html/index.html, as in +# 11_carla_longest6_v2.sh). scripts/eval/tidy_carla_run.py then leaves $OUT/results.csv, run_info.txt, obs_html/, +# routes/ (flat: route__s.result.json of every route run, .telemetry.csv + .world_log.npz of the failed +# ones) and details/ (timing.csv, workers.log, evaluators.log). Speed: wall and game time of every route go to +# 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 @@ -31,25 +34,41 @@ # --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), +# ROUTE_SUBSET ("0 5 17": only these route ids), WORKERS, NUM_GPUS, STAGGER_S (5: seconds between worker launches), PORT_BASE +# (12000: CARLA rpc port of worker 0; worker w uses PORT_BASE+50w, TM PORT_BASE+6000+50w; a range with a port already +# listening is skipped for 12000, then 22000, see run_info.txt port_base), OMP_NUM_THREADS (policy server +# threads; 2 from 16 workers on, else 4), LOGGING (1, default: telemetry, world log and obs replay per route; 0: scores +# only), CARLA_VIEW (0, default: no chase-cam; 1: chase-cam mp4 per route, needs LOGGING=1), OBS_HTML (1 = failures, +# default: the interactive obs replay pages of the route runs with a flagged infraction or a failed status, rendered +# into $OUT/obs_html at the end; all = every route run, ~100 MB per page; 0 = none; needs LOGGING=1), REPORT (0, +# default; 1 = the nuPlan-style HTML report of scripts/eval/analyze_carla_cosim.py in $OUT/report, needs LOGGING=1), +# TIDY (1, default: scripts/eval/tidy_carla_run.py flattens $OUT at the end, see above; 0 = keep every per-route file), +# 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: vehicles' stopped-time +# obs held at 0, walkers keep theirs; 0 = all real), # 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 +# strictly larger than the smallest training vehicle; 0 = true boxes), COSIM_OBS_SLOTS_PARTNERS_N (default 40: partner slots +# in the policy's obs, GIGAFLOW's eval value; the slots are max-pooled, so a checkpoint trained with 20 transfers; empty = +# the checkpoint's training value), COSIM_WALKER_SPEED_FLOOR_MPS (default 0; the diagnostic walker speed floor of +# shadow_ego.py, tags the folder _walkerfloor), 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). +# system libtiff.so.6 under that name; 0: never link), WANDB_REPORT (1, default: the DS of results.csv becomes the training +# run's final_eval/longest6_v1_DS in wandb, replacing an earlier report; 0 for ablations, so the run keeps the number of +# its standard eval; never reported: a run with missing route records (a completion reports once every route has one), +# the scenario-free variant and a walker speed floor run). +# COSIM_DYNAMICS_SOURCE (carla_pedals, default: CARLA's physics moves the ego, driven only by the VehicleControl the agent +# returns; tags the folder _carla_pedals; pufferdrive: the ego is teleported along PufferDrive's own dynamics). +# Ablations of carla_pedals: COSIM_PEDAL_REVERSE (0, default: the car never backs up; 1 = it reverses, tags the folder _reverse) and +# COSIM_PEDAL_YELLOW_MARGIN (1, default; 0 = the policy sees CARLA's remaining yellow, tags the folder _noyellowmargin). # # 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 +# resubmit; ROUTE_SUBSET narrows the candidates, deleting a route run's result.json forces it). SCENARIOS, SEEDS, REPETITIONS, +# COSIM_VRU_SIZE_PAD_M and COSIM_OBS_SLOTS_PARTNERS_N 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. +# line, and results.csv, the report and the gallery are rebuilt over all routes (also when no route is left to rerun). +# A requeued job (SLURM_RESTART_COUNT > 0 after a preemption or node failure) completes its own folder this way by itself. # 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 @@ -58,12 +77,19 @@ export GARAGE_PY=${GARAGE_PY:-$HOME/miniconda3/envs/garage/bin/python} # cp37: 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} +RUN_DIR=${RUN_DIR:-$PD/experiments/k_scaled_0071_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:-} +# a requeued job (preemption, node failure) resumes its own run folder instead of starting over +if [ -z "$COMPLETE_RUN" ] && [ "${SLURM_RESTART_COUNT:-0}" -gt 0 ]; then + for candidate in "$RUN_DIR"/eval/carla_longest6_v1*_"${SLURM_JOB_ID:-none}"; do # sorted by start time: the newest wins + [ -d "$candidate/routes" ] && [ -f "$candidate/run_info.txt" ] && COMPLETE_RUN=$candidate + done + [ -n "$COMPLETE_RUN" ] && echo "requeued job $SLURM_JOB_ID (restart $SLURM_RESTART_COUNT): resuming $COMPLETE_RUN" +fi 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"; } @@ -74,6 +100,12 @@ if [ -n "$COMPLETE_RUN" ]; then # 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)} + COSIM_OBS_SLOTS_PARTNERS_N=${COSIM_OBS_SLOTS_PARTNERS_N-$(run_info_value obs_slots_partners_n)} # a run from before the knob: empty = the checkpoint's value + COSIM_WALKER_SPEED_FLOOR_MPS=${COSIM_WALKER_SPEED_FLOOR_MPS:-$(run_info_value walker_speed_floor_mps)} + COSIM_DYNAMICS_SOURCE=${COSIM_DYNAMICS_SOURCE:-$(run_info_value dynamics_source)} + COSIM_DYNAMICS_SOURCE=${COSIM_DYNAMICS_SOURCE:-pufferdrive} # a run from before the knob was driven with pufferdrive + COSIM_PEDAL_REVERSE=${COSIM_PEDAL_REVERSE:-$(run_info_value pedal_reverse)} # a pedal run from before these two knobs has no such lines: set both by hand + COSIM_PEDAL_YELLOW_MARGIN=${COSIM_PEDAL_YELLOW_MARGIN:-$(run_info_value pedal_yellow_margin)} fi SCENARIOS=${SCENARIOS:-1} if [ "$SCENARIOS" = "1" ]; then @@ -87,17 +119,32 @@ for seed in $SEEDS; do [[ "$seed" =~ ^[0-9]+$ ]] || { echo "SEEDS=\"$SEEDS\" mus 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 +unit_record() { # the route run's record: in its folder while it drives, flat once tidied + local nested=$OUT/routes/$(unit_dir "$1" "$2")/result.json + [ -f "$nested" ] && echo "$nested" || echo "$OUT/routes/$(unit_dir "$1" "$2").result.json" +} ROUTE_SUBSET=${ROUTE_SUBSET:-} -PORT_BASE=${PORT_BASE:-2000} +PORT_BASE=${PORT_BASE:-12000} MAX_ATTEMPTS=${MAX_ATTEMPTS:-3} -WORKERS=${WORKERS:-16} -STAGGER_S=${STAGGER_S:-15} +WORKERS=${WORKERS:-32} +STAGGER_S=${STAGGER_S:-5} 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} +OBS_HTML=${OBS_HTML:-1} # 1|failures: replay pages of the failed route runs; all: every route run; 0: none +case "$OBS_HTML" in 1) OBS_HTML=failures ;; 0|failures|all) ;; *) echo "OBS_HTML=\"$OBS_HTML\" must be 0, 1, failures or all"; exit 1 ;; esac +REPORT=${REPORT:-0} +TIDY=${TIDY:-1} +WANDB_REPORT=${WANDB_REPORT:-1} export COSIM_VRU_SIZE_PAD_M=${COSIM_VRU_SIZE_PAD_M:-0.8} +export COSIM_OBS_SLOTS_PARTNERS_N=${COSIM_OBS_SLOTS_PARTNERS_N-40} # GIGAFLOW's eval value; empty = the checkpoint's training value +[[ -z "$COSIM_OBS_SLOTS_PARTNERS_N" || "$COSIM_OBS_SLOTS_PARTNERS_N" =~ ^[1-9][0-9]*$ ]] || { echo "COSIM_OBS_SLOTS_PARTNERS_N=\"$COSIM_OBS_SLOTS_PARTNERS_N\" must be a positive integer or empty"; exit 1; } +export COSIM_WALKER_SPEED_FLOOR_MPS=${COSIM_WALKER_SPEED_FLOOR_MPS:-0} # diagnostic only, see shadow_ego.py +[[ "$COSIM_WALKER_SPEED_FLOOR_MPS" =~ ^[0-9]+([.][0-9]+)?$ ]] || { echo "COSIM_WALKER_SPEED_FLOOR_MPS=\"$COSIM_WALKER_SPEED_FLOOR_MPS\" must be a non-negative number"; exit 1; } +export COSIM_DYNAMICS_SOURCE=${COSIM_DYNAMICS_SOURCE:-carla_pedals} # carla_pedals: CARLA's physics moves the ego, see shadow_ego.py +[[ "$COSIM_DYNAMICS_SOURCE" =~ ^(pufferdrive|carla_pedals)$ ]] || { echo "COSIM_DYNAMICS_SOURCE=\"$COSIM_DYNAMICS_SOURCE\" must be pufferdrive or carla_pedals"; exit 1; } +export COSIM_PEDAL_REVERSE=${COSIM_PEDAL_REVERSE:-0} COSIM_PEDAL_YELLOW_MARGIN=${COSIM_PEDAL_YELLOW_MARGIN:-1} # carla_pedals ablations, see shadow_ego.py +[[ "$COSIM_PEDAL_REVERSE" =~ ^[01]$ && "$COSIM_PEDAL_YELLOW_MARGIN" =~ ^[01]$ ]] || { echo "COSIM_PEDAL_REVERSE=\"$COSIM_PEDAL_REVERSE\" and COSIM_PEDAL_YELLOW_MARGIN=\"$COSIM_PEDAL_YELLOW_MARGIN\" must be 0 or 1"; exit 1; } 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 @@ -114,9 +161,35 @@ echo "Evaluating checkpoint: $CKPT" 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 +command -v flock >/dev/null 2>&1 || { echo "flock (util-linux) is needed for the route queue"; exit 1; } + +# Every port of the worker layout (rpc, streaming = rpc+1, TM) must be free of listeners; the first clear base wins. +# Bases stay below the ephemeral range (32768+); the old 2000 put node_exporter (9100) and dcgm-exporter (9400) inside it. +pick_port_base() { + local listening base w port busy + listening=" $(ss -ltun 2>/dev/null | awk 'NR > 1 {sub(/.*:/, "", $5); print $5}' | sort -u | tr '\n' ' ') " + for base in "$@"; do + busy=0 + for ((w = 0; w < WORKERS; w++)); do + for port in $((base + 50 * w)) $((base + 50 * w + 1)) $((base + 6000 + 50 * w)); do + [[ "$listening" == *" $port "* ]] && { busy=$port; break 2; } + done + done + [ "$busy" = 0 ] && { echo "$base"; return 0; } + echo "port base $base: port $busy is in use, trying the next range" >&2 + done + return 1 +} +PORT_BASE=$(pick_port_base "$PORT_BASE" 12000 22000) || { echo "no free port range for $WORKERS workers (ss -ltun)"; exit 1; } +echo "port base $PORT_BASE: rpc $PORT_BASE+50w, TM $((PORT_BASE + 6000))+50w" 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") +TAG=$TAG$([ "$COSIM_VRU_SIZE_PAD_M" = "0" ] || echo "_vrupad$COSIM_VRU_SIZE_PAD_M")${COSIM_OBS_SLOTS_PARTNERS_N:+_partners$COSIM_OBS_SLOTS_PARTNERS_N} +TAG=$TAG$([ "$COSIM_WALKER_SPEED_FLOOR_MPS" = "0" ] || echo "_walkerfloor$COSIM_WALKER_SPEED_FLOOR_MPS") +TAG=$TAG$([ "$COSIM_DYNAMICS_SOURCE" = "pufferdrive" ] || echo "_$COSIM_DYNAMICS_SOURCE") +if [ "$COSIM_DYNAMICS_SOURCE" = "carla_pedals" ]; then + TAG=$TAG$([ "$COSIM_PEDAL_REVERSE" = "0" ] || echo "_reverse")$([ "$COSIM_PEDAL_YELLOW_MARGIN" = "1" ] || echo "_noyellowmargin") +fi # 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" @@ -126,6 +199,10 @@ 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 "obs_slots_partners_n $COSIM_OBS_SLOTS_PARTNERS_N"; echo "walker_speed_floor_mps $COSIM_WALKER_SPEED_FLOOR_MPS" + echo "dynamics_source $COSIM_DYNAMICS_SOURCE" + echo "pedal_reverse $COSIM_PEDAL_REVERSE"; echo "pedal_yellow_margin $COSIM_PEDAL_YELLOW_MARGIN" + echo "port_base $PORT_BASE" echo "git $GIT_STATE" } > "$OUT/run_info.txt" else @@ -167,7 +244,6 @@ 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} @@ -206,43 +282,45 @@ 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") + route_done "$(unit_record "${unit%%:*}" "${unit##*:}")" || 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 + UNITS=(${PENDING[@]+"${PENDING[@]}"}) 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" + if [ "${#UNITS[@]}" -eq 0 ]; then + echo "every route run of $OUT has a valid record: only results.csv, the report and the gallery are rebuilt" + else + echo "completing $OUT: ${#UNITS[@]} of $TOTAL_UNITS route runs without a valid record: ${UNITS[*]}" + # what the completion replaces stays readable next to the live results + 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/" + for f in "$OUT/routes/$route_name".*; do [ -e "$f" ] && mv "$f" "$PREVIOUS/routes/"; done # a tidied run's flat files + rm -f "$(unit_json "${unit%%:*}" "${unit##*:}")" + done + for f in "$OUT"/worker_w*.log; do [ -e "$f" ] && mv "$f" "$PREVIOUS/"; done + fi + echo "completion $COMPLETION_STAMP on $(hostname): route runs ${UNITS[*]:-none}; checkpoint $CKPT; scenarios $SCENARIO_FILE; seeds $SEEDS; repetitions $REPETITIONS; vru_size_pad_m $COSIM_VRU_SIZE_PAD_M; obs_slots_partners_n $COSIM_OBS_SLOTS_PARTNERS_N; walker_speed_floor_mps $COSIM_WALKER_SPEED_FLOOR_MPS; dynamics_source $COSIM_DYNAMICS_SOURCE; pedal_reverse $COSIM_PEDAL_REVERSE; pedal_yellow_margin $COSIM_PEDAL_YELLOW_MARGIN; 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' +# one shared queue, longest route first; length = the xml's waypoint polyline (an earlier run's duration would read a +# route that failed early as short) +QUEUE=$OUT/queue +"$SERVER_PY" - "$ROUTE_DIR" "${UNITS[@]}" > "$QUEUE" <<'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 +route_dir, units = sys.argv[1], sys.argv[2:] # unit = ":" 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") + print(unit) EOF -) +[ "$(wc -l < "$QUEUE")" -eq "${#UNITS[@]}" ] || { echo "ERROR: route queue $QUEUE holds $(wc -l < "$QUEUE") of ${#UNITS[@]} route runs"; exit 1; } +: > "$QUEUE.lock" 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[@]}" +echo " queue, longest first: $(tr '\n' ' ' < "$QUEUE")" [ -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)) @@ -289,7 +367,8 @@ except Exception: sleep 5 done if [ $up -ne 0 ]; then - pkill -9 -P "$server_pid" 2>/dev/null; kill -9 "$server_pid" 2>/dev/null + # reaped here: a dead unreaped server makes the next attempt's wait -n return at once and kill its evaluator + pkill -9 -P "$server_pid" 2>/dev/null; kill -9 "$server_pid" 2>/dev/null; wait "$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 @@ -297,7 +376,7 @@ except Exception: 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") + [ "$OBS_HTML" != "0" ] && 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" \ @@ -320,11 +399,19 @@ except Exception: 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 +next_unit() { # pops the longest remaining route run off $QUEUE under its lock; status 1 once the queue is empty + ( + flock 9 || exit 1 + unit=$(head -n 1 "$QUEUE") + [ -n "$unit" ] || exit 1 + tail -n +2 "$QUEUE" > "$QUEUE.next" && mv "$QUEUE.next" "$QUEUE" || exit 1 + echo "$unit" + ) 9>>"$QUEUE.lock" +} + +run_worker() { # $1 worker: pulls route runs from the queue until it is empty, each retried up to MAX_ATTEMPTS on the next GPU + local w=$1 unit route_id seed attempt gpu + while unit=$(next_unit); do route_id=${unit%%:*}; seed=${unit##*:} for ((attempt = 1; attempt <= MAX_ATTEMPTS; attempt++)); do gpu=$(( (w + attempt - 1) % NUM_GPUS )) @@ -333,16 +420,9 @@ run_worker() { # $1 worker: its assigned route runs, sequentially, each retried echo "[w$w] route $route_id${seed:+ seed $seed} attempt $attempt failed" done done + echo "[w$w] queue empty $(date +%H:%M:%S)" } -# 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=() @@ -352,9 +432,10 @@ for ((w = 0; w < WORKERS; w++)); do run_worker "$w" > "$OUT/worker_w$w.log" 2>&1 & WORKER_PIDS+=($!) done -wait "${WORKER_PIDS[@]}" +wait ${WORKER_PIDS[@]+"${WORKER_PIDS[@]}"} ROUTES_WALL_S=$(( $(date +%s) - ROUTES_T0 )) -cat "$OUT"/worker_w*.log +rm -f "$QUEUE" "$QUEUE.lock" +cat "$OUT"/worker_w*.log 2>/dev/null "$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]) @@ -373,8 +454,9 @@ 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")") +for f in "$OUT"/routes/route_*/result.json "$OUT"/routes/route_*.result.json; do # route folders, or the flat records of a tidied run + [ -e "$f" ] || continue + name=$(basename "$f" .result.json); [ "$name" = "result.json" ] && 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")" @@ -416,18 +498,30 @@ if rows: 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 +# the repo .venv holds wandb; the script reads results.csv and the route records and never imports pufferlib +if [ "$WANDB_REPORT" = "1" ] && [ "$SCENARIOS" = "1" ] && [ "$COSIM_WALKER_SPEED_FLOOR_MPS" = "0" ]; then + env -u PYTHONPATH "$PD/.venv/bin/python" "$PD/pufferlib/ocean/evaluation_utils/eval_wandb_report.py" \ + longest6_v1 "$OUT/aggregate/results.csv" --run-dir "$(dirname "$CKPT")" --expected-episode-count $((TOTAL_UNITS * REPETITIONS)) \ + || echo "wandb report failed (results above are unaffected)" +fi REPORT_ARGS=() -[ "$OBS_HTML" = "1" ] && REPORT_ARGS+=(--obs-html-dir "$OUT/obs_html") +[ "$OBS_HTML" != "0" ] && 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" \ +if [ "$OBS_HTML" != "0" ] && [ "$LOGGING" = "1" ]; then + env -u PYTHONPATH "$SERVER_PY" "$PD/scripts/eval/render_carla_obs_html.py" "$OUT" --routes "$ROUTES_XML" --select "$OBS_HTML" \ || echo "obs replay gallery failed (results above are unaffected)" fi +# the pages are rendered: flatten the route folders, keep the logs of the failed (OBS_HTML=all: every) route runs +if [ "$TIDY" = "1" ]; then + env -u PYTHONPATH "$SERVER_PY" "$PD/scripts/eval/tidy_carla_run.py" "$OUT" --keep "$([ "$OBS_HTML" = "all" ] && echo all || echo failures)" \ + || echo "tidy 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" +for results in "$OUT/results.csv" "$OUT/aggregate/results.csv"; do [ -f "$results" ] && { echo "results -> $results"; break; }; done echo "Done -> $OUT" diff --git a/scripts/kesai/17_carla_longest6_v1_complete_local.sh b/scripts/kesai/17_carla_longest6_v1_complete_local.sh index 1fd936a55d..62a8ba029e 100755 --- a/scripts/kesai/17_carla_longest6_v1_complete_local.sh +++ b/scripts/kesai/17_carla_longest6_v1_complete_local.sh @@ -10,8 +10,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. +# interactive server), STAGGER_S, ROUTE_SUBSET, and every other variable of script 14 (SCENARIOS, SEEDS, REPETITIONS, +# COSIM_VRU_SIZE_PAD_M, COSIM_OBS_SLOTS_PARTNERS_N and COSIM_WALKER_SPEED_FLOOR_MPS 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 } diff --git a/scripts/kesai/18_train_16GPU_gigaflow_slurm.sh b/scripts/kesai/18_train_8GPU_3seeds_slurm.sh similarity index 69% rename from scripts/kesai/18_train_16GPU_gigaflow_slurm.sh rename to scripts/kesai/18_train_8GPU_3seeds_slurm.sh index aae73a75fa..96e76a3b2c 100755 --- a/scripts/kesai/18_train_16GPU_gigaflow_slurm.sh +++ b/scripts/kesai/18_train_8GPU_3seeds_slurm.sh @@ -6,9 +6,10 @@ #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 --output experiments/logs/train_3seeds_%a_%A.out +#SBATCH --error experiments/logs/train_3seeds_%a_%A.err #SBATCH --partition dev +#SBATCH --array=0-1 source scripts/kesai/cluster_env.sh || exit 1 @@ -22,9 +23,19 @@ echo "Master node: ${MASTER_ADDR}:${MASTER_PORT}" echo "Total nodes: ${WORLD_SIZE}" start=$(date +%s) -export SEED=1000 +SEEDS=(1000 2000 1000 2000) +MINIBATCH_SIZES=(262144 262144 393216 393216) +RUN_PREFIXES=(k_scaled_0073 k_scaled_0073 k_scaled_0074 k_scaled_0074) -export RUN_NAME=${RUN_NAME:-k_scaled_0060_${SEED}} +export SEED=${SEEDS[$SLURM_ARRAY_TASK_ID]} +MINIBATCH_SIZE=${MINIBATCH_SIZES[$SLURM_ARRAY_TASK_ID]} +if [ -z "${SEED}" ]; then + echo "SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID} has no run configured; aborting." + exit 1 +fi +echo "SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID} -> train.seed=${SEED} minibatch=${MINIBATCH_SIZE}" + +export RUN_NAME=${RUN_PREFIXES[$SLURM_ARRAY_TASK_ID]}_${SEED} echo ${RUN_NAME} export DATA_DIR=$PD/experiments/${RUN_NAME} @@ -46,24 +57,12 @@ 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} +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}" +# 0073 / 0074 = the 0072 config (3_train_64GPU_multinode_slurm.sh) on one node per training, minibatch 262144 / 393216, seeds 1000 and 2000. Default 0073; 0074: sbatch --array=2-3 scripts/kesai/18_train_8GPU_3seeds_slurm.sh # 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} \ @@ -78,25 +77,17 @@ PUFFER_SLOT_ENCODER_BACKEND=${SLOT_ENCODER_BACKEND} srun --cpus-per-task="$SLURM 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.map_dir=$PD/pufferlib/resources/drive/binaries/carla_128_affine_zones_stop_ndshoulders_v20 \ 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 \ + env.obs_traffic_light_yellow_remaining=true \ + env.scripted_pedestrian_ratio=0.3 \ 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 \ + vec.num_envs=128 \ 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.max_minibatch_size=${MINIBATCH_SIZE} \ + train.minibatch_size=${MINIBATCH_SIZE} \ train.final_model_name=${FINAL_MODEL_NAME} \ train.seed=${SEED} \ tb=True @@ -123,14 +114,15 @@ 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. +# parallel_eval runs 8 srun shards per node (one GPU and an eighth of the CPUs each, 60 blocking workers like +# training) so the eval uses every GPU and thread of the allocation; 60 workers x 5 scenarios = 300 scenarios per 9000-step episode. 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 \ + --total-scenarios 80000 \ --num-nodes 1 \ + --shards-per-node 8 \ env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ - vec.num_envs=128 \ + vec.num_envs=60 \ eval.reward_comfort=0.0 \ eval.reward_lane_center=0.0075 \ env.eval_perceived_size_margin_m=0.2 \ diff --git a/scripts/kesai/21_selfplay_eval_slurm.sh b/scripts/kesai/21_selfplay_eval_slurm.sh new file mode 100755 index 0000000000..f4ba02625b --- /dev/null +++ b/scripts/kesai/21_selfplay_eval_slurm.sh @@ -0,0 +1,64 @@ +#!/bin/bash +#SBATCH --job-name selfplay_eval +#SBATCH --nodes 4 +#SBATCH --ntasks-per-node 1 +#SBATCH --cpus-per-task 480 # whole kesai node without --exclusive (480 schedulable CPUs) +#SBATCH --gres gpu:8 +#SBATCH --mem=0 +#SBATCH --time 2:00:00 +#SBATCH --output experiments/logs/selfplay_eval_%j.out +#SBATCH --error experiments/logs/selfplay_eval_%j.err +#SBATCH --partition dev + +source scripts/kesai/cluster_env.sh || exit 1 + +# CARLA self-play eval only (no co-sim follow-ups), same overrides as the post-training eval in 3_train. +# Defaults below evaluate k_scaled_0059 on the 9000-step layout; override per job with sbatch --export, e.g. the 400 s A/B: +# sbatch --export=ALL,SCENARIO_LENGTH=6000,EVAL_OUTPUT_DIR_NAME=eval80k_400s,WANDB_REPORT=0 scripts/kesai/21_selfplay_eval_slurm.sh +# EVAL_OUTPUT_DIR_NAME names the output folder, so re-evals never overwrite the training run's eval on disk. WANDB_REPORT=1 +# (default) makes this eval's km/infraction the run's final_eval/selfplay_km_per_infraction in wandb; 0 for A/B variants. +echo "START TIME: $(date)" +start=$(date +%s) + +export RUN_NAME=${RUN_NAME:-k_scaled_0059_1000} +export MODEL_PATH=$PD/experiments/${RUN_NAME}/final_model.pt +if [ ! -f "${MODEL_PATH}" ]; then + echo "missing checkpoint ${MODEL_PATH}" >&2 + exit 1 +fi +EVAL_OUTPUT_DIR_NAME=${EVAL_OUTPUT_DIR_NAME:-eval80k} +TOTAL_SCENARIOS=${TOTAL_SCENARIOS:-80000} +echo "run ${RUN_NAME} output dir ${EVAL_OUTPUT_DIR_NAME} scenarios ${TOTAL_SCENARIOS} scenario_length ${SCENARIO_LENGTH:-benchmark default}" + +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 + +PUFFER_SLOT_ENCODER_BACKEND=${SLOT_ENCODER_BACKEND:-cuda} .venv/bin/python scripts/parallel_eval.py carla \ + --total-scenarios ${TOTAL_SCENARIOS} \ + --num-nodes ${SLURM_JOB_NUM_NODES} \ + --shards-per-node 8 \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ + vec.num_envs=60 \ + 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 \ + ${SCENARIO_LENGTH:+env.scenario_length=${SCENARIO_LENGTH}} \ + eval.output_dir_name=${EVAL_OUTPUT_DIR_NAME} \ + eval.output_name=${RUN_NAME} \ + load_model_path=${MODEL_PATH} \ + wandb=${WANDB_REPORT:-1} + +end=$(date +%s) +echo "END TIME: $(date)" +echo "Runtime: $((end-start))" diff --git a/scripts/kesai/22_ensemble_eval_slurm.sh b/scripts/kesai/22_ensemble_eval_slurm.sh new file mode 100755 index 0000000000..2cfae78ef0 --- /dev/null +++ b/scripts/kesai/22_ensemble_eval_slurm.sh @@ -0,0 +1,106 @@ +#!/bin/bash +#SBATCH --job-name ensemble_eval +#SBATCH --nodes 1 # the self-play eval uses every node of the allocation; more nodes for bigger ensembles +#SBATCH --ntasks-per-node 1 +#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 1-00:00 +#SBATCH --output experiments/logs/ensemble_eval_%j.out +#SBATCH --error experiments/logs/ensemble_eval_%j.err +#SBATCH --partition dev + +source scripts/kesai/cluster_env.sh || exit 1 + +# Evaluates an ensemble of trained runs with the evals of 3_train / 18_train. scripts/eval/build_ensemble_checkpoint.py +# packs the members' final_model.pt into $PD/experiments/$ENSEMBLE_NAME/final_model.pt next to the first member's +# config.yaml, so every eval takes the folder like a training run: at every policy step the ensemble executes the mean of +# the members' mean actions (one forward pass per member; mode / sampled actions and the logged probabilities come from +# the members' averaged distribution). +# The ensemble has no wandb run: the chained evals get WANDB_REPORT=0 (12_alpasim.sh has no such switch; its report step +# fails harmlessly) and the self-play eval runs with wandb=False; results live in experiments/$ENSEMBLE_NAME/eval/. +echo "START TIME: $(date)" +start=$(date +%s) + +ENSEMBLE_NAME=k_scaled_0073_0074_ensemble +MEMBER_RUNS=(k_scaled_0073_1000 k_scaled_0073_2000 k_scaled_0074_1000 k_scaled_0074_2000) # run folders under $PD/experiments, each with a final_model.pt +echo "ensemble ${ENSEMBLE_NAME} of ${MEMBER_RUNS[*]}" + +DATA_DIR=$PD/experiments/${ENSEMBLE_NAME} +MODEL_PATH=${DATA_DIR}/final_model.pt +MEMBER_CHECKPOINTS=() +for member_run in "${MEMBER_RUNS[@]}"; do + member_checkpoint=$PD/experiments/${member_run}/final_model.pt + if [ ! -f "${member_checkpoint}" ]; then + echo "missing member checkpoint ${member_checkpoint}" + exit 1 + fi + MEMBER_CHECKPOINTS+=("${member_checkpoint}") +done + +# 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=cuda +echo "slot encoder backend: ${SLOT_ENCODER_BACKEND}" + +source .venv/bin/activate +bash scripts/kesai/build_ext_if_changed.sh $PD || exit 1 + +echo "Building ${MODEL_PATH}" +.venv/bin/python scripts/eval/build_ensemble_checkpoint.py \ + --output-dir ${DATA_DIR} \ + --run-name ${ENSEMBLE_NAME} \ + "${MEMBER_CHECKPOINTS[@]}" || exit 1 + +# 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 "Submitting nuPlan reactive eval for ${MODEL_PATH}" +WANDB_REPORT=0 RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh \ + || echo "nuPlan eval submission failed; run by hand: WANDB_REPORT=0 RUN_DIR=${DATA_DIR} sbatch scripts/kesai/9_nuPlan.sh" +echo "Submitting nuPlan non-reactive eval for ${MODEL_PATH}" +WANDB_REPORT=0 RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh \ + || echo "nuPlan non-reactive eval submission failed; run by hand: WANDB_REPORT=0 RUN_DIR=${DATA_DIR} sbatch scripts/kesai/13_nuPlan_nonreactive.sh" +echo "Submitting longest6 v2 eval for ${MODEL_PATH}" +WANDB_REPORT=0 RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6_v2.sh \ + || echo "longest6 v2 eval submission failed; run by hand: WANDB_REPORT=0 RUN_DIR=${DATA_DIR} sbatch scripts/kesai/11_carla_longest6_v2.sh" +echo "Submitting longest6 v1 eval for ${MODEL_PATH}" +WANDB_REPORT=0 RUN_DIR=${DATA_DIR} sbatch scripts/kesai/14_carla_longest6_v1.sh \ + || echo "longest6 v1 eval submission failed; run by hand: WANDB_REPORT=0 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 runs 8 srun shards per node (one GPU and an eighth of the CPUs each, 60 blocking workers like +# training) so the eval uses every GPU and thread of the allocation; 60 workers x 5 scenarios = 300 scenarios per 9000-step episode. +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 80000 \ + --num-nodes ${SLURM_JOB_NUM_NODES} \ + --shards-per-node 8 \ + env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ + vec.num_envs=60 \ + 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=${ENSEMBLE_NAME} \ + load_model_path=${MODEL_PATH} \ + wandb=False + +end=$(date +%s) +runtime=$((end-start)) +echo "END TIME: $(date)" +echo "Runtime: ${runtime}" diff --git a/scripts/kesai/3_train_64GPU_multinode_slurm.sh b/scripts/kesai/3_train_64GPU_multinode_slurm.sh index 1e99d19a95..87e428dce0 100644 --- a/scripts/kesai/3_train_64GPU_multinode_slurm.sh +++ b/scripts/kesai/3_train_64GPU_multinode_slurm.sh @@ -24,7 +24,7 @@ start=$(date +%s) export SEED=1000 -export RUN_NAME=${RUN_NAME:-k_scaled_0059_${SEED}} +export RUN_NAME=${RUN_NAME:-k_scaled_0076_${SEED}} echo ${RUN_NAME} export DATA_DIR=$PD/experiments/${RUN_NAME} @@ -51,11 +51,31 @@ echo "slot encoder backend: ${SLOT_ENCODER_BACKEND}" source .venv/bin/activate 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. +# Blocking env workers: kesai nodes give 60 threads per rank (--exclusive); vec.num_envs=60 is one worker per thread (0053/0059, +# ~25 M SPS vs ~21.5 M with 32). 0063 ran 256 workers per rank (4.3x oversubscribed, 75 GB per GPU): +18% SPS from the worker count +# but ~55% sample efficiency on goals/speed at 1e12 steps, so 60 stays the default. GPU memory ~21 GB fixed + 0.22 GB per 1024-agent worker (fp16). +# A second instance needs its own RUN_NAME (plain VAR=... sbatch did not propagate on kesai): sbatch --export=ALL,RUN_NAME=k_scaled_0076_1000 scripts/kesai/3_train_64GPU_multinode_slurm.sh +# 0067 = 0066 (the 0059 config, puffer_drive.yaml defaults, float16 rollout buffer, static footprint red light test) with the +# training light cycle green U(3,10) / yellow U(2,3) s instead of GIGAFLOW's U(1,10) / U(1.5,2.25) and no training rear margin: +# 0066's training red-light rate sat at 3.7% because committed slow crossings could not clear CARLA's tail band within a +# 1.5 s yellow (8-map replay sweep: 17 -> 5 onsets). 0068 = 0067 with 128 workers per rank. +# 0069 = 0068 + env.obs_traffic_light_yellow_remaining: every traffic-control row carries the seconds the light stays +# yellow / 3 s (the training max; co-sims clip their timers to it). The evals below read the flag from the run's config.yaml. +# 0070 = 0069 + env.scripted_pedestrian_ratio=0.1: scripted pedestrian crossings (CARLA Scenario3 geometry: ceil(0.1 x agents) +# walker slots per env pop in at the right kerb 12-20 m ahead of a moving car and cross at 2-4.2 m/s as 1.0-1.5 m boxes). +# Ablation = the same job with env.scripted_pedestrian_ratio=0. The eval profiles pin the ratio to 0. +# 0071 = 0070 with the Leaderboard 1.0 walker profile (drive.h, 2026-11-04): walkers appear standing 12 m + 1.5 m per extra +# lane ahead (+-2 m), half a lane outside the kerb, wait 0-0.6 s and ramp up over 0.3-0.7 s to 3.4 + 0.4 m/s per lane (+-1); +# 0070's walkers set off at 2-4.2 m/s from their first frame, so its policy never learned to brake for a standing walker. +# 0072 = 0071 with five pedestrian behaviours (constants.h PEDESTRIAN_*, 2026-11-05): 40% the 0071 pop-in crossers, 15% slow +# crossers from either kerb, 20% standers, 15% kerb walkers, 10% kerb steppers. 0071 waited for every standing or passing +# pedestrian (longest6 v1 timeouts at stuck walkers, nuPlan CLS -0.02 / -0.035). Ratio 0.1 -> 0.3 keeps the pop-in crossings +# per car where 0071 had them (1.8 per episode); the walker slots are visible 48% of the time. +# 0075 = 0072 with 0074's per-rank minibatch (393216, fills the GPU; 0072 ran 131072) on the 2026-11-09 audit fixes: map +# layouts differ per rank and seed, failed init spawns keep their slot, light phases round up to whole steps (yellow +# 2.1-3.0 s at dt 0.3), offroad is detected up to 40 m/s, spawns use the red-light detector, reset no longer consumes goals. +# 0076 = 0075 with the 0072 minibatch (131072): 393216 on 32 ranks is 12.6M rows per optimizer step and gave 0075 62k optimizer steps against 175k in 0072. +# Other overrides are run identity, logging, maps and hardware. # 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} \ @@ -70,13 +90,15 @@ PUFFER_SLOT_ENCODER_BACKEND=${SLOT_ENCODER_BACKEND} srun --cpus-per-task="$SLURM 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.map_dir=$PD/pufferlib/resources/drive/binaries/carla_128_affine_zones_stop_ndshoulders_v20 \ env.num_maps=128 \ + env.obs_traffic_light_yellow_remaining=true \ + env.scripted_pedestrian_ratio=0.3 \ train.name=${RUN_NAME} \ run_name=${RUN_NAME} \ train.total_timesteps=1000000000000 \ - vec.num_envs=60 \ - ${ROLLOUT_DTYPE:+train.rollout_dtype=$ROLLOUT_DTYPE} \ + vec.num_envs=128 \ + train.rollout_dtype=float16 \ train.max_minibatch_size=131072 \ train.minibatch_size=131072 \ train.final_model_name=${FINAL_MODEL_NAME} \ @@ -105,14 +127,15 @@ 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. +# parallel_eval runs 8 srun shards per node (one GPU and an eighth of the CPUs each, 60 blocking workers like +# training) so the eval uses every GPU and thread of the allocation; 60 workers x 5 scenarios = 300 scenarios per 9000-step episode. 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 \ + --total-scenarios 80000 \ --num-nodes 4 \ + --shards-per-node 8 \ env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ - vec.num_envs=128 \ + vec.num_envs=60 \ eval.reward_comfort=0.0 \ eval.reward_lane_center=0.0075 \ env.eval_perceived_size_margin_m=0.2 \ diff --git a/scripts/kesai/7_eval_8node_slurm.sh b/scripts/kesai/7_eval_8node_slurm.sh index d04d767a4f..4b956d5809 100644 --- a/scripts/kesai/7_eval_8node_slurm.sh +++ b/scripts/kesai/7_eval_8node_slurm.sh @@ -44,14 +44,15 @@ if [ ! -f ${MODEL_PATH} ]; then exit 1 fi -# 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. +# parallel_eval runs 8 srun shards per node (one GPU and an eighth of the CPUs each, 60 blocking workers like +# training) so the eval uses every GPU and thread of the allocation; 60 workers x 5 scenarios = 300 scenarios per 9000-step episode. echo "Evaluating ${MODEL_PATH}" .venv/bin/python scripts/parallel_eval.py carla \ - --total-scenarios 40000 \ + --total-scenarios 80000 \ --num-nodes 2 \ + --shards-per-node 8 \ env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ - vec.num_envs=128 \ + vec.num_envs=60 \ eval.reward_comfort=0.0 \ eval.reward_lane_center=0.0075 \ env.eval_perceived_size_margin_m=0.2 \ diff --git a/scripts/kesai/8_train_64GPU_multinode_slurm.sh b/scripts/kesai/8_train_64GPU_multinode_slurm.sh index b9ad229177..7e09dc0efa 100644 --- a/scripts/kesai/8_train_64GPU_multinode_slurm.sh +++ b/scripts/kesai/8_train_64GPU_multinode_slurm.sh @@ -86,14 +86,15 @@ if [ ! -f ${MODEL_PATH} ]; then exit 1 fi -# parallel_eval places one shard per allocated node via srun, so each shard's 64 -# env workers get a full node's cores instead of sharing the batch host. +# parallel_eval runs 8 srun shards per node (one GPU and an eighth of the CPUs each, 60 blocking workers like +# training) so the eval uses every GPU and thread of the allocation; 60 workers x 5 scenarios = 300 scenarios per 9000-step episode. echo "Training done, evaluating ${MODEL_PATH}" .venv/bin/python scripts/parallel_eval.py carla \ - --total-scenarios 40000 \ + --total-scenarios 80000 \ --num-nodes 8 \ + --shards-per-node 8 \ env.map_dir=$PD/pufferlib/resources/drive/binaries/carla \ - vec.num_envs=64 \ + vec.num_envs=60 \ eval.reward_comfort=0.0 \ eval.reward_lane_center=0.0075 \ env.eval_perceived_size_margin_m=0.2 \ diff --git a/scripts/kesai/9_nuPlan.sh b/scripts/kesai/9_nuPlan.sh index a39aa8428c..10fc653170 100644 --- a/scripts/kesai/9_nuPlan.sh +++ b/scripts/kesai/9_nuPlan.sh @@ -13,20 +13,25 @@ set -u source scripts/kesai/cluster_env.sh || exit 1 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:-$PD/experiments/k_scaled_0045_1000} +# Overridable env: RUN_DIR (3_train_64GPU_multinode_slurm.sh submits this job with RUN_DIR=), CHALLENGES, +# WANDB_REPORT (1, default: the CLS becomes the training run's final_eval/nuplan_{reactive,nonreactive}_CLS in wandb, +# replacing an earlier report of that challenge; 0 for ablations, so the run keeps the number of its standard eval) +RUN_DIR=${RUN_DIR:-$PD/experiments/k_scaled_0059_1000} +case "$RUN_DIR" in /*) ;; *) RUN_DIR=$PD/$RUN_DIR ;; esac # relative to the checkout: $PD is only defined inside the job # 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; } # reactive = IDM agents (CLS-R), nonreactive = log-replay agents (CLS-NR) export CHALLENGES=${CHALLENGES:-closed_loop_reactive_agents_pufferdrive} case "$CHALLENGES" in - closed_loop_reactive_agents_pufferdrive) CHALLENGE_TAG="" ;; - closed_loop_nonreactive_agents_pufferdrive) CHALLENGE_TAG="_nonreactive" ;; + closed_loop_reactive_agents_pufferdrive) CHALLENGE_TAG=""; WANDB_EVAL_TYPE=nuplan_reactive ;; + closed_loop_nonreactive_agents_pufferdrive) CHALLENGE_TAG="_nonreactive"; WANDB_EVAL_TYPE=nuplan_nonreactive ;; *) echo "CHALLENGES must be closed_loop_reactive_agents_pufferdrive or closed_loop_nonreactive_agents_pufferdrive"; exit 1 ;; esac +WANDB_REPORT=${WANDB_REPORT:-1} +VAL14_SCENARIO_COUNT=1118 # a run that scored fewer scenarios (LIMIT_TOTAL_SCENARIOS, failed simulations) is not reported echo "Evaluating checkpoint: $CKPT" -export CITY_BIN_DIR=$SHARED_DATA/nuplan/PufferDrive +export CITY_BIN_DIR=${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=$SHARED_DATA/nuplan export NUPLAN_MAPS_ROOT=$SHARED_DATA/nuplan/maps @@ -39,13 +44,16 @@ export PYTHONPATH=$PD:$CARL_DEVKIT_ROOT:$NUPLAN_DEVKIT_ROOT${INHERITED_PYTHONPAT cd "$PD" || exit 1 export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 NUMEXPR_NUM_THREADS=1 export COSIM_OBS_HTML=infractions # obs replays (what the policy saw) only for collision/offroad/wrong-way/no-progress scenarios; all|failures|infractions|0 -# carl_visualization_callback: nuPlan ground-truth video per scenario -> $GROUP/simulation///visualization -export CALLBACKS="[simulation_log_callback, carl_visualization_callback]" +# one simulation log per scenario; the tidy keeps those of the COSIM_OBS_HTML scenarios in simulation_log/. Adding +# carl_visualization_callback renders nuPlan's ground-truth video of every scenario (1118 .avi, 14 GB, a matplotlib +# frame per step; set OPENCV_FOR_THREADS_NUM=1 against its thread pool) for the few the tidy keeps as details/videos/. +export CALLBACKS="[simulation_log_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. -# 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)} +# 128 fits the 144-CPU scaleway nodes. A kesai RTX node gives the job 480 CPUs on 256 physical cores and a worker plateaus +# at ~2.0 GB once it has cached all four city maps (256 workers = 509 GB, 30 Oct): 384 fills half the hyperthreads and +# stays under ~770 GB of the 1 TB, leaving the rest to the page cache of the logs. +export THREADS_PER_NODE=${THREADS_PER_NODE:-$([ "$CLUSTER" = kesai ] && echo 384 || 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 @@ -55,6 +63,10 @@ export STARTUP_ACCEL_JERK_CAP=2 export STARTUP_BRAKE_JERK_CAP=2 # 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 +# lane width = 2 x distance to the nearer GPKG boundary (city bins store left + right, 6 m on nuPlan's parking-strip lanes +# whose baseline sits 1.7 m from the oncoming side); cached copy of the bin in NUPLAN_EXP_ROOT/pufferdrive_bin_cache. +# Default off: alone it scored 0.9257 vs 0.9314 for the plain clip (run 295, 2026-10-29): TTC and progress worse, lanes read 2.85-3.4 m. +export SYMMETRIC_LANE_WIDTHS=${SYMMETRIC_LANE_WIDTHS:-false} # 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). @@ -62,15 +74,79 @@ 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 +# Lane-width clip [m]: the policy observes lane widths clipped to the training range (CARLA maps 2.85-4.4 m; nuPlan +# city bins have 22% of segments wider than 5 m, up to 50 m). LANE_WIDTH_CLIP_MAX_M=0 turns it off (pre-2026-10-29 evals). +export LANE_WIDTH_CLIP_MIN_M=${LANE_WIDTH_CLIP_MIN_M:-2.85} +export LANE_WIDTH_CLIP_MAX_M=${LANE_WIDTH_CLIP_MAX_M:-4.4} +# Ego plant (default since 2026-10-30): nuplan_bicycle = nuPlan's KinematicBicycleModel via CaRL's one_stage_controller, +# driven exactly as CaRL drives it, with the model's own actuator lag (BICYCLE_UNDO_LOW_PASS=false: it realises +# dt/(dt+tau) of each command; run 431 = 93.01). EGO_DYNAMICS=pufferdrive = the pre-2026-10-30 perfect-tracking eval +# (shadow jerk model integrates the ego, nuPlan replays the pose; hides yaw rate and the center-accel term from the +# comfort metrics; 93.12). Ablated: BICYCLE_MIN_SPEED_MPS -2 (Drive's floor) costs ~0.9 on every safety metric; +# BICYCLE_UNDO_LOW_PASS=true (feedforward inverse of the lag using the plant's time constants) scores 93.48 (run 428) +# by trading +2.7 comfort for 6 collisions, 4 offroads and 1.2 TTC; kept off to stay faithful to the CaRL plant. +export EGO_DYNAMICS=${EGO_DYNAMICS:-nuplan_bicycle} +# EGO_DYNAMICS=nuplan_tracker = GIGAFLOW's nuPlan evaluation: every step the planner rolls the scene out for the 8 s +# trajectory horizon in ROLLOUT_STEP_SECONDS steps (the policy drives the ego and the nearest ROLLOUT_POLICY_VEHICLES +# vehicles and bicycles, the rest keeps its velocity) and nuPlan's two_stage_controller tracks the ego's rolled-out +# trajectory. EGO_TRACKER: lqr = nuPlan's default tracker, ilqr = its iLQR tracker without the wall-clock solve cap. +# ROLLOUT_ALLOW_REVERSE=false holds an agent the policy would back up at rest (as BICYCLE_MIN_SPEED_MPS=0). +# ROLLOUT_POLICY_VEHICLES=0: everyone else keeps its observed velocity. 32 (runs 2566/2567, LQR 0.897 / iLQR 0.870 vs +# 0.923 for the bicycle plant) forecast queue leaders pulling up to the stop line and the ego drove into the gap. +# Set WANDB_REPORT=0 for a tracker run unless its CLS should replace the standard eval's number in wandb. +export EGO_TRACKER=${EGO_TRACKER:-lqr} +export ROLLOUT_STEP_SECONDS=${ROLLOUT_STEP_SECONDS:-0.3} +export ROLLOUT_POLICY_VEHICLES=${ROLLOUT_POLICY_VEHICLES:-0} +export ROLLOUT_ALLOW_REVERSE=${ROLLOUT_ALLOW_REVERSE:-false} +export BICYCLE_MIN_SPEED_MPS=${BICYCLE_MIN_SPEED_MPS:-0} +export BICYCLE_UNDO_LOW_PASS=${BICYCLE_UNDO_LOW_PASS:-false} +# Comfort: nuPlan scores the center accel (rear + d*yaw_rate^2 + d*yaw_accel), which the bicycle plant exposes. +# BICYCLE_CENTER_ACCEL_CAP=true reduces the plant's accel while turning (CaRL clip_angular_adjustment; 248 -> 20 +# positive lon-accel fails), MARGIN 0.95 takes the Savitzky-Golay overshoot residual (+0.2 CLS, run 426). +# Dead ends: EGO_ACCEL_CAP_MPS2=2.0 (global cap, costs progress, run 421) and BICYCLE_MAX_YAW_JERK=2 (steering +# inertia the policy never saw: offroad 5%, run 427). +export BICYCLE_CENTER_ACCEL_CAP=${BICYCLE_CENTER_ACCEL_CAP:-true} +export BICYCLE_CENTER_ACCEL_CAP_MARGIN=${BICYCLE_CENTER_ACCEL_CAP_MARGIN:-0.95} +export BICYCLE_MAX_YAW_JERK=${BICYCLE_MAX_YAW_JERK:-0} +export EGO_ACCEL_CAP_MPS2=${EGO_ACCEL_CAP_MPS2:-} +# Run-426 anatomy arms: PARTNER_SPLIT_MAX_LENGTH_M=7 (buses up to 12.7 m vs training max 7 m), STANDSTILL_JERK_DEADBAND=0 +# (five stalls were the policy's 0-1.5 m/s^3 start intent being zeroed; the floor now does the deadband's anti-reverse +# job). Drive's jerk model keeps its axle at the box centre while nuPlan's sits 1.46 m behind (four inside-rear corner +# collisions in run 426): a training-side item, no eval-side conversion. +# true since 2026-10-30 evening (tag _clamp): the policy's integrated accel is zeroed while the plant sits on the floor +# (9 of run 426's 18 stalls had it wound down to -5 at standstill); false = run 426 behaviour. +export BICYCLE_FLOOR_ACCEL_CLAMP=${BICYCLE_FLOOR_ACCEL_CLAMP:-true} +export PARTNER_SPLIT_MAX_LENGTH_M=${PARTNER_SPLIT_MAX_LENGTH_M:-} +export STANDSTILL_JERK_DEADBAND=${STANDSTILL_JERK_DEADBAND:-} 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}" [ "$VRU_SIZE_PAD_M" != "0" ] && ABLATION_TAG="${ABLATION_TAG}_vrupad${VRU_SIZE_PAD_M}m" +[ "$SYMMETRIC_LANE_WIDTHS" = "true" ] && ABLATION_TAG="${ABLATION_TAG}_symw" +[ "$CITY_BIN_DIR" != "$SHARED_DATA/nuplan/PufferDrive" ] && ABLATION_TAG="${ABLATION_TAG}_bins$(basename "$CITY_BIN_DIR")" [ -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 +[ "$LANE_WIDTH_CLIP_MAX_M" != "0" ] && ABLATION_TAG="${ABLATION_TAG}_lanew${LANE_WIDTH_CLIP_MIN_M}-${LANE_WIDTH_CLIP_MAX_M}" +if [ "$EGO_DYNAMICS" = "nuplan_bicycle" ]; then + ABLATION_TAG="${ABLATION_TAG}_bicycle_vmin${BICYCLE_MIN_SPEED_MPS}" + [ "$BICYCLE_UNDO_LOW_PASS" = "true" ] && ABLATION_TAG="${ABLATION_TAG}_nolag" + if [ "$BICYCLE_CENTER_ACCEL_CAP" = "true" ]; then + ABLATION_TAG="${ABLATION_TAG}_centercap" + [ "$BICYCLE_CENTER_ACCEL_CAP_MARGIN" != "1.0" ] && ABLATION_TAG="${ABLATION_TAG}${BICYCLE_CENTER_ACCEL_CAP_MARGIN}" + fi + [ "$BICYCLE_MAX_YAW_JERK" != "0" ] && ABLATION_TAG="${ABLATION_TAG}_yawjerk${BICYCLE_MAX_YAW_JERK}" + [ "$BICYCLE_FLOOR_ACCEL_CLAMP" = "true" ] && ABLATION_TAG="${ABLATION_TAG}_clamp" +fi +if [ "$EGO_DYNAMICS" = "nuplan_tracker" ]; then + ABLATION_TAG="${ABLATION_TAG}_tracker_${EGO_TRACKER}_step${ROLLOUT_STEP_SECONDS}_veh${ROLLOUT_POLICY_VEHICLES}" + [ "$ROLLOUT_ALLOW_REVERSE" = "true" ] && ABLATION_TAG="${ABLATION_TAG}_reverse" +fi +[ -n "$EGO_ACCEL_CAP_MPS2" ] && ABLATION_TAG="${ABLATION_TAG}_acap${EGO_ACCEL_CAP_MPS2}" +[ -n "$PARTNER_SPLIT_MAX_LENGTH_M" ] && ABLATION_TAG="${ABLATION_TAG}_split${PARTNER_SPLIT_MAX_LENGTH_M}" +[ -n "$STANDSTILL_JERK_DEADBAND" ] && ABLATION_TAG="${ABLATION_TAG}_deadband${STANDSTILL_JERK_DEADBAND}" +# results: the model's own eval folder, next to the PufferDrive benchmark evals (overview.csv, scenarios.csv, obs_html/, simulation_log/, details/) export GROUP=$RUN_DIR/eval/nuplan_val14${CHALLENGE_TAG}_${GOAL_SOURCE}${ABLATION_TAG}_$(date +%Y%m%d_%H%M%S)_${SLURM_JOB_ID:-local} # NO_TRAIN skips the training-only CUDA extension (_C, never imported by the planner): the cp39 nuPlan env builds only @@ -85,3 +161,16 @@ case "$IMPORTED_PLANNER" in esac bash "$PD/pufferlib/ocean/cosim/nuplan/run_nuplan_planner.sh" +RUNNER_STATUS=$? +# the repo .venv holds wandb; the script reads overview.csv and never imports pufferlib +if [ "$WANDB_REPORT" = "1" ] && [ $RUNNER_STATUS -eq 0 ]; then + env -u PYTHONPATH "$PD/.venv/bin/python" "$PD/pufferlib/ocean/evaluation_utils/eval_wandb_report.py" \ + "$WANDB_EVAL_TYPE" "$GROUP/overview.csv" --run-dir "$RUN_DIR" --expected-episode-count "$VAL14_SCENARIO_COUNT" \ + || echo "wandb report failed (results above are unaffected)" +fi +# the job's own log next to the results (the SBATCH --output of this script or of 13_nuPlan_nonreactive.sh) +JOB_LOGS=(experiments/logs/nuPlan_*_"${SLURM_JOB_ID:-none}".out experiments/logs/nuPlan_*_"${SLURM_JOB_ID:-none}".err) +for job_log in "${JOB_LOGS[@]}"; do + [ -f "$job_log" ] && [ -d "$GROUP/details" ] && cp "$job_log" "$GROUP/details/" +done +exit $RUNNER_STATUS diff --git a/scripts/kesai/cluster_env.sh b/scripts/kesai/cluster_env.sh index 096c458f4d..a564473bba 100644 --- a/scripts/kesai/cluster_env.sh +++ b/scripts/kesai/cluster_env.sh @@ -29,6 +29,8 @@ if [ ! -f "$PD/setup.py" ]; then return 1 2>/dev/null || exit 1 fi export SLURM_CPUS_PER_TASK=${SLURM_CPUS_ON_NODE:-${SLURM_CPUS_PER_TASK:-$(nproc)}} +# pufferlib opens 4 pipe fds per env worker; the job's soft nofile limit (1024 on kesai) caps workers near 200 per rank +ulimit -n "$(ulimit -Hn)" 2>/dev/null || echo "cluster_env.sh: could not raise the open-file limit (soft $(ulimit -Sn), hard $(ulimit -Hn))" >&2 export CLUSTER HOME_BASE SHARED_DATA PD echo "cluster: $CLUSTER checkout: $PD shared data: ${SHARED_DATA:-none} cpus per task: $SLURM_CPUS_PER_TASK" diff --git a/scripts/parallel_eval.py b/scripts/parallel_eval.py index 97312f0c71..e9f9f64b64 100644 --- a/scripts/parallel_eval.py +++ b/scripts/parallel_eval.py @@ -2,16 +2,19 @@ Each shard runs an independent `puffer eval` process over a disjoint scenario window (eval.scenario_offset), then the per-shard episode_metrics.csv files are -merged into a single report and optionally attached to the training wandb run. ---num-gpus runs one shard per local GPU; --num-nodes runs one shard per node of -the surrounding SLURM allocation (each claiming a full node's CPUs and one GPU -via srun), so every shard gets its own node's cores. +merged into a single report; wandb=True writes its km/infraction into the summary +of the training wandb run (self-play benchmark only). scripts/eval/tidy_selfplay_eval.py +then folds the shard folders into the aggregate dir (evaluation_summary.json, +obs_html/ with every shard's failure pages, details/); --keep-shards leaves them. +--num-gpus runs one shard per local GPU; --num-nodes runs the shards as srun steps +of the surrounding SLURM allocation, --shards-per-node of them per node (default 1), +each with one GPU and an equal share of the node's CPUs (SLURM_CPUS_PER_TASK). Usage (from the repo root, venv active): python scripts/parallel_eval.py carla \ - --total-scenarios 40000 --num-gpus 8 \ + --total-scenarios 80000 --num-nodes 4 --shards-per-node 8 \ load_model_path=experiments/run/final_model.pt \ - eval.output_name=run vec.num_envs=16 wandb=True + eval.output_name=run vec.num_envs=60 wandb=True """ import argparse @@ -33,8 +36,21 @@ def parse_arguments(): parser.add_argument("--total-scenarios", type=int, required=True) shard_group = parser.add_mutually_exclusive_group(required=True) shard_group.add_argument("--num-gpus", type=int, help="Run one shard per local GPU") - shard_group.add_argument("--num-nodes", type=int, help="Run one shard per SLURM node via srun, 1 GPU each") + shard_group.add_argument( + "--num-nodes", type=int, help="Run shards as srun steps over this many SLURM nodes, 1 GPU each" + ) + parser.add_argument("--keep-shards", action="store_true", help="Leave the per-shard folders untidied") + parser.add_argument( + "--shards-per-node", + type=int, + default=1, + help="With --num-nodes: steps per node, one GPU each, the node's CPUs split evenly between them", + ) arguments, overrides = parser.parse_known_args() + if arguments.shards_per_node < 1: + parser.error("--shards-per-node must be at least 1") + if arguments.shards_per_node > 1 and arguments.num_nodes is None: + parser.error("--shards-per-node requires --num-nodes") for override in overrides: if override.startswith("--"): parser.error(f"Unknown flag {override}; Hydra overrides use key=value syntax") @@ -83,7 +99,7 @@ def print_log_tail(log_path): def main(): arguments, overrides = parse_arguments() use_srun = arguments.num_nodes is not None - num_shards = arguments.num_nodes if use_srun else arguments.num_gpus + num_shards = arguments.num_nodes * arguments.shards_per_node if use_srun else arguments.num_gpus if num_shards < 1: raise SystemExit("--num-gpus/--num-nodes must be at least 1") if arguments.total_scenarios < num_shards: @@ -91,9 +107,20 @@ def main(): if use_srun: if "SLURM_JOB_ID" not in os.environ: raise SystemExit("--num-nodes requires running inside a SLURM allocation") - cpus_per_shard = os.environ.get("SLURM_CPUS_PER_TASK") - if cpus_per_shard is None: + cpus_per_node = os.environ.get("SLURM_CPUS_PER_TASK") + if cpus_per_node is None or not cpus_per_node.isdigit(): raise SystemExit("--num-nodes requires SLURM_CPUS_PER_TASK (set #SBATCH --cpus-per-task)") + cpus_per_shard = int(cpus_per_node) // arguments.shards_per_node + if cpus_per_shard < 1: + raise SystemExit(f"--shards-per-node {arguments.shards_per_node} exceeds the {cpus_per_node} CPUs per node") + numa_note = ", in-process NUMA bind off" if arguments.shards_per_node > 1 else "" + print( + f"{num_shards} shards: {arguments.num_nodes} nodes x {arguments.shards_per_node} per node, " + f"{cpus_per_shard} CPUs each{numa_note}" + ) + num_envs_override = override_value(overrides, "vec.num_envs") + if num_envs_override is not None and num_envs_override.isdigit() and int(num_envs_override) > cpus_per_shard: + print(f"Note: vec.num_envs={num_envs_override} blocking workers exceed a shard's {cpus_per_shard} CPUs") for override in overrides: if override.partition("=")[0] in RESERVED_OVERRIDE_KEYS: raise SystemExit(f"Override '{override}' is managed by parallel_eval.py; remove it") @@ -105,7 +132,13 @@ def main(): report_wandb = str(override_value(overrides, "wandb")).lower() in ("true", "1") shard_overrides = [o for o in overrides if o.partition("=")[0] != "wandb"] + ["wandb=False"] - from pufferlib.ocean.evaluation_utils import evaluation_utils + from pufferlib.ocean.evaluation_utils import eval_wandb_report, evaluation_utils + + if report_wandb and arguments.benchmark != eval_wandb_report.SELFPLAY_BENCHMARK_NAME: + raise SystemExit( + f"wandb=True reports only the {eval_wandb_report.SELFPLAY_BENCHMARK_NAME} benchmark's headline metric; " + f"remove it for {arguments.benchmark}" + ) run_dir = evaluation_utils.resolve_run_dir(model_path) output_dir_name = override_value(overrides, "eval.output_dir_name") or "eval" @@ -134,17 +167,22 @@ def main(): f"eval.output_subdir={shard_subdir}", ] if use_srun: - # Full-node CPU request forces SLURM to place each shard step on its own node. + # --exact keeps the step to its CPU share (default: the whole node); --mem=0 shares the job's memory between steps command = [ "srun", "--nodes=1", "--ntasks=1", + "--exact", f"--cpus-per-task={cpus_per_shard}", + "--mem=0", "--gres=gpu:1", *command, ] shard_env = dict(os.environ) - placement = "own node via srun" + if arguments.shards_per_node > 1: + # the --exact CPU slice is the placement; a socket bind inside it could leave fewer CPUs than workers + shard_env.setdefault("PUFFER_NUMA_BIND", "0") + placement = f"srun step, 1 GPU, {cpus_per_shard} CPUs" else: shard_env = dict(os.environ, CUDA_VISIBLE_DEVICES=gpu_ids[shard_idx]) placement = f"GPU {gpu_ids[shard_idx]}" @@ -183,20 +221,16 @@ def main(): episode_summaries = pd.concat(shard_frames, ignore_index=True).to_dict("records") summary = evaluation_utils._write_eval_reports(episode_summaries, aggregate_dir, arguments.total_scenarios) print(f"Aggregated {summary['num_episodes']} episodes from {len(processes)} shards into {aggregate_dir}") + if not arguments.keep_shards: + tidy_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "eval", "tidy_selfplay_eval.py") + if subprocess.run([sys.executable, tidy_script, aggregate_dir], check=False).returncode != 0: + print("tidy failed (results above are unaffected)") if report_wandb: - import pufferlib.pufferl as pufferl - - sys.argv = [sys.argv[0]] + [o for o in overrides if o.partition("=")[0] != "wandb"] - config_args = pufferl.load_config(arguments.env_name) - config_args["wandb"] = True - checkpoint_config_path = os.path.join(run_dir, "config.yaml") - wandb_run_identity = evaluation_utils.load_checkpoint_run_identity(checkpoint_config_path) - pufferl.report_eval_to_wandb( - config_args, - {arguments.benchmark: {"summary": summary}}, - wandb_run_identity, - output_dir_name, + eval_wandb_report.report_headline_metric( + eval_wandb_report.EVAL_TYPE_SELFPLAY, + os.path.join(aggregate_dir, eval_wandb_report.SELFPLAY_SUMMARY_FILE_NAME), + run_dir, ) diff --git a/tests/drive/Makefile b/tests/drive/Makefile index cf7b5479bc..3d3d01cbf8 100644 --- a/tests/drive/Makefile +++ b/tests/drive/Makefile @@ -19,7 +19,8 @@ UNIT_SRCS := \ test_drive_map_cache.c \ test_drive_idm.c \ test_drive_speed_limits.c \ - test_drive_cosim_setters.c + test_drive_cosim_setters.c \ + test_drive_pedestrians.c SMOKE_SRCS := test_drive_env_smoke.c PERF_SRCS := test_perf.c diff --git a/tests/drive/include/drive_fixture.h b/tests/drive/include/drive_fixture.h index 7aa6cc6393..229e03524c 100644 --- a/tests/drive/include/drive_fixture.h +++ b/tests/drive/include/drive_fixture.h @@ -38,6 +38,7 @@ static inline Agent drive_test_agent(float x, float y, float heading) { agent.sim_valid = 1; agent.current_lane_idx = 1; agent.previous_lane_idx = 1; + agent.red_light_latched_idx = -1; agent.wheelbase = 2.7f; copy_pose_to_prev(&agent); update_agent_radius(&agent); @@ -91,6 +92,9 @@ static inline Drive drive_test_env_config( env.obs_partner_relative_velocity = 0; env.obs_slots_traffic_controls_n = 4; env.traffic_lights_enabled = 1; + env.traffic_light_green_min_s = 3.0f; + env.traffic_light_yellow_min_s = 2.0f; + env.traffic_light_yellow_max_s = 3.0f; env.stop_signs_enabled = 0; env.yield_signs_enabled = 0; env.dt = 0.1f; @@ -123,6 +127,8 @@ static inline Drive drive_test_env_config( env.obs_norm_veh_width_m = 10.0f; env.obs_norm_road_seg_length_m = 5.0f; env.obs_norm_road_seg_width_m = 5.0f; + env.obs_lane_width_clip_min_m = 0.0f; + env.obs_lane_width_clip_max_m = 0.0f; env.obs_norm_z_m = 10.0f; env.obs_range_traffic_control_m = 100.0f; env.obs_range_partner_m = 100.0f; diff --git a/tests/drive/test_drive_cosim_setters.c b/tests/drive/test_drive_cosim_setters.c index f8b97415f6..ef0ce16861 100644 --- a/tests/drive/test_drive_cosim_setters.c +++ b/tests/drive/test_drive_cosim_setters.c @@ -81,7 +81,7 @@ static int test_set_traffic_light_states_writes_current_and_next_timestep_for_li env.timestep = 2; int new_states[2] = {3, 7}; // one entry per traffic element, indexed like env->traffic_elements - EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states), 0); + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, NULL), 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 @@ -108,26 +108,109 @@ static int test_set_traffic_light_states_rejects_out_of_range_timestep_and_state int new_states[2] = {1, 1}; env.timestep = 10; // out of range for state_size=3 - EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states), -1); + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, NULL), -1); EXPECT_EQ_INT(light_states[0], 99); EXPECT_EQ_INT(light_states[1], 99); EXPECT_EQ_INT(light_states[2], 99); env.timestep = 1; int bad_states[2] = {1, TRAFFIC_CONTROL_STATE_OFF + 1}; - EXPECT_EQ_INT(c_set_traffic_light_states(&env, bad_states), -1); + EXPECT_EQ_INT(c_set_traffic_light_states(&env, bad_states, NULL), -1); EXPECT_EQ_INT(light_states[1], 99); - EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states), 0); + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, NULL), 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(c_set_traffic_light_states(&env, last_states, NULL), 0); EXPECT_EQ_INT(light_states[2], 3); return 0; } +#define TIMING_STATE_COUNT 12 + +static Drive make_timing_env(TrafficControlElement *elements, int element_count, int (*storage)[TIMING_STATE_COUNT]) { + for (int i = 0; i < element_count; i++) { + for (int t = 0; t < TIMING_STATE_COUNT; t++) { + storage[i][t] = 99; + } + int type = i == element_count - 1 ? TRAFFIC_CONTROL_TYPE_STOP_SIGN : TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT; + elements[i] = make_traffic_element(type, TIMING_STATE_COUNT, storage[i]); + } + Drive env = {0}; + env.traffic_elements = elements; + env.num_traffic_elements = element_count; + env.dt = 0.1f; + env.traffic_light_yellow_max_s = 3.0f; + return env; +} + +static int test_set_traffic_light_states_yellow_timing_patches_schedule(void) { + // A yellow light with r seconds left stays yellow through ts + ceil(r / dt) - 1 and turns red right after; + // the held step ts + 1 is always yellow, and non-yellow lights or non-light elements get no patch. + int storage[3][TIMING_STATE_COUNT]; + TrafficControlElement elements[3]; + Drive env = make_timing_env(elements, 3, storage); + int *yellow_light = storage[0], *green_light = storage[1], *stop_sign = storage[2]; + + env.timestep = 2; + int new_states[3] = {TRAFFIC_CONTROL_STATE_YELLOW, TRAFFIC_CONTROL_STATE_GREEN, 7}; + float remaining_s[3] = {0.5f, 0.0f, 0.0f}; // 5 steps: yellow at 2..6, red at 7 + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, remaining_s), 0); + EXPECT_EQ_INT(yellow_light[1], 99); + for (int t = 2; t <= 6; t++) { + EXPECT_EQ_INT(yellow_light[t], TRAFFIC_CONTROL_STATE_YELLOW); + } + EXPECT_EQ_INT(yellow_light[7], TRAFFIC_CONTROL_STATE_RED); + EXPECT_EQ_INT(yellow_light[8], 99); + EXPECT_EQ_INT(green_light[2], TRAFFIC_CONTROL_STATE_GREEN); + EXPECT_EQ_INT(green_light[3], TRAFFIC_CONTROL_STATE_GREEN); + EXPECT_EQ_INT(green_light[4], 99); + for (int t = 0; t < TIMING_STATE_COUNT; t++) { + EXPECT_EQ_INT(stop_sign[t], 99); + } + + env.timestep = 0; // under one step left: the held step still counts, red right after it + float short_remaining_s[3] = {0.05f, 0.0f, 0.0f}; + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, short_remaining_s), 0); + EXPECT_EQ_INT(yellow_light[0], TRAFFIC_CONTROL_STATE_YELLOW); + EXPECT_EQ_INT(yellow_light[1], TRAFFIC_CONTROL_STATE_YELLOW); + EXPECT_EQ_INT(yellow_light[2], TRAFFIC_CONTROL_STATE_RED); + + env.timestep = 9; // the schedule end truncates the patch + float long_remaining_s[3] = {3.0f, 0.0f, 0.0f}; + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, long_remaining_s), 0); + for (int t = 9; t < TIMING_STATE_COUNT; t++) { + EXPECT_EQ_INT(yellow_light[t], TRAFFIC_CONTROL_STATE_YELLOW); + } + return 0; +} + +static int test_set_traffic_light_states_yellow_timing_validation(void) { + // The obs flag makes the timing mandatory; NaN, negative or above-the-training-max seconds fail the call. + int storage[2][TIMING_STATE_COUNT]; + TrafficControlElement elements[2]; + Drive env = make_timing_env(elements, 2, storage); + env.timestep = 1; + int new_states[2] = {TRAFFIC_CONTROL_STATE_YELLOW, 7}; + + env.obs_traffic_light_yellow_remaining = 1; + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, NULL), -1); + EXPECT_EQ_INT(storage[0][1], 99); + float at_max_s[2] = {3.0f, 0.0f}; + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, at_max_s), 0); + EXPECT_EQ_INT(storage[0][1], TRAFFIC_CONTROL_STATE_YELLOW); + + float nan_s[2] = {NAN, 0.0f}; + float negative_s[2] = {-0.1f, 0.0f}; + float above_max_s[2] = {3.5f, 0.0f}; + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, nan_s), -1); + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, negative_s), -1); + EXPECT_EQ_INT(c_set_traffic_light_states(&env, new_states, above_max_s), -1); + return 0; +} + // --------------------------------------------------------------------------- // c_set_agent_states // --------------------------------------------------------------------------- @@ -296,6 +379,261 @@ static int test_set_agent_goals_rejects_out_of_range_agent_idx(void) { return 0; } +static int test_set_agent_goal_progress_moves_the_current_goal_and_rejects_bad_input(void) { + // Undoing a forecast rollout puts the consumed-goal count back: the current-goal alias follows it, a + // fully consumed window (idx == goal_count) is a legal state, and anything outside 0..goal_count or an + // unknown agent fails without touching the window. + Drive env = {0}; + Agent agent = {0}; + env.agents = &agent; + env.num_total_agents = 1; + float gx[3] = {10.0f, 30.0f, 50.0f}; + float gy[3] = {1.0f, 2.0f, 3.0f}; + float gz[3] = {0.0f, 0.0f, 0.0f}; + float gdir[3] = {0.0f, 0.0f, 0.0f}; + EXPECT_EQ_INT(c_set_agent_goals(&env, 0, 3, gx, gy, gz, gdir, gdir), 0); + + EXPECT_EQ_INT(c_set_agent_goal_progress(&env, 0, 2), 0); + EXPECT_EQ_INT(agent.current_goal_idx, 2); + EXPECT_NEAR(agent.current_goal_x, 50.0f, 1e-6f); + EXPECT_NEAR(agent.current_goal_y, 3.0f, 1e-6f); + EXPECT_EQ_INT(c_set_agent_goal_progress(&env, 0, 3), 0); + EXPECT_EQ_INT(agent.current_goal_idx, 3); + EXPECT_EQ_INT(c_set_agent_goal_progress(&env, 0, 4), -1); + EXPECT_EQ_INT(c_set_agent_goal_progress(&env, 0, -1), -1); + EXPECT_EQ_INT(c_set_agent_goal_progress(&env, 1, 0), -1); + EXPECT_EQ_INT(agent.current_goal_idx, 3); + EXPECT_EQ_INT(c_set_agent_goal_progress(&env, 0, 0), 0); + EXPECT_NEAR(agent.current_goal_x, 10.0f, 1e-6f); + return 0; +} + +// --------------------------------------------------------------------------- +// c_set_agent_lane_goals +// --------------------------------------------------------------------------- + +#define LANE_GOAL_TEST_MAX_POINTS 12 + +typedef struct { + float x[LANE_GOAL_TEST_MAX_POINTS]; + float y[LANE_GOAL_TEST_MAX_POINTS]; + float z[LANE_GOAL_TEST_MAX_POINTS]; + float headings[LANE_GOAL_TEST_MAX_POINTS]; + float cum_lengths[LANE_GOAL_TEST_MAX_POINTS]; + int exit_lanes[2]; +} LaneGoalTestStorage; + +static RoadMapElement make_test_lane(LaneGoalTestStorage *storage, int point_count, int num_exits) { + RoadMapElement lane = {0}; + lane.type = LANE_SURFACE_STREET; + lane.segment_size = point_count; + lane.x = storage->x; + lane.y = storage->y; + lane.z = storage->z; + lane.headings = storage->headings; + lane.cum_lengths = storage->cum_lengths; + lane.exit_lanes = storage->exit_lanes; + lane.num_exits = num_exits; + storage->cum_lengths[0] = 0.0f; + for (int point_idx = 0; point_idx + 1 < point_count; point_idx++) { + float dx = storage->x[point_idx + 1] - storage->x[point_idx]; + float dy = storage->y[point_idx + 1] - storage->y[point_idx]; + storage->headings[point_idx] = atan2f(dy, dx); + storage->cum_lengths[point_idx + 1] = storage->cum_lengths[point_idx] + sqrtf(dx * dx + dy * dy); + } + storage->headings[point_count - 1] = storage->headings[point_count - 2]; + lane.length = storage->cum_lengths[point_count - 1]; + return lane; +} + +// Lane 0 runs 30 m along +x and forks into lane 1 (left turn, listed first) and lane 2 (100 m straight on). +static void make_fork_map(RoadMapElement *lanes, LaneGoalTestStorage *storage) { + for (int point_idx = 0; point_idx < 4; point_idx++) { + storage[0].x[point_idx] = 10.0f * (float) point_idx; + } + storage[0].exit_lanes[0] = 1; + storage[0].exit_lanes[1] = 2; + lanes[0] = make_test_lane(&storage[0], 4, 2); + + const float turn_x[4] = {30.0f, 36.0f, 39.0f, 40.0f}; + const float turn_y[4] = {0.0f, 1.0f, 4.0f, 10.0f}; + for (int point_idx = 0; point_idx < 4; point_idx++) { + storage[1].x[point_idx] = turn_x[point_idx]; + storage[1].y[point_idx] = turn_y[point_idx]; + } + lanes[1] = make_test_lane(&storage[1], 4, 0); + + for (int point_idx = 0; point_idx < 11; point_idx++) { + storage[2].x[point_idx] = 30.0f + 10.0f * (float) point_idx; + } + lanes[2] = make_test_lane(&storage[2], 11, 0); +} + +static int test_set_agent_lane_goals_follows_the_straightest_exit(void) { + // A forecast agent has no route: its goals run along its own lane and, at a fork, the exit that turns + // least -- not the first exit listed, and without drawing from the env RNG (the forecast must not + // change the random stream of anything else). + RoadMapElement lanes[3]; + LaneGoalTestStorage storage[3] = {0}; + make_fork_map(lanes, storage); + Drive env = {0}; + env.road_elements = lanes; + env.num_road_elements = 3; + env.num_goals = 3; + Agent agent = drive_test_agent(5.0f, 0.0f, 0.0f); + agent.current_lane_idx = 0; + agent.goal_count = 99; + env.agents = &agent; + env.num_total_agents = 1; + Rng rng_before = env.rng_state; + + int idx[1] = {0}; + int placed[1] = {-1}; + EXPECT_EQ_INT(c_set_agent_lane_goals(&env, 1, idx, 20.0f, placed), 0); + + EXPECT_EQ_INT(placed[0], 3); + EXPECT_EQ_INT(agent.goal_count, 3); + EXPECT_EQ_INT(agent.current_goal_idx, 0); + EXPECT_NEAR(agent.lane_goal_spacing_m, 20.0f, 1e-6f); + const float expected_x[3] = {25.0f, 45.0f, 65.0f}; + const int expected_lane[3] = {0, 2, 2}; + for (int goal_idx = 0; goal_idx < 3; goal_idx++) { + EXPECT_NEAR(agent.list_goal_x[goal_idx], expected_x[goal_idx], 1e-4f); + EXPECT_NEAR(agent.list_goal_y[goal_idx], 0.0f, 1e-4f); + EXPECT_EQ_INT(agent.list_goal_lane[goal_idx], expected_lane[goal_idx]); + } + EXPECT_NEAR(agent.current_goal_x, 25.0f, 1e-4f); + EXPECT_TRUE(memcmp(&rng_before, &env.rng_state, sizeof(Rng)) == 0); + return 0; +} + +static int test_set_agent_lane_goals_stops_where_the_lane_graph_ends(void) { + // Goals are never invented past a dead end: the call reports how many it placed, 0 for an agent with no + // lane under it or with less than one spacing of lane left, and a stale window is cleared either way. + RoadMapElement lanes[3]; + LaneGoalTestStorage storage[3] = {0}; + make_fork_map(lanes, storage); + Drive env = {0}; + env.road_elements = lanes; + env.num_road_elements = 3; + env.num_goals = 3; + Agent agents[3]; + agents[0] = drive_test_agent(85.0f, 0.0f, 0.0f); + agents[0].current_lane_idx = 2; + agents[1] = drive_test_agent(120.0f, 0.0f, 0.0f); + agents[1].current_lane_idx = 2; + agents[1].goal_count = 3; + agents[2] = drive_test_agent(500.0f, 500.0f, 0.0f); + agents[2].current_lane_idx = -1; + agents[2].goal_count = 3; + env.agents = agents; + env.num_total_agents = 3; + + int idx[3] = {0, 1, 2}; + int placed[3] = {-1, -1, -1}; + EXPECT_EQ_INT(c_set_agent_lane_goals(&env, 3, idx, 20.0f, placed), 0); + + EXPECT_EQ_INT(placed[0], 2); + EXPECT_EQ_INT(agents[0].goal_count, 2); + EXPECT_NEAR(agents[0].list_goal_x[0], 105.0f, 1e-4f); + EXPECT_NEAR(agents[0].list_goal_x[1], 125.0f, 1e-4f); + EXPECT_EQ_INT(agents[0].list_goal_lane[2], -1); + EXPECT_EQ_INT(placed[1], 0); + EXPECT_EQ_INT(agents[1].goal_count, 0); + EXPECT_EQ_INT(placed[2], 0); + EXPECT_EQ_INT(agents[2].goal_count, 0); + return 0; +} + +static int test_set_agent_lane_goals_rejects_bad_input_untouched(void) { + // External input: an index out of range, an agent without a pose or a spacing that is not a positive + // finite number fails the whole call before any agent is touched. + RoadMapElement lanes[3]; + LaneGoalTestStorage storage[3] = {0}; + make_fork_map(lanes, storage); + Drive env = {0}; + env.road_elements = lanes; + env.num_road_elements = 3; + env.num_goals = 3; + Agent agents[2]; + agents[0] = drive_test_agent(5.0f, 0.0f, 0.0f); + agents[0].current_lane_idx = 0; + agents[0].goal_count = 7; + agents[1] = drive_test_agent(5.0f, 0.0f, 0.0f); + agents[1].sim_valid = 0; + env.agents = agents; + env.num_total_agents = 2; + + int placed[2] = {-1, -1}; + int out_of_range[2] = {0, 2}; + EXPECT_EQ_INT(c_set_agent_lane_goals(&env, 2, out_of_range, 20.0f, placed), -1); + int without_pose[2] = {0, 1}; + EXPECT_EQ_INT(c_set_agent_lane_goals(&env, 2, without_pose, 20.0f, placed), -1); + int idx[1] = {0}; + EXPECT_EQ_INT(c_set_agent_lane_goals(&env, 1, idx, 0.0f, placed), -1); + EXPECT_EQ_INT(c_set_agent_lane_goals(&env, 1, idx, -5.0f, placed), -1); + EXPECT_EQ_INT(c_set_agent_lane_goals(&env, 1, idx, NAN, placed), -1); + EXPECT_EQ_INT(c_set_agent_lane_goals(&env, 1, idx, INFINITY, placed), -1); + EXPECT_EQ_INT(agents[0].goal_count, 7); + EXPECT_NEAR(agents[0].lane_goal_spacing_m, 0.0f, 1e-6f); + EXPECT_EQ_INT(placed[0], -1); + return 0; +} + +static int test_set_agent_goals_ends_lane_following(void) { + // An externally pushed route window takes the agent out of lane-following, so c_step stops renewing it. + Drive env = {0}; + Agent agent = {0}; + agent.lane_goal_spacing_m = 20.0f; + env.agents = &agent; + env.num_total_agents = 1; + + float gx[1] = {1.0f}, gy[1] = {1.0f}, gz[1] = {0.0f}, gdir[1] = {0.0f}; + EXPECT_EQ_INT(c_set_agent_goals(&env, 0, 1, gx, gy, gz, gdir, gdir), 0); + EXPECT_NEAR(agent.lane_goal_spacing_m, 0.0f, 1e-6f); + return 0; +} + +static int test_step_renews_lane_goals_when_one_is_reached(void) { + // With external goals c_step never regenerates a window itself, except for lane-following agents: when + // one reaches a goal, a full window ahead of its new pose replaces the old one, and the observation + // returned by that step already shows it. + Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); + env.goal_source = GOAL_SOURCE_EXTERNAL; + allocate(&env); + c_reset(&env); + Agent *agent = &env.agents[env.active_agent_indices[0]]; + EXPECT_EQ_INT(agent->goal_count, 0); + + int idx[1] = {env.active_agent_indices[0]}; + int placed[1] = {0}; + EXPECT_EQ_INT(c_set_agent_lane_goals(&env, 1, idx, 20.0f, placed), 0); + EXPECT_EQ_INT(placed[0], env.num_goals); + float first_goal_x = agent->list_goal_x[0]; + float first_goal_y = agent->list_goal_y[0]; + float first_goal_m = sqrtf(powf(first_goal_x - agent->sim_x, 2.0f) + powf(first_goal_y - agent->sim_y, 2.0f)); + EXPECT_TRUE(first_goal_m > 15.0f && first_goal_m < 20.5f); + + agent->current_goal_x = agent->list_goal_x[0] = agent->sim_x; + agent->current_goal_y = agent->list_goal_y[0] = agent->sim_y; + drive_set_neutral_actions(&env); + c_step(&env); + + EXPECT_EQ_INT(agent->current_goal_idx, 0); + EXPECT_EQ_INT(agent->goal_count, env.num_goals); + EXPECT_NEAR(agent->list_goal_x[0], first_goal_x, 1e-3f); + EXPECT_NEAR(agent->list_goal_y[0], first_goal_y, 1e-3f); + float goal_dx = agent->current_goal_x - agent->sim_x; + float goal_dy = agent->current_goal_y - agent->sim_y; + float goal_norm = fmaxf(env.obs_norm_goal_offset_m, sqrtf(goal_dx * goal_dx + goal_dy * goal_dy)); + EXPECT_NEAR( + env.observations[EGO_FEATURES], + (goal_dx * agent->cos_heading + goal_dy * agent->sin_heading) / goal_norm, + 1e-6f); + free_allocated(&env); + return 0; +} + // --------------------------------------------------------------------------- // c_set_stop_signs // --------------------------------------------------------------------------- @@ -414,6 +752,8 @@ int main(void) { RUN_TEST(test_set_agent_sizes_rejects_out_of_range_index_and_bad_size); 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_traffic_light_states_yellow_timing_patches_schedule); + RUN_TEST(test_set_traffic_light_states_yellow_timing_validation); 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); @@ -423,5 +763,11 @@ int main(void) { 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); + RUN_TEST(test_set_agent_goal_progress_moves_the_current_goal_and_rejects_bad_input); + RUN_TEST(test_set_agent_lane_goals_follows_the_straightest_exit); + RUN_TEST(test_set_agent_lane_goals_stops_where_the_lane_graph_ends); + RUN_TEST(test_set_agent_lane_goals_rejects_bad_input_untouched); + RUN_TEST(test_set_agent_goals_ends_lane_following); + RUN_TEST(test_step_renews_lane_goals_when_one_is_reached); return test_summary(failures); } diff --git a/tests/drive/test_drive_env_smoke.c b/tests/drive/test_drive_env_smoke.c index e90b5405c4..4f16de664d 100644 --- a/tests/drive/test_drive_env_smoke.c +++ b/tests/drive/test_drive_env_smoke.c @@ -38,6 +38,59 @@ static int test_carla_gigaflow_load_step_log(void) { return run_case("carla-gigaflow", drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 32); } +static int test_failed_init_spawn_keeps_policy_slot(void) { + // more agents than Town01 has room for: failed spawns stay addressable, removed and masked + int requested_agent_count = 3000; + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, requested_agent_count, 0); + EXPECT_EQ_INT(env.active_agent_count, requested_agent_count); + int removed_agent_count = 0; + for (int i = 0; i < env.active_agent_count; i++) { + Agent *agent = &env.agents[env.active_agent_indices[i]]; + removed_agent_count += agent->removed; + EXPECT_EQ_INT(agent->sim_valid, !agent->removed); + EXPECT_EQ_INT(env.masks[i], !agent->removed); + } + EXPECT_TRUE(removed_agent_count > 0); + free_allocated(&env); + return 0; +} + +static int test_ground_truth_trajectories_fit_the_buffer_width(void) { + Drive env = drive_test_make_env(drive_nuplan_map(), SIMULATION_MODE_REPLAY, 1, 0); + int agent_count = env.active_agent_count; + Agent *agent = &env.agents[env.active_agent_indices[0]]; + int buffer_widths[2] = {agent->trajectory_size / 2, agent->trajectory_size + 7}; + for (int w = 0; w < 2; w++) { + int num_timesteps = buffer_widths[w]; + int cell_count = agent_count * num_timesteps; + float *xyz_heading = (float *) calloc(4 * cell_count, sizeof(float)); + int *valid = (int *) malloc(cell_count * sizeof(int)); + int *ids = (int *) calloc(2 * agent_count, sizeof(int)); + for (int k = 0; k < cell_count; k++) { + valid[k] = -1; + } + c_get_global_ground_truth_trajectories( + &env, + num_timesteps, + xyz_heading, + xyz_heading + cell_count, + xyz_heading + 2 * cell_count, + xyz_heading + 3 * cell_count, + valid, + ids, + ids + agent_count); + for (int t = 0; t < num_timesteps; t++) { + int expected_valid = t < agent->trajectory_size ? agent->log_valid[t] : 0; + EXPECT_EQ_INT(valid[t], expected_valid); + } + free(xyz_heading); + free(valid); + free(ids); + } + free_allocated(&env); + return 0; +} + static int test_nuplan_gigaflow_load_step_log(void) { return run_case("nuplan-gigaflow", drive_nuplan_map(), SIMULATION_MODE_GIGAFLOW, 32); } @@ -256,5 +309,7 @@ int main(void) { RUN_TEST(test_short_early_reset_flags_and_logs); RUN_TEST(test_stagger_first_episode); RUN_TEST(test_eval_spawn_keeps_edge_clearance); + RUN_TEST(test_failed_init_spawn_keeps_policy_slot); + RUN_TEST(test_ground_truth_trajectories_fit_the_buffer_width); return test_summary(failures); } diff --git a/tests/drive/test_drive_geometry.c b/tests/drive/test_drive_geometry.c index bb2ca4dbcd..f41d1ac889 100644 --- a/tests/drive/test_drive_geometry.c +++ b/tests/drive/test_drive_geometry.c @@ -182,92 +182,6 @@ static int test_swept_line_box_stationary_collinear_segment(void) { return 0; } -static int test_stop_line_red_light_gating(void) { - Drive env = {0}; - TrafficControlElement tc = {0}; - int red_state[1] = {TRAFFIC_CONTROL_STATE_RED}; - int state[1] = {TRAFFIC_CONTROL_STATE_RED}; - int controlled_lane[1] = {1}; - Agent agent = drive_test_agent(1.0f, 0.0f, 0.0f); - - agent.prev_x = -3.0f; - tc.type = TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT; - tc.state_size = 1; - tc.states = state; - tc.stop_line[0] = 0.0f; - tc.stop_line[1] = -1.0f; - tc.stop_line[3] = 0.0f; - tc.stop_line[4] = 1.0f; - tc.heading = 0.0f; - tc.num_controlled_lanes = 1; - tc.controlled_lanes = controlled_lane; - env.num_traffic_elements = 1; - env.traffic_elements = &tc; - env.timestep = 0; - - EXPECT_TRUE(check_agent_on_stop_line(&env, &agent, false)); - - state[0] = TRAFFIC_CONTROL_STATE_GREEN; - EXPECT_FALSE(check_agent_on_stop_line(&env, &agent, false)); - state[0] = TRAFFIC_CONTROL_STATE_YELLOW; - EXPECT_FALSE(check_agent_on_stop_line(&env, &agent, false)); - EXPECT_TRUE(check_agent_on_stop_line(&env, &agent, true)); // yellow counts when spawn-mode flag is set - state[0] = TRAFFIC_CONTROL_STATE_OFF; - EXPECT_FALSE(check_agent_on_stop_line(&env, &agent, false)); - state[0] = TRAFFIC_CONTROL_STATE_RED; - - env.timestep = 1; - EXPECT_FALSE(check_agent_on_stop_line(&env, &agent, false)); - env.timestep = 0; - - controlled_lane[0] = 2; - EXPECT_FALSE(check_agent_on_stop_line(&env, &agent, false)); - controlled_lane[0] = 1; - - agent = drive_test_agent(1.0f, 0.0f, (float) M_PI); - agent.prev_x = -3.0f; - EXPECT_FALSE(check_agent_on_stop_line(&env, &agent, false)); - - agent = drive_test_agent(11.0f, 0.0f, 0.0f); - agent.prev_x = -1.0f; - EXPECT_FALSE(check_agent_on_stop_line(&env, &agent, false)); - - agent = drive_test_agent(1.0f, 1.4f, 0.0f); - agent.prev_x = -3.0f; - agent.prev_y = 1.4f; - agent.sim_width = 0.4f; - update_agent_radius(&agent); - tc.states = red_state; - EXPECT_TRUE(check_agent_on_stop_line(&env, &agent, false)); - - return 0; -} - -static int test_stop_line_stationary_on_red(void) { - Drive env = {0}; - TrafficControlElement tc = {0}; - int state[1] = {TRAFFIC_CONTROL_STATE_RED}; - int controlled_lane[1] = {1}; - Agent agent = drive_test_agent(0.0f, 0.0f, 0.0f); - - tc.type = TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT; - tc.state_size = 1; - tc.states = state; - tc.stop_line[0] = 0.0f; - tc.stop_line[1] = -1.0f; - tc.stop_line[3] = 0.0f; - tc.stop_line[4] = 1.0f; - tc.heading = 0.0f; - tc.num_controlled_lanes = 1; - tc.controlled_lanes = controlled_lane; - env.num_traffic_elements = 1; - env.traffic_elements = &tc; - env.timestep = 0; - - EXPECT_TRUE(check_agent_on_stop_line(&env, &agent, false)); - return 0; -} - int main(void) { int failures = 0; RUN_TEST(test_heading_normalization); @@ -280,7 +194,5 @@ int main(void) { RUN_TEST(test_segment_aabb_cases); RUN_TEST(test_swept_line_box_cases); RUN_TEST(test_swept_line_box_stationary_collinear_segment); - RUN_TEST(test_stop_line_red_light_gating); - RUN_TEST(test_stop_line_stationary_on_red); return test_summary(failures); } diff --git a/tests/drive/test_drive_pedestrians.c b/tests/drive/test_drive_pedestrians.c new file mode 100644 index 0000000000..b3c8042665 --- /dev/null +++ b/tests/drive/test_drive_pedestrians.c @@ -0,0 +1,471 @@ +#include "include/drive_fixture.h" +#include "include/test.h" + +static Drive make_pedestrian_env(int num_agents, int pedestrian_slots, float spawn_initial_speed) { + Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, num_agents, 0); + env.scripted_pedestrian_slots = pedestrian_slots; + env.spawn_initial_speed = spawn_initial_speed; + env.scenario_length = 400; + allocate(&env); + c_reset(&env); + return env; +} + +static int first_visible_pedestrian(Drive *env) { + for (int k = 0; k < env->scripted_pedestrian_slots; k++) { + if (env->pedestrians[k].state != PEDESTRIAN_STATE_IDLE) { + return k; + } + } + return -1; +} + +// Every slot appears in `mode` as soon as a car takes it +static int step_until_appearance(Drive *env, int mode, int max_steps) { + for (int k = 0; k < env->scripted_pedestrian_slots; k++) { + env->pedestrians[k].timer_s = 0.0f; + env->pedestrians[k].mode = mode; + } + for (int t = 0; t < max_steps; t++) { + drive_set_neutral_actions(env); + c_step(env); + int visible_k = first_visible_pedestrian(env); + if (visible_k != -1) { + return visible_k; + } + } + return -1; +} + +// Road edges between a point and the nearest lane vertex: an odd count means the point lies outside the road +static int road_edges_to_nearest_lane(Drive *env, float px, float py) { + float nearest_sq = 1e18f; + float lane_x = px; + float lane_y = py; + for (int i = 0; i < env->num_road_elements; i++) { + RoadMapElement *element = &env->road_elements[i]; + if (!is_road_lane(element->type)) { + continue; + } + for (int v = 0; v < element->segment_size; v++) { + float dist_sq = (element->x[v] - px) * (element->x[v] - px) + (element->y[v] - py) * (element->y[v] - py); + if (dist_sq < nearest_sq) { + nearest_sq = dist_sq; + lane_x = element->x[v]; + lane_y = element->y[v]; + } + } + } + float reach_m = sqrtf(nearest_sq); + float dir_x = (lane_x - px) / reach_m; + float dir_y = (lane_y - py) / reach_m; + int edge_count = 0; + for (int i = 0; i < env->num_road_elements; i++) { + RoadMapElement *element = &env->road_elements[i]; + if (!is_road_edge(element->type)) { + continue; + } + for (int v = 0; v + 1 < element->segment_size; v++) { + float t; + bool hit = ray_hits_segment( + px, + py, + dir_x, + dir_y, + element->x[v], + element->y[v], + element->x[v + 1], + element->y[v + 1], + &t); + edge_count += hit && t < reach_m; + } + } + return edge_count; +} + +static int target_slot_of(Drive *env, const ScriptedPedestrian *pedestrian) { + for (int j = 0; j < env->active_agent_count; j++) { + if (env->active_agent_indices[j] == pedestrian->target_agent_idx) { + return j; + } + } + return -1; +} + +static int test_ray_hits_segment(void) { + float t = -1.0f; + EXPECT_TRUE(ray_hits_segment(0.0f, 0.0f, 1.0f, 0.0f, 5.0f, -1.0f, 5.0f, 1.0f, &t)); + EXPECT_NEAR(t, 5.0f, 1e-5f); + EXPECT_FALSE(ray_hits_segment(0.0f, 0.0f, -1.0f, 0.0f, 5.0f, -1.0f, 5.0f, 1.0f, &t)); + EXPECT_FALSE(ray_hits_segment(0.0f, 0.0f, 1.0f, 0.0f, 5.0f, 1.0f, 5.0f, 3.0f, &t)); + EXPECT_FALSE(ray_hits_segment(0.0f, 0.0f, 1.0f, 0.0f, 2.0f, 0.0f, 8.0f, 0.0f, &t)); + EXPECT_TRUE(ray_hits_segment(1.0f, 1.0f, 0.0f, 1.0f, -3.0f, 4.0f, 3.0f, 4.0f, &t)); + EXPECT_NEAR(t, 3.0f, 1e-5f); + return 0; +} + +static int test_pedestrian_slots_are_hidden_static_partners(void) { + srand(7); + Drive env = make_pedestrian_env(16, 3, 0.0f); + EXPECT_EQ_INT(env.static_agent_count, 3); + EXPECT_EQ_INT(env.num_total_agents, 16 + 3); + EXPECT_EQ_INT(env.num_agents, env.active_agent_count + 3); + for (int k = 0; k < 3; k++) { + ScriptedPedestrian *pedestrian = &env.pedestrians[k]; + Agent *agent = &env.agents[pedestrian->agent_idx]; + EXPECT_EQ_INT(env.static_agent_indices[k], pedestrian->agent_idx); + EXPECT_EQ_INT(agent->type, PEDESTRIAN); + EXPECT_EQ_INT(agent->sim_valid, 0); + EXPECT_EQ_INT(agent->active_agent, 0); + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_IDLE); + EXPECT_EQ_INT(pedestrian->target_agent_idx, -1); + EXPECT_TRUE(pedestrian->timer_s >= PEDESTRIAN_COOLDOWN_MIN_S); + EXPECT_TRUE(pedestrian->timer_s <= PEDESTRIAN_COOLDOWN_MAX_S); + EXPECT_TRUE(pedestrian->mode >= PEDESTRIAN_MODE_POP_IN_CROSS && pedestrian->mode <= PEDESTRIAN_MODE_KERB_STEP); + } + free_allocated(&env); + return 0; +} + +static int test_pedestrian_slots_keep_initial_spawns(void) { + srand(7); + Drive without = make_pedestrian_env(16, 0, 0.0f); + srand(7); + Drive with = make_pedestrian_env(16, 2, 0.0f); + EXPECT_EQ_INT(with.active_agent_count, without.active_agent_count); + for (int i = 0; i < without.active_agent_count; i++) { + Agent *a = &without.agents[without.active_agent_indices[i]]; + Agent *b = &with.agents[with.active_agent_indices[i]]; + EXPECT_NEAR(b->sim_x, a->sim_x, 0.0f); + EXPECT_NEAR(b->sim_y, a->sim_y, 0.0f); + EXPECT_NEAR(b->sim_heading, a->sim_heading, 0.0f); + EXPECT_NEAR(b->sim_length, a->sim_length, 0.0f); + } + free_allocated(&without); + free_allocated(&with); + return 0; +} + +static int test_pop_in_crossing_geometry_and_lifecycle(void) { + srand(7); + Drive env = make_pedestrian_env(16, 2, 5.0f); + EXPECT_TRUE(env.active_agent_count >= 4); + int crossing_k = step_until_appearance(&env, PEDESTRIAN_MODE_POP_IN_CROSS, 60); + EXPECT_TRUE(crossing_k != -1); + ScriptedPedestrian *pedestrian = &env.pedestrians[crossing_k]; + Agent *agent = &env.agents[pedestrian->agent_idx]; + EXPECT_TRUE(pedestrian->target_agent_idx >= 0); + Agent *car = &env.agents[pedestrian->target_agent_idx]; + + EXPECT_EQ_INT(agent->sim_valid, 1); + EXPECT_EQ_INT(agent->type, PEDESTRIAN); + EXPECT_TRUE(agent->sim_length >= PEDESTRIAN_SIZE_MIN_M && agent->sim_length <= PEDESTRIAN_SIZE_MAX_M); + EXPECT_TRUE(agent->sim_width >= PEDESTRIAN_SIZE_MIN_M && agent->sim_width <= PEDESTRIAN_SIZE_MAX_M); + EXPECT_NEAR(agent->sim_height, PEDESTRIAN_HEIGHT_M, 0.0f); + float min_speed_mps + = PEDESTRIAN_LB1_SPEED_BASE_MPS + PEDESTRIAN_LB1_SPEED_PER_LANE_MPS - PEDESTRIAN_SPEED_JITTER_MPS; + float max_speed_mps + = PEDESTRIAN_LB1_SPEED_BASE_MPS + 3.0f * PEDESTRIAN_LB1_SPEED_PER_LANE_MPS + PEDESTRIAN_SPEED_JITTER_MPS; + EXPECT_TRUE(pedestrian->speed_mps >= min_speed_mps && pedestrian->speed_mps <= max_speed_mps); + EXPECT_NEAR(agent->sim_speed, 0.0f, 0.0f); + EXPECT_TRUE(pedestrian->timer_s >= PEDESTRIAN_WAIT_MIN_S && pedestrian->timer_s <= PEDESTRIAN_WAIT_MAX_S); + EXPECT_TRUE(pedestrian->ramp_s >= PEDESTRIAN_RAMP_MIN_S && pedestrian->ramp_s <= PEDESTRIAN_RAMP_MAX_S); + EXPECT_NEAR(pedestrian->dir_x * pedestrian->dir_x + pedestrian->dir_y * pedestrian->dir_y, 1.0f, 1e-4f); + EXPECT_NEAR(agent->cos_heading, pedestrian->dir_x, 1e-5f); + EXPECT_NEAR(agent->sin_heading, pedestrian->dir_y, 1e-5f); + float min_crossing_m = PEDESTRIAN_KERB_OFFSET_MIN_M + PEDESTRIAN_FAR_KERB_OVERSHOOT_M; + float max_crossing_m + = PEDESTRIAN_KERB_OFFSET_MAX_M + PEDESTRIAN_FAR_KERB_OVERSHOOT_M + 2.0f * PEDESTRIAN_KERB_SEARCH_MAX_M; + EXPECT_TRUE(pedestrian->path_m > min_crossing_m && pedestrian->path_m < max_crossing_m); + EXPECT_TRUE(pedestrian->remaining_m > 0.0f && pedestrian->remaining_m <= pedestrian->path_m); + + // Walks across the target car's path, starting ahead of it on its right + float along_car = car->cos_heading * pedestrian->dir_x + car->sin_heading * pedestrian->dir_y; + EXPECT_TRUE(fabsf(along_car) < 0.5f); + float rel_x, rel_y; + project_point_to_ego_frame(car, agent->sim_x, agent->sim_y, &rel_x, &rel_y); + float max_appear_m + = PEDESTRIAN_LB1_APPEAR_BASE_M + 2.0f * PEDESTRIAN_LB1_APPEAR_PER_LANE_M + PEDESTRIAN_APPEAR_JITTER_M; + EXPECT_TRUE(rel_x > 0.0f && rel_x < max_appear_m + 5.0f); + EXPECT_TRUE(rel_y < 0.0f); + + int target_slot = target_slot_of(&env, pedestrian); + EXPECT_TRUE(target_slot != -1); + EXPECT_NEAR(env.logs[target_slot].pedestrian_crossings, 1.0f, 0.0f); + EXPECT_NEAR(env.logs[target_slot].pedestrian_bystanders, 0.0f, 0.0f); + EXPECT_TRUE(pedestrian->stand_s >= PEDESTRIAN_LINGER_MIN_S && pedestrian->stand_s <= PEDESTRIAN_LINGER_MAX_S); + + int saw_standing = 0; + for (int t = 0; t < 250 && pedestrian->state != PEDESTRIAN_STATE_IDLE; t++) { + drive_set_neutral_actions(&env); + c_step(&env); + saw_standing |= pedestrian->state == PEDESTRIAN_STATE_STANDING; + } + EXPECT_TRUE(saw_standing); + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_IDLE); + EXPECT_EQ_INT(agent->sim_valid, 0); + EXPECT_EQ_INT(pedestrian->target_agent_idx, -1); + EXPECT_TRUE(pedestrian->timer_s >= PEDESTRIAN_COOLDOWN_MIN_S); + free_allocated(&env); + return 0; +} + +static int test_car_hitting_pedestrian_logs_pedestrian_collision(void) { + srand(7); + Drive env = make_pedestrian_env(16, 2, 5.0f); + int crossing_k = step_until_appearance(&env, PEDESTRIAN_MODE_POP_IN_CROSS, 60); + EXPECT_TRUE(crossing_k != -1); + ScriptedPedestrian *pedestrian = &env.pedestrians[crossing_k]; + Agent *agent = &env.agents[pedestrian->agent_idx]; + + // Put the walker mid-crossing and park a policy car on it, aligned with the road + float mid_offset_m = pedestrian->remaining_m - 0.5f * pedestrian->path_m; + float mid_x = agent->sim_x + pedestrian->dir_x * mid_offset_m; + float mid_y = agent->sim_y + pedestrian->dir_y * mid_offset_m; + agent->sim_x = mid_x; + agent->sim_y = mid_y; + copy_pose_to_prev(agent); + + int car_slot = 0; + int car_idx = env.active_agent_indices[car_slot]; + Agent *car = &env.agents[car_idx]; + car->sim_x = mid_x; + car->sim_y = mid_y; + car->sim_z = agent->sim_z; + car->sim_heading = atan2f(-pedestrian->dir_x, pedestrian->dir_y); + car->cos_heading = cosf(car->sim_heading); + car->sin_heading = sinf(car->sim_heading); + copy_pose_to_prev(car); + car->stopped = 0; + car->removed = 0; + car->first_collision_partner_idx = -1; + env.logs[car_slot] = (Log) {0}; + + compute_metrics(&env, car_idx, car_slot); + EXPECT_NEAR(car->metrics_array[OFFROAD_IDX], 0.0f, 0.0f); + EXPECT_NEAR(car->metrics_array[COLLISION_IDX], 1.0f, 0.0f); + EXPECT_EQ_INT(car->first_collision_partner_idx, pedestrian->agent_idx); + compute_rewards(&env, car_slot); + EXPECT_NEAR(env.logs[car_slot].collision_rate, 1.0f, 0.0f); + EXPECT_NEAR(env.logs[car_slot].pedestrian_collision_rate, 1.0f, 0.0f); + free_allocated(&env); + return 0; +} + +static int test_pop_in_crosser_stands_then_ramps_to_speed(void) { + srand(7); + Drive env = make_pedestrian_env(16, 2, 5.0f); + int crossing_k = step_until_appearance(&env, PEDESTRIAN_MODE_POP_IN_CROSS, 60); + EXPECT_TRUE(crossing_k != -1); + ScriptedPedestrian *pedestrian = &env.pedestrians[crossing_k]; + Agent *agent = &env.agents[pedestrian->agent_idx]; + float start_x = agent->sim_x; + float start_y = agent->sim_y; + int standing_steps = 0; + while (pedestrian->timer_s > 0.0f && standing_steps < 10) { + drive_set_neutral_actions(&env); + c_step(&env); + standing_steps++; + EXPECT_NEAR(agent->sim_x, start_x, 0.0f); + EXPECT_NEAR(agent->sim_y, start_y, 0.0f); + EXPECT_NEAR(agent->sim_speed, 0.0f, 0.0f); + } + EXPECT_TRUE(standing_steps <= (int) ceilf(PEDESTRIAN_WAIT_MAX_S / env.dt) + 1); + float previous_speed_mps = 0.0f; + int ramp_steps = (int) ceilf(pedestrian->ramp_s / env.dt); + for (int t = 0; t < ramp_steps; t++) { + drive_set_neutral_actions(&env); + c_step(&env); + EXPECT_TRUE(agent->sim_speed > previous_speed_mps || previous_speed_mps >= pedestrian->speed_mps - 1e-4f); + EXPECT_TRUE(agent->sim_speed <= pedestrian->speed_mps + 1e-4f); + previous_speed_mps = agent->sim_speed; + } + EXPECT_NEAR(agent->sim_speed, pedestrian->speed_mps, 1e-4f); + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_WALKING); + EXPECT_TRUE(hypotf(agent->sim_x - start_x, agent->sim_y - start_y) > 0.0f); + free_allocated(&env); + return 0; +} + +static int test_pedestrian_mode_shares(void) { + srand(7); + Drive env = make_pedestrian_env(16, 1, 0.0f); + const int draw_count = 20000; + int mode_counts[PEDESTRIAN_MODE_KERB_STEP + 1] = {0}; + for (int i = 0; i < draw_count; i++) { + mode_counts[sample_pedestrian_mode(&env)]++; + } + float kerb_step_prob = 1.0f - PEDESTRIAN_POP_IN_CROSS_PROB - PEDESTRIAN_SLOW_CROSS_PROB - PEDESTRIAN_STAND_PROB + - PEDESTRIAN_KERB_WALK_PROB; + EXPECT_TRUE(kerb_step_prob > 0.0f); + EXPECT_NEAR((float) mode_counts[PEDESTRIAN_MODE_POP_IN_CROSS] / draw_count, PEDESTRIAN_POP_IN_CROSS_PROB, 0.02f); + EXPECT_NEAR((float) mode_counts[PEDESTRIAN_MODE_SLOW_CROSS] / draw_count, PEDESTRIAN_SLOW_CROSS_PROB, 0.02f); + EXPECT_NEAR((float) mode_counts[PEDESTRIAN_MODE_STAND] / draw_count, PEDESTRIAN_STAND_PROB, 0.02f); + EXPECT_NEAR((float) mode_counts[PEDESTRIAN_MODE_KERB_WALK] / draw_count, PEDESTRIAN_KERB_WALK_PROB, 0.02f); + EXPECT_NEAR((float) mode_counts[PEDESTRIAN_MODE_KERB_STEP] / draw_count, kerb_step_prob, 0.02f); + free_allocated(&env); + return 0; +} + +static int test_slow_crosser_walks_from_a_kerb_across_the_road(void) { + srand(7); + Drive env = make_pedestrian_env(16, 1, 5.0f); + int visible_k = step_until_appearance(&env, PEDESTRIAN_MODE_SLOW_CROSS, 100); + EXPECT_TRUE(visible_k != -1); + ScriptedPedestrian *pedestrian = &env.pedestrians[visible_k]; + Agent *agent = &env.agents[pedestrian->agent_idx]; + Agent *car = &env.agents[pedestrian->target_agent_idx]; + + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_WALKING); + EXPECT_TRUE(pedestrian->speed_mps >= PEDESTRIAN_SLOW_SPEED_MIN_MPS); + EXPECT_TRUE(pedestrian->speed_mps <= PEDESTRIAN_SLOW_SPEED_MAX_MPS); + EXPECT_TRUE(pedestrian->timer_s >= 0.0f && pedestrian->timer_s <= PEDESTRIAN_SLOW_WAIT_MAX_S); + EXPECT_NEAR(pedestrian->dir_x * pedestrian->dir_x + pedestrian->dir_y * pedestrian->dir_y, 1.0f, 1e-4f); + EXPECT_TRUE(hypotf(agent->sim_x - car->sim_x, agent->sim_y - car->sim_y) > 0.5f * PEDESTRIAN_SLOW_APPEAR_MIN_M); + int target_slot = target_slot_of(&env, pedestrian); + EXPECT_NEAR(env.logs[target_slot].pedestrian_crossings, 1.0f, 0.0f); + + // Kerb, road, kerb + float end_x = agent->sim_x + pedestrian->dir_x * pedestrian->path_m; + float end_y = agent->sim_y + pedestrian->dir_y * pedestrian->path_m; + EXPECT_EQ_INT(road_edges_to_nearest_lane(&env, agent->sim_x, agent->sim_y) % 2, 1); + EXPECT_EQ_INT(road_edges_to_nearest_lane(&env, end_x, end_y) % 2, 1); + EXPECT_EQ_INT(road_edges_to_nearest_lane(&env, 0.5f * (agent->sim_x + end_x), 0.5f * (agent->sim_y + end_y)), 0); + + for (int t = 0; t < 350 && pedestrian->state == PEDESTRIAN_STATE_WALKING; t++) { + drive_set_neutral_actions(&env); + c_step(&env); + EXPECT_TRUE(agent->sim_speed <= pedestrian->speed_mps + 1e-4f); + } + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_STANDING); + EXPECT_NEAR(agent->sim_x, end_x, 1e-2f); + EXPECT_NEAR(agent->sim_y, end_y, 1e-2f); + free_allocated(&env); + return 0; +} + +static int test_stander_stays_beside_the_road_with_a_running_stopped_time(void) { + srand(7); + Drive env = make_pedestrian_env(16, 1, 5.0f); + int visible_k = step_until_appearance(&env, PEDESTRIAN_MODE_STAND, 100); + EXPECT_TRUE(visible_k != -1); + ScriptedPedestrian *pedestrian = &env.pedestrians[visible_k]; + Agent *agent = &env.agents[pedestrian->agent_idx]; + + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_STANDING); + EXPECT_TRUE(pedestrian->timer_s >= PEDESTRIAN_BYSTANDER_STAND_MIN_S); + EXPECT_TRUE(pedestrian->timer_s <= PEDESTRIAN_BYSTANDER_STAND_MAX_S); + EXPECT_TRUE(agent->seconds_stopped >= 0.0f && agent->seconds_stopped <= MAX_STOPPED_SECONDS + env.dt); + EXPECT_EQ_INT(road_edges_to_nearest_lane(&env, agent->sim_x, agent->sim_y) % 2, 1); + int target_slot = target_slot_of(&env, pedestrian); + EXPECT_NEAR(env.logs[target_slot].pedestrian_bystanders, 1.0f, 0.0f); + EXPECT_NEAR(env.logs[target_slot].pedestrian_crossings, 0.0f, 0.0f); + + float start_x = agent->sim_x; + float start_y = agent->sim_y; + float stopped_s = agent->seconds_stopped; + int standing_steps = (int) (PEDESTRIAN_BYSTANDER_STAND_MIN_S / env.dt) - 2; + for (int t = 0; t < standing_steps; t++) { + drive_set_neutral_actions(&env); + c_step(&env); + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_STANDING); + EXPECT_NEAR(agent->sim_x, start_x, 0.0f); + EXPECT_NEAR(agent->sim_y, start_y, 0.0f); + EXPECT_NEAR(agent->sim_speed, 0.0f, 0.0f); + } + EXPECT_NEAR(agent->seconds_stopped, stopped_s + (float) standing_steps * env.dt, 1e-2f); + free_allocated(&env); + return 0; +} + +static int test_kerb_walker_stays_off_the_road(void) { + srand(7); + Drive env = make_pedestrian_env(16, 1, 5.0f); + int visible_k = step_until_appearance(&env, PEDESTRIAN_MODE_KERB_WALK, 150); + EXPECT_TRUE(visible_k != -1); + ScriptedPedestrian *pedestrian = &env.pedestrians[visible_k]; + Agent *agent = &env.agents[pedestrian->agent_idx]; + + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_WALKING); + EXPECT_TRUE(pedestrian->speed_mps >= PEDESTRIAN_KERB_WALK_SPEED_MIN_MPS); + EXPECT_TRUE(pedestrian->speed_mps <= PEDESTRIAN_KERB_WALK_SPEED_MAX_MPS); + EXPECT_NEAR(agent->sim_speed, pedestrian->speed_mps, 1e-4f); + EXPECT_NEAR(agent->seconds_stopped, 0.0f, 0.0f); + EXPECT_TRUE(pedestrian->path_m > PEDESTRIAN_KERB_WALK_MIN_M - 2.0f * PEDESTRIAN_KERB_WALK_BEND_TOLERANCE_M); + EXPECT_TRUE(pedestrian->path_m < PEDESTRIAN_KERB_WALK_MAX_M + 2.0f * PEDESTRIAN_KERB_WALK_BEND_TOLERANCE_M); + int target_slot = target_slot_of(&env, pedestrian); + EXPECT_NEAR(env.logs[target_slot].pedestrian_bystanders, 1.0f, 0.0f); + + float start_x = agent->sim_x; + float start_y = agent->sim_y; + float start_z = agent->sim_z; + float path_m = pedestrian->path_m; + EXPECT_NEAR(pedestrian->climb_per_m, 0.0f, 0.0f); // Town01 is flat + pedestrian->climb_per_m = 0.05f; + for (int t = 0; t < 380 && pedestrian->state == PEDESTRIAN_STATE_WALKING; t++) { + EXPECT_EQ_INT(road_edges_to_nearest_lane(&env, agent->sim_x, agent->sim_y) % 2, 1); + drive_set_neutral_actions(&env); + c_step(&env); + } + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_STANDING); + EXPECT_NEAR(hypotf(agent->sim_x - start_x, agent->sim_y - start_y), path_m, 1e-2f); + EXPECT_NEAR(agent->sim_z - start_z, 0.05f * path_m, 1e-2f); + EXPECT_EQ_INT(road_edges_to_nearest_lane(&env, agent->sim_x, agent->sim_y) % 2, 1); + free_allocated(&env); + return 0; +} + +static int test_kerb_stepper_pops_in_steps_and_stands_off_the_road(void) { + srand(7); + Drive env = make_pedestrian_env(16, 1, 5.0f); + int visible_k = step_until_appearance(&env, PEDESTRIAN_MODE_KERB_STEP, 60); + EXPECT_TRUE(visible_k != -1); + ScriptedPedestrian *pedestrian = &env.pedestrians[visible_k]; + Agent *agent = &env.agents[pedestrian->agent_idx]; + Agent *car = &env.agents[pedestrian->target_agent_idx]; + + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_WALKING); + EXPECT_TRUE(pedestrian->timer_s >= PEDESTRIAN_WAIT_MIN_S && pedestrian->timer_s <= PEDESTRIAN_WAIT_MAX_S); + EXPECT_TRUE(pedestrian->path_m >= 0.0f && pedestrian->path_m <= PEDESTRIAN_KERB_STEP_MAX_M); + EXPECT_TRUE(pedestrian->stand_s >= PEDESTRIAN_BYSTANDER_STAND_MIN_S); + EXPECT_TRUE(pedestrian->stand_s <= PEDESTRIAN_BYSTANDER_STAND_MAX_S); + float rel_x, rel_y; + project_point_to_ego_frame(car, agent->sim_x, agent->sim_y, &rel_x, &rel_y); + float max_appear_m + = PEDESTRIAN_LB1_APPEAR_BASE_M + 2.0f * PEDESTRIAN_LB1_APPEAR_PER_LANE_M + PEDESTRIAN_APPEAR_JITTER_M; + EXPECT_TRUE(rel_x > 0.0f && rel_x < max_appear_m + 5.0f); + EXPECT_TRUE(rel_y < 0.0f); + int target_slot = target_slot_of(&env, pedestrian); + EXPECT_NEAR(env.logs[target_slot].pedestrian_bystanders, 1.0f, 0.0f); + EXPECT_NEAR(env.logs[target_slot].pedestrian_crossings, 0.0f, 0.0f); + + float start_x = agent->sim_x; + float start_y = agent->sim_y; + float path_m = pedestrian->path_m; + for (int t = 0; t < 40 && pedestrian->state == PEDESTRIAN_STATE_WALKING; t++) { + drive_set_neutral_actions(&env); + c_step(&env); + } + EXPECT_EQ_INT(pedestrian->state, PEDESTRIAN_STATE_STANDING); + EXPECT_NEAR(hypotf(agent->sim_x - start_x, agent->sim_y - start_y), path_m, 1e-3f); + EXPECT_EQ_INT(road_edges_to_nearest_lane(&env, agent->sim_x, agent->sim_y) % 2, 1); + EXPECT_NEAR(agent->sim_speed, 0.0f, 0.0f); + free_allocated(&env); + return 0; +} + +int main(void) { + int failures = 0; + RUN_TEST(test_ray_hits_segment); + RUN_TEST(test_pedestrian_slots_are_hidden_static_partners); + RUN_TEST(test_pedestrian_slots_keep_initial_spawns); + RUN_TEST(test_pedestrian_mode_shares); + RUN_TEST(test_pop_in_crossing_geometry_and_lifecycle); + RUN_TEST(test_pop_in_crosser_stands_then_ramps_to_speed); + RUN_TEST(test_car_hitting_pedestrian_logs_pedestrian_collision); + RUN_TEST(test_slow_crosser_walks_from_a_kerb_across_the_road); + RUN_TEST(test_stander_stays_beside_the_road_with_a_running_stopped_time); + RUN_TEST(test_kerb_walker_stays_off_the_road); + RUN_TEST(test_kerb_stepper_pops_in_steps_and_stands_off_the_road); + return test_summary(failures); +} diff --git a/tests/drive/test_drive_red_light.c b/tests/drive/test_drive_red_light.c index 7791179fd3..5e44eb05fc 100644 --- a/tests/drive/test_drive_red_light.c +++ b/tests/drive/test_drive_red_light.c @@ -4,6 +4,7 @@ // Stop line spans y in [-2, 2] at x = 0, traffic flows +x. Controlled lanes only // matter for the spawn overlap check; violation detection is purely geometric. #define LIGHT_STATE_COUNT 16 +#define TEST_HALF_LENGTH_M 2.0f // drive_test_agent is 4 m long: front edge = center + 2 m at heading 0 static int light_states[LIGHT_STATE_COUNT]; static int light_lanes[2] = {100, 101}; @@ -34,10 +35,11 @@ static Drive drive_test_light_env(TrafficControlElement *tc) { return env; } -// Violation = vehicle center crossing the line; stationary/backward motion never flags. -static Agent drive_test_driving_agent(float prev_x, float cur_x, float y) { - Agent agent = drive_test_agent(prev_x, y, 0.0f); - agent.sim_x = cur_x; +// Violation = on red the line lies between the tail clearance point (rear - 1 m) and the front edge, moving or +// not; arguments are front-edge x (rear = front - 4 m). +static Agent drive_test_driving_agent(float prev_front_x, float cur_front_x, float y) { + Agent agent = drive_test_agent(prev_front_x - TEST_HALF_LENGTH_M, y, 0.0f); + agent.sim_x = cur_front_x - TEST_HALF_LENGTH_M; agent.current_lane_idx = 100; agent.previous_lane_idx = 100; return agent; @@ -53,12 +55,23 @@ static int test_red_crossing_flags(void) { return 0; } -static int test_straddle_then_proceed_on_red_flags(void) { +static int test_front_past_line_on_red_flags(void) { TrafficControlElement tc = drive_test_light(); Drive env = drive_test_light_env(&tc); set_light(TRAFFIC_CONTROL_STATE_RED); - // nose already over the line when the light turned red; proceeding completes entry - Agent agent = drive_test_driving_agent(-1.5f, 0.5f, 0.0f); + // nose in the junction, rear before the line: creeping on red is a violation however the front got there + Agent agent = drive_test_driving_agent(0.5f, 2.5f, 0.0f); + env.agents = &agent; + EXPECT_TRUE(check_red_light_violation(&env, 0)); + return 0; +} + +static int test_stationary_straddle_flags(void) { + TrafficControlElement tc = drive_test_light(); + Drive env = drive_test_light_env(&tc); + set_light(TRAFFIC_CONTROL_STATE_RED); + // parked with the tail on the line when the light turns red + Agent agent = drive_test_driving_agent(1.5f, 1.5f, 0.0f); env.agents = &agent; EXPECT_TRUE(check_red_light_violation(&env, 0)); return 0; @@ -68,35 +81,59 @@ static int test_green_crossing_then_red_inside_no_flag(void) { TrafficControlElement tc = drive_test_light(); Drive env = drive_test_light_env(&tc); set_light(TRAFFIC_CONTROL_STATE_GREEN); - Agent agent = drive_test_driving_agent(1.5f, 2.5f, 0.0f); + Agent agent = drive_test_driving_agent(-0.5f, 0.5f, 0.0f); env.agents = &agent; - EXPECT_TRUE(!check_red_light_violation(&env, 0)); - // light turns red once the agent is fully inside the junction + EXPECT_FALSE(check_red_light_violation(&env, 0)); + // light turns red once the rear (front - 4 m) is more than the 1 m clearance past the line + set_light(TRAFFIC_CONTROL_STATE_RED); + agent = drive_test_driving_agent(5.5f, 6.5f, 0.0f); + EXPECT_FALSE(check_red_light_violation(&env, 0)); + return 0; +} + +static int test_tail_clearance_on_red(void) { + TrafficControlElement tc = drive_test_light(); + Drive env = drive_test_light_env(&tc); set_light(TRAFFIC_CONTROL_STATE_RED); - agent = drive_test_driving_agent(4.0f, 5.0f, 0.0f); - EXPECT_TRUE(!check_red_light_violation(&env, 0)); + // rear 0.9 m past the line: the clearance point still sits before it + Agent agent = drive_test_driving_agent(4.9f, 4.9f, 0.0f); + env.agents = &agent; + EXPECT_TRUE(check_red_light_violation(&env, 0)); + agent = drive_test_driving_agent(5.1f, 5.1f, 0.0f); + EXPECT_FALSE(check_red_light_violation(&env, 0)); return 0; } -static int test_stationary_straddle_no_flag(void) { +static int test_rear_crossing_on_red_flags(void) { TrafficControlElement tc = drive_test_light(); Drive env = drive_test_light_env(&tc); set_light(TRAFFIC_CONTROL_STATE_RED); - Agent agent = drive_test_driving_agent(-0.5f, -0.5f, 0.0f); + // entered on yellow, red while straddling: the rear (front - 4 m) clears the line on red + Agent agent = drive_test_driving_agent(3.5f, 4.5f, 0.0f); env.agents = &agent; - EXPECT_TRUE(!check_red_light_violation(&env, 0)); - // the overlap check (spawn semantics) still sees the straddle - EXPECT_TRUE(check_agent_on_stop_line(&env, &agent, false)); + EXPECT_TRUE(check_red_light_violation(&env, 0)); return 0; } -static int test_crossing_fires_once(void) { +static int test_violation_fires_once_per_stay(void) { TrafficControlElement tc = drive_test_light(); Drive env = drive_test_light_env(&tc); set_light(TRAFFIC_CONTROL_STATE_RED); - Agent agent = drive_test_driving_agent(2.5f, 3.5f, 0.0f); + Agent agent = drive_test_driving_agent(1.5f, 1.5f, 0.0f); env.agents = &agent; - EXPECT_TRUE(!check_red_light_violation(&env, 0)); + EXPECT_TRUE(check_red_light_violation(&env, 0)); + EXPECT_FALSE(check_red_light_violation(&env, 0)); + // a green phase spent on the line does not re-arm the latch + set_light(TRAFFIC_CONTROL_STATE_GREEN); + EXPECT_FALSE(check_red_light_violation(&env, 0)); + set_light(TRAFFIC_CONTROL_STATE_RED); + EXPECT_FALSE(check_red_light_violation(&env, 0)); + // rear 3.5 m past the line: footprint off it, re-armed + agent.sim_x = 7.5f - TEST_HALF_LENGTH_M; + EXPECT_FALSE(check_red_light_violation(&env, 0)); + // reversed back onto it + agent.sim_x = 1.5f - TEST_HALF_LENGTH_M; + EXPECT_TRUE(check_red_light_violation(&env, 0)); return 0; } @@ -117,7 +154,7 @@ static int test_crossing_beyond_virtual_extension_no_flag(void) { set_light(TRAFFIC_CONTROL_STATE_RED); Agent agent = drive_test_driving_agent(-0.5f, 0.5f, 18.0f); env.agents = &agent; - EXPECT_TRUE(!check_red_light_violation(&env, 0)); + EXPECT_FALSE(check_red_light_violation(&env, 0)); return 0; } @@ -129,7 +166,7 @@ static int test_exiting_direction_no_flag(void) { Agent agent = drive_test_agent(1.0f, 0.0f, (float) M_PI); agent.sim_x = -1.0f; env.agents = &agent; - EXPECT_TRUE(!check_red_light_violation(&env, 0)); + EXPECT_FALSE(check_red_light_violation(&env, 0)); return 0; } @@ -137,9 +174,9 @@ static int test_diagonal_entry_flags(void) { TrafficControlElement tc = drive_test_light(); Drive env = drive_test_light_env(&tc); set_light(TRAFFIC_CONTROL_STATE_RED); - // 1.2 rad off the light heading is still entering (< pi/2 gate) - Agent agent = drive_test_agent(-0.5f, 0.0f, 1.2f); - agent.sim_x = 0.5f; + // 1.2 rad off the light heading is still entering (< pi/2 gate); the front edge sits at x 0.22 + Agent agent = drive_test_agent(-1.5f, 0.0f, 1.2f); + agent.sim_x = -0.5f; env.agents = &agent; EXPECT_TRUE(check_red_light_violation(&env, 0)); return 0; @@ -153,7 +190,7 @@ static int test_perpendicular_crossing_no_flag(void) { Agent agent = drive_test_agent(-0.5f, 0.0f, 1.8f); agent.sim_x = 0.0f; env.agents = &agent; - EXPECT_TRUE(!check_red_light_violation(&env, 0)); + EXPECT_FALSE(check_red_light_violation(&env, 0)); return 0; } @@ -164,7 +201,7 @@ static int test_different_z_level_no_flag(void) { Agent agent = drive_test_driving_agent(-0.5f, 0.5f, 0.0f); agent.sim_z = 6.0f; env.agents = &agent; - EXPECT_TRUE(!check_red_light_violation(&env, 0)); + EXPECT_FALSE(check_red_light_violation(&env, 0)); return 0; } @@ -180,13 +217,85 @@ static int test_no_lane_assignment_still_flags(void) { return 0; } +static int test_training_margin_flags_before_line(void) { + TrafficControlElement tc = drive_test_light(); + Drive env = drive_test_light_env(&tc); + env.red_light_train_front_margin_m = 0.3f; + set_light(TRAFFIC_CONTROL_STATE_RED); + Agent agent = drive_test_driving_agent(-0.5f, -0.4f, 0.0f); + env.agents = &agent; + EXPECT_FALSE(check_red_light_violation(&env, 0)); + agent = drive_test_driving_agent(-0.5f, -0.2f, 0.0f); + EXPECT_TRUE(check_red_light_violation(&env, 0)); + return 0; +} + +static int test_training_rear_margin_extends_tail_clearance(void) { + TrafficControlElement tc = drive_test_light(); + Drive env = drive_test_light_env(&tc); + env.red_light_train_rear_margin_m = 1.0f; + set_light(TRAFFIC_CONTROL_STATE_RED); + // rear 2.5 m past the line clears 1 m + margin + Agent agent = drive_test_driving_agent(6.5f, 6.5f, 0.0f); + env.agents = &agent; + EXPECT_FALSE(check_red_light_violation(&env, 0)); + agent = drive_test_driving_agent(5.9f, 5.9f, 0.0f); + EXPECT_TRUE(check_red_light_violation(&env, 0)); + // eval semantics: the real 1 m clearance + env.eval_mode = 1; + agent = drive_test_driving_agent(5.9f, 5.9f, 0.0f); + EXPECT_FALSE(check_red_light_violation(&env, 0)); + agent = drive_test_driving_agent(4.9f, 4.9f, 0.0f); + EXPECT_TRUE(check_red_light_violation(&env, 0)); + return 0; +} + +static int test_eval_semantics_drop_training_margin(void) { + TrafficControlElement tc = drive_test_light(); + Drive env = drive_test_light_env(&tc); + env.red_light_train_front_margin_m = 0.3f; + env.eval_mode = 1; + set_light(TRAFFIC_CONTROL_STATE_RED); + Agent agent = drive_test_driving_agent(-0.5f, -0.2f, 0.0f); + env.agents = &agent; + EXPECT_FALSE(check_red_light_violation(&env, 0)); + agent = drive_test_driving_agent(-0.1f, 0.1f, 0.0f); + EXPECT_TRUE(check_red_light_violation(&env, 0)); + // rendering the training distribution keeps the training margin + env.eval_training_render = 1; + agent = drive_test_driving_agent(-0.5f, -0.2f, 0.0f); + EXPECT_TRUE(check_red_light_violation(&env, 0)); + // co-sim shadow envs use eval semantics without eval_mode + env.eval_mode = 0; + env.eval_training_render = 0; + env.cosim_eval_semantics = 1; + agent = drive_test_driving_agent(-0.5f, -0.2f, 0.0f); + EXPECT_FALSE(check_red_light_violation(&env, 0)); + return 0; +} + +static int test_perceived_size_margin_never_inflates_detection(void) { + TrafficControlElement tc = drive_test_light(); + Drive env = drive_test_light_env(&tc); + env.eval_mode = 1; + env.eval_perceived_size_margin_m = 0.1f; + set_light(TRAFFIC_CONTROL_STATE_RED); + // eval observes the ego 0.2 m longer; the detector keeps the real front edge, which stops 5 cm short + Agent agent = drive_test_driving_agent(-0.15f, -0.05f, 0.0f); + env.agents = &agent; + EXPECT_FALSE(check_red_light_violation(&env, 0)); + return 0; +} + int main(void) { int failures = 0; RUN_TEST(test_red_crossing_flags); - RUN_TEST(test_straddle_then_proceed_on_red_flags); + RUN_TEST(test_front_past_line_on_red_flags); + RUN_TEST(test_stationary_straddle_flags); RUN_TEST(test_green_crossing_then_red_inside_no_flag); - RUN_TEST(test_stationary_straddle_no_flag); - RUN_TEST(test_crossing_fires_once); + RUN_TEST(test_tail_clearance_on_red); + RUN_TEST(test_rear_crossing_on_red_flags); + RUN_TEST(test_violation_fires_once_per_stay); RUN_TEST(test_opposing_lane_bypass_flags); RUN_TEST(test_crossing_beyond_virtual_extension_no_flag); RUN_TEST(test_exiting_direction_no_flag); @@ -194,5 +303,9 @@ int main(void) { RUN_TEST(test_perpendicular_crossing_no_flag); RUN_TEST(test_different_z_level_no_flag); RUN_TEST(test_no_lane_assignment_still_flags); + RUN_TEST(test_training_margin_flags_before_line); + RUN_TEST(test_training_rear_margin_extends_tail_clearance); + RUN_TEST(test_eval_semantics_drop_training_margin); + RUN_TEST(test_perceived_size_margin_never_inflates_detection); return test_summary(failures); } diff --git a/tests/drive/test_drive_traffic_lights.c b/tests/drive/test_drive_traffic_lights.c index 9c7b3c088c..011d336a34 100644 --- a/tests/drive/test_drive_traffic_lights.c +++ b/tests/drive/test_drive_traffic_lights.c @@ -195,6 +195,59 @@ static int test_schedule_covers_terminal_observation(void) { return 0; } +// The optional yellow-remaining column (before type/state) reads the consecutive yellow steps ahead in +// seconds over traffic_light_yellow_max_s, saturating at 1; every other state reads 0 +static int test_yellow_remaining_obs_counts_down_and_saturates(void) { + Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); + env.eval_mode = 1; + int plain_obs_size = compute_observation_size(&env); + env.obs_traffic_light_yellow_remaining = 1; + int row_features = TRAFFIC_CONTROL_FEATURES + TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES; + EXPECT_EQ_INT(compute_observation_size(&env), plain_obs_size + env.obs_slots_traffic_controls_n); + allocate(&env); + c_reset(&env); + TrafficControlElement *probe_light = NULL; + for (int i = 0; i < env.num_traffic_elements && probe_light == NULL; i++) { + if (env.traffic_elements[i].type == TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT) { + probe_light = &env.traffic_elements[i]; + } + } + EXPECT_TRUE(probe_light != NULL); + int yellow_start = 10; + int yellow_steps = 40; // 4 s at dt 0.1: longer than the 3 s training max + for (int t = 0; t < probe_light->state_size; t++) { + probe_light->states[t] = TRAFFIC_CONTROL_STATE_RED; + } + for (int t = yellow_start; t < yellow_start + yellow_steps; t++) { + probe_light->states[t] = TRAFFIC_CONTROL_STATE_YELLOW; + } + 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); + int valid_base = compute_observation_size(&env) - OBS_VALID_COUNT_FEATURES; + int traffic_base = valid_base - env.obs_slots_traffic_controls_n * row_features; + int yellow_column = traffic_base + row_features - 3; + int probe_timesteps[4] = {yellow_start, yellow_start + 20, yellow_start + 39, yellow_start + 40}; + float expected[4] = {1.0f, 2.0f / 3.0f, 0.1f / 3.0f, 0.0f}; + int expected_state[4] + = {TRAFFIC_CONTROL_STATE_YELLOW, + TRAFFIC_CONTROL_STATE_YELLOW, + TRAFFIC_CONTROL_STATE_YELLOW, + TRAFFIC_CONTROL_STATE_RED}; + for (int k = 0; k < 4; k++) { + env.timestep = probe_timesteps[k]; + compute_observations(&env); + EXPECT_TRUE(env.observations[valid_base + 3] >= 1.0f); + EXPECT_TRUE(fabsf(env.observations[yellow_column] - expected[k]) < 1e-5f); + EXPECT_EQ_INT((int) env.observations[traffic_base + row_features - 2], TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT); + EXPECT_EQ_INT((int) env.observations[traffic_base + row_features - 1], expected_state[k]); + } + free_allocated(&env); + return 0; +} + int main(void) { int failures = 0; RUN_TEST(test_junction_lights_never_green_together); @@ -205,5 +258,6 @@ int main(void) { RUN_TEST(test_training_group_removal_turns_off_whole_junction); RUN_TEST(test_eval_never_removes_lights); RUN_TEST(test_schedule_covers_terminal_observation); + RUN_TEST(test_yellow_remaining_obs_counts_down_and_saturates); return test_summary(failures); } diff --git a/tests/eval/test_eval.py b/tests/eval/test_eval.py index ae5c261ead..81357b608c 100644 --- a/tests/eval/test_eval.py +++ b/tests/eval/test_eval.py @@ -2,6 +2,7 @@ import random import struct import sys +import types import zlib from pathlib import Path from unittest.mock import patch @@ -32,6 +33,9 @@ TRAIN_HORIZON = 16 TRAIN_ENV_COUNT = 2 TRAIN_AGENTS_PER_ENV = 8 +WANDB_RUN_NAME = "k_exp_0002_1000" +WANDB_PROJECT = "nightly-multi-long" +WANDB_TRAINING_SUMMARY = {"environment/collision_rate": 0.0123} REPO_ROOT = Path(__file__).resolve().parents[2] CARLA_MAP_DIR = REPO_ROOT / "pufferlib/resources/drive/binaries/carla" REPLAY_MAP_DIR = REPO_ROOT / "pufferlib/resources/drive/binaries/sdc_replay_test" @@ -297,6 +301,70 @@ def test_eval_report_reduces_distance_per_infraction_from_raw_totals(): assert summary["metrics_mean"]["avg_distance_per_infraction"] == pytest.approx(200.0 / 11.0) +def _run_standalone_eval_with_wandb(tmp_path, monkeypatch, benchmark_name): + """Evaluate a checkpoint folder with wandb=True against a stub wandb -> (run dir, stored summary, run paths).""" + run_dir = tmp_path / WANDB_RUN_NAME + run_dir.mkdir() + benchmark_config_path = _write_benchmark_config( + tmp_path / "benchmark.yaml", + name=benchmark_name, + map_dir=CARLA_MAP_DIR, + simulation_mode="gigaflow", + num_maps=CARLA_MAP_COUNT, + num_scenarios=CARLA_SCENARIO_COUNT, + scenario_length=CARLA_SCENARIO_LENGTH, + max_agents_per_env=8, + control_mode="control_vehicles", + use_neighbor_cache=True, + ) + args = _standalone_eval_args(benchmark_config_path) + args["eval"]["benchmarks"] = benchmark_name + args["run_name"] = WANDB_RUN_NAME + args["wandb_project"] = WANDB_PROJECT + (run_dir / "config.yaml").write_text(yaml.safe_dump(json.loads(json.dumps(args)))) + (run_dir / "final_model.pt").write_bytes(b"") + args["load_model_path"] = str(run_dir / "final_model.pt") + args["wandb"] = True + + stored_summary = dict(WANDB_TRAINING_SUMMARY) + requested_run_paths = [] + + def fetch_run(run_path): + requested_run_paths.append(run_path) + return types.SimpleNamespace( + state="finished", + summary_metrics=dict(stored_summary), + summary=types.SimpleNamespace(update=stored_summary.update), + ) + + stub_wandb = types.ModuleType("wandb") + stub_wandb.Api = lambda: types.SimpleNamespace(run=fetch_run) + monkeypatch.setitem(sys.modules, "wandb", stub_wandb) + + pufferl.eval(env_name="puffer_drive", args=args, policy=ZeroPolicy(action_count=63)) + return run_dir, stored_summary, requested_run_paths + + +def test_standalone_selfplay_eval_reports_its_km_per_infraction_to_wandb(tmp_path, monkeypatch): + run_dir, stored_summary, requested_run_paths = _run_standalone_eval_with_wandb(tmp_path, monkeypatch, "carla") + + summary_path = next(run_dir.glob("eval/carla/*/evaluation_summary.json")) + distance_per_infraction_meters = json.loads(summary_path.read_text())["metrics_mean"]["avg_distance_per_infraction"] + assert set(requested_run_paths) == {f"{WANDB_PROJECT}/{WANDB_RUN_NAME}"} + assert stored_summary == { + **WANDB_TRAINING_SUMMARY, + "final_eval/selfplay_km_per_infraction": distance_per_infraction_meters / 1000.0, + "final_eval/selfplay_km_per_infraction_source": str(summary_path.relative_to(run_dir)), + } + + +def test_standalone_eval_of_another_benchmark_reports_nothing_to_wandb(tmp_path, monkeypatch): + _, stored_summary, requested_run_paths = _run_standalone_eval_with_wandb(tmp_path, monkeypatch, "carla_fast") + + assert requested_run_paths == [] + assert stored_summary == WANDB_TRAINING_SUMMARY + + def test_seed_replay_writes_exactly_identical_metrics(carla_evaluation): args = carla_evaluation["args"] benchmark_config_path = carla_evaluation["benchmark_config_path"] diff --git a/tests/smoke_tests/data/drive_rollout_golden.json b/tests/smoke_tests/data/drive_rollout_golden.json index 5b983ca767..d282277546 100644 --- a/tests/smoke_tests/data/drive_rollout_golden.json +++ b/tests/smoke_tests/data/drive_rollout_golden.json @@ -1,43 +1,46 @@ { "env": { - "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, + "avg_distance_per_infraction": 13.32529857158661, + "avg_speed_per_agent": 1.3309484645724297, + "collision_rate": 0.017125687561929226, + "comfort_violation_count": 0.611062702909112, + "dnf_rate": 0.5401545438915492, "early_reset_short": 0.0, - "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, + "episode_length": 14.6875, + "episode_return": -1.1725615710020065, + "lane_center_rate": 0.26990433819591997, + "n": 14.5125, + "num_goals_reached": 0.008311012107878924, + "offroad_rate": 0.42702610287815335, + "pedestrian_bystanders": 0.0, + "pedestrian_collision_rate": 0.0, + "pedestrian_crossings": 0.0, + "red_light_violation_rate": 0.009838026948273182, "reward_components/ade": 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.011113312932415098, + "reward_components/collision": -0.03214438278228045, + "reward_components/comfort": -0.44294892586767676, + "reward_components/goal": 0.008311012107878924, + "reward_components/lane_align": -0.026388277858495714, + "reward_components/lane_center": -0.011815103818662465, + "reward_components/offroad": -0.6299905642867089, + "reward_components/overspeed": -0.02166639088682132, + "reward_components/red_light": -0.005256565727177076, + "reward_components/reverse": -0.01083271369279828, "reward_components/stop_sign": 0.0, - "reward_components/timestep": -9.126150063268511e-05, - "reward_components/velocity": 0.00021098649279027292, + "reward_components/timestep": -8.983786779026559e-05, + "reward_components/velocity": 0.00026019296847152874, "score": 0.0, "spawn_failed": 0.0, "spawn_failed_goal": 0.0, - "spawn_reject_collision": 0.06356702546907377, + "spawn_reject_collision": 0.04464400317519903, "spawn_reject_empty_cell": 0.0, - "spawn_reject_offroad": 0.40240530132115643, - "spawn_reject_stop_line": 0.020339523783967466, + "spawn_reject_offroad": 0.4232177251949906, + "spawn_reject_stop_line": 0.016670101415365936, "stop_sign_violation_rate": 0.0, "stopped_at_reset": 0.0, - "total_distance_travelled_sum": 85.30317799049088, - "total_infraction_count": 6.544303797468355, - "velocity_progress_sum": 0.019357881364353662 + "total_distance_travelled_sum": 85.11123423576355, + "total_infraction_count": 6.4875, + "velocity_progress_sum": 0.02372938830521889 }, "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 df6979618d..7a0f60a7fb 100644 --- a/tests/smoke_tests/data/drive_smoke_golden.json +++ b/tests/smoke_tests/data/drive_smoke_golden.json @@ -1,65 +1,68 @@ { "env": { - "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, + "avg_distance_per_infraction": 14.478009234417925, + "avg_speed_per_agent": 1.3486858678586555, + "collision_rate": 0.025252526343771904, + "comfort_violation_count": 0.6310719248622355, + "distance_per_agent": 6.654034547131471, + "dnf_rate": 0.5252525335309481, "early_reset_short": 0.0, - "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, + "episode_length": 16.424242424242426, + "episode_return": -1.2947040934755345, + "infractions_per_agent": 0.4595959595959596, + "lane_center_rate": 0.28529129591253066, + "n": 198.0, + "num_goals_reached": 0.010101010326785271, + "obs/max": 1.0, + "obs/mean": 0.12931787461275235, "obs/min": -1.0000000009313226, - "offroad_rate": 0.45238095649651117, + "offroad_rate": 0.4444444488846894, + "pedestrian_bystanders": 0.0, + "pedestrian_collision_rate": 0.0, + "pedestrian_crossings": 0.0, "red_light_violation_rate": 0.0, "reward_components/ade": 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/collision": -0.045137426768890536, + "reward_components/comfort": -0.5125306743865061, + "reward_components/goal": 0.010101010326785271, + "reward_components/lane_align": -0.032201278023421764, + "reward_components/lane_center": -0.01303829498720741, + "reward_components/offroad": -0.6757460800987302, + "reward_components/overspeed": -0.012211139871259078, "reward_components/red_light": 0.0, - "reward_components/reverse": -0.012323847098187322, + "reward_components/reverse": -0.013902073166323732, "reward_components/stop_sign": 0.0, - "reward_components/timestep": -9.978571310750253e-05, - "reward_components/velocity": 0.000197039804022227, + "reward_components/timestep": -0.00010208333274343691, + "reward_components/velocity": 6.394532847663886e-05, "score": 0.0, "spawn_failed": 0.0, "spawn_failed_goal": 0.0, - "spawn_reject_collision": 0.07142857298964546, + "spawn_reject_collision": 0.040404041871579005, "spawn_reject_empty_cell": 0.0, - "spawn_reject_offroad": 0.40476190976443743, - "spawn_reject_stop_line": 0.014285714285714285, + "spawn_reject_offroad": 0.4040404089321994, + "spawn_reject_stop_line": 0.0252525263814011, "stop_sign_violation_rate": 0.0, "stopped_at_reset": 0.0, - "total_distance_travelled": 1363.0418014526367, - "total_infractions": 95.0, - "velocity_progress_sum": 0.01484889875476559 + "total_distance_travelled": 1317.4988403320312, + "total_infractions": 91.0, + "velocity_progress_sum": 0.005202303671118135 }, "losses": { - "actor_grad_norm": 1.2954603433609009, - "approx_kl": 0.011499513871967793, - "clipfrac": 0.22194667160511017, - "critic_grad_norm": 0.14509320259094238, - "ema_max": 3.0478071570396423, + "actor_grad_norm": 1.365248680114746, + "approx_kl": 0.007157053332775831, + "clipfrac": 0.14219732582569122, + "critic_grad_norm": 0.7422468066215515, + "ema_max": 2.855304539203644, "ent_coef": 0.01, - "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 + "entropy": 2.457249879837036, + "explained_variance": 0.08638936281204224, + "filter_threshold": 0.028553045392036437, + "filtered_fraction": 0.08272667107875575, + "kept_fraction": 0.9172733289212442, + "masked_fraction": 0.26220703125, + "old_approx_kl": 0.011373384855687618, + "policy_loss": -0.004838721826672554, + "value_loss": 0.15589582920074463 }, "meta": { "bptt_horizon": 64, diff --git a/tests/unit_tests/test_carla_pedal_map.py b/tests/unit_tests/test_carla_pedal_map.py new file mode 100644 index 0000000000..929d593038 --- /dev/null +++ b/tests/unit_tests/test_carla_pedal_map.py @@ -0,0 +1,180 @@ +import math + +import numpy as np +import pytest + +from pufferlib.ocean.cosim.carla.controller import ( + PEDAL_GAINS_BY_BLUEPRINT, + REVERSE_GEAR, + PedalMap, + pedal_gains_for, + pedals, +) + + +GAINS = PEDAL_GAINS_BY_BLUEPRINT["vehicle.lincoln.mkz_2020"] +SECOND_GEAR = GAINS.gears[1] +TICK_SECONDS = 0.05 +CRUISE_SPEED_MPS = 8.0 +MAX_SPEED_MPS = 30.0 + + +def _holding_throttle(gear_gains, speed_mps): + return float(np.interp(speed_mps, gear_gains.speeds_mps, gear_gains.throttle_holding_speed)) + + +def _accel_per_throttle(gear_gains, speed_mps): + return float(np.interp(speed_mps, gear_gains.speeds_mps, gear_gains.accel_per_throttle_mps2)) + + +def _engine_braking(gear_gains, speed_mps): + steady = _holding_throttle(gear_gains, speed_mps) * _accel_per_throttle(gear_gains, speed_mps) + return min(steady, GAINS.max_engine_braking_mps2) + + +def test_zero_acceleration_is_the_speed_holding_throttle_of_the_gear(): + throttle, brake = pedals( + SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, 0.0, CRUISE_SPEED_MPS + ) + assert brake == 0.0 + assert math.isclose(throttle, _holding_throttle(SECOND_GEAR, CRUISE_SPEED_MPS)) + + +def test_an_extra_mps2_costs_more_throttle_in_a_higher_gear(): + def throttle_per_mps2(gear_gains): + slow = pedals(gear_gains, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, 1.0, CRUISE_SPEED_MPS)[0] + fast = pedals(gear_gains, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, 2.0, CRUISE_SPEED_MPS)[0] + return fast - slow + + assert throttle_per_mps2(GAINS.gears[3]) > throttle_per_mps2(GAINS.gears[0]) + + +def test_gentle_deceleration_only_lifts_the_throttle(): + engine_braking = _engine_braking(SECOND_GEAR, CRUISE_SPEED_MPS) + throttle, brake = pedals( + SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, -0.5 * engine_braking, CRUISE_SPEED_MPS + ) + assert brake == 0.0 + assert 0.0 < throttle < _holding_throttle(SECOND_GEAR, CRUISE_SPEED_MPS) + + +def test_the_brake_supplies_what_engine_braking_does_not(): + engine_braking = _engine_braking(SECOND_GEAR, CRUISE_SPEED_MPS) + _, brake = pedals( + SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, -engine_braking - 1.0, CRUISE_SPEED_MPS + ) + assert math.isclose(brake, GAINS.brake_per_mps2) + + +def test_while_the_gearbox_shifts_the_brake_supplies_the_whole_deceleration(): + _, brake = pedals(SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, True, -1.0, CRUISE_SPEED_MPS) + assert math.isclose(brake, GAINS.brake_per_mps2) + + +def test_pedals_stay_within_zero_and_one(): + assert pedals(SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, 100.0, 0.0) == (1.0, 0.0) + assert pedals(SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, -100.0, 0.0) == (0.0, 1.0) + + +@pytest.mark.parametrize("blueprint", sorted(PEDAL_GAINS_BY_BLUEPRINT)) +def test_every_gear_has_usable_numbers_at_increasing_speeds(blueprint): + gains = pedal_gains_for(blueprint) + for gear_gains in (*gains.gears, gains.reverse): + assert all(low < high for low, high in zip(gear_gains.speeds_mps, gear_gains.speeds_mps[1:])) + assert ( + len(gear_gains.speeds_mps) + == len(gear_gains.throttle_holding_speed) + == len(gear_gains.accel_per_throttle_mps2) + ) + assert all(0.0 < throttle < 1.0 for throttle in gear_gains.throttle_holding_speed) + assert all(accel > 1.0 for accel in gear_gains.accel_per_throttle_mps2) + assert gains.gears[0].speeds_mps[0] == 0.0 + assert gains.reverse.speeds_mps[0] == 0.0 + assert gains.rear_axle_to_center_m > 0.0 + assert gains.curvature_per_steer_at_rest + gains.curvature_per_steer_per_mps * MAX_SPEED_MPS > 0.2 + + +def test_outside_its_speed_span_a_gear_keeps_the_numbers_of_the_nearest_end(): + low_speed_mps, high_speed_mps = SECOND_GEAR.speeds_mps[0], SECOND_GEAR.speeds_mps[-1] + below = pedals(SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, 1.0, low_speed_mps - 3.0) + assert below == pedals(SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, 1.0, low_speed_mps) + above = pedals(SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, 1.0, high_speed_mps + 3.0) + assert above == pedals(SECOND_GEAR, GAINS.max_engine_braking_mps2, GAINS.brake_per_mps2, False, 1.0, high_speed_mps) + + +def test_steer_follows_the_curvature_sign_weakens_with_speed_and_clips(): + pedal_map = PedalMap(GAINS) + slow = pedal_map.step(0.0, 0.1, 2.0, 1, False, False)[0] + fast = pedal_map.step(0.0, 0.1, 20.0, 4, False, False)[0] + assert math.isclose(slow, 0.1 / (GAINS.curvature_per_steer_at_rest + GAINS.curvature_per_steer_per_mps * 2.0)) + assert fast > slow + assert pedal_map.step(0.0, -0.1, 5.0, 2, False, False)[0] < 0.0 + assert pedal_map.step(0.0, 10.0, 5.0, 2, False, False)[0] == 1.0 + + +def test_a_shift_keeps_the_throttle_line_of_the_gear_last_engaged(): + pedal_map = PedalMap(GAINS) + in_third = pedal_map.step(1.0, 0.0, CRUISE_SPEED_MPS, 3, False, False) + shifting = pedal_map.step(1.0, 0.0, CRUISE_SPEED_MPS, 0, False, False) + assert shifting == in_third + + +def test_standstill_hold_is_full_brake_and_is_not_scored(): + pedal_map = PedalMap(GAINS) + assert pedal_map.step(0.0, 0.0, 0.0, 1, True, False) == (0.0, 0.0, 1.0, False) + pedal_map.observe(0.0, 0.0, TICK_SECONDS) + assert pedal_map.stats() == {} + + +def test_a_backward_acceleration_in_reverse_presses_the_throttle_by_the_reverse_gains(): + pedal_map = PedalMap(GAINS) + backward_accel_mps2 = 1.0 + steer, throttle, brake, reverse = pedal_map.step(-backward_accel_mps2, 0.1, 0.0, 1, False, True) + expected_throttle = _holding_throttle(GAINS.reverse, 0.0) + backward_accel_mps2 / _accel_per_throttle( + GAINS.reverse, 0.0 + ) + assert reverse + assert brake == 0.0 + assert math.isclose(throttle, expected_throttle) + assert math.isclose(steer, 0.1 / GAINS.curvature_per_steer_at_rest) + + +def test_slowing_down_a_car_that_rolls_backwards_uses_the_brake(): + pedal_map = PedalMap(GAINS) + backward_speed_mps = 1.5 + slowing_accel_mps2 = _engine_braking(GAINS.reverse, backward_speed_mps) + 1.0 + _, throttle, brake, reverse = pedal_map.step( + slowing_accel_mps2, 0.0, -backward_speed_mps, REVERSE_GEAR, False, True + ) + assert reverse + assert throttle == 0.0 + assert math.isclose(brake, GAINS.brake_per_mps2) + + +def test_reverse_ticks_are_scored_under_the_reverse_gear(): + pedal_map = PedalMap(GAINS) + pedal_map.step(-1.0, 0.0, 0.0, 1, False, True) + pedal_map.observe(-1.0 * TICK_SECONDS, 0.0, TICK_SECONDS) + assert pedal_map.stats()["accel_bias_by_gear_mps2"] == {REVERSE_GEAR: (1, 0.0)} + + +def test_stats_report_the_acceleration_bias_and_the_yaw_rate_gain(): + pedal_map = PedalMap(GAINS) + speed_mps = 5.0 + hard_braking_mps2 = -_engine_braking(SECOND_GEAR, speed_mps) - 2.0 + for target_accel_mps2 in (1.0, 1.0, hard_braking_mps2, hard_braking_mps2): + pedal_map.step(target_accel_mps2, 0.1, speed_mps, 2, False, False) + speed_mps += (target_accel_mps2 - 0.5) * TICK_SECONDS + pedal_map.observe(speed_mps, 0.5 * 0.1 * speed_mps, TICK_SECONDS) + stats = pedal_map.stats() + assert stats["ticks"] == 4 + assert math.isclose(stats["accel_bias_throttle_mps2"], -0.5, abs_tol=1e-9) + assert math.isclose(stats["accel_bias_brake_mps2"], -0.5, abs_tol=1e-9) + assert math.isclose(stats["yaw_rate_gain"], 0.5, abs_tol=1e-9) + + +def test_an_unmeasured_car_is_refused(): + with pytest.raises(ValueError): + pedal_gains_for("vehicle.tesla.model3") + with pytest.raises(ValueError): + pedal_gains_for(None) diff --git a/tests/unit_tests/test_drive_config.py b/tests/unit_tests/test_drive_config.py index 364da1755a..155922db29 100644 --- a/tests/unit_tests/test_drive_config.py +++ b/tests/unit_tests/test_drive_config.py @@ -87,6 +87,23 @@ def test_lane_speed_limit_obs_widens_lane_rows(self): with_limit.num_obs - plain.num_obs, plain.obs_slots_lane_kept * binding.LANE_SPEED_LIMIT_FEATURES ) + @patch("sys.argv", ["pufferl.py", "env.obs_traffic_light_yellow_remaining=true"]) + def test_yellow_remaining_cli_override(self): + args = load_config("puffer_drive") + self.assertTrue(args["env"]["obs_traffic_light_yellow_remaining"]) + + def test_yellow_remaining_obs_widens_traffic_control_rows(self): + plain = Drive(config_only=True) + with_yellow = Drive(config_only=True, obs_traffic_light_yellow_remaining=True) + self.assertEqual( + with_yellow.traffic_control_features, + plain.traffic_control_features + binding.TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES, + ) + self.assertEqual( + with_yellow.num_obs - plain.num_obs, + plain.obs_slots_traffic_controls_n * binding.TRAFFIC_CONTROL_YELLOW_REMAINING_FEATURES, + ) + @patch("sys.argv", ["pufferl.py", "train.learning_rate=0.5"]) def test_cli_override(self): """Test that Hydra CLI overrides win over the config file values.""" diff --git a/tests/unit_tests/test_drive_cosim_lane_goals.py b/tests/unit_tests/test_drive_cosim_lane_goals.py new file mode 100644 index 0000000000..070a2f67f2 --- /dev/null +++ b/tests/unit_tests/test_drive_cosim_lane_goals.py @@ -0,0 +1,87 @@ +"""Co-sim forecast goals through the Python binding: agents without a known route get goals along their lane +(Drive.set_agent_lane_goals) and a rollout's goal progress can be put back (Drive.set_agent_goal_progress).""" + +import numpy as np +import pytest + +from pufferlib.ocean.drive.drive import Drive + + +MAP_DIR = "pufferlib/resources/drive/binaries/carla" +POLICY_AGENTS = 3 +PARTNER_SLOTS = 2 +GOAL_SPACING_M = 20.0 + + +@pytest.fixture +def env(): + drive = Drive( + num_agents=POLICY_AGENTS, + min_agents_per_env=POLICY_AGENTS, + max_agents_per_env=POLICY_AGENTS, + cosim_partner_slots=PARTNER_SLOTS, + goal_source="external", + num_maps=1, + allow_map_subset=True, + map_dir=MAP_DIR, + simulation_mode="gigaflow", + action_type="continuous", + dynamics_model="jerk", + scenario_length=300, + resample_frequency=0, + termination_mode=0, + stagger_first_episode=False, + collision_behavior="ignore", + offroad_behavior="ignore", + seed=3, + ) + drive.reset(seed=3) + yield drive + drive.close() + + +def test_lane_goals_are_placed_ahead_of_spawned_agents(env): + assert env.get_agent_goal_progress(0) == (0, 0) # external goal source: spawning assigns no goals + counts = env.set_agent_lane_goals(np.arange(POLICY_AGENTS), GOAL_SPACING_M) + assert counts.dtype == np.int32 and counts.shape == (POLICY_AGENTS,) + assert np.all(counts >= 1) and np.all(counts <= env.num_goals) # spawned on a lane with road ahead + for agent_idx in range(POLICY_AGENTS): + assert env.get_agent_goal_progress(agent_idx) == (0, int(counts[agent_idx])) + + +def test_lane_goals_survive_a_rollout_of_coasting_agents(env): + env.set_agent_lane_goals(np.arange(POLICY_AGENTS), GOAL_SPACING_M) + actions = np.zeros_like(env.actions) + for _ in range(40): + env.step(actions) + for agent_idx in range(POLICY_AGENTS): + current_goal_idx, goal_count = env.get_agent_goal_progress(agent_idx) + assert 0 <= current_goal_idx <= goal_count <= env.num_goals + assert np.all(np.isfinite(env.observations)) + + +def test_goal_progress_can_be_set_back_and_forth(env): + counts = env.set_agent_lane_goals(np.array([0]), GOAL_SPACING_M) + goal_count = int(counts[0]) + env.set_agent_goal_progress(0, goal_count) + assert env.get_agent_goal_progress(0) == (goal_count, goal_count) + env.set_agent_goal_progress(0, 0) + assert env.get_agent_goal_progress(0) == (0, goal_count) + with pytest.raises(ValueError, match="goal_count"): + env.set_agent_goal_progress(0, goal_count + 1) + with pytest.raises(ValueError, match="goal_count"): + env.set_agent_goal_progress(0, -1) + with pytest.raises(ValueError, match="out of range"): + env.set_agent_goal_progress(POLICY_AGENTS + PARTNER_SLOTS, 0) + + +def test_lane_goals_reject_bad_input(env): + with pytest.raises(ValueError, match="spacing_m"): + env.set_agent_lane_goals(np.array([0]), 0.0) + with pytest.raises(ValueError, match="spacing_m"): + env.set_agent_lane_goals(np.array([0]), float("nan")) + with pytest.raises(ValueError, match="out of range"): + env.set_agent_lane_goals(np.array([POLICY_AGENTS + PARTNER_SLOTS]), GOAL_SPACING_M) + with pytest.raises(ValueError, match="without a pose"): + env.set_agent_lane_goals(np.array([POLICY_AGENTS]), GOAL_SPACING_M) # a partner slot nothing was put in + assert env.set_agent_lane_goals(np.zeros(0, np.int32), GOAL_SPACING_M).shape == (0,) diff --git a/tests/unit_tests/test_ensemble_policy.py b/tests/unit_tests/test_ensemble_policy.py new file mode 100644 index 0000000000..847a10b56a --- /dev/null +++ b/tests/unit_tests/test_ensemble_policy.py @@ -0,0 +1,271 @@ +"""Evaluation-time ensembles: averaged member predictions, loading from an ensemble checkpoint, the checkpoint builder.""" + +import copy +from pathlib import Path + +import pytest +import torch +import yaml +from torch import nn + +import pufferlib.pytorch as P +from pufferlib import ensemble + + +ACTION_TABLE = torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]) +ACTION_TABLE_PHYSICAL = torch.tensor([[-4.0, -1.0], [0.0, 0.0], [2.0, 1.0]]) +LONG_NEG_SCALE = 4.0 +LONG_POS_SCALE = 2.0 +BATCH_SIZE = 4 +OBS_DIM = 2 +RUN_CONFIG = { + "run_name": "member", + "wandb_project": "nightly", + "wandb": True, + "policy_name": "Drive", + "rnn_name": None, + "rnn": {"input_size": 128, "hidden_size": 128}, + "policy": {"actor_hidden_size": 8, "fused_slot_encoder": True}, + "env": {"dt": 0.1, "num_maps": 128}, +} + + +class _TinyPolicy(nn.Module): + """Discrete head on a continuous env with Drive's piecewise longitudinal scale; constant outputs, checkable by hand.""" + + is_continuous = False + + def __init__(self, logits=(0.0, 0.0, 0.0), value=0.0): + super().__init__() + self.logit_bias = nn.Parameter(torch.tensor(logits)) + self.value_bias = nn.Parameter(torch.tensor([value])) + + def forward_eval(self, observations, state=None): + batch = observations.shape[0] + return self.logit_bias.expand(batch, -1), self.value_bias.expand(batch, -1) + + def discrete_actions_to_continuous(self, actions): + return ACTION_TABLE[actions.long()] + + def discrete_probs_to_continuous_mean(self, probs): + mean_physical = probs @ ACTION_TABLE_PHYSICAL.to(probs.dtype) + mean_long = mean_physical[..., 0] + mean_long_norm = torch.where(mean_long < 0.0, mean_long / LONG_NEG_SCALE, mean_long / LONG_POS_SCALE) + return torch.stack([mean_long_norm, mean_physical[..., 1]], dim=-1) + + +def _observations(): + return torch.zeros(BATCH_SIZE, OBS_DIM) + + +def _ensemble_checkpoint(*members): + return { + ensemble.ENSEMBLE_CHECKPOINT_KEY: { + "format_version": ensemble.ENSEMBLE_FORMAT_VERSION, + "member_checkpoints": [f"member_{idx}.pt" for idx in range(len(members))], + "member_state_dicts": [member.state_dict() for member in members], + } + } + + +def _write_run_dir(root, name, state_dict, config_overrides=None): + run_dir = root / name + run_dir.mkdir() + config = {**copy.deepcopy(RUN_CONFIG), "run_name": name, **(config_overrides or {})} + (run_dir / "config.yaml").write_text(yaml.safe_dump(config)) + torch.save(state_dict, run_dir / "final_model.pt") + return str(run_dir / "final_model.pt") + + +def test_discrete_average_is_the_mean_of_member_probabilities_not_a_vote(): + confident = _TinyPolicy(logits=(4.0, 0.0, 0.0)) + hesitant = _TinyPolicy(logits=(0.0, 1.0, 0.5)) + policy = ensemble.EnsemblePolicy([confident, hesitant]) + + logits, _ = policy.forward_eval(_observations()) + member_probs = torch.stack([torch.softmax(member.logit_bias, -1) for member in (confident, hesitant)]) + + assert torch.allclose(torch.softmax(logits, -1), member_probs.mean(0).expand(BATCH_SIZE, -1), atol=1e-6) + mode, _, _, _ = P.sample_logits(logits, action_selection=P.ACTION_SELECT_MODE, env_continuous=True, policy=policy) + assert mode.tolist() == [0] * BATCH_SIZE + + +def test_mean_action_selection_averages_the_member_mean_actions_in_physical_units(): + braking = _TinyPolicy(logits=(9.0, 0.0, 0.0)) + accelerating = _TinyPolicy(logits=(0.0, 0.0, 9.0)) + policy = ensemble.EnsemblePolicy([braking, accelerating]) + + logits, _ = policy.forward_eval(_observations()) + _, _, _, mean_action = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MEAN, env_continuous=True, policy=policy + ) + mixture_expectation = braking.discrete_probs_to_continuous_mean(torch.softmax(logits, -1)) + + # members execute normalised long -1 and +1; the brake scale is the larger one, so their physical mean brakes + assert torch.allclose(mean_action[:, 0], torch.full((BATCH_SIZE,), -0.25), atol=1e-3) + assert torch.allclose(mean_action, mixture_expectation, atol=1e-6) + + +def test_mean_action_needs_the_member_probabilities_of_the_same_batch(): + policy = ensemble.EnsemblePolicy([_TinyPolicy(), _TinyPolicy()]) + uniform = torch.full((1, BATCH_SIZE, 3), 1.0 / 3.0) + + with pytest.raises(RuntimeError, match="forward_eval"): + policy.discrete_probs_to_continuous_mean(uniform) + policy.forward_eval(_observations()) + with pytest.raises(RuntimeError, match="same batch"): + policy.discrete_probs_to_continuous_mean(uniform[:, : BATCH_SIZE - 1]) + assert policy.discrete_probs_to_continuous_mean(uniform).shape == (1, BATCH_SIZE, 2) + + +def test_values_are_averaged(): + policy = ensemble.EnsemblePolicy([_TinyPolicy(value=1.0), _TinyPolicy(value=3.0)]) + + _, value = policy.forward_eval(_observations()) + + assert value.shape == (BATCH_SIZE, 1) + assert torch.allclose(value, torch.full((BATCH_SIZE, 1), 2.0)) + + +def test_normal_average_matches_the_mixture_moments(): + members = [ + torch.distributions.Normal(torch.zeros(BATCH_SIZE, 2), torch.ones(BATCH_SIZE, 2)), + torch.distributions.Normal(torch.full((BATCH_SIZE, 2), 2.0), torch.ones(BATCH_SIZE, 2)), + ] + + mixture = ensemble.average_member_predictions(members) + + assert torch.allclose(mixture.loc, torch.ones(BATCH_SIZE, 2)) + assert torch.allclose(mixture.scale, torch.full((BATCH_SIZE, 2), 2.0**0.5)) + + +def test_multi_discrete_outputs_average_per_component(): + member_a = (torch.tensor([[4.0, 0.0]]), torch.tensor([[0.0, 0.0, 0.0]])) + member_b = (torch.tensor([[0.0, 4.0]]), torch.tensor([[0.0, 0.0, 0.0]])) + + mixture = ensemble.average_member_predictions([member_a, member_b]) + + assert isinstance(mixture, tuple) and len(mixture) == 2 + assert torch.allclose(torch.softmax(mixture[0], -1), torch.tensor([[0.5, 0.5]]), atol=1e-6) + assert torch.allclose(torch.softmax(mixture[1], -1), torch.full((1, 3), 1.0 / 3.0), atol=1e-6) + + +def test_load_policy_weights_keeps_a_trained_checkpoint_as_the_same_policy(): + trained = _TinyPolicy(logits=(1.0, 2.0, 3.0), value=5.0) + checkpoint = {f"_orig_mod.module.{key}": value for key, value in trained.state_dict().items()} + policy = _TinyPolicy() + + loaded = ensemble.load_policy_weights(policy, checkpoint) + + assert loaded is policy + assert torch.equal(loaded.logit_bias, trained.logit_bias) + assert torch.equal(loaded.value_bias, trained.value_bias) + + +def test_load_policy_weights_builds_an_ensemble_of_independent_copies(): + trained = [_TinyPolicy(logits=(1.0, 0.0, 0.0), value=1.0), _TinyPolicy(logits=(0.0, 1.0, 0.0), value=3.0)] + + loaded = ensemble.load_policy_weights(_TinyPolicy(), _ensemble_checkpoint(*trained)) + + assert isinstance(loaded, ensemble.EnsemblePolicy) + assert loaded.member_count == 2 + for member, source in zip(loaded.members, trained): + assert torch.equal(member.logit_bias, source.logit_bias) + with torch.no_grad(): + loaded.members[0].logit_bias.fill_(7.0) + assert torch.equal(loaded.members[1].logit_bias, trained[1].logit_bias) + _, value = loaded.forward_eval(_observations()) + assert torch.allclose(value, torch.full((BATCH_SIZE, 1), 2.0)) + + +def test_ensemble_checkpoints_of_another_format_or_a_single_member_are_refused(): + single = _ensemble_checkpoint(_TinyPolicy()) + with pytest.raises(ValueError, match="at least"): + ensemble.member_state_dicts(single) + other_format = _ensemble_checkpoint(_TinyPolicy(), _TinyPolicy()) + other_format[ensemble.ENSEMBLE_CHECKPOINT_KEY]["format_version"] = 99 + with pytest.raises(ValueError, match="format"): + ensemble.member_state_dicts(other_format) + + +def test_recurrent_members_are_refused(): + recurrent = _TinyPolicy() + recurrent.lstm = nn.LSTM(OBS_DIM, OBS_DIM) + + with pytest.raises(ValueError, match="recurrent"): + ensemble.EnsemblePolicy([recurrent, _TinyPolicy()]) + + +def test_write_ensemble_checkpoint_round_trip(tmp_path): + trained = [_TinyPolicy(logits=(1.0, 0.0, 0.0), value=1.0), _TinyPolicy(logits=(0.0, 2.0, 0.0), value=3.0)] + member_paths = [ + _write_run_dir(tmp_path, "run_1000", trained[0].state_dict()), + _write_run_dir( + tmp_path, + "run_2000", + trained[1].state_dict(), + {"policy": {**RUN_CONFIG["policy"], "fused_slot_encoder": False}}, + ), + ] + output_dir = tmp_path / "ensemble_dir" + + checkpoint_path, config_path = ensemble.write_ensemble_checkpoint(member_paths, str(output_dir), "my_ensemble") + + assert checkpoint_path == str(output_dir / "final_model.pt") and config_path == str(output_dir / "config.yaml") + checkpoint = torch.load(checkpoint_path, map_location="cpu") + assert ensemble.is_ensemble_checkpoint(checkpoint) + assert checkpoint[ensemble.ENSEMBLE_CHECKPOINT_KEY]["member_checkpoints"] == member_paths + state_dicts = ensemble.member_state_dicts(checkpoint) + assert len(state_dicts) == 2 + assert torch.equal(state_dicts[1]["logit_bias"], trained[1].logit_bias) + config = yaml.safe_load(Path(config_path).read_text()) + assert config["run_name"] == "my_ensemble" + assert config["wandb"] is False + assert config["policy"] == RUN_CONFIG["policy"] + assert config[ensemble.ENSEMBLE_CONFIG_KEY] == { + "member_checkpoints": member_paths, + "member_run_names": ["run_1000", "run_2000"], + } + loaded = ensemble.load_policy_weights(_TinyPolicy(), checkpoint) + _, value = loaded.forward_eval(_observations()) + assert torch.allclose(value, torch.full((BATCH_SIZE, 1), 2.0)) + + ensemble.write_ensemble_checkpoint(member_paths, str(output_dir), "my_ensemble") # a rebuild replaces the ensemble + + +def test_write_ensemble_checkpoint_refuses_members_that_do_not_fit_one_architecture(tmp_path): + first = _write_run_dir(tmp_path, "first", _TinyPolicy().state_dict()) + other_shape = _write_run_dir(tmp_path, "other_shape", _TinyPolicy(logits=(0.0, 0.0)).state_dict()) + other_policy = _write_run_dir( + tmp_path, + "other_policy", + _TinyPolicy().state_dict(), + {"policy": {**RUN_CONFIG["policy"], "actor_hidden_size": 16}}, + ) + other_env = _write_run_dir(tmp_path, "other_env", _TinyPolicy().state_dict(), {"env": {"dt": 0.2, "num_maps": 128}}) + recurrent = _write_run_dir(tmp_path, "recurrent", _TinyPolicy().state_dict(), {"rnn_name": "Recurrent"}) + output_dir = str(tmp_path / "out") + + with pytest.raises(ValueError, match="parameter shapes.*logit_bias"): + ensemble.write_ensemble_checkpoint([first, other_shape], output_dir, "ens") + with pytest.raises(ValueError, match="policy config"): + ensemble.write_ensemble_checkpoint([first, other_policy], output_dir, "ens") + with pytest.raises(ValueError, match=r"env config.*\['dt'\]"): + ensemble.write_ensemble_checkpoint([first, other_env], output_dir, "ens") + ensemble.write_ensemble_checkpoint([first, other_env], output_dir, "ens", allow_env_mismatch=True) + assert yaml.safe_load((tmp_path / "out" / "config.yaml").read_text())["env"]["dt"] == 0.1 + with pytest.raises(ValueError, match="rnn_name"): + ensemble.write_ensemble_checkpoint([first, recurrent], output_dir, "ens") + with pytest.raises(ValueError, match="at least"): + ensemble.write_ensemble_checkpoint([first], output_dir, "ens") + with pytest.raises(ValueError, match="repeat"): + ensemble.write_ensemble_checkpoint([first, first], output_dir, "ens") + + +def test_write_ensemble_checkpoint_never_overwrites_a_trained_checkpoint(tmp_path): + first = _write_run_dir(tmp_path, "first", _TinyPolicy().state_dict()) + second = _write_run_dir(tmp_path, "second", _TinyPolicy().state_dict()) + + with pytest.raises(FileExistsError, match="trained checkpoint"): + ensemble.write_ensemble_checkpoint([first, second], str(tmp_path / "first"), "ens") + assert torch.load(first, map_location="cpu").keys() == _TinyPolicy().state_dict().keys() diff --git a/tests/unit_tests/test_eval_wandb_reporting.py b/tests/unit_tests/test_eval_wandb_reporting.py index 28c1b41bf4..d1b457d1b3 100644 --- a/tests/unit_tests/test_eval_wandb_reporting.py +++ b/tests/unit_tests/test_eval_wandb_reporting.py @@ -1,162 +1,417 @@ -"""Unit tests for reporting standalone eval results onto the training wandb run. - -A post-training `puffer eval` runs in its own process, so to land on the run that -produced the checkpoint it must resolve that run's identity from the checkpoint's -config.yaml rather than from its own config defaults. It also must not disturb what -training already published: no config overwrite, no model re-upload, and metrics are -logged step-free so wandb appends them after training's step history. All of that is -asserted here against a stub wandb module, with no env, C-sim, GPU, or network. -""" +"""Unit tests for writing an evaluation's headline metric into the summary of the training wandb run.""" +import csv +import json +import os +import subprocess import sys import types from collections import defaultdict +from pathlib import Path import pytest import yaml -import pufferlib from pufferlib import pufferl +from pufferlib.ocean.evaluation_utils import eval_wandb_report from pufferlib.ocean.evaluation_utils import evaluation_utils as drive_benchmark -RUN_IDENTITY = { - "run_name": "k_exp_0002_1000", - "wandb_project": "nightly-multi-long", - "wandb_group": "emerge_", -} +RUN_IDENTITY = {"run_name": "k_exp_0002_1000", "wandb_project": "nightly-multi-long"} +RUN_PATH = "nightly-multi-long/k_exp_0002_1000" +TRAINING_SUMMARY = {"environment/collision_rate": 0.0123, "_step": 1000190509056, "early_stop": False} +ROUTE_COUNT = 3 +SEED_COUNT = 2 +SCENARIO_COUNT = 5 +SCENE_COUNT = 4 +STUB_WANDB_MODULE = """ +import json +import os +import sys -def _write_checkpoint_config(tmp_path, **overrides): - config = {**RUN_IDENTITY, "policy_name": "Drive", "rnn_name": None} - config.update(overrides) - config_path = tmp_path / "config.yaml" - config_path.write_text(yaml.safe_dump(config)) - return str(config_path) +class _Summary: + def update(self, values): + with open(os.environ["STUB_WANDB_SUMMARY_PATH"]) as summary_file: + stored_summary = json.load(summary_file) + stored_summary.update(values) + with open(os.environ["STUB_WANDB_SUMMARY_PATH"], "w") as summary_file: + json.dump(stored_summary, summary_file) -def _benchmark_result(**metrics_mean): - return { - "episodes": [], - "summary": {"num_scenarios": 1000, "num_episodes": 998, "metrics_mean": metrics_mean}, - } +class _Run: + state = "finished" + summary = _Summary() + + @property + def summary_metrics(self): + with open(os.environ["STUB_WANDB_SUMMARY_PATH"]) as summary_file: + return json.load(summary_file) + + +class Api: + def run(self, run_path): + assert "pufferlib" not in sys.modules + return _Run() +""" + + +class _StubSummary: + def __init__(self, stub): + self._stub = stub + + def update(self, values): + self._stub.updates.append(dict(values)) + self._stub.store_update(values) class _StubRun: - def __init__(self): - self.summary = {} - self.id = "stub_run_id" + def __init__(self, stub): + self.state = stub.run_state + self.summary_metrics = dict(stub.stored_summary) + self.summary = _StubSummary(stub) + + +class _StubApi: + def __init__(self, stub): + self._stub = stub - def log_artifact(self, artifact): - raise AssertionError("eval reporting must not upload a model artifact") + def run(self, run_path): + self._stub.requested_run_paths.append(run_path) + return _StubRun(self._stub) @pytest.fixture def stub_wandb(monkeypatch): - """Install a stub `wandb` module and record how the logger drives it.""" + """A wandb module that only has the public API: an attempt to open a session or log history fails.""" stub = types.ModuleType("wandb") - stub.run = _StubRun() - stub.init_kwargs = None - stub.logged = [] - stub.finished = False - - def init(**kwargs): - stub.init_kwargs = kwargs - - def log(metrics, step=None): - stub.logged.append((metrics, step)) - - stub.init = init - stub.log = log - stub.Settings = lambda **kwargs: kwargs - stub.finish = lambda: setattr(stub, "finished", True) + stub.run_state = "finished" + stub.stored_summary = dict(TRAINING_SUMMARY) + stub.updates = [] + stub.requested_run_paths = [] + stub.store_update = stub.stored_summary.update + stub.Api = lambda: _StubApi(stub) monkeypatch.setitem(sys.modules, "wandb", stub) + monkeypatch.setattr(eval_wandb_report, "SUMMARY_VERIFY_BACKOFF_SECONDS", 0.0) return stub @pytest.fixture -def eval_args(): - """Args shaped like load_config's return: a defaultdict whose optional top-level - keys (no_model_upload) are absent and materialize on access.""" - return defaultdict( - dict, - { - "wandb": True, - "run_name": "default_run", - "wandb_project": "pufferlib", - "wandb_group": "debug", - "tag": None, - }, - ) +def run_dir(tmp_path): + (tmp_path / "config.yaml").write_text(yaml.safe_dump({**RUN_IDENTITY, "wandb_group": "emerge_"})) + return tmp_path + + +def _write_longest6_results(run_dir, route_scores, driving_score): + """results.csv plus one route record per score, in the layout of the longest6 scripts' aggregate folder.""" + aggregate_dir = run_dir / "eval" / "carla_longest6_v1_seeds0-1_20261105_085352_2176" / "aggregate" + for score_idx, route_score in enumerate(route_scores): + seed_dir = aggregate_dir / f"s{score_idx % SEED_COUNT}" + seed_dir.mkdir(parents=True, exist_ok=True) + record = { + "route_id": f"RouteScenario_{score_idx}", + "status": "Completed", + "scores": {"score_composed": route_score}, + } + (seed_dir / f"longest_weathers_{score_idx}.json").write_text(json.dumps({"_checkpoint": {"records": [record]}})) + with open(aggregate_dir / "results.csv", "w", newline="") as results_file: + writer = csv.writer(results_file) + writer.writerows([["sep=,"], ["Avg. driving score", driving_score], ["Avg. route completion", 93.3]]) + return aggregate_dir / "results.csv" + + +def _write_nuplan_results(run_dir, closed_loop_score, scenario_count): + group_dir = run_dir / "eval" / "nuplan_val14_roadblock_20261105_085444_2173" + group_dir.mkdir(parents=True) + (group_dir / "overview.csv").write_text(f"CLS,NAC,DAC\n{closed_loop_score},0.98,0.99\n") + scenario_rows = "".join(f"token{scenario_idx},log,1.0\n" for scenario_idx in range(scenario_count)) + (group_dir / "scenarios.csv").write_text("scenario,log_name,score\n" + scenario_rows) + return group_dir / "overview.csv" + + +def _write_selfplay_summary(run_dir, distance_per_infraction_meters, num_episodes): + aggregate_dir = run_dir / "eval" / "carla_k_exp_0002_1000" / "20261105-081916" + aggregate_dir.mkdir(parents=True) + summary = { + "num_scenarios": SCENARIO_COUNT, + "num_episodes": num_episodes, + "metrics_mean": {"avg_distance_per_infraction": distance_per_infraction_meters, "collision_rate": 2.7e-05}, + } + (aggregate_dir / "evaluation_summary.json").write_text(json.dumps(summary)) + return aggregate_dir / "evaluation_summary.json" + + +def _write_alpasim_results(run_dir, score_line): + out_dir = run_dir / "eval" / "alpasim_gt_time5_r2_final_model_20261102_013615_909" + out_dir.mkdir(parents=True) + (out_dir / "results.txt").write_text(score_line + "\ncollision_any mean 0.0925 n 400\n") + return out_dir / "results.txt" + + +def _complete_result(run_dir, eval_type): + """A complete result of the eval type -> (result path, expected episode count, expected headline value).""" + if eval_type in (eval_wandb_report.EVAL_TYPE_LONGEST6_V1, eval_wandb_report.EVAL_TYPE_LONGEST6_V2): + route_scores = [100.0, 60.0, 80.0, 90.0, 70.0, 80.0] + return _write_longest6_results(run_dir, route_scores, 80.00001), ROUTE_COUNT * SEED_COUNT, 80.00001 + if eval_type in (eval_wandb_report.EVAL_TYPE_NUPLAN_REACTIVE, eval_wandb_report.EVAL_TYPE_NUPLAN_NONREACTIVE): + return _write_nuplan_results(run_dir, 0.9028013951163436, SCENARIO_COUNT), SCENARIO_COUNT, 0.9028013951163436 + if eval_type == eval_wandb_report.EVAL_TYPE_ALPASIM: + score_line = f"official_score 0.7832 scenes {SCENE_COUNT} failed 0" + return _write_alpasim_results(run_dir, score_line), SCENE_COUNT, 0.7832 + return _write_selfplay_summary(run_dir, 244218962.33697918, SCENARIO_COUNT), None, 244218.96233697918 + + +@pytest.mark.parametrize("eval_type", sorted(eval_wandb_report.HEADLINE_METRIC_KEYS)) +def test_each_eval_type_writes_its_headline_value_and_source(stub_wandb, run_dir, eval_type): + result_path, expected_episode_count, expected_value = _complete_result(run_dir, eval_type) + + eval_wandb_report.report_headline_metric(eval_type, str(result_path), str(run_dir), expected_episode_count) + + metric_key = eval_wandb_report.HEADLINE_METRIC_KEYS[eval_type] + reported_values = {metric_key: expected_value, f"{metric_key}_source": str(result_path.relative_to(run_dir))} + assert stub_wandb.requested_run_paths[0] == RUN_PATH + assert stub_wandb.updates == [reported_values] + # Training's own summary values stay exactly as they were. + assert stub_wandb.stored_summary == {**TRAINING_SUMMARY, **reported_values} + + +def test_headline_keys_are_one_per_eval_type(): + assert eval_wandb_report.HEADLINE_METRIC_KEYS == { + "longest6_v1": "final_eval/longest6_v1_DS", + "longest6_v2": "final_eval/longest6_v2_DS", + "nuplan_reactive": "final_eval/nuplan_reactive_CLS", + "nuplan_nonreactive": "final_eval/nuplan_nonreactive_CLS", + "selfplay": "final_eval/selfplay_km_per_infraction", + "alpasim": "final_eval/alpasim_score", + } -def test_identity_comes_from_checkpoint_not_eval_config(stub_wandb, eval_args): - pufferl.report_eval_to_wandb(eval_args, {"carla": _benchmark_result(collision_rate=0.11)}, RUN_IDENTITY, "eval") +def test_repeated_report_replaces_the_value_instead_of_adding_one(stub_wandb, run_dir): + first_result = _write_selfplay_summary(run_dir, 691229300.0, SCENARIO_COUNT) + eval_wandb_report.report_headline_metric(eval_wandb_report.EVAL_TYPE_SELFPLAY, str(first_result), str(run_dir)) - assert stub_wandb.init_kwargs["id"] == RUN_IDENTITY["run_name"] - assert stub_wandb.init_kwargs["name"] == RUN_IDENTITY["run_name"] - assert stub_wandb.init_kwargs["project"] == RUN_IDENTITY["wandb_project"] - assert stub_wandb.init_kwargs["group"] == RUN_IDENTITY["wandb_group"] - # A missing target run must fail loudly rather than open a stray eval-only run. - assert stub_wandb.init_kwargs["resume"] == "must" + second_result = run_dir / "eval" / "carla_k_exp_0002_1000_redlight" / "evaluation_summary.json" + second_result.parent.mkdir(parents=True) + second_result.write_text(first_result.read_text().replace("691229300.0", "200820400.0")) + eval_wandb_report.report_headline_metric(eval_wandb_report.EVAL_TYPE_SELFPLAY, str(second_result), str(run_dir)) + # One value per run and eval type, and the source names the eval the value came from. + assert stub_wandb.stored_summary == { + **TRAINING_SUMMARY, + "final_eval/selfplay_km_per_infraction": 200820.4, + "final_eval/selfplay_km_per_infraction_source": "eval/carla_k_exp_0002_1000_redlight/evaluation_summary.json", + } -def test_training_config_and_artifacts_are_left_alone(stub_wandb, eval_args): - pufferl.report_eval_to_wandb(eval_args, {"carla": _benchmark_result(collision_rate=0.11)}, RUN_IDENTITY, "eval") - assert stub_wandb.init_kwargs["config"] is None - assert stub_wandb.finished is True - # The caller's args are reused after eval returns, so they must not be rewritten. - assert eval_args["run_name"] == "default_run" - assert eval_args["wandb_project"] == "pufferlib" +def _incomplete_longest6(run_dir): + return _write_longest6_results(run_dir, [100.0, 60.0, 80.0, 90.0, 70.0], 80.0), ROUTE_COUNT * SEED_COUNT -def test_metrics_are_logged_step_free_under_the_final_eval_prefix(stub_wandb, eval_args): - results = { - "carla": _benchmark_result(collision_rate=0.11), - "nuplan_single": _benchmark_result(collision_rate=0.04, offroad_rate=0.01), - } - pufferl.report_eval_to_wandb(eval_args, results, RUN_IDENTITY, "eval") - - expected_metrics = { - "final_eval_carla/num_scenarios": 1000, - "final_eval_carla/num_episodes": 998, - "final_eval_carla/collision_rate": 0.11, - "final_eval_nuplan_single/num_scenarios": 1000, - "final_eval_nuplan_single/num_episodes": 998, - "final_eval_nuplan_single/collision_rate": 0.04, - "final_eval_nuplan_single/offroad_rate": 0.01, - } - assert stub_wandb.logged == [(expected_metrics, None)] +def _incomplete_nuplan(run_dir): + return _write_nuplan_results(run_dir, 0.9028, SCENARIO_COUNT - 1), SCENARIO_COUNT + + +def _incomplete_selfplay(run_dir): + return _write_selfplay_summary(run_dir, 244218962.3, SCENARIO_COUNT - 1), None + + +def _alpasim_with_failed_scenes(run_dir): + return _write_alpasim_results(run_dir, f"official_score 0.7832 scenes {SCENE_COUNT} failed 1"), SCENE_COUNT + + +def _alpasim_subset(run_dir): + return _write_alpasim_results(run_dir, "official_score 0.7832 scenes 2 failed 0"), SCENE_COUNT + + +def _alpasim_without_score(run_dir): + score_line = "official_score n/a (incomplete: 398 scenes finished, 2 failed)" + return _write_alpasim_results(run_dir, score_line), SCENE_COUNT + + +@pytest.mark.parametrize( + ("eval_type", "write_result", "error_match"), + [ + ("longest6_v1", _incomplete_longest6, "5 of 6 route runs"), + ("nuplan_reactive", _incomplete_nuplan, "4 of 5 scenarios"), + ("selfplay", _incomplete_selfplay, "4 of 5 scenarios"), + ("alpasim", _alpasim_with_failed_scenes, "1 failed"), + ("alpasim", _alpasim_subset, "2 of 4 scenes"), + ("alpasim", _alpasim_without_score, "no official score"), + ], +) +def test_incomplete_evals_are_refused_before_wandb_is_touched( + stub_wandb, run_dir, eval_type, write_result, error_match +): + result_path, expected_episode_count = write_result(run_dir) + + with pytest.raises(eval_wandb_report.EvalReportError, match=error_match): + eval_wandb_report.report_headline_metric(eval_type, str(result_path), str(run_dir), expected_episode_count) + assert stub_wandb.requested_run_paths == [] + assert stub_wandb.stored_summary == TRAINING_SUMMARY -def test_no_metrics_opens_no_wandb_session(stub_wandb, eval_args): - pufferl.report_eval_to_wandb(eval_args, {"carla": {"episodes": [], "summary": None}}, RUN_IDENTITY, "eval") - assert stub_wandb.init_kwargs is None - assert stub_wandb.logged == [] +def test_stale_longest6_results_csv_is_refused(stub_wandb, run_dir): + """results.csv left over from before a completion no longer matches the route records next to it.""" + result_path = _write_longest6_results(run_dir, [100.0, 60.0, 80.0, 90.0, 70.0, 80.0], 85.6) + with pytest.raises(eval_wandb_report.EvalReportError, match="stale"): + eval_wandb_report.report_headline_metric( + "longest6_v1", str(result_path), str(run_dir), ROUTE_COUNT * SEED_COUNT + ) -@pytest.mark.parametrize("missing_key", sorted(drive_benchmark.CHECKPOINT_RUN_IDENTITY_KEYS)) -def test_incomplete_checkpoint_identity_is_rejected(tmp_path, missing_key): + assert stub_wandb.updates == [] + + +def test_missing_episode_count_is_refused(stub_wandb, run_dir): + result_path, _, _ = _complete_result(run_dir, eval_wandb_report.EVAL_TYPE_NUPLAN_REACTIVE) + + with pytest.raises(eval_wandb_report.EvalReportError, match="episode count"): + eval_wandb_report.report_headline_metric("nuplan_reactive", str(result_path), str(run_dir)) + + assert stub_wandb.updates == [] + + +def test_missing_result_file_is_refused(stub_wandb, run_dir): + with pytest.raises(eval_wandb_report.EvalReportError, match="missing result file"): + eval_wandb_report.report_headline_metric("alpasim", str(run_dir / "results.txt"), str(run_dir), SCENE_COUNT) + + +def test_running_training_run_is_refused_without_a_write(stub_wandb, run_dir): + stub_wandb.run_state = "running" + result_path, _, _ = _complete_result(run_dir, eval_wandb_report.EVAL_TYPE_SELFPLAY) + + with pytest.raises(eval_wandb_report.EvalReportError, match="still running"): + eval_wandb_report.report_headline_metric(eval_wandb_report.EVAL_TYPE_SELFPLAY, str(result_path), str(run_dir)) + + assert stub_wandb.updates == [] + + +def test_summary_damaged_by_the_update_is_reported(stub_wandb, run_dir): + def store_update_and_drop_a_training_value(values): + stub_wandb.stored_summary.update(values) + del stub_wandb.stored_summary["environment/collision_rate"] + + stub_wandb.store_update = store_update_and_drop_a_training_value + result_path, _, _ = _complete_result(run_dir, eval_wandb_report.EVAL_TYPE_SELFPLAY) + + with pytest.raises(eval_wandb_report.EvalReportError, match="environment/collision_rate"): + eval_wandb_report.report_headline_metric(eval_wandb_report.EVAL_TYPE_SELFPLAY, str(result_path), str(run_dir)) + + +def test_server_float_formatting_is_not_a_mismatch(stub_wandb, run_dir): + def store_update_with_16_digits(values): + for key, value in values.items(): + stub_wandb.stored_summary[key] = float(f"{value:.16g}") if isinstance(value, float) else value + + stub_wandb.store_update = store_update_with_16_digits + result_path, _, expected_value = _complete_result(run_dir, eval_wandb_report.EVAL_TYPE_SELFPLAY) + + eval_wandb_report.report_headline_metric(eval_wandb_report.EVAL_TYPE_SELFPLAY, str(result_path), str(run_dir)) + + stored_value = stub_wandb.stored_summary["final_eval/selfplay_km_per_infraction"] + assert stored_value != expected_value + assert stored_value == pytest.approx(expected_value, rel=1e-12) + + +@pytest.mark.parametrize("missing_key", sorted(eval_wandb_report.RUN_IDENTITY_KEYS)) +def test_incomplete_run_identity_is_rejected(tmp_path, missing_key): config = {key: value for key, value in RUN_IDENTITY.items() if key != missing_key} - config_path = tmp_path / "config.yaml" - config_path.write_text(yaml.safe_dump(config)) + (tmp_path / "config.yaml").write_text(yaml.safe_dump(config)) + + with pytest.raises(eval_wandb_report.EvalReportError, match=missing_key): + eval_wandb_report.load_run_identity(str(tmp_path / "config.yaml")) + + +def test_blank_run_identity_is_rejected(tmp_path): + (tmp_path / "config.yaml").write_text(yaml.safe_dump({**RUN_IDENTITY, "run_name": " "})) + + with pytest.raises(eval_wandb_report.EvalReportError, match="run_name"): + eval_wandb_report.load_run_identity(str(tmp_path / "config.yaml")) + + +def test_command_line_reports_without_importing_pufferlib(tmp_path, run_dir): + """The eval jobs run the file by path in interpreters that must not pay for, or break on, the training stack.""" + stub_package_dir = tmp_path / "stub_site" + stub_package_dir.mkdir() + (stub_package_dir / "wandb.py").write_text(STUB_WANDB_MODULE) + stored_summary_path = tmp_path / "stored_summary.json" + stored_summary_path.write_text(json.dumps(TRAINING_SUMMARY)) + result_path, expected_episode_count, expected_value = _complete_result(run_dir, "nuplan_nonreactive") + + completed = subprocess.run( + [ + sys.executable, + str(Path(eval_wandb_report.__file__)), + "nuplan_nonreactive", + str(result_path), + "--run-dir", + str(run_dir), + "--expected-episode-count", + str(expected_episode_count), + ], + env={**os.environ, "PYTHONPATH": str(stub_package_dir), "STUB_WANDB_SUMMARY_PATH": str(stored_summary_path)}, + capture_output=True, + text=True, + check=False, + ) - with pytest.raises(pufferlib.APIUsageError, match=missing_key): - drive_benchmark.load_checkpoint_run_identity(str(config_path)) + assert completed.returncode == 0, completed.stderr + assert json.loads(stored_summary_path.read_text()) == { + **TRAINING_SUMMARY, + "final_eval/nuplan_nonreactive_CLS": expected_value, + "final_eval/nuplan_nonreactive_CLS_source": str(result_path.relative_to(run_dir)), + } -def test_blank_checkpoint_identity_is_rejected(tmp_path): - config_path = _write_checkpoint_config(tmp_path, wandb_group=" ") +def test_command_line_exits_with_the_reason_for_a_refused_report(run_dir): + result_path, expected_episode_count = _incomplete_nuplan(run_dir) + + completed = subprocess.run( + [ + sys.executable, + str(Path(eval_wandb_report.__file__)), + "nuplan_reactive", + str(result_path), + "--run-dir", + str(run_dir), + "--expected-episode-count", + str(expected_episode_count), + ], + capture_output=True, + text=True, + check=False, + ) - with pytest.raises(pufferlib.APIUsageError, match="wandb_group"): - drive_benchmark.load_checkpoint_run_identity(config_path) + assert completed.returncode == 1 + assert "eval_wandb_report: incomplete nuPlan eval: 4 of 5 scenarios" in completed.stderr + + +def test_training_logger_opens_the_run_with_its_config(monkeypatch): + """Args shaped like load_config's return: optional top-level keys (no_model_upload) materialize on access.""" + stub = types.ModuleType("wandb") + stub.run = types.SimpleNamespace(id="stub_run_id") + stub.init_kwargs = None + stub.init = lambda **kwargs: setattr(stub, "init_kwargs", kwargs) + stub.Settings = lambda **kwargs: kwargs + monkeypatch.setitem(sys.modules, "wandb", stub) + args = defaultdict(dict, {**RUN_IDENTITY, "wandb_group": "emerge_", "tag": None}) + + logger = pufferl.WandbLogger(args) + + assert logger.run_id == "stub_run_id" + assert stub.init_kwargs["id"] == RUN_IDENTITY["run_name"] + assert stub.init_kwargs["name"] == RUN_IDENTITY["run_name"] + assert stub.init_kwargs["project"] == RUN_IDENTITY["wandb_project"] + assert stub.init_kwargs["group"] == "emerge_" + assert stub.init_kwargs["resume"] == "allow" + assert stub.init_kwargs["config"] is args def test_summarize_benchmark_metrics_skips_benchmarks_without_a_summary(): + summary = {"num_scenarios": 1000, "num_episodes": 998, "metrics_mean": {"collision_rate": 0.11}} results = { - "carla": _benchmark_result(collision_rate=0.11), + "carla": {"episodes": [], "summary": summary}, "nuplan_single": {"episodes": [], "summary": None}, } @@ -164,3 +419,23 @@ def test_summarize_benchmark_metrics_skips_benchmarks_without_a_summary(): assert set(metrics) == {"eval_carla/num_scenarios", "eval_carla/num_episodes", "eval_carla/collision_rate"} assert drive_benchmark.summarize_benchmark_metrics({}, "eval_") == {} + + +def test_tidied_longest6_run_reads_its_flat_route_records(tmp_path): + """After tidy_carla_run.py results.csv sits at the run's top and the records are routes/*.result.json; other + json files under the run are not records.""" + run = tmp_path / "carla_longest6_v2_seeds0-1_20261109_120000_7" + (run / "routes").mkdir(parents=True) + (run / "details").mkdir() + (run / "details" / "not_a_record.json").write_text("{}") + scores = [100.0, 60.0, 80.0] + for idx, score in enumerate(scores): + record = {"route_id": f"RouteScenario_{idx}", "status": "Completed", "scores": {"score_composed": score}} + record_path = run / "routes" / f"route_{idx:02d}_s0.result.json" + record_path.write_text(json.dumps({"_checkpoint": {"records": [record]}})) + with open(run / "results.csv", "w", newline="") as results_file: + csv.writer(results_file).writerows([["sep=,"], ["Avg. driving score", 80.0]]) + + assert eval_wandb_report.read_longest6_driving_score(str(run / "results.csv"), len(scores)) == 80.0 + with pytest.raises(eval_wandb_report.EvalReportError, match="incomplete"): + eval_wandb_report.read_longest6_driving_score(str(run / "results.csv"), len(scores) + 1) diff --git a/tests/unit_tests/test_light_timing.py b/tests/unit_tests/test_light_timing.py new file mode 100644 index 0000000000..d9a9ccc2da --- /dev/null +++ b/tests/unit_tests/test_light_timing.py @@ -0,0 +1,34 @@ +"""YellowCountdown (co-sim remaining-yellow fallback) and the CARLA phase-timer helper.""" + +import pytest + +from pufferlib.ocean.cosim import carla_bridge as cb +from pufferlib.ocean.cosim.light_timing import YellowCountdown + + +def test_countdown_runs_from_onset_and_floors_at_one_step(): + countdown = YellowCountdown(assumed_total_s=2.0, min_remaining_s=0.1) + assert countdown.tick({7}, 10.0) == {7: 2.0} + assert countdown.tick({7}, 10.5) == {7: pytest.approx(1.5)} + assert countdown.tick({7}, 12.5) == {7: pytest.approx(0.1)} # past the assumed total: one step stays + assert countdown.tick(set(), 12.6) == {} + assert countdown.tick({7}, 20.0) == {7: 2.0} # restarts after a non-yellow tick + + +def test_countdown_tracks_lights_independently(): + countdown = YellowCountdown(2.0, 0.1) + countdown.tick({1}, 0.0) + assert countdown.tick({1, 2}, 1.0) == {1: pytest.approx(1.0), 2: 2.0} + + +def test_countdown_rejects_nonpositive_durations(): + with pytest.raises(ValueError): + YellowCountdown(0.0, 0.1) + with pytest.raises(ValueError): + YellowCountdown(2.0, 0.0) + + +def test_carla_timer_gives_remaining_or_none_when_frozen(): + assert cb.carla_yellow_remaining_s(elapsed_s=0.5, yellow_time_s=3.1) == pytest.approx(2.6) + assert cb.carla_yellow_remaining_s(elapsed_s=3.3, yellow_time_s=3.1) == 0.0 + assert cb.carla_yellow_remaining_s(elapsed_s=0.5, yellow_time_s=1e9) is None diff --git a/tests/unit_tests/test_map_load_failure.py b/tests/unit_tests/test_map_load_failure.py new file mode 100644 index 0000000000..d00eabe7da --- /dev/null +++ b/tests/unit_tests/test_map_load_failure.py @@ -0,0 +1,36 @@ +"""A map bin that cannot be loaded fails env creation with an exception instead of crashing on first use.""" + +import pytest + +from pufferlib.ocean.drive.drive import Drive + +SOURCE_BIN = "pufferlib/resources/drive/binaries/carla/opendrive__Town01.bin" +AGENTS_PER_ENV = 8 + + +def _kwargs(map_dir, simulation_mode): + return { + "num_agents": AGENTS_PER_ENV, + "min_agents_per_env": AGENTS_PER_ENV, + "max_agents_per_env": AGENTS_PER_ENV, + "num_maps": 1, + "allow_map_subset": True, + "map_dir": str(map_dir), + "simulation_mode": simulation_mode, + "action_type": "continuous", + "dynamics_model": "jerk", + "scenario_length": 50, + "resample_frequency": 0, + "termination_mode": 0, + "stagger_first_episode": False, + "seed": 3, + } + + +@pytest.mark.parametrize("simulation_mode", ["gigaflow", "replay"]) +def test_truncated_map_bin_raises(tmp_path, simulation_mode): + with open(SOURCE_BIN, "rb") as source: + map_bytes = source.read() + (tmp_path / "truncated.bin").write_bytes(map_bytes[: len(map_bytes) // 2]) + with pytest.raises(RuntimeError): + Drive(**_kwargs(tmp_path, simulation_mode)) diff --git a/tests/unit_tests/test_nuplan_bicycle_command.py b/tests/unit_tests/test_nuplan_bicycle_command.py new file mode 100644 index 0000000000..b766eceb68 --- /dev/null +++ b/tests/unit_tests/test_nuplan_bicycle_command.py @@ -0,0 +1,269 @@ +"""Shadow-env accel/steering targets -> CaRL ActionTrajectory commands for nuPlan's KinematicBicycleModel.""" + +import math +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from pufferlib.ocean.cosim.nuplan_bicycle import ( + BicycleCommandConfig, + bicycle_commands, + bicycle_response, + center_accel_shift, + floor_clamped_intent_accel, + low_pass_gain, + low_pass_response, +) + +PACIFICA_WHEEL_BASE_M = 3.089 +PACIFICA_REAR_AXLE_TO_CENTER_M = 1.461 +ACCEL_CAP_MPS2 = 2.304 + +DT_S = 0.1 + + +def devkit_step(cfg, accel, steering, speed, accel_cmd, steering_rate_cmd): + """KinematicBicycleModel._update_commands + forward Euler on the longitudinal state, written out.""" + accel_gain = low_pass_gain(DT_S, cfg.accel_time_constant_s) + steering_gain = low_pass_gain(DT_S, cfg.steering_time_constant_s) + accel_next = accel_gain * (accel_cmd - accel) + accel + steering_next = steering_gain * (DT_S * steering_rate_cmd + steering - steering) + steering + return accel_next, steering_next, speed + accel_next * DT_S + + +def test_undo_low_pass_lands_on_the_targets_after_one_devkit_step(): + cfg = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=-2.0, undo_low_pass=True) + accel_cmd, rate_cmd = bicycle_commands(cfg, 1.5, -0.2, 0.3, 0.1, 5.0) + accel_next, steering_next, _ = devkit_step(cfg, 0.3, 0.1, 5.0, accel_cmd, rate_cmd) + assert accel_next == pytest.approx(1.5) + assert steering_next == pytest.approx(-0.2) + + +def test_without_undo_the_model_realises_the_lag_fraction_of_the_step(): + cfg = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=-2.0, undo_low_pass=False) + accel_cmd, rate_cmd = bicycle_commands(cfg, 1.5, -0.2, 0.3, 0.1, 5.0) + assert accel_cmd == pytest.approx(1.5) + assert rate_cmd == pytest.approx((-0.2 - 0.1) / DT_S) + accel_next, steering_next, _ = devkit_step(cfg, 0.3, 0.1, 5.0, accel_cmd, rate_cmd) + assert accel_next == pytest.approx(0.3 + (1.5 - 0.3) / 3.0) + assert steering_next == pytest.approx(0.1 + (-0.2 - 0.1) * 2.0 / 3.0) + + +@pytest.mark.parametrize("undo_low_pass", [False, True]) +@pytest.mark.parametrize("min_speed_mps, speed_mps", [(0.0, 0.05), (-2.0, -1.95)]) +def test_speed_floor_binds_on_the_realised_acceleration(undo_low_pass, min_speed_mps, speed_mps): + cfg = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=min_speed_mps, undo_low_pass=undo_low_pass) + accel_cmd, _ = bicycle_commands(cfg, -3.0, 0.0, -1.0, 0.0, speed_mps) + _, _, speed_next = devkit_step(cfg, -1.0, 0.0, speed_mps, accel_cmd, 0.0) + assert speed_next == pytest.approx(min_speed_mps) + + +def test_floor_zero_holds_a_standstill_against_a_braking_target(): + cfg = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=0.0, undo_low_pass=False) + accel_cmd, _ = bicycle_commands(cfg, -2.0, 0.0, -0.5, 0.0, 0.0) + accel_next, _, speed_next = devkit_step(cfg, -0.5, 0.0, 0.0, accel_cmd, 0.0) + assert accel_next == pytest.approx(0.0) + assert speed_next == pytest.approx(0.0) + + +@pytest.mark.parametrize("undo_low_pass", [False, True]) +@pytest.mark.parametrize("min_speed_mps", [0.0, -2.0]) +def test_shadow_stop_snap_brings_the_model_to_exactly_zero_speed(undo_low_pass, min_speed_mps): + cfg = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=min_speed_mps, undo_low_pass=undo_low_pass) + accel_cmd, _ = bicycle_commands(cfg, 0.0, 0.0, -1.8, 0.0, 0.15, stop_this_step=True) + _, _, speed_next = devkit_step(cfg, -1.8, 0.0, 0.15, accel_cmd, 0.0) + assert speed_next == pytest.approx(0.0, abs=1e-12) + + +def test_floor_leaves_a_non_crossing_command_alone(): + cfg = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=0.0, undo_low_pass=False) + accel_cmd, _ = bicycle_commands(cfg, -3.0, 0.0, -1.0, 0.0, 8.0) + assert accel_cmd == pytest.approx(-3.0) + + +def test_response_matches_the_devkit_step_for_the_actual_elapsed_time(): + cfg = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=-2.0, undo_low_pass=True) + accel_cmd, _ = bicycle_commands(cfg, 1.0, 0.0, 0.0, 0.0, 3.0) + accel_next, _, speed_next = devkit_step(cfg, 0.0, 0.0, 3.0, accel_cmd, 0.0) + assert bicycle_response(cfg, 0.0, accel_cmd, 3.0, DT_S) == pytest.approx((accel_next, speed_next)) + + +@pytest.mark.parametrize( + "kwargs", + [ + dict(dt_s=0.0, min_speed_mps=0.0, undo_low_pass=False), + dict(dt_s=math.nan, min_speed_mps=0.0, undo_low_pass=False), + dict(dt_s=DT_S, min_speed_mps=0.5, undo_low_pass=False), + dict(dt_s=DT_S, min_speed_mps=0.0, undo_low_pass=False, accel_time_constant_s=-0.1), + ], +) +def test_config_rejects_invalid_values(kwargs): + with pytest.raises(ValueError): + BicycleCommandConfig(**kwargs) + + +def _center_cap_cfg(undo_low_pass): + return BicycleCommandConfig( + dt_s=DT_S, + min_speed_mps=0.0, + undo_low_pass=undo_low_pass, + center_accel_cap_mps2=ACCEL_CAP_MPS2, + wheel_base_m=PACIFICA_WHEEL_BASE_M, + rear_axle_to_center_m=PACIFICA_REAR_AXLE_TO_CENTER_M, + ) + + +@pytest.mark.parametrize("undo_low_pass", [False, True]) +def test_center_accel_cap_binds_only_when_the_yaw_terms_push_the_center_over_the_cap(undo_low_pass): + cfg = _center_cap_cfg(undo_low_pass) + accel_gain = low_pass_gain(DT_S, cfg.accel_time_constant_s) + speed, accel = 0.5, 2.0 + straight, _ = bicycle_commands(cfg, ACCEL_CAP_MPS2, 0.0, accel, 0.0, speed) + unlimited = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=0.0, undo_low_pass=undo_low_pass) + assert straight == pytest.approx(bicycle_commands(unlimited, ACCEL_CAP_MPS2, 0.0, accel, 0.0, speed)[0]) + left_turn, rate = bicycle_commands(cfg, ACCEL_CAP_MPS2, 0.3, accel, 0.3, speed, yaw_rate_radps=0.0) + realised = low_pass_response(accel, left_turn, accel_gain) + assert realised < low_pass_response(accel, straight, accel_gain) + steering_next = low_pass_response(0.3, 0.3 + rate * DT_S, low_pass_gain(DT_S, cfg.steering_time_constant_s)) + assert realised + center_accel_shift(cfg, realised, speed, steering_next, 0.0) <= ACCEL_CAP_MPS2 + 1e-9 + right_turn, _ = bicycle_commands(cfg, ACCEL_CAP_MPS2, -0.3, accel, -0.3, speed, yaw_rate_radps=0.0) + assert right_turn == pytest.approx(straight) + + +def _yaw_rate(speed, steering): + return speed * math.tan(steering) / PACIFICA_WHEEL_BASE_M + + +@pytest.mark.parametrize("undo_low_pass", [False, True]) +def test_yaw_jerk_bound_spreads_a_steering_reversal_over_several_steps(undo_low_pass): + max_yaw_jerk = 1.5 + cfg = BicycleCommandConfig( + dt_s=DT_S, + min_speed_mps=0.0, + undo_low_pass=undo_low_pass, + max_yaw_jerk_radps3=max_yaw_jerk, + wheel_base_m=PACIFICA_WHEEL_BASE_M, + ) + steering_gain = low_pass_gain(DT_S, cfg.steering_time_constant_s) + speed, steering = 3.0, 0.2 + yaw_rate_prev = _yaw_rate(speed, steering - 0.06) # the wheel was moving left at the 0.6 rad/s limit + yaw_rate = _yaw_rate(speed, steering) + _, rate_cmd = bicycle_commands( + cfg, 0.0, steering - 0.06, 0.0, steering, speed, yaw_rate_radps=yaw_rate, yaw_rate_prev_radps=yaw_rate_prev + ) + steering_next = low_pass_response(steering, steering + rate_cmd * DT_S, steering_gain) + yaw_accel_change = (_yaw_rate(speed, steering_next) - 2 * yaw_rate + yaw_rate_prev) / DT_S / DT_S + assert abs(yaw_accel_change) <= max_yaw_jerk + 1e-9 + assert steering_next > steering - 0.06 # the full reversal was not allowed in one step + unbounded = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=0.0, undo_low_pass=undo_low_pass) + gentle, _ = bicycle_commands(unbounded, 0.0, steering + 0.005, 0.0, steering, speed, yaw_rate, yaw_rate) + bounded, _ = bicycle_commands(cfg, 0.0, steering + 0.005, 0.0, steering, speed, yaw_rate, yaw_rate) + assert bounded == pytest.approx(gentle) + + +def test_yaw_jerk_bound_is_inactive_at_standstill_and_on_the_first_step(): + cfg = BicycleCommandConfig( + dt_s=DT_S, min_speed_mps=0.0, undo_low_pass=True, max_yaw_jerk_radps3=1.5, wheel_base_m=PACIFICA_WHEEL_BASE_M + ) + free = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=0.0, undo_low_pass=True) + assert bicycle_commands(cfg, 0.0, -0.3, 0.0, 0.3, 0.0, 0.0, 0.0) == pytest.approx( + bicycle_commands(free, 0.0, -0.3, 0.0, 0.3, 0.0, 0.0, 0.0) + ) + assert bicycle_commands(cfg, 0.0, -0.3, 0.0, 0.3, 5.0, 0.0, None) == pytest.approx( + bicycle_commands(free, 0.0, -0.3, 0.0, 0.3, 5.0, 0.0, None) + ) + + +@pytest.mark.parametrize("min_speed_mps", [0.0, -2.0]) +def test_intent_accel_is_clamped_to_zero_only_while_braking_on_the_floor(min_speed_mps): + cfg = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=min_speed_mps, undo_low_pass=True) + assert floor_clamped_intent_accel(cfg, min_speed_mps, -3.0) == 0.0 + assert floor_clamped_intent_accel(cfg, min_speed_mps + 5e-7, -3.0) == 0.0 + assert floor_clamped_intent_accel(cfg, min_speed_mps, 1.2) == 1.2 + assert floor_clamped_intent_accel(cfg, min_speed_mps + 0.5, -3.0) == -3.0 + + +def test_center_accel_cap_config_needs_the_vehicle_geometry(): + with pytest.raises(ValueError): + BicycleCommandConfig(dt_s=DT_S, min_speed_mps=0.0, undo_low_pass=False, center_accel_cap_mps2=2.0) + + +def test_center_accel_cap_holds_in_the_devkit_kinematic_bicycle_model(): + pytest.importorskip("nuplan") + from nuplan.common.actor_state.dynamic_car_state import DynamicCarState + from nuplan.common.actor_state.ego_state import EgoState + from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimePoint + from nuplan.common.actor_state.vehicle_parameters import get_pacifica_parameters + from nuplan.planning.simulation.controller.motion_model.kinematic_bicycle import KinematicBicycleModel + + params = get_pacifica_parameters() + cfg = BicycleCommandConfig( + dt_s=DT_S, + min_speed_mps=0.0, + undo_low_pass=True, + center_accel_cap_mps2=ACCEL_CAP_MPS2, + wheel_base_m=params.wheel_base, + rear_axle_to_center_m=params.rear_axle_to_center, + ) + model = KinematicBicycleModel(params) + speed, accel, steering = 0.5, 2.0, 0.3 + yaw_rate = speed * math.tan(steering) / params.wheel_base + state = EgoState.build_from_rear_axle( + rear_axle_pose=StateSE2(0.0, 0.0, 0.0), + rear_axle_velocity_2d=StateVector2D(speed, 0.0), + rear_axle_acceleration_2d=StateVector2D(accel, 0.0), + tire_steering_angle=steering, + time_point=TimePoint(0), + vehicle_parameters=params, + angular_vel=yaw_rate, + ) + accel_cmd, rate_cmd = bicycle_commands( + cfg, ACCEL_CAP_MPS2, steering, accel, steering, speed, yaw_rate_radps=yaw_rate + ) + command = DynamicCarState.build_from_rear_axle( + rear_axle_to_center_dist=state.car_footprint.rear_axle_to_center_dist, + rear_axle_velocity_2d=state.dynamic_car_state.rear_axle_velocity_2d, + rear_axle_acceleration_2d=StateVector2D(accel_cmd, 0.0), + tire_steering_rate=rate_cmd, + ) + propagated = model.propagate_state(state, command, TimePoint(int(DT_S * 1e6))) + assert propagated.dynamic_car_state.rear_axle_acceleration_2d.x < ACCEL_CAP_MPS2 + center_accel = propagated.dynamic_car_state.center_acceleration_2d.x + assert ACCEL_CAP_MPS2 - 0.1 < center_accel <= ACCEL_CAP_MPS2 + 1e-9 # one-sided clip: under the cap, close to it + + +def test_commands_drive_the_devkit_kinematic_bicycle_model_onto_the_targets(): + pytest.importorskip("nuplan") + from nuplan.common.actor_state.dynamic_car_state import DynamicCarState + from nuplan.common.actor_state.ego_state import EgoState + from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimePoint + from nuplan.common.actor_state.vehicle_parameters import get_pacifica_parameters + from nuplan.planning.simulation.controller.motion_model.kinematic_bicycle import KinematicBicycleModel + + params = get_pacifica_parameters() + model = KinematicBicycleModel(params) + accel, steering, speed = 0.3, 0.1, 5.0 + state = EgoState.build_from_rear_axle( + rear_axle_pose=StateSE2(10.0, 20.0, 0.5), + rear_axle_velocity_2d=StateVector2D(speed, 0.0), + rear_axle_acceleration_2d=StateVector2D(accel, 0.0), + tire_steering_angle=steering, + time_point=TimePoint(0), + vehicle_parameters=params, + ) + cfg = BicycleCommandConfig(dt_s=DT_S, min_speed_mps=-2.0, undo_low_pass=True) + accel_cmd, rate_cmd = bicycle_commands(cfg, 1.5, -0.2, accel, steering, speed) + command = DynamicCarState.build_from_rear_axle( + rear_axle_to_center_dist=state.car_footprint.rear_axle_to_center_dist, + rear_axle_velocity_2d=state.dynamic_car_state.rear_axle_velocity_2d, + rear_axle_acceleration_2d=StateVector2D(accel_cmd, 0.0), + tire_steering_rate=rate_cmd, + ) + propagated = model.propagate_state(state, command, TimePoint(int(DT_S * 1e6))) + assert propagated.dynamic_car_state.rear_axle_acceleration_2d.x == pytest.approx(1.5) + assert propagated.tire_steering_angle == pytest.approx(-0.2) + expected = bicycle_response(cfg, accel, accel_cmd, speed, DT_S) + assert propagated.dynamic_car_state.rear_axle_velocity_2d.x == pytest.approx(expected[1]) diff --git a/tests/unit_tests/test_nuplan_bin_symmetric_widths.py b/tests/unit_tests/test_nuplan_bin_symmetric_widths.py new file mode 100644 index 0000000000..6b84f7d2ee --- /dev/null +++ b/tests/unit_tests/test_nuplan_bin_symmetric_widths.py @@ -0,0 +1,60 @@ +"""symmetric_lane_widths: bin lane widths become twice the distance to the nearer nuPlan boundary.""" + +import numpy as np +import pytest + +shapely = pytest.importorskip("shapely") +from shapely.geometry import LineString # noqa: E402 + +from pufferlib.ocean.cosim import nuplan_bridge as nb # noqa: E402 + + +class _Polyline: + def __init__(self, coords): + self.linestring = LineString(coords) + + +class _MapLane: + def __init__(self, baseline, left, right): + self.baseline_path = _Polyline(baseline) + self.left_boundary = _Polyline(left) + self.right_boundary = _Polyline(right) + + +def _lane_road(tf, utm_points, widths): + xs, ys = zip(*(tf.loc_to_bin(x, y) for x, y in utm_points)) + return {"type": 2, "S": len(utm_points), "x": tuple(xs), "y": tuple(ys), "widths": tuple(widths)} + + +def test_parking_strip_lane_shrinks_to_the_narrow_side_and_centered_lane_keeps_its_width(): + tf = nb.NuPlanTransform(1000.0, 2000.0) + parking = _MapLane([(0, 0), (20, 0)], [(0, 1.7), (20, 1.7)], [(0, -4.3), (20, -4.3)]) + centered = _MapLane([(0, 10), (20, 10)], [(0, 11.5), (20, 11.5)], [(0, 8.5), (20, 8.5)]) + roads = [ + _lane_road(tf, [(0, 0), (10, 0), (20, 0)], [6.0, 6.0, 6.0]), + _lane_road(tf, [(0, 10), (10, 10), (20, 10)], [3.0, 3.0, 3.0]), + {"type": 15, "S": 2, "x": (0.0, 1.0), "y": (0.0, 0.0)}, + ] + patched, unmatched = nb.symmetric_lane_widths(roads, tf, [parking, centered]) + assert (patched, unmatched) == (2, 0) + np.testing.assert_allclose(roads[0]["widths"], [3.4, 3.4, 3.4]) + np.testing.assert_allclose(roads[1]["widths"], [3.0, 3.0, 3.0]) + assert "widths" not in roads[2] + + +def test_lane_without_a_nearby_baseline_keeps_its_widths(): + tf = nb.NuPlanTransform(0.0, 0.0) + lane = _MapLane([(0, 0), (20, 0)], [(0, 1.7), (20, 1.7)], [(0, -4.3), (20, -4.3)]) + roads = [_lane_road(tf, [(0, 50), (20, 50)], [5.5, 5.5]), _lane_road(tf, [(0, 1.0), (20, 1.0)], [6.0, 6.0])] + patched, unmatched = nb.symmetric_lane_widths(roads, tf, [lane]) + assert (patched, unmatched) == (0, 2) # 50 m away: no candidate; 1 m off the baseline: beyond the match tolerance + assert roads[0]["widths"] == (5.5, 5.5) and roads[1]["widths"] == (6.0, 6.0) + + +def test_width_is_floored_where_the_baseline_touches_a_boundary(): + tf = nb.NuPlanTransform(0.0, 0.0) + tapered = _MapLane([(0, 0), (20, 0)], [(0, 0.1), (10, 0.1), (10, 1.7), (20, 1.7)], [(0, -3.0), (20, -3.0)]) + roads = [_lane_road(tf, [(0, 0), (20, 0)], [4.0, 4.0])] + nb.symmetric_lane_widths(roads, tf, [tapered]) + assert roads[0]["widths"][0] == nb.SYMMETRIC_WIDTH_MIN_M + assert roads[0]["widths"][1] == pytest.approx(3.4) diff --git a/tests/unit_tests/test_nuplan_partner_filter.py b/tests/unit_tests/test_nuplan_partner_filter.py index 4539a5e0ee..d8bc5282c5 100644 --- a/tests/unit_tests/test_nuplan_partner_filter.py +++ b/tests/unit_tests/test_nuplan_partner_filter.py @@ -1,4 +1,4 @@ -"""Static nuPlan clutter (cones, poles, barriers) reaches the shadow env only when it stands inside a lane +"""Static nuPlan clutter (cones, poles, barriers) reaches the shadow env only when its box overlaps a lane or lane-connector polygon; intersection corners and crosswalks (drivable area) do not count.""" from enum import Enum @@ -7,6 +7,7 @@ import pytest nuplan = pytest.importorskip("nuplan") # noqa: F841 (partner_tracked_objects imports nuPlan datatypes) +from shapely.geometry import LineString, box # noqa: E402 from pufferlib.ocean.cosim import nuplan_bridge as nb # noqa: E402 @@ -16,31 +17,60 @@ class FakeType(Enum): PEDESTRIAN = 1 BICYCLE = 2 TRAFFIC_CONE = 3 + BARRIER = 4 GENERIC_OBJECT = 6 EGO = 7 -def _obj(type_name, token, x=0.0): +def _obj(type_name, token, x=0.0, y=0.0, length=0.4, width=0.4): + """Axis-aligned box: length along x, width along y.""" + geometry = box(x - length / 2, y - width / 2, x + length / 2, y + width / 2) return SimpleNamespace( - tracked_object_type=FakeType[type_name], track_token=token, center=SimpleNamespace(x=x, y=0.0) + tracked_object_type=FakeType[type_name], + track_token=token, + center=SimpleNamespace(x=x, y=y), + box=SimpleNamespace(length=length, width=width, geometry=geometry), ) +class ConnectorWithoutPolygon: + """Devkit lane connector whose scaled-width polygon row is missing: .polygon asserts.""" + + def __init__(self, centerline): + self.centerline = centerline + + @property + def polygon(self): + raise AssertionError("Could not find the desired key") + + class FakeMap: - def __init__(self, lane_x, connector_x=()): - self.lane_x = lane_x - self.connector_x = set(connector_x) + """Devkit proximity semantics: the square query patch is matched against lane polygons but against lane + connector centerlines.""" + + def __init__(self, lane_polygons, connectors=()): + self.lanes = [SimpleNamespace(polygon=p) for p in lane_polygons] + self.connectors = [ + c if isinstance(c, ConnectorWithoutPolygon) else SimpleNamespace(polygon=c[0], centerline=c[1]) + for c in connectors + ] self.queries = 0 - def is_in_layer(self, point, layer): + def get_proximal_map_objects(self, point, radius, layers): self.queries += 1 - assert layer.name == "LANE", f"unexpected layer {layer}" - return point.x in self.lane_x + patch = box(point.x - radius, point.y - radius, point.x + radius, point.y + radius) + result = {} + for layer in layers: + if layer.name == "LANE": + result[layer] = [lane for lane in self.lanes if lane.polygon.intersects(patch)] + else: + result[layer] = [c for c in self.connectors if c.centerline.intersects(patch)] + return result - def get_all_map_objects(self, point, layer): - self.queries += 1 - assert layer.name == "LANE_CONNECTOR", f"unexpected layer {layer}" - return ["connector"] if point.x in self.connector_x else [] + +LANE = box(-10.0, -2.0, 0.0, 2.0) # right lane edge at x = 0 +CONNECTOR = box(19.0, -1.0, 21.0, 1.0) +CONNECTOR_CENTERLINE = LineString([(19.0, 0.0), (21.0, 0.0)]) def test_moving_agents_kept_static_off_road_dropped(): @@ -48,25 +78,53 @@ def test_moving_agents_kept_static_off_road_dropped(): _obj("VEHICLE", "v1", 50.0), _obj("PEDESTRIAN", "p1", 60.0), _obj("BICYCLE", "b1", 70.0), - _obj("TRAFFIC_CONE", "c_on_lane", 1.0), - _obj("TRAFFIC_CONE", "c_off", 2.0), - _obj("GENERIC_OBJECT", "g_off", 3.0), - _obj("GENERIC_OBJECT", "g_on_connector", 4.0), + _obj("TRAFFIC_CONE", "c_on_lane", -5.0), + _obj("GENERIC_OBJECT", "g_protruding", 1.0, length=2.3, width=1.7), # center 1 m off the lane, box reaches in + _obj("TRAFFIC_CONE", "c_off", 1.0), + _obj("BARRIER", "bar_end_reaches", 1.9, length=4.0, width=0.5), # only the far end overlaps the lane + _obj("GENERIC_OBJECT", "g_on_connector", 20.0), _obj("EGO", "ego", 0.0), ] - map_api = FakeMap(lane_x={1.0}, connector_x={4.0}) + map_api = FakeMap([LANE], [(CONNECTOR, CONNECTOR_CENTERLINE)]) cache = {} kept = nb.partner_tracked_objects(objs, map_api, cache) - assert [o.track_token for o in kept] == ["v1", "p1", "b1", "c_on_lane", "g_on_connector"] - assert cache == {"c_on_lane": True, "c_off": False, "g_off": False, "g_on_connector": True} - assert map_api.queries == 7 # lane hit: 1 query; connector hit or miss: 2 + assert [o.track_token for o in kept] == [ + "v1", + "p1", + "b1", + "c_on_lane", + "g_protruding", + "bar_end_reaches", + "g_on_connector", + ] + assert cache == { + "c_on_lane": True, + "g_protruding": True, + "c_off": False, + "bar_end_reaches": True, + "g_on_connector": True, + } + assert map_api.queries == 5 def test_static_verdict_is_cached_per_track(): - objs = [_obj("TRAFFIC_CONE", "c1", 1.0), _obj("GENERIC_OBJECT", "g1", 9.0)] - map_api = FakeMap(lane_x={1.0}) + objs = [_obj("TRAFFIC_CONE", "c1", -1.0), _obj("GENERIC_OBJECT", "g1", 9.0)] + map_api = FakeMap([LANE]) cache = {} nb.partner_tracked_objects(objs, map_api, cache) nb.partner_tracked_objects(objs, map_api, cache) - assert map_api.queries == 3 + assert map_api.queries == 2 assert [o.track_token for o in nb.partner_tracked_objects(objs, map_api, cache)] == ["c1"] + + +def test_static_inside_a_connector_away_from_its_centerline_is_kept(): + """A 0.3 m patch around a cone 0.8 m off the centerline never reached the connector (run 390 regression).""" + objs = [_obj("TRAFFIC_CONE", "c_edge", 20.0, 0.8), _obj("TRAFFIC_CONE", "c_outside", 20.0, 1.5)] + kept = nb.partner_tracked_objects(objs, FakeMap([], [(CONNECTOR, CONNECTOR_CENTERLINE)]), {}) + assert [o.track_token for o in kept] == ["c_edge"] + + +def test_connector_without_polygon_row_is_skipped(): + objs = [_obj("TRAFFIC_CONE", "c1", 30.0)] + map_api = FakeMap([], [ConnectorWithoutPolygon(LineString([(29.0, 0.0), (31.0, 0.0)]))]) + assert nb.partner_tracked_objects(objs, map_api, {}) == [] diff --git a/tests/unit_tests/test_nuplan_partner_split.py b/tests/unit_tests/test_nuplan_partner_split.py new file mode 100644 index 0000000000..100dfce5a2 --- /dev/null +++ b/tests/unit_tests/test_nuplan_partner_split.py @@ -0,0 +1,51 @@ +"""Long nuPlan partners (buses) are shown to the policy as training-sized boxes along their heading.""" + +import math +import os +import sys + +import numpy as np +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from pufferlib.ocean.cosim import nuplan_bridge as nb # noqa: E402 + + +def _partners(): + x = np.array([10.0, 50.0], np.float32) + y = np.array([0.0, 5.0], np.float32) + z = np.zeros(2, np.float32) + h = np.array([0.0, math.pi / 2], np.float32) + vx = np.array([3.0, 0.0], np.float32) + vy = np.array([0.0, 1.0], np.float32) + types = np.array([1, 1], np.int32) + lengths = np.array([4.5, 12.7], np.float32) + widths = np.array([1.9, 2.4], np.float32) + return x, y, z, h, vx, vy, types, lengths, widths + + +def test_only_partners_over_the_limit_are_split_into_equal_pieces_along_their_heading(): + x, y, z, h, vx, vy, types, lengths, widths = nb.split_long_partners(*_partners(), max_length_m=7.0) + assert len(x) == 3 + assert (x[0], y[0], lengths[0]) == (10.0, 0.0, pytest.approx(4.5)) + assert lengths[1] == pytest.approx(6.35) and lengths[2] == pytest.approx(6.35) + assert x[1] == pytest.approx(50.0) and x[2] == pytest.approx(50.0) + assert sorted((y[1], y[2])) == pytest.approx([5.0 - 3.175, 5.0 + 3.175]) + assert list(widths) == pytest.approx([1.9, 2.4, 2.4]) + assert list(vy) == pytest.approx([0.0, 1.0, 1.0]) and list(types) == [1, 1, 1] + assert h.dtype == np.float32 and types.dtype == np.int32 + + +def test_split_is_the_identity_above_the_longest_partner(): + before = _partners() + after = nb.split_long_partners(*before, max_length_m=20.0) + for a, b in zip(before, after): + assert np.array_equal(a, b) + + +def test_nearest_partner_slots_keeps_the_devkit_order_until_the_slots_run_out(): + x = np.array([30.0, 5.0, 20.0], np.float32) + y = np.zeros(3, np.float32) + assert list(nb.nearest_partner_slots(x, y, 0.0, 0.0, 3)) == [0, 1, 2] + assert list(nb.nearest_partner_slots(x, y, 0.0, 0.0, 2)) == [1, 2] diff --git a/tests/unit_tests/test_nuplan_planner_log_state.py b/tests/unit_tests/test_nuplan_planner_log_state.py new file mode 100644 index 0000000000..1d423bcc48 --- /dev/null +++ b/tests/unit_tests/test_nuplan_planner_log_state.py @@ -0,0 +1,28 @@ +"""The planner's pickled state (nuPlan writes it into every simulation log) carries no policy, env or goal window.""" + +import os +import pickle +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +pytest.importorskip("nuplan") +pytest.importorskip("carl_nuplan") + +from pufferlib.ocean.cosim.nuplan.planner import LOG_EXCLUDED_STATE, PufferDrivePlanner + +WEIGHT_BYTES = 1 << 20 + + +def test_pickled_planner_drops_its_runtime_state(): + planner = PufferDrivePlanner.__new__(PufferDrivePlanner) + planner.__dict__.update({"_checkpoint_path": "model.pt", "_arch": {"num_agents": 2}}) + planner.__dict__.update({name: {"weights": bytes(WEIGHT_BYTES)} for name in LOG_EXCLUDED_STATE}) + pickled = pickle.dumps(planner) + assert len(pickled) < WEIGHT_BYTES + restored = pickle.loads(pickled) + assert all(getattr(restored, name) is None for name in LOG_EXCLUDED_STATE) + assert restored._checkpoint_path == "model.pt" and restored._arch == {"num_agents": 2} + assert all(getattr(planner, name) is not None for name in LOG_EXCLUDED_STATE) diff --git a/tests/unit_tests/test_nuplan_planner_rollout.py b/tests/unit_tests/test_nuplan_planner_rollout.py new file mode 100644 index 0000000000..234b4f846d --- /dev/null +++ b/tests/unit_tests/test_nuplan_planner_rollout.py @@ -0,0 +1,355 @@ +"""nuplan_tracker: the planner's forecast rollout. Slot assignment, what the policy drives, what moves at its observed +velocity, the no-reverse hold and the trajectory handed to nuPlan's tracker.""" + +import os +import sys +from types import SimpleNamespace + +import numpy as np +import pytest + + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +pytest.importorskip("nuplan") +pytest.importorskip("carl_nuplan") + +from nuplan.common.actor_state.ego_state import EgoState +from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimePoint +from nuplan.common.actor_state.vehicle_parameters import get_pacifica_parameters + +from pufferlib.ocean.cosim.nuplan.planner import ( + DUMMY_ROLLOUT_EGO_ACTION, + EGO_DYNAMICS_NUPLAN_TRACKER, + EGO_OBS_ACCEL_LONG_IDX, + EGO_OBS_STEERING_IDX, + PufferDrivePlanner, + startup_jerk_action_bounds, +) +from pufferlib.ocean.cosim.nuplan_bridge import LIGHT_STATE_GREEN, LIGHT_STATE_RED, LIGHT_STATE_YELLOW +from pufferlib.ocean.drive import binding + + +POLICY_VEHICLES = 2 +OBS_DIM = 12 +OBS_NORM_SPEED_MPS = 50.0 +STEP_SECONDS = 0.3 +ROLLOUT_STEPS = 4 +EGO_STEP_M = 1.5 # the scripted env moves the ego this far along x every step + + +class _FakeEnv: + """Records the co-sim calls; step() moves the ego along x and writes its speed, accel and steering obs.""" + + def __init__(self, agent_count=1 + POLICY_VEHICLES): + self.num_agents = agent_count + self.single_action_space = SimpleNamespace(shape=(2,)) + self.obs_norm_speed_mps = OBS_NORM_SPEED_MPS + self.obs_traffic_light_yellow_remaining = False + self.observations = np.zeros((agent_count, OBS_DIM), np.float32) + self.x = np.zeros(agent_count, np.float32) + self.lane_goal_counts = {} + self.state_calls = [] + self.light_calls = [] + self.step_actions = [] + self.recompute_count = 0 + self.step_speed_mps = 5.0 + + def set_agent_states( + self, idx, x, y, z, heading, vx, vy, yaw_rate, accel_long, seconds_stopped=None, steering=None + ): + call = {"idx": idx, "x": x, "y": y, "heading": heading, "vx": vx, "vy": vy} + call = {name: np.array(values) for name, values in call.items()} + call["steering"] = None if steering is None else np.array(steering) + self.state_calls.append(call) + + def set_agent_sizes(self, idx, length, width): + pass + + def set_agent_lane_goals(self, idx, spacing_m): + return np.array([self.lane_goal_counts.get(int(agent_idx), 3) for agent_idx in idx], np.int32) + + def set_traffic_light_states(self, states, yellow_remaining_s=None): + self.light_calls.append( + (np.array(states), None if yellow_remaining_s is None else np.array(yellow_remaining_s)) + ) + + def recompute_observations(self): + self.recompute_count += 1 + return self.observations + + def step(self, actions): + self.step_actions.append(np.array(actions)) + self.x[0] += EGO_STEP_M + self.observations[0, 0] = self.step_speed_mps / OBS_NORM_SPEED_MPS + self.observations[0, EGO_OBS_ACCEL_LONG_IDX] = 1.0 / binding.ACCEL_LONG_NORM + self.observations[0, EGO_OBS_STEERING_IDX] = 0.1 / binding.STEERING_LIMIT + + def get_global_agent_state(self): + zeros = np.zeros(self.num_agents, np.float32) + return {"x": self.x.copy(), "y": zeros.copy(), "z": zeros.copy(), "heading": zeros.copy()} + + +class _FakeGoalWindow: + def __init__(self, push_on_sync=()): + self.sync_poses = [] + self.push_on_sync = set(push_on_sync) + self.restored = None + + def sync(self, x, y, heading): + self.sync_poses.append(x) + return len(self.sync_poses) in self.push_on_sync + + def snapshot(self): + return ("snapshot after sync", len(self.sync_poses)) + + def restore(self, snapshot): + self.restored = snapshot + + +def _planner(env, **state): + planner = PufferDrivePlanner.__new__(PufferDrivePlanner) + planner.__dict__.update( + { + "_env": env, + "_policy": None, + "_obs_replay": None, + "_num_agents": 6, + "_goal_spacing": 20.0, + "_rollout_policy_vehicles": POLICY_VEHICLES, + "_rollout_step_seconds": STEP_SECONDS, + "_rollout_steps": ROLLOUT_STEPS, + "_rollout_allow_reverse": False, + "_rollout_policy_rows": np.zeros(1, np.int64), + "_forecast_objects": None, + "_planning_step_s": 0.1, + "_startup_action_bounds": None, + "_startup_jerk_cap_seconds": 1.5, + "_num_traffic": 0, + "_last_light_states": None, + "_last_yellow_remaining_s": None, + "_transform": SimpleNamespace(bin_to_loc=lambda x, y: (x + 1000.0, y + 2000.0)), + "_goal_window": _FakeGoalWindow(), + } + ) + planner.__dict__.update(state) + return planner + + +def _ego_state(time_us=5_000_000): + return EgoState.build_from_rear_axle( + rear_axle_pose=StateSE2(1000.0, 2000.0, 0.0), + rear_axle_velocity_2d=StateVector2D(4.0, 0.0), + rear_axle_acceleration_2d=StateVector2D(0.5, 0.0), + tire_steering_angle=0.0, + time_point=TimePoint(time_us), + vehicle_parameters=get_pacifica_parameters(), + ) + + +def _track(type_name, x, speed_mps): + return SimpleNamespace( + tracked_object_type=SimpleNamespace(name=type_name), + track_token=type_name, + center=SimpleNamespace(x=x, y=2000.0, heading=0.0), + velocity=SimpleNamespace(x=speed_mps, y=0.0), + box=SimpleNamespace(length=2.0, width=1.0), + ) + + +def test_sync_lets_the_policy_drive_vehicles_and_bicycles_but_not_pedestrians(): + env = _FakeEnv() + planner = _planner( + env, + _ego_dynamics=EGO_DYNAMICS_NUPLAN_TRACKER, + _ego_stopped_s=0.0, + _partner_stopped_s={}, + _initialization=SimpleNamespace(map_api=None), + _static_on_lane={}, + _vru_size_pad_m=0.0, + _partner_split_max_length_m=0.0, + _transform=SimpleNamespace(loc_to_bin=lambda x, y: (x - 1000.0, y - 2000.0)), + ) + tracks = [_track("PEDESTRIAN", 1003.0, 1.5), _track("BICYCLE", 1010.0, 4.0), _track("VEHICLE", 1020.0, 8.0)] + planner._sync(_ego_state(), SimpleNamespace(tracked_objects=tracks), None) + partners = env.state_calls[1] + slot_at_x = dict(zip(partners["x"].tolist(), partners["idx"].tolist())) + assert slot_at_x == {3.0: 3, 10.0: 1, 20.0: 2} # the pedestrian is nearest, yet it takes the first static slot + assert planner._rollout_policy_rows.tolist() == [0, 1, 2] + assert planner._forecast_objects[0].tolist() == [3] # the pedestrian walks on at its observed velocity + + +def test_nearest_vehicles_get_the_policy_slots_and_the_rest_static_slots(): + planner = _planner(_FakeEnv()) + x = np.array([30.0, 2.0, 10.0, 5.0, 60.0], np.float32) + eligible = np.array([True, False, True, True, True]) # row 1: a pedestrian next to the ego + idx = planner._rollout_partner_slots(x, np.zeros(5, np.float32), eligible, 0.0, 0.0) + assert idx.tolist() == [3, 4, 1, 2, 5] # vehicles at 10 m and 5 m take slots 1 and 2, statics start at 3 + + +def test_policy_drives_only_vehicles_with_a_lane_and_statics_keep_their_velocity(): + env = _FakeEnv() + env.lane_goal_counts = {2: 0} # the vehicle in slot 2 has no lane to follow + planner = _planner(env) + idx = np.array([1, 3, 2, 4], np.int32) + x = np.array([10.0, 20.0, 30.0, 40.0], np.float32) + zeros = np.zeros(4, np.float32) + vx = np.array([5.0, 1.5, 0.0, 0.0], np.float32) # slot 3: a walking pedestrian, slot 4: a cone + planner._arm_rollout_partners(idx, x, zeros, zeros, zeros, vx, zeros) + assert planner._rollout_policy_rows.tolist() == [0, 1] + forecast_idx, forecast_x, _, _, _, forecast_vx, _ = planner._forecast_objects + assert forecast_idx.tolist() == [3] and forecast_x.tolist() == [20.0] and forecast_vx.tolist() == [1.5] + + +def test_forecast_moves_statics_at_their_velocity_and_lets_a_yellow_run_out(): + env = _FakeEnv() + # one pedestrian in static slot 4 at (20, 1), walking at (1.5, -2) m/s + pedestrian = [np.array([value], np.float32) for value in (20.0, 1.0, 0.0, 0.5, 1.5, -2.0)] + planner = _planner( + env, + _forecast_objects=(np.array([4], np.int32), *pedestrian), + _num_traffic=2, + _last_light_states=np.array([LIGHT_STATE_GREEN, LIGHT_STATE_YELLOW], np.int32), + _last_yellow_remaining_s=np.array([0.0, 0.5], np.float32), + ) + planner._stamp_forecast(0.6) + stamped = env.state_calls[-1] + assert stamped["idx"].tolist() == [4] + np.testing.assert_allclose(stamped["x"], [20.9], atol=1e-5) + np.testing.assert_allclose(stamped["y"], [-0.2], atol=1e-5) + np.testing.assert_allclose(stamped["vx"], [1.5]) + states, remaining = env.light_calls[-1] + assert states.tolist() == [LIGHT_STATE_GREEN, LIGHT_STATE_RED] + assert remaining is None # the checkpoint has no yellow countdown column + + env.obs_traffic_light_yellow_remaining = True + planner._stamp_forecast(0.3) + states, remaining = env.light_calls[-1] + assert states.tolist() == [LIGHT_STATE_GREEN, LIGHT_STATE_YELLOW] + np.testing.assert_allclose(remaining, [0.0, 0.2], atol=1e-6) + + +def test_a_policy_agent_that_would_back_up_is_held_where_it_stood(): + env = _FakeEnv(agent_count=4) + planner = _planner(env) + obs = np.zeros((4, OBS_DIM), np.float32) + obs[:, 0] = [0.02, -0.01, -0.03, 0.1] # agent 1: policy-driven, backing up; agent 2: coasting backwards + obs[1, EGO_OBS_STEERING_IDX] = 0.3 / binding.STEERING_LIMIT + before = {key: np.array([1.0, 2.0, 3.0, 4.0], np.float32) for key in ("x", "y", "z", "heading")} + after = {key: values + 0.5 for key, values in before.items()} + assert planner._hold_reversing_agents(np.array([0, 1, 3]), obs, before, after) is True + held = env.state_calls[-1] + assert held["idx"].tolist() == [1] + assert held["x"].tolist() == [2.0] and held["heading"].tolist() == [2.0] + assert held["vx"].tolist() == [0.0] and held["vy"].tolist() == [0.0] + np.testing.assert_allclose(held["steering"], [0.3], atol=1e-6) + assert after["x"].tolist() == [1.5, 2.0, 3.5, 4.5] # only the held agent is back at its old pose + + calls = len(env.state_calls) + obs[1, 0] = 0.0 + assert planner._hold_reversing_agents(np.array([0, 1, 3]), obs, before, after) is False + assert len(env.state_calls) == calls + + +def test_rollout_returns_the_ego_states_as_a_trajectory_and_undoes_its_goal_progress(): + env = _FakeEnv() + goal_window = _FakeGoalWindow(push_on_sync={3}) + planner = _planner(env, _goal_window=goal_window, _rollout_policy_rows=np.array([0, 2])) + ego_state = _ego_state() + + trajectory = planner._rollout_trajectory(ego_state, iteration_index=0) + + states = trajectory.get_sampled_trajectory() + assert len(states) == 1 + ROLLOUT_STEPS + assert states[0] is ego_state + step_us = round(STEP_SECONDS * 1e6) + assert [s.time_point.time_us for s in states] == [5_000_000 + k * step_us for k in range(1 + ROLLOUT_STEPS)] + for step_idx, state in enumerate(states[1:], start=1): + assert state.center.x == pytest.approx(1000.0 + EGO_STEP_M * step_idx) # bin frame -> nuPlan frame + assert state.center.y == pytest.approx(2000.0) + assert state.dynamic_car_state.center_velocity_2d.x == pytest.approx(5.0, abs=1e-5) + assert state.dynamic_car_state.center_acceleration_2d.x == pytest.approx(1.0, abs=1e-5) + assert state.tire_steering_angle == pytest.approx(0.1, abs=1e-6) + assert trajectory.end_time.time_us - trajectory.start_time.time_us == ROLLOUT_STEPS * step_us + + assert len(env.step_actions) == ROLLOUT_STEPS + for actions in env.step_actions: # no checkpoint: the ego gets the wiring-test action, everything else coasts + assert actions[0].tolist() == list(DUMMY_ROLLOUT_EGO_ACTION) + assert not actions[1:].any() + assert goal_window.sync_poses == pytest.approx([EGO_STEP_M * k for k in range(1, 1 + ROLLOUT_STEPS)]) + assert goal_window.restored == ("snapshot after sync", 1) # the first step's progress is the one that is kept + assert env.recompute_count == 2 # once before the rollout, once after the goal window pushed on step 3 + + +def test_rollout_caps_the_ego_jerk_only_inside_the_startup_window(): + env = _FakeEnv() + bounds = startup_jerk_action_bounds(2.0, 2.0) + planner = _planner(env, _startup_action_bounds=bounds, _startup_jerk_cap_seconds=0.5) + planner._rollout_trajectory(_ego_state(), iteration_index=0) + ego_jerk_actions = [float(actions[0, 0]) for actions in env.step_actions] + assert ego_jerk_actions == pytest.approx([bounds[1], bounds[1], 1.0, 1.0]) # steps at 0.0 and 0.3 s are capped + + env.step_actions.clear() + planner._rollout_trajectory(_ego_state(), iteration_index=4) # 0.4 s into the scenario: only the first step + ego_jerk_actions = [float(actions[0, 0]) for actions in env.step_actions] + assert ego_jerk_actions == pytest.approx([bounds[1], 1.0, 1.0, 1.0]) + + +class _FakeObsReplay: + def __init__(self): + self.captures = [] + + def policy_outputs(self, obs_tensor, logits, value, entropy): + return {"obs_rows": obs_tensor.shape[0], "logit_rows": logits.shape[0], "value": value, "entropy": entropy} + + def capture(self, obs, actions, aux, action_index=None): + # like ObsReplayCapture: the action array is kept by reference until the replay is written + self.captures.append((np.array(obs), np.asarray(actions, dtype=np.float32), aux, action_index)) + + +@pytest.mark.parametrize("policy_rows", [[0], [0, 2]]) +def test_obs_replay_gets_the_ego_row_of_the_first_rollout_step(policy_rows): + torch = pytest.importorskip("torch") + env = _FakeEnv() + replay = _FakeObsReplay() + planner = _planner( + env, + _policy=SimpleNamespace(is_continuous=False), + _obs_replay=replay, + _rollout_policy_rows=np.array(policy_rows), + ) + row_count = len(policy_rows) + + calls = [] + + def policy_forward(obs): + # shapes as sample_logits returns them: a one-row batch has its entropy squeezed to 0-dim + entropy = torch.full((row_count,), 0.25).squeeze(0) if row_count == 1 else torch.full((row_count,), 0.25) + first_call = not calls + calls.append(len(obs)) + act = np.tile(np.array([[0.5, -0.5]] if first_call else [[-1.0, 1.0]], np.float32), (row_count, 1)) + action = torch.arange(7, 7 + row_count) + return act, torch.as_tensor(obs), torch.zeros(row_count, 12), torch.ones(row_count, 1), entropy, action + + planner._policy_forward = policy_forward + planner._rollout_trajectory(_ego_state(), iteration_index=0) + + assert len(replay.captures) == 1 # only the step taken from the real state is a replay frame + obs, actions, aux, action_index = replay.captures[0] + assert obs.shape == (1, OBS_DIM) + assert actions.tolist() == [[0.5, -0.5]] # the first step's action, not what later steps wrote into the buffer + assert calls == [row_count] * ROLLOUT_STEPS + assert aux["obs_rows"] == 1 and aux["logit_rows"] == 1 + assert tuple(aux["value"].shape) == (1, 1) and tuple(aux["entropy"].shape) == (1,) + assert action_index.tolist() == [7] + + +def test_tracker_mode_rejects_more_policy_vehicles_than_partners(): + with pytest.raises(ValueError, match="rollout_policy_vehicles"): + PufferDrivePlanner("dummy", "/tmp", num_agents=8, ego_dynamics="nuplan_tracker", rollout_policy_vehicles=8) + with pytest.raises(ValueError, match="rollout_step_seconds"): + PufferDrivePlanner("dummy", "/tmp", ego_dynamics="nuplan_tracker", rollout_step_seconds=0.0) + planner = PufferDrivePlanner( + "dummy", "/tmp", num_agents=8, ego_dynamics="nuplan_tracker", rollout_policy_vehicles=7 + ) + assert planner._rollout_policy_vehicles == 7 diff --git a/tests/unit_tests/test_nuplan_planner_startup_jerk_cap.py b/tests/unit_tests/test_nuplan_planner_startup_jerk_cap.py index dcfe1474af..389aaf5dcd 100644 --- a/tests/unit_tests/test_nuplan_planner_startup_jerk_cap.py +++ b/tests/unit_tests/test_nuplan_planner_startup_jerk_cap.py @@ -8,6 +8,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) nuplan = pytest.importorskip("nuplan") +pytest.importorskip("carl_nuplan") from pufferlib.ocean.cosim.nuplan.planner import startup_jerk_action_bounds from pufferlib.ocean.drive import binding diff --git a/tests/unit_tests/test_nuplan_planner_stopped_time.py b/tests/unit_tests/test_nuplan_planner_stopped_time.py new file mode 100644 index 0000000000..edfb7bec72 --- /dev/null +++ b/tests/unit_tests/test_nuplan_planner_stopped_time.py @@ -0,0 +1,54 @@ +"""A nuPlan partner's stopped time follows its track token, not the shadow-env slot it happens to occupy.""" + +import os +import sys +from types import SimpleNamespace + +import pytest + + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +pytest.importorskip("nuplan") +pytest.importorskip("carl_nuplan") + +from pufferlib.ocean.cosim.nuplan.planner import PufferDrivePlanner + + +STEP_SECONDS = 0.1 + + +def _planner(): + planner = PufferDrivePlanner.__new__(PufferDrivePlanner) + planner.__dict__.update({"_planning_step_s": STEP_SECONDS, "_partner_stopped_s": {}}) + return planner + + +def _track(track_token, speed_mps): + return SimpleNamespace(track_token=track_token, velocity=SimpleNamespace(x=speed_mps, y=0.0)) + + +def _static_object(track_token): + return SimpleNamespace(track_token=track_token) + + +def test_stopped_time_survives_a_change_of_list_order(): + planner = _planner() + planner._advance_partner_stopped_s([_track("walker", 0.0), _track("car", 5.0)]) + planner._advance_partner_stopped_s([_track("car", 5.0), _track("new", 3.0), _track("walker", 0.05)]) + assert planner._partner_stopped_s == pytest.approx({"walker": 2 * STEP_SECONDS, "car": 0.0, "new": 0.0}) + + +def test_moving_restarts_the_count_and_static_objects_always_count(): + planner = _planner() + for walker_speed_mps in (0.0, 0.0, 1.0, 0.0): + planner._advance_partner_stopped_s([_track("walker", walker_speed_mps), _static_object("cone")]) + assert planner._partner_stopped_s == pytest.approx({"walker": STEP_SECONDS, "cone": 4 * STEP_SECONDS}) + + +def test_a_track_missing_from_one_observation_starts_over(): + planner = _planner() + planner._advance_partner_stopped_s([_track("walker", 0.0)]) + planner._advance_partner_stopped_s([]) + planner._advance_partner_stopped_s([_track("walker", 0.0)]) + assert planner._partner_stopped_s == pytest.approx({"walker": STEP_SECONDS}) diff --git a/tests/unit_tests/test_nuplan_rollout.py b/tests/unit_tests/test_nuplan_rollout.py new file mode 100644 index 0000000000..89db60ade8 --- /dev/null +++ b/tests/unit_tests/test_nuplan_rollout.py @@ -0,0 +1,129 @@ +"""Forecast rollout helpers of the nuPlan planner's nuplan_tracker mode: step count, slot split, light forecast, +and the goal window snapshot that lets a rollout be undone.""" + +import os +import sys + +import numpy as np +import pytest + + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from pufferlib.ocean.cosim.goals import RouteGoalWindow, route_goals_from_xy +from pufferlib.ocean.cosim.nuplan_bridge import LIGHT_STATE_GREEN, LIGHT_STATE_RED, LIGHT_STATE_YELLOW +from pufferlib.ocean.cosim.nuplan_rollout import forecast_light_states, rollout_step_count, split_rollout_slots + + +NUM_GOALS = 3 +ROUTE_XY = np.array([[20.0, 0.0], [40.0, 0.0], [60.0, 0.0], [80.0, 0.0], [100.0, 0.0]]) + + +class _FakeGoalEnv: + def __init__(self): + self.num_goals = NUM_GOALS + self.goal_radius = 10.0 + self.window = [] + self.current_goal_idx = 0 + + def set_agent_goals(self, agent_idx, gx, gy, gz, gdir_x=None, gdir_y=None): + self.window = np.column_stack([gx, gy]).tolist() + self.current_goal_idx = 0 + + def get_agent_goal_progress(self, agent_idx): + return self.current_goal_idx, len(self.window) + + def set_agent_goal_progress(self, agent_idx, current_goal_idx): + assert 0 <= current_goal_idx <= len(self.window) + self.current_goal_idx = current_goal_idx + + +def test_step_count_reaches_the_horizon_without_a_rounding_step(): + assert rollout_step_count(8.0, 0.3) == 27 # 26 steps end at 7.8 s + assert rollout_step_count(8.1, 0.3) == 27 # 8.1 / 0.3 is 27.000000000000004 in floats + assert rollout_step_count(0.3, 0.3) == 1 + assert rollout_step_count(0.2, 0.3) == 1 + assert rollout_step_count(8.0, 0.1) == 80 + with pytest.raises(ValueError): + rollout_step_count(8.0, 0.0) + with pytest.raises(ValueError): + rollout_step_count(0.0, 0.3) + + +def test_nearest_eligible_partners_take_the_policy_slots(): + distance_sq = np.array([400.0, 1.0, 25.0, 100.0, 9.0]) + eligible = np.array([True, False, True, True, True]) # row 1: a pedestrian right next to the ego + policy_rows, static_rows = split_rollout_slots(distance_sq, eligible, 2) + assert policy_rows.tolist() == [2, 4] # the two nearest vehicles + assert static_rows.tolist() == [0, 1, 3] # the pedestrian and the vehicles that did not fit + everything, nothing = split_rollout_slots(distance_sq, eligible, 10) + assert everything.tolist() == [0, 2, 3, 4] and nothing.tolist() == [1] + none, all_static = split_rollout_slots(distance_sq, eligible, 0) + assert none.tolist() == [] and all_static.tolist() == [0, 1, 2, 3, 4] + + +def test_slot_split_of_an_empty_scene(): + policy_rows, static_rows = split_rollout_slots(np.zeros(0), np.zeros(0, bool), 4) + assert len(policy_rows) == 0 and len(static_rows) == 0 + + +def test_forecast_lights_hold_except_a_yellow_that_runs_out(): + states = np.array([LIGHT_STATE_GREEN, LIGHT_STATE_RED, LIGHT_STATE_YELLOW, LIGHT_STATE_YELLOW], np.int32) + remaining = np.array([0.0, 0.0, 0.5, 2.0], np.float32) + early_states, early_remaining = forecast_light_states(states, remaining, 0.3) + assert early_states.tolist() == states.tolist() + np.testing.assert_allclose(early_remaining, [0.0, 0.0, 0.2, 1.7], atol=1e-6) + late_states, late_remaining = forecast_light_states(states, remaining, 0.6) + assert late_states.tolist() == [LIGHT_STATE_GREEN, LIGHT_STATE_RED, LIGHT_STATE_RED, LIGHT_STATE_YELLOW] + np.testing.assert_allclose(late_remaining, [0.0, 0.0, 0.0, 1.4], atol=1e-6) + assert late_states.dtype == np.int32 and late_remaining.dtype == np.float32 + assert states.tolist()[2] == LIGHT_STATE_YELLOW # the synced arrays are reused for every rollout step + + +def test_restore_undoes_the_goals_a_rollout_consumed_in_a_sliding_window(): + env = _FakeGoalEnv() + window = RouteGoalWindow(env, route_goals_from_xy(ROUTE_XY), sliding=True) + assert window.sync(0.0, 0.0, 0.0) is True # first push + assert window.sync(1.0, 0.0, 0.0) is False # nothing consumed: the observations stay valid + env.current_goal_idx = 1 # the first rollout step reaches the first goal: this one counts + assert window.sync(12.0, 0.0, 0.0) is True + snapshot = window.snapshot() + for ego_x in (32.0, 52.0): # later rollout steps run two more goals down + env.current_goal_idx = 1 + window.sync(ego_x, 0.0, 0.0) + assert env.window == [[80.0, 0.0], [100.0, 0.0]] + window.restore(snapshot) + assert env.window == [[40.0, 0.0], [60.0, 0.0], [80.0, 0.0]] + assert (window.window_start, window.current_index, env.current_goal_idx) == (1, 1, 0) + + +def test_restore_keeps_the_consumed_count_of_a_batch_window(): + env = _FakeGoalEnv() + window = RouteGoalWindow(env, route_goals_from_xy(ROUTE_XY)) + window.sync(0.0, 0.0, 0.0) + env.current_goal_idx = 1 + window.sync(12.0, 0.0, 0.0) + snapshot = window.snapshot() # one of three goals consumed, window not replaced + env.current_goal_idx = 3 + window.sync(62.0, 0.0, 0.0) # the rollout exhausts the window: the next one is pushed + assert env.window == [[80.0, 0.0], [100.0, 0.0]] + window.restore(snapshot) + assert env.window == [[20.0, 0.0], [40.0, 0.0], [60.0, 0.0]] + assert env.current_goal_idx == 1 and window.window_start == 0 + + +def test_restore_keeps_an_exhausted_route_exhausted(): + env = _FakeGoalEnv() + window = RouteGoalWindow(env, route_goals_from_xy(ROUTE_XY[:1]), sliding=True) + window.sync(0.0, 0.0, 0.0) + env.current_goal_idx = 1 # the only goal is consumed: nothing left to push + assert window.sync(15.0, 0.0, 0.0) is False + snapshot = window.snapshot() + window.restore(snapshot) + assert env.get_agent_goal_progress(0) == (1, 1) # still consumed: the goal must not reappear + + +def test_snapshot_needs_a_pushed_window(): + window = RouteGoalWindow(_FakeGoalEnv(), route_goals_from_xy(ROUTE_XY)) + with pytest.raises(RuntimeError): + window.snapshot() diff --git a/tests/unit_tests/test_obs_lane_width_clip.py b/tests/unit_tests/test_obs_lane_width_clip.py new file mode 100644 index 0000000000..970a8b1e3b --- /dev/null +++ b/tests/unit_tests/test_obs_lane_width_clip.py @@ -0,0 +1,106 @@ +"""obs_lane_width_clip_{min,max}_m clip only the lane-row width column of the observation.""" + +import os +import shutil +import sys + +import numpy as np +import pytest + +from pufferlib.ocean.drive.drive import Drive + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, os.path.join(REPO_ROOT, "data_utils")) +from mirror_map_bin import _is_lane, read_bin, write_bin # noqa: E402 + +SOURCE_BIN = os.path.join( + REPO_ROOT, + "pufferlib", + "resources", + "drive", + "binaries", + "sdc_replay_test", + "nuplan__c2a99be0-5b3e-5e67-88e6-6fd0f22e5419.bin", +) +LANE_WIDTH_COLUMN = 6 +CLIP_MIN_M = 2.85 +CLIP_MAX_M = 4.4 +ROAD_SEG_WIDTH_NORM_M = 5.0 + + +def _write_uniform_width_bin(out_dir, lane_width_m): + data = read_bin(SOURCE_BIN) + for road in data["roads"]: + if _is_lane(road["type"]): + road["widths"] = tuple([float(lane_width_m)] * road["S"]) + data["has_phase_section"] = True # the loader reads the tail sections in order: phases, widths, zones + data["has_width_section"] = True + os.makedirs(out_dir) + write_bin(data, os.path.join(out_dir, os.path.basename(SOURCE_BIN))) + return out_dir + + +def _lane_width_obs(map_dir, **clip_kwargs): + env = Drive( + num_agents=1, + min_agents_per_env=1, + max_agents_per_env=1, + num_maps=1, + map_dir=map_dir, + simulation_mode="replay", + goal_source="gt", + control_mode="control_sdc_only", + sdc_controller="replay", + non_sdc_controller="replay", + scenario_length=400, + resample_frequency=1_000_000, + termination_mode=False, + terminate_on_goal=False, + obs_dropout_lane=0.0, + obs_dropout_boundary=0.0, + obs_norm_road_seg_width_m=ROAD_SEG_WIDTH_NORM_M, + **clip_kwargs, + ) + obs, _ = env.reset() + lane_start = ( + env.ego_features + env.num_reward_coefs + env.goal_dim + env.obs_slots_partners_n * env.partner_features + ) + lane_rows = obs[:, lane_start : lane_start + env.obs_slots_lane_kept * env.lane_features] + lane_rows = lane_rows.reshape(len(obs), env.obs_slots_lane_kept, env.lane_features) + lane_counts = obs[:, -4].astype(int) + env.close() + widths = [lane_rows[agent, :count, LANE_WIDTH_COLUMN] for agent, count in enumerate(lane_counts)] + widths = np.concatenate(widths) + assert len(widths) > 0 + return widths + + +@pytest.mark.parametrize("lane_width_m,expected_clipped_m", [(12.0, CLIP_MAX_M), (1.0, CLIP_MIN_M)]) +def test_lane_width_clipped_to_range(tmp_path, lane_width_m, expected_clipped_m): + map_dir = _write_uniform_width_bin(str(tmp_path / "bins"), lane_width_m) + unclipped = _lane_width_obs(map_dir) + clipped = _lane_width_obs(map_dir, obs_lane_width_clip_min_m=CLIP_MIN_M, obs_lane_width_clip_max_m=CLIP_MAX_M) + np.testing.assert_allclose(unclipped, lane_width_m / ROAD_SEG_WIDTH_NORM_M, rtol=1e-5) + np.testing.assert_allclose(clipped, expected_clipped_m / ROAD_SEG_WIDTH_NORM_M, rtol=1e-5) + shutil.rmtree(map_dir) + + +def test_lane_width_clip_disabled_by_default_and_by_zero_max(tmp_path): + map_dir = _write_uniform_width_bin(str(tmp_path / "bins"), 12.0) + default = _lane_width_obs(map_dir) + zero_max = _lane_width_obs(map_dir, obs_lane_width_clip_min_m=CLIP_MIN_M, obs_lane_width_clip_max_m=0.0) + np.testing.assert_allclose(default, 12.0 / ROAD_SEG_WIDTH_NORM_M, rtol=1e-5) + np.testing.assert_array_equal(zero_max, default) + + +@pytest.mark.parametrize( + "kwargs", + [ + dict(obs_lane_width_clip_min_m=-1.0), + dict(obs_lane_width_clip_max_m=float("nan")), + dict(obs_lane_width_clip_min_m=5.0, obs_lane_width_clip_max_m=4.4), + ], +) +def test_lane_width_clip_validation(kwargs): + with pytest.raises(ValueError, match="obs_lane_width_clip"): + Drive(config_only=True, **kwargs) diff --git a/tests/unit_tests/test_render_infraction_steps.py b/tests/unit_tests/test_render_infraction_steps.py index 9a56b7b4e7..419a11f514 100644 --- a/tests/unit_tests/test_render_infraction_steps.py +++ b/tests/unit_tests/test_render_infraction_steps.py @@ -7,6 +7,7 @@ import pandas as pd from pufferlib import viz +from pufferlib.ocean.cosim.carla.shadow_ego import TELEMETRY_HEADER sys.path.insert( @@ -21,6 +22,9 @@ TICK_DT = 0.05 OFFSET = (10.0, 20.0) COLLISION_FRAME = 600 +RED_LIGHT_FRAME = 450 +# the light itself, 18 m off the ego's path +RED_LIGHT_MESSAGE = "Agent ran a red light 365 at (x=310.0, y=-62.0, z=0.3)" def _replay(ego_xy, collision_frames=()): @@ -84,34 +88,93 @@ def test_read_replay_zlib_chunk_subset_matches_full_read(tmp_path): assert chunks_full["agent_f32"].nbytes > viz.REPLAY_HEADER_PROBE_BYTES -def test_carla_infraction_steps_link_nearest_ego_frame(tmp_path): +def test_carla_infractions_are_placed_at_the_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) + frames, frame_count = carla_render.infraction_frames(record, str(path), "Town01", log_meta) collision, deviation, outside = ( record["infractions"][key][0] for key in ("collisions_vehicle", "route_dev", "outside_route_lanes") ) + assert frame_count == FRAMES + assert frames == {collision: COLLISION_FRAME, deviation: FRAMES} + assert carla_render.infraction_frames(record, str(path), "", {}) == ({}, FRAMES) + + steps = carla_render.step_links(frames, frame_count, TICK_DT) 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_ego_is_marked_from_every_infraction_placed_in_the_replay(tmp_path): + path = tmp_path / "route.replay.zlib" + viz.save_interactive_replay_zlib({}, _replay(_straight_ego_xy()), str(path)) + record = _carla_record() + frames, frame_count = carla_render.infraction_frames(record, str(path), "Town01", {"offset": list(OFFSET)}) + + marks = carla_render.ego_marks(record, frames, frame_count) + assert marks == [(carla_render.EGO_AGENT_IDX, COLLISION_FRAME, carla_render.MARK_FRAMES, "LEADERBOARD COLLISION")] + script = viz.replay_agent_marks(marks) + mark_end_frame = COLLISION_FRAME + carla_render.MARK_FRAMES + assert script == ( + f'' + ) + + 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) + message = f"Agent got blocked at (x={300.0 - tx:.3f}, y={-(100.0 - ty):.3f}, z=0.0)" + record = {"infractions": {"vehicle_blocked": [message]}} + frames, frame_count = carla_render.infraction_frames(record, str(path), "Town01", {}) + assert frames == {message: COLLISION_FRAME} + assert carla_render.step_links(frames, frame_count, None)[message] == viz.replay_step_link(COLLISION_FRAME) + + +def _write_telemetry(path, row_count, red_light_rows): + columns = TELEMETRY_HEADER.strip().split(",") + with open(path, "w") as telemetry_file: + telemetry_file.write(TELEMETRY_HEADER) + for row in range(row_count): + values = {"step": row + 1, carla_render.RED_LIGHT_FLAG_COLUMN: int(row in red_light_rows)} + telemetry_file.write(",".join(str(values.get(column, 0)) for column in columns) + "\n") + + +def test_carla_red_lights_are_placed_at_the_step_the_telemetry_flagged(tmp_path): + path = tmp_path / "route.replay.zlib" + viz.save_interactive_replay_zlib({}, _replay(_straight_ego_xy()), str(path)) + telemetry = tmp_path / "route.csv" + _write_telemetry(telemetry, FRAMES + 100, {RED_LIGHT_FRAME, FRAMES + 50}) + after_the_replay = "Agent ran a red light 366 at (x=1500.0, y=-62.0, z=0.3)" + record = {"infractions": {"red_light": [RED_LIGHT_MESSAGE, after_the_replay]}} + log_meta = {"offset": list(OFFSET), "tick_dt": TICK_DT} + + frames, frame_count = carla_render.infraction_frames(record, str(path), "Town01", log_meta, str(telemetry)) + assert frames == {RED_LIGHT_MESSAGE: RED_LIGHT_FRAME, after_the_replay: FRAMES + 50} + steps = carla_render.step_links(frames, frame_count, TICK_DT) + assert steps[RED_LIGHT_MESSAGE] == viz.replay_step_link(RED_LIGHT_FRAME, RED_LIGHT_FRAME * TICK_DT) + assert steps[after_the_replay] == "beyond the replay" + assert carla_render.ego_marks(record, frames, frame_count) == [ + (carla_render.EGO_AGENT_IDX, RED_LIGHT_FRAME, carla_render.MARK_FRAMES, "LEADERBOARD RED LIGHT") + ] + + +def test_carla_red_lights_are_not_placed_without_one_telemetry_flag_each(tmp_path): + path = tmp_path / "route.replay.zlib" + viz.save_interactive_replay_zlib({}, _replay(_straight_ego_xy()), str(path)) + record = {"infractions": {"red_light": [RED_LIGHT_MESSAGE]}} + log_meta = {"offset": list(OFFSET), "tick_dt": TICK_DT} + assert carla_render.infraction_frames(record, str(path), "Town01", log_meta) == ({}, FRAMES) + + telemetry = tmp_path / "route.csv" + _write_telemetry(telemetry, FRAMES, {RED_LIGHT_FRAME, RED_LIGHT_FRAME + 200}) + assert carla_render.infraction_frames(record, str(path), "Town01", log_meta, str(telemetry)) == ({}, FRAMES) def _write_series(metrics_dir, name, values, extra=None): @@ -179,3 +242,53 @@ def test_nuplan_score_panel_links_violated_metrics_only(): 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 + + +def _aggregator_csv(group_dir, scenarios): + """scenarios: token -> (score, at-fault collision value); every other scored metric compliant.""" + sim = group_dir / "simulation" / "closed_loop" / "2026.01.01.00.00.00" + (sim / "aggregator_metric").mkdir(parents=True, exist_ok=True) + (sim / "metrics").mkdir(exist_ok=True) + table = [] + for token, (score, collision) in scenarios.items(): + row = {"scenario": token, "log_name": "log", "scenario_type": "turn", "score": score} + row.update({metric: 1.0 for metric, _, _ in nuplan_render.SCORED_METRICS}) + row[nuplan_render.COLLISION_METRIC] = collision + table.append(row) + table.append({"scenario": "final_score", "scenario_type": "final_score", "score": 0.5}) + pd.DataFrame(table).to_csv(sim / "aggregator_metric" / "closed_loop_weighted_average_metrics_2026.csv", index=False) + + +def test_page_payload_matches_the_replay_file(tmp_path): + replay_path = tmp_path / "tok.replay.zlib" + viz.save_interactive_replay_zlib({}, _replay(_straight_ego_xy(), collision_frames=(7,)), str(replay_path)) + page = tmp_path / "tok.html" + viz.render_interactive_replay_zlib(str(replay_path), str(page)) + assert viz.read_replay_source(page) == replay_path.read_bytes() + header, chunks = viz.read_replay_zlib(page, chunk_names=("metrics_f32",)) + header_from_file, chunks_from_file = viz.read_replay_zlib(replay_path) + assert header == header_from_file + assert np.array_equal(chunks["metrics_f32"], chunks_from_file["metrics_f32"]) + assert nuplan_render.replay_collision_step(page) == 7 + + +def test_render_prune_deletes_every_replay_and_pages_re_render_from_themselves(tmp_path, monkeypatch): + colliding, clean = "aaaa0000000000ad", "aaaa00000000900d" + obs_dir = tmp_path / "obs_html" + obs_dir.mkdir() + _aggregator_csv(tmp_path, {colliding: (0.0, 0.0), clean: (1.0, 1.0)}) + for token in (colliding, clean): + replay = _replay(np.zeros((FRAMES, 2)), collision_frames=(40,)) + viz.save_interactive_replay_zlib({}, replay, str(obs_dir / f"{token}.replay.zlib")) + argv = ["render_obs_html", str(tmp_path), "--max-score", "0", "--metric", nuplan_render.COLLISION_METRIC, "--prune"] + monkeypatch.setattr(sys, "argv", argv) + nuplan_render.main() + page = obs_dir / f"s0.00_COL_turn_{colliding}.html" + assert sorted(p.name for p in obs_dir.iterdir()) == ["index.html", page.name] + assert nuplan_render.replay_sources(obs_dir) == {colliding: page} + assert nuplan_render.replay_collision_step(page) == 40 + payload = viz.read_replay_source(page) + nuplan_render.main() + assert sorted(p.name for p in obs_dir.iterdir()) == ["index.html", page.name] + assert viz.read_replay_source(page) == payload + assert 'onclick="step=40;' in page.read_text() diff --git a/tests/unit_tests/test_scripted_pedestrians.py b/tests/unit_tests/test_scripted_pedestrians.py new file mode 100644 index 0000000000..c645cefbcd --- /dev/null +++ b/tests/unit_tests/test_scripted_pedestrians.py @@ -0,0 +1,105 @@ +"""Scripted pedestrians (env.scripted_pedestrian_ratio): config validation, per-env slot rounding, and a +short gigaflow rollout in which walkers appear as static partners and get logged.""" + +import numpy as np +import pytest + +from pufferlib.ocean.drive.drive import Drive + +MAP_DIR = "pufferlib/resources/drive/binaries/carla" +AGENTS_PER_ENV = 32 + + +def _kwargs(**overrides): + kwargs = { + "num_agents": AGENTS_PER_ENV, + "min_agents_per_env": AGENTS_PER_ENV, + "max_agents_per_env": AGENTS_PER_ENV, + "num_maps": 1, + "allow_map_subset": True, + "map_dir": MAP_DIR, + "simulation_mode": "gigaflow", + "action_type": "continuous", + "dynamics_model": "jerk", + "scenario_length": 300, + "resample_frequency": 0, + "termination_mode": 0, + "stagger_first_episode": False, + "spawn_initial_speed": 5.0, + "seed": 3, + } + kwargs.update(overrides) + return kwargs + + +def test_ratio_validation(): + with pytest.raises(ValueError, match="scripted_pedestrian_ratio"): + Drive(**_kwargs(scripted_pedestrian_ratio=-0.1, config_only=True)) + with pytest.raises(ValueError, match="scripted_pedestrian_ratio"): + Drive(**_kwargs(scripted_pedestrian_ratio=0.1, simulation_mode="replay", config_only=True)) + with pytest.raises(ValueError, match="scripted_pedestrian_ratio"): + Drive(**_kwargs(scripted_pedestrian_ratio=0.1, cosim_partner_slots=4, config_only=True)) + + +def test_slot_rounding_matches_ceil_of_ratio_times_agents(): + env = Drive(**_kwargs(scripted_pedestrian_ratio=0.1, config_only=True)) + assert [env._scripted_pedestrian_slots(n) for n in (1, 10, 11, 30, 150)] == [1, 1, 2, 3, 15] + assert Drive(**_kwargs(config_only=True))._scripted_pedestrian_slots(150) == 0 + assert env._env_init_kwargs("unused.bin", 30)["scripted_pedestrian_slots"] == 3 + + +def test_rollout_spawns_pedestrians_as_static_partners(): + env = Drive(**_kwargs(scripted_pedestrian_ratio=0.125)) + try: + env.reset(seed=3) + assert env.num_envs == 1 + expected_slots = 4 + state = env.get_global_agent_state(include_static=True) + assert len(state["x"]) == env.num_agents + expected_slots + + actions = np.zeros_like(env.actions) + walkers_seen = 0 + appearances_logged = 0.0 + for _ in range(env.scenario_length + 1): + obs, _, _, _, infos = env.step(actions) + assert np.all(np.isfinite(obs)) + state = env.get_global_agent_state(include_static=True) + walker_x = state["x"][env.num_agents :] + walkers_seen += int(np.sum(walker_x > -9000.0)) + for info in infos: + appearances_logged += float(info["pedestrian_crossings"]) + float(info["pedestrian_bystanders"]) + assert walkers_seen > 0 + assert appearances_logged > 0.0 + finally: + env.close() + + +def _walker_trace(seed, num_steps): + env = Drive(**_kwargs(scripted_pedestrian_ratio=0.125, seed=seed)) + try: + env.reset(seed=seed) + actions = np.zeros_like(env.actions) + trace = [] + for _ in range(num_steps): + env.step(actions) + trace.append(env.get_global_agent_state(include_static=True)["x"][env.num_agents :].copy()) + return np.asarray(trace) + finally: + env.close() + + +def test_walkers_are_deterministic_per_seed(): + first = _walker_trace(5, 120) + second = _walker_trace(5, 120) + assert np.array_equal(first, second) + assert np.any(first > -9000.0) + + +def test_zero_ratio_adds_no_slots(): + env = Drive(**_kwargs()) + try: + env.reset(seed=3) + state = env.get_global_agent_state(include_static=True) + assert len(state["x"]) == env.num_agents + finally: + env.close() diff --git a/tests/unit_tests/test_shadow_ego_plant_state.py b/tests/unit_tests/test_shadow_ego_plant_state.py new file mode 100644 index 0000000000..f23d75aa92 --- /dev/null +++ b/tests/unit_tests/test_shadow_ego_plant_state.py @@ -0,0 +1,112 @@ +import math +from types import SimpleNamespace + +from pufferlib.ocean.cosim.carla.controller import pedal_gains_for +from pufferlib.ocean.cosim.carla.shadow_ego import DRIVE_WHEELBASE_PER_LENGTH, ShadowEgo + + +GAINS = pedal_gains_for("vehicle.lincoln.mkz2017") +EGO_HALF_LENGTH_M = 2.45 +HEADING_RAD = 0.3 +RESIDUAL_SPEED_MPS = 0.02 + + +def _shadow(intent_accel_mps2=0.0, held=False, moving_direction=1.0, intent_steering_rad=0.0, pedal_reverse=True): + return SimpleNamespace( + controller=SimpleNamespace(gains=GAINS, observe=lambda speed_mps, yaw_rate_radps, tick_dt: None), + tick_dt=0.05, + ego_extent=[EGO_HALF_LENGTH_M, 1.0, 0.75], + pedal_reverse=pedal_reverse, + _standing=False, + _intent_accel_mps2=intent_accel_mps2, + _intent_steering_rad=intent_steering_rad, + _held=held, + _moving_direction=moving_direction, + _plant_speed_mps=0.0, + ) + + +def _plant_state(shadow, forward_speed_mps): + vx, vy = forward_speed_mps * math.cos(HEADING_RAD), forward_speed_mps * math.sin(HEADING_RAD) + return ShadowEgo._plant_state(shadow, (1.0, 2.0, 0.0, HEADING_RAD, vx, vy, 0.0, 0.0)) + + +def test_a_moving_car_keeps_its_velocity_and_the_integrator_acceleration(): + shadow = _shadow(intent_accel_mps2=-2.0, held=True) + _, _, _, _, vx, vy, _, accel_mps2 = _plant_state(shadow, 3.0) + assert math.isclose(math.hypot(vx, vy), 3.0) + assert accel_mps2 == -2.0 + assert not shadow._held + + +def test_braking_to_a_stop_holds_the_car_and_starts_the_acceleration_from_zero(): + shadow = _shadow(intent_accel_mps2=-2.0) + _, _, _, _, vx, vy, _, accel_mps2 = _plant_state(shadow, RESIDUAL_SPEED_MPS) + assert (vx, vy, accel_mps2) == (0.0, 0.0, 0.0) + assert shadow._held + + +def test_stopping_from_reverse_holds_the_car_too(): + shadow = _shadow(intent_accel_mps2=1.5, moving_direction=-1.0) + _, _, _, _, vx, vy, _, accel_mps2 = _plant_state(shadow, -RESIDUAL_SPEED_MPS) + assert (vx, vy, accel_mps2) == (0.0, 0.0, 0.0) + assert shadow._held + + +def test_a_launch_from_the_hold_keeps_its_acceleration_in_either_direction(): + for launch_accel_mps2 in (0.4, -0.4): + shadow = _shadow(intent_accel_mps2=launch_accel_mps2, held=True) + accel_mps2 = _plant_state(shadow, 0.0)[7] + assert accel_mps2 == launch_accel_mps2 + assert shadow._held + + +def test_without_reverse_a_request_to_go_backwards_never_builds_up(): + shadow = _shadow(intent_accel_mps2=-0.4, held=True, pedal_reverse=False) + accel_mps2 = _plant_state(shadow, 0.0)[7] + assert accel_mps2 == 0.0 + assert shadow._held + + +def test_without_reverse_a_car_that_rolls_backwards_is_held_like_a_standing_one(): + shadow = _shadow(intent_accel_mps2=-1.0, pedal_reverse=False) + _, _, _, _, vx, vy, _, accel_mps2 = _plant_state(shadow, -0.5) + assert (vx, vy, accel_mps2) == (0.0, 0.0, 0.0) + assert shadow._standing and shadow._held + + +def test_with_reverse_a_car_that_rolls_backwards_counts_as_moving(): + shadow = _shadow(intent_accel_mps2=-1.0) + _, _, _, _, vx, vy, _, accel_mps2 = _plant_state(shadow, -0.5) + assert math.isclose(math.hypot(vx, vy), 0.5) + assert accel_mps2 == -1.0 + assert not shadow._standing and shadow._moving_direction == -1.0 + + +def test_a_standing_car_reads_stopped_even_when_it_creeps_against_the_launch(): + shadow = _shadow(intent_accel_mps2=-0.4, held=True) + _, _, _, _, vx, vy, _, accel_mps2 = _plant_state(shadow, RESIDUAL_SPEED_MPS) + assert (vx, vy) == (0.0, 0.0) + assert accel_mps2 == -0.4 + + +def test_a_car_that_creeps_in_the_direction_of_its_launch_keeps_its_velocity(): + shadow = _shadow(intent_accel_mps2=0.4, held=True) + _, _, _, _, vx, vy, _, _ = _plant_state(shadow, RESIDUAL_SPEED_MPS) + assert math.isclose(math.hypot(vx, vy), RESIDUAL_SPEED_MPS) + + +def test_a_car_that_is_still_accelerating_when_it_stalls_is_not_held(): + shadow = _shadow(intent_accel_mps2=1.0) + accel_mps2 = _plant_state(shadow, RESIDUAL_SPEED_MPS)[7] + assert accel_mps2 == 1.0 + assert not shadow._held + + +def test_the_reported_heading_is_the_direction_the_centre_travels_in(): + steering_rad = 0.3 + wheelbase_m = DRIVE_WHEELBASE_PER_LENGTH * 2.0 * EGO_HALF_LENGTH_M + slip_rad = math.atan(GAINS.rear_axle_to_center_m * math.tan(steering_rad) / wheelbase_m) + assert math.isclose(_plant_state(_shadow(intent_steering_rad=steering_rad), 3.0)[3], HEADING_RAD + slip_rad) + assert math.isclose(_plant_state(_shadow(intent_steering_rad=-steering_rad), 3.0)[3], HEADING_RAD - slip_rad) + assert _plant_state(_shadow(), 3.0)[3] == HEADING_RAD diff --git a/tests/unit_tests/test_shadow_ego_stopped_time.py b/tests/unit_tests/test_shadow_ego_stopped_time.py new file mode 100644 index 0000000000..6b4ee146c1 --- /dev/null +++ b/tests/unit_tests/test_shadow_ego_stopped_time.py @@ -0,0 +1,69 @@ +from types import SimpleNamespace + +import numpy as np + +from pufferlib.ocean.cosim.carla.shadow_ego import ( + PARTNER_ACTOR_ID, + PARTNER_IS_WALKER, + PARTNER_ROW_COLUMNS, + STATE_VX, + ShadowEgo, +) + + +STEP_SECONDS = 0.05 +VEHICLE_ROW, WALKER_ROW = 0, 1 + + +def _state_to_bin(x, y, z, yaw_deg, vx, vy, yaw_rate, accel_x, accel_y, box_z, extent_z): + return x, y, z, 0.0, vx, vy, yaw_rate, accel_x + + +def _shadow(zero_partner_stopped_time): + return SimpleNamespace( + transform=SimpleNamespace(state_to_bin=_state_to_bin), + step_count=0, + tick_dt=STEP_SECONDS, + dt=STEP_SECONDS, + walker_speed_floor_mps=0.0, + zero_partner_stopped_time=zero_partner_stopped_time, + _partner_stopped_speed_threshold=0.2, + _partner_stopped_s={}, + ) + + +def _partners(walker_speed_mps=0.0): + partners = np.zeros((2, PARTNER_ROW_COLUMNS), np.float32) + partners[VEHICLE_ROW, PARTNER_ACTOR_ID] = 11 + partners[WALKER_ROW, PARTNER_ACTOR_ID] = 12 + partners[WALKER_ROW, PARTNER_IS_WALKER] = 1.0 + partners[WALKER_ROW, STATE_VX] = walker_speed_mps + return partners + + +def _stopped_seconds(shadow, partners): + return ShadowEgo._read_states(shadow, partners)[-1] + + +def test_zeroing_holds_vehicles_at_zero_and_keeps_the_walker_counter(): + shadow = _shadow(True) + for _ in range(4): + stopped_seconds = _stopped_seconds(shadow, _partners()) + assert stopped_seconds[VEHICLE_ROW] == 0.0 + assert abs(stopped_seconds[WALKER_ROW] - 4 * STEP_SECONDS) < 1e-6 + + +def test_walker_counter_restarts_once_the_walker_moves(): + shadow = _shadow(True) + for _ in range(4): + _stopped_seconds(shadow, _partners()) + assert _stopped_seconds(shadow, _partners(walker_speed_mps=1.0))[WALKER_ROW] == 0.0 + assert abs(_stopped_seconds(shadow, _partners())[WALKER_ROW] - STEP_SECONDS) < 1e-6 + + +def test_without_zeroing_every_partner_keeps_its_counter(): + shadow = _shadow(False) + for _ in range(4): + stopped_seconds = _stopped_seconds(shadow, _partners()) + assert abs(stopped_seconds[VEHICLE_ROW] - 4 * STEP_SECONDS) < 1e-6 + assert abs(stopped_seconds[WALKER_ROW] - 4 * STEP_SECONDS) < 1e-6 diff --git a/tests/unit_tests/test_shadow_ego_walker_floor.py b/tests/unit_tests/test_shadow_ego_walker_floor.py new file mode 100644 index 0000000000..fac0500e80 --- /dev/null +++ b/tests/unit_tests/test_shadow_ego_walker_floor.py @@ -0,0 +1,24 @@ +import math +from types import SimpleNamespace + +from pufferlib.ocean.cosim.carla.shadow_ego import WALKER_SPEED_FLOOR_WINDOW_S, ShadowEgo + + +def _shadow(floor_mps): + return SimpleNamespace(walker_speed_floor_mps=floor_mps, _walker_first_seen_s={}) + + +def test_fresh_slow_walker_reads_floor_speed_along_its_facing(): + shadow = _shadow(3.8) + vx, vy = ShadowEgo._floored_walker_velocity(shadow, 7, 10.0, math.pi / 2, 0.0, 0.0) + assert abs(vx) < 1e-6 and abs(vy - 3.8) < 1e-6 + vx, vy = ShadowEgo._floored_walker_velocity(shadow, 7, 10.5, math.pi / 2, 0.0, 1.0) + assert abs(vx) < 1e-6 and abs(vy - 3.8) < 1e-6 + + +def test_fast_walker_and_expired_window_keep_raw_velocity(): + shadow = _shadow(3.8) + assert ShadowEgo._floored_walker_velocity(shadow, 7, 10.0, 0.0, 4.0, 0.5) == (4.0, 0.5) + late_s = 10.0 + WALKER_SPEED_FLOOR_WINDOW_S + 0.1 + assert ShadowEgo._floored_walker_velocity(shadow, 7, late_s, 0.0, 0.0, 0.0) == (0.0, 0.0) + assert shadow._walker_first_seen_s == {7: 10.0} diff --git a/tests/unit_tests/test_shadow_ego_yellow_margin.py b/tests/unit_tests/test_shadow_ego_yellow_margin.py new file mode 100644 index 0000000000..444b1d5c29 --- /dev/null +++ b/tests/unit_tests/test_shadow_ego_yellow_margin.py @@ -0,0 +1,34 @@ +import math +from types import SimpleNamespace + +from pufferlib.ocean.cosim import carla_bridge as cb +from pufferlib.ocean.cosim.carla.shadow_ego import ShadowEgo + + +YELLOW_LIGHT_ID = 7 +YELLOW_ELEMENT_IDX = 1 + + +def _reported_yellow_seconds(yellow_margin_s, carla_remaining_s): + shadow = SimpleNamespace( + num_traffic=2, + light_bin_indices_by_id={YELLOW_LIGHT_ID: [YELLOW_ELEMENT_IDX]}, + _yellow_margin_s=yellow_margin_s, + _yellow_remaining_by_id=lambda snapshot, state_by_id: {YELLOW_LIGHT_ID: carla_remaining_s}, + ) + snapshot = {"light_states": [[YELLOW_LIGHT_ID, "Yellow"]], "ego_light_id": YELLOW_LIGHT_ID} + states, yellow_remaining_s = ShadowEgo._read_light_states(shadow, snapshot) + assert states[YELLOW_ELEMENT_IDX] == cb.TRAFFIC_LIGHT_STATE_YELLOW + return float(yellow_remaining_s[YELLOW_ELEMENT_IDX]) + + +def test_without_a_margin_the_policy_sees_the_yellow_time_carla_has_left(): + assert _reported_yellow_seconds(0.0, 2.0) == 2.0 + + +def test_the_margin_shortens_the_yellow_time_the_policy_sees(): + assert math.isclose(_reported_yellow_seconds(0.5, 2.0), 1.5) + + +def test_the_reported_yellow_time_never_goes_negative(): + assert _reported_yellow_seconds(0.5, 0.2) == 0.0 diff --git a/tests/unit_tests/test_tidy_alpasim_run.py b/tests/unit_tests/test_tidy_alpasim_run.py new file mode 100644 index 0000000000..eece73167b --- /dev/null +++ b/tests/unit_tests/test_tidy_alpasim_run.py @@ -0,0 +1,288 @@ +"""tidy_alpasim_run.py flattens an AlpaSim run to results.txt, run_info.txt, obs_html/ and details/, migrates the +2026-10-12 layout and leaves an unaggregated run's wizard tree; render_alpasim_obs_html.py pages the imperfect +scenes by default and re-renders a pruned run from its pages.""" + +import json +import os +import sys + +import numpy as np +import pandas as pd +import yaml + +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_alpasim_obs_html as render +import tidy_alpasim_run as tidy + + +FRAMES, AGENT_CAP = 20, 2 +FAIL, SLOW, CLEAN = ( + "clipgt-aaaaaaaa-0000-0000-0000-000000000001", + "clipgt-aaaaaaaa-0000-0000-0000-000000000002", + "clipgt-aaaaaaaa-0000-0000-0000-000000000003", +) +# scene -> (score, status, at-fault collision, progress score) +SCENES = {FAIL: (0.0, "fail", 1.0, 1.0), SLOW: (0.5, "pass", 0.0, 0.5), CLEAN: (1.0, "pass", 0.0, 1.0)} +SESSION = {FAIL: "a" * 32, SLOW: "b" * 32, CLEAN: "c" * 32} +SERIES_SECONDS = np.arange(FRAMES) * 0.1 +INFO_LINE = b"runtime-0-1 | 2026-11-08 08:30:00 INFO: sim event\n" +HOST_LINE = b"[2026-11-08 08:29:11][alpasim_wizard][WARNING] - host line\n" + + +def _replay(): + 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, 0] = np.arange(FRAMES) + return replay + + +def _row(clipgt): + score, status, collision, progress = SCENES[clipgt] + metrics = { + "collision_any": collision, + "collision_front": collision, + "collision_lateral": 0.0, + "collision_rear": 0.0, + "wrong_lane": 0.0, + "dist_traveled_m": 100.0 * progress, + "gt_dist_traveled_m": 100.0, + "dist_to_gt_trajectory": 1.0, + "lateral_dist_to_gt_trajectory": 0.5, + } + score_metrics = { + "collision_at_fault": collision, + "offroad": 0.0, + "left_corridor_laterally": 0.0, + "progress_score": progress, + "progress_clipped_rel": progress * 0.8, + } + return { + "clipgt_id": clipgt, + "rollout_id": SESSION[clipgt], + "score": score, + "status": status, + "failure_reason": "collision_at_fault" if status == "fail" else "", + "metrics": metrics, + "score_metrics": score_metrics, + } + + +def _summary(): + return { + "rollouts": [_row(c) for c in SCENES], + "score_criteria": {"collision_at_fault": "== 0"}, + "metrics_results": [], + } + + +def _metrics_frame(): + rows = [] + for clipgt in SCENES: + for name in render.PER_STEP_SERIES: + for step, t in enumerate(SERIES_SECONDS): + value = 0.0 + if name == "eval_relevant" and step >= 2: + value = 1.0 + if name in ("collision_any", "collision_front") and clipgt == FAIL and step >= 10: + value = 1.0 + if name == "dist_to_gt_trajectory": + value = 1.0 + rows.append((name, t * 1e6, value, clipgt, SESSION[clipgt])) + return pd.DataFrame(rows, columns=["name", "timestamps_us", "values", "clipgt_id", "rollout_id"]) + + +def _write_fresh_run(root, aggregated=True, source_exists=True): + """A finished 12_alpasim.sh run before the tidy: the host's alpagym/ tree with the wizard underneath.""" + run = root / "alpasim_gt_time5_r2_final_model_20261109_120000_7" + alpagym, wizard = run / "alpagym", run / "alpagym" / "alpasim" / "wizard_0" + (alpagym / "logs").mkdir(parents=True) + (alpagym / "config.json").write_text("{}") + (alpagym / "evaluation.yaml").write_text("scenes: []\n") + (alpagym / "logs" / "evaluation.log").write_text("evaluation\n") + obs = alpagym / "obs_replays" + obs.mkdir() + for clipgt in SCENES: + viz.save_interactive_replay_zlib({}, _replay(), str(obs / f"{clipgt}.replay.zlib")) + (obs / f"{clipgt}.html").write_text("page the run rendered") + for clipgt in SCENES: + rollout = wizard / "rollouts" / clipgt / SESSION[clipgt] + rollout.mkdir(parents=True) + (rollout / "metrics.parquet").write_bytes(b"per-rollout metrics") + (rollout / "_complete").write_text("") + (rollout / "rollout.asl").write_bytes(b"rollout log " + clipgt.encode()) + (wizard / "controller").mkdir() + (wizard / "controller" / "alpasim_controller_x.csv").write_text("t,steer\n") + (wizard / "txt-logs").mkdir() + (wizard / "txt-logs" / "out-7-attempt-0-physics-0-log.txt").write_bytes(b"[INFO] tick\nphysics warning\n") + (wizard / "txt-logs" / "runtime_worker_0.log").write_bytes(b"worker copy\n" * 10) + (wizard / "prometheus").mkdir() + (wizard / "prometheus" / "dcgm-counters.csv").write_text("x") + for name in ("eval-config.yaml", "wizard-config.yaml", "docker-compose.yaml", "generated-user-config-0.yaml"): + (wizard / name).write_text(yaml.safe_dump({"aggregation_modifiers": {"max_dist_to_gt_trajectory": 4.0}})) + (wizard / "failed_rollouts.json").write_text("[]") + if aggregated: + (wizard / "aggregate").mkdir() + (wizard / "aggregate" / "results-summary.json").write_text(json.dumps(_summary())) + _metrics_frame().to_parquet(wizard / "aggregate" / "metrics_unprocessed.parquet") + (wizard / "aggregate" / "metrics_results.png").write_bytes(b"png") + source = root / "k_scaled_0099_1000" / "final_model.pt" + source.parent.mkdir(exist_ok=True) + if source_exists: + source.write_bytes(b"weights") + (run / "final_model.pt").write_bytes(b"weights") + (run / "config.yaml").write_text("train: {}\n") + (run / "run_info.txt").write_text(f"checkpoint {source} (evaluated copy: {run / 'final_model.pt'})\nscenes 3\n") + (run / "results.txt").write_text("official_score 0.5000 scenes 3 failed 0\n") + (run / "nre_rc.local").write_text('exec "$@"\n') + (run / "evaluate.log").write_bytes(HOST_LINE + INFO_LINE * 5 + HOST_LINE) + return run + + +def _tree(root): + return sorted(str(path.relative_to(root)) for path in root.rglob("*")) + + +def _tidy(run, *args, monkeypatch): + monkeypatch.setattr(sys, "argv", ["tidy_alpasim_run", str(run), *args]) + tidy.main() + + +def _render(run, *args, monkeypatch): + monkeypatch.setattr(sys, "argv", ["render_alpasim_obs_html", str(run), *args]) + render.main() + + +def test_tidy_flattens_a_fresh_run(tmp_path, monkeypatch): + run = _write_fresh_run(tmp_path) + _tidy(run, monkeypatch=monkeypatch) + assert sorted(path.name for path in run.iterdir()) == ["details", "obs_html", "results.txt", "run_info.txt"] + assert sorted(path.name for path in (run / "details").iterdir()) == [ + "config.json", + "eval-config.yaml", + "evaluate.log", + "evaluation.log", + "evaluation.yaml", + "metrics_unprocessed.parquet", + "results-summary.json", + "services.log", + "wizard-config.yaml", + ] + assert (run / "details" / "evaluate.log").read_bytes() == HOST_LINE + HOST_LINE + services = (run / "details" / "services.log").read_text() + assert "==== out-7-attempt-0-physics-0-log.txt ====" in services and "physics warning" in services + assert "[INFO]" not in services and "worker copy" not in services + assert sorted(path.name for path in (run / "obs_html").iterdir()) == sorted( + [f"{c}.replay.zlib" for c in SCENES] + [f"{c}.html" for c in SCENES] + ) + before = _tree(run) + _tidy(run, monkeypatch=monkeypatch) + assert _tree(run) == before + + +def test_tidy_keeps_the_checkpoint_copy_when_its_source_is_gone(tmp_path, monkeypatch): + run = _write_fresh_run(tmp_path, source_exists=False) + _tidy(run, monkeypatch=monkeypatch) + assert (run / "final_model.pt").is_file() and (run / "config.yaml").is_file() + + +def test_tidy_keeps_the_selected_rollout_logs_where_the_offline_tools_look(tmp_path, monkeypatch): + run = _write_fresh_run(tmp_path) + _tidy(run, "--keep_rollout_logs", "imperfect", monkeypatch=monkeypatch) + kept = sorted(str(path.relative_to(run)) for path in run.rglob("rollout.asl")) + assert kept == [f"alpasim/wizard_0/rollouts/{c}/{SESSION[c]}/rollout.asl" for c in (FAIL, SLOW)] + + +def test_tidy_keeps_the_wizard_tree_of_an_unaggregated_run(tmp_path, monkeypatch, capsys): + run = _write_fresh_run(tmp_path, aggregated=False) + _tidy(run, monkeypatch=monkeypatch) + assert "never aggregated" in capsys.readouterr().out + wizard = run / "alpasim" / "wizard_0" + assert len(list(wizard.rglob("metrics.parquet"))) == 3 and (wizard / "eval-config.yaml").is_file() + assert not (wizard / "controller").exists() and not (wizard / "txt-logs").exists() + assert (run / "details" / "services.log").is_file() and not (run / "details" / "results-summary.json").exists() + + +def test_tidy_migrates_the_2026_10_layout(tmp_path, monkeypatch): + run = tmp_path / "alpasim_gt_time5_r2_final_model_20261012_120000_1" + (run / "obs_replays").mkdir(parents=True) + (run / "obs_replays" / f"s0.00_COL_{FAIL}.html").write_text("page") + (run / "logs" / "txt-logs").mkdir(parents=True) + for name in ("evaluate.log", "config.json", "evaluation.yaml", "evaluation.log", "run_info.txt"): + (run / "logs" / name).write_text(name) + (run / "logs" / "txt-logs" / "runtime_worker_3.log").write_text("copy") + (run / "logs" / "txt-logs" / "out-1-attempt-0-renderer-0-log.txt").write_text("renderer line\n") + (run / "alpasim").mkdir() + for name in ("results-summary.json", "metrics_unprocessed.parquet", "eval-config.yaml", "wizard-config.yaml"): + (run / "alpasim" / name).write_text(name) + (run / "results.txt").write_text("official_score 0.7\n") + _tidy(run, monkeypatch=monkeypatch) + assert sorted(path.name for path in run.iterdir()) == ["details", "obs_html", "results.txt", "run_info.txt"] + assert sorted(path.name for path in (run / "details").iterdir()) == [ + "config.json", + "eval-config.yaml", + "evaluate.log", + "evaluation.log", + "evaluation.yaml", + "metrics_unprocessed.parquet", + "results-summary.json", + "services.log", + "wizard-config.yaml", + ] + assert "renderer line" in (run / "details" / "services.log").read_text() + + +def test_gallery_pages_the_imperfect_scenes_and_prunes_the_replays(tmp_path, monkeypatch): + run = _write_fresh_run(tmp_path) + _tidy(run, monkeypatch=monkeypatch) + _render(run, "--prune", monkeypatch=monkeypatch) + assert sorted(path.name for path in (run / "obs_html").iterdir()) == [ + "index.html", + f"s0.00_COL_{FAIL}.html", + f"s0.50_PROG_{SLOW}.html", + ] + page = (run / "obs_html" / f"s0.00_COL_{FAIL}.html").read_text() + assert "AlpaSim score 0.00" in page and "page the run rendered" not in page + # a pruned run renders again from its pages + _render(run, monkeypatch=monkeypatch) + assert sorted(path.name for path in (run / "obs_html").iterdir()) == [ + "index.html", + f"s0.00_COL_{FAIL}.html", + f"s0.50_PROG_{SLOW}.html", + ] + + +def test_gallery_select_failed_and_all(tmp_path, monkeypatch): + run = _write_fresh_run(tmp_path) + _tidy(run, monkeypatch=monkeypatch) + _render(run, "--select", "failed", monkeypatch=monkeypatch) + assert sorted(path.name for path in (run / "obs_html").glob("s*.html")) == [f"s0.00_COL_{FAIL}.html"] + assert not list((run / "obs_html").glob("clipgt-*.html")) + other = _write_fresh_run(tmp_path / "other") + _tidy(other, monkeypatch=monkeypatch) + _render(other, "--select", "all", monkeypatch=monkeypatch) + assert len(list((other / "obs_html").glob("s*.html"))) == 3 diff --git a/tests/unit_tests/test_tidy_carla_run.py b/tests/unit_tests/test_tidy_carla_run.py new file mode 100644 index 0000000000..f3840d1a04 --- /dev/null +++ b/tests/unit_tests/test_tidy_carla_run.py @@ -0,0 +1,251 @@ +"""tidy_carla_run.py flattens a longest6 run to results.csv, run_info.txt, obs_html/, routes/ (flat) and details/, +keeping the telemetry and world logs of the failed route runs only; render_carla_obs_html.py pages the failed route +runs by default and renders a tidied run from its pages; carla_run_layout resolves both layouts.""" + +import json +import os +import sys + +import numpy as np + +from pufferlib import viz +from pufferlib.ocean.cosim.carla.shadow_ego import TELEMETRY_HEADER + + +sys.path.insert( + 0, os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "scripts", "eval") +) + +import carla_run_layout as layout +import render_carla_obs_html as render +import tidy_carla_run as tidy + + +FRAMES, AGENT_CAP = 40, 2 +# route run -> (status, infractions): a collision, a perfect route, a min-speed-only route, a crashed route +ROUTE_RUNS = { + "route_03_s0": ( + "Completed", + { + "collisions_vehicle": [ + "Agent collided against object with type=vehicle.audi.tt and id=42 at (x=1.0, y=2.0, z=0.0)" + ] + }, + ), + "route_03_s1": ("Completed", {}), + "route_07_s0": ("Completed", {"min_speed_infractions": ["Agent drove too slow"] * 3}), + "route_07_s1": ("Failed - Simulation crashed", {}), +} +SCORES = {"route_03_s0": 50.0, "route_03_s1": 100.0, "route_07_s0": 80.5, "route_07_s1": 0.0} +TOWNS = {3: "Town01", 7: "Town02"} + + +def _record(route_run): + status, infractions = ROUTE_RUNS[route_run] + route_id = int(route_run.split("_")[1]) + return { + "route_id": f"RouteScenario_{route_id}", + "status": status, + "infractions": infractions, + "scores": {"score_composed": SCORES[route_run], "score_route": 100.0, "score_penalty": SCORES[route_run] / 100}, + "meta": {"route_length": 1000.0, "duration_game": 300.0, "duration_system": 150.0}, + } + + +def _replay(): + 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, 0] = np.arange(FRAMES) + return replay + + +def _write_attempt(route_dir, stamp, with_replay=True): + """The files one evaluator attempt leaves: telemetry, world log, replay (unless it crashed first) and logs.""" + for sub in ("telemetry", "world_log", "obs_html"): + (route_dir / sub).mkdir(exist_ok=True) + (route_dir / "telemetry" / f"route_{stamp}.csv").write_text(TELEMETRY_HEADER + "0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0\n") + meta = json.dumps({"town": "Town01", "tick_dt": 0.05, "offset": [0.0, 0.0]}) + np.savez_compressed( + route_dir / "world_log" / f"route_{stamp}.npz", ego=np.zeros((FRAMES, 7), np.float32), meta=meta + ) + if with_replay: + viz.save_interactive_replay_zlib({}, _replay(), str(route_dir / "obs_html" / f"route_{stamp}.replay.zlib")) + + +def _write_run(root): + """A longest6 run as the scripts leave it before the gallery and the tidy: two routes x two seeds.""" + run = root / "carla_longest6_v2_seeds0-1_20261109_120000_7" + for route_run in ROUTE_RUNS: + route_dir = run / "routes" / route_run + route_dir.mkdir(parents=True) + _write_attempt(route_dir, "20261109_120100_000001") + (route_dir / "evaluator_attempt1.log").write_text(f"[puffer_agent] route done {route_run}\n") + (route_dir / "carla_server_attempt1.log").write_text("4.24.3 UE4\n") + (route_dir / "result.json").write_text(json.dumps({"_checkpoint": {"records": [_record(route_run)]}})) + # the crashed route was retried: a first attempt without a replay, server and evaluator logs of both + retried = run / "routes" / "route_07_s1" + _write_attempt(retried, "20261109_120000_000000", with_replay=False) + (retried / "evaluator_attempt2.log").write_text("second attempt\n") + (retried / "carla_server_attempt2.log").write_text("4.24.3 UE4\n") + aggregate = run / "aggregate" + for route_run in ROUTE_RUNS: + seed_dir = aggregate / f"s{route_run[-1]}" + seed_dir.mkdir(parents=True, exist_ok=True) + (seed_dir / f"{route_run[:-3]}.json").write_text((run / "routes" / route_run / "result.json").read_text()) + (aggregate / "results.csv").write_text('"sep=,"\nAvg. driving score,57.625\n') + (run / "run_info.txt").write_text("checkpoint /x/final_model.pt\nseeds 0 1\n") + (run / "timing.csv").write_text( + "route,worker,gpu,attempt,start,wall_s,game_s,pace,status,driving_score,seed\n3,0,0,1,12:01:00,100,300.0,3.0,Completed,50.00,0\n" + ) + (run / "server_pids").write_text("1234\n") + for worker in (0, 1, 10): + (run / f"worker_w{worker}.log").write_text(f"[w{worker}] queue empty\n") + (run / "compat_lib").mkdir() + (run / "compat_lib" / "libtiff.so.5").symlink_to("/usr/lib/libtiff.so.6") + return run + + +def _tree(root): + return sorted(str(path.relative_to(root)) for path in root.rglob("*")) + + +def _render(run, *args, monkeypatch): + monkeypatch.setattr(sys, "argv", ["render_carla_obs_html", str(run), *args]) + render.main() + + +def _tidy(run, *args, monkeypatch): + monkeypatch.setattr(sys, "argv", ["tidy_carla_run", str(run), *args]) + tidy.main() + + +def test_route_failed_flags_infractions_and_failed_statuses_only(): + assert layout.route_failed(_record("route_03_s0")) + assert layout.route_failed(_record("route_07_s1")) + assert not layout.route_failed(_record("route_03_s1")) + assert not layout.route_failed(_record("route_07_s0")) # min-speed infractions alone + + +def test_flat_stems_name_repetitions_before_the_seed(): + assert layout.flat_stem("route_08_s1", 0) == "route_08_s1" + assert layout.flat_stem("route_08_s1", 2) == "route_08_rep2_s1" + assert layout.flat_stem("route_08", 1) == "route_08_rep1" + assert layout.seed_of("route_08_s1") == "1" and layout.seed_of("route_08") is None + + +def test_layout_pairs_the_newest_attempt_files_with_the_records(tmp_path): + run = _write_run(tmp_path) + runs = dict(layout.route_runs(run)) + assert sorted(runs) == sorted(ROUTE_RUNS) + retried = runs["route_07_s1"] + assert layout.record_stems(retried, 1) == ["route_20261109_120100_000001"] + assert layout.record_file(retried, "obs_html", ".replay.zlib", "route_20261109_120100_000001").is_file() + assert layout.record_stems(retried, 3) == [None, "route_20261109_120000_000000", "route_20261109_120100_000001"] + + +def test_gallery_pages_the_failed_route_runs_only(tmp_path, monkeypatch): + run = _write_run(tmp_path) + _render(run, monkeypatch=monkeypatch) + pages = sorted(path.name for path in (run / "obs_html").iterdir()) + assert pages == ["ds000.0_Town01_route07_s1.html", "ds050.0_COL_Town01_route03_s0.html", "index.html"] + _render(run, "--select", "all", monkeypatch=monkeypatch) + assert len(list((run / "obs_html").glob("ds*.html"))) == 4 + + +def test_tidy_flattens_the_run_and_keeps_the_failed_routes_logs(tmp_path, monkeypatch): + run = _write_run(tmp_path) + _render(run, monkeypatch=monkeypatch) + _tidy(run, monkeypatch=monkeypatch) + assert sorted(path.name for path in run.iterdir()) == [ + "details", + "obs_html", + "results.csv", + "routes", + "run_info.txt", + ] + assert sorted(path.name for path in (run / "routes").iterdir()) == [ + "route_03_s0.result.json", + "route_03_s0.telemetry.csv", + "route_03_s0.world_log.npz", + "route_03_s1.result.json", + "route_07_s0.result.json", + "route_07_s1.result.json", + "route_07_s1.telemetry.csv", + "route_07_s1.world_log.npz", + ] + assert sorted(path.name for path in (run / "details").iterdir()) == ["evaluators.log", "timing.csv", "workers.log"] + evaluators = (run / "details" / "evaluators.log").read_text() + assert "==== route_07_s1 evaluator_attempt1.log ====" in evaluators and "second attempt" in evaluators + assert (run / "details" / "workers.log").read_text().index("[w1]") < ( + run / "details" / "workers.log" + ).read_text().index("[w10]") + assert (run / "results.csv").read_text().startswith('"sep=,"') + assert len(list((run / "obs_html").glob("ds*.html"))) == 2 + # the layout resolves the flat run; the kept world log pairs with its record + runs = dict(layout.route_runs(run)) + assert sorted(runs) == sorted(ROUTE_RUNS) + assert layout.record_stems(runs["route_03_s0"], 1) == ["route_03_s0"] + assert ( + layout.record_file(runs["route_03_s0"], "world_log", ".npz", "route_03_s0").name == "route_03_s0.world_log.npz" + ) + assert layout.record_file(runs["route_03_s1"], "world_log", ".npz", "route_03_s1") is None + before = _tree(run) + _tidy(run, monkeypatch=monkeypatch) + assert _tree(run) == before + + +def test_tidy_keep_all_and_none(tmp_path, monkeypatch): + run = _write_run(tmp_path) + _tidy(run, "--keep", "all", monkeypatch=monkeypatch) + assert len(list((run / "routes").glob("*.world_log.npz"))) == 4 + other = _write_run(tmp_path / "other") + _tidy(other, "--keep", "none", monkeypatch=monkeypatch) + assert sorted(path.name for path in (other / "routes").iterdir()) == [ + f"{r}.result.json" for r in sorted(ROUTE_RUNS) + ] + + +def test_tidied_run_renders_again_from_its_pages(tmp_path, monkeypatch): + run = _write_run(tmp_path) + _render(run, monkeypatch=monkeypatch) + _tidy(run, monkeypatch=monkeypatch) + page = run / "obs_html" / "ds050.0_COL_Town01_route03_s0.html" + page.write_text(page.read_text().replace("", "", 1)) + _render(run, monkeypatch=monkeypatch) + assert "" not in page.read_text() + assert sorted(path.name for path in (run / "obs_html").iterdir()) == [ + "ds000.0_Town01_route07_s1.html", + "ds050.0_COL_Town01_route03_s0.html", + "index.html", + ] + + +def test_a_route_driven_again_replaces_its_old_page(tmp_path, monkeypatch): + run = _write_run(tmp_path) + _render(run, monkeypatch=monkeypatch) + record = _record("route_03_s0") + record["infractions"] = {} + record["status"] = "Failed - Agent got blocked" + (run / "routes" / "route_03_s0" / "result.json").write_text(json.dumps({"_checkpoint": {"records": [record]}})) + _render(run, monkeypatch=monkeypatch) + pages = sorted(path.name for path in (run / "obs_html").iterdir()) + assert pages == ["ds000.0_Town01_route07_s1.html", "ds050.0_Town01_route03_s0.html", "index.html"] diff --git a/tests/unit_tests/test_tidy_nuplan_run.py b/tests/unit_tests/test_tidy_nuplan_run.py new file mode 100644 index 0000000000..b98342be43 --- /dev/null +++ b/tests/unit_tests/test_tidy_nuplan_run.py @@ -0,0 +1,187 @@ +"""tidy_nuplan_run.py flattens a nuPlan run to overview.csv, scenarios.csv, obs_html/, simulation_log/ and details/ +and keeps the logs (and videos, when rendered) of the selected scenarios only; nuplan_run_layout resolves the fresh +and the tidied layout.""" + +import os +import sys + +import pandas as pd + +sys.path.insert( + 0, os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "scripts", "eval") +) + +import nuplan_run_layout +import tidy_nuplan_run as tidy + +METRICS = ( + "no_ego_at_fault_collisions", + "drivable_area_compliance", + "driving_direction_compliance", + "ego_is_making_progress", + "time_to_collision_within_bound", + "ego_progress_along_expert_route", + "speed_limit_compliance", + "ego_is_comfortable", +) +# token -> (score, at-fault collision metric): an infraction, a failure without infraction, a perfect scenario +COLLIDING, SLOW, CLEAN = "aaaa000000000001", "aaaa000000000002", "aaaa000000000003" +SCENARIOS = {COLLIDING: (0.0, 0.0), SLOW: (0.5, 1.0), CLEAN: (1.0, 1.0)} +CHALLENGE = "closed_loop_reactive_agents" +TIMESTAMP = "2026.01.01.00.00.00" + + +def _aggregator_table(): + rows = [] + for token, (score, collision) in SCENARIOS.items(): + row = {"scenario": token, "log_name": "log", "scenario_type": "turn", "score": score, "num_scenarios": None} + row.update({metric: 1.0 for metric in METRICS}) + row["no_ego_at_fault_collisions"] = collision + rows.append(row) + final = {"scenario": "final_score", "log_name": None, "scenario_type": "final_score", "score": 0.5} + rows.append({**final, "num_scenarios": len(SCENARIOS), **{metric: 1.0 for metric in METRICS}}) + return pd.DataFrame(rows) + + +def _write_run(group, videos=False): + """nuPlan's tree after a run with simulation_log_callback; videos=True adds carl_visualization_callback's .avi.""" + sim = group / "simulation" / f"{CHALLENGE}_pufferdrive" / TIMESTAMP + aggregator = sim / "aggregator_metric" + aggregator.mkdir(parents=True) + table = _aggregator_table() + table.to_csv(aggregator / f"{CHALLENGE}_weighted_average_metrics_{TIMESTAMP}.csv", index=False) + table.to_parquet(aggregator / f"{CHALLENGE}_weighted_average_metrics_{TIMESTAMP}.parquet") + (aggregator / "overview.csv").write_text("CLS\n0.5\n") + (aggregator / "scenarios.csv").write_text( + "scenario,score\n" + "".join(f"{t},{s}\n" for t, (s, _) in SCENARIOS.items()) + ) + (sim / "metrics").mkdir() + comfort = pd.DataFrame({"scenario_name": list(SCENARIOS), "metric_score": [0.0, 1.0, 1.0]}) + comfort.to_parquet(sim / "metrics" / "ego_is_comfortable.parquet") + for token in SCENARIOS: + log_dir = sim / "simulation_log" / "PufferDrivePlanner" / "turn" / "log" / token + log_dir.mkdir(parents=True) + (log_dir / f"{token}.msgpack.xz").write_bytes(b"log " + token.encode()) + if videos: + (sim / "visualization").mkdir() + for token in SCENARIOS: + (sim / "visualization" / f"{token}.avi").write_bytes(b"avi " + token.encode()) + (sim / "code" / "hydra").mkdir(parents=True) + for name in ("config.yaml", "overrides.yaml", "hydra.yaml"): + (sim / "code" / "hydra" / name).write_text(name) + (sim / "runner_report.parquet").write_bytes(b"report") + (sim / "log.txt").write_text("log") + (sim / "nuboard_1.nuboard").write_bytes(b"nuboard") + (sim / "logs").mkdir() + obs = group / "obs_html" + obs.mkdir() + (obs / f"s0.00_COL_turn_{COLLIDING}.html").write_text("page") + (obs / f"{COLLIDING}.replay.zlib").write_bytes(b"replay") + (obs / f"{SLOW}.replay.zlib").write_bytes(b"replay") + (obs / "index.html").write_text("index") + return sim + + +def _copy_instead_of_transcode(avi_path, mp4_path, ffmpeg): + mp4_path.parent.mkdir(parents=True, exist_ok=True) + mp4_path.write_bytes(avi_path.read_bytes()) + + +def _tree(root): + return sorted(str(path.relative_to(root)) for path in root.rglob("*")) + + +def _run(group, *args, monkeypatch): + monkeypatch.setattr(tidy, "transcode_video", _copy_instead_of_transcode) + monkeypatch.setattr(sys, "argv", ["tidy_nuplan_run", str(group), *args]) + tidy.main() + + +def test_layout_resolves_the_fresh_run(tmp_path): + sim = _write_run(tmp_path) + assert nuplan_run_layout.details_dir(tmp_path) == sim + assert nuplan_run_layout.metrics_dir(tmp_path) == sim / "metrics" + assert nuplan_run_layout.aggregated_metrics_csv(tmp_path).parent == sim / "aggregator_metric" + assert nuplan_run_layout.scenarios_csv(tmp_path) == sim / "aggregator_metric" / "scenarios.csv" + assert nuplan_run_layout.simulation_log_dir(tmp_path) == sim / "simulation_log" + assert set(nuplan_run_layout.simulation_logs(tmp_path)) == set(SCENARIOS) + + +def test_tidy_keeps_infractions_only(tmp_path, monkeypatch): + _write_run(tmp_path) + _run(tmp_path, monkeypatch=monkeypatch) + details = tmp_path / "details" + assert sorted(p.name for p in tmp_path.iterdir()) == [ + "details", + "obs_html", + "overview.csv", + "scenarios.csv", + "simulation_log", + ] + assert sorted(p.name for p in details.iterdir()) == [ + "aggregated_metrics.csv", + "config.yaml", + "metrics", + "overrides.yaml", + "runner_report.parquet", + ] + assert [p.name for p in (tmp_path / "simulation_log").iterdir()] == [f"{COLLIDING}.msgpack.xz"] + assert (details / "metrics" / "ego_is_comfortable.parquet").is_file() + assert (tmp_path / "overview.csv").read_text() == "CLS\n0.5\n" + # the replay of the rendered page goes, the replay without a page stays + assert sorted(p.name for p in (tmp_path / "obs_html").iterdir()) == [ + f"{SLOW}.replay.zlib", + "index.html", + f"s0.00_COL_turn_{COLLIDING}.html", + ] + assert nuplan_run_layout.details_dir(tmp_path) == details + assert nuplan_run_layout.metrics_dir(tmp_path) == details / "metrics" + assert nuplan_run_layout.aggregated_metrics_csv(tmp_path) == details / "aggregated_metrics.csv" + assert nuplan_run_layout.scenarios_csv(tmp_path) == tmp_path / "scenarios.csv" + assert nuplan_run_layout.simulation_log_dir(tmp_path) == tmp_path / "simulation_log" + assert list(nuplan_run_layout.simulation_logs(tmp_path)) == [COLLIDING] + + +def test_tidy_failures_mode_keeps_low_scores_and_their_videos(tmp_path, monkeypatch): + _write_run(tmp_path, videos=True) + _run(tmp_path, "--keep", "failures", monkeypatch=monkeypatch) + assert set(nuplan_run_layout.simulation_logs(tmp_path)) == {COLLIDING, SLOW} + assert sorted(p.stem for p in (tmp_path / "details" / "videos").iterdir()) == sorted([COLLIDING, SLOW]) + assert not (tmp_path / "simulation").exists() + + +def test_layout_resolves_a_run_tidied_into_details_before_2026_11(tmp_path): + old_log_dir = tmp_path / "details" / "simulation_log" + old_log_dir.mkdir(parents=True) + (old_log_dir / f"{COLLIDING}.msgpack.xz").write_bytes(b"log") + assert nuplan_run_layout.simulation_log_dir(tmp_path) == old_log_dir + assert list(nuplan_run_layout.simulation_logs(tmp_path)) == [COLLIDING] + + +def test_tidy_leaves_unknown_files_and_reports_them(tmp_path, monkeypatch, capsys): + sim = _write_run(tmp_path) + (sim / "extra.bin").write_bytes(b"?") + _run(tmp_path, monkeypatch=monkeypatch) + assert _tree(tmp_path / "simulation") == [ + f"{CHALLENGE}_pufferdrive", + f"{CHALLENGE}_pufferdrive/{TIMESTAMP}", + f"{CHALLENGE}_pufferdrive/{TIMESTAMP}/extra.bin", + ] + assert "left in place, unknown" in capsys.readouterr().out + + +def test_tidy_is_a_no_op_the_second_time_and_on_an_unscored_run(tmp_path, monkeypatch, capsys): + _write_run(tmp_path) + _run(tmp_path, monkeypatch=monkeypatch) + before = _tree(tmp_path) + _run(tmp_path, monkeypatch=monkeypatch) + assert _tree(tmp_path) == before + assert "already tidied" in capsys.readouterr().out + unscored = tmp_path / "unscored" + sim = _write_run(unscored) + for path in (sim / "aggregator_metric").iterdir(): + path.unlink() + before = _tree(unscored) + _run(unscored, monkeypatch=monkeypatch) + assert _tree(unscored) == before + assert "never scored" in capsys.readouterr().out diff --git a/tests/unit_tests/test_tidy_selfplay_eval.py b/tests/unit_tests/test_tidy_selfplay_eval.py new file mode 100644 index 0000000000..44a99c6736 --- /dev/null +++ b/tests/unit_tests/test_tidy_selfplay_eval.py @@ -0,0 +1,205 @@ +"""tidy_selfplay_eval.py folds the shard folders of a parallel self-play eval into evaluation_summary.json, +obs_html/ and details/; EvalReplayCapture keeps the observations of the infracting agents only when asked.""" + +import json +import os +import pickle +import sys +import zlib + +import numpy as np +import pandas as pd +import pytest + +import pufferlib +from pufferlib import viz +from pufferlib.ocean.evaluation_utils import eval_replay + + +sys.path.insert( + 0, os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "scripts", "eval") +) + +import tidy_selfplay_eval as tidy + + +STAMP = "20261109-120000" +# shard -> failing (map, seed) pairs it re-simulated and paged +SHARD_FAILURES = { + "shard00": [("opendrive__Town02", 11)], + "shard01": [("opendrive__Town07", 21), ("opendrive__Town07", 22)], +} +METRIC_COLUMNS = ("collision_rate", "offroad_rate", "episode_length") + + +def _episode_rows(pairs, shard_dir=None): + rows = [] + for episode_id, (map_name, seed) in enumerate(pairs): + row = { + "map_name": map_name, + "scenario_id": map_name.split("__")[1], + "seed": seed, + "collision_rate": 0.02, + "offroad_rate": 0.0, + "episode_length": 9000.0, + "agents_per_batch": 50, + } + if shard_dir is not None: + stem = f"{map_name}__seed_{seed}__episode_{episode_id:06d}" + row["has_replay"] = 1 + row["replay_path"] = f"/cephfs/{shard_dir}/failures/replays_zlib/{stem}.replay.zlib" + rows.append(row) + return pd.DataFrame(rows) + + +def _write_parallel_eval(root): + """The aggregate dir and two shard folders as parallel_eval.py leaves them.""" + benchmark_dir = root / "eval" / "carla_k_scaled_0099_1000" + run_dir = benchmark_dir / STAMP + run_dir.mkdir(parents=True) + (run_dir / "evaluation_summary.json").write_text(json.dumps({"num_scenarios": 20, "num_episodes": 20})) + all_rows = pd.concat( + [ + _episode_rows([("opendrive__Town02", s) for s in range(10)]), + _episode_rows([("opendrive__Town07", s) for s in range(10)]), + ] + ) + all_rows.to_csv(run_dir / "episode_metrics.csv", index=False) + for shard, pairs in SHARD_FAILURES.items(): + (run_dir / f"{shard}.log").write_text(f"{shard} log line\n") + shard_dir = benchmark_dir / f"{STAMP}_{shard}" + failures = shard_dir / "failures" + (failures / "rendered_replays").mkdir(parents=True) + (shard_dir / "resolved_benchmark.yaml").write_text(f"shard: {shard}\n") + _episode_rows([("opendrive__Town02", s) for s in range(10)]).to_csv( + shard_dir / "episode_metrics.csv", index=False + ) + (shard_dir / "evaluation_summary.json").write_text("{}") + (failures / "evaluation_summary.json").write_text("{}") + _episode_rows(pairs).to_csv(failures / "selected_failures.csv", index=False) + _episode_rows(pairs, shard_dir=shard).to_csv(failures / "episode_metrics.csv", index=False) + (failures / "rendered_replays" / "index.html").write_text("shard index") + for episode_id, (map_name, seed) in enumerate(pairs): + (failures / "rendered_replays" / f"{map_name}__seed_{seed}__episode_{episode_id:06d}.html").write_text( + "page" + ) + return run_dir + + +def _tidy(run_dir, monkeypatch): + monkeypatch.setattr(sys, "argv", ["tidy_selfplay_eval", str(run_dir)]) + tidy.main() + + +def test_tidy_folds_the_shards_into_the_aggregate_dir(tmp_path, monkeypatch): + run_dir = _write_parallel_eval(tmp_path) + _tidy(run_dir, monkeypatch) + assert sorted(path.name for path in run_dir.parent.iterdir()) == [STAMP] + assert sorted(path.name for path in run_dir.iterdir()) == ["details", "evaluation_summary.json", "obs_html"] + assert sorted(path.name for path in (run_dir / "details").iterdir()) == [ + "episode_metrics.csv", + "failure_episodes.csv", + "resolved_benchmark.yaml", + "selected_failures.csv", + "shards.log", + ] + assert sorted(path.name for path in (run_dir / "obs_html").iterdir()) == [ + "index.html", + "opendrive__Town02__seed_11__episode_000000.html", + "opendrive__Town07__seed_21__episode_000000.html", + "opendrive__Town07__seed_22__episode_000001.html", + ] + assert len(pd.read_csv(run_dir / "details" / "episode_metrics.csv")) == 20 + failures = pd.read_csv(run_dir / "details" / "failure_episodes.csv") + assert failures["shard"].tolist() == ["shard00", "shard01", "shard01"] + assert (run_dir / "details" / "resolved_benchmark.yaml").read_text() == "shard: shard00\n" + shards_log = (run_dir / "details" / "shards.log").read_text() + assert "==== shard00.log ====" in shards_log and "shard01 log line" in shards_log + index = (run_dir / "obs_html" / "index.html").read_text() + assert "opendrive__Town07__seed_22__episode_000001" in index and "shard index" not in index + + +def test_tidy_is_a_no_op_the_second_time(tmp_path, monkeypatch, capsys): + run_dir = _write_parallel_eval(tmp_path) + _tidy(run_dir, monkeypatch) + before = sorted(str(path) for path in run_dir.rglob("*")) + _tidy(run_dir, monkeypatch) + assert sorted(str(path) for path in run_dir.rglob("*")) == before + assert "already tidied" in capsys.readouterr().out + + +def test_tidy_handles_a_single_process_eval_dir(tmp_path, monkeypatch): + run_dir = tmp_path / "eval" / "carla" / STAMP + failures = run_dir / "failures" / "rendered_replays" + failures.mkdir(parents=True) + (run_dir / "evaluation_summary.json").write_text("{}") + (run_dir / "resolved_benchmark.yaml").write_text("single\n") + _episode_rows([("opendrive__Town02", s) for s in range(3)]).to_csv(run_dir / "episode_metrics.csv", index=False) + _episode_rows([("opendrive__Town02", 1)], shard_dir="x").to_csv( + run_dir / "failures" / "episode_metrics.csv", index=False + ) + (failures / "opendrive__Town02__seed_1__episode_000000.html").write_text("page") + (failures / "index.html").write_text("old index") + _tidy(run_dir, monkeypatch) + assert sorted(path.name for path in run_dir.iterdir()) == ["details", "evaluation_summary.json", "obs_html"] + assert (run_dir / "details" / "episode_metrics.csv").is_file() and ( + run_dir / "details" / "resolved_benchmark.yaml" + ).is_file() + assert sorted(path.name for path in (run_dir / "obs_html").iterdir()) == [ + "index.html", + "opendrive__Town02__seed_1__episode_000000.html", + ] + + +def _capture(tmp_path, observation_agents): + args = { + "eval": {"observation_replay_writer_count": 1, "capture_observations_agents": observation_agents}, + "env": {}, + } + return eval_replay.EvalReplayCapture(args, None, str(tmp_path / "replays"), True, 1, 2, 4, 0) + + +def _bundle_summary(colliding_agent): + metrics = np.zeros((4, 2, len(viz.METRIC_LABELS)), np.float32) + metrics[2, colliding_agent, viz.METRIC_LABELS.index("collision")] = 1.0 + bundle = { + "schema": "interactive_replay_environment_v1", + "metadata": {"episode_length": 4, "worker_idx": 0, "active_agent_offset": 0, "active_agent_count": 2}, + "scenario": {}, + "frames": {"metrics_f32": metrics}, + } + return { + "replay_environment_bundle": zlib.compress(pickle.dumps(bundle)), + "map_name": "x/opendrive__Town01.bin", + "seed": 7, + } + + +def test_infracting_mode_keeps_only_the_infracting_agents_observations(tmp_path): + capture = _capture(tmp_path, "infracting") + capture.policy_history = { + "obs": np.ones((4, 2, 3), np.float16), + "pool_lane": np.ones((4, 2, 5), np.int16), + "value": np.ones((4, 2), np.float32), + } + capture.policy_history_frame_count = 4 + capture.queue_replay(_bundle_summary(colliding_agent=1), 0) + replay = capture.pending_replays[0][1] + assert replay["obs_agents"] == [1] + assert not replay["obs"][:, 0].any() and replay["obs"][:, 1].all() + assert not replay["pool_lane"][:, 0].any() and replay["pool_lane"][:, 1].all() + assert replay["value"].all() # the small per-agent chunks stay + + +def test_all_mode_keeps_every_agents_observations(tmp_path): + capture = _capture(tmp_path, "all") + capture.policy_history = {"obs": np.ones((4, 2, 3), np.float16)} + capture.policy_history_frame_count = 4 + capture.queue_replay(_bundle_summary(colliding_agent=1), 0) + replay = capture.pending_replays[0][1] + assert replay["obs"].all() and "obs_agents" not in replay + + +def test_unknown_observation_agents_mode_is_refused(tmp_path): + with pytest.raises(pufferlib.APIUsageError, match="capture_observations_agents"): + _capture(tmp_path, "nobody") diff --git a/tests/unit_tests/test_traffic_light_yellow_remaining.py b/tests/unit_tests/test_traffic_light_yellow_remaining.py new file mode 100644 index 0000000000..146e75c791 --- /dev/null +++ b/tests/unit_tests/test_traffic_light_yellow_remaining.py @@ -0,0 +1,100 @@ +"""obs_traffic_light_yellow_remaining: the co-sim setter's timing reaches the traffic-control obs column.""" + +import numpy as np +import pytest + +from pufferlib.ocean.drive import binding +from pufferlib.ocean.drive.drive import Drive + +MAP_DIR = "pufferlib/resources/drive/binaries/carla" +YELLOW = binding.TRAFFIC_CONTROL_STATE_YELLOW +RED = binding.TRAFFIC_CONTROL_STATE_RED +YELLOW_MAX_S = 3.0 +DT = 0.1 + + +def _make_env(**overrides): + kwargs = { + "num_agents": 16, + "min_agents_per_env": 1, + "max_agents_per_env": 16, + "num_maps": 1, + "allow_map_subset": True, + "map_dir": MAP_DIR, + "simulation_mode": "gigaflow", + "action_type": "continuous", + "dynamics_model": "jerk", + "scenario_length": 120, + "resample_frequency": 1_000_000, + "termination_mode": 0, + "seed": 3, + "dt": DT, + "traffic_lights_enabled": True, + "obs_slots_traffic_controls_n": 4, + "traffic_light_yellow_min_s": 2.0, + "traffic_light_yellow_max_s": YELLOW_MAX_S, + "obs_traffic_light_yellow_remaining": True, + } + kwargs.update(overrides) + env = Drive(**kwargs) + env.reset() + return env + + +def _num_traffic(env): + state = env.get_state() + scenario = state[0] if isinstance(state, list) else state + return int(scenario["num_traffic_elements"]) + + +def _light_rows(env): + """(rows, width) traffic-control rows of every agent that hold a traffic light.""" + slots, width = env.obs_slots_traffic_controls_n, env.traffic_control_features + end = env.num_obs - env.obs_valid_count_features + rows = np.asarray(env.recompute_observations())[:, end - slots * width : end].reshape(-1, slots, width) + return rows[rows[:, :, -2] == binding.TRAFFIC_CONTROL_TYPE_TRAFFIC_LIGHT] + + +def test_setter_timing_reaches_obs_column_clipped_to_training_max(): + env = _make_env() + num_traffic = _num_traffic(env) + yellow = np.full(num_traffic, YELLOW, np.int32) + + env.set_traffic_light_states(yellow, np.full(num_traffic, 5.0)) # above the training max: clipped to 3 s + rows = _light_rows(env) + yellow_rows = rows[rows[:, -1] == YELLOW] + assert len(yellow_rows) > 0 + assert np.allclose(yellow_rows[:, -3], 1.0) + assert np.all(rows[rows[:, -1] != YELLOW][:, -3] == 0.0) + + env.set_traffic_light_states(yellow, np.full(num_traffic, 1.0)) # 10 steps of 0.1 s + rows = _light_rows(env) + assert np.allclose(rows[rows[:, -1] == YELLOW][:, -3], 1.0 / YELLOW_MAX_S, atol=1e-6) + + env.set_traffic_light_states(np.full(num_traffic, RED, np.int32), np.zeros(num_traffic)) + rows = _light_rows(env) + assert len(rows) > 0 + assert np.all(rows[:, -3] == 0.0) + + +def test_setter_rejects_missing_or_invalid_timing_when_observed(): + env = _make_env() + num_traffic = _num_traffic(env) + yellow = np.full(num_traffic, YELLOW, np.int32) + with pytest.raises(ValueError): + env.set_traffic_light_states(yellow) + with pytest.raises(ValueError): + env.set_traffic_light_states(yellow, np.full(num_traffic, np.nan)) + with pytest.raises(ValueError): + env.set_traffic_light_states(yellow, np.full(num_traffic, -1.0)) + with pytest.raises(ValueError): + env.set_traffic_light_states(yellow, np.zeros(num_traffic - 1)) + + +def test_flag_off_keeps_the_row_width_and_accepts_timing(): + env = _make_env(obs_traffic_light_yellow_remaining=False) + num_traffic = _num_traffic(env) + assert env.traffic_control_features == binding.TRAFFIC_CONTROL_FEATURES + env.set_traffic_light_states(np.full(num_traffic, YELLOW, np.int32)) + env.set_traffic_light_states(np.full(num_traffic, YELLOW, np.int32), np.full(num_traffic, 1.0)) + assert len(_light_rows(env)) > 0 diff --git a/tests/unit_tests/test_vector_driver.py b/tests/unit_tests/test_vector_driver.py index 8b82eedf78..65a78d4eea 100644 --- a/tests/unit_tests/test_vector_driver.py +++ b/tests/unit_tests/test_vector_driver.py @@ -182,3 +182,19 @@ def test_dead_worker_raises_instead_of_hanging(monkeypatch): vecenv.recv() finally: vecenv.close() + + +class SeedRecordingEnv(MinimalEnv): + def __init__(self, buf=None, seed=0): + self.constructor_seed = seed + super().__init__(buf=buf, seed=seed) + + +def test_make_forwards_seed_to_backend(): + constructor_seeds = {} + for vec_seed in (1, 2): + vecenv = pufferlib.vector.make(SeedRecordingEnv, backend="Serial", num_envs=2, seed=vec_seed) + constructor_seeds[vec_seed] = [env.constructor_seed for env in vecenv.envs] + vecenv.close() + assert constructor_seeds[1] != constructor_seeds[2] + assert len(set(constructor_seeds[1])) == 2