src/mcss/mechanism_pilot/rgbd_long_train_carrier.py

"""V17: four times longer training of the unchanged V11 C1 recipe (training length, on hit).

PLAN_BEFORE_V17_TRAINING.md fixed V17 before any V17 training. The single variant L is the V11
C1 recipe (width 32, 32^3, CONTEXT_DEPTH bounds, the frozen V2 C1 loss, the V11 extra-view rule,
TRAIN72) trained for 24000 instead of 6000 steps. The learning rate is constant and every random
stream has its own seeded generator, so the first 6000 steps of each seed must equal V11 C1
bit-for-bit (a built-in machinery control). The computation is V14's C0 path unchanged; V14
already resolves the variant's own TRAIN set here instead of V8's variant table.
"""

from __future__ import annotations

from dataclasses import asdict
from math import prod
from pathlib import Path

import torch

from mcss.dynamic.config import CarrierConfig
from mcss.dynamic.types import OnlineObservation, hash_value
from mcss.mechanism_pilot import rgbd_completion_carrier as v11
from mcss.mechanism_pilot import rgbd_evidence_carrier as v5
from mcss.mechanism_pilot import rgbd_resolution_carrier as v8
from mcss.mechanism_pilot.dense_evidence_carrier import _Materializer
from mcss.mechanism_pilot.small_training import sha
from mcss.training.grounded import build_grounded_state

EXPERIMENT = "EXP-3D-RGBD-LONG-TRAIN-V17"
SCHEMA = "mcss.rgbd_long_train_carrier.v17"
GRID = v11.GRID
EXTRA_MAX = v11.EXTRA_MAX
EXTRA_SEED_OFFSET = v11.EXTRA_SEED_OFFSET
VARIANT_SPECS = {
    "L": {"label": "TRAIN72_24K", "hidden": 32, "expansion": 64, "model": "C1", "train": "TRAIN72"},
}
VARIANTS = tuple(VARIANT_SPECS)
PRIMARY_PAIR = ("L",)
V8_LOSS_VARIANT = v11.V8_LOSS_VARIANT
extra_count = v11.extra_count
with_extra_views = v11.with_extra_views
parameter_hash = v5.parameter_hash
load_v11_checkpoint = v11.load_checkpoint

__all__ = [
    "EXPERIMENT",
    "EXTRA_MAX",
    "EXTRA_SEED_OFFSET",
    "GRID",
    "PRIMARY_PAIR",
    "VARIANTS",
    "VARIANT_SPECS",
    "build_state",
    "carrier_config",
    "extra_count",
    "load_checkpoint",
    "load_v11_checkpoint",
    "make_carrier",
    "parameter_hash",
    "save_checkpoint",
    "scene_data",
    "train_records",
    "training_loss",
    "with_extra_views",
]


def _check_config(config):
    widths = {(8, s["hidden"], s["expansion"]) for s in VARIANT_SPECS.values()}
    if (
        tuple(config.grid_size) != (GRID, GRID, GRID)
        or (config.feature_dim, config.hidden_dim, config.expansion_dim) not in widths
        or config.token_count != prod(config.grid_size)
    ):
        raise ValueError("The registered V17 width at 32^3 (one candidate per voxel) is required")


def train_records(manifest, config, variant):
    """Exactly the variant's own frozen TRAIN set (V17 table), sorted by scene id."""
    names = set(config["train_sets"][VARIANT_SPECS[variant]["train"]])
    records = sorted(
        (r for r in manifest["scenes"] if r["split"] == "TRAIN" and r["scene_id"] in names),
        key=lambda r: r["scene_id"],
    )
    if len(records) != len(names) or len(records) < 24:
        raise PermissionError("The variant's frozen TRAIN set is incomplete")
    return records


def carrier_config(variant):
    spec = VARIANT_SPECS[variant]
    return CarrierConfig(
        grid_size=(GRID, GRID, GRID),
        token_count=GRID**3,
        hidden_dim=spec["hidden"],
        expansion_dim=spec["expansion"],
    )


def make_carrier(variant, seed, device, near=None, far=None):
    """The V8 resolution carrier at width 32; for V17 L this is the V11 C1 draw."""
    del near, far
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    return v8.ResolutionRGBDCarrier(carrier_config(variant)).to(device)


def build_state(carrier, context, bounds, episode_id):
    """The V8 build_state under the width guard of the registered V17 width."""
    if not isinstance(context, v5.RGBDContext):
        raise TypeError("V17 states are built from an RGBDContext")
    if not isinstance(carrier, v8.ResolutionRGBDCarrier):
        raise TypeError("V8 resolution carrier required")
    _check_config(carrier.config)
    observations = tuple(context)
    if not observations or any(type(o) is not OnlineObservation for o in observations):
        raise TypeError("Only RGB-camera OnlineObservation values may construct state")
    bounds = torch.as_tensor(bounds, device=carrier._bounds.device, dtype=carrier._bounds.dtype)
    if bounds.shape not in ((2, 3), (1, 2, 3)) or bounds.requires_grad:
        raise ValueError("Frozen nonlearned bounds must have shape [2,3] or [1,2,3]")
    bounds = bounds.reshape(1, 2, 3).detach().clone()
    if not torch.isfinite(bounds).all() or not (bounds[:, 1] > bounds[:, 0]).all():
        raise ValueError("Finite ordered bounds required")
    fractions = (carrier._candidate_normalized_xyz + 1) * 0.5
    points = bounds[:, 0] + fractions * (bounds[:, 1] - bounds[:, 0])
    wrapped = _Materializer(carrier)
    replacements = {"carrier._bounds": bounds, "carrier._candidate_points": points}

    class ContextOnlyAdapter:
        encode = carrier.encode
        initial_fast = carrier.initial_fast

        @staticmethod
        def materialize(cache, fast):
            return torch.func.functional_call(wrapped, replacements, (cache, fast), strict=False)

    carrier._pending_depth = (
        context.frame_ids,
        context.depths.to(device=carrier._bounds.device, dtype=torch.float32),
    )
    try:
        return build_grounded_state(ContextOnlyAdapter(), observations, episode_id)
    finally:
        carrier._pending_depth = None


def training_loss(state, local_cameras, rgb, depth, indices, variant):
    """The frozen V2 C1 loss at 32^3 (surface band at the frozen 16-grid spacing)."""
    if variant not in VARIANT_SPECS:
        raise KeyError(variant)
    return v8.training_loss(state, local_cameras, rgb, depth, indices, V8_LOSS_VARIANT)


def scene_data(variant):
    del variant
    return v8.SCENE_DATA["CONTEXT_DEPTH"]


def save_checkpoint(path, carrier, *, variant, seed, step, lock_sha256):
    payload = {
        "schema": SCHEMA,
        "config": asdict(carrier.config),
        "carrier": {k: v.detach().cpu() for k, v in carrier.state_dict().items()},
        "depth_bypass": True,
        "variant": variant,
        "seed": seed,
        "step": step,
        "lock_sha256": lock_sha256,
        "test_time_input": "RGB+DEPTH+CAMERA",
        "writer_trained": False,
    }
    path = Path(path)
    if path.exists():
        raise FileExistsError(path)
    torch.save(payload, path)
    restored, _ = load_checkpoint(path, "cpu")
    if hash_value(restored.state_dict()) != hash_value(payload["carrier"]):
        raise RuntimeError("Checkpoint restore mismatch")
    return sha(path)


def load_checkpoint(path, device):
    payload = torch.load(path, map_location=device, weights_only=True)
    if payload.get("schema") != SCHEMA or payload.get("writer_trained"):
        raise PermissionError("V17 static RGB-D long-training carrier schema required")
    config = CarrierConfig(**payload["config"])
    _check_config(config)
    if not payload["depth_bypass"] or not any(k.startswith("depth_") for k in payload["carrier"]):
        raise PermissionError("V17 checkpoints always carry the measured-depth bypass")
    model = v8.ResolutionRGBDCarrier(config).to(device)
    model.load_state_dict(payload["carrier"], strict=True)
    return model, payload