Spaces:
Paused
Paused
| from __future__ import annotations | |
| import h5py | |
| import numpy as np | |
| import pytest | |
| from tests._shared.dataset_generation import DatasetCase | |
| pytestmark = pytest.mark.dataset | |
| def _make_case(env_id: str, base_seed: int) -> DatasetCase: | |
| return DatasetCase( | |
| env_id=env_id, | |
| episode=0, | |
| base_seed=base_seed, | |
| difficulty="easy", | |
| save_video=True, | |
| mode_tag="waypoint_phase_isolation", | |
| ) | |
| def _decode_h5_string(raw) -> str: | |
| if isinstance(raw, np.ndarray): | |
| raw = raw.flatten()[0] | |
| if isinstance(raw, (bytes, np.bytes_)): | |
| return raw.decode("utf-8") | |
| return str(raw) | |
| def _collect_records(ep_group: h5py.Group) -> list[dict]: | |
| timestep_keys = sorted( | |
| (k for k in ep_group.keys() if k.startswith("timestep_")), | |
| key=lambda k: int(k.split("_")[1]), | |
| ) | |
| out: list[dict] = [] | |
| for key in timestep_keys: | |
| ts = ep_group[key] | |
| info = ts["info"] | |
| waypoint_action = np.asarray(ts["action"]["waypoint_action"][()]).flatten() | |
| out.append( | |
| { | |
| "timestep": int(key.split("_")[1]), | |
| "is_demo": bool(np.reshape(np.asarray(info["is_video_demo"][()]), -1)[0]), | |
| "subgoal": _decode_h5_string(info["simple_subgoal_online"][()]), | |
| "waypoint_action": waypoint_action, | |
| "is_finite_waypoint": bool( | |
| waypoint_action.shape == (7,) and np.all(np.isfinite(waypoint_action)) | |
| ), | |
| } | |
| ) | |
| return out | |
| def _find_demo_to_non_demo_boundary(records: list[dict]) -> int | None: | |
| for idx in range(1, len(records)): | |
| if records[idx - 1]["is_demo"] and not records[idx]["is_demo"]: | |
| return idx | |
| return None | |
| def _last_finite_demo_waypoint(records: list[dict], boundary_idx: int) -> np.ndarray | None: | |
| for idx in range(boundary_idx - 1, -1, -1): | |
| row = records[idx] | |
| if not row["is_demo"]: | |
| continue | |
| if row["is_finite_waypoint"]: | |
| return np.asarray(row["waypoint_action"]).flatten() | |
| return None | |
| def _unique_finite_waypoints(rows: list[dict]) -> list[np.ndarray]: | |
| uniques: list[np.ndarray] = [] | |
| prev: np.ndarray | None = None | |
| for row in rows: | |
| if not row["is_finite_waypoint"]: | |
| continue | |
| wa = np.asarray(row["waypoint_action"]).flatten() | |
| if prev is None or not np.array_equal(wa, prev): | |
| uniques.append(wa.copy()) | |
| prev = wa.copy() | |
| return uniques | |
| def test_waypoint_isolation_across_demo_phase( | |
| env_id: str, | |
| base_seed: int, | |
| assert_first_segment_midpoint: bool, | |
| dataset_factory, | |
| ): | |
| generated = dataset_factory(_make_case(env_id, base_seed)) | |
| with h5py.File(generated.raw_h5_path, "r") as h5f: | |
| records = _collect_records(h5f["episode_0"]) | |
| assert records, f"{env_id}: episode_0 has no recorded timesteps." | |
| boundary_idx = _find_demo_to_non_demo_boundary(records) | |
| assert boundary_idx is not None, f"{env_id}: missing demo->non-demo boundary." | |
| first_non_demo = records[boundary_idx] | |
| assert not first_non_demo["is_demo"], ( | |
| f"{env_id}: boundary row should be non-demo, got demo at timestep " | |
| f"{first_non_demo['timestep']}." | |
| ) | |
| last_demo_waypoint = _last_finite_demo_waypoint(records, boundary_idx) | |
| if first_non_demo["is_finite_waypoint"] and last_demo_waypoint is not None: | |
| assert not np.array_equal( | |
| np.asarray(first_non_demo["waypoint_action"]).flatten(), | |
| last_demo_waypoint, | |
| ), ( | |
| f"{env_id}: boundary non-demo step consumed demo-phase pending waypoint " | |
| f"(timestep {first_non_demo['timestep']})." | |
| ) | |
| non_demo_rows = [row for row in records if not row["is_demo"]] | |
| assert non_demo_rows, f"{env_id}: no non-demo rows." | |
| all_non_demo_unique = _unique_finite_waypoints(non_demo_rows) | |
| assert all_non_demo_unique, f"{env_id}: non-demo phase has no finite waypoints." | |
| if assert_first_segment_midpoint: | |
| first_subgoal = non_demo_rows[0]["subgoal"] | |
| first_segment_rows: list[dict] = [] | |
| for row in non_demo_rows: | |
| if row["subgoal"] != first_subgoal: | |
| break | |
| first_segment_rows.append(row) | |
| first_segment_unique = _unique_finite_waypoints(first_segment_rows) | |
| assert len(first_segment_unique) >= 2, ( | |
| f"{env_id}: first non-demo subgoal segment lost midpoint waypoint; " | |
| f"expected >=2 unique finite waypoints, got {len(first_segment_unique)}." | |
| ) | |