"""Checkpoint evaluator regression tests for dynamic scene reporting.""" from __future__ import annotations import csv import json from pathlib import Path import numpy as np import pytest import torch from PIL import Image import mcss.evaluation.checkpoint_report as checkpoint_report from mcss.dynamic.carrier import DynamicSceneCarrier from mcss.dynamic.checkpoint import save_dynamic_checkpoint from mcss.dynamic.config import CarrierConfig, WriteConfig from mcss.dynamic.policy import FEATURE_SCHEMA_VERSION, POLICY_FEATURE_DIM, LearnedActionPolicy from mcss.dynamic.policy_checkpoint import save_policy_checkpoint from mcss.dynamic.types import hash_value from mcss.dynamic.write_rule import DirectWriteRule from mcss.training.action_teacher import make_binding from mcss.training.policy_fit import ACTION_TEACHER_SCHEMA, fit_action_policy def test_fit_checkpoint_evaluator_cpu_round_trip(tmp_path: Path) -> None: checkpoint, expected_state_hash = _write_checkpoint(tmp_path) entry, prepared_root = _write_episode(tmp_path, split_id="dev", scene_id="scene_a") binding = make_binding( model_content_hash=expected_state_hash, renderer_samples=2, ray_chunk_size=2048, max_units=1e12, horizon=3, ) binding["teacher_dataset_hash"] = "teacher-dataset" binding["teacher_source_hashes"] = {"teacher.jsonl": "teacher-source"} rows = [] for index, value in enumerate((-0.2, 0.0, 0.2)): features = [0.0] * POLICY_FEATURE_DIM features[0] = value rows.append( { "schema_version": ACTION_TEACHER_SCHEMA, "row_id": f"row-{index}", "episode_id": "train-episode", "scene_id": "train-scene", "split_id": "train", "prefix_step": index, "prefix_frame_id": index, "control_input": {}, "features": features, "actions": ["OFF", "FUSE", "COMPLETE", "ALL"], "feasible_mask": [True, True, True, True], "losses": [1.0, 0.8, 0.6, 0.4], "advantages": [0.0, 0.2, 0.4, 0.6], "costs": [0.0, 1.0, 2.0, 3.0], "horizon_used": 3, "provenance": dict(binding), } ) result = fit_action_policy( rows, hidden_dim=32, epochs=2, holdout_fraction=0.0, seed=5, ) policy_binding = dict(result.binding) policy_binding["teacher_dataset_hash"] = binding["teacher_dataset_hash"] policy_binding["teacher_source_hashes"] = binding["teacher_source_hashes"] policy_checkpoint = tmp_path / "fitted_policy.pt" save_policy_checkpoint( policy_checkpoint, result.policy, binding=policy_binding, provenance=result.checkpoint_provenance(), ) report = checkpoint_report.evaluate_dynamic_checkpoint( [entry], checkpoint, tmp_path / "fit_evaluation", device="cpu", policies=("OFF", "LEARNED"), render_samples=2, ray_chunk_size=2048, policy_checkpoint=policy_checkpoint, prepared_root=prepared_root, ) assert report["status"] == "complete" learned = next(record for record in report["records"] if record["policy"] == "LEARNED") assert learned["trained_policy"] is True assert learned["policy_binding"]["model_content_hash"] == expected_state_hash def test_evaluator_loads_checkpoint_and_writes_scene_policy_records( tmp_path: Path, monkeypatch ) -> None: checkpoint, expected_state_hash = _write_checkpoint(tmp_path) entry, prepared_root = _write_episode(tmp_path, split_id="dev", scene_id="scene_a") calls: list[Path] = [] original_load = checkpoint_report.load_dynamic_checkpoint def tracked_load(path: str | Path, device: str | torch.device = "cpu"): calls.append(Path(path).resolve()) return original_load(path, device) monkeypatch.setattr(checkpoint_report, "load_dynamic_checkpoint", tracked_load) output = tmp_path / "evaluation" report = checkpoint_report.evaluate_dynamic_checkpoint( [entry], checkpoint, output, device="cpu", policies=("OFF", "ALL"), render_samples=2, prepared_root=prepared_root, ) assert calls == [checkpoint.resolve()] assert report["status"] == "complete" assert report["checkpoint"]["sha256"] == _sha256(checkpoint) assert len(report["records"]) == 2 assert {record["policy"] for record in report["records"]} == {"OFF", "ALL"} assert {record["checkpoint_hash"] for record in report["records"]} == {expected_state_hash} assert {record["checkpoint_file_sha256"] for record in report["records"]} == { _sha256(checkpoint) } assert all(record["sealed_state_hash"] for record in report["records"]) assert all("averages" in record["metrics"] for record in report["records"]) assert all( record["actions"] == [item["action"] for item in record["history"]] for record in report["records"] ) assert all( record["budget"]["units_kind"].endswith("not_measured_flops") for record in report["records"] ) persisted = json.loads((output / "report.json").read_text(encoding="utf-8")) assert persisted["records"] == report["records"] with (output / "summary.csv").open(encoding="utf-8", newline="") as handle: rows = list(csv.DictReader(handle)) assert {(row["split_id"], row["policy"]) for row in rows} == { ("dev", "OFF"), ("dev", "ALL"), } def test_max_episodes_does_not_read_unselected_episode_data(tmp_path: Path) -> None: checkpoint, _ = _write_checkpoint(tmp_path) first, prepared_root = _write_episode(tmp_path, split_id="dev", scene_id="scene_a") second, _ = _write_episode(tmp_path, split_id="dev", scene_id="scene_b") second_online = json.loads(Path(second["online_manifest"]).read_text(encoding="utf-8")) Path(second_online["warmup"][0]["rgb"]).write_bytes(b"unselected image must not be decoded") report = checkpoint_report.evaluate_dynamic_checkpoint( [first, second], checkpoint, tmp_path / "subset", device="cpu", policies=("OFF",), render_samples=2, max_episodes=1, prepared_root=prepared_root, ) assert report["actual_subset"] == { "requested_max_episodes": 1, "selected_episodes": 1, "selected_episode_ids": ["dev--scene_a--cam_00"], } assert [record["episode_id"] for record in report["records"]] == [ "dev--scene_a--cam_00" ] def test_evaluator_rejects_train_entry_before_decoding_data(tmp_path: Path) -> None: checkpoint, _ = _write_checkpoint(tmp_path) entry, prepared_root = _write_episode(tmp_path, split_id="train", scene_id="scene_a") online = json.loads(Path(entry["online_manifest"]).read_text(encoding="utf-8")) Path(online["warmup"][0]["rgb"]).write_bytes(b"must not be decoded") with pytest.raises(ValueError, match="only dev entries"): checkpoint_report.evaluate_dynamic_checkpoint( [entry], checkpoint, tmp_path / "train_rejected", device="cpu", policies=("OFF",), render_samples=2, prepared_root=prepared_root, ) def test_evaluator_requires_the_exact_dev_episode_protocol(tmp_path: Path) -> None: checkpoint, _ = _write_checkpoint(tmp_path) entry, prepared_root = _write_episode(tmp_path, split_id="dev", scene_id="scene_a") online_path = Path(entry["online_manifest"]) online = json.loads(online_path.read_text(encoding="utf-8")) online["stream"] = online["stream"][:-1] online["declared_length"] = 11 online_path.write_text(json.dumps(online), encoding="utf-8") entry["stream_steps"] = 7 with pytest.raises(ValueError, match="exactly warmup=4, stream=8, and query=4"): checkpoint_report.evaluate_dynamic_checkpoint( [entry], checkpoint, tmp_path / "bad_protocol", device="cpu", policies=("OFF",), render_samples=2, prepared_root=prepared_root, ) def test_fixed_policy_fallback_to_off_fails_comparison(tmp_path: Path) -> None: checkpoint, _ = _write_checkpoint(tmp_path) entry, prepared_root = _write_episode(tmp_path, split_id="dev", scene_id="scene_a") with pytest.raises(RuntimeError, match="fixed policy FUSE fell back"): checkpoint_report.evaluate_dynamic_checkpoint( [entry], checkpoint, tmp_path / "fallback", device="cpu", policies=("OFF", "FUSE"), render_samples=2, max_units=9704.0, prepared_root=prepared_root, ) def test_learned_policy_requires_checkpoint_and_logs_bound_trajectory(tmp_path: Path) -> None: checkpoint, expected_state_hash = _write_checkpoint(tmp_path) entry, prepared_root = _write_episode(tmp_path, split_id="dev", scene_id="scene_a") policy_checkpoint = _write_policy_checkpoint( tmp_path, model_content_hash=expected_state_hash, render_samples=2, ray_chunk_size=2048, max_units=1e12, ) with pytest.raises(ValueError, match="requires an explicit policy_checkpoint"): checkpoint_report.evaluate_dynamic_checkpoint( [entry], checkpoint, tmp_path / "missing_policy", device="cpu", policies=("OFF", "LEARNED"), render_samples=2, prepared_root=prepared_root, ) report = checkpoint_report.evaluate_dynamic_checkpoint( [entry], checkpoint, tmp_path / "learned", device="cpu", policies=("OFF", "LEARNED"), render_samples=2, ray_chunk_size=2048, policy_checkpoint=policy_checkpoint, prepared_root=prepared_root, ) learned = next(record for record in report["records"] if record["policy"] == "LEARNED") assert learned["trained_policy"] is True assert learned["policy_file_sha256"] == _sha256(policy_checkpoint) assert learned["policy_checkpoint_file_sha256"] == _sha256(policy_checkpoint) assert learned["policy_content_hash"] assert learned["policy_binding"]["model_content_hash"] == expected_state_hash assert len(learned["actions"]) == 8 assert all("control_input" in item for item in learned["history"]) @pytest.mark.parametrize( ("binding_field", "replacement", "message"), [ ("model_content_hash", "wrong-carrier", "binding mismatch"), ( "feature_schema_version", "wrong-schema", "binding mismatch", ), ( "render_protocol", {"renderer": "FixedMeasurementRenderer", "renderer_samples": 4, "ray_chunk_size": 2048}, "binding mismatch", ), ( "budget_protocol", {"max_units": 11.0, "policy_work_units": 1}, "binding mismatch", ), ], ) def test_learned_policy_rejects_mismatched_binding_before_episode_io( tmp_path: Path, binding_field: str, replacement: object, message: str, ) -> None: checkpoint, expected_state_hash = _write_checkpoint(tmp_path) entry, prepared_root = _write_episode(tmp_path, split_id="dev", scene_id="scene_a") policy_checkpoint = _write_policy_checkpoint( tmp_path, model_content_hash=expected_state_hash, render_samples=2, ray_chunk_size=2048, max_units=1e12, binding_field=binding_field, replacement=replacement, ) with pytest.raises(ValueError, match=message): checkpoint_report.evaluate_dynamic_checkpoint( [entry], checkpoint, tmp_path / f"bad_{binding_field}", device="cpu", policies=("OFF", "LEARNED"), render_samples=2, ray_chunk_size=2048, policy_checkpoint=policy_checkpoint, prepared_root=prepared_root, ) def test_learned_evaluation_opens_query_only_after_seal( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: checkpoint, expected_state_hash = _write_checkpoint(tmp_path) entry, prepared_root = _write_episode(tmp_path, split_id="dev", scene_id="scene_a") policy_checkpoint = _write_policy_checkpoint( tmp_path, model_content_hash=expected_state_hash, render_samples=2, ray_chunk_size=2048, max_units=1e12, ) query_path = Path(entry["query_manifest"]) query = json.loads(query_path.read_text(encoding="utf-8")) for frame in query["query"]: Path(frame["rgb"]).write_bytes(b"query RGB must not be opened before seal") Path(frame["depth"]).write_bytes(b"query depth must not be opened before seal") calls: list[tuple[str, tuple[int, ...]]] = [] def fake_evaluate(sealed, *_args, **_kwargs): calls.append(("evaluate", sealed.observed_ids)) return { "per_query": [ { "frame_id": frame_id, "rgb_mse": 0.0, "rgb_psnr": float("inf"), "depth_abs_rel": 0.0, "opacity_mean": 0.0, "coverage": 0.0, } for frame_id in range(12, 16) ], "averages": { "rgb_mse": 0.0, "rgb_psnr": float("inf"), "depth_abs_rel": 0.0, "opacity_mean": 0.0, "coverage": 0.0, }, } monkeypatch.setattr(checkpoint_report, "evaluate_sealed", fake_evaluate) checkpoint_report.evaluate_dynamic_checkpoint( [entry], checkpoint, tmp_path / "sealed_boundary", device="cpu", policies=("OFF", "LEARNED"), render_samples=2, policy_checkpoint=policy_checkpoint, prepared_root=prepared_root, ) assert len(calls) == 2 assert all(observed_ids == tuple(range(12)) for _, observed_ids in calls) def test_scene_aggregation_uses_scene_means_and_matched_off_deltas() -> None: records = [ _record("scene_a", "OFF", 10.0, 20.0, 1.0), _record("scene_a", "FUSE", 7.0, 22.0, 0.8), _record("scene_b", "OFF", 4.0, 24.0, 0.5), _record("scene_b", "FUSE", 1.0, 26.0, 0.3), ] summaries, scene_results, metric_names = checkpoint_report.aggregate_scene_records(records) fuse = next(summary for summary in summaries if summary["policy"] == "FUSE") assert metric_names == ("coverage", "rgb_mse", "rgb_psnr") assert fuse["scene_count"] == 2 assert fuse["metrics_mean"]["rgb_mse"] == 4.0 assert fuse["metric_valid_scene_counts"]["rgb_mse"] == 2 assert fuse["paired_delta_vs_off_mean"]["rgb_mse"] == -3.0 assert {(row["scene_id"], row["policy"]) for row in scene_results} == { ("scene_a", "OFF"), ("scene_a", "FUSE"), ("scene_b", "OFF"), ("scene_b", "FUSE"), } def _write_checkpoint(directory: Path) -> tuple[Path, str]: config = CarrierConfig( feature_dim=2, hidden_dim=2, expansion_dim=2, grid_size=(2, 2, 2), local_bounds_m=((-1.0, -1.0, 0.1), (1.0, 1.0, 2.0)), token_count=4, ) carrier = DynamicSceneCarrier(config) write_rule = DirectWriteRule(config, WriteConfig(learning_rate=0.02, max_update_norm=0.1)) checkpoint = directory / "dynamic.pt" save_dynamic_checkpoint( checkpoint, carrier, write_rule, phase="test", provenance={"split": "train"}, ) state_hash = hash_value( {"carrier": carrier.state_dict(), "write_rule": write_rule.state_dict()} ) return checkpoint, state_hash def _write_policy_checkpoint( directory: Path, *, model_content_hash: str, render_samples: int, ray_chunk_size: int, max_units: float, binding_field: str | None = None, replacement: object = None, ) -> Path: policy = LearnedActionPolicy(hidden_dim=4, seed=11) binding: dict[str, object] = { "model_content_hash": model_content_hash, "feature_schema_version": FEATURE_SCHEMA_VERSION, "render_protocol": { "renderer": "FixedMeasurementRenderer", "renderer_samples": render_samples, "ray_chunk_size": ray_chunk_size, }, "budget_protocol": { "max_units": float(max_units), "policy_work_units": policy.declared_work_units, }, "utility_protocol": { "horizon": 1, "utility": "mean_masked_depth_abs_rel", }, "teacher_dataset_hash": "dataset-test", "teacher_source_hashes": {"rows": "rows-test"}, } serialized_binding_replacement = None if binding_field == "feature_schema_version": serialized_binding_replacement = replacement elif binding_field is not None: binding[binding_field] = replacement path = directory / "policy.pt" save_policy_checkpoint(path, policy, binding=binding, provenance={"split": "train"}) if serialized_binding_replacement is not None: payload = torch.load(path, map_location="cpu", weights_only=True) payload["binding"][binding_field] = serialized_binding_replacement torch.save(payload, path) return path def _write_episode( directory: Path, *, split_id: str, scene_id: str ) -> tuple[dict[str, object], Path]: prepared_root = directory / "prepared" prepared_split = {"train": "train", "dev": "val"}[split_id] scene_root = prepared_root / prepared_split / scene_id rgb_root = scene_root / "rgb" depth_root = scene_root / "depth" rgb_root.mkdir(parents=True) depth_root.mkdir(parents=True) for frame_id in range(16): value = 16 + frame_id * 8 Image.fromarray(np.full((2, 2, 3), value, dtype=np.uint8)).save( rgb_root / f"{frame_id:06d}.png" ) for frame_id in range(12, 16): np.save(depth_root / f"{frame_id:06d}.npy", np.ones((2, 2), dtype=np.float32)) identity = { "episode_id": f"{split_id}--{scene_id}--cam_00", "scene_id": scene_id, "split_id": split_id, "query_vault_id": f"qv-{scene_id}", } online = { "schema_version": "mcss.dynamic.online_episode.v1", **identity, "image_size": [2, 2], "declared_length": 12, "warmup": [ _online_frame(frame_id, rgb_root / f"{frame_id:06d}.png") for frame_id in range(4) ], "stream": [ _online_frame(frame_id, rgb_root / f"{frame_id:06d}.png") for frame_id in range(4, 12) ], } query = { "schema_version": "mcss.dynamic.query_vault.v1", **identity, "image_size": [2, 2], "query": [ _query_frame( frame_id, rgb_root / f"{frame_id:06d}.png", depth_root / f"{frame_id:06d}.npy", ) for frame_id in range(12, 16) ], } artifact_root = directory / "artifacts" artifact_root.mkdir(exist_ok=True) online_path = artifact_root / f"{scene_id}.online.json" query_path = artifact_root / f"{scene_id}.query.json" online_path.write_text(json.dumps(online), encoding="utf-8") query_path.write_text(json.dumps(query), encoding="utf-8") return ( { **identity, "online_manifest": str(online_path), "query_manifest": str(query_path), "stream_steps": 8, "query_count": 4, }, prepared_root, ) def _online_frame(frame_id: int, rgb: Path) -> dict[str, object]: return { "frame_id": frame_id, "intrinsics": [[2.0, 0.0, 0.5], [0.0, 2.0, 0.5], [0.0, 0.0, 1.0]], "c2w": [ [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0], ], "rgb": str(rgb), } def _query_frame(frame_id: int, rgb: Path, depth: Path) -> dict[str, object]: return {**_online_frame(frame_id, rgb), "depth": str(depth)} def _record(scene_id: str, policy: str, rgb_mse: float, rgb_psnr: float, coverage: float) -> dict: return { "episode_id": f"train--{scene_id}--cam_00", "scene_id": scene_id, "split_id": "train", "policy": policy, "metrics": { "averages": { "rgb_mse": rgb_mse, "rgb_psnr": rgb_psnr, "coverage": coverage, } }, } def _sha256(path: Path) -> str: import hashlib return hashlib.sha256(path.read_bytes()).hexdigest()