Scheduled Commit

scripts/eval/trex_ablation_eval.py

@@ -0,0 +1,593 @@
+"""Controlled ablation evaluation for the T-Rex Track-Force cascade.
+
+For each named checkpoint, on identical data:
+
+* fixed-tau forward losses (action / dynamics / track / force flow MSE);
+* open-loop chunk reconstruction: normalized 62-D action MSE of ``sample()``
+  against the ground-truth delta-base chunk (cascade and coarse-only);
+* Stage-1 self-attention mass of action/obs queries over key token groups;
+* Stage-2 force-transformer attention mass of action queries over
+  F6 / VQ-history / deform / coarse-memory tokens;
+* tactile sensitivity: |refine(real tactile) - refine(tactile masked)|.
+
+Usage:
+    python scripts/eval/trex_ablation_eval.py \
+        --dataset-root data/trex_mini_force \
+        --run full=checkpoints/ablate_full_3k/checkpoint-3000 \
+        --out ablation_eval.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+
+from groot.vla.data.schema import DatasetMetadata, EmbodimentTag
+from groot.vla.experiment.trex_eval_utils import TrexEpisode
+from groot.vla.model.trex_track_force.attention import TokenType
+from groot.vla.model.trex_track_force.dataset import (
+    DEFORM_VIDEO_KEYS,
+    eef62_delta_base,
+    nearest_timestamp_indices,
+    uniform_target_times,
+)
+from groot.vla.model.trex_track_force.force import (
+    ACTION_HORIZON,
+    FORCE_HISTORY_FRAMES,
+    FORCE_OFFSETS,
+    euler_flow_step,
+    pad_action_62_to_64,
+)
+from groot.vla.model.trex_track_force.runtime import TrexRuntimeStatistics
+from groot.vla.model.trex_track_force.track import TRACK_HORIZON
+from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA
+
+ACTION_RATE_HZ = 20.0
+TACTILE_RATE_HZ = 5.0
+VIDEO_RATE_HZ = 10.0
+AR_BLOCKS = 4
+VIDEO_FRAMES_PER_BLOCK = 8
+VIDEO_KEYS = ("video.head_left", "video.left_wrist", "video.right_wrist")
+
+
+def _column(episode: TrexEpisode, name: str, dtype=np.float32) -> np.ndarray:
+    values = episode.table.column(name).to_numpy(zero_copy_only=False)
+    return np.stack([np.asarray(row, dtype=dtype) for row in values])
+
+
+def _sample(timestamps, anchor, offsets, rate):
+    return nearest_timestamp_indices(
+        timestamps, uniform_target_times(anchor, offsets, rate)
+    )
+
+
+def _read_frames(episode: TrexEpisode, key: str, indices: np.ndarray) -> np.ndarray:
+    import decord
+
+    reader = decord.VideoReader(episode.video_dirs[key], num_threads=1)
+    return reader.get_batch([int(i) for i in indices]).asnumpy().astype(np.uint8)
+
+
+class ChunkBuilder:
+    """Builds training-format (K-block) raw samples from one episode."""
+
+    def __init__(self, dataset_root: str, episode_index: int = 0) -> None:
+        self.root = dataset_root
+        self.episode = TrexEpisode(dataset_root, episode_index)
+        self.stats = TrexRuntimeStatistics.from_dataset(dataset_root)
+        self.timestamps = _column(self.episode, "timestamp", np.float64).reshape(-1)
+        self.state = _column(self.episode, "observation.state_eef62")
+        self.action_abs = _column(self.episode, "action.eef62_absolute")
+        self.track_xy = _column(self.episode, "observation.track_xy")
+        self.track_vis = _column(self.episode, "observation.track_visibility")
+        self.force = _column(self.episode, "observation.tactile_force").reshape(
+            -1, 10, 6
+        )
+        self._deform: np.ndarray | None = None
+
+    def valid_anchor_times(self, blocks: int) -> list[float]:
+        future = max(
+            blocks * ACTION_HORIZON / ACTION_RATE_HZ,
+            blocks * VIDEO_FRAMES_PER_BLOCK / VIDEO_RATE_HZ,
+        )
+        grid = np.arange(
+            self.timestamps[0], self.timestamps[-1] + 1e-9, 1.0 / ACTION_RATE_HZ
+        )
+        return [float(t) for t in grid if t + future <= self.timestamps[-1]]
+
+    def deform_frames(self, size: int = 96) -> np.ndarray:
+        if self._deform is None:
+            import cv2
+            import decord
+
+            streams = []
+            for key in DEFORM_VIDEO_KEYS:
+                path = (
+                    Path(self.root) / "videos" / "chunk-000"
+                    / f"observation.images.{key}"
+                    / f"episode_{self.episode.episode_index:06d}.mp4"
+                )
+                reader = decord.VideoReader(str(path), num_threads=1)
+                frames = reader.get_batch(range(len(reader))).asnumpy()
+                frames = np.stack(
+                    [cv2.resize(f, (size, size), interpolation=cv2.INTER_AREA)
+                     for f in frames]
+                )
+                streams.append(frames.astype(np.uint8))
+            self._deform = np.stack(streams, axis=1)
+        return self._deform
+
+    def _norm_force(self, selection) -> np.ndarray:
+        values = self.stats.normalize_force(self.force[selection.indices])
+        values[selection.padding_mask] = 0.0
+        return values
+
+    def build(self, anchor: float, *, blocks: int, prompt: str,
+              with_deform: bool, history_only_video: bool = False) -> dict:
+        ts = self.timestamps
+        block_anchors = [
+            anchor + b * ACTION_HORIZON / ACTION_RATE_HZ for b in range(blocks)
+        ]
+        action_sel = _sample(ts, anchor, range(blocks * ACTION_HORIZON), ACTION_RATE_HZ)
+        state_sel = _sample(
+            ts, anchor, range(0, blocks * ACTION_HORIZON, ACTION_HORIZON),
+            ACTION_RATE_HZ,
+        )
+        reference = self.state[state_sel.indices]
+        absolute = self.action_abs[action_sel.indices].reshape(
+            blocks, ACTION_HORIZON, 62
+        )
+        delta = np.stack(
+            [eef62_delta_base(reference[b], absolute[b]) for b in range(blocks)]
+        ).reshape(blocks * ACTION_HORIZON, 62)
+
+        past_sels = [
+            _sample(ts, b, range(-(FORCE_HISTORY_FRAMES - 1), 1), ACTION_RATE_HZ)
+            for b in block_anchors
+        ]
+        future_sels = [
+            _sample(ts, b, range(TRACK_HORIZON), ACTION_RATE_HZ)
+            for b in block_anchors
+        ]
+        force_sels = [
+            [
+                _sample(ts, b + off / ACTION_RATE_HZ,
+                        range(-(FORCE_HISTORY_FRAMES - 1), 1), TACTILE_RATE_HZ)
+                for off in FORCE_OFFSETS
+            ]
+            for b in block_anchors
+        ]
+        force_history = np.stack(
+            [[self._norm_force(sel) for sel in block_sel] for block_sel in force_sels]
+        )
+
+        video_hist = _sample(ts, anchor, range(1), VIDEO_RATE_HZ)
+        if history_only_video:
+            video_indices = video_hist.indices
+        else:
+            video_future = _sample(
+                ts, anchor, range(1, blocks * VIDEO_FRAMES_PER_BLOCK + 1),
+                VIDEO_RATE_HZ,
+            )
+            video_indices = np.concatenate(
+                (video_hist.indices, video_future.indices)
+            )
+
+        raw: dict[str, object] = {
+            key: _read_frames(self.episode, key, video_indices)
+            for key in VIDEO_KEYS
+        }
+        raw.update(
+            {
+                "state.eef62": reference.astype(np.float32),
+                "action.eef62": delta.astype(np.float32),
+                "track_past_xy": np.stack(
+                    [self.track_xy[s.indices] for s in past_sels]
+                ),
+                "track_past_visibility": np.stack(
+                    [self.track_vis[s.indices] * (~s.padding_mask[:, None])
+                     for s in past_sels]
+                ),
+                "track_future_xy": np.stack(
+                    [self.track_xy[s.indices] for s in future_sels]
+                ),
+                "track_future_visibility": np.stack(
+                    [self.track_vis[s.indices] for s in future_sels]
+                ),
+                "current_force": force_history[:, :, -1],
+                "force_history": force_history,
+                "force_history_padding_mask": np.stack(
+                    [[sel.padding_mask for sel in block_sel]
+                     for block_sel in force_sels]
+                ),
+                "annotation.task": prompt,
+            }
+        )
+        if with_deform:
+            deform = self.deform_frames()
+            refresh = np.stack(
+                [[int(sel.indices[-1]) for sel in block_sel]
+                 for block_sel in force_sels]
+            )
+            raw["deform_current"] = deform[refresh]
+        return raw
+
+
+def load_pipeline(checkpoint: Path, *, training: bool):
+    cfg_dir = checkpoint / "experiment_cfg"
+    if not cfg_dir.exists():
+        cfg_dir = checkpoint.parent / "experiment_cfg"
+    cfg = OmegaConf.load(cfg_dir / "conf.yaml")
+    with open(cfg_dir / "metadata.json", "r", encoding="utf-8") as handle:
+        metadata = DatasetMetadata.model_validate(
+            json.load(handle)[EmbodimentTag.TREX.value]
+        )
+    transform = instantiate(cfg.transforms["trex"])
+    transform.set_metadata(metadata)
+    transform.train() if training else transform.eval()
+    collator = instantiate(cfg.data_collator)
+    return transform, collator
+
+
+def to_batch(sample: dict, collator, device, dtype) -> dict:
+    batch = collator([sample])
+    out = {}
+    for key, value in batch.items():
+        if torch.is_tensor(value):
+            value = (
+                value.to(device=device, dtype=dtype)
+                if value.is_floating_point()
+                else value.to(device=device)
+            )
+        out[key] = value
+    return out
+
+
+def stage1_attention_masses(policy, batch, *, layers, tau_value):
+    from groot.vla.model.trex_track_force import blocks as block_module
+
+    records: dict[int, dict] = {}
+
+    def make_patched(layer_index, module):
+        def patched(x, *, layout, rope_frequencies, allow_matrix=None, **_):
+            batch_size, length = x.shape[:2]
+            heads, head_dim = module.num_heads, module.head_dim
+            query = module.norm_q(module.q(x)).view(batch_size, length, heads, head_dim)
+            key = module.norm_k(module.k(x)).view(batch_size, length, heads, head_dim)
+            value = module.v(x).view(batch_size, length, heads, head_dim)
+            query = block_module.apply_multimodal_rope(
+                query, rope_frequencies
+            ).type_as(value)
+            key = block_module.apply_multimodal_rope(
+                key, rope_frequencies
+            ).type_as(value)
+            scores = torch.einsum("blhd,bmhd->bhlm", query, key) * head_dim**-0.5
+            scores = scores.float().masked_fill(
+                ~allow_matrix.view(1, 1, length, length), float("-inf")
+            )
+            attention = scores.softmax(dim=-1)
+            token_type, _ = layout.token_metadata(device=x.device)
+            layer_record = {}
+            for query_group, query_type in (
+                ("action", TokenType.ACTION), ("obs", TokenType.OBS),
+            ):
+                rows = (token_type == int(query_type)).nonzero(as_tuple=True)[0]
+                row_attention = attention[:, :, rows]
+                masses = {}
+                for name, key_type in (
+                    ("cond_obs", TokenType.CONDITIONING_OBS),
+                    ("obs", TokenType.OBS),
+                    ("action", TokenType.ACTION),
+                    ("state", TokenType.STATE),
+                    ("track_past", TokenType.TRACK_PAST),
+                    ("track_future", TokenType.TRACK_FUTURE),
+                ):
+                    cols = (token_type == int(key_type)).nonzero(as_tuple=True)[0]
+                    masses[name] = (
+                        float(row_attention[..., cols].sum(dim=-1).mean().item())
+                        if cols.numel()
+                        else 0.0
+                    )
+                layer_record[query_group] = masses
+            records[layer_index] = layer_record
+            output = torch.einsum(
+                "bhlm,bmhd->blhd", attention.to(value.dtype), value
+            ).reshape(batch_size, length, module.dim)
+            return module.o(output), None
+
+        return patched
+
+    originals = {}
+    for layer_index in layers:
+        module = policy.model.blocks[layer_index].self_attn
+        originals[layer_index] = module.forward
+        module.forward = make_patched(layer_index, module)
+    try:
+        tau = torch.full(
+            (1, AR_BLOCKS), tau_value, device=policy.device, dtype=policy.dtype
+        )
+        with torch.inference_mode():
+            policy.forward_core(batch, tau=tau)
+    finally:
+        for layer_index, forward in originals.items():
+            policy.model.blocks[layer_index].self_attn.forward = forward
+    return records
+
+
+class Stage2Recorder:
+    """Wrap each force-transformer layer to record action-query attention."""
+
+    def __init__(self, force):
+        self.force = force
+        self.groups = {"action": force.action_horizon,
+                       "f6": force.force_sensor_count,
+                       "vq": force.force_sensor_count}
+        if force.use_deform_tactile:
+            self.groups["deform"] = force.force_sensor_count
+        self.records: list[dict[str, float]] = []
+        self._originals = []
+
+    def __enter__(self):
+        for layer in self.force.transformer.layers:
+            self._originals.append((layer, layer.forward))
+            layer.forward = self._make(layer)
+        return self
+
+    def __exit__(self, *exc):
+        for layer, forward in self._originals:
+            layer.forward = forward
+
+    def _make(self, layer):
+        groups = self.groups
+        records = self.records
+
+        def forward(src, src_mask=None, src_key_padding_mask=None, is_causal=False):
+            x = src
+            normed = layer.norm1(x)
+            attn_out, weights = layer.self_attn(
+                normed, normed, normed,
+                attn_mask=src_mask,
+                key_padding_mask=src_key_padding_mask,
+                need_weights=True,
+                average_attn_weights=True,
+            )
+            action_rows = weights[:, : groups["action"]]
+            cursor, masses = 0, {}
+            for name, size in groups.items():
+                masses[name] = float(
+                    action_rows[..., cursor:cursor + size].sum(-1).mean().item()
+                )
+                cursor += size
+            masses["memory"] = float(action_rows[..., cursor:].sum(-1).mean().item())
+            records.append(masses)
+            x = x + layer.dropout1(attn_out)
+            x = x + layer._ff_block(layer.norm2(x))
+            return x
+
+        return forward
+
+
+def refine_with_keep(policy, state_batch, refinement, *, keep: bool,
+                     deform_images=None):
+    """Mirror refine_action_suffix but force the tactile keep mask."""
+
+    force = policy.force_transformer
+    action = pad_action_62_to_64(refinement["coarse_action"]).clone()
+    action[..., policy.config.physical_action_dim:] = 0
+    keep_mask = torch.full(
+        (action.shape[0],), 1.0 if keep else 0.0,
+        device=policy.device, dtype=policy.dtype,
+    )
+    for step in policy.schedule.iter_steps("force"):
+        flow = force(
+            action,
+            step.tau,
+            refinement["current_force"],
+            refinement["history"],
+            coarse_memory=refinement["memory"],
+            update_offset=0,
+            tactile_keep_mask=keep_mask,
+            tactile_history_valid_mask=refinement["history_valid"],
+            deform_images=deform_images,
+        )
+        action = euler_flow_step(action, flow, step.tau, step.tau_next)
+        action[..., policy.config.physical_action_dim:] = 0
+    return action
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset-root", required=True)
+    parser.add_argument("--run", action="append", required=True)
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--forward-anchors", type=int, default=6)
+    parser.add_argument("--openloop-anchors", type=int, default=16)
+    parser.add_argument("--tau-grid", default="0.8,0.6,0.4,0.2,0.05")
+    parser.add_argument("--tau-repeats", type=int, default=4)
+    parser.add_argument("--attn-layers", default="0,14,29")
+    parser.add_argument("--prompt", default=None)
+    args = parser.parse_args()
+
+    device = torch.device("cuda")
+    builder = ChunkBuilder(args.dataset_root)
+    prompt = args.prompt
+    if prompt is None:
+        with open(Path(args.dataset_root) / "meta" / "tasks.jsonl") as handle:
+            prompt = json.loads(handle.readline())["task"]
+
+    anchors_k4 = builder.valid_anchor_times(AR_BLOCKS)
+    forward_anchor_times = [
+        anchors_k4[int(i)]
+        for i in np.linspace(0, len(anchors_k4) - 1, args.forward_anchors)
+    ]
+    anchors_k1 = builder.valid_anchor_times(1)
+    openloop_anchor_times = [
+        anchors_k1[int(i)]
+        for i in np.linspace(0, len(anchors_k1) - 1, args.openloop_anchors)
+    ]
+    tau_grid = [float(v) for v in args.tau_grid.split(",")]
+    attn_layers = [int(v) for v in args.attn_layers.split(",")]
+
+    results: dict[str, dict] = {}
+    for spec in args.run:
+        name, _, checkpoint = spec.partition("=")
+        checkpoint_path = Path(checkpoint)
+        print(f"=== {name}: {checkpoint_path} ===", flush=True)
+        model = TrexTrackForceVLA.load_lora(str(checkpoint_path))
+        model = model.to(device=device, dtype=torch.bfloat16)
+        model.eval()
+        policy = model.action_head
+        use_force = policy.config.use_force
+        use_deform = policy.config.use_deform_tactile
+        transform, collator = load_pipeline(checkpoint_path, training=True)
+        transform_eval, _ = load_pipeline(checkpoint_path, training=False)
+
+        record: dict[str, object] = {
+            "checkpoint": str(checkpoint_path),
+            "use_track": policy.config.use_track,
+            "use_force": use_force,
+            "use_deform_tactile": use_deform,
+        }
+
+        # ---------------- controlled forward losses ----------------------
+        forward_batches = []
+        for anchor in forward_anchor_times:
+            raw = builder.build(anchor, blocks=AR_BLOCKS, prompt=prompt,
+                                with_deform=use_deform)
+            forward_batches.append(
+                to_batch(transform(dict(raw)), collator, device, torch.bfloat16)
+            )
+        tau_table = {}
+        for tau_value in tau_grid:
+            metrics: dict[str, list[float]] = {}
+            for repeat in range(args.tau_repeats):
+                for batch_index, batch in enumerate(forward_batches):
+                    torch.manual_seed(100000 + repeat * 1000 + batch_index * 10)
+                    tau = torch.full((1, AR_BLOCKS), tau_value,
+                                     device=device, dtype=torch.bfloat16)
+                    with torch.inference_mode():
+                        out = policy.forward_core(batch, tau=tau)
+                    for key in ("action_loss", "dynamics_loss",
+                                "track_flow_loss", "force_loss"):
+                        metrics.setdefault(key, []).append(float(out[key]))
+            tau_table[f"{tau_value:g}"] = {
+                key: [float(np.mean(v)), float(np.std(v))]
+                for key, v in metrics.items()
+            }
+        record["forward_tau_losses"] = tau_table
+        print(f"  forward losses done", flush=True)
+
+        # ---------------- open-loop chunk reconstruction -----------------
+        arm_dims = list(range(0, 9)) + list(range(31, 40))
+        hand_dims = list(range(9, 31)) + list(range(40, 62))
+        openloop: dict[str, list[float]] = {}
+        split_metrics: dict[str, list[float]] = {}
+        refinement_state = None
+        for anchor_index, anchor in enumerate(openloop_anchor_times):
+            raw = builder.build(anchor, blocks=1, prompt=prompt,
+                                with_deform=use_deform, history_only_video=True)
+            sample = transform_eval(dict(raw))
+            gt = np.asarray(sample["action"], dtype=np.float32)[..., :62]
+            gt = torch.as_tensor(gt.reshape(-1, 62)[:ACTION_HORIZON])
+            batch = to_batch(sample, collator, device, torch.bfloat16)
+            if use_deform and "deform_current" in batch:
+                batch["deform_current"] = batch["deform_current"][:, :1, :1]
+            modes = [("cascade", True)] if use_force else []
+            modes.append(("coarse_only", False))
+            for mode_name, refine in modes:
+                with torch.inference_mode():
+                    result = policy.sample(
+                        batch, seed=500 + anchor_index,
+                        run_force_refinement=refine,
+                        return_refinement_state=refine and anchor_index == 0,
+                    )
+                if refine and anchor_index == 0:
+                    current, history, valid = policy._extract_force_inputs(batch, 1)
+                    refinement_state = {
+                        "coarse_action": result["coarse_action_at_split"],
+                        "memory": result["coarse_memory"],
+                        "current_force": current[:, 0, 0].to(policy.device,
+                                                             policy.dtype),
+                        "history": history[:, 0, 0].to(policy.device),
+                        "history_valid": (
+                            None if valid is None
+                            else valid[:, 0, 0].to(policy.device)
+                        ),
+                        "deform": (
+                            policy._extract_deform_images(batch, 1, 1)
+                            if use_deform else None
+                        ),
+                    }
+                pred = result["action_pred"][0, :, :62].float().cpu()
+                openloop.setdefault(mode_name, []).append(
+                    float(((pred - gt) ** 2).mean())
+                )
+                prefix = "cascade" if mode_name == "cascade" else "coarse"
+                split_metrics.setdefault(f"{prefix}_arm", []).append(
+                    float(((pred[:, arm_dims] - gt[:, arm_dims]) ** 2).mean())
+                )
+                split_metrics.setdefault(f"{prefix}_hand", []).append(
+                    float(((pred[:, hand_dims] - gt[:, hand_dims]) ** 2).mean())
+                )
+        record["openloop_action_mse"] = {
+            key: [float(np.mean(v)), float(np.std(v)), len(v)]
+            for key, v in openloop.items()
+        }
+        record["openloop_action_mse_split"] = {
+            key: float(np.mean(v)) for key, v in split_metrics.items()
+        }
+        print(f"  open-loop done", flush=True)
+
+        # ---------------- attention masses --------------------------------
+        record["stage1_attention"] = {
+            str(layer): masses
+            for layer, masses in stage1_attention_masses(
+                policy, forward_batches[0], layers=attn_layers, tau_value=0.4
+            ).items()
+        }
+        if use_force and refinement_state is not None:
+            with Stage2Recorder(policy.force_transformer) as recorder:
+                refine_with_keep(policy, None, refinement_state, keep=True,
+                                 deform_images=refinement_state["deform"])
+            per_layer = recorder.records[: len(
+                policy.force_transformer.transformer.layers
+            )]
+            record["stage2_attention"] = per_layer
+
+            refined_real = refine_with_keep(
+                policy, None, refinement_state, keep=True,
+                deform_images=refinement_state["deform"],
+            )
+            refined_masked = refine_with_keep(
+                policy, None, refinement_state, keep=False,
+                deform_images=refinement_state["deform"],
+            )
+            coarse = pad_action_62_to_64(refinement_state["coarse_action"])
+            delta_tactile = (refined_real - refined_masked)[..., :62].float()
+            delta_refine = (refined_real - coarse)[..., :62].float()
+            record["tactile_sensitivity"] = {
+                "mean_abs_delta_vs_masked": float(delta_tactile.abs().mean()),
+                "max_abs_delta_vs_masked": float(delta_tactile.abs().max()),
+                "mean_abs_refinement": float(delta_refine.abs().mean()),
+            }
+        print(f"  attention/sensitivity done", flush=True)
+
+        results[name] = record
+        del model, policy, forward_batches
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    with open(args.out, "w", encoding="utf-8") as handle:
+        json.dump(results, handle, indent=1)
+    print(f"wrote {args.out}")
+
+
+if __name__ == "__main__":
+    main()