"""Core B V3: learned fast-weight writes on the frozen V17 L step-24000 carrier (32^3, width 32). Identical to Core B V2 except for the frozen slow carrier: the V17 L checkpoint at step 24000 of each seed (the V11 C1 recipe, width 32, CONTEXT_DEPTH bounds, trained with 3-7 views for 24000 steps), fixed by PLAN_BEFORE_COREB_V3.md. It is a V8 resolution carrier at 32^3 like V9's, so the stream rule, policies, the DirectWriteRule, its offline TRAIN72 training and the loss are those of Core B V2; OFF_CLAMP3 keeps its V1 meaning (the lifting support count clamped at 3). """ from __future__ import annotations from dataclasses import asdict from pathlib import Path import torch import torch.nn.functional as functional from mcss.dynamic.carrier import _CarrierComputation from mcss.dynamic.config import CarrierConfig, WriteConfig from mcss.dynamic.lifting import lift_cached_observations from mcss.dynamic.types import Action, WriteTrace, hash_value from mcss.dynamic.write_rule import DirectWriteRule, commit_write from mcss.mechanism_pilot import rgbd_evidence_carrier as v5 from mcss.mechanism_pilot import rgbd_long_train_carrier as v17 from mcss.mechanism_pilot import rgbd_resolution_carrier as v8 from mcss.mechanism_pilot import rgbd_stream_write as v1 from mcss.mechanism_pilot.small_training import sha EXPERIMENT = "EXP-3D-RGBD-DYNAMIC-WRITE-V3" GRID = 32 STREAM_LENGTH = v1.STREAM_LENGTH TRAINED_VIEW_COUNT = v1.TRAINED_VIEW_COUNT POLICIES = v1.POLICIES CONTROLS = v1.CONTROLS ACTIONS = v1.ACTIONS WRITE_SCHEMA = "mcss.rgbd_stream_write_rule.v3" V8_LOSS_VARIANT = "C0" # every V8 variant registers the frozen V2 C1 surface loss stream_frame_ids = v1.stream_frame_ids StreamRGBDLoader = v1.StreamRGBDLoader rebuild_with = v1.rebuild_with make_write_rule = v1.make_write_rule class StreamResolutionCarrier(v8.ResolutionRGBDCarrier): """The frozen V17 (V8 resolution) carrier; OFF_CLAMP3 clamps the support-count statistic.""" count_clamp = None def _compute(self, cache, fast): self._validate_fast(cache, fast) if self._pending_depth is None: raise PermissionError("RGB-D carrier states are built only through build_state") frame_ids, depths = self._pending_depth if tuple(entry.observation.frame_id for entry in cache.entries) != frame_ids: raise ValueError("Measured depth maps must align with the cached context views") 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) if self.count_clamp is not None: column = 2 * self.config.feature_dim statistics = torch.cat( ( statistics[:, :column], statistics[:, column : column + 1].clamp(max=float(self.count_clamp)), statistics[:, column + 1 :], ), dim=-1, ) support_weights = lifted.support_weights.to(dtype=statistics.dtype) fuse_activation = functional.gelu(self.fuse_up(statistics)) 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) grid = torch.tensor( tuple(self.config.grid_size), device=self._bounds.device, dtype=self._bounds.dtype ) edge = ((self._bounds[:, 1] - self._bounds[:, 0]) / grid).mean() measured = v5.depth_statistics( [entry.observation for entry in cache.entries], depths, self._candidate_points, self.depth_log_sigma, edge.to(device=self._candidate_points.device, dtype=torch.float32), ).to(dtype=statistics.dtype) fused_hidden = fused_hidden + functional.linear(measured, self.depth_weight) 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 load_frozen_carrier(path, device): """The sealed V17 L step-24000 weights in a StreamResolutionCarrier; every parameter frozen.""" model, payload = v17.load_checkpoint(path, device) if tuple(model.config.grid_size) != (GRID, GRID, GRID) or payload["step"] != 24000: raise PermissionError("Core B V3 runs on the 32^3 V17 L step-24000 carrier only") carrier = StreamResolutionCarrier(model.config).to(device) carrier.load_state_dict(model.state_dict(), strict=True) if hash_value(carrier.state_dict()) != hash_value(model.state_dict()): raise RuntimeError("Frozen carrier transfer changed a tensor") carrier.eval() for parameter in carrier.parameters(): parameter.requires_grad_(False) return carrier, payload def unroll(carrier, rule, warmup, stream, bounds, policy, episode_id): """(final state, anchor c2w, final fast weights, episode) of one fixed-policy episode.""" if policy not in POLICIES: raise ValueError(f"Unregistered policy {policy}") if not isinstance(warmup, v5.RGBDContext) or not isinstance(stream, v5.RGBDContext): raise TypeError("Warmup and stream must be RGBDContext values") if not isinstance(carrier, StreamResolutionCarrier): raise TypeError("Core B V3 requires the StreamResolutionCarrier") episode = v1._Episode(carrier, bounds, episode_id, warmup[0].scene_id) carrier.count_clamp = TRAINED_VIEW_COUNT if policy == "OFF_CLAMP3" else None try: fast = carrier.initial_fast(episode_id) for observation, depth in zip(warmup, warmup.depths, strict=True): episode.arrive(observation, depth) if policy != "NO_STREAM": for observation, depth in zip(stream, stream.depths, strict=True): episode.arrive(observation, depth) action = ACTIONS[policy] if action != Action.OFF: # ALL is committed from one trace, so both matrices see identical old weights. proposal = rule.propose(episode.run(fast, "trace")) fast = commit_write( fast, proposal, action, cache_revision=episode.cache.revision ) state = episode.run(fast, "materialize") finally: carrier.count_clamp = None return state, episode.anchor.clone(), fast, episode def save_write_rule(path, rule, *, seed, step, lock_sha256, carrier_sha256): payload = { "schema": WRITE_SCHEMA, "carrier_config": asdict(rule.carrier_config), "write_config": asdict(rule.write_config), "rule": {k: v.detach().cpu() for k, v in rule.state_dict().items()}, "seed": seed, "step": step, "lock_sha256": lock_sha256, "carrier_checkpoint_sha256": carrier_sha256, "slow_carrier_updated": False, } path = Path(path) if path.exists(): raise FileExistsError(path) torch.save(payload, path) restored, _ = load_write_rule(path, "cpu") if hash_value(restored.state_dict()) != hash_value(payload["rule"]): raise RuntimeError("Write-rule restore mismatch") return sha(path) def load_write_rule(path, device): payload = torch.load(path, map_location=device, weights_only=True) if payload.get("schema") != WRITE_SCHEMA or payload.get("slow_carrier_updated"): raise PermissionError("Core B V3 write-rule schema required") carrier_config = CarrierConfig(**payload["carrier_config"]) rule = DirectWriteRule(carrier_config, WriteConfig(**payload["write_config"])).to(device) rule.load_state_dict(payload["rule"], strict=True) return rule, payload def training_loss(state, local_cameras, rgb, depth, indices): """The frozen V2 C1 loss at 32^3 (the loss the V17 carrier was trained with).""" return v8.training_loss(state, local_cameras, rgb, depth, indices, V8_LOSS_VARIANT)