| """Replayable, bounded source staging for the public ActionCodec2.fit interface.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Mapping |
| from contextlib import ExitStack |
| from pathlib import Path |
| from tempfile import TemporaryDirectory |
|
|
| import numpy as np |
| import torch |
|
|
| from ..physical.boundary import bootstrap_second_order_trajectory |
| from ..routing.routing import ActionCodec2Router |
| from ..frontend.tokenizer import _numpy |
| from .training import ProfileTrainer |
| from ..physical.trajectory import ActiveTrajectory |
|
|
|
|
| class TrajectorySpool: |
| """Stage canonical episodes in two sequential files, without retaining source data. |
| |
| Args: |
| directory: Task-owned scratch directory, removed by the caller. |
| """ |
|
|
| def __init__(self, directory): |
| self.directory = Path(directory) |
| self.counts = {} |
| self._stack = ExitStack() |
| self._writers = {} |
|
|
| def __enter__(self): |
| return self |
|
|
| def __exit__(self, *exc): |
| return self._stack.__exit__(*exc) |
|
|
| def append(self, profile, trajectory): |
| if profile not in self._writers: |
| self._writers[profile] = self._stack.enter_context( |
| (self.directory / f"{profile}.npy").open("wb") |
| ) |
| stream = self._writers[profile] |
| for value in ( |
| trajectory.values, |
| trajectory.active_dims, |
| [] |
| if trajectory.previous_first_order is None |
| else trajectory.previous_first_order, |
| ): |
| np.save(stream, np.asarray(value), allow_pickle=False) |
| self.counts[profile] = self.counts.get(profile, 0) + 1 |
|
|
| def trajectories(self, profile): |
| self._writers[profile].flush() |
| with (self.directory / f"{profile}.npy").open("rb") as stream: |
| for _ in range(self.counts[profile]): |
| values = np.load(stream, allow_pickle=False) |
| dims = tuple(int(d) for d in np.load(stream, allow_pickle=False)) |
| previous = np.load(stream, allow_pickle=False) |
| yield ActiveTrajectory( |
| values, |
| dims, |
| previous_first_order=previous if previous.size else None, |
| ) |
|
|
|
|
| class ActionCodec2Fitter: |
| """Route user layouts and fit replacement profiles without mutating the caller. |
| |
| Args: |
| codec: Public processor supplying action-space definitions and profile configs. |
| bootstrap_context: Explicitly consume two canonical context frames for second order. |
| """ |
|
|
| def __init__(self, codec, *, bootstrap_context=False): |
| if not isinstance(bootstrap_context, bool): |
| raise TypeError("bootstrap_context must be a boolean") |
| self.codec = codec |
| self.router = codec.action_tokenizer |
| self.bootstrap_context = bootstrap_context |
|
|
| def _route(self, record): |
| unknown = set(record) - { |
| "actions", |
| "action_space", |
| "fps", |
| "current_state", |
| "previous_first_order", |
| } |
| if unknown: |
| raise ValueError(f"unknown training record fields: {sorted(unknown)}") |
| if "actions" not in record: |
| raise ValueError("training record requires an 'actions' array") |
| preset = self.router.preset( |
| record.get("action_space", self.codec.action_space.name) |
| ) |
| values = _numpy(torch.as_tensor(record["actions"])) |
| if values.ndim == 2: |
| values = values[None] |
| if values.ndim != 3 or values.shape[0] == 0 or values.shape[1] == 0: |
| raise ValueError( |
| "training actions must have nonempty shape (T,D) or (B,T,D)" |
| ) |
| current = _numpy(record.get("current_state")) |
| if preset.input == "absolute" and current is None: |
| raise ValueError( |
| f"absolute action space {preset.name!r} requires current_state(B,D)" |
| ) |
| if preset.input == "delta" and current is not None: |
| raise ValueError(f"delta action space {preset.name!r}: omit current_state") |
| previous = self.codec._numpy_boundary(record.get("previous_first_order")) |
| kwargs = dict( |
| preset=preset.name, |
| fps=self.codec._source_fps(record.get("fps"), preset), |
| current_state=current, |
| ) |
| if not self.bootstrap_context: |
| return self.router.route_trajectories( |
| values, previous_first_order=previous, **kwargs |
| ) |
| if previous is not None: |
| raise ValueError( |
| "bootstrap_context and previous_first_order are mutually exclusive" |
| ) |
| canonical = preset.canonicalize( |
| values, source_hz=kwargs["fps"], current_state=current |
| ) |
| routed = {} |
| for component in preset.components: |
| processor = self.router.profiles[component.space].processor |
| trajectories = canonical[component.name] |
| if processor.requires_previous_first_order: |
| trajectories = [ |
| bootstrap_second_order_trajectory(t, processor.config) |
| for t in trajectories |
| ] |
| routed.setdefault(component.space, []).extend(trajectories) |
| return routed |
|
|
| def fit(self, episodes, *, work_dir=None, metadata=None, **options): |
| """Fit observed profiles from arrays or records, replaying staged episodes per profile.""" |
| metadata = {k: v for k, v in (metadata or {}).items() if v is not None} |
| if isinstance(episodes, (np.ndarray, torch.Tensor, Mapping)): |
| episodes = (episodes,) |
| reports = {} |
| processors = { |
| name: runtime.processor for name, runtime in self.router.profiles.items() |
| } |
| records = 0 |
| with TemporaryDirectory(prefix="actioncodec2-fit-", dir=work_dir) as directory: |
| with TrajectorySpool(directory) as spool: |
| for episode in episodes: |
| record = ( |
| dict(episode) |
| if isinstance(episode, Mapping) |
| else {"actions": episode} |
| ) |
| duplicates = set(record).intersection(metadata) |
| if duplicates: |
| raise ValueError( |
| f"training arguments provided twice: {sorted(duplicates)}" |
| ) |
| record.update(metadata) |
| for profile, trajectories in self._route(record).items(): |
| for trajectory in trajectories: |
| spool.append(profile, trajectory) |
| records += 1 |
| if not spool.counts: |
| raise ValueError("fit requires at least one nonempty episode") |
| for name in spool.counts: |
| processors[name], reports[name] = ProfileTrainer( |
| processors[name].config |
| ).fit(spool.trajectories(name), **options) |
| router = ActionCodec2Router.from_processors( |
| processors, |
| token_budgets={ |
| name: runtime.spec.token_budget |
| for name, runtime in self.router.profiles.items() |
| }, |
| presets=list(self.router.action_spaces.values()), |
| ) |
| return router, { |
| "records": records, |
| "bootstrap_context": self.bootstrap_context, |
| "profiles": reports, |
| } |
|
|
