File size: 5,680 Bytes
5cd1287 | 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 | """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]),
)
|