diff --git a/tests/base/simbackend_contract_baseline.json b/tests/base/simbackend_contract_baseline.json new file mode 100644 index 000000000..1d6c8a465 --- /dev/null +++ b/tests/base/simbackend_contract_baseline.json @@ -0,0 +1,78 @@ +{ + "abstract": [ + "apply_interval_randomization", + "get_actuator_ctrl_range", + "get_base_ang_vel", + "get_base_lin_vel", + "get_base_pos", + "get_base_quat", + "get_body_ang_vel_b", + "get_body_ang_vel_w", + "get_body_ids", + "get_body_lin_vel_b", + "get_body_lin_vel_w", + "get_body_pos_b", + "get_body_pos_w", + "get_body_quat_b", + "get_body_quat_w", + "get_dof_pos", + "get_dof_vel", + "get_dr_capabilities", + "get_init_qvel", + "get_joint_range", + "get_keyframe_qpos", + "get_sensor_data", + "model", + "num_actuators", + "num_dof_vel", + "num_envs", + "set_state", + "step" + ], + "optional": [ + "apply_body_force", + "apply_body_linear_velocity_delta", + "apply_init_randomization", + "capture_video_frame", + "create_hfield_scanner", + "get_actuator_gains", + "get_body_ipos", + "get_body_mass", + "get_body_subtree_ids", + "get_default_qpos", + "get_dof_armature", + "get_geom_body_ids", + "get_geom_contact_masks", + "get_geom_friction", + "get_geom_id", + "get_geom_names", + "get_geom_size", + "get_gravity", + "get_joint_dof_indices", + "get_joint_dof_pos_indices", + "get_joint_dof_vel_indices", + "get_motion_body_ids", + "get_physics_state", + "get_site_ids", + "get_site_jacobian_w", + "init_renderer", + "render", + "resolve_play_render_plan", + "run_playback" + ], + "concrete": [ + "cleanup_scene_assets", + "copy_body_state_w", + "get_body_id", + "get_body_pose_w", + "get_body_pose_w_rows", + "get_body_state_w", + "get_body_vel_w", + "get_play_capabilities", + "get_playback_model", + "get_sensor_data_batch", + "get_sensor_data_rows", + "materialize", + "set_pre_step_control" + ] +} diff --git a/tests/base/test_simbackend_contract.py b/tests/base/test_simbackend_contract.py new file mode 100644 index 000000000..0f6ca0270 --- /dev/null +++ b/tests/base/test_simbackend_contract.py @@ -0,0 +1,80 @@ +"""SimBackend public-method contract inventory (Phase 6 safety net; locks current contract). + +Pins the *name sets* of the three method categories of ``SimBackend``: +- ``abstract`` : ``@abc.abstractmethod`` (required — subclasses must implement) +- ``optional`` : concrete body that ``raise NotImplementedError`` (subclass may override) +- ``concrete`` : real default implementation + +b.md b4: lock the NAME SETS, not just counts — a "delete one required abstract + mis-promote +one optional" swap leaves the count unchanged but breaks the contract. When Phase 6 splits +SimBackend into 6 Mixin ABCs, the composed class must still classify identically, so the +abstract→optional demotion the plan warns against is caught here. + +Baseline: ``simbackend_contract_baseline.json`` (currently 28 / 29 / 13). Regenerate only +on an intentional contract change (see ``_dump_baseline`` below). +""" + +from __future__ import annotations + +import inspect +import json +from pathlib import Path + +from unilab.base.backend import base + +BASELINE = Path(__file__).resolve().parent / "simbackend_contract_baseline.json" + + +def classify_public_methods(cls: type) -> dict[str, list[str]]: + """Classify every public attribute of ``cls`` as abstract / optional / concrete. + + Uses ``getattr_static`` (never triggers descriptors) + ``__isabstractmethod__`` + + source inspection for ``raise NotImplementedError``. Deterministic, so generator and + test agree. + """ + out: dict[str, list[str]] = {"abstract": [], "optional": [], "concrete": []} + for name in dir(cls): + if name.startswith("_"): + continue + attr = inspect.getattr_static(cls, name) + if getattr(attr, "__isabstractmethod__", False): + out["abstract"].append(name) + continue + target = attr + if isinstance(attr, property): + target = attr.fget + elif isinstance(attr, (staticmethod, classmethod)): + target = attr.__func__ + try: + src = inspect.getsource(target) if target is not None else "" + except (TypeError, OSError): + src = "" + bucket = "optional" if "raise NotImplementedError" in src else "concrete" + out[bucket].append(name) + return {key: sorted(values) for key, values in out.items()} + + +def _dump_baseline() -> None: + """Regenerate the baseline after an intentional contract change. Run manually:: + + uv run python -c "from tests.base.test_simbackend_contract import _dump_baseline; _dump_baseline()" + """ + BASELINE.write_text(json.dumps(classify_public_methods(base.SimBackend), indent=2) + "\n") + + +def test_simbackend_contract_unchanged(): + expected = json.loads(BASELINE.read_text()) + current = classify_public_methods(base.SimBackend) + if current == expected: + return + drift = [] + for category in ("abstract", "optional", "concrete"): + added = sorted(set(current[category]) - set(expected[category])) + removed = sorted(set(expected[category]) - set(current[category])) + if added or removed: + drift.append(f" {category}: added={added} removed={removed}") + raise AssertionError( + "SimBackend public-method contract drifted from baseline:\n" + + "\n".join(drift) + + "\n(If intentional, regenerate via _dump_baseline; do not weaken required abstracts.)" + ) diff --git a/tests/config/test_owner_yaml_contract.py b/tests/config/test_owner_yaml_contract.py new file mode 100644 index 000000000..ba63a3db3 --- /dev/null +++ b/tests/config/test_owner_yaml_contract.py @@ -0,0 +1,99 @@ +"""Owner-YAML config contract (Phase 1 safety net) + CLI route guard. + +Scope note: the full task/backend compose matrix (reward populated, ``training.sim_backend`` +matches the owner path) is already covered dynamically by tests/config/test_config_system.py +(``_supported_task_cases`` -> ``test_supported_task_composes``). This file deliberately does +NOT duplicate that; it adds only the contract checks that file does not cover: +- No orphan top-level ``algorithm:`` in any ppo/appo/offpolicy owner YAML (ADR-0003: + algorithm hyperparams live under ``algo:``; train_appo reads ``cfg.algo``, so a top-level + ``algorithm:`` is silently dropped to defaults). +- CLI rejects ``training.sim_backend`` as a route-defining override (cli.py + RESERVED_OVERRIDE_KEYS), so backend switching must go through ``task=/``. + +The one known live violation (APPO go2/motrix puts ``algorithm:`` at top level) is excluded +from the scan and captured by the strict xfail below — which flips to a hard suite failure +once Phase 1 Step 1.1 moves the block under ``algo:``, forcing the xfail to be removed. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from omegaconf import OmegaConf + +REPO = Path(__file__).resolve().parents[2] + +# The single known live orphan-``algorithm:`` violation, excluded from the scan and pinned +# by the strict xfail below (it flips to a suite failure the moment Phase 1 fixes the YAML). +KNOWN_ORPHAN_REL = "appo/task/go2_joystick_flat/motrix.yaml" + + +def _owner_task_yamls() -> list[Path]: + """Every ppo/appo/offpolicy owner task YAML, minus the one known-orphan exception.""" + paths: list[Path] = [] + for algo_dir in ("ppo", "appo", "offpolicy"): + root = REPO / "conf" / algo_dir / "task" + if not root.is_dir(): + continue + for path in sorted(root.glob("**/*.yaml")): + if path.relative_to(REPO / "conf").as_posix() == KNOWN_ORPHAN_REL: + continue + paths.append(path) + return paths + + +def _compose(algo_dir: str, task: str, overrides: tuple[str, ...] = ()): + from hydra import compose, initialize_config_dir + + with initialize_config_dir(config_dir=str(REPO / "conf" / algo_dir), version_base="1.3"): + return compose("config", overrides=[f"task={task}", *overrides]) + + +@pytest.mark.parametrize( + "path", _owner_task_yamls(), ids=lambda p: p.relative_to(REPO / "conf").as_posix() +) +def test_owner_yaml_has_no_orphan_top_level_algorithm(path: Path): + """Raw-file scan: no owner YAML may carry a top-level ``algorithm:`` (ADR-0003). + + Complements (does not duplicate) test_config_system's compose matrix — that proves the + composed cfg resolves; this proves owner-file hygiene at the source. The one known + violation is excluded (see KNOWN_ORPHAN_REL) and covered by the strict xfail below. + """ + container = OmegaConf.to_container(OmegaConf.load(path), resolve=False) + assert isinstance(container, dict) + assert "algorithm" not in container, ( + f"{path.relative_to(REPO / 'conf')}: orphan top-level 'algorithm:' — algorithm " + "hyperparams belong under 'algo:' (ADR-0003); a top-level block is read by nothing." + ) + + +def test_cli_rejects_sim_backend_passthrough_override(): + """b.md b5: backend switch must go through task=/, not a bare override. + + cli.py puts training.sim_backend in RESERVED_OVERRIDE_KEYS, so build_command must reject + it — bare Hydra compose alone can't prove the CLI guard. + """ + from unilab.cli import build_command + + with pytest.raises(SystemExit): + build_command( + mode="train", + algo="ppo", + task="go2_joystick_flat", + sim="mujoco", + overrides=["training.sim_backend=motrix"], + ) + + +@pytest.mark.xfail( + reason="APPO go2/motrix YAML puts algorithm: at top level (orphan); train_appo reads " + "cfg.algo, so num_learning_epochs falls back to the default 5. Phase 1 Step 1.1 moves " + "the block under algo: — this test then XPASSes. strict=True turns that XPASS into a " + "suite failure, forcing this xfail (and KNOWN_ORPHAN_REL) to be removed once fixed.", + strict=True, +) +def test_appo_go2_motrix_algorithm_under_algo(): + cfg = _compose("appo", "go2_joystick_flat/motrix") + assert OmegaConf.select(cfg, "algorithm") is None # no orphan top-level algorithm + assert cfg.algo.algorithm.num_learning_epochs == 25 # YAML intent (not default 5) diff --git a/tests/envs/locomotion/fixtures/go1_joystick_golden.npz b/tests/envs/locomotion/fixtures/go1_joystick_golden.npz new file mode 100644 index 000000000..65dd83cf6 Binary files /dev/null and b/tests/envs/locomotion/fixtures/go1_joystick_golden.npz differ diff --git a/tests/envs/locomotion/fixtures/go1_rough_golden.npz b/tests/envs/locomotion/fixtures/go1_rough_golden.npz new file mode 100644 index 000000000..7f3f7ebe3 Binary files /dev/null and b/tests/envs/locomotion/fixtures/go1_rough_golden.npz differ diff --git a/tests/envs/locomotion/fixtures/go2_joystick_golden.npz b/tests/envs/locomotion/fixtures/go2_joystick_golden.npz new file mode 100644 index 000000000..12207e02f Binary files /dev/null and b/tests/envs/locomotion/fixtures/go2_joystick_golden.npz differ diff --git a/tests/envs/locomotion/fixtures/go2_rough_golden.npz b/tests/envs/locomotion/fixtures/go2_rough_golden.npz new file mode 100644 index 000000000..9d12209a5 Binary files /dev/null and b/tests/envs/locomotion/fixtures/go2_rough_golden.npz differ diff --git a/tests/envs/locomotion/fixtures/go2w_rough_golden.npz b/tests/envs/locomotion/fixtures/go2w_rough_golden.npz new file mode 100644 index 000000000..bcdb92aae Binary files /dev/null and b/tests/envs/locomotion/fixtures/go2w_rough_golden.npz differ diff --git a/tests/envs/locomotion/fixtures/reset_go1_rough_golden.npz b/tests/envs/locomotion/fixtures/reset_go1_rough_golden.npz new file mode 100644 index 000000000..7df5b5934 Binary files /dev/null and b/tests/envs/locomotion/fixtures/reset_go1_rough_golden.npz differ diff --git a/tests/envs/locomotion/fixtures/reset_go2_rough_golden.npz b/tests/envs/locomotion/fixtures/reset_go2_rough_golden.npz new file mode 100644 index 000000000..a2672eaf8 Binary files /dev/null and b/tests/envs/locomotion/fixtures/reset_go2_rough_golden.npz differ diff --git a/tests/envs/locomotion/fixtures/reset_go2w_rough_golden.npz b/tests/envs/locomotion/fixtures/reset_go2w_rough_golden.npz new file mode 100644 index 000000000..dcf9f5cbc Binary files /dev/null and b/tests/envs/locomotion/fixtures/reset_go2w_rough_golden.npz differ diff --git a/tests/envs/locomotion/test_reset_plan_golden.py b/tests/envs/locomotion/test_reset_plan_golden.py new file mode 100644 index 000000000..93adb0379 --- /dev/null +++ b/tests/envs/locomotion/test_reset_plan_golden.py @@ -0,0 +1,94 @@ +"""Golden characterization for DR-provider ``build_reset_plan`` (Phase 4.5 safety net). + +Pins the *current* ResetPlan output of the Go1/Go2/Go2W rough providers — qpos/qvel +spawn pose, randomization payload (motor gains, friction, mass/COM, gravity, ...), +info_updates key set, and active DR terms — before Phase 4.5 converges their 90%-clone +``build_reset_plan`` methods. Reward golden cannot see any of this. + +Also pins the terrain-curriculum *side effect*: build_reset_plan records each episode's +start pose into ``env._spawn`` (``spawn_episode_start_xyz``/``spawn_has_started``), which +update_on_done later reads. The returned ResetPlan hides this, so without the snapshot a +refactor could drop the recording and the qpos/qvel golden would still pass. Go1/Go2 record +it; Go2W's override does not — the generic numeric compare locks whichever is current. + +Imports ``snapshot_reset_plan`` from the generator (single source of truth: same env +build + same provider call), so replay matches fixtures exactly. Not marked ``slow``. +""" + +from __future__ import annotations + +import functools +import sys +from pathlib import Path + +import numpy as np +import pytest + +_TOOLS = Path(__file__).resolve().parent / "tools" +if str(_TOOLS) not in sys.path: + sys.path.insert(0, str(_TOOLS)) +from generate_reset_plan_golden import PROVIDERS, snapshot_reset_plan # noqa: E402 + +FIXTURE_DIR = Path(__file__).resolve().parent / "fixtures" +_META_KEYS = ("info_keys", "rand_terms") # object string arrays, compared as sets + + +@functools.lru_cache(maxsize=None) +def _snapshot(task: str, module: str, cls: str): + return snapshot_reset_plan(task, module, cls) + + +def _fixture_path(fixture: str) -> Path: + return FIXTURE_DIR / f"reset_{fixture}_golden.npz" + + +def test_reset_plan_fixtures_present(): + """Hard-fail (not silent skip) if any DR fixture is missing from the repo.""" + missing = [ + f"reset_{name}_golden.npz" + for _, name, _, _ in PROVIDERS + if not _fixture_path(name).exists() + ] + assert not missing, ( + "Missing DR ResetPlan fixtures: " + + ", ".join(missing) + + " — run `uv run python tests/envs/locomotion/tools/generate_reset_plan_golden.py`" + ) + + +@pytest.mark.parametrize("task,fixture,module,cls", PROVIDERS) +def test_reset_plan_unchanged(task: str, fixture: str, module: str, cls: str): + path = _fixture_path(fixture) + if not path.exists(): + pytest.skip(f"fixture missing: {path.name} (see test_reset_plan_fixtures_present)") + # Fixtures are produced by the sibling generator and committed in-repo; they contain + # only ndarrays (numeric + unicode str arrays, no pickled objects), so allow_pickle stays off. + golden = np.load(path) + snap = _snapshot(task, module, cls) + + # Whole-shape lock: a dropped/added qpos field, info entry, or DR term changes the key set. + assert set(snap) == set(golden.files), ( + f"{fixture}: ResetPlan key set changed: {set(snap) ^ set(golden.files)}" + ) + + # info_updates keys + active DR terms — exact set match. + for meta in _META_KEYS: + if meta in golden.files: + assert sorted(snap[meta].tolist()) == sorted(golden[meta].tolist()), ( + f"{fixture}: {meta} changed" + ) + + # All numeric arrays fully deterministic now: _build_env forces terrain_curriculum.seed, + # so env._spawn picks fixed tiles and absolute spawn (qpos[:,0:3]) is reproducible too + # (previously excluded — the unseeded default_rng(None) made tile choice random per process). + # This now locks the full spawn pose (qpos incl. xy/z), qvel, info_updates, and DR payload. + for key in golden.files: + if key in _META_KEYS: + continue + np.testing.assert_allclose( + np.asarray(snap[key], dtype=np.float64), + golden[key], + rtol=1e-6, + atol=1e-7, + err_msg=f"{fixture} ResetPlan mismatch: {key}", + ) diff --git a/tests/envs/locomotion/test_reward_dispatch.py b/tests/envs/locomotion/test_reward_dispatch.py new file mode 100644 index 000000000..a00efe6a9 --- /dev/null +++ b/tests/envs/locomotion/test_reward_dispatch.py @@ -0,0 +1,168 @@ +"""Edge-case tests for ``run_reward_dispatch`` (Phase 4.3 safety net). + +Total-reward golden cannot isolate the dispatch branches; these pin them directly +against the current implementation (rewards.py): +- ``scale == 0`` / missing-key skip (line 355) +- ``scale * fn(ctx) * ctrl_dt`` reduction (lines 357-360, 366) +- ``only_positive`` clamp before ``ctrl_dt`` (line 364) +- log cadence: components written only when ``step % log_every_n_steps == 0`` (line 347) + +Pure-function — no simulator, no Hydra. Not marked ``slow``. +""" + +from __future__ import annotations + +import numpy as np + +from unilab.envs.locomotion.common.rewards import RewardContext, run_reward_dispatch + + +def _ctx(num_envs: int = 2) -> RewardContext: + """Minimal RewardContext: info/linvel/gyro/dof_pos are required (no defaults).""" + return RewardContext( + info={}, + num_envs=num_envs, + linvel=np.zeros((num_envs, 3), dtype=np.float32), + gyro=np.zeros((num_envs, 3), dtype=np.float32), + dof_pos=np.zeros((num_envs, 12), dtype=np.float32), + ) + + +def _ones(ctx: RewardContext) -> np.ndarray: + return np.ones(ctx.num_envs) + + +class _SpyFn: + """Reward fn that counts calls — proves the dispatch short-circuit, not just the output.""" + + def __init__(self) -> None: + self.calls = 0 + + def __call__(self, ctx: RewardContext) -> np.ndarray: + self.calls += 1 + return np.ones(ctx.num_envs) + + +def test_dispatch_skips_zero_scale_and_missing_key(): + out = run_reward_dispatch( + scales={"a": 0.0, "b": 1.0}, # a: scale 0 -> skip; b: not in fns -> skip + fns={"a": _ones}, + ctx=_ctx(2), + info={"steps": np.zeros(2, dtype=np.uint32)}, + enable_log=True, + ctrl_dt=0.02, + ) + np.testing.assert_array_equal(out, np.zeros(2)) + + +def test_dispatch_zero_scale_does_not_call_fn(): + """rewards.py:355 short-circuits on ``scale == 0`` *before* calling fn — a perf contract. + + Output alone can't prove this (``fn`` * 0 is also 0); the spy proves fn was never run. + """ + spy = _SpyFn() + run_reward_dispatch( + scales={"a": 0.0}, + fns={"a": spy}, + ctx=_ctx(2), + info={"steps": np.zeros(2, dtype=np.uint32)}, + enable_log=True, + ctrl_dt=0.02, + ) + assert spy.calls == 0 # scale==0 -> `continue` before fns[name](ctx) + + +def test_dispatch_applies_scale_and_ctrl_dt(): + out = run_reward_dispatch( + scales={"a": 2.0}, + fns={"a": _ones}, + ctx=_ctx(2), + info={"steps": np.zeros(2, dtype=np.uint32)}, + enable_log=False, + ctrl_dt=0.5, + ) + np.testing.assert_allclose(out, np.full(2, 1.0)) # 1 * 2 * 0.5 + + +def test_dispatch_only_positive_clamps(): + out = run_reward_dispatch( + scales={"a": -1.0}, + fns={"a": _ones}, + ctx=_ctx(2), + info={"steps": np.zeros(2, dtype=np.uint32)}, + enable_log=False, + ctrl_dt=1.0, + only_positive=True, + ) + np.testing.assert_array_equal(out, np.zeros(2)) # max(reward, 0) before ctrl_dt + + +def test_dispatch_logs_on_cadence_step(): + info = {"steps": np.zeros(2, dtype=np.uint32)} # step 0 % 4 == 0 -> logs + run_reward_dispatch( + scales={"a": 1.0}, + fns={"a": _ones}, + ctx=_ctx(2), + info=info, + enable_log=True, + ctrl_dt=1.0, + log_every_n_steps=4, + ) + assert "reward/a" in info["log"] + + +def test_dispatch_log_value_is_pre_ctrl_dt(): + """log stores ``mean(scale*fn)`` BEFORE ctrl_dt (rewards.py:361); the return is *after* + (rewards.py:366). scale=2, fn=1, ctrl_dt=0.5 -> component logged 2.0 but reward 1.0. + + Locks the semantic that dashboard component values are pre-ctrl_dt and intentionally + differ from the summed reward — a refactor moving the log past the ctrl_dt scale breaks it. + """ + info = {"steps": np.zeros(2, dtype=np.uint32)} # step 0 -> logs + out = run_reward_dispatch( + scales={"a": 2.0}, + fns={"a": _ones}, + ctx=_ctx(2), + info=info, + enable_log=True, + ctrl_dt=0.5, + ) + np.testing.assert_allclose(out, np.full(2, 1.0)) # 1 * 2 * 0.5 (post-ctrl_dt) + assert info["log"]["reward/a"] == 2.0 # 1 * 2 (pre-ctrl_dt mean) + + +def test_dispatch_skips_log_off_cadence_step(): + info = {"steps": np.ones(2, dtype=np.uint32)} # step 1 % 4 != 0 -> no fresh component + run_reward_dispatch( + scales={"a": 1.0}, + fns={"a": _ones}, + ctx=_ctx(2), + info=info, + enable_log=True, + ctrl_dt=1.0, + log_every_n_steps=4, + ) + assert "reward/a" not in info.get("log", {}) + + +def test_dispatch_off_cadence_preserves_prior_log(): + """Off-cadence keeps the prior ``info["log"]`` untouched (rewards.py:352 carry-forward). + + ``run_env_trajectory`` relies on this: components written on a cadence step persist + across the off-cadence steps in between. A refactor that cleared the log every step + (or appended fresh components off-cadence) would break that carry-forward. + """ + info = { + "steps": np.ones(2, dtype=np.uint32), # step 1 % 4 != 0 -> off-cadence + "log": {"reward/old": 5.0}, # component from a prior cadence step + } + run_reward_dispatch( + scales={"a": 1.0}, + fns={"a": _ones}, + ctx=_ctx(2), + info=info, + enable_log=True, + ctrl_dt=1.0, + log_every_n_steps=4, + ) + assert info["log"] == {"reward/old": 5.0} # preserved verbatim, no fresh "reward/a" diff --git a/tests/envs/locomotion/test_reward_golden.py b/tests/envs/locomotion/test_reward_golden.py new file mode 100644 index 000000000..44313b947 --- /dev/null +++ b/tests/envs/locomotion/test_reward_golden.py @@ -0,0 +1,123 @@ +"""Golden regression tests for locomotion reward numerical equivalence. + +These pin reward values *before* Phase 4 refactoring. Any change to reward values +(even 1e-7) fails here, so reward pure-function extraction / joystick base class / +DR provider convergence cannot silently drift algorithm behavior. + +Fixtures live in ``fixtures/*_golden.npz`` and are produced by +``tools/generate_reward_golden.py``. This test imports the generator's +``run_env_trajectory`` so the replay matches fixture generation exactly (same seed, +same non-zero action schedule, same env construction) — single source of truth, no drift. + +Not marked ``slow`` — this is the safety net, so it must run inside ``make test-all`` +(4 envs x 10 steps is light enough). Fixtures + tool + this test ship in one PR. +""" + +from __future__ import annotations + +import functools +import sys +from pathlib import Path + +import numpy as np +import pytest + +# Import the generator's trajectory runner (single source of truth for env build + +# action schedule). Insert the sibling tools/ dir rather than relying on a package chain. +_TOOLS = Path(__file__).resolve().parent / "tools" +if str(_TOOLS) not in sys.path: + sys.path.insert(0, str(_TOOLS)) +from generate_reward_golden import run_env_trajectory # noqa: E402 + +FIXTURE_DIR = Path(__file__).resolve().parent / "fixtures" + +# (task override, fixture basename) — mirrors MUJOCO_TASKS in the generator. +TASKS: list[tuple[str, str]] = [ + ("go2_joystick_rough/mujoco", "go2_rough"), + ("go1_joystick_rough/mujoco", "go1_rough"), + ("go2w_joystick_rough/mujoco", "go2w_rough"), + ("go2_joystick_flat/mujoco", "go2_joystick"), + ("go1_joystick_flat/mujoco", "go1_joystick"), +] + + +@functools.lru_cache(maxsize=None) +def _trajectory(task: str): + """Run + cache one env trajectory per task (avoids rebuilding MuJoCo env per test).""" + return run_env_trajectory(task) + + +def _fixture_path(fixture: str) -> Path: + return FIXTURE_DIR / f"{fixture}_golden.npz" + + +def test_golden_fixtures_present(): + """Hard-fail (not silent skip) if any expected fixture is missing from the repo. + + The golden tests below skip when a fixture is absent so a partial checkout gives a + clear reason rather than a confusing error; this test makes 'fixtures must be + committed with the test' explicit so the safety net cannot silently degrade to no-op. + """ + missing = [f"{name}_golden.npz" for _, name in TASKS if not _fixture_path(name).exists()] + assert not missing, ( + "Missing golden fixtures: " + + ", ".join(missing) + + " — run `uv run python tests/envs/locomotion/tools/generate_reward_golden.py`" + + " and commit fixtures alongside this test." + ) + + +@pytest.mark.parametrize("task,fixture", TASKS) +def test_reward_trajectory(task: str, fixture: str): + """Full step0..9 reward trajectory matches golden (not just the last step).""" + path = _fixture_path(fixture) + if not path.exists(): + pytest.skip(f"fixture missing: {path.name} (see test_golden_fixtures_present)") + golden = np.load(path) + rewards, _ = _trajectory(task) + n = sum(1 for k in golden.files if k.startswith("reward_step")) + assert len(rewards) == n, f"{fixture}: trajectory length {len(rewards)} != golden {n}" + for i, reward in enumerate(rewards): + np.testing.assert_allclose( + reward, + golden[f"reward_step{i}"], + rtol=1e-6, + atol=1e-7, + err_msg=f"{fixture} reward mismatch at step {i}", + ) + + +@pytest.mark.parametrize("task,fixture", TASKS) +def test_reward_components(task: str, fixture: str): + """Per-component reward log: key set unchanged + each component value matches. + + Catches refactor errors that total reward can hide — a dropped reward key, a renamed + log entry, or one term's scale/gating wrong but offset by another in the sum. + """ + path = _fixture_path(fixture) + if not path.exists(): + pytest.skip(f"fixture missing: {path.name} (see test_golden_fixtures_present)") + golden = np.load(path) + _, final_log = _trajectory(task) + + expected_keys = {k[len("comp_") :] for k in golden.files if k.startswith("comp_reward/")} + assert expected_keys, f"{fixture}: no reward components captured in fixture" + # Exact set (a.md/b.md): catch ADDED reward keys too, not just dropped — a newly logged + # reward/* component is also a behavior change Phase 4 must not introduce silently. + actual_keys = {k for k in final_log if k.startswith("reward/")} + assert actual_keys == expected_keys, ( + f"{fixture}: reward component key set changed: " + f"added={sorted(actual_keys - expected_keys)} dropped={sorted(expected_keys - actual_keys)}" + ) + + for key in golden.files: + if not key.startswith("comp_reward/"): + continue + name = key[len("comp_") :] + np.testing.assert_allclose( + float(final_log[name]), + float(golden[key]), + rtol=1e-5, + atol=1e-6, + err_msg=f"{fixture} component mismatch: {name}", + ) diff --git a/tests/envs/locomotion/tools/generate_reset_plan_golden.py b/tests/envs/locomotion/tools/generate_reset_plan_golden.py new file mode 100644 index 000000000..5c57e0333 --- /dev/null +++ b/tests/envs/locomotion/tools/generate_reset_plan_golden.py @@ -0,0 +1,131 @@ +"""Generate golden ResetPlan fixtures for DR-provider characterization (Phase 4.5). + +Reward golden cannot prove DR ``build_reset_plan`` behavior — qpos/qvel spawn pose, +randomization payload (motor gains, friction, mass/COM, gravity, ...), info_updates, +and subset ``env_ids`` handling are invisible to reward values. Phase 4.5 converges the +Go1/Go2/Go2W rough providers' ``build_reset_plan`` (90% clones); these fixtures pin its +*current* output so the convergence cannot silently change reset behavior. + +Deterministic: ``_build_env`` seeds (``apply_training_seed(42)``) then constructs the env +the same way every run, so the RNG state at ``build_reset_plan`` time is reproducible. + +Run from the repo root:: + + uv run python tests/envs/locomotion/tools/generate_reset_plan_golden.py +""" + +from __future__ import annotations + +import importlib +from pathlib import Path + +import numpy as np + +# Reuse the single source of truth for env construction (same tools/ dir on sys.path). +from generate_reward_golden import REPO, _build_env # noqa: E402 + +# (task, fixture basename, provider module, provider class). The 90%-clone rough providers. +PROVIDERS: list[tuple[str, str, str, str]] = [ + ( + "go2_joystick_rough/mujoco", + "go2_rough", + "unilab.envs.locomotion.go2.rough", + "Go2JoystickRoughDomainRandomizationProvider", + ), + ( + "go1_joystick_rough/mujoco", + "go1_rough", + "unilab.envs.locomotion.go1.rough", + "Go1JoystickRoughDomainRandomizationProvider", + ), + ( + "go2w_joystick_rough/mujoco", + "go2w_rough", + "unilab.envs.locomotion.go2w.rough", + "Go2WJoystickRoughDomainRandomizationProvider", + ), +] + +# Subset reset — exercises env_ids handling, not just full reset. +ENV_IDS = np.array([0, 2], dtype=np.int64) + +# ResetRandomizationPayload fields (dr/types.py:91), all np.ndarray | None. +RAND_FIELDS = [ + "base_mass_delta", + "base_com_offset", + "gravity", + "body_iquat", + "body_inertia", + "body_ipos", + "body_mass", + "dof_armature", + "geom_friction", + "kp", + "kd", +] + + +def snapshot_reset_plan( + task: str, provider_module: str, provider_cls: str +) -> dict[str, np.ndarray]: + """Build env, call provider.build_reset_plan(env, ENV_IDS), flatten ResetPlan to arrays.""" + provider_type = getattr(importlib.import_module(provider_module), provider_cls) + env = _build_env(task, 4) + try: + env.init_state() + plan = provider_type().build_reset_plan(env, ENV_IDS) + save: dict[str, np.ndarray] = { + "env_ids": np.asarray(plan.env_ids), + "qpos": np.asarray(plan.qpos, dtype=np.float64), + "qvel": np.asarray(plan.qvel, dtype=np.float64), + # info_updates key set (lock against dropped/renamed keys). + # Unicode str array (not dtype=object) so fixtures load without allow_pickle. + "info_keys": np.array(sorted(plan.info_updates)), + } + for key, value in plan.info_updates.items(): + if isinstance(value, np.ndarray): + save[f"info_{key}"] = np.asarray(value, dtype=np.float64) + if plan.randomization is not None: + # active DR terms (lock which randomizations fire) + save["rand_terms"] = np.array(sorted(plan.randomization.requested_terms())) + for field in RAND_FIELDS: + value = getattr(plan.randomization, field) + if value is not None: + save[f"rand_{field}"] = np.asarray(value, dtype=np.float64) + # Terrain-curriculum side effect: build_reset_plan mutates env._spawn in-place + # (record_episode_start -> _episode_start_xyz/_has_started, terrain_spawn.py:185-187), + # later read by update_on_done for level progression. The returned ResetPlan does NOT + # show this, so a refactor could drop the call and the qpos/qvel golden would still pass. + # Snapshot the post-call state for ENV_IDS so the side effect is locked too. Go1/Go2 rough + # record it (has_started True); Go2W's override does NOT — both truths are pinned per-fixture. + spawn = getattr(env, "_spawn", None) + start_xyz = getattr(spawn, "_episode_start_xyz", None) + has_started = getattr(spawn, "_has_started", None) + if start_xyz is not None and has_started is not None: + save["spawn_episode_start_xyz"] = np.asarray(start_xyz, dtype=np.float64)[ENV_IDS] + save["spawn_has_started"] = np.asarray(has_started, dtype=np.float64)[ENV_IDS] + finally: + env.close() + return save + + +def generate(task: str, fixture: str, provider_module: str, provider_cls: str) -> None: + save = snapshot_reset_plan(task, provider_module, provider_cls) + out = REPO / "tests" / "envs" / "locomotion" / "fixtures" / f"reset_{fixture}_golden.npz" + out.parent.mkdir(parents=True, exist_ok=True) + np.savez(out, **save) + n_rand = sum(k.startswith("rand_") and k != "rand_terms" for k in save) + print( + f" saved {out.name}: qpos{save['qpos'].shape} qvel{save['qvel'].shape} " + f"{len(save['info_keys'])} info keys + {n_rand} rand fields" + ) + + +def main() -> None: + print("Generating DR ResetPlan golden fixtures (Phase 4.5 targets)...") + for task, fixture, module, cls in PROVIDERS: + generate(task, fixture, module, cls) + + +if __name__ == "__main__": + main() diff --git a/tests/envs/locomotion/tools/generate_reward_golden.py b/tests/envs/locomotion/tools/generate_reward_golden.py new file mode 100644 index 000000000..9421be15d --- /dev/null +++ b/tests/envs/locomotion/tools/generate_reward_golden.py @@ -0,0 +1,160 @@ +"""Generate golden reward fixtures for locomotion regression testing. + +These fixtures pin the *current* reward behavior of the locomotion envs before +any Phase 4 refactoring (reward pure-function extraction, joystick base class, +DR provider convergence). Any refactor that changes reward values — even by +1e-7 — fails the golden tests in ``test_reward_golden.py``. + +Run from the repo root:: + + uv run python tests/envs/locomotion/tools/generate_reward_golden.py + +This regenerates every ``*_golden.npz`` under ``tests/envs/locomotion/fixtures/``. +Fixtures are committed alongside the test + this tool in the same PR, so a clean +``make test-all`` checkout has the safety net available. + +``run_env_trajectory`` / ``deterministic_actions`` are the single source of truth: +``test_reward_golden.py`` imports them so the test replays the env exactly the way +the fixtures were generated (same seed, same action schedule, same construction). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np + +# Repo root: tools -> locomotion -> envs -> tests -> +REPO = Path(__file__).resolve().parents[4] + +# (task override, fixture basename). mujoco only — the actual Phase 4 targets. +MUJOCO_TASKS: list[tuple[str, str]] = [ + ("go2_joystick_rough/mujoco", "go2_rough"), + ("go1_joystick_rough/mujoco", "go1_rough"), + ("go2w_joystick_rough/mujoco", "go2w_rough"), # mixed gated+env-bound table + ("go2_joystick_flat/mujoco", "go2_joystick"), + ("go1_joystick_flat/mujoco", "go1_joystick"), +] +# motrix dual-backend coverage (M1). Generated best-effort: skipped if the motrix +# backend is unavailable in this environment (logged, not silently dropped). +MOTRIX_TASKS: list[tuple[str, str]] = [ + ("go2_joystick_flat/motrix", "go2_joystick_motrix"), +] + + +def deterministic_actions( + num_envs: int, action_dim: int, steps: int, *, seed: int = 0 +) -> np.ndarray: + """Seeded non-zero clipped-random action schedule, shape (steps, num_envs, action_dim). + + Non-zero actions are required to exercise action-dependent rewards + (``action_rate`` reads ``info["current_actions"]`` / ``info["last_actions"]`` — + an all-zero schedule leaves it identically 0 and the golden would not protect it). + """ + rng = np.random.RandomState(seed) + acts = rng.uniform(-1.0, 1.0, size=(steps, num_envs, action_dim)).astype(np.float32) + return np.clip(acts, -1.0, 1.0) + + +def _build_env(task: str, num_envs: int, seed: int = 42): + """Construct an env via the real training path (Hydra compose + BackendAdapter). + + Mirrors scripts/train_*.py so reward_config injection matches actual training. + ``seed`` drives env/reset/command/DR sampling (paired with the same seed for the action + schedule), so the trajectory is reproducible — fixes the seed-API inconsistency where a + non-default ``seed`` only changed actions, not env construction. + """ + from hydra import compose, initialize_config_dir + + from unilab.training.backend_adapter import BackendAdapter + from unilab.training.common import create_env, ensure_registries + from unilab.training.seed import apply_training_seed + + ensure_registries() + # Seed before construction — reset/command/DR all sample random numbers. + apply_training_seed(seed, torch_runtime=True, cuda=False) + + from omegaconf import OmegaConf + + with initialize_config_dir(config_dir=str(REPO / "conf" / "ppo"), version_base="1.3"): + cfg = compose("config", overrides=[f"task={task}"]) + + # Pin terrain spawn determinism on rough (terrain) tasks. terrain_curriculum.seed defaults + # to None (terrain_spawn.py:59) → np.random.default_rng(None) (terrain_spawn.py:98) draws a + # RANDOM spawn tile per process. default_rng is a SEPARATE RNG from np.random.seed(), so + # apply_training_seed() never controls it — the real source of cross-process golden + # flakiness (NOT "MuJoCo float chaos"; MuJoCo is bit-deterministic for identical inputs). + # force_add because some rough YAMLs (e.g. go2w) don't materialize env.terrain_curriculum — + # only the typed cfg default does. Gated on "rough" so flat-task cfg classes that lack the + # field (e.g. Go1JoystickCfg) stay untouched; flat tasks use BaseSpawnManager (deterministic). + if "rough" in task: + OmegaConf.update(cfg, "env.terrain_curriculum.seed", int(seed), force_add=True) + + adapter = BackendAdapter(cfg, root_dir=REPO) + env_cfg_override = adapter.build_task_env_cfg_override() + return create_env(cfg, num_envs=num_envs, env_cfg_override=env_cfg_override) + + +def run_env_trajectory( + task: str, *, num_envs: int = 4, steps: int = 10, seed: int = 42 +) -> tuple[list[np.ndarray], dict[str, Any]]: + """Run a fixed env trajectory; return (per-step rewards, final reward-component log). + + init_state() rather than reset(env_indices) avoids double-init (reset does not set + self._state, so the first step would re-init). ``seed`` drives env construction, the + action schedule, AND the forced terrain_curriculum.seed (see _build_env), so the whole + trajectory is bit-reproducible across processes — verified: two builds give identical + spawn levels and 0.0 reward diff over all 10 steps. 10 steps is safe now that the + spawn-tile randomization is pinned (was the only nondeterminism source). + """ + env = _build_env(task, num_envs, seed) + try: + env.init_state() + action_dim = int(env.action_space.shape[0]) + actions = deterministic_actions(num_envs, action_dim, steps, seed=seed) + rewards: list[np.ndarray] = [] + final_log: dict[str, Any] = {} + for i in range(steps): + state = env.step(actions[i]) + rewards.append(np.asarray(state.reward, dtype=np.float64).copy()) + # info["log"] carries the most recent logged components forward + # (dispatch reuses the prior dict on non-log steps), so the last + # step always exposes the latest reward/* values. + final_log = dict(state.info.get("log", {})) + finally: + env.close() + return rewards, final_log + + +def generate(task: str, fixture: str, *, num_envs: int = 4, steps: int = 10) -> int: + rewards, final_log = run_env_trajectory(task, num_envs=num_envs, steps=steps) + save: dict[str, np.ndarray] = {f"reward_step{i}": r for i, r in enumerate(rewards)} + for key, value in final_log.items(): + if key.startswith("reward/"): + save[f"comp_{key}"] = np.asarray(value, dtype=np.float64) + out = REPO / "tests" / "envs" / "locomotion" / "fixtures" / f"{fixture}_golden.npz" + out.parent.mkdir(parents=True, exist_ok=True) + np.savez(out, **save) + print( + f" saved {out.name}: {len(rewards)} reward steps + " + f"{sum(k.startswith('comp_') for k in save)} components" + ) + return len(save) + + +def main() -> None: + print("Generating mujoco golden fixtures (Phase 4 refactor targets)...") + for task, fixture in MUJOCO_TASKS: + generate(task, fixture) + + # M1 motrix dual-backend golden is deferred to a dedicated lane: motrixsim is an + # optional extra (default `uv sync` may omit it), so a motrix golden cannot live in + # the default not-slow `make test-all` net. It needs its own importorskip-guarded + # test + marker before generation is wired in. MOTRIX_TASKS is kept above for that + # follow-up; not generated by default to avoid committing an unconsumed fixture. + print("Skipping motrix golden (deferred to importorskip lane — see MOTRIX_TASKS).") + + +if __name__ == "__main__": + main()