src/mcss/dynamic/runner.py
"""Causal single-episode deployment orchestration; no label or teacher imports."""
from dataclasses import asdict
from time import perf_counter
import torch
from mcss.dynamic.budget import WorkBudget
from mcss.dynamic.cache import ObservationCache
from mcss.dynamic.feedback import anchored_observation, preupdate_feedback
from mcss.dynamic.types import (
Action,
ControlInput,
SealedScene,
clone_scene_state,
hash_fast_content,
hash_scene_state,
hash_value,
)
from mcss.dynamic.write_rule import commit_write
class StreamingRunner:
def __init__(self, carrier, write_rule, renderer, policy, *, max_units=1e12):
self.carrier = carrier
self.write_rule = write_rule
self.renderer = renderer
self.policy = policy
self.max_units = max_units
self.active = False
for module in (carrier, write_rule, renderer):
module.eval()
module.requires_grad_(False)
if isinstance(policy, torch.nn.Module):
policy.eval()
policy.requires_grad_(False)
if list(renderer.parameters()):
raise ValueError("Dynamic main method requires a parameter-free renderer")
self.checkpoint_hash = hash_value(
{"carrier": carrier.state_dict(), "write_rule": write_rule.state_dict()}
)
@torch.no_grad()
def reset(
self, warmup, *, episode_id, scene_id, split_id, query_vault_id, stream_steps, query_count=4
):
self.active = False
if not warmup or stream_steps < 0 or query_count < 0:
raise ValueError("Invalid episode protocol")
if any(observation.scene_id != scene_id for observation in warmup):
raise ValueError("Warmup scene identity mismatch")
image_size = warmup[0].camera.image_size
if any(observation.camera.image_size != image_size for observation in warmup):
raise ValueError("All observations must share declared resolution")
self.anchor_c2w = warmup[0].camera.c2w.detach().clone()
self.episode_id, self.scene_id = episode_id, scene_id
self.split_id, self.query_vault_id = split_id, query_vault_id
self.remaining_steps, self.query_count = stream_steps, query_count
self.image_size = image_size
self.cache = ObservationCache(episode_id, scene_id)
self.fast = self.carrier.initial_fast(episode_id)
self.budget = WorkBudget(
self.carrier.config,
image_size,
self.renderer.n_samples,
self.max_units,
policy=self.policy,
)
minimum = (
len(warmup) * self.budget.encode_units
+ self.budget.materialize_units(len(warmup))
+ self.budget.base_future_units(len(warmup), stream_steps, query_count)
)
if minimum > self.max_units:
raise ValueError("Budget is insufficient before episode start")
self.history = []
self.previous_action = Action.OFF
self.previous_camera = None
for observation in warmup:
arrived = anchored_observation(observation, self.anchor_c2w)
features = self.carrier.encode(arrived)
self.budget.record("encode", self.budget.encode_units)
self.cache = self.cache.append(arrived, features)
self.previous_camera = arrived.camera
self.scene_state = self.carrier.materialize(self.cache, self.fast)
self.budget.record(
"materialize",
self.budget.materialize_units(self.cache.revision),
views=self.cache.revision,
)
self.active = True
self.sealed = False
@torch.no_grad()
def step(self, observation):
if not self.active or self.sealed or self.remaining_steps == 0:
raise ValueError("Episode is not accepting observations")
if (
observation.scene_id != self.scene_id
or observation.camera.image_size != self.image_size
):
raise ValueError("Observation scene/resolution differs from episode")
if observation.frame_id <= self.cache.entries[-1].observation.frame_id:
raise ValueError("Observations must arrive in strictly increasing frame order")
if observation.rgb.is_cuda:
torch.cuda.synchronize(observation.rgb.device)
started = perf_counter()
arrived = anchored_observation(observation, self.anchor_c2w)
features = self.carrier.encode(arrived)
self.budget.record("encode", self.budget.encode_units)
# Current RGB is used only for comparison here; the cache and fast state are still old.
feedback = preupdate_feedback(
self.scene_state,
arrived,
self.renderer,
image_feature=features,
previous_camera=self.previous_camera,
)
self.budget.record("feedback_render", self.budget.render_units)
control = ControlInput(
**feedback,
current_image_mean=float(arrived.rgb.mean()),
density_mean=float(self.scene_state.density_logits.mean()),
fast_norm=float(self.fast.delta_fuse.norm() + self.fast.delta_complete.norm()),
observed_count=self.cache.revision,
previous_action=self.previous_action,
remaining_budget=self.budget.remaining,
fast_state_stats=(
float(self.fast.delta_fuse.norm()),
float(self.fast.delta_complete.norm()),
float(self.fast.delta_fuse.norm() + self.fast.delta_complete.norm()),
float(
torch.maximum(
self.fast.delta_fuse.abs().max(), self.fast.delta_complete.abs().max()
)
),
),
remaining_steps=self.remaining_steps,
)
self.budget.record("policy", self.budget.policy_units)
feasible = self.budget.feasible_actions(
self.cache.revision + 1, self.remaining_steps - 1, self.query_count
)
action = Action(self.policy.choose(control, feasible))
if action not in feasible:
raise ValueError("Policy selected an infeasible action")
self.cache = self.cache.append(arrived, features)
if action != Action.OFF:
trace = self.carrier.trace(self.cache, self.fast)
self.budget.record(
"trace",
self.budget.materialize_units(self.cache.revision),
views=self.cache.revision,
)
proposal = self.write_rule.propose(trace)
c = self.carrier.config
self.budget.record("proposal_pair", 2 * c.token_count * c.hidden_dim * c.expansion_dim)
self.fast = commit_write(
self.fast, proposal, action, cache_revision=self.cache.revision
)
self.scene_state = self.carrier.materialize(self.cache, self.fast)
self.budget.record(
"materialize",
self.budget.materialize_units(self.cache.revision),
views=self.cache.revision,
)
self.remaining_steps -= 1
self.previous_action = action
self.previous_camera = arrived.camera
if observation.rgb.is_cuda:
torch.cuda.synchronize(observation.rgb.device)
record = {
"frame_id": observation.frame_id,
"action": str(action),
"feedback": feedback,
"control_input": control.to_serializable(),
"observed_count": self.cache.revision,
"fast_step": self.fast.step,
"seconds": perf_counter() - started,
"fast_fuse_norm": float(self.fast.delta_fuse.norm()),
"fast_complete_norm": float(self.fast.delta_complete.norm()),
"budget": self.budget.report(),
}
self.history.append(record)
return record
def seal(self):
if not self.active or self.remaining_steps:
raise ValueError("Cannot seal before the declared stream is consumed")
state = clone_scene_state(self.scene_state)
self.sealed = True
config = {
"carrier": asdict(self.carrier.config),
"write": asdict(self.write_rule.write_config),
"render_samples": self.renderer.n_samples,
"budget": self.max_units,
}
return SealedScene(
self.episode_id,
self.scene_id,
self.split_id,
self.query_vault_id,
state,
hash_scene_state(state),
tuple(e.observation.frame_id for e in self.cache.entries),
self.checkpoint_hash,
hash_value(config),
hash_fast_content(self.fast),
self.anchor_c2w.detach().clone(),
)