feat: 集成ACT episode采集节点

2026-08-10 18:22:35 +08:00
parent c1ea1a2816
commit 6008e34b5d
3 changed files with 1009 additions and 1 deletions
@@ -1,3 +1,7 @@
+import queue
+import threading
+from types import SimpleNamespace
+
 import numpy as np
 import pytest

@@ -6,13 +10,16 @@ try:
 except ImportError:
    h5py = None

+from xr_rm_interfaces.msg import ActControlSample, XrController
 from xr_rm_teleop.act_episode_recorder import (
    NO_ACTION,
+    ActEpisodeRecorder,
    ButtonTracker,
    CameraBuffer,
    CameraFrame,
    EpisodeMetadata,
    EpisodeStore,
+    EpisodeWriter,
    QualityError,
    QualityLimits,
    RecordingSession,
@@ -598,3 +605,257 @@ def test_validate_episode_reports_qp_failures_without_rejecting(tmp_path):
    assert report.metrics["qp_failure_ratio"] == pytest.approx(2.0 / 3.0)
    assert report.metrics["qp_longest_failure_streak"] == 2
    assert report.metrics["target_clamped_count"] == 2
+
+
+class _StatusPublisher:
+    def __init__(self):
+        self.messages = []
+
+    def publish(self, message):
+        self.messages.append(message.data)
+
+
+class _Logger:
+    def info(self, *_args, **_kwargs):
+        pass
+
+    def warn(self, *_args, **_kwargs):
+        pass
+
+    def error(self, *_args, **_kwargs):
+        pass
+
+
+def _control_message(seq, control_ns, *, grip=True):
+    message = ActControlSample()
+    message.control_seq = seq
+    message.control_monotonic_ns = control_ns
+    message.feedback_monotonic_ns = control_ns - 1_000_000
+    message.action_monotonic_ns = control_ns + 1_000_000 if grip else -1
+    message.q_actual = [0.0] * 7
+    message.q_qp_raw = [0.0] * 7
+    message.q_target = [0.0] * 7
+    message.joint_lower_limits = [-3.0] * 7
+    message.joint_upper_limits = [3.0] * 7
+    message.tcp_current.orientation.w = 1.0
+    message.tcp_raw_target.orientation.w = 1.0
+    message.tcp_target.orientation.w = 1.0
+    message.pico_pose.orientation.w = 1.0
+    message.pico_grip = grip
+    message.gripper_target_open = True
+    message.gripper_state_open = True
+    message.gripper_state_known = True
+    message.teleop_active = grip
+    message.feedback_valid = True
+    message.action_valid = True
+    message.command_sent = grip
+    message.qp_attempted = grip
+    message.qp_success = grip
+    return message
+
+
+def _seed_camera(start_ns, first_number):
+    camera = SimpleNamespace(
+        buffer=CameraBuffer(maxlen=4),
+        last_error=None,
+    )
+    image = _image(1)
+    for index in range(151):
+        camera.buffer.push(
+            CameraFrame(
+                image,
+                first_number + index,
+                index * (1000.0 / 30.0),
+                start_ns
+                - 5_020_000_000
+                + round(index * 5_000_000_000 / 150),
+            )
+        )
+    return camera
+
+
+def _recorder_for_test(tmp_path):
+    now_ns = 10_000_000_000
+    recorder = object.__new__(ActEpisodeRecorder)
+    recorder._task_dir = tmp_path / "tomato_pick"
+    recorder._task_dir.mkdir(parents=True)
+    recorder._quality_limits = _quality_limits()
+    recorder._min_free_space_bytes = 4 * 1024**3
+    recorder._controller_timeout_ns = 500_000_000
+    recorder._camera_warmup_ns = 5_000_000_000
+    recorder._gripper_completion_timeout_ns = 3_000_000_000
+    recorder._writer_queue_size = 8
+    recorder._max_camera_age_ms = 50.0
+    recorder._max_camera_skew_ms = 50.0
+    recorder._session = RecordingSession(max_samples=100)
+    recorder._button_tracker = ButtonTracker(hold_ns=1_000_000_000)
+    recorder._latest_control = _control_message(99, now_ns - 1_000_000)
+    recorder._latest_control_received_ns = now_ns - 1_000_000
+    recorder._right_controller_received_ns = now_ns - 1_000_000
+    recorder._left_controller_received_ns = now_ns - 1_000_000
+    recorder._left_secondary = False
+    recorder._high_camera = _seed_camera(now_ns, 850)
+    recorder._wrist_camera = _seed_camera(now_ns, 1850)
+    recorder._camera_start_error = None
+    recorder._writer = None
+    recorder._store = None
+    recorder._partial_path = None
+    recorder._camera_baselines = None
+    recorder._saving_deadline_ns = None
+    recorder._status_pub = _StatusPublisher()
+    recorder._now_ns = lambda: now_ns
+    recorder._disk_usage = lambda _path: SimpleNamespace(free=5 * 1024**3)
+    recorder.get_logger = lambda: _Logger()
+    return recorder
+
+
+@requires_h5py
+def test_preflight_requires_open_gripper_fresh_inputs_and_disk_space(tmp_path):
+    recorder = _recorder_for_test(tmp_path)
+
+    assert recorder._run_preflight() is None
+
+    recorder._latest_control.gripper_state_open = False
+    assert recorder._run_preflight() == "gripper_not_open"
+    recorder._latest_control.gripper_state_open = True
+    recorder._right_controller_received_ns = 0
+    assert recorder._run_preflight() == "right_controller_stale"
+    recorder._right_controller_received_ns = recorder._now_ns()
+    recorder._disk_usage = lambda _path: SimpleNamespace(free=1024)
+    assert recorder._run_preflight() == "insufficient_disk_space"
+
+
+def _push_recording_frames(recorder, control_ns, frame_number):
+    recorder._high_camera.buffer.push(
+        CameraFrame(
+            _image(2),
+            frame_number,
+            float(frame_number),
+            control_ns - 2_000_000,
+        )
+    )
+    recorder._wrist_camera.buffer.push(
+        CameraFrame(
+            _image(3),
+            frame_number,
+            float(frame_number),
+            control_ns - 3_000_000,
+        )
+    )
+
+
+@requires_h5py
+def test_end_to_end_fake_episode_saves_and_returns_idle(tmp_path):
+    recorder = _recorder_for_test(tmp_path)
+    recorder._handle_right_b(grip=False)
+
+    assert recorder.state is RecordingState.ARMED
+
+    start_ns = recorder._now_ns()
+    frame_number = 1001
+    for offset, seq in enumerate(range(100, 107)):
+        control_ns = start_ns + offset * 11_111_111
+        if (seq - 100) % 3 == 0:
+            _push_recording_frames(recorder, control_ns, frame_number)
+            frame_number += 1
+        recorder._on_control_sample(_control_message(seq, control_ns))
+
+    recorder._on_control_sample(
+        _control_message(107, start_ns + 7 * 11_111_111, grip=False)
+    )
+    recorder._handle_right_b(grip=False)
+    final_ns = start_ns + 8 * 11_111_111
+    _push_recording_frames(recorder, final_ns, frame_number)
+    recorder._on_control_sample(_control_message(108, final_ns, grip=False))
+
+    assert (tmp_path / "tomato_pick" / "episode_0.hdf5").is_file()
+    assert recorder.state is RecordingState.IDLE
+    assert "SAVING" in recorder._status_pub.messages
+    assert recorder._status_pub.messages[-2:] == ["SAVED", "IDLE"]
+
+
+@requires_h5py
+def test_discard_and_interrupt_only_process_current_partial(tmp_path):
+    recorder = _recorder_for_test(tmp_path)
+    recorder._handle_right_b(grip=False)
+    partial = recorder._partial_path
+
+    recorder._discard_current()
+
+    assert not partial.exists()
+    assert recorder.state is RecordingState.IDLE
+
+    recorder._handle_right_b(grip=False)
+    recorder.interrupt_recording("interrupted")
+    rejected = list((recorder._task_dir / "rejected").glob("*.hdf5"))
+    assert len(rejected) == 1
+    with h5py.File(rejected[0], "r") as root:
+        assert root.attrs["reject_reason"] == "interrupted"
+        assert root.attrs["interrupted"] == np.bool_(True)
+
+
+@requires_h5py
+def test_a_button_and_camera_error_reject_without_robot_commands(tmp_path):
+    recorder = _recorder_for_test(tmp_path)
+    recorder._handle_right_b(grip=False)
+    recorder._button_tracker.on_right(
+        False,
+        False,
+        False,
+        recorder._now_ns(),
+        RecordingState.RECORDING,
+    )
+    recorder._session.state = RecordingState.RECORDING
+    controller = XrController()
+    controller.primary = True
+
+    recorder._on_right_controller(controller)
+
+    rejected = list((recorder._task_dir / "rejected").glob("*.hdf5"))
+    assert len(rejected) == 1
+    with h5py.File(rejected[0], "r") as root:
+        assert root.attrs["reject_reason"] == (
+            "initial_pose_command_during_episode"
+        )
+
+    recorder._handle_right_b(grip=False)
+    recorder._high_camera.last_error = RuntimeError("usb")
+    recorder._on_control_sample(
+        _control_message(100, recorder._now_ns())
+    )
+    rejected = list((recorder._task_dir / "rejected").glob("*.hdf5"))
+    assert len(rejected) == 2
+    camera_rejected = next(
+        path for path in rejected if "camera_error" in path.name
+    )
+    with h5py.File(camera_rejected, "r") as root:
+        assert root.attrs["reject_reason"] == "camera_error"
+
+
+def test_episode_writer_reports_queue_backlog_and_write_error():
+    started = threading.Event()
+    release = threading.Event()
+
+    class BlockingStore:
+        def append(self, _sample):
+            started.set()
+            assert release.wait(timeout=1.0)
+
+    writer = EpisodeWriter(BlockingStore(), queue_size=1)
+    writer.submit({})
+    assert started.wait(timeout=1.0)
+    writer.submit({})
+    with pytest.raises(queue.Full):
+        writer.submit({})
+    release.set()
+    writer.finish()
+    assert writer.error is None
+
+    class FailingStore:
+        def append(self, _sample):
+            raise OSError("disk full")
+
+    writer = EpisodeWriter(FailingStore(), queue_size=1)
+    writer.submit({})
+    writer.finish()
+    assert isinstance(writer.error, OSError)