tests/test_budget_state_optimization.py
"""Frozen budget trajectories preserve previous state semantics and optimizer prefix."""
import json
from dataclasses import replace
import pytest
import torch
from mcss.dynamic.types import hash_scene_state
from mcss.mechanism_pilot.budget_state_optimization import plateau_diagnostic, run_budget_trajectory
from mcss.mechanism_pilot.direct_capacity_contracts import ObservationBatch
from mcss.mechanism_pilot.direct_capacity_optimization import (
DirectOptimizationConfig,
optimize_state,
)
from mcss.types import Cameras
def batch(device="cpu"):
camera = Cameras(
torch.tensor([[[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]]]], device=device),
torch.eye(4, device=device)[None, None],
(2, 2),
)
value = ObservationBatch(
"synthetic",
(0,),
torch.full((1, 1, 3, 2, 2), 0.2, device=device),
camera,
torch.full((1, 1, 1, 2, 2), 2.0, device=device),
"CONTEXT_ONLY_RGBD",
)
return value, torch.tensor([[-2.0, -2.0, 0.1], [2.0, 2.0, 4.0]], device=device)
def test_single_trajectory_prefix_and_context_only_selection(tmp_path):
torch.set_num_threads(1)
observed, bounds = batch()
config = DirectOptimizationConfig(steps=20, batch_rays=8, checkpoint_interval=10)
report = run_budget_trajectory(
observed,
bounds,
tmp_path / "trajectory",
seed=5,
budgets=(10, 20),
config=config,
engineering_fixture=True,
)
old_state, old_summary = optimize_state(
observed, bounds, 16, 64, tmp_path / "old", seed=5, config=replace(config, steps=10)
)
selected = torch.load(
tmp_path / "trajectory/budget_10/CONTEXT_SELECTED/state.pt", weights_only=False
)
assert hash_scene_state(selected) == hash_scene_state(old_state)
assert report["outputs"][1]["selected_step"] == old_summary["selected_step"]
assert [c["step"] for c in report["full_objective_checkpoints"]] == [0, 10, 20]
for budget in (10, 20):
eligible = [c for c in report["full_objective_checkpoints"] if c["step"] <= budget]
best = min(eligible, key=lambda c: (c["context_objective"], c["step"]))
selected_summary = json.loads(
(tmp_path / f"trajectory/budget_{budget}/CONTEXT_SELECTED/summary.json").read_text()
)
assert selected_summary["selected_step"] == best["step"]
assert selected_summary["selected_context_metrics"]["step"] == best["step"]
fixed = torch.load(
tmp_path / f"trajectory/budget_{budget}/FIXED_BUDGET/state.pt", weights_only=False
)
snapshot = torch.load(
tmp_path / f"trajectory/snapshots/step_{budget}.pt", weights_only=False
)
assert hash_scene_state(fixed) == hash_scene_state(snapshot)
assert fixed.density_logits.shape == (1, 1, 16, 16, 16)
assert report["single_trajectory"] and not report["query_selection"]
with pytest.raises(TypeError):
run_budget_trajectory(observed, bounds, tmp_path / "query", seed=5, query_score=0.0)
def test_1000_step_cpu_prefix_matches_previous_kernel_exactly(tmp_path):
torch.set_num_threads(1)
observed, bounds = batch("cpu")
config = DirectOptimizationConfig(steps=1000, batch_rays=8)
report = run_budget_trajectory(
observed,
bounds,
tmp_path / "trajectory",
seed=20260927,
budgets=(1000,),
config=config,
engineering_fixture=True,
)
old_state, old_summary = optimize_state(
observed, bounds, 16, 64, tmp_path / "old", seed=20260927, config=config
)
state = torch.load(
tmp_path / "trajectory/budget_1000/CONTEXT_SELECTED/state.pt", weights_only=False
)
assert hash_scene_state(state) == hash_scene_state(old_state)
assert report["outputs"][1]["selected_step"] == old_summary["selected_step"]
for current, prior in zip(
report["full_objective_checkpoints"], old_summary["full_objective_checkpoints"], strict=True
):
assert all(current[key] == value for key, value in prior.items())
def test_formal_optimizer_and_budget_changes_are_forbidden(tmp_path):
observed, bounds = batch()
with pytest.raises(PermissionError):
run_budget_trajectory(
observed,
bounds,
tmp_path / "lr",
seed=1,
config=DirectOptimizationConfig(steps=10000, learning_rate=0.02),
)
with pytest.raises(ValueError):
run_budget_trajectory(observed, bounds, tmp_path / "budget", seed=1, budgets=(1000, 2000))
with pytest.raises(ValueError):
run_budget_trajectory(
observed, bounds, tmp_path / "mutable", seed=1, budgets=[1000, 3000, 10000]
)
def test_plateau_frozen_last_ten_percent_median_gradient_rule():
checks = [{"step": 900, "context_objective": 1.0}, {"step": 1000, "context_objective": 0.9995}]
trace = [{"step": step, "gradient_norm_before_clip": 1.0} for step in range(900, 1001, 10)]
result = plateau_diagnostic(
checks, 1000, trace=trace, objective_threshold=0.001, gradient_threshold=0.10
)
assert result["plateau"] and result["n_trace_gradients"] == 11
trace[-5:] = [{"step": r["step"], "gradient_norm_before_clip": 0.5} for r in trace[-5:]]
result = plateau_diagnostic(
checks, 1000, trace=trace, objective_threshold=0.001, gradient_threshold=0.10
)
assert not result["plateau"] and not result["gradient_stable"]
assert not plateau_diagnostic(
checks, 1000, trace=trace[:2], objective_threshold=0.001, gradient_threshold=0.10
)["plateau"]
def test_secondary_single_budget_only_explicit_rgb_only(tmp_path):
observed, bounds = batch()
with pytest.raises(PermissionError, match="RGB_ONLY"):
run_budget_trajectory(
observed, bounds, tmp_path / "denied", seed=1, budgets=(10000,), secondary_sanity=True
)
with pytest.raises(ValueError, match="frozen main"):
run_budget_trajectory(observed, bounds, tmp_path / "main", seed=1, budgets=(10000,))
rgb_only = replace(observed, depth=None, supervision="CONTEXT_ONLY_RGB_ONLY")
result = run_budget_trajectory(
rgb_only,
bounds,
tmp_path / "secondary",
seed=1,
budgets=(10,),
config=DirectOptimizationConfig(steps=10, batch_rays=8, checkpoint_interval=10),
engineering_fixture=True,
secondary_sanity=True,
)
assert len(result["outputs"]) == 2 and not result["outputs"][0]["has_depth"]