feat: 流式保存ACT HDF5数据

xr_rm_teleop/test/test_act_episode_recorder.py

@@ -1,14 +1,26 @@
 import numpy as np
 import pytest
 
+try:
+    import h5py
+except ImportError:
+    h5py = None
+
 from xr_rm_teleop.act_episode_recorder import (
     NO_ACTION,
     ButtonTracker,
     CameraBuffer,
     CameraFrame,
+    EpisodeMetadata,
+    EpisodeStore,
     QualityError,
     RecordingSession,
     RecordingState,
+    TaskDirectoryLock,
+    discard_partial,
+    next_episode_index,
+    publish_without_overwrite,
+    recover_partial_files,
     select_camera_pair,
     select_frame,
 )
@@ -273,3 +285,195 @@ def test_camera_buffer_is_bounded_and_counts_dropped_frames():
     assert stats.dropped_frames == 1
     assert stats.drop_ratio == pytest.approx(1.0 / 6.0)
     assert stats.fps == pytest.approx(800.0)
+
+
+requires_h5py = pytest.mark.skipif(
+    h5py is None,
+    reason="h5py is not installed",
+)
+
+
+def _metadata():
+    return EpisodeMetadata(
+        joint_names=tuple(f"joint_{index}" for index in range(1, 8))
+        + ("gripper",),
+        joint_lower_limits=np.full(7, -3.0, dtype=np.float64),
+        joint_upper_limits=np.full(7, 3.0, dtype=np.float64),
+    )
+
+
+def _episode_sample(seq):
+    control_ns = seq * 11_111_111
+    qpos = np.zeros(8, dtype=np.float32)
+    qpos[-1] = 1.0
+    action = qpos.copy()
+    return {
+        "observations/qpos": qpos,
+        "action": action,
+        "observations/images/cam_high": _image(seq % 255),
+        "observations/images/cam_right_wrist": _image((seq + 1) % 255),
+        "debug/timestamps/control_monotonic_ns": control_ns,
+        "debug/timestamps/feedback_monotonic_ns": control_ns - 1_000_000,
+        "debug/timestamps/action_monotonic_ns": control_ns,
+        "debug/timestamps/cam_high_host_monotonic_ns": control_ns - 2_000_000,
+        "debug/timestamps/cam_wrist_host_monotonic_ns": control_ns - 3_000_000,
+        "debug/timestamps/cam_high_hardware_ms": float(seq),
+        "debug/timestamps/cam_wrist_hardware_ms": float(seq),
+        "debug/timestamps/cam_high_age_ms": 2.0,
+        "debug/timestamps/cam_wrist_age_ms": 3.0,
+        "debug/timestamps/inter_camera_skew_ms": 1.0,
+        "debug/cameras/cam_high_frame_number": seq,
+        "debug/cameras/cam_wrist_frame_number": seq,
+        "debug/control/control_seq": seq,
+        "debug/control/teleop_active": 1,
+        "debug/control/action_valid": 1,
+        "debug/control/command_sent": 1,
+        "debug/control/target_clamped": 0,
+        "debug/control/control_fault": 0,
+        "debug/qp/raw_target": np.zeros(7, dtype=np.float32),
+        "debug/qp/attempted": 1,
+        "debug/qp/success": 1,
+        "debug/qp/duration_ms": 1.0,
+        "debug/tcp/current_pose": np.zeros(7, dtype=np.float32),
+        "debug/tcp/raw_target_pose": np.zeros(7, dtype=np.float32),
+        "debug/tcp/final_target_pose": np.zeros(7, dtype=np.float32),
+        "debug/tcp/command_velocity": np.zeros(6, dtype=np.float32),
+        "debug/pico/right_pose": np.zeros(7, dtype=np.float32),
+        "debug/pico/right_inputs": np.zeros(6, dtype=np.float32),
+        "debug/pico/left_secondary": 0,
+        "debug/gripper/target_open": 1,
+        "debug/gripper/state_open": 1,
+        "debug/gripper/command_pending": 0,
+        "debug/gripper/command_failed": 0,
+    }
+
+
+@requires_h5py
+def test_episode_store_writes_act_core_schema(tmp_path):
+    store = EpisodeStore.create(
+        tmp_path / "episode_0.partial.hdf5",
+        _metadata(),
+    )
+    for seq in (100, 103, 106):
+        store.append(_episode_sample(seq))
+    store.close()
+
+    with h5py.File(store.path, "r") as root:
+        assert root.attrs["sim"] == np.bool_(False)
+        assert root.attrs["action_alignment"] == "same_step_causal"
+        assert root["observations/qpos"].shape == (3, 8)
+        assert root["observations/qpos"].dtype == np.float32
+        assert root["action"].shape == (3, 8)
+        assert root["action"].dtype == np.float32
+        assert root["observations/images/cam_high"].shape == (
+            3,
+            480,
+            640,
+            3,
+        )
+        assert root["observations/images/cam_high"].dtype == np.uint8
+        assert root["observations/images/cam_right_wrist"].shape == (
+            3,
+            480,
+            640,
+            3,
+        )
+        assert "observations/qvel" not in root
+        assert "observations/effort" not in root
+        assert "compress_len" not in root
+
+
+@requires_h5py
+def test_episode_store_truncates_every_time_axis_dataset(tmp_path):
+    store = EpisodeStore.create(
+        tmp_path / "episode_0.partial.hdf5",
+        _metadata(),
+    )
+    for seq in (100, 103, 106):
+        store.append(_episode_sample(seq))
+
+    store.truncate(2)
+    store.close()
+
+    lengths = []
+    with h5py.File(store.path, "r") as root:
+        root.visititems(
+            lambda _name, item: lengths.append(item.shape[0])
+            if isinstance(item, h5py.Dataset)
+            else None
+        )
+    assert lengths
+    assert set(lengths) == {2}
+
+
+def test_next_index_uses_max_saved_episode_and_ignores_rejected(tmp_path):
+    (tmp_path / "episode_2.hdf5").touch()
+    (tmp_path / "episode_9.hdf5").touch()
+    rejected = tmp_path / "rejected"
+    rejected.mkdir()
+    (rejected / "episode_20_bad_20260810.hdf5").touch()
+
+    assert next_episode_index(tmp_path) == 10
+
+
+def test_publish_never_overwrites_existing_episode(tmp_path):
+    partial = tmp_path / "episode_1.partial.hdf5"
+    partial.write_bytes(b"new")
+    final = tmp_path / "episode_1.hdf5"
+    final.write_bytes(b"old")
+
+    with pytest.raises(FileExistsError):
+        publish_without_overwrite(partial, final)
+
+    assert final.read_bytes() == b"old"
+    assert partial.read_bytes() == b"new"
+
+
+@requires_h5py
+def test_recover_marks_readable_and_preserves_unreadable_partial(tmp_path):
+    readable = tmp_path / "episode_1.partial.hdf5"
+    store = EpisodeStore.create(readable, _metadata())
+    store.append(_episode_sample(100))
+    store.close()
+    unreadable = tmp_path / "episode_2.partial.hdf5"
+    unreadable.write_bytes(b"not hdf5")
+
+    recovered = recover_partial_files(tmp_path, timestamp="20260810T120000")
+
+    assert len(recovered) == 2
+    rejected = tmp_path / "rejected" / (
+        "episode_1_crash_recovered_20260810T120000.hdf5"
+    )
+    assert rejected in recovered
+    with h5py.File(rejected, "r") as root:
+        assert root.attrs["episode_status"] == "rejected"
+        assert root.attrs["reject_reason"] == "crash_recovered"
+        assert root.attrs["interrupted"] == np.bool_(True)
+    assert not readable.exists()
+    assert any(path.name.endswith(".partial.hdf5") for path in recovered)
+    assert unreadable not in recovered
+    assert not unreadable.exists()
+
+
+def test_discard_partial_removes_only_current_file(tmp_path):
+    current = tmp_path / "episode_1.partial.hdf5"
+    current.write_bytes(b"current")
+    saved = tmp_path / "episode_0.hdf5"
+    saved.write_bytes(b"saved")
+
+    discard_partial(current)
+
+    assert not current.exists()
+    assert saved.read_bytes() == b"saved"
+    assert next_episode_index(tmp_path) == 1
+
+
+def test_task_directory_lock_rejects_second_recorder(tmp_path):
+    first = TaskDirectoryLock(tmp_path)
+    second = TaskDirectoryLock(tmp_path)
+    first.acquire()
+    try:
+        with pytest.raises(BlockingIOError):
+            second.acquire()
+    finally:
+        first.release()