"""Runtime-only values: no dataset, teacher, or query readers are imported here.""" from dataclasses import dataclass, field, fields, is_dataclass from enum import StrEnum from hashlib import sha256 from math import isfinite from numbers import Integral, Real from uuid import uuid4 import torch from torch import Tensor from mcss.types import Cameras, SceneState class Action(StrEnum): OFF = "OFF" FUSE = "FUSE" COMPLETE = "COMPLETE" ALL = "ALL" ACTION_ORDER: tuple[Action, ...] = (Action.OFF, Action.FUSE, Action.COMPLETE, Action.ALL) @dataclass(frozen=True) class OnlineObservation: scene_id: str frame_id: int rgb: Tensor camera: Cameras def __post_init__(self): if not self.scene_id or not isinstance(self.frame_id, int) or self.frame_id < 0: raise ValueError("Observation identity is invalid") if self.rgb.shape != (3, *self.camera.image_size) or self.camera.leading_shape: raise ValueError("One observation requires RGB [3,H,W] and one unbatched camera") if self.rgb.device != self.camera.device or self.rgb.dtype != self.camera.dtype: raise ValueError("Observation RGB and camera must share device and dtype") if not torch.isfinite(self.rgb).all() or ((self.rgb < 0) | (self.rgb > 1)).any(): raise ValueError("Observation RGB must be finite and within [0,1]") @dataclass(frozen=True) class FastWeights: episode_id: str delta_fuse: Tensor delta_complete: Tensor step: int = 0 branch_id: str = field(default_factory=lambda: uuid4().hex) state_id: str = field(default_factory=lambda: uuid4().hex) def __post_init__(self): if not self.episode_id or not self.branch_id or not self.state_id or self.step < 0: raise ValueError("Fast state identity/version is invalid") if self.delta_fuse.ndim != 2 or self.delta_complete.shape != self.delta_fuse.shape: raise ValueError("Fast matrices must have matching [d,m] shapes") if self.delta_fuse.device != self.delta_complete.device: raise ValueError("Fast matrices must share device") if not all(torch.isfinite(t).all() for t in (self.delta_fuse, self.delta_complete)): raise ValueError("Nonfinite fast weights") def fork(self) -> "FastWeights": """Create a private trajectory branch; clones retain offline autograd paths.""" return FastWeights( self.episode_id, self.delta_fuse.clone(), self.delta_complete.clone(), self.step ) @dataclass(frozen=True) class WriteTrace: episode_id: str cache_revision: int fast_step: int fuse_activation: Tensor complete_activation: Tensor observed_feature_statistics: Tensor support_weights: Tensor candidate_ids: Tensor branch_id: str source_state_id: str @dataclass(frozen=True) class WriteProposal: episode_id: str cache_revision: int fast_step: int delta_fuse: Tensor delta_complete: Tensor branch_id: str source_state_id: str @dataclass(frozen=True) class ControlInput: """Values available to the policy before the current observation is cached. The first eight fields are the original runtime contract. New fields are optional so old callers can keep constructing a control input positionally. Tensors are accepted at the boundary for local use, but the runner records the serializable summary returned by :meth:`to_serializable`. """ rgb_mse: float opacity_mean: float current_image_mean: float density_mean: float fast_norm: float observed_count: int previous_action: Action remaining_budget: float image_feature: Tensor | None = None image_feature_stats: Tensor | tuple[float, ...] | None = None rgb_residual_4x4: Tensor | tuple[float, ...] | None = None coverage: float = 0.0 valid_count: int = 0 camera_change: Tensor | tuple[float, ...] | None = None fast_state_stats: Tensor | tuple[float, ...] | None = None remaining_steps: int = 0 def __post_init__(self): object.__setattr__(self, "previous_action", Action(self.previous_action)) for name in ( "rgb_mse", "opacity_mean", "current_image_mean", "density_mean", "fast_norm", "remaining_budget", "coverage", ): value = getattr(self, name) if isinstance(value, bool) or not isinstance(value, Real) or not isfinite(value): raise ValueError(f"ControlInput {name} must be finite") if self.remaining_budget < 0: raise ValueError("ControlInput remaining_budget must be nonnegative") for name in ("observed_count", "valid_count", "remaining_steps"): value = getattr(self, name) if isinstance(value, bool) or not isinstance(value, Integral) or value < 0: raise ValueError(f"ControlInput {name} must be a nonnegative integer") def to_serializable(self) -> dict[str, object]: """Return a JSON-compatible prefix record without the raw feature map.""" def values(value): if value is None: return None if isinstance(value, Tensor): value = value.detach().cpu().reshape(-1).tolist() return [float(item) for item in value] return { "rgb_mse": float(self.rgb_mse), "opacity_mean": float(self.opacity_mean), "current_image_mean": float(self.current_image_mean), "density_mean": float(self.density_mean), "fast_norm": float(self.fast_norm), "observed_count": int(self.observed_count), "previous_action": str(self.previous_action), "remaining_budget": float(self.remaining_budget), "image_feature_stats": values(self.image_feature_stats), "rgb_residual_4x4": values(self.rgb_residual_4x4), "coverage": float(self.coverage), "valid_count": int(self.valid_count), "camera_change": values(self.camera_change), "fast_state_stats": values(self.fast_state_stats), "remaining_steps": int(self.remaining_steps), } @dataclass(frozen=True) class SealedScene: episode_id: str scene_id: str split_id: str query_vault_id: str scene_state: SceneState state_hash: str observed_ids: tuple[int, ...] checkpoint_hash: str config_hash: str fast_state_hash: str anchor_c2w: Tensor def hash_value(value) -> str: """Hash content, shape and dtype; device location is not part of tensor identity.""" digest = sha256() def visit(item): if isinstance(item, Tensor): data = item.detach().contiguous().cpu() digest.update(f"Tensor:{data.dtype}:{tuple(data.shape)}:".encode()) digest.update(data.reshape(-1).view(torch.uint8).numpy().tobytes()) elif is_dataclass(item): digest.update(type(item).__name__.encode()) for field in fields(item): digest.update(field.name.encode()) visit(getattr(item, field.name)) elif isinstance(item, dict): for key in sorted(item): visit(key) visit(item[key]) elif isinstance(item, (list, tuple)): digest.update(f"sequence:{len(item)}:".encode()) for child in item: visit(child) else: digest.update(f"{type(item).__name__}:{item!r};".encode()) visit(value) return digest.hexdigest() def hash_scene_state(state: SceneState) -> str: return hash_value(state) def hash_fast_content(fast: FastWeights) -> str: """Reproducible numerical identity, excluding the runtime-only branch nonce.""" return hash_value( {"delta_fuse": fast.delta_fuse, "delta_complete": fast.delta_complete, "step": fast.step} ) def clone_scene_state(state: SceneState) -> SceneState: def clone(item): if isinstance(item, Tensor): return item.detach().clone() if is_dataclass(item): return type(item)( **{field.name: clone(getattr(item, field.name)) for field in fields(item)} ) return item return clone(state)