File size: 4,761 Bytes
06c11b0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | 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
@pytest.mark.parametrize(
"env_id,base_seed,assert_first_segment_midpoint",
[
("PatternLock", 15001, False),
("RouteStick", 16000, True),
],
)
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)}."
)
|