"""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.n1_5.sim_policy import unsqueeze_dict_values 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) values = np.asarray(values) while values.dtype == object: values = np.stack([np.stack(row) for row in values]) return values.astype(dtype) 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 @torch.inference_mode() 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) delta = np.asarray(raw.pop("action.eef62"), dtype=np.float32) scale = builder.stats.action_q99 - builder.stats.action_q01 gt = np.clip( 2.0 * (delta - builder.stats.action_q01) / np.where(scale == 0, 1.0, scale) - 1.0, -1.0, 1.0, ) gt = np.where(scale == 0, delta, gt) gt = torch.as_tensor(gt[:ACTION_HORIZON]) collated = transform_eval(unsqueeze_dict_values(dict(raw))) batch = {} for key, value in collated.items(): if torch.is_tensor(value): value = ( value.to(device=device, dtype=torch.bfloat16) if value.is_floating_point() else value.to(device=device) ) batch[key] = value 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_raw"] = openloop 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()