File size: 14,174 Bytes
7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 | """Action-space-bound, arbitrary-horizon ActionCodec2 frontend."""
from __future__ import annotations
import sys
import warnings
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import TextIO
import numpy as np
import torch
from ...backends.actioncodec2 import (
ActionCodec2Router,
PresetConfig,
)
from ...backends.actioncodec2.native_backend import native_kernel
from ...backends.actioncodec2.semantic import resampled_horizon
from .._types import ActionIdxBatch, TokenIds
from ..base_action_tokenizer import BaseActionTokenizer
from ..layout import ActionSpaceMixin
from ..conversion import _numpy, _numpy_token_rows, restore_source_layout
class ActionCodec2Tokenizer(ActionSpaceMixin, BaseActionTokenizer):
'ActionCodec2 frontend bound to physical tensor semantics, not a horizon.'
encode_kwargs = frozenset({"fps", "current_state", "previous_first_order"})
def __init__(self, backend: ActionCodec2Router, *, action_space: str) -> None:
if not isinstance(backend, ActionCodec2Router):
raise TypeError(
f"ActionCodec2Tokenizer requires ActionCodec2Router, got {type(backend).__name__}"
)
preset = backend.preset(action_space)
if preset.target_hz != 15.0:
raise ValueError(
f"ActionCodec2 frontend artifacts require target_hz=15; got {preset.target_hz}. "
"Pass the source data rate via fps; use a physical profile calibrated at 15 Hz. "
"Changing target_hz alone does not recalibrate physical bounds."
)
self.action_tokenizer = backend
self._action_space = preset
self._action_vocab_size = int(backend.vocab_size)
self._new_action_tokens = list(backend.new_action_tokens)
self.action_dim = preset.input_dim
self._warned_default_fps = False
self.requires_previous_first_order = bool(backend.second_order_boundaries(preset))
self._encoded_action_dim = sum(
len(component.active_dims) for component in preset.components
)
self.max_action_token_len = None
self.fixed_action_token_len = None
super().__init__(append_bar_sentinels=False)
@classmethod
def from_pretrained(
cls,
pretrained_model_name_or_path: str | Path,
*,
action_space: str | Path | PresetConfig | Mapping[str, object],
**kwargs: object,
) -> ActionCodec2Tokenizer:
'Load a local/Hub router and bind one explicit physical action space.'
router = ActionCodec2Router.from_pretrained(pretrained_model_name_or_path, **kwargs)
return cls._from_router(router, action_space=action_space)
@classmethod
def _from_router(cls, router: ActionCodec2Router, *, action_space):
if isinstance(action_space, (PresetConfig, Mapping, Path)):
name = router.register_action_space(action_space).name
elif isinstance(action_space, str):
candidate = Path(action_space)
name = (
router.register_action_space(candidate).name
if candidate.is_file()
else action_space
)
else:
raise TypeError(
"action_space must be a registered name, PresetConfig, mapping, or YAML path"
)
return cls(router, action_space=name)
def configure_encoding(
self,
*,
spatial_search: bool = False,
temporal_search: bool = False,
search_window: int = 32,
search_beam: int = 4,
) -> ActionCodec2Tokenizer:
'Opt into native packing search; both independent stages default off.'
self.action_tokenizer.configure_encoding(
spatial_search=spatial_search,
temporal_search=temporal_search,
search_window=search_window,
search_beam=search_beam,
)
return self
def _describe_action_spaces(self, action_dim: int | None = None) -> str:
return self.action_tokenizer.describe_action_spaces(
selected=self.action_space.name, action_dim=action_dim
)
def print_action_spaces(
self, *, action_dim: int | None = None, file: TextIO | None = None
) -> None:
"""Print all layouts or candidates matching ``action_dim`` without changing the binding."""
description = self._describe_action_spaces(action_dim)
print(description, file=sys.stdout if file is None else file)
def for_action_space(self, name: str) -> ActionCodec2Tokenizer:
"""Return a new tokenizer sharing weights but bound to another option."""
return type(self)(self.action_tokenizer, action_space=name)
def _source_fps(self, fps, preset=None):
preset = self.action_space if preset is None else preset
if fps is None:
if not self._warned_default_fps:
warnings.warn(
f"fps was omitted; assuming {preset.source_hz:g} Hz from action_space "
f"{preset.name!r}. Pass the actual fps explicitly to avoid changing action timing.",
UserWarning,
stacklevel=3,
)
self._warned_default_fps = True
return preset.source_hz
return fps
def max_tokens_for_horizon(self, horizon: int, *, fps: float | None = None) -> int:
'Return the token bound after converting a source chunk to the codec clock.'
source_fps = self._source_fps(fps)
internal = resampled_horizon(
horizon, source_hz=source_fps, target_hz=self.action_space.target_hz
)
return internal * self._encoded_action_dim
def grammar(self, horizon: int, *, fps: float | None = None):
'Build a grammar for the codec steps implied by source chunk geometry.'
source_fps = self._source_fps(fps)
internal = resampled_horizon(
horizon, source_hz=source_fps, target_hz=self.action_space.target_hz
)
return self.action_tokenizer.grammar(preset=self.action_space.name, horizon=internal)
def __call__(
self,
action: torch.Tensor | np.ndarray,
*,
fps: float | None = None,
current_state: torch.Tensor | np.ndarray | None = None,
previous_first_order: Mapping[str, torch.Tensor | np.ndarray] | None = None,
encode_kwargs: dict | None = None,
) -> ActionIdxBatch:
kwargs = self._check_encode_kwargs(encode_kwargs)
direct = {
"fps": fps,
"current_state": current_state,
"previous_first_order": previous_first_order,
}
duplicates = sorted(
key for key, value in direct.items() if value is not None and key in kwargs
)
if duplicates:
raise ValueError(f"encode arguments provided twice: {duplicates}")
kwargs.update({key: value for key, value in direct.items() if value is not None})
reference = kwargs.get("current_state")
if self.action_space.input == "absolute" and reference is None:
raise ValueError(
f"action space {self.action_space.name!r} is absolute; encode requires "
"current_state with shape [B, D] to preserve the first transition"
)
if self.action_space.input == "delta" and reference is not None:
raise ValueError(
f"action space {self.action_space.name!r} is delta; omit current_state"
)
values = torch.as_tensor(action)
if values.ndim == 2:
values = values.unsqueeze(0)
return self._encode_action_indices(values, encode_kwargs=kwargs)
encode = __call__
def decode(
self,
action_token_ids: Sequence[TokenIds],
*,
fps: float | None = None,
current_state: torch.Tensor | np.ndarray | None = None,
previous_first_order: Mapping[str, torch.Tensor | np.ndarray] | None = None,
) -> torch.Tensor:
'Decode a batch, deriving sequence length from tokens and output rate.'
return self.decode_token_ids_to_actions_batch(
action_token_ids,
action_dim=self.action_dim,
decode_kwargs={
"fps": fps,
"current_state": current_state,
"previous_first_order": previous_first_order,
},
)
def decode_token_ids_to_actions_batch(
self,
action_token_ids: Sequence[TokenIds],
*,
time_horizon: int | None = None,
action_dim: int | None = None,
decode_kwargs: dict | None = None,
) -> torch.Tensor:
'Decode integer rows to ``(B,T,D)``; the native router validates ranges.'
if native_kernel("route_decode_batch") is None:
return super().decode_token_ids_to_actions_batch(
action_token_ids,
time_horizon=time_horizon,
action_dim=action_dim,
decode_kwargs=decode_kwargs,
)
actions, _ = self._decode_action_indices(
_numpy_token_rows(action_token_ids),
time_horizon=time_horizon,
action_dim=action_dim,
decode_kwargs=decode_kwargs,
)
return torch.from_numpy(actions)
def decode_prefix(
self,
action_token_ids,
*,
horizon: int,
steps: int,
fps: float | None = None,
current_state=None,
previous_first_order=None,
) -> torch.Tensor:
'Decode a covered streaming prefix to ``(B,steps,D)`` on the 15 Hz grid.'
if self.action_space.input == "absolute" and current_state is None:
raise ValueError("absolute prefix decode requires current_state")
if self.action_space.input == "delta" and current_state is not None:
raise ValueError("delta prefix decode must omit current_state")
decoded = self.action_tokenizer.decode_prefix_batch(
[self._checked_decode_row(row).tolist() for row in action_token_ids],
preset=self.action_space.name,
horizon=horizon,
steps=steps,
fps=float(self._source_fps(fps)),
current_state=_numpy(current_state),
previous_first_order=self._numpy_boundary(previous_first_order),
)
return torch.from_numpy(self._restore_source_layout(decoded))
def final_first_order(
self,
action_token_ids: Sequence[TokenIds],
*,
previous_first_order: Mapping[str, torch.Tensor | np.ndarray],
executed_steps: int | None = None,
) -> dict[str, np.ndarray]:
'Return reconstructed component states after a canonical-step prefix.'
return self.action_tokenizer.final_first_order(
[self._checked_decode_row(row).tolist() for row in action_token_ids],
preset=self.action_space.name,
previous_first_order=self._numpy_boundary(previous_first_order),
executed_steps=executed_steps,
)
def _encode_action_indices(
self, action: torch.Tensor, encode_kwargs: dict | None = None
) -> ActionIdxBatch:
kwargs = dict(encode_kwargs or {})
fps = float(self._source_fps(kwargs.get("fps")))
current_state = _numpy(kwargs.get("current_state"))
values = _numpy(action)
if values.ndim != 3 or values.shape[2] != self.action_space.input_dim:
raise ValueError(
f"actions must have shape [B, T_source, {self.action_space.input_dim}]"
)
return self.action_tokenizer.encode_batch(
values,
preset=self.action_space.name,
fps=fps,
current_state=current_state,
previous_first_order=self._numpy_boundary(kwargs.get("previous_first_order")),
)
def _decode_action_indices(
self,
action_indices: ActionIdxBatch,
*,
time_horizon: int | None = None,
action_dim: int | None = None,
decode_kwargs: dict | None = None,
) -> tuple[np.ndarray, list[set[str]]]:
if time_horizon is not None:
raise ValueError(
"ActionCodec2 decode infers sequence length from tokens; "
"do not provide time_horizon"
)
if action_dim is not None and int(action_dim) != self.action_dim:
raise ValueError(
f"action_dim={action_dim} disagrees with action space "
f"{self.action_space.name!r} (D={self.action_dim})"
)
kwargs = dict(decode_kwargs or {})
unknown = sorted(set(kwargs) - self.encode_kwargs)
if unknown:
raise ValueError(f"ActionCodec2Tokenizer does not consume decode_kwargs {unknown}")
raw_fps = kwargs.get("fps")
fps = float(self._source_fps(raw_fps))
current_state = _numpy(kwargs.get("current_state"))
if self.action_space.input == "absolute" and current_state is None:
raise ValueError(
f"action space {self.action_space.name!r} is absolute; decode requires "
"current_state with shape [B, D] to restore its external reference"
)
if self.action_space.input == "delta" and current_state is not None:
raise ValueError(
f"action space {self.action_space.name!r} is delta; omit current_state "
"to decode in the registered input semantics"
)
decoded = self.action_tokenizer.decode_batch(
action_indices,
preset=self.action_space.name,
fps=fps,
current_state=current_state,
previous_first_order=self._numpy_boundary(kwargs.get("previous_first_order")),
)
actions = self._restore_source_layout(decoded)
return actions, [set() for _ in range(actions.shape[0])]
def _restore_source_layout(self, decoded: Mapping[str, np.ndarray]) -> np.ndarray:
return restore_source_layout(self.action_space, decoded)
|