File size: 3,630 Bytes
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 | """Explicit tokenizer pretraining and cleaning shared by both token layers."""
from __future__ import annotations
from pathlib import Path
from typing import Any
import numpy as np
from .fitting import ProfileFitter
class TrainingMixin:
def preprocess(
self,
values: np.ndarray,
*,
profile: str | None = None,
gripper_column: int | None = None,
**kwargs: Any,
) -> np.ndarray:
"""Clean a canonical ``(T,D_active)`` profile trajectory for explicit tokenizer pretraining."""
if profile is None:
profiles = {component.space for component in self.action_space.components}
if len(profiles) != 1:
raise ValueError("profile is required for layouts combining joint and EEF")
profile = profiles.pop()
if profile not in self.preprocessing:
raise ValueError(f"no recorded preprocessing for profile {profile!r}")
values = np.asarray(values)
if "ease_limits" not in kwargs and values.ndim == 2:
kwargs["ease_limits"] = self._joint_ease_limits(
profile, values.shape[1], gripper_column, kwargs.get("active_dims")
)
return self.preprocessing[profile].apply(values, gripper_column=gripper_column, **kwargs)
def _joint_ease_limits(self, profile, width, gripper_column, active_dims):
processor = self.action_tokenizer.profiles[profile].processor
component = processor.config.physical.components.get("generalized_joint_position")
if profile != "joint" or component is None or component.primitive_order != 1:
return None
domain = component.constraints["order_1"]
limits = np.full(processor.action_dim, np.inf)
limits[list(component.dims)] = np.minimum(np.abs(domain.minimum), np.abs(domain.maximum))
dims = tuple(range(width)) if active_dims is None else tuple(active_dims)
limits = limits[list(dims)]
if gripper_column is not None:
limits[gripper_column] = np.inf
return limits
def fit(
self,
episodes,
*,
fps: float | None = None,
current_state=None,
previous_first_order=None,
bootstrap_context: bool = False,
apply_preprocessing: bool = False,
preprocessing=None,
backend: str = "auto",
threads: int = 0,
batch_size: int = 256,
progress: bool = False,
max_cell_bytes: int | None = None,
work_dir: str | Path | None = None,
) -> TrainingMixin:
"""Fit vocabularies from episodes in declared source layouts."""
fitter = ProfileFitter(self, bootstrap_context=bootstrap_context,
apply_preprocessing=apply_preprocessing, preprocessing=preprocessing)
router, report = fitter.fit(
episodes,
work_dir=work_dir,
metadata=dict(
fps=fps, current_state=current_state, previous_first_order=previous_first_order
),
backend=backend,
threads=threads,
batch_size=batch_size,
progress=progress,
max_cell_bytes=max_cell_bytes,
)
replacement = type(self)(router, action_space=self.action_space.name)
replacement._fitted_profiles = self.fitted_profiles | report["profiles"].keys()
replacement.training_report = report
replacement.preprocessing = dict(fitter.preprocessing)
replacement._warned_default_fps = self._warned_default_fps
self.__dict__.update(replacement.__dict__)
return self
|