| |
| """Build APE weights on the frozen Florence-Bridge valid-anchor support.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from collections import OrderedDict, defaultdict |
| from pathlib import Path |
|
|
| import numpy as np |
| import pyarrow.parquet as parquet |
| import torch |
| from torch.utils.data import DataLoader, TensorDataset |
|
|
| from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest |
|
|
| from external_baselines.keyframe_focused_adapted.keyframe_cache import ( |
| ActionCopycat, |
| history_target, |
| q99_normalize, |
| ) |
|
|
|
|
| def episode_path(data_root: Path, dataset_name: str, trajectory_id: int) -> Path: |
| return ( |
| data_root |
| / dataset_name |
| / "data" |
| / f"chunk-{trajectory_id // 1000:03d}" |
| / f"episode_{trajectory_id:06d}.parquet" |
| ) |
|
|
|
|
| def read_actions( |
| data_root: Path, |
| dataset_name: str, |
| trajectory_id: int, |
| action_dim: int | None, |
| ) -> np.ndarray: |
| path = episode_path(data_root, dataset_name, trajectory_id) |
| values = parquet.read_table(path, columns=["action"])["action"].to_pylist() |
| actions = np.asarray(values, dtype=np.float32) |
| if actions.ndim != 2 or (action_dim is not None and actions.shape[1] != action_dim): |
| raise RuntimeError(f"unexpected action shape {actions.shape} in {path}") |
| return actions |
|
|
|
|
| def resolve_statistics_key( |
| dataset_name: str, |
| statistics_payload: dict, |
| explicit_mapping: dict[str, str], |
| ) -> str: |
| """Resolve Bridge aliases and arbitrary RoboTwin dataset statistics.""" |
|
|
| aliases = { |
| "bridge_orig_lerobot": "oxe_bridge", |
| "fractal20220817_data_lerobot": "oxe_rt1", |
| } |
| candidates = ( |
| explicit_mapping.get(dataset_name), |
| aliases.get(dataset_name), |
| dataset_name, |
| ) |
| for candidate in candidates: |
| if candidate and candidate in statistics_payload: |
| return candidate |
| if len(statistics_payload) == 1: |
| return next(iter(statistics_payload)) |
| raise KeyError( |
| f"no statistics entry for {dataset_name!r}; provide --statistics-key-map" |
| ) |
|
|
|
|
| def requested_support(manifest, args): |
| if args.scope == "full": |
| return { |
| dataset.name: OrderedDict( |
| (int(episode.trajectory_id), tuple(int(x) for x in episode.anchors)) |
| for episode in dataset.episodes |
| ) |
| for dataset in manifest.datasets |
| } |
| selected = defaultdict(lambda: defaultdict(set)) |
| for draw_id in range(args.draw_start, args.draw_start + args.draw_count): |
| trace = manifest.sample(args.draw_seed, LogicalDraw(0, draw_id)) |
| selected[trace.dataset_name][int(trace.trajectory_id)].add(int(trace.base_index)) |
| return { |
| dataset.name: OrderedDict( |
| ( |
| int(episode.trajectory_id), |
| tuple(sorted(selected[dataset.name][int(episode.trajectory_id)])), |
| ) |
| for episode in dataset.episodes |
| if selected[dataset.name].get(int(episode.trajectory_id)) |
| ) |
| for dataset in manifest.datasets |
| } |
|
|
|
|
| def fitting_support(manifest, max_anchors: int | None): |
| result = {} |
| for dataset in manifest.datasets: |
| episodes = OrderedDict() |
| count = 0 |
| for episode in dataset.episodes: |
| anchors = tuple(int(x) for x in episode.anchors if int(x) >= 9) |
| if max_anchors is not None: |
| anchors = anchors[: max(0, max_anchors - count)] |
| if anchors: |
| episodes[int(episode.trajectory_id)] = anchors |
| count += len(anchors) |
| if max_anchors is not None and count >= max_anchors: |
| break |
| if count == 0: |
| raise RuntimeError(f"no copycat fitting anchors for {dataset.name}") |
| result[dataset.name] = episodes |
| return result |
|
|
|
|
| def examples_for_support( |
| data_root, |
| dataset_name, |
| support, |
| statistics, |
| *, |
| action_dim, |
| binary_action_indices, |
| history=9, |
| ): |
| inputs, targets = [], [] |
| for trajectory_id, anchors in support.items(): |
| actions = q99_normalize( |
| read_actions(data_root, dataset_name, trajectory_id, action_dim), |
| statistics, |
| binary_action_indices=binary_action_indices, |
| ) |
| for anchor in anchors: |
| x, y = history_target(actions, anchor, history=history) |
| inputs.append(x) |
| targets.append(y) |
| return np.asarray(inputs, dtype=np.float32), np.asarray(targets, dtype=np.float32) |
|
|
|
|
| def train_copycat(inputs, targets, *, action_dim, epochs, device, seed=0): |
| generator = torch.Generator().manual_seed(seed) |
| permutation = torch.randperm(len(inputs), generator=generator) |
| split = int(len(inputs) * 0.8) |
| train_indices = permutation[:split] |
| val_indices = permutation[split:] |
| x = torch.from_numpy(inputs) |
| y = torch.from_numpy(targets) |
| train = TensorDataset(x[train_indices], y[train_indices]) |
| val = TensorDataset(x[val_indices], y[val_indices]) |
| train_loader = DataLoader( |
| train, batch_size=512, shuffle=True, generator=generator, num_workers=0 |
| ) |
| val_loader = DataLoader(val, batch_size=512, shuffle=False, num_workers=0) |
| torch.manual_seed(seed) |
| model = ActionCopycat(action_dim=action_dim).to(device) |
| optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999)) |
| history = [] |
| for epoch in range(1, epochs + 1): |
| model.train() |
| losses = [] |
| for batch_x, batch_y in train_loader: |
| batch_x = batch_x.to(device) |
| batch_y = batch_y.to(device) |
| optimizer.zero_grad(set_to_none=True) |
| loss = (model(batch_x) - batch_y).square().mean(dim=-1).mean() |
| loss.backward() |
| optimizer.step() |
| losses.append(float(loss.detach().cpu())) |
| model.eval() |
| val_losses = [] |
| with torch.no_grad(): |
| for batch_x, batch_y in val_loader: |
| batch_x = batch_x.to(device) |
| batch_y = batch_y.to(device) |
| loss = (model(batch_x) - batch_y).square().mean(dim=-1) |
| val_losses.extend(loss.cpu().tolist()) |
| history.append( |
| { |
| "epoch": epoch, |
| "train_loss": float(np.mean(losses)), |
| "validation_loss": float(np.mean(val_losses)), |
| } |
| ) |
| print(json.dumps(history[-1], sort_keys=True), flush=True) |
| if epoch % max(1, epochs // 3) == 0: |
| for group in optimizer.param_groups: |
| group["lr"] *= 0.1 |
| return model, history |
|
|
|
|
| def prediction_errors(model, inputs, targets, device): |
| errors = [] |
| loader = DataLoader( |
| TensorDataset(torch.from_numpy(inputs), torch.from_numpy(targets)), |
| batch_size=4096, |
| shuffle=False, |
| num_workers=0, |
| ) |
| model.eval() |
| with torch.no_grad(): |
| for batch_x, batch_y in loader: |
| batch_x = batch_x.to(device) |
| batch_y = batch_y.to(device) |
| errors.append( |
| (model(batch_x) - batch_y).square().mean(dim=-1).cpu().numpy() |
| ) |
| return np.concatenate(errors) |
|
|
|
|
| def score_support( |
| model, |
| data_root, |
| dataset_name, |
| support, |
| statistics, |
| device, |
| *, |
| action_dim, |
| binary_action_indices, |
| ): |
| trajectory_ids, offsets, anchors_out, ape_out = [], [0], [], [] |
| model.eval() |
| with torch.no_grad(): |
| for trajectory_id, anchors in support.items(): |
| actions = q99_normalize( |
| read_actions(data_root, dataset_name, trajectory_id, action_dim), |
| statistics, |
| binary_action_indices=binary_action_indices, |
| ) |
| batch_x, batch_y = [], [] |
| for anchor in anchors: |
| x, y = history_target(actions, anchor, history=9) |
| batch_x.append(x) |
| batch_y.append(y) |
| x = torch.from_numpy(np.asarray(batch_x, dtype=np.float32)).to(device) |
| y = torch.from_numpy(np.asarray(batch_y, dtype=np.float32)).to(device) |
| ape = (model(x) - y).square().mean(dim=-1).cpu().numpy() |
| trajectory_ids.append(trajectory_id) |
| anchors_out.extend(anchors) |
| ape_out.extend(ape.tolist()) |
| offsets.append(len(anchors_out)) |
| return { |
| "trajectory_ids": np.asarray(trajectory_ids, dtype=np.int64), |
| "offsets": np.asarray(offsets, dtype=np.int64), |
| "anchors": np.asarray(anchors_out, dtype=np.int32), |
| "ape": np.asarray(ape_out, dtype=np.float32), |
| } |
|
|
|
|
| def verify_reference_overlap(reference, dataset_index, scored, weights): |
| prefix = f"d{dataset_index}_" |
| reference_trajectory_ids = reference[prefix + "trajectory_ids"] |
| reference_offsets = reference[prefix + "offsets"] |
| reference_anchors = reference[prefix + "anchors"] |
| reference_ape = reference[prefix + "ape"] |
| reference_weights = reference[prefix + "weights"] |
| new_trajectory_positions = { |
| int(trajectory_id): position |
| for position, trajectory_id in enumerate(scored["trajectory_ids"]) |
| } |
| checked = 0 |
| max_ape_difference = 0.0 |
| for reference_position, trajectory_id in enumerate(reference_trajectory_ids): |
| new_position = new_trajectory_positions.get(int(trajectory_id)) |
| if new_position is None: |
| raise RuntimeError( |
| f"full cache is missing reference trajectory {int(trajectory_id)}" |
| ) |
| reference_slice = slice( |
| int(reference_offsets[reference_position]), |
| int(reference_offsets[reference_position + 1]), |
| ) |
| new_slice = slice( |
| int(scored["offsets"][new_position]), |
| int(scored["offsets"][new_position + 1]), |
| ) |
| new_anchors = scored["anchors"][new_slice] |
| positions = np.searchsorted(new_anchors, reference_anchors[reference_slice]) |
| if np.any(positions >= len(new_anchors)) or not np.array_equal( |
| new_anchors[positions], reference_anchors[reference_slice] |
| ): |
| raise RuntimeError( |
| f"full cache does not contain the reference anchor support for " |
| f"trajectory {int(trajectory_id)}" |
| ) |
| new_ape = scored["ape"][new_slice][positions] |
| ape_difference = np.abs(new_ape - reference_ape[reference_slice]) |
| if len(ape_difference): |
| max_ape_difference = max(max_ape_difference, float(ape_difference.max())) |
| if not np.array_equal( |
| weights[new_slice][positions], reference_weights[reference_slice] |
| ): |
| raise RuntimeError( |
| f"full cache changed reference weights for trajectory " |
| f"{int(trajectory_id)}" |
| ) |
| checked += len(positions) |
| if max_ape_difference > 1e-6: |
| raise RuntimeError( |
| f"full cache changed reference APE scores: max_abs_diff={max_ape_difference}" |
| ) |
| return { |
| "checked_anchor_count": checked, |
| "weights_exact_match": True, |
| "max_ape_abs_difference": max_ape_difference, |
| } |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--manifest", type=Path, required=True) |
| parser.add_argument("--data-root", type=Path, required=True) |
| parser.add_argument("--statistics", type=Path, required=True) |
| parser.add_argument( |
| "--statistics-key-map", |
| type=Path, |
| help="Optional JSON mapping from manifest dataset name to statistics key.", |
| ) |
| parser.add_argument( |
| "--action-dim", |
| type=int, |
| choices=(7, 14), |
| help="Expected action dimension; inferred from statistics when omitted.", |
| ) |
| parser.add_argument( |
| "--binary-action-indices", |
| help=( |
| "Comma-separated zero-based binary action indices. Defaults to index 6 " |
| "for 7D Bridge and none for 14D RoboTwin absolute qpos." |
| ), |
| ) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--scope", choices=("full", "draw-range"), default="full") |
| parser.add_argument("--draw-seed", type=int, default=117465) |
| parser.add_argument("--draw-start", type=int, default=1_280_000) |
| parser.add_argument("--draw-count", type=int, default=128) |
| parser.add_argument("--fit-max-anchors", type=int) |
| parser.add_argument("--epochs", type=int, default=300) |
| parser.add_argument("--device", default="cuda") |
| parser.add_argument("--reference-cache", type=Path) |
| args = parser.parse_args() |
|
|
| manifest = ValidAnchorManifest.load(args.manifest) |
| statistics_payload = json.loads(args.statistics.read_text(encoding="utf-8")) |
| statistics_key_map = ( |
| json.loads(args.statistics_key_map.read_text(encoding="utf-8")) |
| if args.statistics_key_map is not None |
| else {} |
| ) |
| reference = None |
| reference_metadata = None |
| if args.reference_cache is not None: |
| if args.scope != "full": |
| raise ValueError("--reference-cache requires --scope full") |
| reference = np.load(args.reference_cache, allow_pickle=False) |
| reference_metadata = json.loads(str(reference["metadata_json"])) |
| if int(reference_metadata["copycat_epochs"]) != args.epochs: |
| raise RuntimeError("reference cache copycat epoch count does not match") |
| scoring = requested_support(manifest, args) |
| fitting = fitting_support(manifest, args.fit_max_anchors) |
| arrays = {} |
| dataset_reports = [] |
| for dataset_index, dataset in enumerate(manifest.datasets): |
| dataset_name = dataset.name |
| statistics_key = resolve_statistics_key( |
| dataset_name, statistics_payload, statistics_key_map |
| ) |
| action_statistics = statistics_payload[statistics_key]["action"] |
| inferred_action_dim = len(action_statistics["q99"]) |
| action_dim = args.action_dim or inferred_action_dim |
| if inferred_action_dim != action_dim: |
| raise RuntimeError( |
| f"statistics action dimension {inferred_action_dim} != {action_dim} " |
| f"for {dataset_name}" |
| ) |
| if args.binary_action_indices is None: |
| binary_action_indices = (6,) if action_dim == 7 else () |
| else: |
| binary_action_indices = tuple( |
| int(value) |
| for value in args.binary_action_indices.split(",") |
| if value.strip() |
| ) |
| inputs, targets = examples_for_support( |
| args.data_root, |
| dataset_name, |
| fitting[dataset_name], |
| action_statistics, |
| action_dim=action_dim, |
| binary_action_indices=binary_action_indices, |
| ) |
| print( |
| f"training copycat for {dataset_name}: {len(inputs)} anchors", |
| flush=True, |
| ) |
| model, train_history = train_copycat( |
| inputs, |
| targets, |
| action_dim=action_dim, |
| epochs=args.epochs, |
| device=torch.device(args.device), |
| ) |
| fitted_ape = prediction_errors( |
| model, inputs, targets, torch.device(args.device) |
| ) |
| scored = score_support( |
| model, |
| args.data_root, |
| dataset_name, |
| scoring[dataset_name], |
| action_statistics, |
| torch.device(args.device), |
| action_dim=action_dim, |
| binary_action_indices=binary_action_indices, |
| ) |
| if reference_metadata is None: |
| threshold_values = scored["ape"] if args.scope == "full" else fitted_ape |
| threshold = float(np.quantile(threshold_values, 0.9)) |
| else: |
| reference_report = reference_metadata["datasets"][dataset_index] |
| if reference_report["name"] != dataset_name: |
| raise RuntimeError("reference cache dataset order does not match") |
| if int(reference_report["fit_anchor_count"]) != len(inputs): |
| raise RuntimeError("reference cache fitting support does not match") |
| threshold = float(reference_report["ape_threshold_90th_percentile"]) |
| weights = np.where(scored["ape"] >= threshold, 5.0, 1.0).astype(np.float32) |
| overlap_report = None |
| if reference is not None: |
| overlap_report = verify_reference_overlap( |
| reference, dataset_index, scored, weights |
| ) |
| for key, value in scored.items(): |
| arrays[f"d{dataset_index}_{key}"] = value |
| arrays[f"d{dataset_index}_weights"] = weights |
| dataset_reports.append( |
| { |
| "name": dataset_name, |
| "statistics_key": statistics_key, |
| "action_dim": action_dim, |
| "binary_action_indices": list(binary_action_indices), |
| "fit_anchor_count": len(inputs), |
| "scored_anchor_count": len(scored["anchors"]), |
| "ape_threshold_90th_percentile": threshold, |
| "upweighted_count": int((weights > 1).sum()), |
| "copycat_training": train_history, |
| "reference_overlap": overlap_report, |
| } |
| ) |
|
|
| metadata = { |
| "schema_version": 1, |
| "method": "Keyframe-Focused-adapted", |
| "history_length": 9, |
| "copycat_hidden_units": [300], |
| "copycat_optimizer": {"name": "Adam", "lr": 1e-3, "betas": [0.9, 0.999]}, |
| "copycat_epochs": args.epochs, |
| "ape_metric": "mean squared error over normalized action dimensions", |
| "boundary_padding": "repeat earliest physical action", |
| "threshold_scope": "within dataset source", |
| "threshold_percentile": 90, |
| "high_weight": 5.0, |
| "low_weight": 1.0, |
| "cache_scope": args.scope, |
| "threshold_source": ( |
| "reference_cache" if args.reference_cache is not None else "computed" |
| ), |
| "dataset_names": [dataset.name for dataset in manifest.datasets], |
| "datasets": dataset_reports, |
| } |
| arrays["metadata_json"] = np.asarray(json.dumps(metadata, sort_keys=True)) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| temporary = args.output.with_suffix(args.output.suffix + ".tmp") |
| with temporary.open("wb") as destination: |
| np.savez_compressed(destination, **arrays) |
| temporary.replace(args.output) |
| print(f"wrote {args.output}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
