| """Action-space routing over ActionCodec2Sparse profiles.""" |
|
|
| from __future__ import annotations |
|
|
| from ..presets import canonical_preset_name |
| from ..routing import _schedule |
|
|
| from collections.abc import Mapping, Sequence |
| from dataclasses import dataclass |
| from pathlib import Path |
| from types import MappingProxyType |
|
|
| 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 ConfigError, _integer, _mapping, _reject_unknown, _UniqueKeyLoader |
| from ..hf import save_huggingface_metadata |
| from ..presets import ComponentSpec, PresetConfig, format_action_spaces |
| from ..routing import ( |
| _PROFILE_WIDTHS, |
| _PROFILES, |
| ProfileSpec, |
| _ProfileRuntime, |
| _second_order_boundaries, |
| _second_order_dims, |
| ) |
| from ..semantic import resampled_horizon |
| from ..trajectory import ActiveTrajectory |
| from .processor import ActionCodec2SparseProcessor |
|
|
| SPARSE_ROUTER_MODEL_TYPE = "actioncodec2_sparse_router" |
| SPARSE_ROUTER_CONFIG_NAME = "sparse_router_config.yaml" |
| SPARSE_ROUTER_CONFIG_VERSION = 1 |
|
|
|
|
| RoutedState = tuple[tuple[bool, ...], tuple] |
|
|
|
|
| @dataclass(frozen=True) |
| class _Segment: |
| component: ComponentSpec |
| spec: ProfileSpec |
| processor: ActionCodec2SparseProcessor |
|
|
|
|
| class _FrontTracker: |
| def __init__(self, segment: _Segment, horizon: int) -> None: |
| tables = segment.processor.vocabulary.tables |
| self.emask, self.offset = tables.emask, segment.spec.token_offset |
| self.vocab, self.horizon = len(tables), horizon |
| self.front, self.closed = 0, 0 |
|
|
| def push(self, token: int) -> None: |
| self.closed |= self.emask[token] << self.front |
| while self.front < self.horizon and self.closed >> self.front & 1: |
| self.front += 1 |
|
|
|
|
| class RoutedSparseGrammar: |
| """Constrained-generation grammar of one preset and codec-frame horizon.""" |
|
|
| def __init__(self, router: ActionCodec2SparseRouter, *, preset: str, horizon: int) -> None: |
| self.horizon = _integer(horizon, "horizon", positive=True) |
| self.segments = router._segments(router.preset(preset)) |
| self.time_major = router.preset(preset).token_order == "time_major" |
| self.vocab_size, self.end_token = router.vocab_size, router.end_token |
| self._grammars = tuple( |
| s.processor.grammar(self.horizon, active_dims=s.component.active_dims) |
| for s in self.segments |
| ) |
| self._initial: RoutedState = ( |
| (False,) * len(self.segments), |
| tuple(g.initial() for g in self._grammars), |
| ) |
| self._cache = BoundedCache(256) |
|
|
| def _scheduled(self, state: RoutedState) -> int | None: |
| fronts = [g.front(s) for g, s in zip(self._grammars, state[1], strict=True)] |
| return _schedule(fronts, state[0], self.time_major) |
|
|
| def _advance(self, state: RoutedState, token: int) -> RoutedState | None: |
| finished, states = state |
| index = self._scheduled(state) |
| if index is None: |
| return None |
| if token == self.end_token: |
| if finished.count(False) < 2 or not self._grammars[index].is_complete(states[index]): |
| return None |
| return (*finished[:index], True, *finished[index + 1 :]), states |
| local = token - self.segments[index].spec.token_offset |
| updated = self._grammars[index].advance(states[index], local) |
| if updated is None: |
| return None |
| return finished, (*states[:index], updated, *states[index + 1 :]) |
|
|
| def _consume(self, tokens: Sequence[int]) -> RoutedState | None: |
| checked = [] |
| for token in tokens: |
| if isinstance(token, bool) or not isinstance(token, (int, np.integer)): |
| return None |
| checked.append(int(token)) |
| return cached_prefix_state(tuple(checked), self._cache, self._initial, self._advance) |
|
|
| def reserve_states(self, rows: int) -> None: |
| """Size the prefix cache for ``rows`` streams advanced in lockstep.""" |
| self._cache.reserve(4 * _integer(rows, "rows", positive=True)) |
|
|
| def _mask_from_state(self, state: RoutedState) -> np.ndarray: |
| output = np.zeros(self.vocab_size, dtype=bool) |
| index = self._scheduled(state) |
| if index is None: |
| return output |
| finished, states = state |
| mask = self._grammars[index].mask(states[index]) |
| start = self.segments[index].spec.token_offset |
| output[start : start + len(mask)] = mask |
| if finished.count(False) >= 2 and self._grammars[index].is_complete(states[index]): |
| output[self.end_token] = True |
| return output |
|
|
| def _state_complete(self, state: RoutedState) -> bool: |
| finished, states = state |
| open_ = [i for i, done in enumerate(finished) if not done] |
| return len(open_) == 1 and self._grammars[open_[0]].is_complete(states[open_[0]]) |
|
|
| def next_token_mask(self, tokens: Sequence[int]) -> np.ndarray: |
| """Boolean ``(vocab_size,)`` mask of legal next routed tokens.""" |
| state = self._consume(tokens) |
| if state is None: |
| raise ValueError("token prefix is not legal for this sparse preset") |
| return self._mask_from_state(state) |
|
|
| def is_complete(self, tokens: Sequence[int]) -> bool: |
| """Whether the stream may end here: one component left and its chunk complete.""" |
| state = self._consume(tokens) |
| return state is not None and self._state_complete(state) |
|
|
| def covered_steps(self, tokens: Sequence[int]) -> int: |
| """Leading codec frames that no further token can change.""" |
| state = self._consume(tokens) |
| if state is None: |
| raise ValueError("token prefix is not legal for this sparse preset") |
| return min( |
| self.horizon if done else grammar.front(local) |
| for grammar, local, done in zip(self._grammars, state[1], state[0], strict=True) |
| ) |
|
|
| def accepts(self, tokens: Sequence[int]) -> bool: |
| return self.is_complete(tokens) |
|
|
|
|
| class ActionCodec2SparseRouter: |
| """Route action-space tensors through ActionCodec2Sparse profile namespaces.""" |
|
|
| model_type = SPARSE_ROUTER_MODEL_TYPE |
|
|
| def __init__( |
| self, |
| profiles: Sequence[tuple[ProfileSpec, ActionCodec2SparseProcessor]], |
| presets: Sequence[PresetConfig], |
| ) -> None: |
| runtimes: dict[str, _ProfileRuntime] = {} |
| offset = 0 |
| for spec, processor in profiles: |
| if spec.name in runtimes or spec.name not in _PROFILES: |
| raise ConfigError(f"invalid or duplicate sparse profile {spec.name!r}") |
| if spec.token_offset != 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} fingerprint does not match its artifact") |
| if processor.action_dim != _PROFILE_WIDTHS[spec.name]: |
| raise ConfigError( |
| f"profile {spec.name!r} must have width {_PROFILE_WIDTHS[spec.name]}" |
| ) |
| runtimes[spec.name] = _ProfileRuntime(spec, processor) |
| offset += spec.token_budget |
| if not runtimes: |
| raise ConfigError("a sparse router needs at least one profile") |
| self.profiles = MappingProxyType(runtimes) |
| self._second_order_dims = { |
| name: _second_order_dims(runtime.processor.config) for name, runtime in runtimes.items() |
| } |
| self.end_token = offset |
| self.vocab_size = offset + 1 |
| self._grammars = BoundedCache(8) |
| self._action_spaces: dict[str, PresetConfig] = {} |
| for preset in presets: |
| self.register_action_space(preset) |
|
|
| @classmethod |
| def from_processors( |
| cls, |
| processors: Mapping[str, ActionCodec2SparseProcessor], |
| *, |
| token_budgets: Mapping[str, int], |
| presets: Sequence[PresetConfig], |
| ) -> ActionCodec2SparseRouter: |
| """Assign stable contiguous namespaces in canonical joint/EEF order.""" |
| if set(processors) != set(token_budgets) or set(processors) - _PROFILES: |
| raise ConfigError("processors and token_budgets must share joint/eef keys") |
| entries, offset = [], 0 |
| for name in ("joint", "eef"): |
| if name in processors: |
| budget = _integer(token_budgets[name], f"token_budgets.{name}", positive=True) |
| spec = ProfileSpec(name, offset, budget, processors[name].artifact_fingerprint) |
| entries.append((spec, processors[name])) |
| offset += budget |
| return cls(entries, presets) |
|
|
| @property |
| def action_spaces(self) -> Mapping[str, PresetConfig]: |
| return MappingProxyType(self._action_spaces) |
|
|
| def preset(self, name: str) -> PresetConfig: |
| 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: |
| """Register a validated preset, mapping, or YAML path without changing token IDs.""" |
| 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 YAML path") |
| missing = {c.space for c 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 second_order_boundaries(self, preset) -> 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 _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." |
| ) |
|
|
| @staticmethod |
| def _boundary(previous_first_order, segment: _Segment): |
| return ( |
| None |
| if previous_first_order is None |
| else previous_first_order.get(segment.component.name) |
| ) |
|
|
| def describe_action_spaces(self, *, selected=None, action_dim=None) -> str: |
| return format_action_spaces( |
| self.action_spaces, |
| selected=selected, |
| action_dim=action_dim, |
| codec_name="ActionCodec2Sparse", |
| chunk_horizon_decode=True, |
| ) |
|
|
| def __str__(self) -> str: |
| return self.describe_action_spaces() |
|
|
| def to(self, device: object) -> ActionCodec2SparseRouter: |
| """Return ``self`` because routed processors use CPU NumPy arrays.""" |
| del device |
| return self |
|
|
| def _segments(self, preset: PresetConfig) -> tuple[_Segment, ...]: |
| return tuple( |
| _Segment(c, self.profiles[c.space].spec, self.profiles[c.space].processor) |
| for c in preset.components |
| ) |
|
|
| def grammar(self, *, preset: str, horizon: int) -> RoutedSparseGrammar: |
| key = (preset, int(horizon)) |
| grammar = self._grammars.get(key) |
| if grammar is None: |
| grammar = RoutedSparseGrammar(self, preset=preset, horizon=horizon) |
| self._grammars[key] = grammar |
| return grammar |
|
|
| def encode_batch( |
| self, |
| inputs, |
| *, |
| preset: str, |
| fps=None, |
| current_state=None, |
| previous_first_order=None, |
| threads: int = 0, |
| ) -> list[list[int]]: |
| """Encode ``(B,T,D)`` source chunks at ``fps`` into routed token rows.""" |
| config = self.preset(preset) |
| self._validate_boundary_keys(previous_first_order, config) |
| canonical = config._canonicalize_batch(inputs, source_hz=fps, current_state=current_state) |
| segments = self._segments(config) |
| streams = [ |
| s.processor.encode_batch( |
| canonical[s.component.name], |
| active_dims=s.component.active_dims, |
| previous_first_order=self._boundary(previous_first_order, s), |
| threads=threads, |
| ) |
| for s in segments |
| ] |
| horizon = next(iter(canonical.values())).shape[1] |
| time_major = config.token_order == "time_major" |
| return [ |
| self._interleave([rows[i] for rows in streams], segments, horizon, time_major) |
| for i in range(len(streams[0])) |
| ] |
|
|
| def _interleave(self, streams, segments, horizon, time_major) -> list[int]: |
| trackers = [_FrontTracker(s, horizon) for s in segments] |
| cursors, finished, output = [0] * len(streams), [False] * len(streams), [] |
| while True: |
| index = _schedule([t.front for t in trackers], finished, time_major) |
| if index is None: |
| return output |
| if cursors[index] == len(streams[index]): |
| if finished.count(False) == 1: |
| return output |
| finished[index] = True |
| output.append(self.end_token) |
| continue |
| token = int(streams[index][cursors[index]]) |
| cursors[index] += 1 |
| trackers[index].push(token) |
| output.append(token + segments[index].spec.token_offset) |
|
|
| def _split(self, row, segments, horizon, time_major) -> list[list[int]]: |
| trackers = [_FrontTracker(s, horizon) for s in segments] |
| finished, output = [False] * len(segments), [[] for _ in segments] |
| for position, raw in enumerate(row): |
| if isinstance(raw, bool) or not isinstance(raw, (int, np.integer)): |
| raise ValueError(f"tokens[{position}] must be an integer") |
| index = _schedule([t.front for t in trackers], finished, time_major) |
| if index is None: |
| raise ValueError("token after every component is finished") |
| token = int(raw) |
| if token == self.end_token: |
| if finished.count(False) < 2: |
| raise ValueError("the end token cannot close the last open component") |
| finished[index] = True |
| continue |
| local = token - segments[index].spec.token_offset |
| if not 0 <= local < trackers[index].vocab: |
| raise ValueError( |
| f"tokens[{position}] is outside component {segments[index].component.name!r}" |
| ) |
| if trackers[index].front >= horizon: |
| raise ValueError("token after the last open frame") |
| trackers[index].push(local) |
| output[index].append(local) |
| return output |
|
|
| def _split_rows(self, rows, config, horizon): |
| horizon = _integer(horizon, "horizon", positive=True) |
| rows = list(rows) |
| if not rows: |
| raise ValueError("rows must not be empty") |
| segments = self._segments(config) |
| time_major = config.token_order == "time_major" |
| return segments, [self._split(row, segments, horizon, time_major) for row in rows] |
|
|
| def decode_batch( |
| self, |
| rows, |
| *, |
| preset: str, |
| horizon: int, |
| fps=None, |
| current_state=None, |
| previous_first_order=None, |
| ) -> dict[str, np.ndarray]: |
| """Decode rows of one codec-frame ``horizon`` to ``{component: (B,T_out,D_c)}``.""" |
| config = self.preset(preset) |
| self._validate_boundary_keys(previous_first_order, config) |
| segments, splits = self._split_rows(rows, config, horizon) |
| state = optional_finite_array(current_state, shape=(len(splits), config.input_dim)) |
| output_fps = config.source_hz if fps is None else fps |
| steps = resampled_horizon(horizon, source_hz=config.target_hz, target_hz=output_fps) |
| output = {} |
| for index, segment in enumerate(segments): |
| decoded = segment.processor.decode_batch( |
| [split[index] for split in splits], |
| horizon=horizon, |
| active_dims=segment.component.active_dims, |
| previous_first_order=self._boundary(previous_first_order, segment), |
| ) |
| output[segment.component.name] = externalize_batch( |
| decoded, |
| segment.component, |
| source_hz=config.target_hz, |
| target_hz=output_fps, |
| target_length=steps, |
| method=config.interpolation, |
| current_state=state, |
| ) |
| return output |
|
|
| def final_first_order( |
| self, |
| rows, |
| *, |
| preset: str, |
| horizon: int, |
| previous_first_order, |
| executed_steps: int | None = None, |
| ) -> dict[str, np.ndarray]: |
| """Component ``(B,D_active)`` boundaries after ``executed_steps`` codec frames.""" |
| 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, config) |
| segments, splits = self._split_rows(rows, config, horizon) |
| output = {} |
| for index, segment in enumerate(segments): |
| if not segment.processor.requires_previous_first_order: |
| continue |
| output[segment.component.name] = segment.processor.final_first_order( |
| [split[index] for split in splits], |
| horizon=horizon, |
| previous_first_order=self._boundary(previous_first_order, segment), |
| active_dims=segment.component.active_dims, |
| executed_steps=executed_steps, |
| ) |
| return output |
|
|
| 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 |
|
|
| @property |
| def new_action_tokens(self) -> list[str]: |
| """Stable token names: profile namespaces (budget slots included), then the end token.""" |
| names = [ |
| f"<action_{name}_{index:04d}>" |
| for name, runtime in self.profiles.items() |
| for index in range(runtime.spec.token_budget) |
| ] |
| return [*names, "<action_end>"] |
|
|
| def save_pretrained(self, directory) -> list[str]: |
| root = Path(directory) |
| files = [] |
| for name, runtime in self.profiles.items(): |
| files.extend(runtime.processor.save_pretrained(root / "profiles" / name)) |
| manifest = { |
| "model_type": self.model_type, |
| "config_version": SPARSE_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 / SPARSE_ROUTER_CONFIG_NAME |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8") |
| return [str(path), *files] |
|
|
| @staticmethod |
| def read_manifest(directory) -> tuple[tuple[ProfileSpec, ...], tuple[PresetConfig, ...]]: |
| """Read profile specs and action spaces without loading numerical arrays.""" |
| path = Path(directory) / SPARSE_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") != SPARSE_ROUTER_MODEL_TYPE: |
| raise ConfigError(f"router.model_type must be {SPARSE_ROUTER_MODEL_TYPE!r}") |
| if payload.get("config_version") != SPARSE_ROUTER_CONFIG_VERSION: |
| raise ConfigError( |
| f"unsupported sparse router config_version {payload.get('config_version')!r}" |
| ) |
| profiles, presets = payload.get("profiles"), payload.get("presets") |
| if not isinstance(profiles, list) or not isinstance(presets, list): |
| raise ConfigError("router.profiles and router.presets must be lists") |
| return ( |
| tuple( |
| ProfileSpec.from_dict(p, path=f"router.profiles[{i}]") |
| for i, p in enumerate(profiles) |
| ), |
| tuple( |
| PresetConfig.from_dict(p, path=f"router.presets[{i}]") |
| for i, p in enumerate(presets) |
| ), |
| ) |
|
|
| @classmethod |
| def from_pretrained(cls, directory) -> ActionCodec2SparseRouter: |
| root = Path(directory) |
| specs, presets = cls.read_manifest(root) |
| profiles = [ |
| (spec, ActionCodec2SparseProcessor.from_pretrained(root / "profiles" / spec.name)) |
| for spec in specs |
| ] |
| return cls(profiles, presets) |
|
|
|
|
| def save_router_metadata(directory, **processor_config) -> tuple[Path, Path]: |
| """Hugging Face discovery files for a saved sparse router or bound frontend.""" |
| return save_huggingface_metadata( |
| directory, |
| model_type="actioncodec2_sparse", |
| processor_class="ActionCodec2Sparse", |
| **processor_config, |
| ) |
|
|
