| """Profile-level ActionCodec2Sparse fitting and default profile derivation.""" |
|
|
| from __future__ import annotations |
|
|
| import time |
| from collections.abc import Iterable, Sequence |
|
|
| import numpy as np |
|
|
| from ..configuration import ActionCodec2Config, ConfigError |
| from ..physical_codec import PhysicalCodec |
| from ..primitives import PrimitiveStore |
| from ..training import _quantize_trajectories, build_primitives |
| from ..trajectory import TrajectoryLike |
| from .layout import SparseLayout, training_chunks |
| from .processor import ActionCodec2SparseProcessor |
| from .training import fit_sparse_vocabulary |
|
|
| DEFAULT_VOCAB_SIZE = 4096 |
| DEFAULT_CHUNK_HORIZON = 32 |
| DEFAULT_STATIC_RUN_CAP = 8 |
|
|
|
|
| def sparse_profile_config( |
| config: ActionCodec2Config, |
| *, |
| vocab_size: int = DEFAULT_VOCAB_SIZE, |
| chunk_horizon: int = DEFAULT_CHUNK_HORIZON, |
| static_run_cap: int = DEFAULT_STATIC_RUN_CAP, |
| ) -> ActionCodec2Config: |
| """Derive an ActionCodec2Sparse profile from a first- or second-order ActionCodec2 profile.""" |
| if config.train is None: |
| raise ConfigError("a sparse profile is derived from a config with a train section") |
| payload = config.to_dict() |
| payload["set_bpe"]["scheme"] = "sparse_set_bpe" |
| for component in payload["physical"]["components"].values(): |
| if component["space"] == "so3" and component["primitive_order"] == 1: |
| component["primitive_selection"] = "momentum_lazy" |
| layout = SparseLayout.from_config(config, build_primitives(config)) |
| payload["train"]["bpe"] = { |
| "spatial_vocab_size": sum(layout.sparse_bins), |
| "vocab_size": int(vocab_size), |
| } |
| payload["train"]["sparse"] = { |
| "chunk_horizon": int(chunk_horizon), |
| "static_run_cap": int(static_run_cap), |
| } |
| return ActionCodec2Config.from_dict(payload) |
|
|
|
|
| class SparseProfileTrainer: |
| """Fit one ActionCodec2Sparse profile on complete canonical trajectories.""" |
|
|
| def __init__(self, config: ActionCodec2Config) -> None: |
| if config.set_bpe.scheme != "sparse_set_bpe" or config.train is None: |
| raise ConfigError("SparseProfileTrainer needs a sparse_set_bpe config with train") |
| self.config = config |
|
|
| def fit( |
| self, |
| trajectories: Iterable[TrajectoryLike], |
| *, |
| backend: str = "auto", |
| threads: int = 0, |
| batch_size: int = 256, |
| progress: bool = False, |
| max_cell_bytes: int | None = None, |
| ) -> tuple[ActionCodec2SparseProcessor, dict]: |
| """Quantize then fit; see :meth:`fit_cells` for the token-layer stage.""" |
| started = time.perf_counter() |
| primitives = build_primitives(self.config) |
| cells, actions, active_values, arities = _quantize_trajectories( |
| trajectories, |
| PhysicalCodec(self.config, primitives), |
| batch_size=batch_size, |
| max_cell_bytes=max_cell_bytes, |
| ) |
| quantized = time.perf_counter() |
| processor, report = self.fit_cells( |
| cells, primitives=primitives, backend=backend, threads=threads, progress=progress |
| ) |
| report.update( |
| trajectories=len(cells), |
| actions=actions, |
| active_values=active_values, |
| trajectory_arities={str(k): v for k, v in arities.items()}, |
| ) |
| report["seconds"]["quantize"] = quantized - started |
| return processor, report |
|
|
| def fit_cells( |
| self, |
| cells: Sequence[np.ndarray], |
| *, |
| primitives: PrimitiveStore | None = None, |
| backend: str = "auto", |
| threads: int = 0, |
| progress: bool = False, |
| ) -> tuple[ActionCodec2SparseProcessor, dict]: |
| """Fit the token layer on quantized cell trajectories ``[(T_i, D)]`` (``-1`` inactive).""" |
| config = self.config |
| primitives = build_primitives(config) if primitives is None else primitives |
| layout = SparseLayout.from_config(config, primitives) |
| options = config.train.sparse |
| chunks = training_chunks( |
| cells, layout, horizon=options.chunk_horizon, static_run_cap=options.static_run_cap |
| ) |
| vocabulary, report = fit_sparse_vocabulary( |
| layout, |
| chunks, |
| spatial_vocab_size=config.train.bpe.spatial_vocab_size, |
| vocab_size=config.train.bpe.vocab_size, |
| lookahead=config.set_bpe.lookahead, |
| backend=backend, |
| threads=threads, |
| progress=progress, |
| ) |
| report["corpus"] = { |
| "trajectories": len(cells), |
| "source_frames": int(sum(len(c) for c in cells)), |
| "training_frames": report["frames"], |
| "chunk_horizon": options.chunk_horizon, |
| "static_run_cap": options.static_run_cap, |
| } |
| report["tokenization_scheme"] = "sparse_set_bpe" |
| return ActionCodec2SparseProcessor(config, primitives, vocabulary), report |
|
|
