src/mcss/mechanism_pilot/static_residual.py
"""Frozen-carrier residual control; training-side labels only, no dynamic writes."""
from __future__ import annotations
import json
import time
from dataclasses import asdict, dataclass
from pathlib import Path
import numpy as np
import torch
from mcss.dynamic.checkpoint import load_dynamic_checkpoint
from mcss.dynamic.types import hash_scene_state, hash_value
from mcss.geometry import generate_rays, transform_cameras
from mcss.measurements import FixedMeasurementRenderer
from mcss.mechanism_pilot.contracts import ResidualReadoutControl
from mcss.mechanism_pilot.small_training import AccessLog, TrainScene, sha, write_json
from mcss.training.grounded import build_grounded_state
from mcss.types import SceneState
SCHEMA = "mcss.static_residual.v1"
ZERO_STATE_DEFINITION = (
"same bounds; density_logits=-100; color/features/log_variance=0; no evidence payload"
)
@dataclass(frozen=True)
class ResidualTrainingConfig:
seed: int = 20260927
steps: int = 3000
learning_rate: float = 1e-3
gradient_clip: float = 1.0
validation_interval: int = 250
renderer_samples: int = 64
ray_chunk_size: int = 2048
def validate_training_manifest(manifest, training_scenes):
"""Identity firewall runs before loading any RGB/depth from manifest."""
rows = manifest["scenes"]
actual = [row["scene_id"] for row in rows]
if len(actual) != len(set(actual)) or set(actual) != set(training_scenes):
raise PermissionError(
"Residual training scene identities must equal checkpoint train allowlist"
)
for row in rows:
if row.get("split") != "train":
raise PermissionError("Residual trainer rejects unseen/dev/test scene records")
if row["roles"]["query"] != [12, 13, 14, 15]:
raise ValueError("Locked residual train/validation query split changed")
context = set(row["roles"]["context_a"]) | set(row["roles"]["context_b"])
if context.intersection([8, 9, 12, 13, 14, 15]):
raise PermissionError("Query entered state construction")
def zero_state_like(state):
"""DIAGNOSTIC ONLY: empty near-transparent volume, identical coordinate bounds."""
return SceneState(
torch.full_like(state.density_logits, -100),
torch.zeros_like(state.color),
torch.zeros_like(state.log_variance),
state.bounds.detach().clone(),
features=None if state.features is None else torch.zeros_like(state.features),
)
def cache_readout(model, state, cameras):
"""Cache detached immutable fixed measurement inputs; no GT accepted."""
state_hash, camera_hash = hash_scene_state(state), hash_value(cameras)
with torch.no_grad():
base = model.renderer(state, cameras, {"rgb", "depth", "visibility"})
features = model.renderer.render_features(state, cameras)
origins, directions = generate_rays(cameras)
rays = torch.cat((origins, directions), dim=-1).permute(0, 1, 4, 2, 3)
combined = torch.cat((features, rays), dim=2).detach()
if state_hash != hash_scene_state(state) or camera_hash != hash_value(cameras):
raise RuntimeError("Cache construction mutated shared inputs")
return {
"base": base,
"combined": combined,
"state_hash": state_hash,
"camera_hash": camera_hash,
}
def cached_forward(model, cache):
combined, base = cache["combined"], cache["base"]
batch, views, channels, height, width = combined.shape
residual = model.head(combined.reshape(batch * views, channels, height, width))
residual = residual.reshape(batch, views, 4, height, width)
result = dict(base)
result["rgb"] = (base["rgb"] + residual[:, :, :3]).clamp(0, 1)
result["depth"] = (base["depth"] + residual[:, :, 3:4]).clamp_min(0)
return result
def supervised_loss(prediction, rgb, depth):
valid = torch.isfinite(depth) & (depth > 0)
if not valid.any():
raise ValueError("No valid train depth")
rgb_mse = (prediction["rgb"] - rgb).square().mean()
absrel = ((prediction["depth"][valid] - depth[valid]).abs() / depth[valid]).mean()
return rgb_mse + absrel, {
"rgb_mse": float(rgb_mse.detach()),
"depth_absrel": float(absrel.detach()),
}
def load_residual_checkpoint(path, *, carrier_sha256, device="cpu"):
payload = torch.load(path, map_location=device, weights_only=True)
if payload.get("schema") != SCHEMA or payload["carrier_sha256"] != carrier_sha256:
raise ValueError("Residual schema/carrier checkpoint mismatch")
config = ResidualTrainingConfig(**payload["config"])
renderer = FixedMeasurementRenderer(
n_samples=config.renderer_samples, ray_chunk_size=config.ray_chunk_size
).to(device)
model = ResidualReadoutControl(payload["feature_dim"], renderer).to(device)
model.head.load_state_dict(payload["head_state_dict"], strict=True)
model.eval().requires_grad_(False)
model.training_status = "TRAINED_FROZEN_TRAIN_SIDE_SELECTED"
return model, payload
def train_residual(manifest_path, carrier_path, output_dir, *, device="cuda", config=None):
config = config or ResidualTrainingConfig()
if min(config.steps, config.validation_interval) <= 0:
raise ValueError("Positive training and validation budgets required")
manifest_path, carrier_path, output = map(Path, (manifest_path, carrier_path, output_dir))
if output.exists() and any(output.iterdir()):
raise FileExistsError("Preserve previous residual results")
output.mkdir(parents=True, exist_ok=True)
manifest = json.loads(manifest_path.read_text())
carrier, _, metadata = load_dynamic_checkpoint(carrier_path, device=device)
training_scenes = metadata["provenance"]["training_scenes"]
validate_training_manifest(manifest, training_scenes)
torch.manual_seed(config.seed)
np.random.seed(config.seed)
carrier.eval().requires_grad_(False)
renderer = FixedMeasurementRenderer(
n_samples=config.renderer_samples, ray_chunk_size=config.ray_chunk_size
).to(device)
model = ResidualReadoutControl(carrier.config.feature_dim, renderer).to(device)
carrier_before = hash_value(carrier.state_dict())
log = AccessLog(output / "label_access.jsonl")
write_json(
output / "lock.json",
{
"config": asdict(config),
"carrier_sha256": sha(carrier_path),
"manifest_sha256": sha(manifest_path),
"training_scenes": training_scenes,
"train_query_ids": [12, 13, 14],
"validation_query_ids": [15],
"validation_caveat": (
"15 was seen by frozen carrier; held out only from residual optimizer"
),
"primary_query_ids_never_loaded": [8, 9],
"zero_state_definition": ZERO_STATE_DEFINITION,
"selection": "lowest equal scene/context mean RGB MSE + depth AbsRel on query 15",
"architecture": "existing ResidualReadoutControl unchanged",
},
)
train, validation, cached_checks = [], [], []
start = time.perf_counter()
for row in manifest["scenes"]:
scene = TrainScene(row, manifest["image_size"], manifest_path.parent, device, log)
for role in ("context_a", "context_b"):
with torch.no_grad():
state, anchor = build_grounded_state(
carrier, scene.observations(role), f"residual:{row['scene_id']}:{role}"
)
for slot in range(4):
frame_id, camera, rgb, depth = scene.query(
slot, "residual_validation" if slot == 3 else "residual_optimizer_training"
)
local = transform_cameras(camera, torch.linalg.inv(anchor))
cache = cache_readout(model, state, local)
with torch.no_grad():
direct, cached = model(state, local), cached_forward(model, cache)
exact = all(torch.equal(direct[k], cached[k]) for k in direct)
if not exact:
raise RuntimeError("Cached readout differs from original forward")
item = {
"scene_id": row["scene_id"],
"context": role,
"query_id": frame_id,
"cache": cache,
"rgb": rgb,
"depth": depth,
"state": state,
"camera": local,
}
(validation if slot == 3 else train).append(item)
cached_checks.append(
{
"scene_id": row["scene_id"],
"context": role,
"query_id": frame_id,
"initial_exact": exact,
}
)
cache_seconds = time.perf_counter() - start
optimizer = torch.optim.Adam(model.head.parameters(), lr=config.learning_rate)
best, best_step = float("inf"), None
curves = []
optimizer_start = time.perf_counter()
def evaluate(step):
nonlocal best, best_step
with torch.no_grad():
losses = [
float(supervised_loss(cached_forward(model, r["cache"]), r["rgb"], r["depth"])[0])
for r in validation
]
training_loss = np.mean(
[
float(
supervised_loss(cached_forward(model, r["cache"]), r["rgb"], r["depth"])[0]
)
for r in train
]
)
val_loss = float(np.mean(losses))
curves.append(
{
"step": step,
"train_mean_loss": float(training_loss),
"validation_mean_loss": val_loss,
"validation_per_item": losses,
}
)
if val_loss < best:
best, best_step = val_loss, step
torch.save(
{
"schema": SCHEMA,
"head_state_dict": model.head.state_dict(),
"feature_dim": carrier.config.feature_dim,
"config": asdict(config),
"carrier_sha256": sha(carrier_path),
"manifest_sha256": sha(manifest_path),
"training_scenes": training_scenes,
"selected_step": step,
"validation_loss": val_loss,
},
output / "best.pt",
)
write_json(output / "validation_curve.json", curves)
evaluate(0)
with (output / "training.jsonl").open("w") as handle:
for step in range(config.steps):
row = train[step % len(train)]
optimizer.zero_grad(set_to_none=True)
loss, terms = supervised_loss(
cached_forward(model, row["cache"]), row["rgb"], row["depth"]
)
loss.backward()
norm = torch.nn.utils.clip_grad_norm_(model.head.parameters(), config.gradient_clip)
if not torch.isfinite(norm) or not torch.isfinite(loss):
raise RuntimeError("Nonfinite residual optimization")
optimizer.step()
handle.write(
json.dumps(
{
"step": step + 1,
"scene_id": row["scene_id"],
"context": row["context"],
"query_id": row["query_id"],
"loss": float(loss.detach()),
"gradient_norm": float(norm),
**terms,
}
)
+ "\n"
)
if (step + 1) % config.validation_interval == 0 or step + 1 == config.steps:
evaluate(step + 1)
optimizer_seconds = time.perf_counter() - optimizer_start
model, _ = load_residual_checkpoint(
output / "best.pt", carrier_sha256=sha(carrier_path), device=device
)
with torch.no_grad():
for row in train + validation:
direct, cached = model(row["state"], row["camera"]), cached_forward(model, row["cache"])
if not all(torch.equal(direct[k], cached[k]) for k in direct):
raise RuntimeError("Trained checkpoint cache/direct mismatch")
if (
hash_value(carrier.state_dict()) != carrier_before
or sha(carrier_path) != metadata["sha256"]
):
raise RuntimeError("Frozen carrier modified")
parameters = sum(p.numel() for p in model.head.parameters())
height, width = manifest["image_size"]
sync = (lambda: torch.cuda.synchronize()) if str(device).startswith("cuda") else lambda: None
row = train[0]
timings = {}
with torch.no_grad():
for name, call in (
("head_only", lambda: cached_forward(model, row["cache"])),
(
"fixed",
lambda: renderer(row["state"], row["camera"], {"rgb", "depth", "visibility"}),
),
("residual_full", lambda: model(row["state"], row["camera"])),
):
for _ in range(3):
call()
sync()
clock = time.perf_counter()
for _ in range(20):
call()
sync()
timings[name + "_seconds_per_query"] = (time.perf_counter() - clock) / 20
recent = curves[-min(3, len(curves)) :]
improvement = (recent[0]["train_mean_loss"] - recent[-1]["train_mean_loss"]) / max(
recent[0]["train_mean_loss"], 1e-12
)
summary = {
"status": "TRAINED_FROZEN",
"selected_step": best_step,
"validation_loss": best,
"steps_completed": config.steps,
"parameters": parameters,
"head_linear_flops_per_query": 2
* height
* width
* ((carrier.config.feature_dim + 6) * 16 + 16 * 4),
"flops_definition": "multiply+add=2; excludes SiLU/bias/clamp/renderer",
"cache_seconds": cache_seconds,
"optimizer_seconds": optimizer_seconds,
**timings,
"training_convergence": "NOT_ESTABLISHED"
if improvement > 0.01
else "RECENT_TRAIN_LOSS_PLATEAU",
"recent_relative_train_loss_improvement": improvement,
"convergence_note": (
"fixed budget; plateau is not proof of global optimum or sufficient capacity"
),
"carrier_unchanged": True,
"cache_exact_initial_and_selected": True,
"checkpoint_sha256": sha(output / "best.pt"),
"dynamic_ttt_run": False,
"zero_state_definition": ZERO_STATE_DEFINITION,
}
write_json(output / "cache_equivalence.json", cached_checks)
write_json(output / "summary.json", summary)
return summary