src/mcss/oracle_runner.py
"""Reproducible runner for diagnostic-only appearance capacity oracles."""
from __future__ import annotations
import argparse
import hashlib
import json
import math
import os
import time
from collections.abc import Mapping, Sequence
from pathlib import Path
import torch
from torch import Tensor
from mcss.appearance_oracle import AppearanceOracle, OracleField, OracleVariant
from mcss.config import load_config
from mcss.data.collate import collate_scene_examples
from mcss.engine import Trainer
from mcss.measurements import FixedMeasurementRenderer
from mcss.metrics import compute_metrics
from mcss.types import SceneBatch, SceneState, StateAppearance
NATIVE_VARIANTS: tuple[OracleVariant, ...] = (
"native_shared",
"highres_shared",
"native_per_view",
)
TYPED_VARIANTS: tuple[OracleVariant, ...] = ("typed_shared", "typed_per_view")
def main(argv: Sequence[str] | None = None) -> int:
parser = argparse.ArgumentParser(
description="Estimate the RGB capacity ceiling of fixed scene-state appearance fields."
)
parser.add_argument("--config", required=True, type=Path)
parser.add_argument("--checkpoint", required=True, type=Path)
parser.add_argument("--output-dir", required=True, type=Path)
parser.add_argument("--steps", type=int, default=1000)
parser.add_argument("--learning-rate", type=float, default=0.05)
parser.add_argument("--seed", type=int, default=23)
parser.add_argument("--field", choices=("native", "typed_appearance"), default="native")
parser.add_argument("--overwrite", action="store_true")
args = parser.parse_args(argv)
if args.steps < 1:
parser.error("--steps must be positive")
if not math.isfinite(args.learning_rate) or args.learning_rate <= 0:
parser.error("--learning-rate must be positive and finite")
output_dir: Path = args.output_dir
report_path = output_dir / "oracle.json"
if report_path.exists() and not args.overwrite:
raise FileExistsError(f"refusing to overwrite existing oracle report: {report_path}")
output_dir.mkdir(parents=True, exist_ok=True)
log_path = output_dir / "oracle.jsonl"
log_path.write_text("", encoding="utf-8")
torch.manual_seed(args.seed)
config = load_config(args.config)
trainer = Trainer(config)
trainer.load_checkpoint(args.checkpoint, load_optimizer=False)
batch = collate_scene_examples([trainer.dataset[0]]).to(trainer.device)
trainer.model.eval()
with torch.no_grad():
model_output = trainer.model(
batch.context_rgb,
batch.context_cameras,
batch.target_cameras,
batch.bounds,
)
field: OracleField = args.field
reference_state = _strip_state(model_output.state, field=field)
metric_targets = _rgb_targets(batch)
baseline = _rgb_metrics(model_output.predictions["rgb"], metric_targets)
renderer = FixedMeasurementRenderer(
n_samples=config.model.n_samples,
ray_chunk_size=config.model.ray_chunk_size,
).to(trainer.device)
if sum(parameter.numel() for parameter in renderer.parameters()) != 0:
raise RuntimeError("appearance oracle requires a parameter-free renderer")
del model_output
trainer.model.to("cpu")
del trainer
if reference_state.density_logits.device.type == "cuda":
torch.cuda.empty_cache()
variants_to_run = NATIVE_VARIANTS if field == "native" else TYPED_VARIANTS
variants: dict[str, dict[str, object]] = {}
for variant in variants_to_run:
result = _run_variant(
variant,
reference_state,
renderer,
batch,
field=field,
steps=args.steps,
learning_rate=args.learning_rate,
)
variants[variant] = result
_append_jsonl(log_path, {"field": field, "variant": variant, **result})
report = {
"schema_version": "mcss.appearance_oracle.v1",
"diagnostic_only": True,
"target_labels_used": True,
"field": field,
"native_state_shape": list(reference_state.spatial_shape),
"optimized_field_shape": list(reference_state.appearance_shape),
"seed": args.seed,
"source": {
"config_path": str(args.config.resolve()),
"config_sha256": _sha256(args.config),
"checkpoint_path": str(args.checkpoint.resolve()),
"checkpoint_sha256": _sha256(args.checkpoint),
"scene_ids": list(batch.scene_ids),
},
"baseline": baseline,
"selection": _selection_summary(variants, field=field),
"variants": variants,
}
_atomic_json(report_path, report)
print(json.dumps(report, sort_keys=True))
return 0
def _run_variant(
variant: OracleVariant,
reference_state: SceneState,
renderer: FixedMeasurementRenderer,
batch: SceneBatch,
*,
field: OracleField,
steps: int,
learning_rate: float,
) -> dict[str, object]:
oracle = AppearanceOracle(
reference_state,
target_views=batch.target_rgb.shape[1],
variant=variant,
field=field,
).to(batch.target_rgb.device)
optimizer = torch.optim.Adam(oracle.parameters(), lr=learning_rate)
targets = _rgb_targets(batch)
started = time.perf_counter()
trace: list[dict[str, float | int]] = []
with torch.no_grad():
initial_rgb = oracle.render_rgb(renderer, batch.target_cameras)
initial = _rgb_metrics(initial_rgb, targets)
best_loss = float(_masked_rgb_mse(initial_rgb, batch.target_rgb, batch.target_support))
best_logits = oracle.color_logits.detach().clone()
log_every = max(1, steps // 20)
for step in range(1, steps + 1):
optimizer.zero_grad(set_to_none=True)
rgb = oracle.render_rgb(renderer, batch.target_cameras)
loss = _masked_rgb_mse(rgb, batch.target_rgb, batch.target_support)
if not torch.isfinite(loss):
raise RuntimeError(f"{variant} produced non-finite RGB loss at step {step}")
loss_value = float(loss.detach())
if loss_value < best_loss:
best_loss = loss_value
best_logits = oracle.color_logits.detach().clone()
loss.backward()
gradient = oracle.color_logits.grad
if gradient is None or not torch.isfinite(gradient).all():
raise RuntimeError(f"{variant} produced invalid color gradients at step {step}")
optimizer.step()
if step == 1 or step == steps or step % log_every == 0:
trace.append({"step": step, "rgb/mse_before_update": loss_value})
with torch.no_grad():
final_rgb = oracle.render_rgb(renderer, batch.target_cameras)
final_loss = float(_masked_rgb_mse(final_rgb, batch.target_rgb, batch.target_support))
final = _rgb_metrics(final_rgb, targets)
if final_loss < best_loss:
best_loss = final_loss
best_logits = oracle.color_logits.detach().clone()
oracle.color_logits.copy_(best_logits)
best_rgb = oracle.render_rgb(renderer, batch.target_cameras)
best = _rgb_metrics(best_rgb, targets)
return {
"steps": steps,
"learning_rate": learning_rate,
"parameter_count": oracle.parameter_count,
"spatial_shape": list(oracle.spatial_shape),
"runtime_seconds": time.perf_counter() - started,
"initial": initial,
"best": best,
"final": final,
"tail_improvement_fraction": _tail_improvement_fraction(trace),
"trace": trace,
}
def _rgb_targets(batch: SceneBatch) -> dict[str, Tensor]:
targets = {"rgb": batch.target_rgb}
if batch.target_support is not None:
targets["support"] = batch.target_support
return targets
def _rgb_metrics(prediction: Tensor, targets: Mapping[str, Tensor]) -> dict[str, float]:
metrics = compute_metrics({"rgb": prediction}, targets)
return {name: value for name, value in metrics.items() if name.startswith("rgb/")}
def _masked_rgb_mse(prediction: Tensor, target: Tensor, support: Tensor | None) -> Tensor:
valid = torch.isfinite(target).all(dim=-3, keepdim=True)
if support is not None:
valid = valid & torch.isfinite(support) & (support > 0.5)
mask = valid.expand_as(target).to(prediction.dtype)
return ((prediction - target).square() * mask).sum() / mask.sum().clamp_min(1.0)
def _strip_state(state: SceneState, *, field: OracleField) -> SceneState:
appearance = None
if field == "typed_appearance":
if state.appearance is None:
raise ValueError("typed_appearance oracle requires SceneState.StateAppearance")
source = state.appearance
appearance = StateAppearance(
source.confidence.detach().clone(),
source.unknown_probability.detach().clone(),
source.completion_gate.detach().clone(),
source.provenance.detach().clone(),
source.base_color.detach().clone(),
source.color_logit_residual.detach().clone(),
)
return SceneState(
state.density_logits.detach().clone(),
state.color.detach().clone(),
state.log_variance.detach().clone(),
state.bounds.detach().clone(),
appearance=appearance,
)
def _tail_improvement_fraction(trace: list[dict[str, float | int]]) -> float:
if len(trace) < 2:
return 0.0
tail_start = max(0, math.floor(0.8 * (len(trace) - 1)))
start = float(trace[tail_start]["rgb/mse_before_update"])
end = min(float(record["rgb/mse_before_update"]) for record in trace[tail_start:])
return max(0.0, (start - end) / max(start, 1e-12))
def _selection_summary(
variants: Mapping[str, Mapping[str, object]],
*,
field: OracleField,
) -> dict[str, object]:
if field == "typed_appearance":
shared = float(_best_metrics(variants["typed_shared"])["rgb/psnr"])
per_view = float(_best_metrics(variants["typed_per_view"])["rgb/psnr"])
per_view_gain = per_view - shared
return {
"provisional_branch": (
"optimization_only" if per_view_gain < 0.25 else "directional_basis"
),
"per_view_gain_db_over_shared": per_view_gain,
}
native = float(_best_metrics(variants["native_shared"])["rgb/psnr"])
highres = float(_best_metrics(variants["highres_shared"])["rgb/psnr"])
per_view = float(_best_metrics(variants["native_per_view"])["rgb/psnr"])
highres_gain = highres - native
per_view_gain = per_view - native
if max(highres_gain, per_view_gain) < 0.25:
branch = "optimization_only"
elif highres_gain + 0.5 >= per_view_gain:
branch = "highres_shared"
else:
branch = "directional_basis"
return {
"provisional_branch": branch,
"highres_gain_db_over_native": highres_gain,
"per_view_gain_db_over_native": per_view_gain,
"tie_margin_db": 0.5,
}
def _best_metrics(result: Mapping[str, object]) -> Mapping[str, object]:
best = result["best"]
if not isinstance(best, Mapping):
raise TypeError("oracle best metrics must be a mapping")
return best
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _append_jsonl(path: Path, payload: Mapping[str, object]) -> None:
with path.open("a", encoding="utf-8") as handle:
handle.write(json.dumps(payload, sort_keys=True) + "\n")
def _atomic_json(path: Path, payload: Mapping[str, object]) -> None:
temporary = path.with_name(path.name + ".part")
temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
os.replace(temporary, path)