meiqi/TG3
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: keep RGB-D captures only on workstation

Browse Source
This commit is contained in:
meiqi
2026-08-10 17:14:44 +08:00
parent d02b99d2aa
commit 67c9b65768
15 changed files with 2817 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files

221
tg3_local_teleop/test_data_collection.py
View File

@@ -15,12 +15,45 @@ from typing import Any
from data_collection import (
DataRecorderManager,
OptionalTopicGroupConfig,
RecorderConfig,
RecordingToggleGate,
left_joystick_pressed,
)
HEAD_CAMERA_TOPICS = (
"/ob_camera_head/color/image_raw/compressed",
"/ob_camera_head/color/camera_info",
"/ob_camera_head/color/metadata",
"/ob_camera_head/depth/image_raw/compressedDepth",
"/ob_camera_head/depth/camera_info",
"/ob_camera_head/depth/metadata",
)
WAIST_CAMERA_TOPICS = tuple(
topic.replace("_head", "_waist") for topic in HEAD_CAMERA_TOPICS
)
def optional_camera_groups() -> dict[str, OptionalTopicGroupConfig]:
return {
"head_rgbd": OptionalTopicGroupConfig(
topics=HEAD_CAMERA_TOPICS,
minimum_topic_rates_hz={
HEAD_CAMERA_TOPICS[0]: 20.0,
HEAD_CAMERA_TOPICS[3]: 20.0,
},
),
"waist_rgbd": OptionalTopicGroupConfig(
topics=WAIST_CAMERA_TOPICS,
minimum_topic_rates_hz={
WAIST_CAMERA_TOPICS[0]: 20.0,
WAIST_CAMERA_TOPICS[3]: 20.0,
},
),
}
class LeftJoystickParserTest(unittest.TestCase):
def test_accepts_only_live_bool_or_binary_integer(self) -> None:
self.assertIs(
@@ -291,6 +324,10 @@ class DataRecorderManagerTest(unittest.TestCase):
topics: Sequence[str] = ("/joint_states", "/tf"),
required_topics: Sequence[str] = ("/joint_states", "/tf"),
minimum_topic_rates_hz: Mapping[str, float] | None = None,
optional_topic_groups: Mapping[
str, OptionalTopicGroupConfig
] | None = None,
retain_failed_episodes: bool = True,
) -> tuple[DataRecorderManager, Any]:
if factory is None:
factory = FakeProcessFactory(**(process_options or {}))
@@ -305,6 +342,8 @@ class DataRecorderManagerTest(unittest.TestCase):
topics=topics,
required_topics=required_topics,
minimum_topic_rates_hz=minimum_topic_rates_hz or {},
optional_topic_groups=optional_topic_groups or {},
retain_failed_episodes=retain_failed_episodes,
minimum_free_bytes=minimum_free,
max_duration_seconds=max_duration,
poll_interval_seconds=0.005,
@@ -330,6 +369,36 @@ class DataRecorderManagerTest(unittest.TestCase):
self.fail("condition did not become true before timeout")
time.sleep(0.002)
def record_with_optional_cameras(
self, topic_counts: Mapping[str, int]
) -> tuple[dict[str, Any], dict[str, Any]]:
topics = (
"/joint_states",
"/tf",
*HEAD_CAMERA_TOPICS,
*WAIST_CAMERA_TOPICS,
)
manager, _ = self.make_manager(
topics=topics,
required_topics=("/joint_states", "/tf"),
optional_topic_groups=optional_camera_groups(),
process_options={
"topic_counts": topic_counts,
"duration_nanoseconds": 1_000_000_000,
},
)
self.assertTrue(manager.request_start("ca" * 16))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
manifest = json.loads(
(
Path(status["last_episode_directory"]) / "manifest.json"
).read_text()
)
return status, manifest
def test_manual_stop_uses_sigint_and_atomically_creates_ready_manifest(
self,
) -> None:
@@ -627,6 +696,103 @@ class DataRecorderManagerTest(unittest.TestCase):
0 if head_count is None else head_count,
)
def test_optional_cameras_absent_do_not_block_core_ready(self) -> None:
status, manifest = self.record_with_optional_cameras(
{"/joint_states": 25, "/tf": 50}
)
self.assertEqual(status["last_result"], "ready")
self.assertEqual(manifest["data_quality_warnings"], [])
self.assertEqual(
{
name: observation["state"]
for name, observation in manifest[
"optional_topic_groups"
].items()
},
{"head_rgbd": "absent", "waist_rgbd": "absent"},
)
self.assertEqual(
manifest["optional_topic_groups"]["head_rgbd"][
"below_minimum_rate_topics"
],
[],
)
self.assertTrue((Path(status["last_episode_directory"]) / "READY").exists())
def test_each_active_optional_camera_is_observed_as_healthy(self) -> None:
for active_name, active_topics in (
("head_rgbd", HEAD_CAMERA_TOPICS),
("waist_rgbd", WAIST_CAMERA_TOPICS),
):
with self.subTest(active_name=active_name):
counts = {"/joint_states": 25, "/tf": 50}
counts.update({topic: 1 for topic in active_topics})
counts[active_topics[0]] = 25
counts[active_topics[3]] = 25
status, manifest = self.record_with_optional_cameras(counts)
self.assertEqual(status["last_result"], "ready")
observation = manifest["optional_topic_groups"][active_name]
self.assertEqual(observation["state"], "healthy")
self.assertEqual(
observation["observed_topic_rates_hz"],
{active_topics[0]: 25.0, active_topics[3]: 25.0},
)
inactive_name = (
"waist_rgbd"
if active_name == "head_rgbd"
else "head_rgbd"
)
self.assertEqual(
manifest["optional_topic_groups"][inactive_name]["state"],
"absent",
)
self.assertEqual(manifest["data_quality_warnings"], [])
def test_partial_optional_camera_warns_but_remains_ready(self) -> None:
counts = {"/joint_states": 25, "/tf": 50}
counts.update(
{
HEAD_CAMERA_TOPICS[0]: 25,
HEAD_CAMERA_TOPICS[1]: 1,
HEAD_CAMERA_TOPICS[2]: 1,
}
)
status, manifest = self.record_with_optional_cameras(counts)
self.assertEqual(status["last_result"], "ready")
observation = manifest["optional_topic_groups"]["head_rgbd"]
self.assertEqual(observation["state"], "partial")
self.assertEqual(
observation["zero_message_topics"],
list(HEAD_CAMERA_TOPICS[3:]),
)
self.assertTrue(manifest["data_quality_warnings"])
self.assertIn("is partial", manifest["data_quality_warnings"][0])
self.assertTrue((Path(status["last_episode_directory"]) / "READY").exists())
def test_low_rate_optional_camera_warns_but_remains_ready(self) -> None:
counts = {"/joint_states": 25, "/tf": 50}
counts.update({topic: 1 for topic in WAIST_CAMERA_TOPICS})
counts[WAIST_CAMERA_TOPICS[0]] = 5
counts[WAIST_CAMERA_TOPICS[3]] = 10
status, manifest = self.record_with_optional_cameras(counts)
self.assertEqual(status["last_result"], "ready")
observation = manifest["optional_topic_groups"]["waist_rgbd"]
self.assertEqual(observation["state"], "low_rate")
self.assertEqual(
observation["observed_topic_rates_hz"],
{
WAIST_CAMERA_TOPICS[0]: 5.0,
WAIST_CAMERA_TOPICS[3]: 10.0,
},
)
self.assertEqual(
observation["below_minimum_rate_topics"],
[WAIST_CAMERA_TOPICS[0], WAIST_CAMERA_TOPICS[3]],
)
self.assertTrue(manifest["data_quality_warnings"])
self.assertIn("below", manifest["data_quality_warnings"][0])
self.assertTrue((Path(status["last_episode_directory"]) / "READY").exists())
def test_topic_that_stops_mid_episode_fails_average_rate(self) -> None:
bag_info = FakeBagInfoRunner()
manager, _ = self.make_manager(
@@ -779,6 +945,20 @@ class DataRecorderManagerTest(unittest.TestCase):
)
self.assertEqual(failure["status"], "failed")
def test_failed_payload_is_discarded_when_retention_is_disabled(self) -> None:
manager, factory = self.make_manager(
free_bytes=lambda _path: 0,
retain_failed_episodes=False,
)
self.assertTrue(manager.request_start("d0" * 16))
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "failed")
self.assertIsNone(status["last_episode_directory"])
self.assertEqual(factory.processes, [])
self.assertEqual(list((self.base / "active").iterdir()), [])
self.assertEqual(list((self.base / "failed").iterdir()), [])
def test_configuration_rejects_unsafe_or_ambiguous_values(self) -> None:
with self.assertRaises(ValueError):
RecorderConfig(self.base, ())
@@ -790,6 +970,8 @@ class DataRecorderManagerTest(unittest.TestCase):
RecorderConfig(self.base, ("/ok",), minimum_free_bytes=-1)
with self.assertRaises(ValueError):
RecorderConfig(self.base, ("/ok",), validate_bag_info=1)
with self.assertRaises(ValueError):
RecorderConfig(self.base, ("/ok",), retain_failed_episodes=1)
with self.assertRaises(ValueError):
RecorderConfig(self.base, ("/ok",), bag_info_timeout_seconds=0)
with self.assertRaises(ValueError):
@@ -827,6 +1009,45 @@ class DataRecorderManagerTest(unittest.TestCase):
required_topics=("/recorded",),
minimum_topic_rates_hz=[], # type: ignore[arg-type]
)
with self.assertRaises(ValueError):
OptionalTopicGroupConfig(topics=())
with self.assertRaises(ValueError):
OptionalTopicGroupConfig(
topics=("/camera/image",),
minimum_topic_rates_hz={"/other": 20.0},
)
camera_group = OptionalTopicGroupConfig(
topics=("/camera/image",),
minimum_topic_rates_hz={"/camera/image": 20.0},
)
with self.assertRaises(ValueError):
RecorderConfig(
self.base,
("/recorded",),
optional_topic_groups={"camera": camera_group},
)
with self.assertRaises(ValueError):
RecorderConfig(
self.base,
("/recorded", "/camera/image"),
required_topics=("/camera/image",),
optional_topic_groups={"camera": camera_group},
)
with self.assertRaises(ValueError):
RecorderConfig(
self.base,
("/recorded", "/camera/image"),
optional_topic_groups={
"camera_a": camera_group,
"camera_b": camera_group,
},
)
with self.assertRaises(ValueError):
RecorderConfig(
self.base,
("/recorded", "/camera/image"),
optional_topic_groups={"bad name": camera_group},
)
if __name__ == "__main__":