| """APE model and compact cache used by the isolated Keyframe baseline.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
|
|
| class ActionCopycat(nn.Module): |
| """Official one-hidden-layer action-only copycat architecture.""" |
|
|
| def __init__(self, action_dim: int = 7, history: int = 9, hidden: int = 300): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(action_dim * history, hidden), |
| nn.LeakyReLU(0.2, inplace=True), |
| nn.Linear(hidden, action_dim), |
| ) |
|
|
| def forward(self, action_history: torch.Tensor) -> torch.Tensor: |
| return self.net(action_history) |
|
|
|
|
| def q99_normalize( |
| actions: np.ndarray, |
| statistics: dict, |
| *, |
| binary_action_indices: tuple[int, ...] = (), |
| ) -> np.ndarray: |
| """Apply q01/q99 normalization with an explicit binary-action contract. |
| |
| A zero percentile span is not sufficient evidence that an action dimension is |
| binary. This matters for RoboTwin absolute qpos, where a joint can be constant |
| in a dataset shard while remaining a continuous joint. Such continuous, |
| zero-span dimensions are mapped to zero; only explicitly declared indices use |
| the binary threshold. |
| """ |
|
|
| actions = np.asarray(actions, dtype=np.float32) |
| q01 = np.asarray(statistics["q01"], dtype=np.float32) |
| q99 = np.asarray(statistics["q99"], dtype=np.float32) |
| if actions.shape[-1] != len(q01) or q01.shape != q99.shape: |
| raise ValueError( |
| f"action/statistics shape mismatch: actions={actions.shape}, " |
| f"q01={q01.shape}, q99={q99.shape}" |
| ) |
| binary = np.zeros(len(q01), dtype=bool) |
| for index in binary_action_indices: |
| if index < 0 or index >= len(q01): |
| raise ValueError(f"binary action index {index} is out of range") |
| binary[index] = True |
| continuous = ~binary |
| scaled = continuous & (q01 != q99) |
| constant = continuous & (q01 == q99) |
| normalized = np.zeros_like(actions, dtype=np.float32) |
| normalized[..., scaled] = ( |
| 2.0 |
| * (actions[..., scaled] - q01[scaled]) |
| / (q99[scaled] - q01[scaled]) |
| - 1.0 |
| ) |
| normalized[..., scaled] = np.clip( |
| normalized[..., scaled], -2.2, 2.2 |
| ) |
| normalized[..., constant] = 0.0 |
| normalized[..., binary] = ( |
| actions[..., binary] > 0.5 |
| ).astype(np.float32) |
| return normalized |
|
|
|
|
| def history_target(actions: np.ndarray, anchor: int, history: int = 9): |
| """Repeat the earliest physical action at a trajectory boundary.""" |
|
|
| indices = np.arange(anchor - history, anchor, dtype=np.int64) |
| indices = np.clip(indices, 0, len(actions) - 1) |
| return actions[indices].reshape(-1), actions[anchor] |
|
|
|
|
| @dataclass(frozen=True) |
| class CacheRecord: |
| ape: float |
| weight: float |
|
|
|
|
| class KeyframeWeightCache: |
| """Read a compact NPZ cache without altering sampling.""" |
|
|
| def __init__(self, path: str | Path): |
| self.path = Path(path) |
| payload = np.load(self.path, allow_pickle=False) |
| self.metadata = json.loads(str(payload["metadata_json"].item())) |
| if self.metadata.get("method") != "Keyframe-Focused-adapted": |
| raise RuntimeError("unexpected keyframe cache method") |
| self._datasets = {} |
| for index, name in enumerate(self.metadata["dataset_names"]): |
| trajectory_ids = payload[f"d{index}_trajectory_ids"].astype(np.int64) |
| offsets = payload[f"d{index}_offsets"].astype(np.int64) |
| anchors = payload[f"d{index}_anchors"].astype(np.int64) |
| ape = payload[f"d{index}_ape"].astype(np.float32) |
| weights = payload[f"d{index}_weights"].astype(np.float32) |
| if not (len(anchors) == len(ape) == len(weights)): |
| raise RuntimeError(f"cache array lengths differ for {name}") |
| if len(offsets) != len(trajectory_ids) + 1 or int(offsets[-1]) != len(anchors): |
| raise RuntimeError(f"cache offsets are invalid for {name}") |
| if np.any(~np.isfinite(ape)) or np.any(~np.isfinite(weights)): |
| raise RuntimeError(f"cache contains non-finite values for {name}") |
| if np.any(weights <= 0): |
| raise RuntimeError(f"cache contains non-positive weights for {name}") |
| self._datasets[name] = { |
| "trajectory_to_row": { |
| int(trajectory_id): row |
| for row, trajectory_id in enumerate(trajectory_ids) |
| }, |
| "offsets": offsets, |
| "anchors": anchors, |
| "ape": ape, |
| "weights": weights, |
| } |
|
|
| def lookup(self, dataset_name: str, trajectory_id: int, base_index: int) -> CacheRecord: |
| try: |
| dataset = self._datasets[dataset_name] |
| row = dataset["trajectory_to_row"][int(trajectory_id)] |
| except KeyError as exc: |
| raise KeyError( |
| f"keyframe cache has no trajectory {(dataset_name, trajectory_id)}" |
| ) from exc |
| start, end = dataset["offsets"][row : row + 2] |
| local = dataset["anchors"][start:end] |
| position = int(np.searchsorted(local, int(base_index))) |
| if position >= len(local) or int(local[position]) != int(base_index): |
| raise KeyError( |
| f"keyframe cache has no anchor {(dataset_name, trajectory_id, base_index)}" |
| ) |
| absolute = int(start + position) |
| return CacheRecord( |
| ape=float(dataset["ape"][absolute]), |
| weight=float(dataset["weights"][absolute]), |
| ) |
|
|
