tests/dataset/test_record_stick.py

"""
test_record_stick.py
====================
Verify Stick environment (PatternLock) and non-Stick environment (PickXtimes)
when RecordWrapper records HDF5, whether the following four dimensions are correctly aligned:

1. gripper_state  : Stick → [0.0, 0.0]；non-Stick → shape==(2,)
2. joint_action   : Stick → shape==(8,) and [-1] == -1.0；non-Stick → shape==(8,)
3. eef_action     : Stick → shape==(7,) and [-1] == -1.0；non-Stick → shape==(7,)
4. waypoint_action: shape==(7,)；finite Stick → [-1] == -1.0，non-Stick → ±1.0；
                    non-finite (NaN/Inf) is treated as a "no keypoint" placeholder, skip sign assertion

Test method: refer to generate-dataset-control-seed-readJson-advanceV3.py,
run a complete episode for each test case using FailAware Planner + screw→RRT* retry patch
(with seed retry), then open the generated HDF5 file and assert item by item.

Run (requires display / headless GPU):
    cd /data/hongzefu/robomme_benchmark
    uv run python tests/dataset/test_record_stick.py
"""

from __future__ import annotations

import sys
import tempfile
import traceback
from pathlib import Path

import h5py
import numpy as np
import pytest

from tests._shared.dataset_generation import DatasetCase, DatasetFactoryCache
from tests._shared.repo_paths import find_repo_root

pytestmark = pytest.mark.dataset

# ── Ensure robomme package can be found (compatible with direct main() run) ──────────────────────────────────
_PROJECT_ROOT = find_repo_root(__file__)
sys.path.insert(0, str(_PROJECT_ROOT / "src"))


# ────────────────────────────────────────────────────────────────────────────
# Assertion functions
# ────────────────────────────────────────────────────────────────────────────

def _verify_stick(h5_path: Path, env_id: str):
    """Verify Stick environment HDF5 data assertions."""
    print(f"\n  [Verify Stick] Opening {h5_path.name}")
    with h5py.File(h5_path, "r") as f:
        episode_keys = [k for k in f.keys() if k.startswith("episode_")]
        assert len(episode_keys) > 0, "No episode group in HDF5 file"
        ep_grp = f[episode_keys[0]]
        ts_keys = [k for k in ep_grp.keys() if k.startswith("timestep_")]
        assert len(ts_keys) > 0, "No timestep in episode group"

        for ts_key in ts_keys:
            ts = ep_grp[ts_key]

            # 1. gripper_state → [0.0, 0.0]
            gs = np.array(ts["obs"]["gripper_state"])
            assert gs.shape == (2,), \
                f"[{env_id}/{ts_key}] gripper_state shape={gs.shape} expected (2,)"
            assert np.allclose(gs, 0.0), \
                f"[{env_id}/{ts_key}] gripper_state={gs} expected [0.0, 0.0]"

            # 2. joint_action → 8D, last bit == -1.0
            ja = np.array(ts["action"]["joint_action"]).flatten()
            assert ja.shape == (8,), \
                f"[{env_id}/{ts_key}] joint_action shape={ja.shape} expected (8,)"
            assert float(ja[-1]) == -1.0, \
                f"[{env_id}/{ts_key}] joint_action[-1]={ja[-1]} expected -1.0"

            # 3. eef_action → 7D, last bit == -1.0
            ea = np.array(ts["action"]["eef_action"]).flatten()
            assert ea.shape == (7,), \
                f"[{env_id}/{ts_key}] eef_action shape={ea.shape} expected (7,)"
            assert float(ea[-1]) == -1.0, \
                f"[{env_id}/{ts_key}] eef_action[-1]={ea[-1]} expected -1.0"

            # 4. waypoint_action → 7D; non-finite as no keypoint placeholder, finite then verify sign
            wa = np.array(ts["action"]["waypoint_action"]).flatten()
            assert wa.shape == (7,), \
                f"[{env_id}/{ts_key}] waypoint_action shape={wa.shape} expected (7,)"
            if np.all(np.isfinite(wa)):
                assert float(wa[-1]) == -1.0, \
                    f"[{env_id}/{ts_key}] waypoint_action[-1]={wa[-1]} expected -1.0"

    print(f"  [Verify Stick ✓] {env_id} all assertions passed, total {len(ts_keys)} timesteps")


def _verify_non_stick(h5_path: Path, env_id: str):
    """Verify non-Stick environment HDF5 data assertions (original logic not broken)."""
    print(f"\n  [Verify Non-Stick] Opening {h5_path.name}")
    with h5py.File(h5_path, "r") as f:
        episode_keys = [k for k in f.keys() if k.startswith("episode_")]
        assert len(episode_keys) > 0, "No episode group in HDF5 file"
        ep_grp = f[episode_keys[0]]
        ts_keys = [k for k in ep_grp.keys() if k.startswith("timestep_")]
        assert len(ts_keys) > 0, "No timestep in episode group"

        for ts_key in ts_keys:
            ts = ep_grp[ts_key]

            # 1. gripper_state shape == (2,)
            gs = np.array(ts["obs"]["gripper_state"])
            assert gs.shape == (2,), \
                f"[{env_id}/{ts_key}] gripper_state shape={gs.shape} expected (2,)"

            # 2. joint_action → 8D
            ja = np.array(ts["action"]["joint_action"]).flatten()
            assert ja.shape == (8,), \
                f"[{env_id}/{ts_key}] joint_action shape={ja.shape} expected (8,)"

            # 3. eef_action → 7D
            ea = np.array(ts["action"]["eef_action"]).flatten()
            assert ea.shape == (7,), \
                f"[{env_id}/{ts_key}] eef_action shape={ea.shape} expected (7,)"

            # 4. waypoint_action → 7D; non-finite as no keypoint placeholder, finite then verify sign
            wa = np.array(ts["action"]["waypoint_action"]).flatten()
            assert wa.shape == (7,), \
                f"[{env_id}/{ts_key}] waypoint_action shape={wa.shape} expected (7,)"
            if np.all(np.isfinite(wa)):
                assert float(wa[-1]) in (-1.0, 1.0), \
                    f"[{env_id}/{ts_key}] waypoint_action[-1]={wa[-1]} should be ±1.0"

    print(f"  [Verify Non-Stick ✓] {env_id} all assertions passed, total {len(ts_keys)} timesteps")


# ────────────────────────────────────────────────────────────────────────────
# Test case configuration
# ────────────────────────────────────────────────────────────────────────────

# (env_id, is_stick, episode, base_seed, difficulty)
# base_seed is unrelated to the seed corresponding to SOURCE_METADATA_ROOT in the V3 script,
# here directly use the SEED_OFFSET rule of generate_dataset.py
TEST_CASES = [
    ("PatternLock", True,  0, 510001, "easy"),
    ("PickXtimes",  False, 0, 504101, "easy"),
]


def _make_case(env_id: str, episode: int, base_seed: int, difficulty: str | None) -> DatasetCase:
    return DatasetCase(
        env_id=env_id,
        episode=episode,
        base_seed=base_seed,
        difficulty=difficulty,
        save_video=True,
        mode_tag="stick_record_replay",
    )


@pytest.mark.parametrize("env_id,is_stick,episode,base_seed,difficulty", TEST_CASES)
def test_record_stick_case(
    env_id: str,
    is_stick: bool,
    episode: int,
    base_seed: int,
    difficulty: str | None,
    dataset_factory,
):
    generated = dataset_factory(_make_case(env_id, episode, base_seed, difficulty))
    if is_stick:
        _verify_stick(generated.raw_h5_path, env_id)
    else:
        _verify_non_stick(generated.raw_h5_path, env_id)


def main():
    all_pass = True
    results = []

    with tempfile.TemporaryDirectory(prefix="test_record_shared_cache_") as tmpdir:
        cache = DatasetFactoryCache(Path(tmpdir))
        for env_id, is_stick, episode, base_seed, difficulty in TEST_CASES:
            print(f"\n{'='*60}")
            print(f"Test case: {env_id}  (is_stick={is_stick}, ep={episode}, base_seed={base_seed})")
            print(f"{'='*60}")
            try:
                generated = cache.get(_make_case(env_id, episode, base_seed, difficulty))
                if is_stick:
                    _verify_stick(generated.raw_h5_path, env_id)
                else:
                    _verify_non_stick(generated.raw_h5_path, env_id)
                results.append((env_id, "PASS", None))
            except AssertionError as exc:
                results.append((env_id, "FAIL", str(exc)))
                all_pass = False
                print(f"\n  [Assertion failed] {exc}")
                traceback.print_exc()
            except Exception as exc:
                results.append((env_id, "ERROR", str(exc)))
                all_pass = False
                print(f"\n  [Error] {exc}")
                traceback.print_exc()

    print(f"\n{'='*60}")
    print("Test results summary")
    print(f"{'='*60}")
    for env_id, status, msg in results:
        marker = "✓" if status == "PASS" else "✗"
        suffix = f"  ({msg})" if msg else ""
        print(f"  {marker} [{status}] {env_id}{suffix}")

    if all_pass:
        print("\n✓ ALL ASSERTIONS PASSED")
        sys.exit(0)
    else:
        print("\n✗ SOME ASSERTIONS FAILED")
        sys.exit(1)


if __name__ == "__main__":
    main()