diff --git a/diffusion_planner/diffusion_planner/config/train_config.py b/diffusion_planner/diffusion_planner/config/train_config.py index eb12c4de5..b44822d4a 100644 --- a/diffusion_planner/diffusion_planner/config/train_config.py +++ b/diffusion_planner/diffusion_planner/config/train_config.py @@ -18,6 +18,35 @@ class TrainConfig(ClosedLoopConfig, ScenarioOpenLoopConfig, ModelConfig): train_set_list: str = cli("JSON list of training NPZ paths", path=True, default="") valid_set_list: str = cli("JSON list of validation NPZ paths", path=True, default="") + # --------------------------------------------------------- + # H5 dataset (new-architecture frame shards) + # --------------------------------------------------------- + # The H5 frames are re-assembled into the same canonical model inputs as the NPZ files, so + # nothing downstream of the dataset changes and each split can be chosen independently -- + # an H5 training split validated against an NPZ split is a supported combination. + h5_train_index: str = cli( + "Parquet frame index of the H5 training split; replaces train_set_list when set", + path=True, + default="", + ) + h5_valid_index: str = cli( + "Parquet frame index of the H5 validation split; replaces valid_set_list when set. " + "Leave empty to keep validating on the NPZ valid_set_list, which is what the " + "replan-consistency evaluation requires", + path=True, + default="", + ) + h5_file_capacity: int = cli( + "open H5 shards cached per DataLoader worker", + default=8, + ) + h5_converter_param_path: str = cli( + "data-converter parameter JSON supplying the true ego_wheel_base per project id; " + "without it the wheelbase is estimated from the H5 ego shape", + path=True, + default="", + ) + # --------------------------------------------------------- # Run output # --------------------------------------------------------- diff --git a/diffusion_planner/diffusion_planner/train.py b/diffusion_planner/diffusion_planner/train.py index 459bf31ca..b6d3dba61 100644 --- a/diffusion_planner/diffusion_planner/train.py +++ b/diffusion_planner/diffusion_planner/train.py @@ -27,6 +27,7 @@ frenet_augmenter_from_args, ) from diffusion_planner.utils.dataset import DiffusionPlannerData, DiffusionPlannerPairData +from diffusion_planner.utils.h5_dataset import H5FrameData, load_wheel_base_by_project from diffusion_planner.utils.lr_schedule import CosineAnnealingWarmUpRestarts, final_phase_lr from diffusion_planner.utils.normalizer import ObservationNormalizer, StateNormalizer from diffusion_planner.utils.onnx_export import export_checkpoint_onnx_guarded @@ -278,10 +279,32 @@ def model_training(args: TrainConfig): aug = None # prepare dataset - train_set = DiffusionPlannerData(args.train_set_list) - valid_set = DiffusionPlannerData(args.valid_set_list) + # The two splits are chosen independently: both sources hand over the same canonical model + # inputs, so an H5 training split can be validated against an NPZ split. That is worth + # keeping available, because the NPZ validation split is what the closed-loop, open-loop + # and replan-consistency evaluations are built around. + wheel_base_by_project = ( + load_wheel_base_by_project(args.h5_converter_param_path) + if args.h5_converter_param_path + else None + ) + + def build_dataset(h5_index: str, path_list: str): + if h5_index: + return H5FrameData( + h5_index, + file_capacity=args.h5_file_capacity, + wheel_base_by_project=wheel_base_by_project, + ) + return DiffusionPlannerData(path_list) - train_set.data_list = train_set.data_list[:: args.train_subsample_step] + train_set = build_dataset(args.h5_train_index, args.train_set_list) + valid_set = build_dataset(args.h5_valid_index, args.valid_set_list) + + if args.h5_train_index: + train_set.subsample(args.train_subsample_step) + else: + train_set.data_list = train_set.data_list[:: args.train_subsample_step] train_sampler = DistributedSampler( train_set, num_replicas=ddp.get_world_size(), rank=global_rank, shuffle=True @@ -311,6 +334,15 @@ def model_training(args: TrainConfig): valid_pair_loader = None if args.enable_replan_consistency_eval: + if args.h5_valid_index: + # The pair dataset walks consecutive NPZ frame paths to find replan pairs; H5 + # frames are addressed by (shard, index) and carry no equivalent path ordering. + # Only the validation split matters here, so an H5 training split can still be + # paired with an NPZ validation split and keep this evaluation. + raise ValueError( + "enable_replan_consistency_eval needs an NPZ validation split; " + "it is not supported with h5_valid_index" + ) expected_gap = args.replan_consistency_expected_gap or None valid_pair_set = DiffusionPlannerPairData(args.valid_set_list, expected_gap=expected_gap) if len(valid_pair_set) > 0: @@ -417,7 +449,14 @@ def model_training(args: TrainConfig): # this function creates dataset artifacts and associate them with wandb run # if wandb_run_id is given, the input artifact is assumed to be created externally and will not be executed if args.use_wandb and args.wandb_run_id is None: - log_dataset_artifact(wandb.run, args.exp_name, args.train_set_list, args.valid_set_list) + # Whichever manifest each split actually came from is what the lineage artifact + # records: the Parquet frame index for an H5 split, the path list for an NPZ one. + log_dataset_artifact( + wandb.run, + args.exp_name, + args.h5_train_index or args.train_set_list, + args.h5_valid_index or args.valid_set_list, + ) if args.ddp: torch.distributed.barrier() diff --git a/diffusion_planner/diffusion_planner/utils/h5_dataset.py b/diffusion_planner/diffusion_planner/utils/h5_dataset.py new file mode 100644 index 000000000..e4a410a65 --- /dev/null +++ b/diffusion_planner/diffusion_planner/utils/h5_dataset.py @@ -0,0 +1,428 @@ +"""Train from the new-architecture H5 frame dataset without changing the model. + +The new architecture (``new-architecture/main``) stores one ``frames.h5`` shard per rosbag +plus a split-level Parquet frame index, documented in ``docs/h5_dataset_schema.md`` of that +branch. Its scene dimensions are identical to the ones this branch's model was built for -- +31 past steps, 80 future steps, 320 neighbours, 140 lanes / 25 route lanes of 20 points, 10 +intersection polygons of 40 points, 60 line strings of 20 points -- but the tensors are +factored differently: boundary types and traffic lights live in their own arrays, neighbour +shape and class are separate from the neighbour tracks, and the ego pose carries cos/sin +instead of a raw yaw. + +So the H5 layout is *re-assembled* here into the canonical NPZ layout the rest of this branch +already consumes, rather than teaching the encoder a second input contract. That keeps the +model, the ONNX export, the ROS 2 node and existing checkpoints untouched: the only thing that +changes is where a training batch comes from. + +Fields H5 does not carry are emitted as all-zero tensors (``static_objects``) or as zeros in +the unused slots (neighbour vx/vy, ego lateral velocity / acceleration / steering). All-zero +is the padding convention on both sides, and ``ObservationNormalizer`` restores exact zeros +after normalizing, so a padded row stays padded end to end. Every attribute written onto a +geometry row is therefore masked by that row's own validity -- writing a type one-hot onto a +padded lane point would make the point look valid to the encoder. +""" + +from __future__ import annotations + +import json +from collections import OrderedDict +from pathlib import Path +from typing import Any + +import h5py +import hdf5plugin # noqa: F401 (registers the zstd filter used by the shards) +import numpy as np +import pyarrow.parquet as pq +from numpy.typing import NDArray +from torch.utils.data import Dataset + +from diffusion_planner import dimensions as dim + +H5_FORMAT = "diffusion_planner_frame_dataset" +H5_FORMAT_VERSION = 4 +REQUIRED_INDEX_COLUMNS = ("h5_path", "frame_index") + +# Frame tensors this converter reads. A shard missing any of them is rejected up front +# instead of failing inside a DataLoader worker. +REQUIRED_H5_KEYS = ( + "ego_agent_past", + "ego_agent_future", + "neighbor_agents_past", + "neighbor_agents_future", + "agent_shape", + "agent_label", + "lanes", + "lane_types", + "lanes_speed_limit", + "lane_traffic_light_past", + "route_lanes", + "route_lane_types", + "route_lanes_speed_limit", + "route_traffic_light_past", + "intersection_area", + "stop_lines", + "road_borders", + "goal_pose", + "ego_shape", + "turn_indicators", +) + +# H5 traffic-light channels: [green, amber, red, unknown_or_unavailable, white_or_no_light, +# arrow_flag]. This branch's 5-wide encoding is [green, yellow, red, white, no_traffic_light]. +# "white" has no separate H5 channel (it is folded into white_or_no_light), and the arrow flag +# has no slot here at all, so both H5 "unknown" and "white or no light" land on +# no_traffic_light and the arrow flag is dropped. +_H5_TRAFFIC_LIGHT_GREEN = 0 +_H5_TRAFFIC_LIGHT_AMBER = 1 +_H5_TRAFFIC_LIGHT_RED = 2 +_H5_TRAFFIC_LIGHT_UNKNOWN = 3 +_H5_TRAFFIC_LIGHT_WHITE_OR_NONE = 4 + +_NEIGHBOR_WIDTH_INDEX = 6 +_NEIGHBOR_LENGTH_INDEX = 7 +_NEIGHBOR_TYPE_START = 8 + + +def _pose_yaw(pose: NDArray[Any]) -> NDArray[np.float32]: + """Convert a trailing [x, y, cos, sin, ...] pose to the canonical [x, y, yaw] layout. + + ``StatePerturbation.interpolation_future_trajectory`` reads column 2 of the ego future as a + raw heading angle, so the ego future must be handed over in the 3-wide layout even though + H5 stores cos/sin. An all-zero padded row maps to an all-zero row because atan2(0, 0) is 0. + """ + yaw = np.arctan2(pose[..., 3], pose[..., 2]) + return np.concatenate([pose[..., :2], yaw[..., None]], axis=-1).astype(np.float32) + + +def _traffic_light_to_canonical(state: NDArray[Any]) -> NDArray[np.float32]: + """Map the trailing H5 6-wide traffic-light encoding onto this branch's 5-wide one.""" + out = np.zeros((*state.shape[:-1], dim.TRAFFIC_LIGHT_ONE_HOT_DIM), dtype=np.float32) + out[..., 0] = state[..., _H5_TRAFFIC_LIGHT_GREEN] + out[..., 1] = state[..., _H5_TRAFFIC_LIGHT_AMBER] + out[..., 2] = state[..., _H5_TRAFFIC_LIGHT_RED] + # index 3 (white) stays zero: H5 cannot distinguish it from "no light". + out[..., 4] = np.maximum( + state[..., _H5_TRAFFIC_LIGHT_UNKNOWN], state[..., _H5_TRAFFIC_LIGHT_WHITE_OR_NONE] + ) + return out + + +def _ego_current_state(ego_agent_past: NDArray[Any]) -> NDArray[np.float32]: + """Rebuild the 10-wide ego current state from the last H5 past step. + + H5 keeps [x, y, cos, sin, velocity, yaw_rate] per step, so vy / ax / ay / steering have no + source and stay zero. The augmenter re-derives the steering angle from the yaw rate and the + wheelbase anyway, and ``compute_training_loss`` only reads the pose and the longitudinal + velocity. + """ + current = ego_agent_past[-1] + state = np.zeros(dim.EGO_CURRENT_STATE_SHAPE[-1], dtype=np.float32) + state[dim.EGOSTATE.X] = current[0] + state[dim.EGOSTATE.Y] = current[1] + state[dim.EGOSTATE.COS] = current[2] + state[dim.EGOSTATE.SIN] = current[3] + state[dim.EGOSTATE.VX] = current[4] + state[dim.EGOSTATE.YAW_RATE] = current[5] + return state + + +def _neighbor_agents_past( + tracks: NDArray[Any], shape: NDArray[Any], label: NDArray[Any] +) -> NDArray[np.float32]: + """Fold the separate H5 neighbour track / shape / class arrays into one 11-wide tensor. + + Layout is [x, y, cos, sin, vx, vy, width, length, vehicle, pedestrian, bicycle]. vx and vy + stay zero because H5 has no neighbour velocity -- and ``NeighborEncoder`` zeroes those two + channels before encoding regardless of what is fed in, so nothing is lost. Shape and class + are per agent in H5 but per timestep here, so they are broadcast over time and then masked + by each step's own validity. + """ + num_agents, num_steps, _ = tracks.shape + out = np.zeros((num_agents, num_steps, dim.NEIGHBOR_SHAPE[-1]), dtype=np.float32) + out[..., : dim.POSE_DIM] = tracks + valid = np.any(tracks != 0.0, axis=-1)[..., None] + out[..., _NEIGHBOR_WIDTH_INDEX : _NEIGHBOR_LENGTH_INDEX + 1] = np.where( + valid, shape[:, None, :], 0.0 + ) + out[..., _NEIGHBOR_TYPE_START:] = np.where(valid, label[:, None, :], 0.0) + return out + + +def _lane_tensor( + points: NDArray[Any], types: NDArray[Any], traffic_light_past: NDArray[Any] +) -> NDArray[np.float32]: + """Rebuild a 33-wide lane/route tensor from H5 geometry, boundary types and lights. + + H5 stores [centre_x, centre_y, left_dx, left_dy, right_dx, right_dy] per point, boundary + types once per segment, and the traffic light as a history; this branch wants + [x, y, dx, dy, left_dx, left_dy, right_dx, right_dy, light(5), left_type(10), + right_type(10)] per point. The boundary offsets already match, the forward difference is + derived from the centreline, and the light is taken at the current step (the last history + entry) because there is only one slot for it here. + """ + num_segments, num_points, _ = points.shape + valid = np.any(points != 0.0, axis=-1) + out = np.zeros((num_segments, num_points, dim.SEGMENT_POINT_DIM), dtype=np.float32) + + centre = points[..., 0:2] + out[..., dim.X : dim.Y + 1] = centre + + # dx, dy point at the next centreline sample; the last valid point of a segment has no + # successor and keeps zeros, matching how the line encoders pad their own differences. + delta = np.zeros_like(centre) + delta[:, :-1] = centre[:, 1:] - centre[:, :-1] + pair_valid = np.zeros_like(valid) + pair_valid[:, :-1] = valid[:, :-1] & valid[:, 1:] + out[..., dim.dX : dim.dY + 1] = np.where(pair_valid[..., None], delta, 0.0) + + out[..., dim.LB_X : dim.LB_Y + 1] = points[..., 2:4] + out[..., dim.RB_X : dim.RB_Y + 1] = points[..., 4:6] + + light = _traffic_light_to_canonical(traffic_light_past[:, -1, :]) + attribute = np.concatenate([light, types.astype(np.float32)], axis=-1) + out[..., dim.TRAFFIC_LIGHT :] = np.where(valid[..., None], attribute[:, None, :], 0.0) + + # Guard the padding invariant: a padded point must stay all-zero across the full row, or + # ObservationNormalizer will shift it off zero and the encoder will read it as valid. + return np.where(valid[..., None], out, 0.0) + + +def _speed_limit(speed_limit: NDArray[Any]) -> tuple[NDArray[np.float32], NDArray[np.bool_]]: + """Split the H5 speed limit into a value and an availability flag (zero means unknown).""" + limit = speed_limit.astype(np.float32) + return limit, limit > 0.0 + + +def _polygons(intersection_area: NDArray[Any]) -> NDArray[np.float32]: + """Turn H5 intersection areas into polygon points carrying their one-wide type flag.""" + points = intersection_area.astype(np.float32) + out = np.zeros((*points.shape[:2], 2 + dim.POLYGON_TYPE_NUM), dtype=np.float32) + out[..., :2] = points + out[..., 2] = np.any(points != 0.0, axis=-1) + return out + + +def _line_strings(stop_lines: NDArray[Any], road_borders: NDArray[Any]) -> NDArray[np.float32]: + """Merge H5 stop lines and road borders into one flagged line-string tensor. + + This branch keeps both in a single array whose columns 2 and 3 flag which one a row is; + H5 splits them, with stop lines carrying only their two end points. The two H5 arrays add + up to ``NUM_LINE_STRINGS`` rows, and stop lines are zero-padded out to the common length. + """ + stop = stop_lines.astype(np.float32) + borders = road_borders.astype(np.float32) + num_stop = stop.shape[0] + num_rows = num_stop + borders.shape[0] + if num_rows != dim.NUM_LINE_STRINGS: + raise ValueError( + f"H5 stop_lines + road_borders is {num_rows} rows, " + f"expected NUM_LINE_STRINGS={dim.NUM_LINE_STRINGS}" + ) + out = np.zeros( + (num_rows, dim.POINTS_PER_LINE_STRING, 2 + dim.LINE_STRING_TYPE_NUM), dtype=np.float32 + ) + for offset, points, flag in ( + (0, stop, dim.LINESTRING.STOP_LINE_FLAG), + (num_stop, borders, dim.LINESTRING.ROAD_BORDER_FLAG), + ): + length = min(points.shape[1], dim.POINTS_PER_LINE_STRING) + rows = slice(offset, offset + points.shape[0]) + out[rows, :length, :2] = points[:, :length] + out[rows, :length, flag] = np.any(points[:, :length] != 0.0, axis=-1) + return out + + +def _ego_shape(ego_shape: NDArray[Any], wheel_base: float | None) -> NDArray[np.float32]: + """Convert the H5 ego shape to this branch's [wheelbase, length, width] layout. + + H5 stores [base_link_to_front, length, width]. Slot 0 is a genuine wheelbase here: the + augmenter turns a yaw rate into a steering angle with it, and the collision loss shifts the + bounding box forward by half of it. When ``wheel_base`` is unknown it is approximated as + twice the base_link-to-box-centre distance, which makes the bounding box exact and leaves + only the augmenter's bicycle model on an estimate. Pass the real per-project value in to + remove that estimate. + """ + base_link_to_front, length, width = (float(value) for value in ego_shape[:3]) + if wheel_base is None: + wheel_base = 2.0 * (base_link_to_front - 0.5 * length) + return np.array([wheel_base, length, width], dtype=np.float32) + + +def convert_h5_frame( + frame: dict[str, NDArray[Any]], wheel_base: float | None = None +) -> dict[str, NDArray[Any]]: + """Convert one H5 frame into the canonical model-input dictionary of this branch.""" + ego_agent_past = frame["ego_agent_past"].astype(np.float32) + lanes_speed_limit, lanes_has_speed_limit = _speed_limit(frame["lanes_speed_limit"]) + route_speed_limit, route_has_speed_limit = _speed_limit(frame["route_lanes_speed_limit"]) + return { + # H5 already carries cos/sin here, and heading_to_cos_sin passes a 4-wide pose through + # unchanged, so these need no yaw round trip. + "ego_agent_past": ego_agent_past[:, : dim.POSE_DIM], + "ego_current_state": _ego_current_state(ego_agent_past), + "neighbor_agents_past": _neighbor_agents_past( + frame["neighbor_agents_past"].astype(np.float32), + frame["agent_shape"].astype(np.float32), + frame["agent_label"].astype(np.float32), + ), + # H5 has no static-object channel; an all-zero block is fully masked by StaticEncoder. + "static_objects": np.zeros( + (dim.NUM_STATIC_OBJECTS, dim.STATIC_OBJECTS_SHAPE[-1]), dtype=np.float32 + ), + "lanes": _lane_tensor( + frame["lanes"].astype(np.float32), + frame["lane_types"], + frame["lane_traffic_light_past"].astype(np.float32), + ), + "lanes_speed_limit": lanes_speed_limit, + "lanes_has_speed_limit": lanes_has_speed_limit, + "route_lanes": _lane_tensor( + frame["route_lanes"].astype(np.float32), + frame["route_lane_types"], + frame["route_traffic_light_past"].astype(np.float32), + ), + "route_lanes_speed_limit": route_speed_limit, + "route_lanes_has_speed_limit": route_has_speed_limit, + "polygons": _polygons(frame["intersection_area"]), + "line_strings": _line_strings(frame["stop_lines"], frame["road_borders"]), + "goal_pose": frame["goal_pose"].astype(np.float32)[: dim.POSE_DIM], + "ego_shape": _ego_shape(frame["ego_shape"], wheel_base), + "turn_indicators": frame["turn_indicators"].astype(np.float32), + # The ego future goes out 3-wide: StatePerturbation's quintic refinement reads its + # column 2 as a heading angle. The neighbour future keeps cos/sin, which is what the + # canonical contract fixes it to and what centric_transform rotates directly. + "ego_agent_future": _pose_yaw(frame["ego_agent_future"].astype(np.float32)), + "neighbor_agents_future": frame["neighbor_agents_future"].astype(np.float32)[ + ..., : dim.POSE_DIM + ], + } + + +def load_wheel_base_by_project(path: str | Path) -> dict[str, float]: + """Read per-project wheelbases from a data-converter parameter JSON. + + The file is the one the NPZ pipeline already uses: a mapping of project id to a dict with + an ``ego_wheel_base`` entry. Projects without that entry are skipped. + """ + data = json.loads(Path(path).expanduser().read_text(encoding="utf-8")) + return { + str(project): float(params["ego_wheel_base"]) + for project, params in data.items() + if isinstance(params, dict) and "ego_wheel_base" in params + } + + +class H5FrameData(Dataset): + """Serve canonical model inputs from H5 shards addressed by a Parquet frame index. + + H5 handles are opened lazily and cached per process so DataLoader workers each keep their + own file descriptors; ``__getstate__`` drops them so the dataset can be pickled to workers. + """ + + def __init__( + self, + index_path: str | Path, + file_capacity: int = 8, + wheel_base_by_project: dict[str, float] | None = None, + ) -> None: + self._index_path = Path(index_path).expanduser().resolve() + if not self._index_path.is_file(): + raise FileNotFoundError(f"Parquet frame index not found: {self._index_path}") + if file_capacity < 1: + raise ValueError(f"file_capacity must be at least 1: {file_capacity}") + + table = pq.read_table(self._index_path) + missing = [name for name in REQUIRED_INDEX_COLUMNS if name not in table.column_names] + if missing: + raise ValueError(f"Parquet index is missing columns: {', '.join(missing)}") + if table.num_rows == 0: + raise ValueError(f"Parquet frame index is empty: {self._index_path}") + + # h5_path is stored relative to the index directory so a dataset stays portable. + relative = table["h5_path"].combine_chunks().to_numpy(zero_copy_only=False) + self.data_list = [str(self._index_path.parent / str(value)) for value in relative] + self._frame_indices = ( + table["frame_index"].combine_chunks().to_numpy(zero_copy_only=False).astype(np.int64) + ) + if np.any(self._frame_indices < 0): + raise ValueError(f"Parquet index has a negative frame_index: {self._index_path}") + + self._file_capacity = file_capacity + self._wheel_base_by_project = dict(wheel_base_by_project or {}) + self._files: OrderedDict[str, h5py.File] = OrderedDict() + self._wheel_base_cache: dict[str, float | None] = {} + + def __len__(self) -> int: + return len(self._frame_indices) + + def subsample(self, step: int) -> None: + """Keep every ``step``-th frame, in place. + + The NPZ dataset is thinned by slicing ``data_list`` directly, which cannot work here: + a frame is addressed by a (shard, frame index) pair, so both arrays have to be sliced + together or the index desynchronizes from the paths. + """ + if step < 1: + raise ValueError(f"subsample step must be at least 1: {step}") + if step == 1: + return + self.data_list = self.data_list[::step] + self._frame_indices = self._frame_indices[::step] + + def __getitem__(self, idx: int) -> dict[str, NDArray[Any]]: + path = self.data_list[idx] + file = self._file_for(path) + frame_index = int(self._frame_indices[idx]) + num_frames = int(file.attrs["num_frames"]) + if not 0 <= frame_index < num_frames: + raise IndexError(f"frame_index {frame_index} outside {path} ({num_frames} frames)") + frames = file["frames"] + frame = {key: np.asarray(frames[key][frame_index]) for key in REQUIRED_H5_KEYS} + return convert_h5_frame(frame, self._wheel_base_cache[path]) + + def _file_for(self, path: str) -> h5py.File: + file = self._files.pop(path, None) + if file is not None: + self._files[path] = file + return file + + file = h5py.File(path, "r") + try: + self._validate(file, path) + except BaseException: + file.close() + raise + project = str(file.attrs.get("project_id", "")) + self._wheel_base_cache[path] = self._wheel_base_by_project.get(project) + self._files[path] = file + while len(self._files) > self._file_capacity: + _, evicted = self._files.popitem(last=False) + evicted.close() + return file + + @staticmethod + def _validate(file: h5py.File, path: str) -> None: + if file.attrs.get("format") != H5_FORMAT: + raise ValueError(f"Not a diffusion-planner H5 shard: {path}") + version = int(file.attrs.get("format_version", -1)) + if version != H5_FORMAT_VERSION: + raise ValueError( + f"H5 format version {version} is unsupported (expected {H5_FORMAT_VERSION}): {path}" + ) + if "num_frames" not in file.attrs or "frames" not in file: + raise ValueError(f"Incomplete H5 shard: {path}") + absent = [key for key in REQUIRED_H5_KEYS if key not in file["frames"]] + if absent: + raise ValueError(f"H5 shard is missing tensors {', '.join(absent)}: {path}") + + def close(self) -> None: + for file in self._files.values(): + file.close() + self._files.clear() + + def __getstate__(self) -> dict[str, Any]: + return {**self.__dict__, "_files": OrderedDict()} + + def __del__(self) -> None: + if getattr(self, "_files", None): + self.close() diff --git a/diffusion_planner/pyproject.toml b/diffusion_planner/pyproject.toml index 77b4b1444..7454a66d7 100644 --- a/diffusion_planner/pyproject.toml +++ b/diffusion_planner/pyproject.toml @@ -5,6 +5,8 @@ requires-python = ">=3.10,<3.11" dependencies = [ "einops", "gradio", + "h5py>=3.16.0", + "hdf5plugin>=7.0.0", "matplotlib", "natsort", "numpy>=1.26,<2", @@ -16,6 +18,7 @@ dependencies = [ "peft", "planner-metrics", "protobuf", + "pyarrow>=24.0.0", "pydantic", "pytorch-lightning", "scipy", diff --git a/diffusion_planner/tests/test_h5_dataset.py b/diffusion_planner/tests/test_h5_dataset.py new file mode 100644 index 000000000..1e067c342 --- /dev/null +++ b/diffusion_planner/tests/test_h5_dataset.py @@ -0,0 +1,548 @@ +"""Contract tests for training off the new-architecture H5 frame dataset.""" + +import json + +import h5py +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +import torch +from diffusion_planner.utils.h5_dataset import ( + H5_FORMAT, + H5_FORMAT_VERSION, + REQUIRED_H5_KEYS, + H5FrameData, + convert_h5_frame, + load_wheel_base_by_project, +) + +from diffusion_planner import dimensions as dim + +PAST_STEPS = dim.INPUT_T + 1 +NUM_STOP_LINES = 30 +NUM_ROAD_BORDERS = 30 + + +def _h5_frame(rng: np.random.Generator, num_valid_neighbors: int = 3) -> dict: + """Build one synthetic H5 frame with realistic padding and one-hot attributes.""" + frame = { + "ego_agent_past": np.zeros((PAST_STEPS, 6), np.float32), + "ego_agent_future": np.zeros((dim.OUTPUT_T, 6), np.float32), + "neighbor_agents_past": np.zeros((dim.MAX_NUM_NEIGHBORS, PAST_STEPS, 4), np.float32), + "neighbor_agents_future": np.zeros((dim.MAX_NUM_NEIGHBORS, dim.OUTPUT_T, 4), np.float32), + "agent_shape": np.zeros((dim.MAX_NUM_NEIGHBORS, 2), np.float32), + "agent_label": np.zeros((dim.MAX_NUM_NEIGHBORS, 3), np.float32), + "lanes": np.zeros((dim.NUM_SEGMENTS_IN_LANE, dim.POINTS_PER_LANELET, 6), np.float32), + "lane_types": np.zeros((dim.NUM_SEGMENTS_IN_LANE, 2 * dim.LINE_TYPE_NUM), np.float32), + "lanes_speed_limit": np.zeros((dim.NUM_SEGMENTS_IN_LANE, 1), np.float32), + "lane_traffic_light_past": np.zeros((dim.NUM_SEGMENTS_IN_LANE, PAST_STEPS, 6), np.float32), + "route_lanes": np.zeros((dim.NUM_SEGMENTS_IN_ROUTE, dim.POINTS_PER_LANELET, 6), np.float32), + "route_lane_types": np.zeros( + (dim.NUM_SEGMENTS_IN_ROUTE, 2 * dim.LINE_TYPE_NUM), np.float32 + ), + "route_lanes_speed_limit": np.zeros((dim.NUM_SEGMENTS_IN_ROUTE, 1), np.float32), + "route_traffic_light_past": np.zeros( + (dim.NUM_SEGMENTS_IN_ROUTE, PAST_STEPS, 6), np.float32 + ), + "intersection_area": np.zeros((dim.NUM_POLYGONS, dim.POINTS_PER_POLYGON, 2), np.float32), + "stop_lines": np.zeros((NUM_STOP_LINES, 2, 2), np.float32), + "road_borders": np.zeros((NUM_ROAD_BORDERS, dim.POINTS_PER_LINE_STRING, 2), np.float32), + "goal_pose": np.zeros(4, np.float32), + "ego_shape": np.array([5.71111, 7.2369, 2.42741], np.float32), + "turn_indicators": np.full(PAST_STEPS, 1.0, np.float32), + } + + yaw = rng.uniform(-np.pi, np.pi, PAST_STEPS).astype(np.float32) + frame["ego_agent_past"][:, 0] = np.linspace(-20.0, 0.0, PAST_STEPS) + frame["ego_agent_past"][:, 1] = rng.normal(0.0, 0.2, PAST_STEPS) + frame["ego_agent_past"][:, 2] = np.cos(yaw) + frame["ego_agent_past"][:, 3] = np.sin(yaw) + frame["ego_agent_past"][:, 4] = 6.0 + frame["ego_agent_past"][:, 5] = 0.03 + + future_yaw = rng.uniform(-np.pi, np.pi, dim.OUTPUT_T).astype(np.float32) + frame["ego_agent_future"][:, 0] = np.linspace(1.0, 40.0, dim.OUTPUT_T) + frame["ego_agent_future"][:, 2] = np.cos(future_yaw) + frame["ego_agent_future"][:, 3] = np.sin(future_yaw) + + for agent in range(num_valid_neighbors): + frame["neighbor_agents_past"][agent, :, 0] = 10.0 + agent + frame["neighbor_agents_past"][agent, :, 1] = -3.0 * agent - 1.0 + frame["neighbor_agents_past"][agent, :, 2] = 1.0 + frame["neighbor_agents_future"][agent, :, 0] = 12.0 + agent + frame["neighbor_agents_future"][agent, :, 2] = 1.0 + frame["agent_shape"][agent] = (1.9, 4.5) + frame["agent_label"][agent, agent % 3] = 1.0 + + # One valid lane whose points walk forward, plus a red light and both boundary types. + frame["lanes"][0, :, 0] = np.linspace(0.0, 19.0, dim.POINTS_PER_LANELET) + frame["lanes"][0, :, 1] = 0.5 + frame["lanes"][0, :, 2:4] = (0.0, 1.75) + frame["lanes"][0, :, 4:6] = (0.0, -1.75) + frame["lane_types"][0, dim.LINE_TYPE_NUM - 1] = 1.0 + frame["lane_types"][0, dim.LINE_TYPE_NUM] = 1.0 + frame["lanes_speed_limit"][0, 0] = 13.89 + frame["lane_traffic_light_past"][0, :, 2] = 1.0 + + frame["route_lanes"][0] = frame["lanes"][0] + frame["route_lane_types"][0] = frame["lane_types"][0] + frame["route_lanes_speed_limit"][0, 0] = 13.89 + frame["route_traffic_light_past"][0, :, 0] = 1.0 + + frame["intersection_area"][0, :, 0] = np.linspace(5.0, 45.0, dim.POINTS_PER_POLYGON) + frame["intersection_area"][0, :, 1] = 2.0 + frame["stop_lines"][0] = ((8.0, -2.0), (8.0, 2.0)) + frame["road_borders"][0, :, 0] = np.linspace(0.0, 19.0, dim.POINTS_PER_LINE_STRING) + frame["road_borders"][0, :, 1] = 4.0 + frame["goal_pose"][:] = (40.0, 0.0, 1.0, 0.0) + return frame + + +def _write_shard(directory, frame: dict, project_id: str = "x2_dev"): + """Write a one-frame H5 shard plus the Parquet index addressing it.""" + shard_dir = directory / "x2_dev" / "0001_area" / "train" / "2026-09-03" / "10-00-00" + shard_dir.mkdir(parents=True) + shard_path = shard_dir / "frames.h5" + with h5py.File(shard_path, "w") as file: + file.attrs["format"] = H5_FORMAT + file.attrs["format_version"] = H5_FORMAT_VERSION + file.attrs["num_frames"] = 1 + file.attrs["project_id"] = project_id + frames = file.create_group("frames") + for key, value in frame.items(): + frames.create_dataset(key, data=value[None, ...]) + + index_dir = directory / "indexes" + index_dir.mkdir() + index_path = index_dir / "train.parquet" + relative = shard_path.relative_to(index_dir.parent) + pq.write_table( + pa.table( + { + "h5_path": pa.array([f"../{relative.as_posix()}"]), + "frame_index": pa.array([0], pa.int64()), + "frame_time_ns": pa.array([1], pa.int64()), + } + ), + index_path, + ) + return index_path + + +@pytest.fixture +def frame(): + return _h5_frame(np.random.default_rng(0)) + + +def test_converted_shapes_match_the_canonical_contract(frame): + """Every emitted tensor matches the shape contract the NPZ converter is held to. + + The contract is imported rather than restated so that a change to the model's tensor + dimensions shows up here instead of drifting silently. It lives in ``rlvr``, which is a + sibling workspace package, so the assertion is skipped where only this package is present. + """ + scene_features = pytest.importorskip("rlvr.autoresearch.scene_features") + converted = convert_h5_frame(frame) + expected = scene_features._canonical_expected_shapes() + + assert set(expected).issubset(converted) + for key, expected_shape in expected.items(): + actual = converted[key].shape + assert len(actual) == len(expected_shape), f"{key}: {actual} vs {expected_shape}" + for axis, (got, want) in enumerate(zip(actual, expected_shape)): + if want is None: + continue + allowed = want if isinstance(want, tuple) else (want,) + assert got in allowed, f"{key} axis {axis}: {got} not in {allowed}" + + +def test_ego_state_and_futures_use_the_layouts_downstream_code_reads(frame): + """The ego current state and both futures follow the NPZ layout, not the H5 one.""" + converted = convert_h5_frame(frame) + + current = converted["ego_current_state"] + assert current.shape == (10,) + np.testing.assert_allclose(current[: dim.EGOSTATE.SIN + 1], frame["ego_agent_past"][-1, :4]) + assert current[dim.EGOSTATE.VX] == pytest.approx(6.0) + assert current[dim.EGOSTATE.YAW_RATE] == pytest.approx(0.03) + # H5 carries no lateral velocity, acceleration or steering angle. + assert current[dim.EGOSTATE.VY] == 0.0 + assert current[dim.EGOSTATE.AX] == 0.0 + assert current[dim.EGOSTATE.AY] == 0.0 + assert current[dim.EGOSTATE.STEERING] == 0.0 + + # StatePerturbation reads column 2 of the ego future as a raw heading; the neighbour + # future is fixed at cos/sin by the canonical contract. + assert converted["ego_agent_future"].shape[-1] == 3 + assert converted["neighbor_agents_future"].shape[-1] == dim.POSE_DIM + np.testing.assert_allclose(converted["neighbor_agents_future"], frame["neighbor_agents_future"]) + h5_future = frame["ego_agent_future"] + np.testing.assert_allclose( + converted["ego_agent_future"][:, 2], + np.arctan2(h5_future[:, 3], h5_future[:, 2]), + atol=1e-6, + ) + # The past and the goal keep cos/sin, which heading_to_cos_sin passes through. + assert converted["ego_agent_past"].shape[-1] == dim.POSE_DIM + assert converted["goal_pose"].shape[-1] == dim.POSE_DIM + + +def test_neighbor_channels_carry_shape_and_class_only_where_valid(frame): + """Neighbour width/length/class are broadcast over time and masked by each step.""" + converted = convert_h5_frame(frame) + neighbors = converted["neighbor_agents_past"] + + np.testing.assert_allclose(neighbors[..., : dim.POSE_DIM], frame["neighbor_agents_past"]) + # vx and vy have no H5 source, and NeighborEncoder zeroes them regardless. + assert not neighbors[..., 4:6].any() + np.testing.assert_allclose(neighbors[0, -1, 6:8], frame["agent_shape"][0]) + np.testing.assert_allclose(neighbors[0, -1, 8:11], frame["agent_label"][0]) + # A padded slot stays all-zero across the full row. + assert not neighbors[100].any() + + +def test_lane_rows_are_reassembled_and_padding_stays_zero(frame): + """Lane geometry, lights and boundary types land in this branch's 33-wide layout.""" + converted = convert_h5_frame(frame) + lanes = converted["lanes"] + source = frame["lanes"] + + assert lanes.shape[-1] == dim.SEGMENT_POINT_DIM + np.testing.assert_allclose(lanes[0, :, dim.X : dim.Y + 1], source[0, :, 0:2]) + np.testing.assert_allclose(lanes[0, :, dim.LB_X : dim.LB_Y + 1], source[0, :, 2:4]) + np.testing.assert_allclose(lanes[0, :, dim.RB_X : dim.RB_Y + 1], source[0, :, 4:6]) + + # dx, dy step to the next centreline sample; the final point has no successor. + np.testing.assert_allclose(lanes[0, :-1, dim.dX], np.ones(dim.POINTS_PER_LANELET - 1)) + assert lanes[0, -1, dim.dX] == 0.0 + assert lanes[0, -1, dim.dY] == 0.0 + + # A red H5 light becomes a red light here, and the arrow flag has no slot. + light = lanes[0, 0, dim.TRAFFIC_LIGHT : dim.TRAFFIC_LIGHT + dim.TRAFFIC_LIGHT_ONE_HOT_DIM] + np.testing.assert_allclose(light, [0.0, 0.0, 1.0, 0.0, 0.0]) + + left = lanes[0, 0, dim.LINE_TYPE_LEFT_START : dim.LINE_TYPE_LEFT_START + dim.LINE_TYPE_NUM] + right = lanes[0, 0, dim.LINE_TYPE_RIGHT_START :] + np.testing.assert_allclose(left, frame["lane_types"][0, : dim.LINE_TYPE_NUM]) + np.testing.assert_allclose(right, frame["lane_types"][0, dim.LINE_TYPE_NUM :]) + + # Padded segments must stay all-zero, or the encoder reads them as valid lanes. + assert not lanes[1:].any() + assert converted["route_lanes"][1:].any() == False # noqa: E712 + + +def test_traffic_light_unknown_and_no_light_share_one_slot(frame): + """H5 'unknown' and 'white or no light' both map onto no_traffic_light.""" + frame["lane_traffic_light_past"][0, :, :] = 0.0 + frame["lane_traffic_light_past"][0, :, 3] = 1.0 + unknown = convert_h5_frame(frame)["lanes"][0, 0, dim.TRAFFIC_LIGHT :][ + : dim.TRAFFIC_LIGHT_ONE_HOT_DIM + ] + assert unknown[dim.TRAFFIC_LIGHT_NO_TRAFFIC_LIGHT - dim.TRAFFIC_LIGHT] == 1.0 + + frame["lane_traffic_light_past"][0, :, :] = 0.0 + frame["lane_traffic_light_past"][0, :, 4] = 1.0 + no_light = convert_h5_frame(frame)["lanes"][0, 0, dim.TRAFFIC_LIGHT :][ + : dim.TRAFFIC_LIGHT_ONE_HOT_DIM + ] + np.testing.assert_allclose(no_light, unknown) + + +def test_speed_limit_flag_is_boolean_and_keyed_on_availability(frame): + """A zero H5 speed limit means unknown, which is what the flag has to say.""" + converted = convert_h5_frame(frame) + flag = converted["lanes_has_speed_limit"] + + assert flag.dtype == np.bool_ + assert bool(flag[0, 0]) + assert not bool(flag[1, 0]) + # torch.where wants a bool condition; uint8 only survives as a deprecation warning. + assert torch.from_numpy(flag).dtype == torch.bool + + +def test_stop_lines_and_road_borders_merge_into_flagged_line_strings(frame): + """Both H5 map arrays land in one tensor whose flags say which is which.""" + line_strings = convert_h5_frame(frame)["line_strings"] + + assert line_strings.shape == ( + dim.NUM_LINE_STRINGS, + dim.POINTS_PER_LINE_STRING, + 2 + dim.LINE_STRING_TYPE_NUM, + ) + np.testing.assert_allclose(line_strings[0, :2, :2], frame["stop_lines"][0]) + assert line_strings[0, :2, dim.LINESTRING.STOP_LINE_FLAG].all() + assert not line_strings[0, :2, dim.LINESTRING.ROAD_BORDER_FLAG].any() + # Stop lines carry two points, so the rest of the row stays padding. + assert not line_strings[0, 2:].any() + + border = line_strings[NUM_STOP_LINES] + np.testing.assert_allclose(border[:, :2], frame["road_borders"][0]) + assert border[:, dim.LINESTRING.ROAD_BORDER_FLAG].all() + # compute_road_border_penalty selects borders on exactly this column. + assert bool((line_strings[..., 3] > 0.5).any(axis=-1)[NUM_STOP_LINES]) + + +def test_polygons_get_their_type_flag_only_on_real_points(frame): + polygons = convert_h5_frame(frame)["polygons"] + + assert polygons.shape == (dim.NUM_POLYGONS, dim.POINTS_PER_POLYGON, 2 + dim.POLYGON_TYPE_NUM) + np.testing.assert_allclose(polygons[0, :, :2], frame["intersection_area"][0]) + assert polygons[0, :, 2].all() + assert not polygons[1:].any() + + +def test_static_objects_are_emitted_as_maskable_padding(frame): + """H5 has no static objects, and an all-zero block is what StaticEncoder masks out.""" + static = convert_h5_frame(frame)["static_objects"] + assert static.shape == (dim.NUM_STATIC_OBJECTS, dim.STATIC_OBJECTS_SHAPE[-1]) + assert not static.any() + + +def test_wheel_base_slot_is_exact_when_supplied_and_geometric_otherwise(frame): + """Slot 0 is a real wheelbase here, so a supplied value must win over the estimate.""" + base_link_to_front, length, width = frame["ego_shape"] + + estimated = convert_h5_frame(frame)["ego_shape"] + assert estimated[dim.EGOSHAPE.WHEEL_BASE] == pytest.approx( + 2.0 * (base_link_to_front - 0.5 * length), abs=1e-5 + ) + assert estimated[dim.EGOSHAPE.LENGTH] == pytest.approx(length) + assert estimated[dim.EGOSHAPE.WIDTH] == pytest.approx(width) + + supplied = convert_h5_frame(frame, wheel_base=4.76012)["ego_shape"] + assert supplied[dim.EGOSHAPE.WHEEL_BASE] == pytest.approx(4.76012) + + +def test_load_wheel_base_by_project_reads_the_converter_param_file(tmp_path): + path = tmp_path / "data_converter_param.json" + path.write_text( + '{"x2_dev": {"ego_wheel_base": 4.76012, "ego_length": 7.2369},' + ' "no_wheel_base": {"ego_length": 1.0}}', + encoding="utf-8", + ) + assert load_wheel_base_by_project(path) == {"x2_dev": 4.76012} + + +def test_dataset_reads_a_shard_through_its_parquet_index(tmp_path, frame): + index_path = _write_shard(tmp_path, frame) + dataset = H5FrameData(index_path) + try: + assert len(dataset) == 1 + item = dataset[0] + assert set(convert_h5_frame(frame)) == set(item) + np.testing.assert_allclose( + item["ego_agent_past"], frame["ego_agent_past"][:, : dim.POSE_DIM] + ) + # Without a converter param file the wheelbase falls back to the estimate. + assert item["ego_shape"][dim.EGOSHAPE.WHEEL_BASE] == pytest.approx(2.09266 * 2, abs=1e-3) + finally: + dataset.close() + + +def test_dataset_applies_the_per_project_wheel_base(tmp_path, frame): + index_path = _write_shard(tmp_path, frame, project_id="x2_dev") + dataset = H5FrameData(index_path, wheel_base_by_project={"x2_dev": 4.76012}) + try: + assert dataset[0]["ego_shape"][dim.EGOSHAPE.WHEEL_BASE] == pytest.approx(4.76012) + finally: + dataset.close() + + +def test_subsample_keeps_paths_and_frame_indices_aligned(tmp_path, frame): + """Thinning must slice the shard paths and the frame indices together.""" + index_path = _write_shard(tmp_path, frame) + index_dir = index_path.parent + shard = (index_dir.parent / "x2_dev/0001_area/train/2026-09-03/10-00-00/frames.h5").resolve() + with h5py.File(shard, "r+") as file: + file.attrs["num_frames"] = 4 + for key in REQUIRED_H5_KEYS: + data = np.repeat(file["frames"][key][...], 4, axis=0) + del file["frames"][key] + file["frames"].create_dataset(key, data=data) + relative = f"../{shard.relative_to(index_dir.parent).as_posix()}" + pq.write_table( + pa.table( + { + "h5_path": pa.array([relative] * 4), + "frame_index": pa.array([0, 1, 2, 3], pa.int64()), + "frame_time_ns": pa.array([0, 1, 2, 3], pa.int64()), + } + ), + index_path, + ) + + dataset = H5FrameData(index_path) + try: + dataset.subsample(2) + assert len(dataset) == 2 + assert len(dataset.data_list) == 2 + assert dataset[1] is not None + finally: + dataset.close() + + +def test_shard_with_the_wrong_format_version_is_rejected(tmp_path, frame): + index_path = _write_shard(tmp_path, frame) + shard = index_path.parent.parent / "x2_dev/0001_area/train/2026-09-03/10-00-00/frames.h5" + with h5py.File(shard, "r+") as file: + file.attrs["format_version"] = H5_FORMAT_VERSION - 1 + + dataset = H5FrameData(index_path) + with pytest.raises(ValueError, match="format version"): + dataset[0] + + +def test_missing_parquet_index_fails_before_any_worker_starts(tmp_path): + with pytest.raises(FileNotFoundError): + H5FrameData(tmp_path / "absent.parquet") + + +def test_converted_batch_trains_the_unmodified_model(tmp_path): + """A converted H5 batch runs the real training step with no change to the network. + + This is the point of re-assembling the H5 tensors instead of teaching the encoder a second + input contract: the augmenter, the observation normalizer, the encoder and the decoder all + take the batch as-is, and gradients reach the parameters. + """ + from diffusion_planner.config.train_config import TrainConfig + from diffusion_planner.model.diffusion_planner import Diffusion_Planner + from diffusion_planner.model.module.decoder import compute_training_loss + from diffusion_planner.train_epoch import heading_to_cos_sin + from diffusion_planner.utils.data_augmentation import StatePerturbation + from diffusion_planner.utils.normalizer import ObservationNormalizer, StateNormalizer + + normalization_path = _write_normalization(tmp_path) + rng = np.random.default_rng(0) + frames = [convert_h5_frame(_h5_frame(rng, num_valid_neighbors=4)) for _ in range(2)] + inputs = { + key: torch.stack([torch.from_numpy(np.asarray(item[key])) for item in frames]) + for key in frames[0] + } + + args = TrainConfig(exp_name="h5_smoke") + args.device = "cpu" + args.normalization_file_path = str(normalization_path) + args.state_normalizer = StateNormalizer.from_json(args) + args.observation_normalizer = ObservationNormalizer.from_json(str(normalization_path)) + + model = Diffusion_Planner(args) + model.train() + + inputs["ego_agent_past"] = heading_to_cos_sin(inputs["ego_agent_past"]) + inputs["goal_pose"] = heading_to_cos_sin(inputs["goal_pose"]) + ego_future = inputs["ego_agent_future"] + neighbors_future = inputs["neighbor_agents_future"] + + aug = StatePerturbation( + augment_prob=1.0, + device="cpu", + num_refine=args.num_refine, + ego_past_noise_std=args.ego_past_noise_std, + use_smoothing_future_trajectory=args.use_smoothing_future_trajectory, + ) + inputs, ego_future, neighbors_future = aug(inputs, ego_future, neighbors_future) + + ego_future = heading_to_cos_sin(ego_future) + mask = torch.sum(torch.ne(neighbors_future[..., :3], 0), dim=-1) == 0 + neighbors_future = heading_to_cos_sin(neighbors_future) + neighbors_future[mask] = 0.0 + inputs = args.observation_normalizer(inputs) + + loss = compute_training_loss(model, inputs, (ego_future, neighbors_future, mask), args) + total = ( + args.alpha_neighbor_loss * loss["neighbor_prediction_loss"] + + args.alpha_planning_loss * loss["ego_planning_loss"] + + loss["turn_indicator_loss"] + ) + assert torch.isfinite(total) + total.backward() + assert any( + parameter.grad is not None and torch.isfinite(parameter.grad).all() + for parameter in model.parameters() + ) + + +def _write_normalization(directory): + """Write a normalization JSON covering every key the observation normalizer touches.""" + widths = { + "ego": dim.POSE_DIM, + "neighbor": dim.POSE_DIM, + "ego_agent_past": dim.POSE_DIM, + "ego_current_state": dim.EGO_CURRENT_STATE_SHAPE[-1], + "neighbor_agents_past": dim.NEIGHBOR_SHAPE[-1], + "lanes": dim.SEGMENT_POINT_DIM, + "lanes_speed_limit": 1, + "route_lanes": dim.SEGMENT_POINT_DIM, + "route_lanes_speed_limit": 1, + "polygons": 2 + dim.POLYGON_TYPE_NUM, + "line_strings": 2 + dim.LINE_STRING_TYPE_NUM, + "goal_pose": dim.POSE_DIM, + } + path = directory / "normalization.json" + payload = {key: {"mean": [0.0] * width, "std": [1.0] * width} for key, width in widths.items()} + path.write_text(json.dumps(payload), encoding="utf-8") + return path + + +def test_auxiliary_losses_read_the_converted_columns(frame): + """The road-border and neighbour-collision penalties fire on converted tensors. + + Both losses reach into specific columns this converter has to get right -- the road-border + flag in line_strings column 3, and neighbour width/length in columns 6/7 -- and both are + hinges, so a silently mislaid column shows up as a penalty that is always zero rather than + as a crash. Driving the ego onto the feature and then far away pins down both ends. + """ + from diffusion_planner.loss import ( + compute_ego_edge_points, + compute_neighbor_collision_penalty, + compute_road_border_penalty, + ) + + converted = convert_h5_frame(frame) + ego_shape = torch.from_numpy(converted["ego_shape"])[None] + line_strings = torch.from_numpy(converted["line_strings"])[None] + neighbors_past = torch.from_numpy(converted["neighbor_agents_past"])[None] + + steps = 8 + far = torch.zeros(1, steps, 4) + far[..., 0] = torch.linspace(0.0, 19.0, steps) + far[..., 1] = 400.0 + far[..., 2] = 1.0 + edge_far = compute_ego_edge_points(far, ego_shape, n_interp=2) + + # The synthetic road border runs along y = 4.0. + on_border = far.clone() + on_border[..., 1] = 4.0 + edge_border = compute_ego_edge_points(on_border, ego_shape, n_interp=2) + assert compute_road_border_penalty(edge_border, line_strings, margin=0.5).sum() > 0.0 + assert compute_road_border_penalty(edge_far, line_strings, margin=0.5).sum() == 0.0 + + # Neighbour 0 sits at (10, -1) with the width and length taken from agent_shape. + neighbors_future = torch.zeros(1, dim.MAX_NUM_NEIGHBORS, steps, dim.POSE_DIM) + neighbors_future[..., 2] = 1.0 + neighbors_future[0, :4, :, 0] = 10.0 + neighbors_future[0, :4, :, 1] = -1.0 + valid = torch.zeros(1, dim.MAX_NUM_NEIGHBORS, steps, dtype=torch.bool) + valid[0, :4] = True + on_agent = far.clone() + on_agent[..., 0] = 10.0 + on_agent[..., 1] = -1.0 + edge_agent = compute_ego_edge_points(on_agent, ego_shape, n_interp=2) + margins = { + "margin_vehicle": 0.3, + "margin_pedestrian": 0.3, + "margin_bicycle": 0.3, + } + assert ( + compute_neighbor_collision_penalty( + edge_agent, neighbors_future, valid, neighbors_past, **margins + ).sum() + > 0.0 + ) + assert ( + compute_neighbor_collision_penalty( + edge_far, neighbors_future, valid, neighbors_past, **margins + ).sum() + == 0.0 + ) diff --git a/docs/h5_dataset_training.md b/docs/h5_dataset_training.md new file mode 100644 index 000000000..75e9b3b58 --- /dev/null +++ b/docs/h5_dataset_training.md @@ -0,0 +1,145 @@ +# Training from the new-architecture H5 dataset + +`tier4-main` trains from one NPZ per frame addressed by a JSON path list. The +`new-architecture/main` branch replaced that with one `frames.h5` shard per rosbag plus a +split-level Parquet frame index. This page describes the option that lets `tier4-main` train +from those H5 shards directly. + +## Usage + +```bash +python diffusion_planner/train_run.py \ + --exp_name h5_run \ + --h5_train_index /mnt/nvme/dataset/hdf5_dataset/indexes/train.parquet \ + --h5_valid_index /mnt/nvme/dataset/hdf5_dataset/indexes/valid.parquet \ + --h5_converter_param_path dataset/data_converter_param.json +``` + +| Option | Meaning | +|---|---| +| `--h5_train_index` | Parquet frame index of the training split. Replaces `--train_set_list`. | +| `--h5_valid_index` | Parquet frame index of the validation split. Replaces `--valid_set_list`. | +| `--h5_file_capacity` | Open H5 shards cached per DataLoader worker (default 8). | +| `--h5_converter_param_path` | Data-converter parameter JSON supplying the true `ego_wheel_base` per project id. Optional; see *Ego shape* below. | + +The two splits are chosen independently, because both sources hand over the same canonical +model inputs. When neither index is set, training reads NPZ path lists exactly as before, so +existing commands are unaffected. + +Training on H5 while validating on the existing NPZ split is the useful mixed case: the +validation split is what the closed-loop, open-loop and replan-consistency evaluations are +built around, and keeping it on NPZ keeps all three available and comparable with earlier runs. + +```bash +python diffusion_planner/train_run.py \ + --exp_name h5_train_npz_valid \ + --h5_train_index /mnt/nvme/dataset/hdf5_dataset/indexes/train.parquet \ + --h5_converter_param_path dataset/data_converter_param.json \ + --valid_set_list .../path_list_valid_sft_balanced_every_100.json +``` + +`--train_subsample_step` works in both modes. The W&B dataset artifact records the Parquet +index in place of the NPZ path list, keeping lineage intact. + +## Why the network is unchanged + +The H5 schema was built against the same scene dimensions this branch's model already uses: + +| | tier4-main | H5 (format version 4) | +|---|---|---| +| Past steps | `INPUT_T + 1` = 31 | 31 | +| Future steps | `OUTPUT_T` = 80 | 80 | +| Neighbours | `MAX_NUM_NEIGHBORS` = 320 | 320 | +| Lanes × points | 140 × 20 | 140 × 20 | +| Route lanes × points | 25 × 20 | 25 × 20 | +| Polygons × points | 10 × 40 | 10 × 40 (`intersection_area`) | +| Line strings × points | 60 × 20 | 30 stop lines + 30 road borders × 20 | + +What differs is how the tensors are *factored*, not how big the scene is. So the H5 frame is +re-assembled into the canonical NPZ layout in +`diffusion_planner/utils/h5_dataset.py`, and nothing downstream of the dataset changes: the +encoder, the decoder, the augmenter, the ONNX export, the ROS 2 node and existing checkpoints +all stay as they are. Teaching the encoder a second input contract would have forked all of +them for no gain. + +## Field mapping + +| Canonical input | Source in H5 | +|---|---| +| `ego_agent_past` (31, 4) | `ego_agent_past[:, :4]` — already cos/sin | +| `ego_current_state` (10) | last `ego_agent_past` step: pose, `velocity` → vx, `yaw_rate` | +| `neighbor_agents_past` (320, 31, 11) | `neighbor_agents_past` + `agent_shape` (width, length) + `agent_label` (one-hot) | +| `static_objects` (5, 10) | no H5 source; all-zero, fully masked | +| `lanes` / `route_lanes` (·, 20, 33) | `lanes` geometry + centreline forward difference + `lane_types` + `lane_traffic_light_past[:, -1]` | +| `lanes_speed_limit` / `..._has_speed_limit` | `lanes_speed_limit`; the flag is `> 0` | +| `polygons` (10, 40, 3) | `intersection_area` + type flag | +| `line_strings` (60, 20, 4) | `stop_lines` (flag col 2) then `road_borders` (flag col 3) | +| `goal_pose` (4) | `goal_pose` | +| `ego_shape` (3) | `ego_shape`; see below | +| `turn_indicators` (31) | `turn_indicators` | +| `ego_agent_future` (80, 3) | `ego_agent_future[:, :4]` → (x, y, yaw) | +| `neighbor_agents_future` (320, 80, 4) | `neighbor_agents_future` | + +### Layouts that are not a straight copy + +**Ego future is 3-wide, neighbour future is 4-wide.** `StatePerturbation`'s quintic refinement +reads column 2 of the ego future as a raw heading angle, so cos/sin is converted back to a yaw +there. The neighbour future is fixed at cos/sin by the canonical contract and passes through. + +**Lane attributes are per point here, per segment in H5.** `LaneEncoder` reads them from point +0 of each segment, so the segment's boundary types and light are broadcast over its points -- +but only over *valid* points. Writing an attribute onto a padded point would break the +all-zero padding convention, and `ObservationNormalizer` keys padding on the whole row: a row +that is not exactly zero gets mean-shifted, after which the encoder's own validity test (on +columns 0..7) reads the padding as a real lane. + +**Traffic lights lose two distinctions.** H5 encodes +`[green, amber, red, unknown_or_unavailable, white_or_no_light, arrow_flag]` per step; this +branch has `[green, yellow, red, white, no_traffic_light]` and one slot for the current state. +So the last history entry is used, both H5 "unknown" and "white or no light" map to +`no_traffic_light`, the `white` slot stays zero, and the arrow flag is dropped. + +**Neighbour velocities are absent.** H5 has no neighbour vx/vy. This costs nothing, because +`NeighborEncoder.forward` zeroes those two channels before encoding regardless of what is fed +in. + +**Ego lateral state is absent.** H5 has no ego vy, ax, ay or steering angle. The augmenter +re-derives the steering angle from the yaw rate and the wheelbase, and `compute_training_loss` +only reads the pose and the longitudinal velocity, so the zeros are only visible to +`ego_current_state`'s own encoder input. + +### Ego shape + +H5 stores `[base_link_to_front, length, width]`; this branch stores `[wheelbase, length, +width]`. Slot 0 has two genuine consumers: + +- `StatePerturbation.augment` turns a yaw rate into a steering angle with it (a bicycle model, + so it wants the real wheelbase); +- `compute_ego_bbox_corners` shifts the box centre forward by half of it. + +Without `--h5_converter_param_path` the value is approximated as +`2 × (base_link_to_front − length / 2)`, which makes the bounding-box shift exact and leaves +only the augmenter's bicycle model on an estimate. Passing `dataset/data_converter_param.json` +from the meta repository supplies the real per-project `ego_wheel_base`, resolved through each +shard's `project_id` attribute; that removes the estimate and is the recommended setting. + +## Limits + +- `--enable_replan_consistency_eval` is rejected when `--h5_valid_index` is set. The pair + dataset finds replan pairs by walking consecutive NPZ frame *paths*; H5 frames are addressed + by (shard, frame index) and carry no equivalent ordering. Only the validation split matters, + so an H5 training split with an NPZ validation split keeps this evaluation. +- `static_objects` is always empty, so any behaviour that depended on static obstacles is not + learnable from an H5 run. +- The H5 frame stride is a generation-time setting (`frame_interval`, 0.5 s by default) and is + coarser than the NPZ pipeline's. It changes how many frames a rosbag yields, not their + contents. + +## Tests + +`diffusion_planner/tests/test_h5_dataset.py` builds synthetic H5 shards and checks the +conversion against the canonical shape contract, the padding invariant, the traffic-light +collapse, the wheelbase handling, and index/subsampling behaviour. Two of them are end-to-end: +one runs a converted batch through the augmenter, the normalizer and `compute_training_loss` +and asserts gradients arrive, and one drives the ego onto a converted road border and +neighbour box to confirm the auxiliary losses really read the columns this converter fills. diff --git a/uv.lock b/uv.lock index 7cdf7af2c..092b4c86c 100644 --- a/uv.lock +++ b/uv.lock @@ -362,6 +362,8 @@ source = { editable = "diffusion_planner" } dependencies = [ { name = "einops" }, { name = "gradio" }, + { name = "h5py" }, + { name = "hdf5plugin" }, { name = "matplotlib" }, { name = "natsort" }, { name = "numpy" }, @@ -373,6 +375,7 @@ dependencies = [ { name = "peft" }, { name = "planner-metrics" }, { name = "protobuf" }, + { name = "pyarrow" }, { name = "pydantic" }, { name = "pytorch-lightning" }, { name = "scipy" }, @@ -388,6 +391,8 @@ dependencies = [ requires-dist = [ { name = "einops" }, { name = "gradio" }, + { name = "h5py", specifier = ">=3.16.0" }, + { name = "hdf5plugin", specifier = ">=7.0.0" }, { name = "matplotlib" }, { name = "natsort" }, { name = "numpy", specifier = ">=1.26,<2" }, @@ -399,6 +404,7 @@ requires-dist = [ { name = "peft" }, { name = "planner-metrics", editable = "planner_metrics" }, { name = "protobuf" }, + { name = "pyarrow", specifier = ">=24.0.0" }, { name = "pydantic" }, { name = "pytorch-lightning" }, { name = "scipy" }, @@ -716,6 +722,41 @@ wheels = [ { url = "https://pypi.flatt.tech/files/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, ] +[[package]] +name = "h5py" +version = "3.16.0" +source = { registry = "https://pypi.flatt.tech/simple/" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://pypi.flatt.tech/files/packages/db/33/acd0ce6863b6c0d7735007df01815403f5589a21ff8c2e1ee2587a38f548/h5py-3.16.0.tar.gz", hash = "sha256:a0dbaad796840ccaa67a4c144a0d0c8080073c34c76d5a6941d6818678ef2738", size = 446526, upload-time = "2026-03-06T13:49:08.07Z" } +wheels = [ + { url = "https://pypi.flatt.tech/files/packages/3a/6b/231413e58a787a89b316bb0d1777da3c62257e4797e09afd8d17ad3549dc/h5py-3.16.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e06f864bedb2c8e7c1358e6c73af48519e317457c444d6f3d332bb4e8fa6d7d9", size = 3724137, upload-time = "2026-03-06T13:47:35.242Z" }, + { url = "https://pypi.flatt.tech/files/packages/74/f9/557ce3aad0fe8471fb5279bab0fc56ea473858a022c4ce8a0b8f303d64e9/h5py-3.16.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ec86d4fffd87a0f4cb3d5796ceb5a50123a2a6d99b43e616e5504e66a953eca3", size = 3090112, upload-time = "2026-03-06T13:47:37.634Z" }, + { url = "https://pypi.flatt.tech/files/packages/7a/f5/e15b3d0dc8a18e56409a839e6468d6fb589bc5207c917399c2e0706eeb44/h5py-3.16.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:86385ea895508220b8a7e45efa428aeafaa586bd737c7af9ee04661d8d84a10d", size = 4844847, upload-time = "2026-03-06T13:47:39.811Z" }, + { url = "https://pypi.flatt.tech/files/packages/cb/92/a8851d936547efe30cc0ce5245feac01f3ec6171f7899bc3f775c72030b3/h5py-3.16.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:8975273c2c5921c25700193b408e28d6bdd0111c37468b2d4e25dcec4cd1d84d", size = 5065352, upload-time = "2026-03-06T13:47:41.489Z" }, + { url = "https://pypi.flatt.tech/files/packages/2b/ae/f2adc5d0ca9626db3277a3d87516e124cbc5d0eea0bd79bc085702d04f2c/h5py-3.16.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:1677ad48b703f44efc9ea0c3ab284527f81bc4f318386aaaebc5fede6bbae56f", size = 4839173, upload-time = "2026-03-06T13:47:43.586Z" }, + { url = "https://pypi.flatt.tech/files/packages/64/0b/e0c8c69da1d8838da023a50cd3080eae5d475691f7636b35eff20bb6ef20/h5py-3.16.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:7c4dd4cf5f0a4e36083f73172f6cfc25a5710789269547f132a20975bfe2434c", size = 5076216, upload-time = "2026-03-06T13:47:45.315Z" }, + { url = "https://pypi.flatt.tech/files/packages/66/35/d88fd6718832133c885004c61ceeeb24dbd6397ef877dbed6b3a64d6a286/h5py-3.16.0-cp310-cp310-win_amd64.whl", hash = "sha256:bdef06507725b455fccba9c16529121a5e1fbf56aa375f7d9713d9e8ff42454d", size = 3183639, upload-time = "2026-03-06T13:47:47.041Z" }, +] + +[[package]] +name = "hdf5plugin" +version = "7.0.0" +source = { registry = "https://pypi.flatt.tech/simple/" } +dependencies = [ + { name = "h5py" }, +] +sdist = { url = "https://pypi.flatt.tech/files/packages/3d/64/0fc6b68e5bc671e7b81d67b930fbed3a4e8a2a92dc4af0f7282b2a2ff988/hdf5plugin-7.0.0.tar.gz", hash = "sha256:e6e6b1f8b0c4d2ca87e616ddc31d08330b36207e466357040f95269e5e0401c8", size = 68284761, upload-time = "2026-06-25T20:59:40.102Z" } +wheels = [ + { url = "https://pypi.flatt.tech/files/packages/03/5a/00d0f491d420d491b5134ee044cd8ead5103382d88f4c8aee623258a6348/hdf5plugin-7.0.0-py3-none-macosx_10_13_x86_64.whl", hash = "sha256:e0ff0a81e6319575ffc8bf230b8df43f3f19f6fe96843e113e04c5ba9f7f3141", size = 6941923, upload-time = "2026-06-25T20:59:24.517Z" }, + { url = "https://pypi.flatt.tech/files/packages/96/ad/b28f4102e619c262b29bfcd13ccf6799e26f9ff113d2fdce19050ae29448/hdf5plugin-7.0.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:97ea4ff6114223c5e8ccce7e23cc6cd398b58c02f4e988e967aeea914fcb8030", size = 6259223, upload-time = "2026-06-25T20:59:26.246Z" }, + { url = "https://pypi.flatt.tech/files/packages/1f/f4/67263173ff61c49800eec4f16b4d84e6a39170be8d4871d4eb0d540b71ee/hdf5plugin-7.0.0-py3-none-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0f3ba9f4e2a370340b45e7042d1b1b1577ab259c1ddf00956264836fa33052ef", size = 42789778, upload-time = "2026-06-25T20:59:28.438Z" }, + { url = "https://pypi.flatt.tech/files/packages/75/d2/66673e2d0ef8499d08dd7061caa194e4ddec12df73ab512431c60538f675/hdf5plugin-7.0.0-py3-none-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1a4cb2207ec3ac538fc728e4596e9e49aed6c4487a641071fb370c04a361e7bf", size = 45295544, upload-time = "2026-06-25T20:59:31.765Z" }, + { url = "https://pypi.flatt.tech/files/packages/d7/0b/855e50e27eab8338c71c3157672c065cd2a2ab38887b3ea4bf128cbd89a1/hdf5plugin-7.0.0-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4ad4ab0d3367699d132e61b1cc382a0c640a7dba56536e5105508205ebbe8762", size = 45012013, upload-time = "2026-06-25T20:59:35.029Z" }, + { url = "https://pypi.flatt.tech/files/packages/4e/14/32cf2aae083c74b95678875e11d25d0cb9e51a33d27f4218eb9aeacfcd70/hdf5plugin-7.0.0-py3-none-win_amd64.whl", hash = "sha256:2e052af8d7848e8bac92646584617503a08bb9b466cfa810a49ecd93e89b7ffa", size = 3523827, upload-time = "2026-06-25T20:59:37.758Z" }, +] + [[package]] name = "hf-gradio" version = "0.4.1" @@ -1508,6 +1549,21 @@ wheels = [ { url = "https://pypi.flatt.tech/files/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, ] +[[package]] +name = "pyarrow" +version = "25.0.1" +source = { registry = "https://pypi.flatt.tech/simple/" } +sdist = { url = "https://pypi.flatt.tech/files/packages/3d/e3/27f57f80141379d60defe6703eb50a707325706f07fedfd1312c7a751995/pyarrow-25.0.1.tar.gz", hash = "sha256:9150a83248bfed9813ea3c3af74c3856c1984d444aa28e58bf7733b9750ddf6a", size = 1201653, upload-time = "2026-08-10T12:40:53.904Z" } +wheels = [ + { url = "https://pypi.flatt.tech/files/packages/0a/3e/5cd70becb51e1d044c54ba5e627424a6e87df5b98008cbd22cc6abd409ca/pyarrow-25.0.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:0b1edbb2f385a6a65e9711b62ba86ac54a7816a3f8d17bb3e8a5929d65fb2485", size = 35954271, upload-time = "2026-08-10T12:36:33.857Z" }, + { url = "https://pypi.flatt.tech/files/packages/64/be/17599e086df264ea7dc221d1101e3131e181e00da428a2f9bd0358f0d06b/pyarrow-25.0.1-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:a4dd8bf99a8fac133efc0ed6a92f5fddbe2adba0d0f6dd720e39ba9855cea85c", size = 37647543, upload-time = "2026-08-10T12:36:39.486Z" }, + { url = "https://pypi.flatt.tech/files/packages/42/34/e138b451fd3970a6eda4599f68ae3b2b32b661bc958de3239d54a0bf6575/pyarrow-25.0.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:bddd0c4f7630c2a3ddf6347c1bdaa79d97bcf6bd445f9e60c816b7d77c85a5ae", size = 46837120, upload-time = "2026-08-10T12:36:46.58Z" }, + { url = "https://pypi.flatt.tech/files/packages/57/5c/f8fc0eb2de03464a557d5a4d0c15e972d73362414696618833b771f7eddd/pyarrow-25.0.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:a4d6d5e9a3d1879a97c08ded0c797579b7965eafd0f0c26c30b45ccc06db939b", size = 50066460, upload-time = "2026-08-10T12:36:53.702Z" }, + { url = "https://pypi.flatt.tech/files/packages/3f/d1/0dd64fd06de0333b808a02f60981635f067b71aad3a30698a9a104fae778/pyarrow-25.0.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:514ddb60285631af068875550c90eddc181db3e8e63a032b1559be189e82f056", size = 49937892, upload-time = "2026-08-10T12:37:00.349Z" }, + { url = "https://pypi.flatt.tech/files/packages/cb/3c/f89d1bd76d5f3284c2a44d7d7ebbd8204535e5ae2b41f4077069b4ff2ec6/pyarrow-25.0.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:cab40b1edfef0262e0e5251aa2c58d75630f24d06dd7794480243acc001a1d7d", size = 53107240, upload-time = "2026-08-10T12:37:07.205Z" }, + { url = "https://pypi.flatt.tech/files/packages/67/67/b554a8e09f3f3decccf405eb8fbe86696321cbcb5b62d18b4a5057a4c113/pyarrow-25.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:60e89d8f13861a1f7f8d950fa54aebb8023b30734d0ac51ffa80beabe2df4bba", size = 27848683, upload-time = "2026-08-10T12:37:12.058Z" }, +] + [[package]] name = "pydantic" version = "2.13.4"