feat: support H5 evaluation dataset export

.gitignore

@@ -13,6 +13,7 @@ wheels/
 build/
 install/
 log/
+/ros2_ws/src/deps/
 
 node_modules/
 

README.md

@@ -34,12 +34,19 @@ listed in `ros2_ws/diffusion_planner.repos`, build the workspace with `colcon`,
 source it before running ROS-dependent commands:
 
 ```bash
+vcs import ros2_ws/src < ros2_ws/diffusion_planner.repos
+cd ros2_ws
+colcon build --symlink-install
+cd ..
 source ros2_ws/install/setup.bash
 ```
 
 ## Dataset
 
-Generate one `frames.h5` shard per rosbag and a Parquet frame index:
+Generate one `frames.h5` shard per rosbag and a Parquet frame index. Source bags
+are first normalized to the current Autoware message schemas with
+`autoware_msg_bag_converter`; converted bags are cached under
+`converted_bag_root`:
 
 ```bash
 source ros2_ws/install/setup.bash
@@ -50,7 +57,9 @@ uv run python scripts/dataset/create_h5_dataset.py \
 ```
 
 Configuration is defined in `configs/dataset/create_h5_dataset.yaml`. Completed H5
-shards are reused when an interrupted index build is resumed.
+shards are reused when an interrupted index build is resumed. Training and open-loop
+generation defaults to a 0.5 s frame interval. Closed-loop releases must explicitly
+use `frame_interval=0.1` so their frame grid matches the 10 Hz rollout clock.
 
 Inspect loading performance with:
 

configs/dataset/create_h5_dataset.yaml

@@ -7,9 +7,16 @@ defaults:
 root: ???
 # Root under which the source bag directory hierarchy and frames.h5 files are written.
 output_root: ???
+# Cache of bags normalized to the current Autoware message schemas. Every bag
+# processed below passes through autoware_msg_bag_converter first.
+converted_bag_root: ${output_root}/converted_rosbags
 index_output: ${output_root}/indexes/${split}.parquet
 split: train
+# Training and open-loop default. Closed-loop releases must override this to
+# 0.1 because their rollout clock and timeout budget operate at 10 Hz.
 frame_interval: 0.5
+# Number of future label steps emitted for each frame.
+future_steps: 80
 jobs: 200
 limit: null
 
@@ -29,3 +36,6 @@ topic_drop_thresholds:
   traffic_signals: 1.0
 
 min_travel_distance: 1000.0
+# Emit one H5 shard per adjacent same-start-pose route group. Disabled for
+# ordinary train/valid generation; closed-loop dataset builds enable it.
+split_routes: false

configs/dataset/vehicles/tier4_fleet.yaml

@@ -36,3 +36,8 @@ dnjXqf91: # Minibus.
   base_link_to_front: 5.71111
   vehicle_length: 7.2369
   vehicle_width: 2.42741
+
+Aaq3qv2Q:
+  base_link_to_front: 4.2042
+  vehicle_length: 5.1568
+  vehicle_width: 2.4042

docs/h5_dataset_schema.md

@@ -1,6 +1,6 @@
 # H5 Dataset Schema
 
-This document describes H5 format version 3 produced by
+This document describes H5 format version 4 produced by
 `scripts/dataset/create_h5_dataset.py`.
 
 ## Dataset layout
@@ -36,7 +36,7 @@ the ego frame at `frame_time_ns`.
 | `area_map_version_id` | Map version identifier |
 | `split` | `train`, `valid`, or `auto` |
 | `num_frames` | Number of stored frames, `N` |
-| `frame_interval_s` | Interval used when selecting frames |
+| `frame_interval_s` | Interval used when selecting frames. Training/open-loop normally uses `0.5`; closed-loop evaluation requires `0.1`. |
 | `traffic_light_timeout_s` | Maximum accepted traffic-light observation age |
 | `neighbor_observation_timeout_s` | Maximum accepted tracked-object observation age |
 
@@ -157,6 +157,116 @@ has a normalized yaw pair, so an all-zero pose is unambiguously invalid.
 | `turn_indicator` | `(N,)` | `uint8` | Current `TurnIndicatorsReport.report` value |
 | `num_objects` | `(N,)` | `int32` | Number of tracked objects at the frame |
 
+## Segment tag sidecar
+
+Curated tags are not stored inside `frames.h5`. They are mutable dataset metadata,
+whereas H5 shards are versioned preprocessing outputs. A shard may instead have an
+adjacent, authoritative JSON sidecar named `frames.tags.json`:
+
+```text
+<source rosbag hierarchy>/
+├── frames.h5
+└── frames.tags.json
+```
+
+Keeping the sidecar adjacent means tags travel when a bag directory is copied, while
+adding or revising tags never requires rewriting the H5 payload or changing its format
+version. A missing sidecar means that the shard has no curated tags.
+
+### Sidecar format
+
+The initial sidecar format is `diffusion_planner_h5_tags`, version `1`:
+
+```json
+{
+  "format": "diffusion_planner_h5_tags",
+  "format_version": 1,
+  "h5": {
+    "format": "diffusion_planner_frame_dataset",
+    "format_version": 4,
+    "source_bag_path": "x2_dev/example/train/2026-01-15/13-24-21",
+    "num_frames": 873,
+    "first_frame_time_ns": 1736598240000000000,
+    "last_frame_time_ns": 1736598676000000000
+  },
+  "spans": [
+    {
+      "frame_start": 0,
+      "frame_stop": 873,
+      "first_frame_time_ns": 1736598240000000000,
+      "last_frame_time_ns": 1736598676000000000,
+      "tags": ["split:validation", "site:odaiba", "split:train", "project:x2_dev"]
+    },
+    {
+      "frame_start": 120,
+      "frame_stop": 171,
+      "first_frame_time_ns": 1736598300000000000,
+      "last_frame_time_ns": 1736598300000000000,
+      "tags": ["override:lane_change"]
+    },
+    {
+      "frame_start": 125,
+      "frame_stop": 126,
+      "first_frame_time_ns": 1736598300000000000,
+      "last_frame_time_ns": 1736598300000000000,
+      "tags": ["turn_indicator:left"]
+    }
+  ]
+}
+```
+
+`h5` is an identity and compatibility check for the sibling shard, not a path to it.
+The sidecar must not contain an absolute path. `source_bag_path`, H5 format/version,
+frame count, and the first/last source timestamps must match the H5 file before a
+consumer uses its tags.
+
+Each item in `spans` attaches one or more tags to a half-open H5 frame interval
+`[frame_start, frame_stop)`. The required constraints are:
+
+- `0 <= frame_start < frame_stop <= num_frames`.
+- `first_frame_time_ns` equals `metadata/frame_time_ns[frame_start]` and
+  `last_frame_time_ns` equals `metadata/frame_time_ns[frame_stop - 1]`.
+- A one-frame tag is represented by `frame_stop = frame_start + 1`.
+- The complete shard is `[0, num_frames)`; its final real timestamp is
+  `metadata/frame_time_ns[num_frames - 1]`. There is deliberately no inferred
+  timestamp after the final frame, because source frame spacing need not be exact.
+- Spans may overlap. The effective tags at a frame are the union of tags from every
+  span containing that frame.
+- A span has a non-empty, sorted, duplicate-free `tags` array. Identical spans are
+  invalid; overlapping spans with different tags are valid.
+
+Tags use the same open vocabulary as `tag_toolkit`: `dimension:value`, with both
+components matching lowercase `[a-z0-9_]+`. The format does not impose a taxonomy
+allow-list, so projects can introduce new dimensions and values without a schema
+release.
+
+`project_id`, `split`, and other values already represented by H5 root attributes
+should not be duplicated in the sidecar. Tag query tools may expose such values as
+derived tags (for example, `project:x2_dev` or `split:train`) together with the
+curated sidecar tags.
+
+### Derived tag indexes and evaluation output
+
+`frames.tags.json` is the source of truth. A dataset-level SQLite or Parquet index is
+optional and must be rebuildable from sidecars. A normalized index stores one row per
+`span` and tag, with at least:
+
+```text
+h5_path, source_bag_path, frame_start, frame_stop,
+first_frame_time_ns, last_frame_time_ns, tag
+```
+
+`h5_path` in such an index is relative to the index file. Tags must not become
+authoritative columns in the normal training frame index; a tag query may instead
+materialize a separate filtered frame index for a particular training or evaluation
+selection.
+
+Open-loop result rows should include the effective tags at their referenced frame.
+Closed-loop result rows should include the H5 path, global `[frame_start, frame_stop)`
+range, the union of overlapping tags, and per-tag tagged-frame coverage. Coverage
+makes a tag that touches only part of an evaluated segment distinguishable from one
+that covers the full segment.
+
 ## Parquet frame index
 
 Each Parquet row addresses exactly one H5 frame.
@@ -177,6 +287,6 @@ Each Parquet row addresses exactly one H5 frame.
 
 ## Compatibility
 
-Format version 3 is incompatible with version 2 because lane geometry and boundary
-types are stored in separate tensors. Existing H5 shards must be regenerated with
-`overwrite=true`; `resume=true` rejects shards with an older format version.
+Format version 4 is incompatible with version 3 because it adds closed-loop pose
+metadata (`ego_x`, `ego_y`, and `ego_yaw`). Existing H5 shards must be regenerated
+with `overwrite=true`; `resume=true` rejects shards with an older format version.

packages/diffusion_planner/src/diffusion_planner/data/h5_writer.py

@@ -0,0 +1,95 @@
+"""Shared, atomic writer for Diffusion Planner frame-dataset H5 shards."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import h5py
+import hdf5plugin
+import numpy as np
+
+FORMAT_NAME = "diffusion_planner_frame_dataset"
+FORMAT_VERSION = 4
+
+
+def validate_h5_arrays(
+    frames: Mapping[str, Any], metadata: Mapping[str, Any]
+) -> int:
+    """Validate the shared frame axis and return its length."""
+    frame_times = np.asarray(metadata.get("frame_time_ns"))
+    if frame_times.ndim != 1:
+        raise ValueError("metadata/frame_time_ns must be a vector")
+    num_frames = len(frame_times)
+    if num_frames == 0:
+        return 0
+    if not frames:
+        raise ValueError("H5 result contains metadata but no frame tensors")
+    if num_frames > 1 and not np.all(frame_times[1:] > frame_times[:-1]):
+        raise ValueError("frame times must be strictly increasing")
+    for group_name, arrays in (("frames", frames), ("metadata", metadata)):
+        for key, values in arrays.items():
+            array = np.asarray(values)
+            if array.ndim == 0 or len(array) != num_frames:
+                raise ValueError(
+                    f"{group_name}/{key} does not have {num_frames} frame entries"
+                )
+            if np.issubdtype(array.dtype, np.floating) and not np.isfinite(array).all():
+                raise ValueError(f"{group_name}/{key} contains NaN or infinity")
+    return num_frames
+
+
+def _compression_options(compression: str | None) -> dict[str, Any]:
+    if compression is None:
+        return {}
+    if compression == "zstd":
+        return dict(hdf5plugin.Zstd())
+    return {"compression": compression}
+
+
+def write_h5_shard(
+    path: Path | str,
+    *,
+    attributes: Mapping[str, Any],
+    frames: Mapping[str, Any],
+    metadata: Mapping[str, Any],
+    compression: str | None = "zstd",
+) -> int:
+    """Atomically write one schema-v4 H5 shard from batched frame arrays."""
+    path = Path(path)
+    num_frames = validate_h5_arrays(frames, metadata)
+    if num_frames == 0:
+        raise ValueError(f"refusing to write an empty H5 shard: {path}")
+    temporary = path.with_suffix(path.suffix + ".incomplete")
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary.unlink(missing_ok=True)
+    options = _compression_options(compression)
+    try:
+        with h5py.File(temporary, "w") as output:
+            output.attrs["format"] = FORMAT_NAME
+            output.attrs["format_version"] = FORMAT_VERSION
+            output.attrs["num_frames"] = num_frames
+            for key, value in attributes.items():
+                if key in {"format", "format_version", "num_frames"}:
+                    raise ValueError(f"reserved H5 attribute: {key}")
+                output.attrs[key] = value
+            frame_group = output.create_group("frames")
+            metadata_group = output.create_group("metadata")
+            for key, values in frames.items():
+                array = np.asarray(values)
+                frame_group.create_dataset(
+                    key,
+                    data=array,
+                    chunks=(1, *array.shape[1:]),
+                    shuffle=compression is not None,
+                    **options,
+                )
+            for key, values in metadata.items():
+                metadata_group.create_dataset(key, data=np.asarray(values))
+            output.flush()
+        temporary.replace(path)
+    except BaseException:
+        temporary.unlink(missing_ok=True)
+        raise
+    return num_frames

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.10"
 dependencies = [
     "diffusion-planner",
     "diffusion-planner-dashboard",
+    "transforms3d>=0.4.2",
     "wandb>=0.28.2",
     "zarr>=2.18.3,<3",
 ]

ros2_ws/diffusion_planner.repos

@@ -1,9 +1,24 @@
 repositories:
+  # Rosbag message-schema migration used before closed-loop H5 generation.
+  deps/autoware_msg_bag_converter:
+    type: git
+    url: https://github.com/autowarefoundation/autoware_msg_bag_converter.git
+    version: 39fafce0fcf670b53349bc3541670058714d5055
+  deps/autoware_auto_msgs:
+    type: git
+    url: https://github.com/tier4/autoware_auto_msgs.git
+    version: 624aae9feb142eb2a4e16069ed06bf4050ccc3ea
+  deps/ros2_numpy:
+    type: git
+    url: https://github.com/Box-Robotics/ros2_numpy.git
+    version: d5ebf723e1d2457d36fa89bf021b212e340c3c11
   # core
   deps/autoware_msgs:
     type: git
     url: https://github.com/autowarefoundation/autoware_msgs.git
-    version: 1.7.0
+    # autoware_core requires autoware_localization_msgs/InitializeLocalization.
+    # That service was added after 1.7.0; 1.9.0 is the minimum compatible tag.
+    version: 1.9.0
   deps/autoware_adapi_msgs:
     type: git
     url: https://github.com/autowarefoundation/autoware_adapi_msgs.git
@@ -35,8 +50,8 @@ repositories:
   # universe
   deps/autoware_universe:
     type: git
-    url: https://github.com/yhisaki/autoware_universe.git
-    version: feat/diffusion-planner-v2
+    url: https://github.com/tier4/autoware_universe.git
+    version: feat/v0.64/e2e
   deps/tier4_autoware_msgs:
     type: git
     url: https://github.com/tier4/tier4_autoware_msgs.git

ros2_ws/src/autoware_internal_planning_msgs_v1_13

@@ -0,0 +1 @@
+deps/autoware_msg_bag_converter/old_msgs/autoware_internal_planning_msgs_v1_13

ros2_ws/src/autoware_perception_msgs_v1_7

@@ -0,0 +1 @@
+deps/autoware_msg_bag_converter/old_msgs/autoware_perception_msgs_v1_7

ros2_ws/src/ml_planner_data/CMakeLists.txt

@@ -10,6 +10,9 @@ if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
 endif()
 
 find_package(ament_cmake_auto REQUIRED)
+# Ubuntu's oneTBB package provides an official CMake config, while the legacy
+# module injected by mrt_cmake_modules expects the removed tbb_stddef.h header.
+find_package(TBB CONFIG REQUIRED)
 ament_auto_find_build_dependencies()
 
 # Non-ament dependencies