src/mcss/mechanism_pilot/plane_sweep_carrier.py
"""V4: minimal non-learned plane-sweep geometry cue for the dense 16^3 learned carrier.
For every context view taken as reference, raw-RGB photometric variance against the other
context views is computed on 32 fronto-parallel inverse-depth planes (z-depth) spanning the
frozen TRAIN prior [near, far]; all views are area-downsampled to 32x40 and compared at that
resolution with correspondingly scaled intrinsics. A per-pixel softmax over depth
(temperature exp(log_tau), the only learned sweep parameter) gives a depth distribution. Every
voxel reads, per reference, a surface likelihood (D * P at its depth) and a free-space
likelihood (probability that the surface lies beyond it), averaged over valid references,
plus the valid-reference fraction. The three statistics enter the fusion through a
zero-initialized bias-free projection, so a sweep carrier starts exactly equal to the matched
dense carrier. Only context RGB, context cameras and the TRAIN prior are read.
"""
from __future__ import annotations
import hashlib
import math
from dataclasses import asdict
from pathlib import Path
import torch
import torch.nn.functional as functional
from torch import nn
from mcss.dynamic.carrier import DynamicSceneCarrier, _CarrierComputation
from mcss.dynamic.config import CarrierConfig
from mcss.dynamic.lifting import lift_cached_observations
from mcss.dynamic.types import WriteTrace, hash_value
from mcss.geometry import project_world
from mcss.mechanism_pilot import geometry_carrier
from mcss.mechanism_pilot.dense_evidence_carrier import build_state
from mcss.mechanism_pilot.small_training import sha
from mcss.types import Cameras
EXPERIMENT = "EXP-3D-PLANE-SWEEP-CARRIER-V4"
GRID = (16, 16, 16)
TOKEN_COUNT = 4096
PLANES = 32
REF_SIZE = (32, 40)
LOG_TAU0 = math.log(0.01)
VARIANT_SPECS = {
"C0": {"label": "DENSE_SURFACE", "sweep": False, "loss": "C1"},
"C1": {"label": "SWEEP_SURFACE", "sweep": True, "loss": "C1"},
"C2": {"label": "SWEEP_NO_SURFACE", "sweep": True, "loss": "C0"},
}
VARIANTS = tuple(VARIANT_SPECS)
PRIMARY_PAIR = ("C0", "C1")
SCHEMA = "mcss.plane_sweep_carrier.v4"
__all__ = [
"EXPERIMENT",
"PRIMARY_PAIR",
"VARIANTS",
"VARIANT_SPECS",
"PlaneSweepCarrier",
"build_state",
"load_checkpoint",
"make_carrier",
"parameter_hash",
"plane_sweep_statistics",
"save_checkpoint",
"training_loss",
]
def _scaled_camera(camera, size):
"""Intrinsics for an exact area downsample; integer pixel centers are kept consistent."""
height, width = camera.image_size
sy, sx = size[0] / height, size[1] / width
intrinsics = camera.intrinsics.clone()
intrinsics[0, 0] *= sx
intrinsics[1, 1] *= sy
intrinsics[0, 2] = sx * (intrinsics[0, 2] + 0.5) - 0.5
intrinsics[1, 2] = sy * (intrinsics[1, 2] + 0.5) - 0.5
return Cameras(intrinsics, camera.c2w, tuple(size))
def _sample_rgb(rgb, pixels):
"""Bilinear RGB at pixel coordinates [..., 2] (align_corners pixel-center convention)."""
height, width = rgb.shape[-2:]
x = pixels[..., 0] / max(width - 1, 1) * 2.0 - 1.0
y = pixels[..., 1] / max(height - 1, 1) * 2.0 - 1.0
grid = torch.nan_to_num(torch.stack((x, y), -1), nan=2.0, posinf=2.0, neginf=-2.0)
flat = grid.reshape(1, -1, 1, 2)
sampled = functional.grid_sample(
rgb.unsqueeze(0), flat, mode="bilinear", padding_mode="zeros", align_corners=True
)
return sampled.reshape(3, *pixels.shape[:-1])
def _depth_distribution(observations, reference, near, far, log_tau, planes, size):
"""Per-pixel softmax over depth planes for one reference view; None if no source view."""
ref = observations[reference]
sources = [o for i, o in enumerate(observations) if i != reference]
camera = _scaled_camera(ref.camera, size)
device = ref.rgb.device
with torch.no_grad():
ref_rgb = functional.interpolate(
ref.rgb.to(torch.float32).unsqueeze(0), size=size, mode="area"
)[0]
inverse = torch.linspace(1.0 / near, 1.0 / far, planes, device=device)
depths = 1.0 / inverse
v, u = torch.meshgrid(
torch.arange(size[0], device=device, dtype=torch.float32),
torch.arange(size[1], device=device, dtype=torch.float32),
indexing="ij",
)
pixels = torch.stack((u, v, torch.ones_like(u)), -1)
directions = pixels @ torch.linalg.inv(camera.intrinsics.to(torch.float32)).T
local = directions.unsqueeze(0) * depths.view(-1, 1, 1, 1)
c2w = camera.c2w.to(torch.float32)
world = local @ c2w[:3, :3].T + c2w[:3, 3]
total = ref_rgb.unsqueeze(1).expand(3, planes, *size).clone()
squares = total.square()
count = torch.ones(planes, *size, device=device)
for source in sources:
# Every view is compared at the same area-downsampled resolution (same blur).
cam = _scaled_camera(source.camera, size).to(dtype=torch.float32)
small = functional.interpolate(
source.rgb.to(torch.float32).unsqueeze(0), size=size, mode="area"
)[0]
projected, _, inside = project_world(world.reshape(-1, 3), cam)
colors = _sample_rgb(small, projected)
valid = inside.reshape(planes, *size).to(torch.float32)
colors = colors.reshape(3, planes, *size) * valid
total = total + colors
squares = squares + colors.square()
count = count + valid
mean = total / count
variance = (squares / count - mean.square()).clamp_min(0).mean(0)
valid_plane = count >= 2
valid_pixel = valid_plane.any(0)
logits = (-variance / torch.exp(log_tau)).masked_fill(~valid_plane, float("-inf"))
logits = torch.where(valid_pixel.unsqueeze(0), logits, torch.zeros_like(logits))
probability = torch.softmax(logits, dim=0) * valid_pixel.unsqueeze(0)
cumulative = torch.cumsum(probability, dim=0)
return camera, probability, cumulative, valid_pixel
def plane_sweep_statistics(
observations, points, near, far, log_tau, *, planes=PLANES, size=REF_SIZE
):
"""[N, 3] = (surface likelihood, free-space likelihood, valid-reference fraction)."""
observations = tuple(observations)
result = points.new_zeros(points.shape[0], 3)
if len(observations) < 2:
return result
points32 = points.to(torch.float32)
surface = torch.zeros(points.shape[0], device=points.device)
free = torch.zeros_like(surface)
references = torch.zeros_like(surface)
for reference in range(len(observations)):
camera, probability, cumulative, valid_pixel = _depth_distribution(
observations, reference, near, far, log_tau, planes, size
)
with torch.no_grad():
pixels, depth, inside = project_world(points32, camera.to(dtype=torch.float32))
in_range = inside & (depth >= near) & (depth <= far)
index = (1.0 / depth.clamp_min(1e-6) - 1.0 / near) / (1.0 / far - 1.0 / near)
coords = torch.stack(
(
pixels[:, 0] / max(size[1] - 1, 1) * 2 - 1,
pixels[:, 1] / max(size[0] - 1, 1) * 2 - 1,
index * 2 - 1,
),
-1,
)
coords = torch.nan_to_num(coords, nan=2.0, posinf=2.0, neginf=-2.0)
grid = coords.view(1, -1, 1, 1, 3)
pixel_ok = functional.grid_sample(
valid_pixel.to(torch.float32).view(1, 1, 1, *size).expand(1, 1, 2, *size),
grid,
mode="bilinear",
align_corners=True,
).view(-1)
use = in_range & (pixel_ok >= 0.999)
volume = torch.stack((probability, cumulative)).unsqueeze(0)
sampled = functional.grid_sample(
volume, grid, mode="bilinear", padding_mode="border", align_corners=True
).view(2, -1)
weight = use.to(sampled.dtype)
surface = surface + weight * planes * sampled[0]
free = free + weight * (1.0 - sampled[1])
references = references + weight
denominator = references.clamp_min(1.0)
result = torch.stack(
(surface / denominator, free / denominator, references / len(observations)), -1
)
return result.to(points.dtype)
class PlaneSweepCarrier(DynamicSceneCarrier):
"""Dense carrier plus a zero-initialized plane-sweep bypass into fuse_up (49 parameters)."""
def __init__(self, config: CarrierConfig, near: float, far: float) -> None:
super().__init__(config)
# Created without consuming the RNG, so shared parameters match the dense baseline.
self.sweep_weight = nn.Parameter(torch.zeros(config.expansion_dim, 3))
self.sweep_log_tau = nn.Parameter(torch.tensor(LOG_TAU0))
if not 0 < near < far:
raise ValueError("Frozen TRAIN prior requires 0 < near < far")
self.sweep_near, self.sweep_far = float(near), float(far)
def _compute(self, cache, fast):
self._validate_fast(cache, fast)
lifted = lift_cached_observations(
cache,
self._candidate_points,
self._candidate_normalized_xyz,
feature_dim=self.config.feature_dim,
)
statistics = lifted.observed_feature_statistics.to(dtype=self.fuse_up.weight.dtype)
support_weights = lifted.support_weights.to(dtype=statistics.dtype)
sweep = plane_sweep_statistics(
[entry.observation for entry in cache.entries],
self._candidate_points,
self.sweep_near,
self.sweep_far,
self.sweep_log_tau,
).to(dtype=statistics.dtype)
fuse_input = self.fuse_up(statistics) + functional.linear(sweep, self.sweep_weight)
fuse_activation = functional.gelu(fuse_input)
fused_hidden = functional.linear(
fuse_activation,
self.fuse_down.weight + fast.delta_fuse,
self.fuse_down.bias,
)
fused_hidden = fused_hidden * support_weights.unsqueeze(-1)
refined = self.refinement(self._scatter_candidates(fused_hidden))
complete_input = self._gather_candidates(refined)
complete_activation = functional.gelu(self.complete_up(complete_input))
trace = WriteTrace(
episode_id=cache.episode_id,
cache_revision=cache.revision,
fast_step=fast.step,
fuse_activation=fuse_activation,
complete_activation=complete_activation,
observed_feature_statistics=statistics,
support_weights=support_weights,
candidate_ids=self._candidate_ids.clone(),
branch_id=fast.branch_id,
source_state_id=fast.state_id,
)
return _CarrierComputation(trace, refined)
def carrier_config():
return CarrierConfig(grid_size=GRID, token_count=TOKEN_COUNT)
def make_carrier(variant, seed, device, near, far):
"""Identical per-seed draw of every shared parameter; the sweep bypass starts at zero."""
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
if VARIANT_SPECS[variant]["sweep"]:
carrier = PlaneSweepCarrier(carrier_config(), near, far)
else:
carrier = DynamicSceneCarrier(carrier_config())
return carrier.to(device)
def parameter_hash(carrier):
"""Hash of the parameters shared by every variant (the sweep bypass is excluded)."""
digest = hashlib.sha256()
for name, value in carrier.named_parameters():
if name.startswith("sweep_"):
continue
digest.update(name.encode())
digest.update(value.detach().to("cpu", torch.float32).contiguous().numpy().tobytes())
return digest.hexdigest()
def training_loss(state, local_cameras, rgb, depth, indices, variant):
"""Frozen V2 loss family: C0/C1 use the V2 surface recipe, C2 the V2 no-surface loss."""
return geometry_carrier.training_loss(
state, local_cameras, rgb, depth, indices, VARIANT_SPECS[variant]["loss"]
)
def save_checkpoint(path, carrier, *, variant, seed, step, lock_sha256):
sweep = None
if isinstance(carrier, PlaneSweepCarrier):
sweep = {"near": carrier.sweep_near, "far": carrier.sweep_far}
payload = {
"schema": SCHEMA,
"config": asdict(carrier.config),
"carrier": {k: v.detach().cpu() for k, v in carrier.state_dict().items()},
"sweep": sweep,
"variant": variant,
"seed": seed,
"step": step,
"lock_sha256": lock_sha256,
"test_time_input": "RGB+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("V4 static plane-sweep carrier schema required")
config = CarrierConfig(**payload["config"])
if tuple(config.grid_size) != GRID or config.token_count != TOKEN_COUNT:
raise PermissionError("Only the frozen dense 16-grid carrier")
if payload["sweep"] is None:
model = DynamicSceneCarrier(config)
else:
model = PlaneSweepCarrier(config, payload["sweep"]["near"], payload["sweep"]["far"])
model = model.to(device)
model.load_state_dict(payload["carrier"], strict=True)
return model, payload