File size: 7,533 Bytes
fee0e43 | 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 | """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,
}
|