feat: capture head and waist RGB-D data

tg3_local_teleop/data_collection.py

@@ -20,9 +20,10 @@ import threading
 import time
 import uuid
 from collections.abc import Callable, Mapping, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from pathlib import Path
+from types import MappingProxyType
 from typing import Any, Literal, Protocol
 
 import yaml
@@ -201,6 +202,7 @@ class RecorderConfig:
     base_directory: Path | str = Path("/home/nvidia/tg3_data_collection")
     topics: Sequence[str] = ()
     required_topics: Sequence[str] = ()
+    minimum_topic_rates_hz: Mapping[str, float] = field(default_factory=dict)
     minimum_free_bytes: int = 5 * 1024**3
     max_duration_seconds: float = 30 * 60.0
     poll_interval_seconds: float = 0.1
@@ -214,6 +216,7 @@ class RecorderConfig:
         base = Path(self.base_directory).expanduser()
         topics = tuple(self.topics)
         required_topics = tuple(self.required_topics)
+        minimum_topic_rates: dict[str, float] = {}
         if not topics:
             raise ValueError("at least one recording topic is required")
         if any(
@@ -238,6 +241,23 @@ class RecorderConfig:
             raise ValueError("required topics must be absolute ROS topic names")
         if not set(required_topics).issubset(topics):
             raise ValueError("required topics must be a subset of recording topics")
+        if not isinstance(self.minimum_topic_rates_hz, Mapping):
+            raise ValueError("minimum topic rates must be a mapping")
+        for topic, rate in self.minimum_topic_rates_hz.items():
+            if topic not in required_topics:
+                raise ValueError(
+                    "minimum-rate topics must be a subset of required topics"
+                )
+            if (
+                isinstance(rate, bool)
+                or not isinstance(rate, (int, float))
+                or not math.isfinite(rate)
+                or rate <= 0.0
+            ):
+                raise ValueError(
+                    f"minimum topic rate for {topic!r} must be positive and finite"
+                )
+            minimum_topic_rates[topic] = float(rate)
         if (
             isinstance(self.minimum_free_bytes, bool)
             or not isinstance(self.minimum_free_bytes, int)
@@ -260,6 +280,11 @@ class RecorderConfig:
         object.__setattr__(self, "base_directory", base)
         object.__setattr__(self, "topics", topics)
         object.__setattr__(self, "required_topics", required_topics)
+        object.__setattr__(
+            self,
+            "minimum_topic_rates_hz",
+            MappingProxyType(minimum_topic_rates),
+        )
 
 
 class RecorderProcess(Protocol):
@@ -296,7 +321,11 @@ class _Episode:
     node_name: str = ""
     command: tuple[str, ...] = ()
     bag_info_validation: dict[str, Any] | None = None
+    topic_message_counts: dict[str, int] | None = None
     required_topic_message_counts: dict[str, int] | None = None
+    minimum_topic_message_counts: dict[str, int] | None = None
+    observed_topic_rates_hz: dict[str, float] | None = None
+    metadata_duration_nanoseconds: int | None = None
     process: RecorderProcess | None = None
     stdout_stream: Any = None
     stderr_stream: Any = None
@@ -728,8 +757,19 @@ class DataRecorderManager:
             "start_reason": episode.request.reason,
             "stop_reason": stop_reason,
             "topics": list(self.config.topics),
+            "topic_message_counts": episode.topic_message_counts,
             "required_topics": list(self.config.required_topics),
             "required_topic_message_counts": required_topic_counts,
+            "minimum_topic_rates_hz": dict(
+                self.config.minimum_topic_rates_hz
+            ),
+            "minimum_topic_message_counts": (
+                episode.minimum_topic_message_counts
+            ),
+            "observed_topic_rates_hz": episode.observed_topic_rates_hz,
+            "metadata_duration_nanoseconds": (
+                episode.metadata_duration_nanoseconds
+            ),
             "storage_id": "mcap",
             "custom_data": {
                 "capture_id": episode.request.episode_id,
@@ -778,6 +818,9 @@ class DataRecorderManager:
                 raise RecordingError(
                     "bag metadata.yaml lacks topics_with_message_count"
                 )
+            episode.topic_message_counts = {
+                topic: 0 for topic in self.config.topics
+            }
             episode.required_topic_message_counts = {}
             return {}
 
@@ -812,6 +855,9 @@ class DataRecorderManager:
             topic: all_counts.get(topic, 0)
             for topic in self.config.required_topics
         }
+        episode.topic_message_counts = {
+            topic: all_counts.get(topic, 0) for topic in self.config.topics
+        }
         episode.required_topic_message_counts = required_counts
         missing = [
             topic for topic in self.config.required_topics if topic not in all_counts
@@ -828,8 +874,61 @@ class DataRecorderManager:
             problems.append("zero-message required topics: " + ", ".join(empty))
         if problems:
             raise RecordingError("; ".join(problems))
+        self._validate_minimum_topic_rates(episode, information, all_counts)
         return required_counts
 
+    def _validate_minimum_topic_rates(
+        self,
+        episode: _Episode,
+        information: Mapping[str, Any],
+        all_counts: Mapping[str, int],
+    ) -> None:
+        configured = self.config.minimum_topic_rates_hz
+        if not configured:
+            episode.minimum_topic_message_counts = {}
+            episode.observed_topic_rates_hz = {}
+            return
+
+        duration = information.get("duration")
+        if not isinstance(duration, Mapping):
+            raise RecordingError(
+                "bag metadata.yaml lacks duration for minimum-rate validation"
+            )
+        nanoseconds = duration.get("nanoseconds")
+        if (
+            isinstance(nanoseconds, bool)
+            or not isinstance(nanoseconds, int)
+            or nanoseconds <= 0
+        ):
+            raise RecordingError(
+                "bag metadata.yaml duration.nanoseconds must be a positive integer"
+            )
+
+        episode.metadata_duration_nanoseconds = nanoseconds
+        minimum_counts: dict[str, int] = {}
+        observed_rates: dict[str, float] = {}
+        failures: list[str] = []
+        duration_seconds = nanoseconds / 1_000_000_000.0
+        for topic, minimum_rate in configured.items():
+            count = all_counts[topic]
+            minimum_count = math.floor(duration_seconds * minimum_rate)
+            observed_rate = count / duration_seconds
+            minimum_counts[topic] = minimum_count
+            observed_rates[topic] = observed_rate
+            if count < minimum_count:
+                failures.append(
+                    f"{topic}: {count} messages < {minimum_count} required "
+                    f"over {duration_seconds:.6f}s "
+                    f"({observed_rate:.3f} Hz < {minimum_rate:.3f} Hz)"
+                )
+        episode.minimum_topic_message_counts = minimum_counts
+        episode.observed_topic_rates_hz = observed_rates
+        if failures:
+            raise RecordingError(
+                "required topics below minimum average rate: "
+                + "; ".join(failures)
+            )
+
     def _validate_bag_info(self, episode: _Episode) -> dict[str, Any]:
         """Run rosbag2's reader-level validation before publishing READY."""
 
@@ -952,10 +1051,21 @@ class DataRecorderManager:
                 "start_reason": episode.request.reason,
                 "stop_reason": stop_reason,
                 "topics": list(self.config.topics),
+                "topic_message_counts": episode.topic_message_counts,
                 "required_topics": list(self.config.required_topics),
                 "required_topic_message_counts": (
                     episode.required_topic_message_counts
                 ),
+                "minimum_topic_rates_hz": dict(
+                    self.config.minimum_topic_rates_hz
+                ),
+                "minimum_topic_message_counts": (
+                    episode.minimum_topic_message_counts
+                ),
+                "observed_topic_rates_hz": episode.observed_topic_rates_hz,
+                "metadata_duration_nanoseconds": (
+                    episode.metadata_duration_nanoseconds
+                ),
                 "storage_id": "mcap",
                 "custom_data": {
                     "capture_id": episode.request.episode_id,