feat: 添加ACT录制状态机

2026-08-10 18:03:25 +08:00
parent e60d620dfe
commit deeee076d7
2 changed files with 393 additions and 0 deletions
@@ -0,0 +1,186 @@
+from xr_rm_teleop.act_episode_recorder import (
+    NO_ACTION,
+    ButtonTracker,
+    RecordingSession,
+    RecordingState,
+)
+
+
+def _recording_session(*, origin_seq=10, max_samples=1800):
+    session = RecordingSession(max_samples=max_samples)
+    session.arm()
+    decision = session.on_control(
+        origin_seq,
+        grip=True,
+        action_valid=True,
+        command_sent=True,
+    )
+    assert decision.record_sample
+    return session
+
+
+def test_recording_starts_on_first_sent_grip_action_and_samples_every_third_cycle():
+    session = RecordingSession(max_samples=1800)
+    session.arm()
+
+    assert session.on_control(
+        100,
+        grip=False,
+        action_valid=True,
+        command_sent=False,
+    ) == NO_ACTION
+    assert session.on_control(
+        101,
+        grip=True,
+        action_valid=True,
+        command_sent=False,
+    ) == NO_ACTION
+    first = session.on_control(
+        102,
+        grip=True,
+        action_valid=True,
+        command_sent=True,
+    )
+    second = session.on_control(
+        103,
+        grip=True,
+        action_valid=True,
+        command_sent=True,
+    )
+    session.on_control(
+        104,
+        grip=True,
+        action_valid=True,
+        command_sent=True,
+    )
+    third = session.on_control(
+        105,
+        grip=True,
+        action_valid=True,
+        command_sent=True,
+    )
+
+    assert first.record_sample
+    assert not second.record_sample
+    assert third.record_sample
+    assert session.sample_origin_seq == 102
+
+
+def test_final_grip_release_marks_crop_point_but_mid_pause_is_kept():
+    session = _recording_session(origin_seq=10)
+    session.on_control(11, grip=True, action_valid=True, command_sent=True)
+    session.on_control(12, grip=True, action_valid=True, command_sent=True)
+    session.on_control(13, grip=False, action_valid=True, command_sent=False)
+    first_crop = session.candidate_end_count
+
+    assert first_crop == 2
+
+    session.on_control(14, grip=True, action_valid=True, command_sent=False)
+    assert session.candidate_end_count is None
+    session.on_control(15, grip=True, action_valid=True, command_sent=True)
+    session.on_control(16, grip=False, action_valid=True, command_sent=False)
+
+    assert session.candidate_end_count == 3
+
+
+def test_finish_waits_for_control_sample_after_b_request():
+    session = _recording_session(origin_seq=10)
+    session.on_control(11, grip=False, action_valid=True, command_sent=False)
+    session.request_finish()
+
+    decision = session.on_control(
+        12,
+        grip=False,
+        action_valid=True,
+        command_sent=False,
+    )
+
+    assert decision.finish
+    assert session.candidate_end_count == 1
+
+
+def test_missing_control_sequence_rejects_recording():
+    session = _recording_session(origin_seq=10)
+    session.on_control(11, grip=True, action_valid=True, command_sent=True)
+
+    decision = session.on_control(
+        13,
+        grip=True,
+        action_valid=True,
+        command_sent=True,
+    )
+
+    assert decision.reject_reason == "control_sequence_gap"
+
+
+def test_max_samples_rejects_without_stopping_robot():
+    session = _recording_session(origin_seq=10, max_samples=2)
+    session.on_control(11, grip=True, action_valid=True, command_sent=True)
+    session.on_control(12, grip=True, action_valid=True, command_sent=True)
+
+    decision = session.on_control(
+        13,
+        grip=True,
+        action_valid=True,
+        command_sent=True,
+    )
+
+    assert decision.reject_reason == "max_duration"
+
+
+def test_right_b_is_ignored_while_grip_is_pressed():
+    tracker = ButtonTracker(hold_ns=1_000_000_000)
+    tracker.on_right(False, False, False, 0, RecordingState.IDLE)
+
+    events = tracker.on_right(
+        False,
+        True,
+        True,
+        1,
+        RecordingState.IDLE,
+    )
+
+    assert not events.right_b
+
+
+def test_left_y_requires_new_press_inside_active_recording_state():
+    tracker = ButtonTracker(hold_ns=1_000_000_000)
+    tracker.on_left(True, 0, RecordingState.IDLE)
+    assert not tracker.on_left(
+        True,
+        2_000_000_000,
+        RecordingState.RECORDING,
+    ).discard
+    tracker.on_left(False, 2_100_000_000, RecordingState.RECORDING)
+    tracker.on_left(True, 3_000_000_000, RecordingState.RECORDING)
+
+    assert not tracker.on_left(
+        True,
+        3_999_999_999,
+        RecordingState.RECORDING,
+    ).discard
+    assert tracker.on_left(
+        True,
+        4_000_000_000,
+        RecordingState.RECORDING,
+    ).discard
+    assert not tracker.on_left(
+        True,
+        5_000_000_000,
+        RecordingState.RECORDING,
+    ).discard
+
+
+def test_right_a_reports_recording_rejection_event():
+    tracker = ButtonTracker(hold_ns=1_000_000_000)
+    tracker.on_right(False, False, False, 0, RecordingState.RECORDING)
+
+    events = tracker.on_right(
+        True,
+        False,
+        False,
+        1,
+        RecordingState.RECORDING,
+    )
+
+    assert events.reject_reason == "initial_pose_command_during_episode"