feat: capture head and waist RGB-D data

tg3_local_teleop/test_data_collection.py

@@ -165,6 +165,7 @@ class FakeProcess:
         unexpected_returncode: int | None = None,
         ignore_sigint: bool = False,
         topic_counts: Mapping[str, int] | None = None,
+        duration_nanoseconds: int = 1_000_000_000,
     ) -> None:
         self.command = list(command)
         self.signals: list[int] = []
@@ -178,6 +179,7 @@ class FakeProcess:
             metadata = {
                 "rosbag2_bagfile_information": {
                     "storage_identifier": "mcap",
+                    "duration": {"nanoseconds": duration_nanoseconds},
                     "topics_with_message_count": [
                         {
                             "topic_metadata": {"name": topic},
@@ -286,7 +288,9 @@ class DataRecorderManagerTest(unittest.TestCase):
         process_options: dict[str, Any] | None = None,
         command_runner: Callable[..., Any] | None = None,
         validate_bag_info: bool = True,
+        topics: Sequence[str] = ("/joint_states", "/tf"),
         required_topics: Sequence[str] = ("/joint_states", "/tf"),
+        minimum_topic_rates_hz: Mapping[str, float] | None = None,
     ) -> tuple[DataRecorderManager, Any]:
         if factory is None:
             factory = FakeProcessFactory(**(process_options or {}))
@@ -298,8 +302,9 @@ class DataRecorderManagerTest(unittest.TestCase):
         manager = DataRecorderManager(
             RecorderConfig(
                 base_directory=self.base,
-                topics=("/joint_states", "/tf"),
+                topics=topics,
                 required_topics=required_topics,
+                minimum_topic_rates_hz=minimum_topic_rates_hz or {},
                 minimum_free_bytes=minimum_free,
                 max_duration_seconds=max_duration,
                 poll_interval_seconds=0.005,
@@ -346,6 +351,10 @@ class DataRecorderManagerTest(unittest.TestCase):
         manifest = json.loads((ready / "manifest.json").read_text())
         self.assertEqual(manifest["session_id"], "a" * 32)
         self.assertEqual(manifest["topics"], ["/joint_states", "/tf"])
+        self.assertEqual(
+            manifest["topic_message_counts"],
+            {"/joint_states": 25, "/tf": 50},
+        )
         self.assertEqual(
             manifest["required_topic_message_counts"],
             {"/joint_states": 25, "/tf": 50},
@@ -553,6 +562,108 @@ class DataRecorderManagerTest(unittest.TestCase):
         self.assertEqual(bag_info.calls, [])
         self.assertFalse((failed / "READY").exists())
 
+    def test_minimum_average_topic_rates_are_recorded_in_ready_manifest(
+        self,
+    ) -> None:
+        configured_rates = {"/joint_states": 20.0, "/tf": 40.0}
+        manager, _ = self.make_manager(
+            minimum_topic_rates_hz=configured_rates
+        )
+        # RecorderConfig owns an immutable copy, not the caller's dictionary.
+        configured_rates["/tf"] = 1.0
+        with self.assertRaises(TypeError):
+            manager.config.minimum_topic_rates_hz["/tf"] = 2.0  # type: ignore[index]
+
+        self.assertTrue(manager.request_start("b0" * 16))
+        self.wait_for(lambda: manager.status()["recording"])
+        self.assertTrue(manager.request_stop())
+        self.assertTrue(manager.wait_until_idle(2.0))
+        status = manager.status()
+        self.assertEqual(status["last_result"], "ready")
+        ready = Path(status["last_episode_directory"])
+        manifest = json.loads((ready / "manifest.json").read_text())
+        self.assertEqual(
+            manifest["minimum_topic_rates_hz"],
+            {"/joint_states": 20.0, "/tf": 40.0},
+        )
+        self.assertEqual(
+            manifest["minimum_topic_message_counts"],
+            {"/joint_states": 20, "/tf": 40},
+        )
+        self.assertEqual(
+            manifest["observed_topic_rates_hz"],
+            {"/joint_states": 25.0, "/tf": 50.0},
+        )
+        self.assertEqual(
+            manifest["metadata_duration_nanoseconds"], 1_000_000_000
+        )
+
+    def test_optional_topic_counts_show_absent_and_active_publishers(self) -> None:
+        head_topic = "/ob_camera_head/color/image_raw/compressed"
+        topics = ("/joint_states", "/tf", head_topic)
+        for head_count in (None, 17):
+            with self.subTest(head_count=head_count):
+                counts = {"/joint_states": 25, "/tf": 50}
+                if head_count is not None:
+                    counts[head_topic] = head_count
+                manager, _ = self.make_manager(
+                    topics=topics,
+                    process_options={"topic_counts": counts},
+                )
+                self.assertTrue(manager.request_start("c0" * 16))
+                self.wait_for(lambda: manager.status()["recording"])
+                self.assertTrue(manager.request_stop())
+                self.assertTrue(manager.wait_until_idle(2.0))
+                status = manager.status()
+                self.assertEqual(status["last_result"], "ready")
+                manifest = json.loads(
+                    (
+                        Path(status["last_episode_directory"])
+                        / "manifest.json"
+                    ).read_text()
+                )
+                self.assertEqual(
+                    manifest["topic_message_counts"][head_topic],
+                    0 if head_count is None else head_count,
+                )
+
+    def test_topic_that_stops_mid_episode_fails_average_rate(self) -> None:
+        bag_info = FakeBagInfoRunner()
+        manager, _ = self.make_manager(
+            command_runner=bag_info,
+            minimum_topic_rates_hz={"/joint_states": 20.0, "/tf": 20.0},
+            process_options={
+                "duration_nanoseconds": 10_000_000_000,
+                # /tf delivered briefly, then stopped for most of the bag.
+                "topic_counts": {"/joint_states": 250, "/tf": 50},
+            },
+        )
+        self.assertTrue(manager.request_start("b1" * 16))
+        self.wait_for(lambda: manager.status()["recording"])
+        self.assertTrue(manager.request_stop())
+        self.assertTrue(manager.wait_until_idle(2.0))
+        status = manager.status()
+        self.assertEqual(status["last_result"], "failed")
+        self.assertIn(
+            "/tf: 50 messages < 200 required", status["last_error"]
+        )
+        failed = Path(status["last_episode_directory"])
+        manifest = json.loads((failed / "manifest.json").read_text())
+        self.assertEqual(
+            manifest["topic_message_counts"],
+            {"/joint_states": 250, "/tf": 50},
+        )
+        self.assertEqual(
+            manifest["observed_topic_rates_hz"],
+            {"/joint_states": 25.0, "/tf": 5.0},
+        )
+        self.assertEqual(
+            manifest["minimum_topic_message_counts"],
+            {"/joint_states": 200, "/tf": 200},
+        )
+        self.assertEqual(bag_info.calls, [])
+        self.assertFalse((failed / "READY").exists())
+
     def test_bag_info_timeout_preserves_episode_as_failed(self) -> None:
         manager, _ = self.make_manager(
             command_runner=FakeBagInfoRunner(timeout=True)
@@ -693,6 +804,29 @@ class DataRecorderManagerTest(unittest.TestCase):
                 ("/recorded",),
                 required_topics=("/recorded", "/recorded"),
             )
+        with self.assertRaises(ValueError):
+            RecorderConfig(
+                self.base,
+                ("/recorded", "/other"),
+                required_topics=("/recorded",),
+                minimum_topic_rates_hz={"/other": 1.0},
+            )
+        for invalid_rate in (0, -1, float("nan"), float("inf"), True, "20"):
+            with self.subTest(invalid_rate=invalid_rate):
+                with self.assertRaises(ValueError):
+                    RecorderConfig(
+                        self.base,
+                        ("/recorded",),
+                        required_topics=("/recorded",),
+                        minimum_topic_rates_hz={"/recorded": invalid_rate},
+                    )
+        with self.assertRaises(ValueError):
+            RecorderConfig(
+                self.base,
+                ("/recorded",),
+                required_topics=("/recorded",),
+                minimum_topic_rates_hz=[],  # type: ignore[arg-type]
+            )
 
 
 if __name__ == "__main__":