'Required-preset routing over the two shared ActionCodec2 processors.' from __future__ import annotations from .presets import canonical_preset_name import re from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass from pathlib import Path from types import MappingProxyType from typing import TYPE_CHECKING import numpy as np import yaml from ._validation import optional_finite_array from .batch_semantics import externalize_batch from .cache import BoundedCache, cached_prefix_state from .configuration import ( ActionCodec2Config, ConfigError, _integer, _mapping, _reject_unknown, _required, _UniqueKeyLoader, ) from .hf import ( ROUTER_ALLOW_PATTERNS, ROUTER_CONFIG_NAME, ROUTER_MODEL_TYPE, resolve_pretrained_directory, save_huggingface_metadata, save_model_card, ) from .native_backend import native_kernel from .presets import ( ComponentSpec, PresetConfig, format_action_spaces, ) from .processor import ActionCodec2Processor from .semantic import resampled_horizon from .trajectory import ActiveTrajectory if TYPE_CHECKING: from ...frontend.actioncodec2 import ActionCodec2 ROUTER_CONFIG_VERSION = 4 _NAME = re.compile(r"[a-z][a-z0-9_]*\Z") _PROFILES = {"eef", "joint"} _PROFILE_WIDTHS = {"eef": 7, "joint": 8} _ComponentStates = tuple[tuple[int, ...], ...] _RoutedState = tuple[int, _ComponentStates] @dataclass(frozen=True) class ProfileSpec: """Token namespace metadata for one shared physical processor.""" name: str token_offset: int token_budget: int artifact_fingerprint: str def __post_init__(self) -> None: if _NAME.fullmatch(self.name) is None: raise ConfigError(f"profile.name must match {_NAME.pattern!r}") _integer(self.token_offset, "profile.token_offset") _integer(self.token_budget, "profile.token_budget", positive=True) if re.fullmatch(r"[0-9a-f]{64}", self.artifact_fingerprint) is None: raise ConfigError("profile.artifact_fingerprint must be a SHA-256 hex digest") @classmethod def from_dict(cls, payload: object, *, path: str) -> ProfileSpec: values = _mapping(payload, path) allowed = {"name", "token_offset", "token_budget", "artifact_fingerprint"} _reject_unknown(values, allowed, path) return cls( name=str(_required(values, "name", path)), token_offset=_integer(_required(values, "token_offset", path), f"{path}.token_offset"), token_budget=_integer( _required(values, "token_budget", path), f"{path}.token_budget", positive=True, ), artifact_fingerprint=str(_required(values, "artifact_fingerprint", path)), ) def to_dict(self) -> dict[str, object]: return asdict(self) def _second_order_dims(config: ActionCodec2Config) -> frozenset[int]: return frozenset( dimension for component in config.physical.components.values() if component.primitive_order == 2 for dimension in component.dims ) def _second_order_boundaries( preset: PresetConfig, second_order_dims: Mapping[str, frozenset[int]] ) -> dict[str, int]: return { component.name: len(component.active_dims) for component in preset.components if second_order_dims[component.space].intersection(component.active_dims) } @dataclass(frozen=True) class _ProfileRuntime: spec: ProfileSpec processor: ActionCodec2Processor @dataclass(frozen=True) class _Segment: component: ComponentSpec profile: _ProfileRuntime def _schedule(fronts: Sequence[int], finished: Sequence[bool], time_major: bool) -> int | None: open_ = [i for i, done in enumerate(finished) if not done] if not open_: return None return min(open_, key=lambda i: (fronts[i], i)) if time_major else open_[0] class RoutedActionCodec2Grammar: 'Coverage grammar over a shared time grid and namespaced physical components.' def __init__(self, router: ActionCodec2Router, *, preset: str, horizon: int) -> None: horizon = _integer(horizon, "horizon", positive=True) self.segments = router._segments(router.preset(preset)) self.horizon = horizon self.token_order = router.preset(preset).token_order self.vocab_size = router.vocab_size self._grammars = tuple( segment.profile.processor.grammar(horizon, active_dims=segment.component.active_dims) for segment in self.segments ) self._native_grammars = tuple( getattr(grammar, "_native", None) for grammar in self._grammars ) self._native_offsets = tuple(segment.profile.spec.token_offset for segment in self.segments) states = tuple(grammar._initial_state() for grammar in self._grammars) self._initial_state = (self._next_component(states), states) self._state_cache = BoundedCache(256) def _compiled_grammars(self): if native_kernel("route_advance") is None: return None if any(g is None for g in self._native_grammars): self._native_grammars = tuple(g._compiled() for g in self._grammars) return self._native_grammars if all(g is not None for g in self._native_grammars) else None def _next_component(self, states: _ComponentStates) -> int: fronts = [grammar._covered_steps(state) for grammar, state in zip(self._grammars, states, strict=True)] index = _schedule(fronts, [front >= self.horizon for front in fronts], self.token_order == "time_major") return len(fronts) if index is None else index def _advance_routed_state( self, routed_state: _RoutedState, token: int, ) -> _RoutedState | None: segment_index, states = routed_state kernel = native_kernel("route_advance") compiled = self._compiled_grammars() if kernel is not None and compiled is not None: return kernel( compiled, self._native_offsets, self.token_order == "time_major", segment_index, states, token, ) if segment_index == len(self.segments): return None segment = self.segments[segment_index] local_token = token - segment.profile.spec.token_offset if not 0 <= local_token < segment.profile.processor.vocab_size: return None next_state = self._grammars[segment_index]._advance(states[segment_index], local_token) if next_state is None: return None updated = (*states[:segment_index], next_state, *states[segment_index + 1 :]) return self._next_component(updated), updated def _consume(self, tokens: Sequence[int]) -> _RoutedState | None: checked: list[int] = [] for raw_token in tokens: if isinstance(raw_token, bool) or not isinstance(raw_token, (int, np.integer)): return None checked.append(int(raw_token)) return cached_prefix_state( tuple(checked), self._state_cache, self._initial_state, self._advance_routed_state ) def reserve_states(self, rows: int) -> None: 'Size the prefix cache for ``rows`` streams advanced in lockstep.' self._state_cache.reserve(4 * _integer(rows, "rows", positive=True)) def next_token_mask(self, tokens: Sequence[int]) -> np.ndarray: """Return a global namespace mask for the next preset token.""" state = self._consume(tokens) if state is None: raise ValueError("token prefix is not legal for this routed preset") segment_index, states = state output = np.zeros(self.vocab_size, dtype=bool) if segment_index == len(self.segments): return output segment = self.segments[segment_index] mask = self._grammars[segment_index]._next_token_mask_from_state(states[segment_index]) start = segment.profile.spec.token_offset output[start : start + len(mask)] = mask return output def is_complete(self, tokens: Sequence[int]) -> bool: """Whether the prefix completes every segment exactly once.""" state = self._consume(tokens) return state is not None and state[0] == len(self.segments) def covered_steps(self, tokens: Sequence[int]) -> int: """Return the contiguous frame prefix covered by every component.""" state = self._consume(tokens) if state is None: raise ValueError("token prefix is not legal for this routed preset") return min( grammar._covered_steps(local) for grammar, local in zip(self._grammars, state[1], strict=True) ) def is_prefix_complete(self, tokens: Sequence[int], *, steps: int) -> bool: """Whether a legal token prefix fully covers the requested action frames.""" if type(steps) is not int or not 1 <= steps <= self.horizon: raise ValueError("steps must be an integer within the grammar horizon") state = self._consume(tokens) return state is not None and all( grammar._covered_steps(local) >= steps for grammar, local in zip(self._grammars, state[1], strict=True) ) def interleave(self, streams: Sequence[Sequence[int]]) -> list[int]: """Route per-component local token streams without changing their tokens.""" kernel = native_kernel("route_interleave") compiled = self._compiled_grammars() if kernel is not None and compiled is not None: return kernel(compiled, self._native_offsets, self.token_order == "time_major", streams) if len(streams) != len(self.segments): raise ValueError("one token stream is required per component") state = self._initial_state cursors = [0] * len(streams) output = [] while state[0] < len(streams): index = state[0] if cursors[index] == len(streams[index]): raise ValueError("component token stream ends before complete coverage") token = ( int(streams[index][cursors[index]]) + self.segments[index].profile.spec.token_offset ) state = self._advance_routed_state(state, token) if state is None: raise ValueError("component token stream violates the routed grammar") cursors[index] += 1 output.append(token) if any(cursor != len(stream) for cursor, stream in zip(cursors, streams, strict=True)): raise ValueError("extra component tokens after complete coverage") return output def complete_prefix(self, tokens: Sequence[int], *, steps: int) -> list[int]: 'Fill only unreturned future cells with zero-motion primitive atoms.' if type(steps) is not int or not 1 <= steps <= self.horizon: raise ValueError("steps must be an integer within the grammar horizon") if self.covered_steps(tokens) < steps: raise ValueError(f"action prefix does not cover the first {steps} frames") kernel = native_kernel("route_complete") compiled = self._compiled_grammars() if kernel is not None and compiled is not None: zero_atoms = [ ( np.r_[0, np.cumsum(segment.profile.processor.bins[:-1])] + segment.profile.processor._zero_cells ).tolist() for segment in self.segments ] return kernel( compiled, self._native_offsets, self.token_order == "time_major", [int(token) for token in tokens], steps, zero_atoms, ) output = [int(token) for token in tokens] state = self._consume(output) while state[0] < len(self.segments): index, states = state segment = self.segments[index] processor = segment.profile.processor local = states[index] dims = segment.component.active_dims if processor.serializer.scheme == "unified_set_bpe": atom = self._grammars[index].completion_token( local, processor._zero_cells, processor.bins ) token = segment.profile.spec.token_offset + atom state = self._advance_routed_state(state, token) if state is None: raise RuntimeError("zero-motion suffix atom violates the routed grammar") output.append(token) continue if processor.serializer.scheme == "bpe": dimension = dims[local[0] % len(dims)] else: full = sum(1 << d for d in dims) uncovered = next(full & ~mask for mask in local if mask != full) dimension = (uncovered & -uncovered).bit_length() - 1 atom = sum(processor.serializer.bins[:dimension]) + int( processor._zero_cells[dimension] ) token = segment.profile.spec.token_offset + atom state = self._advance_routed_state(state, token) if state is None: raise RuntimeError("zero-motion suffix atom violates the routed grammar") output.append(token) return output class ActionCodec2Router: 'Route heterogeneous dataset states into physical token namespaces.' model_type = ROUTER_MODEL_TYPE def __init__( self, profiles: Sequence[tuple[ProfileSpec, ActionCodec2Processor]], presets: Sequence[PresetConfig], ) -> None: runtimes: dict[str, _ProfileRuntime] = {} expected_offset = 0 for spec, processor in profiles: if spec.name in runtimes: raise ConfigError(f"duplicate physical profile {spec.name!r}") if spec.token_offset != expected_offset: raise ConfigError("profile token namespaces must be contiguous and ordered") if processor.vocab_size > spec.token_budget: raise ConfigError( f"profile {spec.name!r} uses {processor.vocab_size} tokens, " f"exceeding budget {spec.token_budget}" ) if processor.artifact_fingerprint != spec.artifact_fingerprint: raise ConfigError( f"profile {spec.name!r} artifact fingerprint does not match its processor" ) if processor.config.set_bpe.allow_cross_component_merge: raise ConfigError( "allow_cross_component_merge=true is not supported by the v1 router; " "a component-aware Set-BPE serializer and grammar are required" ) runtimes[spec.name] = _ProfileRuntime(spec, processor) expected_offset += spec.token_budget if not runtimes or set(runtimes) - _PROFILES: raise ConfigError("router profiles must be a non-empty subset of joint and eef") for name, runtime in runtimes.items(): if runtime.processor.action_dim != _PROFILE_WIDTHS[name]: raise ConfigError( f"profile {name!r} must have width {_PROFILE_WIDTHS[name]}, " f"got {runtime.processor.action_dim}" ) self.profiles = MappingProxyType(runtimes) self._second_order_dims = { name: _second_order_dims(runtime.processor.config) for name, runtime in runtimes.items() } self._grammar_cache = BoundedCache(8) self._action_spaces: dict[str, PresetConfig] = {} for preset in presets: self.register_action_space(preset) self.vocab_size = expected_offset self._token_cells = np.full(self.vocab_size, -1, dtype=np.int64) for runtime in runtimes.values(): start = runtime.spec.token_offset self._token_cells[start : start + runtime.processor.vocab_size] = ( runtime.processor.serializer.token_cell_counts ) @property def action_spaces(self) -> Mapping[str, PresetConfig]: """Read-only registered built-in and user-created action spaces.""" return MappingProxyType(self._action_spaces) def __str__(self) -> str: return self.describe_action_spaces() def second_order_boundaries(self, preset: str | PresetConfig) -> dict[str, int]: 'Map each component needing ``previous_first_order`` to its active width.' if not isinstance(preset, PresetConfig): preset = self.preset(preset) return _second_order_boundaries(preset, self._second_order_dims) def describe_action_spaces( self, *, selected: str | None = None, action_dim: int | None = None ) -> str: """Explain all layouts, or source-width candidates, using this artifact's state requirements.""" return format_action_spaces( self.action_spaces, selected=selected, action_dim=action_dim, second_order_boundaries={ name: self.second_order_boundaries(preset) for name, preset in self.action_spaces.items() }, ) @classmethod def from_processors( cls, processors: Mapping[str, ActionCodec2Processor], *, token_budgets: Mapping[str, int], presets: Sequence[PresetConfig], ) -> ActionCodec2Router: """Assign stable contiguous namespaces in canonical joint/eef order.""" if set(processors) != set(token_budgets): raise ConfigError("processors and token_budgets must share keys") if not processors or set(processors) - _PROFILES: raise ConfigError("processors must contain a non-empty subset of joint and eef") offset = 0 entries = [] for name in ("joint", "eef"): if name not in processors: continue processor = processors[name] budget = _integer(token_budgets[name], f"token_budgets.{name}", positive=True) spec = ProfileSpec( name, offset, budget, processor.artifact_fingerprint, ) entries.append((spec, processor)) offset += budget return cls(entries, presets) def preset(self, name: str) -> PresetConfig: """Resolve one preset or fail with the available names.""" try: return self._action_spaces[canonical_preset_name(name)] except KeyError as error: raise KeyError( f"unknown preset {name!r}; available: {sorted(self._action_spaces)}" ) from error def register_action_space( self, preset: PresetConfig | Mapping[str, object] | str | Path, ) -> PresetConfig: 'Register a validated action-space config, mapping, or YAML file.' if isinstance(preset, Mapping): preset = PresetConfig.from_dict(preset) elif isinstance(preset, (str, Path)): preset = PresetConfig.from_yaml(preset) if not isinstance(preset, PresetConfig): raise TypeError( "register_action_space expects PresetConfig, a mapping, or a preset YAML path" ) missing = {component.space for component in preset.components} - self.profiles.keys() if missing: raise ConfigError( f"action space {preset.name!r} requires missing profiles: {sorted(missing)}" ) existing = self._action_spaces.get(preset.name) if existing is not None and existing != preset: raise ConfigError( f"action-space name {preset.name!r} is already registered differently" ) self._action_spaces[preset.name] = preset return preset def _segments(self, preset: PresetConfig) -> tuple[_Segment, ...]: return tuple( _Segment(component, self.profiles[component.space]) for component in preset.components ) def configure_encoding( self, *, spatial_search: bool = False, temporal_search: bool = False, search_window: int = 32, search_beam: int = 4, ) -> ActionCodec2Router: """Apply runtime packing options to all profiles without changing artifacts.""" from .packing import EncodingSearch options = EncodingSearch(spatial_search, temporal_search, search_window, search_beam) for profile in self.profiles.values(): serializer = profile.processor.serializer options.validate_codec(serializer.scheme, serializer.bins) for profile in self.profiles.values(): profile.processor.configure_encoding( spatial_search=spatial_search, temporal_search=temporal_search, search_window=search_window, search_beam=search_beam, ) return self def encode_batch( self, inputs: np.ndarray, *, preset: str, fps: float | None = None, current_state: np.ndarray | None = None, previous_first_order: Mapping[str, np.ndarray] | None = None, ) -> list[list[int]]: 'Encode one regular-rate ``[B,T,D]`` tensor at ``fps``.' preset_config = self.preset(preset) self._validate_boundary_keys(previous_first_order, preset_config) canonical = preset_config._canonicalize_batch( inputs, source_hz=fps, current_state=current_state ) batch = len(next(iter(canonical.values()))) component_streams = [] for segment in self._segments(preset_config): values = canonical[segment.component.name] streams = segment.profile.processor.encode_batch( values, active_dims=segment.component.active_dims, previous_first_order=None if previous_first_order is None else previous_first_order.get(segment.component.name), ) component_streams.append(streams) steps = next(iter(canonical.values())).shape[1] grammar = self.grammar(preset=preset, horizon=steps) kernel = native_kernel("route_interleave_batch") compiled = grammar._compiled_grammars() if kernel is not None and compiled is not None: return kernel( compiled, grammar._native_offsets, grammar.token_order == "time_major", component_streams, ) return [ grammar.interleave([streams[row] for streams in component_streams]) for row in range(batch) ] def route_trajectories( self, inputs: np.ndarray, *, preset: str, current_state: np.ndarray | None = None, fps: float | None = None, previous_first_order: Mapping[str, np.ndarray] | None = None, ) -> dict[str, list[ActiveTrajectory]]: 'Expose canonical complete trajectories grouped by physical profile.' preset_config = self.preset(preset) self._validate_boundary_keys(previous_first_order, preset_config) canonical = preset_config.canonicalize(inputs, current_state=current_state, source_hz=fps) output: dict[str, list[ActiveTrajectory]] = {} for component in preset_config.components: trajectories = canonical[component.name] processor = self.profiles[component.space].processor previous = ( None if previous_first_order is None else previous_first_order.get(component.name) ) dense = processor._dense_previous_first_order( previous, component.active_dims, len(trajectories) ) rows = ( [None] * len(trajectories) if dense is None else np.take(dense, component.active_dims, axis=1) ) output.setdefault(component.space, []).extend( ActiveTrajectory(t.values, t.active_dims, previous_first_order=p) for t, p in zip(trajectories, rows, strict=True) ) return output def _validate_boundary_keys( self, previous_first_order: Mapping[str, np.ndarray] | None, preset: PresetConfig ) -> None: if previous_first_order is None: previous_first_order = {} if not isinstance(previous_first_order, Mapping): raise ValueError( "previous_first_order must map component names to canonical [B,D_active] boundaries" ) unknown = set(previous_first_order) - {c.name for c in preset.components} if unknown: raise ValueError(f"unknown previous_first_order component(s): {sorted(unknown)}") required = { name: f"(B, {width})" for name, width in self.second_order_boundaries(preset).items() } missing = set(required) - previous_first_order.keys() if missing: raise ValueError( "second-order tokenization requires previous_first_order with arrays " f"{required}; missing {sorted(missing)}. Use canonical {preset.target_hz:g} Hz " "reconstructed increments. codec.first_order_template(batch_size) explicitly " "creates zeros; use zeros only if the preceding increment is known to be zero." ) def final_first_order( self, streams: Sequence[Sequence[int]], *, preset: str, previous_first_order: Mapping[str, np.ndarray], executed_steps: int | None = None, ) -> dict[str, np.ndarray]: 'Return component ``[B,D_active]`` boundaries at the canonical clock.' preset_config = self.preset(preset) if previous_first_order is None: raise ValueError("final_first_order requires previous_first_order") self._validate_boundary_keys(previous_first_order, preset_config) segments = self._segments(preset_config) _, component_streams = self._decode_streams(streams, preset_config, segments) output = {} for index, segment in enumerate(segments): processor = segment.profile.processor if not processor.requires_previous_first_order: continue output[segment.component.name] = processor.final_first_order( component_streams[index], previous_first_order=previous_first_order.get(segment.component.name), active_dims=segment.component.active_dims, executed_steps=executed_steps, ) return output def _split_stream( self, tokens: Sequence[int], *, preset: PresetConfig, steps: int ) -> list[list[int]]: grammar = self.grammar(preset=preset.name, horizon=steps) kernel = native_kernel("route_split") compiled = grammar._compiled_grammars() if kernel is not None and compiled is not None: checked = [] for cursor, token in enumerate(tokens): if isinstance(token, bool) or not isinstance(token, (int, np.integer)): raise ValueError(f"tokens[{cursor}] must be an integer") checked.append(int(token)) return kernel( compiled, grammar._native_offsets, grammar.token_order == "time_major", checked ) state = grammar._initial_state output = [[] for _ in grammar.segments] for cursor, raw_token in enumerate(tokens): if isinstance(raw_token, bool) or not isinstance(raw_token, (int, np.integer)): raise ValueError(f"tokens[{cursor}] must be an integer") index = state[0] if index == len(grammar.segments): raise ValueError("Action token stream has extra tokens after complete coverage") segment = grammar.segments[index] state = grammar._advance_routed_state(state, int(raw_token)) if state is None: raise ValueError( f"Action token stream has invalid coverage for component {segment.component.name!r}; check the preset token_order and constrain generation with codec.grammar(H, fps=...)." ) output[index].append(int(raw_token) - segment.profile.spec.token_offset) if state[0] != len(grammar.segments): raise ValueError("Incomplete action token stream; discard or regenerate this row") return output def _token_cell_count(self, raw_token: object, *, index: int) -> int: if isinstance(raw_token, bool) or not isinstance(raw_token, (int, np.integer)): raise ValueError(f"tokens[{index}] must be an integer") token = int(raw_token) if 0 <= token < self.vocab_size and self._token_cells[token] >= 0: return int(self._token_cells[token]) raise ValueError(f"tokens[{index}] is outside the routed profile vocabularies") def _decoded_steps(self, tokens: Sequence[int], segments: Sequence[_Segment]) -> int: if not tokens: raise ValueError("token stream must not be empty") cells_per_step = sum(len(segment.component.active_dims) for segment in segments) represented_cells = sum( self._token_cell_count(token, index=index) for index, token in enumerate(tokens) ) steps, remainder = divmod(represented_cells, cells_per_step) if steps < 1 or remainder: raise ValueError( "Incomplete action token stream: " f"{cells_per_step - remainder} more coordinate values are needed to finish a step. " "Constrain generation with codec.grammar(H, fps=...).next_token_mask(prefix) " "and check is_complete(tokens) for the intended H before decoding. " "Use codec.batch_decode(..., errors='return') to isolate invalid rows; " "do not pad truncated tokens with zeros." ) return steps def _batch_steps( self, streams: Sequence[Sequence[int]], segments: Sequence[_Segment], ) -> int: if not streams: raise ValueError("streams must not be empty") lengths = [self._decoded_steps(stream, segments) for stream in streams] if len(set(lengths)) != 1: raise ValueError( f"all token streams must decode to the same length, got {lengths}; use codec.batch_decode(...) for independent variable-length outputs" ) return lengths[0] def decode_batch( self, streams: Sequence[Sequence[int]], *, preset: str, fps: float | None = None, current_state: np.ndarray | None = None, previous_first_order: Mapping[str, np.ndarray] | None = None, ) -> dict[str, np.ndarray]: 'Decode rows to deltas, or absolute components with ``current_state``.' preset_config = self.preset(preset) self._validate_boundary_keys(previous_first_order, preset_config) segments = self._segments(preset_config) internal_steps, component_streams = self._decode_streams(streams, preset_config, segments) state_array = optional_finite_array( current_state, shape=(len(streams), preset_config.input_dim) ) output_fps = preset_config.source_hz if fps is None else fps output_steps = resampled_horizon( internal_steps, source_hz=preset_config.target_hz, target_hz=output_fps, ) output: dict[str, np.ndarray] = {} for segment_index, segment in enumerate(segments): decoded_batch = segment.profile.processor.decode_batch( component_streams[segment_index], active_dims=segment.component.active_dims, previous_first_order=None if previous_first_order is None else previous_first_order.get(segment.component.name), ) output[segment.component.name] = externalize_batch( decoded_batch, segment.component, source_hz=preset_config.target_hz, target_hz=output_fps, target_length=output_steps, method=preset_config.interpolation, current_state=state_array, ) return output def _decode_streams(self, streams, preset, segments): kernel = native_kernel("route_decode_batch") if kernel is not None: key = next( (key for key in reversed(tuple(self._grammar_cache)) if key[0] == preset.name), None ) grammar = self._grammar_cache.get(key) if grammar is None: grammar = self.grammar(preset=preset.name, horizon=1) compiled = grammar._compiled_grammars() if compiled is not None: width = sum(len(segment.component.active_dims) for segment in segments) return kernel( compiled, grammar._native_offsets, grammar.token_order == "time_major", streams, self._token_cells, width, ) steps = self._batch_steps(streams, segments) splits = [self._split_stream(row, preset=preset, steps=steps) for row in streams] return steps, [[row[i] for row in splits] for i in range(len(segments))] def decode_prefix_batch( self, streams, *, preset: str, horizon: int, steps: int, fps: float | None = None, current_state=None, previous_first_order=None, ) -> dict[str, np.ndarray]: 'Decode only fully covered execution frames on the native codec clock.' config = self.preset(preset) output_fps = config.source_hz if fps is None else fps if output_fps != config.target_hz: raise ValueError("prefix decode requires fps equal to the native codec clock") grammar = self.grammar(preset=preset, horizon=horizon) completed = [grammar.complete_prefix(row, steps=steps) for row in streams] decoded = self.decode_batch( completed, preset=preset, fps=output_fps, current_state=current_state, previous_first_order=previous_first_order, ) return {name: values[:, :steps] for name, values in decoded.items()} def grammar(self, *, preset: str, horizon: int) -> RoutedActionCodec2Grammar: """Build the global constrained grammar for one preset and target horizon.""" horizon = _integer(horizon, "horizon", positive=True) self.preset(preset) key = (preset, horizon) grammar = self._grammar_cache.get(key) if grammar is None: grammar = RoutedActionCodec2Grammar(self, preset=preset, horizon=horizon) self._grammar_cache[key] = grammar return grammar @property def new_action_tokens(self) -> list[str]: """Stable physical token names, including reserved profile budget slots.""" return [ f"" for name, runtime in self.profiles.items() for index in range(runtime.spec.token_budget) ] def save_pretrained(self, directory: str | Path) -> list[str]: """Save one manifest plus isolated physical processor subdirectories.""" root = Path(directory) files = [] for runtime in self.profiles.values(): files.extend(runtime.processor.save_pretrained(root / "profiles" / runtime.spec.name)) manifest = { "model_type": self.model_type, "config_version": ROUTER_CONFIG_VERSION, "profiles": [runtime.spec.to_dict() for runtime in self.profiles.values()], "presets": [preset.to_dict() for preset in self.action_spaces.values()], } path = root / ROUTER_CONFIG_NAME path.parent.mkdir(parents=True, exist_ok=True) path.write_text(yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8") metadata = save_huggingface_metadata( root, model_type=self.model_type, processor_class=type(self).__name__, ) card = save_model_card(root, str(self)) return [str(path), str(card), *map(str, metadata), *files] @classmethod def from_pretrained( cls, pretrained_model_name_or_path: str | Path, **kwargs: object, ) -> ActionCodec2Router | ActionCodec2: 'Load a router, or bind its public processor when given ``action_space``.' if "action_space" in kwargs: from ...frontend.actioncodec2 import ActionCodec2 return ActionCodec2.from_pretrained(pretrained_model_name_or_path, **kwargs) root = resolve_pretrained_directory( pretrained_model_name_or_path, allow_patterns=ROUTER_ALLOW_PATTERNS, **kwargs, ) specs, presets = cls.read_manifest(root) profiles = [ (spec, ActionCodec2Processor.from_pretrained(root / "profiles" / spec.name)) for spec in specs ] return cls(profiles, presets) @classmethod def read_manifest( cls, directory: str | Path ) -> tuple[tuple[ProfileSpec, ...], tuple[PresetConfig, ...]]: "Read one artifact's profile specs and action spaces, loading no arrays." path = Path(directory) / ROUTER_CONFIG_NAME payload = _mapping( yaml.load(path.read_text(encoding="utf-8"), Loader=_UniqueKeyLoader), "router" ) _reject_unknown(payload, {"model_type", "config_version", "profiles", "presets"}, "router") if payload.get("model_type") != ROUTER_MODEL_TYPE: raise ConfigError(f"router.model_type must be {ROUTER_MODEL_TYPE!r}") version = payload.get("config_version") if version not in (3, ROUTER_CONFIG_VERSION): raise ConfigError( f"unsupported router config_version {payload.get('config_version')!r}" ) raw_profiles = payload.get("profiles") raw_presets = payload.get("presets") if not isinstance(raw_profiles, list) or not isinstance(raw_presets, list): raise ConfigError("router.profiles and router.presets must be lists") if version == 3: raw_presets = [ dict(_mapping(raw, "router.preset"), token_order="time_major") for raw in raw_presets ] elif any("token_order" not in _mapping(raw, "router.preset") for raw in raw_presets): raise ConfigError("router config_version=4 requires explicit preset.token_order") return ( tuple( ProfileSpec.from_dict(raw, path=f"router.profiles[{index}]") for index, raw in enumerate(raw_profiles) ), tuple( PresetConfig.from_dict(raw, path=f"router.presets[{index}]") for index, raw in enumerate(raw_presets) ), ) @classmethod def describe_pretrained( cls, directory: str | Path, *, selected: str | None = None, action_dim: int | None = None, ) -> str: "Explain a saved artifact's layouts without loading primitives or BPE tables." root = Path(directory) specs, presets = cls.read_manifest(root) second_order_dims = { spec.name: _second_order_dims( ActionCodec2Config.from_pretrained(root / "profiles" / spec.name) ) for spec in specs } return format_action_spaces( {preset.name: preset for preset in presets}, selected=selected, action_dim=action_dim, second_order_boundaries={ preset.name: _second_order_boundaries(preset, second_order_dims) for preset in presets }, ) def eval(self) -> ActionCodec2Router: """Return ``self`` because all constituent processors are stateless.""" return self def to(self, device: object) -> ActionCodec2Router: """Return ``self`` because routed processors use CPU NumPy arrays.""" del device return self