#!/usr/bin/env python3 """ SmolVLA Attention Visualizer — End-to-End ========================================== What this does: Loads a pretrained SmolVLA model and a real-world LeRobot dataset, runs inference on sample frames, and extracts attention heatmaps from the SigLIP vision encoder to show exactly which pixels drive the model's action predictions. Architecture recap: SmolVLA = SigLIP (vision encoder) + SmolLM2 (language decoder) + Action Expert SigLIP is a Vision Transformer. It splits each image into 14x14 pixel patches, then runs self-attention across all patches. The attention weights tell us which patches "look at" which other patches. We extract the attention from the LAST layer of SigLIP, average across all heads, and compute the mean attention each patch receives from all others to produce a spatial heatmap over the original image. (SigLIP has no CLS token — it is a pure patch-based encoder with no special classification token.) Usage: pip install lerobot[smolvla] matplotlib numpy Pillow python smolvla_attention_viz.py # Or with options: python smolvla_attention_viz.py \ --model lerobot/smolvla_base \ --dataset lerobot/svla_so101_pickplace \ --episode 0 \ --num-frames 8 \ --output-dir ./attention_maps \ --device cpu Output: Saves PNG files showing original frames with attention heatmap overlays, plus a grid summary image. Each heatmap shows which image regions the vision encoder attends to most strongly. Interpreting results: - GOOD policy: Attention concentrated on gripper, target object, and goal - OVERFITTING: Attention spread across background (shelves, cables, table grain) - This is the "smoking gun" for the background distribution shift problem """ import argparse import logging import math import os import sys import warnings from pathlib import Path # Preload Homebrew FFmpeg 6 libavdevice so PyAV/av doesn't load its bundled copy; # avoids "Class AVFFrameReceiver is implemented in both..." duplicate symbol warning. _ffmpeg6_lib = "/opt/homebrew/opt/ffmpeg@6/lib" if os.path.isdir(_ffmpeg6_lib): _libavdevice = os.path.join(_ffmpeg6_lib, "libavdevice.60.dylib") if os.path.isfile(_libavdevice): try: import ctypes ctypes.CDLL(_libavdevice) except OSError: pass import numpy as np import torch import torch.nn.functional as F import matplotlib matplotlib.use("Agg") # Non-interactive backend for saving files import matplotlib.pyplot as plt import matplotlib.gridspec as gridspec from PIL import Image from matplotlib.colors import LinearSegmentedColormap _CYAN_CMAP = LinearSegmentedColormap.from_list("cyan", ["black", "cyan", "white"]) # --------------------------------------------------------------------------- # 1. Attention hook — intercepts SigLIP's self-attention weights # --------------------------------------------------------------------------- class SigLIPAttentionCapture: """ Registers forward hooks on the SigLIP vision encoder's attention layers to capture attention weights during inference. SigLIP uses torch.nn.MultiheadAttention, which can return attn weights when called with need_weights=True. Since the LeRobot/transformers code may not pass that flag, we hook into the module and monkey-patch the forward call to force weight capture. """ def __init__(self): self.attention_maps = [] # List of (layer_idx, attn_weights) tuples self.hooks = [] def _make_hook(self, layer_idx): """Create a hook closure for a specific layer.""" def hook_fn(module, input_args, output): """ torch.nn.MultiheadAttention.forward returns: (attn_output, attn_output_weights) when need_weights=True (attn_output, None) when need_weights=False We intercept the output. If weights are None, we re-run the attention computation to get them. """ if isinstance(output, tuple) and len(output) == 2: attn_output, attn_weights = output if attn_weights is not None: self.attention_maps.append((layer_idx, attn_weights.detach().cpu())) return output # Weights not available — eager attention is forced so this # should not happen. Log a warning instead of attempting a # manual computation that would ignore multi-head reshaping. print(f" WARNING: Could not capture attention weights at layer {layer_idx}") return output return hook_fn def register_hooks(self, vision_encoder): """ Walk the vision encoder module tree and hook into all MultiheadAttention or equivalent attention modules. """ self.clear() # Strategy 1: Look for nn.MultiheadAttention modules found_mha = False for name, module in vision_encoder.named_modules(): if isinstance(module, torch.nn.MultiheadAttention): layer_idx = len(self.hooks) hook = module.register_forward_hook(self._make_hook(layer_idx)) self.hooks.append(hook) found_mha = True if found_mha: print(f" Hooked into {len(self.hooks)} MultiheadAttention layers") return # Strategy 2: Look for attention modules by name pattern # (transformers-style SiglipAttention or similar) for name, module in vision_encoder.named_modules(): module_type = type(module).__name__.lower() if "attention" in module_type and "embed" not in module_type: layer_idx = len(self.hooks) hook = module.register_forward_hook(self._make_hook(layer_idx)) self.hooks.append(hook) if self.hooks: print(f" Hooked into {len(self.hooks)} attention modules (by name)") else: print(" WARNING: No attention modules found. Will use gradient-based fallback.") def clear(self): """Remove all hooks and clear captured maps.""" for h in self.hooks: h.remove() self.hooks = [] self.attention_maps = [] def reset_maps(self): """Clear captured maps but keep hooks registered.""" self.attention_maps = [] def get_last_layer_attention(self): """ Return attention weights from the last encoder layer. Shape: (batch, num_patches, num_patches) or (batch, heads, patches, patches) """ if not self.attention_maps: return None # Sort by layer index, take last sorted_maps = sorted(self.attention_maps, key=lambda x: x[0]) _, attn = sorted_maps[-1] return attn def get_all_layer_attentions(self): """Return attention weights from all layers, sorted by layer index.""" if not self.attention_maps: return [] return sorted(self.attention_maps, key=lambda x: x[0]) # --------------------------------------------------------------------------- # 1b. Cross-attention hook — action expert attending to vision tokens # --------------------------------------------------------------------------- class ActionVisionAttentionCapture: """ Captures how the action expert's queries attend to vision tokens in the VLM prefix. SmolVLA has no explicit cross-attention *layers*. Instead the VLM builds a KV cache from the prefix (vision + language + state tokens) and the action expert queries that cache inside ``forward_cross_attn_layer()``. All attention goes through a single ``eager_attention_forward()`` method. We monkey-patch that method to intercept the softmax probabilities whenever expert queries (suffix) attend to prefix keys (detected by Q seq-len != K seq-len). After capture the caller selects just the columns corresponding to vision tokens and reshapes into a spatial heatmap. """ def __init__(self): self.cross_attn_weights = [] # (layer_call_idx, probs) probs: (B, heads, q_len, k_len) self._original_fn = None self._patched_obj = None self._call_idx = 0 # ------------------------------------------------------------------ # patch / unpatch # ------------------------------------------------------------------ def register(self, vlm_with_expert): """Wrap ``eager_attention_forward`` on *vlm_with_expert*.""" import types self.clear() self._original_fn = vlm_with_expert.eager_attention_forward self._patched_obj = vlm_with_expert capture = self def _wrapped(self_model, attention_mask, batch_size, head_dim, query_states, key_states, value_states): # Ensure boolean mask (some code paths produce Long 0/1 masks) if attention_mask.dtype != torch.bool: attention_mask = attention_mask.bool() # ---------- original computation ---------- output = capture._original_fn( attention_mask, batch_size, head_dim, query_states, key_states, value_states, ) # ---------- detect cross-attention ---------- q_len = query_states.shape[1] k_len = key_states.shape[1] if q_len != k_len: # Re-derive attention probs (mirrors eager_attention_forward) num_att_heads = self_model.num_attention_heads num_kv_heads = self_model.num_key_value_heads num_kv_groups = num_att_heads // num_kv_heads seq_len_k = key_states.shape[1] ks = key_states[:, :, :, None, :].expand( batch_size, seq_len_k, num_kv_heads, num_kv_groups, head_dim ).reshape(batch_size, seq_len_k, num_kv_heads * num_kv_groups, head_dim) q = query_states.to(dtype=torch.float32).transpose(1, 2) k = ks.to(dtype=torch.float32).transpose(1, 2) scores = torch.matmul(q, k.transpose(2, 3)) * (head_dim ** -0.5) big_neg = torch.finfo(scores.dtype).min scores = torch.where( attention_mask[:, None, :, :], scores, big_neg ) probs = F.softmax(scores, dim=-1) capture.cross_attn_weights.append( (capture._call_idx, probs.detach().cpu()) ) capture._call_idx += 1 return output vlm_with_expert.eager_attention_forward = types.MethodType( _wrapped, vlm_with_expert ) def clear(self): """Remove the monkey-patch and discard captured data.""" if self._original_fn is not None and self._patched_obj is not None: self._patched_obj.eager_attention_forward = self._original_fn self._original_fn = None self._patched_obj = None self.cross_attn_weights = [] self._call_idx = 0 def reset_maps(self): self.cross_attn_weights = [] self._call_idx = 0 # ------------------------------------------------------------------ # helpers # ------------------------------------------------------------------ def get_vision_token_range(self, policy): """ Return ``(start, end)`` — the column indices in the prefix KV cache that correspond to vision tokens (one image, no special tokens by default). """ add_special = getattr(policy.config, "add_image_special_tokens", False) # Run the connector on a dummy image to get the token count vlm_with_expert = policy.model.vlm_with_expert vision_model = vlm_with_expert.get_vlm_model().vision_model img_size = getattr( getattr(vision_model, "config", None), "image_size", 384 ) dummy = torch.zeros(1, 3, img_size, img_size, device=next(vision_model.parameters()).device, dtype=next(vision_model.parameters()).dtype) with torch.no_grad(): hidden = vision_model(pixel_values=dummy).last_hidden_state after_connector = vlm_with_expert.get_vlm_model().connector(hidden) n_vision = after_connector.shape[1] start = 1 if add_special else 0 end = start + n_vision return start, end, n_vision def get_mean_cross_attention(self, vision_start, vision_end): """ Average captured cross-attention over all layers and heads, then select only the vision-token columns. Returns: Tensor of shape ``(n_vision_tokens,)`` or *None*. """ if not self.cross_attn_weights: return None accum = None count = 0 for _idx, probs in self.cross_attn_weights: # probs: (B, heads, q_len, k_len) # Average over batch, heads, and query tokens avg = probs[0].mean(dim=0).mean(dim=0) # (k_len,) vision_avg = avg[vision_start:vision_end] if accum is None: accum = vision_avg else: accum = accum + vision_avg count += 1 if count == 0: return None return accum / count # --------------------------------------------------------------------------- # 2. Gradient-based attention fallback (GradCAM-style) # --------------------------------------------------------------------------- class GradCAMFallback: """ If we can't extract attention weights directly (e.g., because the attention implementation doesn't expose them), we use GradCAM on the last convolutional/linear layer of the vision encoder. GradCAM: Gradient-weighted Class Activation Mapping - Forward pass through the model - Backward pass from the action output w.r.t. vision features - Weight the feature maps by their gradient importance - Produces a heatmap showing which spatial regions influenced the output """ def __init__(self): self.features = None self.gradients = None self.hook_f = None self.hook_b = None def register(self, target_layer): """Register forward and backward hooks on a target layer.""" def forward_hook(module, input, output): if isinstance(output, tuple): self.features = output[0].detach() else: self.features = output.detach() def backward_hook(module, grad_input, grad_output): self.gradients = grad_output[0].detach() self.hook_f = target_layer.register_forward_hook(forward_hook) self.hook_b = target_layer.register_full_backward_hook(backward_hook) def compute_cam(self): """Compute GradCAM heatmap from stored features and gradients.""" if self.features is None or self.gradients is None: return None # Global average pooling of gradients → channel weights weights = self.gradients.mean(dim=-1, keepdim=True) # (B, T, 1) or (B, C, 1, 1) # Weighted combination of feature maps cam = (weights * self.features).sum(dim=-1) # (B, T) or (B, C) cam = F.relu(cam) # Only positive contributions # Normalize if cam.max() > 0: cam = cam / cam.max() return cam.cpu() def clear(self): if self.hook_f: self.hook_f.remove() if self.hook_b: self.hook_b.remove() # --------------------------------------------------------------------------- # 3. Attention-to-heatmap conversion # --------------------------------------------------------------------------- def attention_to_heatmap(attn_weights, grid_size, image_size): """ Convert attention weights from patch-space to pixel-space heatmap. Args: attn_weights: Tensor of shape (num_patches,) or (H_patches, W_patches) representing per-patch attention scores grid_size: (H_patches, W_patches) — the patch grid dimensions image_size: (H_pixels, W_pixels) — the original image dimensions Returns: heatmap: numpy array of shape (H_pixels, W_pixels) normalized to [0, 1] """ h_patches, w_patches = grid_size h_img, w_img = image_size # Reshape to 2D grid if flat if attn_weights.dim() == 1: expected = h_patches * w_patches if attn_weights.shape[0] != expected: # Try to infer square grid side = int(math.sqrt(attn_weights.shape[0])) if side * side == attn_weights.shape[0]: h_patches, w_patches = side, side else: # Truncate or pad attn_weights = attn_weights[:expected] attn_2d = attn_weights.reshape(h_patches, w_patches) else: attn_2d = attn_weights # Upsample to image resolution using bilinear interpolation attn_2d = attn_2d.float().unsqueeze(0).unsqueeze(0) # (1, 1, H, W) heatmap = F.interpolate(attn_2d, size=(h_img, w_img), mode="bilinear", align_corners=False) heatmap = heatmap.squeeze().numpy() # Normalize to [0, 1] hmin, hmax = heatmap.min(), heatmap.max() if hmax > hmin: heatmap = (heatmap - hmin) / (hmax - hmin) return heatmap def compute_patch_attention_scores(attn_weights, method="mean"): """ Reduce full attention matrix to per-patch importance scores. SigLIP is a pure patch encoder with no CLS token, so we use column-wise mean attention (how much attention each patch receives from all other patches) as the default aggregation. Args: attn_weights: Tensor, typically (num_patches, num_patches) or (heads, num_patches, num_patches) method: How to aggregate: - "mean": Average attention each patch receives from all others Returns: scores: Tensor of shape (num_patches,) """ # If multi-head, average across heads first if attn_weights.dim() == 3: attn_weights = attn_weights.mean(dim=0) # (patches, patches) if attn_weights.dim() == 1: return attn_weights # Already reduced # Column-wise mean = how much attention each patch receives scores = attn_weights.mean(dim=0) return scores def compute_attention_rollout(all_layer_attentions): """ Combine attention across all encoder layers using attention rollout. The method accounts for residual connections by mixing each layer's attention with an identity matrix (the residual stream keeps a copy of each token unchanged): rollout = I for A in layers: A_hat = 0.5 * I + 0.5 * A # residual connection rollout = rollout @ A_hat normalize each row to sum to 1 Args: all_layer_attentions: list of (layer_idx, attn_weights) tuples. Each attn_weights tensor has shape (batch, heads, patches, patches) or (heads, patches, patches) or (patches, patches). Returns: rollout: Tensor of shape (patches, patches) — the accumulated attention from input patches to output patches. """ matrices = [] for _layer_idx, attn in sorted(all_layer_attentions, key=lambda x: x[0]): # Reduce to (patches, patches) while attn.dim() > 3: attn = attn[0] # drop batch if attn.dim() == 3: attn = attn.mean(dim=0) # average heads matrices.append(attn.float()) if not matrices: return None n = matrices[0].shape[0] eye = torch.eye(n) rollout = eye.clone() for A in matrices: A_hat = 0.5 * eye + 0.5 * A rollout = rollout @ A_hat # Row-normalise so each row sums to 1 rollout = rollout / (rollout.sum(dim=-1, keepdim=True) + 1e-8) return rollout def compute_positional_baseline(vision_encoder, attn_capture, device, method): """ Compute the attention pattern produced by a content-free (mean-gray) image. This captures the fixed positional component of attention so it can be subtracted from real frames to reveal the content-dependent signal. Args: vision_encoder: The SigLIP vision encoder module. attn_capture: A ``SigLIPAttentionCapture`` instance with hooks already registered on *vision_encoder*. device: Torch device. method: ``"last-layer"`` or ``"rollout"`` — same aggregation used for real frames so the baseline is comparable. Returns: Tensor of shape ``(num_patches,)`` — per-patch baseline scores. """ # Build a mean-gray image at the encoder's expected resolution img_size = getattr( getattr(vision_encoder, "config", None), "image_size", None, ) or getattr(vision_encoder, "image_size", 384) gray = torch.full((1, 3, img_size, img_size), 0.5, device=device) try: enc_dtype = next(vision_encoder.parameters()).dtype gray = gray.to(enc_dtype) except StopIteration: pass patch_size = getattr(vision_encoder, "patch_size", None) or getattr( getattr(vision_encoder, "config", None), "patch_size", 14, ) n_patches_h = img_size // patch_size n_patches_w = img_size // patch_size patch_mask = torch.ones(1, n_patches_h, n_patches_w, dtype=torch.bool, device=device) # Forward pass through the vision encoder attn_capture.reset_maps() with torch.no_grad(): try: if hasattr(vision_encoder, "embeddings") and hasattr(vision_encoder, "encoder"): embeddings = vision_encoder.embeddings(gray, patch_mask) vision_encoder.encoder(embeddings) elif hasattr(vision_encoder, "forward"): vision_encoder(gray) else: vision_encoder(pixel_values=gray, patch_attention_mask=patch_mask) except Exception as e: print(f" WARNING: Baseline forward pass failed ({e}), skipping correction") return None # Reduce captured attention using the same method as real frames if method == "rollout": all_layers = attn_capture.get_all_layer_attentions() rollout_mat = compute_attention_rollout(all_layers) if rollout_mat is None: return None baseline_scores = rollout_mat.mean(dim=0) else: # "last-layer" or "all-layers" both use last-layer for the summary attn = attn_capture.get_last_layer_attention() if attn is None: return None while attn.dim() > 3: attn = attn[0] if attn.dim() == 3: attn = attn.mean(dim=0) baseline_scores = compute_patch_attention_scores(attn, method="mean") attn_capture.reset_maps() return baseline_scores # --------------------------------------------------------------------------- # 4. Visualization # --------------------------------------------------------------------------- def overlay_heatmap(image_np, heatmap, alpha=0.5, colormap="jet"): """ Overlay a heatmap on an image. Args: image_np: numpy array (H, W, 3) in [0, 255] uint8 or [0, 1] float heatmap: numpy array (H, W) in [0, 1] alpha: blend factor (0 = only image, 1 = only heatmap) colormap: matplotlib colormap name Returns: blended: numpy array (H, W, 3) uint8 """ if image_np.dtype == np.float32 or image_np.dtype == np.float64: if image_np.max() <= 1.0: image_np = (image_np * 255).astype(np.uint8) cmap = plt.get_cmap(colormap) heatmap_colored = cmap(heatmap)[:, :, :3] # (H, W, 3) float in [0, 1] heatmap_colored = (heatmap_colored * 255).astype(np.uint8) blended = ( (1 - alpha) * image_np.astype(np.float32) + alpha * heatmap_colored.astype(np.float32) ).astype(np.uint8) return blended def _resize_heatmap(heatmap, h, w): """Resize a heatmap to (h, w) if necessary.""" if heatmap.shape[0] != h or heatmap.shape[1] != w: return np.array( Image.fromarray((heatmap * 255).astype(np.uint8)).resize((w, h)) ) / 255.0 return heatmap def _frame_to_np(frame): """Convert a frame tensor or array to uint8 numpy (H, W, 3).""" if isinstance(frame, torch.Tensor): frame_np = frame.permute(1, 2, 0).numpy() if frame_np.max() <= 1.0: frame_np = (frame_np * 255).astype(np.uint8) else: frame_np = frame_np.astype(np.uint8) else: frame_np = np.array(frame) return frame_np def create_visualization_grid(frames, heatmaps, actions=None, cross_attn_heatmaps=None, episode_idx=0, output_path="attention_grid.png"): """ Create a grid visualization showing original frames, heatmaps, and overlays. Layout per frame: Row 1: Original image Row 2: Vision encoder self-attention heatmap (colorized) Row 3: Overlay (image + self-attention heatmap blended) When *cross_attn_heatmaps* is provided two extra rows are added: Row 4: Action→Vision cross-attention heatmap Row 5: Dual-color overlay (self-attn blue, cross-attn red) """ n_frames = len(frames) has_cross = cross_attn_heatmaps is not None and len(cross_attn_heatmaps) == n_frames n_rows = 5 if has_cross else 3 fig = plt.figure(figsize=(4 * n_frames, 4 * n_rows)) gs = gridspec.GridSpec(n_rows, n_frames, hspace=0.3, wspace=0.05) for i, (frame, heatmap) in enumerate(zip(frames, heatmaps)): frame_np = _frame_to_np(frame) h, w = frame_np.shape[:2] heatmap_resized = _resize_heatmap(heatmap, h, w) overlay = overlay_heatmap(frame_np, heatmap_resized, alpha=0.45) # Row 1: Original ax1 = fig.add_subplot(gs[0, i]) ax1.imshow(frame_np) ax1.set_title(f"Frame {i}", fontsize=10) ax1.axis("off") if i == 0: ax1.set_ylabel("Original", fontsize=11, rotation=0, labelpad=60, va="center") # Row 2: Self-attention heatmap ax2 = fig.add_subplot(gs[1, i]) ax2.imshow(heatmap_resized, cmap="jet", vmin=0, vmax=1) ax2.axis("off") if i == 0: ax2.set_ylabel("SigLIP\nself-attn", fontsize=11, rotation=0, labelpad=60, va="center") # Row 3: Self-attention overlay ax3 = fig.add_subplot(gs[2, i]) ax3.imshow(overlay) ax3.axis("off") if i == 0: ax3.set_ylabel("Self-attn\noverlay", fontsize=11, rotation=0, labelpad=60, va="center") if has_cross: cross_hm = _resize_heatmap(cross_attn_heatmaps[i], h, w) # Row 4: Cross-attention heatmap ax4 = fig.add_subplot(gs[3, i]) ax4.imshow(cross_hm, cmap="Greens", vmin=0, vmax=1) ax4.axis("off") if i == 0: ax4.set_ylabel("Action\ncross-attn", fontsize=11, rotation=0, labelpad=60, va="center") # Row 5: Co-attention overlay (self-attn × cross-attn) co_attn = heatmap_resized * cross_hm # element-wise product co_attn = co_attn / (co_attn.max() + 1e-8) # renormalize to [0, 1] co_overlay = frame_np.copy() co_overlay = (0.5 * co_overlay.astype(np.float32) + 0.5 * _CYAN_CMAP(co_attn)[:, :, :3] * 255) co_overlay = np.clip(co_overlay, 0, 255).astype(np.uint8) ax5 = fig.add_subplot(gs[4, i]) ax5.imshow(co_overlay) ax5.axis("off") if i == 0: ax5.set_ylabel("Co-attention\noverlay", fontsize=11, rotation=0, labelpad=60, va="center") if has_cross: legend = ("Row 1: Original | Row 2: SigLIP self-attn heatmap | Row 3: Self-attn overlay | " "Row 4: Action cross-attn heatmap | Row 5: Co-attention (self × cross)") dual_legend = "Co-attention: self-attn × cross-attn — bright regions are both visually salient and action-relevant" else: legend = "Row 1: Original | Row 2: SigLIP self-attn heatmap | Row 3: Overlay (heatmap on frame)" dual_legend = None if dual_legend: title = ( f"SmolVLA Attention — Episode {episode_idx}\n\n" f"{legend}\n\n" f"{dual_legend}" ) else: title = ( f"SmolVLA Attention — Episode {episode_idx}\n\n" f"{legend}" ) fig.suptitle(title, fontsize=14, fontweight="bold", y=0.98) plt.savefig(output_path, dpi=150, bbox_inches="tight", facecolor="white") plt.close() print(f" Saved grid: {output_path}") def save_individual_frames(frames, heatmaps, output_dir, episode_idx=0): """Save each frame's overlay as a separate high-res PNG.""" os.makedirs(output_dir, exist_ok=True) for i, (frame, heatmap) in enumerate(zip(frames, heatmaps)): if isinstance(frame, torch.Tensor): frame_np = frame.permute(1, 2, 0).numpy() if frame_np.max() <= 1.0: frame_np = (frame_np * 255).astype(np.uint8) else: frame_np = frame_np.astype(np.uint8) else: frame_np = np.array(frame) h, w = frame_np.shape[:2] if heatmap.shape[0] != h or heatmap.shape[1] != w: heatmap_resized = np.array( Image.fromarray((heatmap * 255).astype(np.uint8)).resize((w, h)) ) / 255.0 else: heatmap_resized = heatmap overlay = overlay_heatmap(frame_np, heatmap_resized, alpha=0.45) # Save overlay path = os.path.join(output_dir, f"ep{episode_idx:03d}_frame{i:04d}_overlay.png") Image.fromarray(overlay).save(path) # Save raw heatmap path_hm = os.path.join(output_dir, f"ep{episode_idx:03d}_frame{i:04d}_heatmap.png") fig_hm, ax_hm = plt.subplots(figsize=(6, 6)) ax_hm.imshow(heatmap_resized, cmap="jet") ax_hm.axis("off") fig_hm.savefig(path_hm, dpi=100, bbox_inches="tight") plt.close(fig_hm) print(f" Saved {len(frames)} individual frames to {output_dir}/") def create_per_head_grid(frame, attn_weights, grid_size, image_size, output_path="per_head_attention.png"): """ Visualise each attention head's pattern individually for a single frame. Useful for identifying specialised heads (e.g. one tracking the gripper, another tracking the object). Args: frame: image tensor (C, H, W) or numpy (H, W, 3) attn_weights: Tensor of shape ``(heads, patches, patches)`` — one attention matrix per head (not yet averaged). grid_size: ``(H_patches, W_patches)`` image_size: ``(H_pixels, W_pixels)`` output_path: where to save the PNG """ if isinstance(frame, torch.Tensor): frame_np = frame.permute(1, 2, 0).numpy() if frame_np.max() <= 1.0: frame_np = (frame_np * 255).astype(np.uint8) else: frame_np = frame_np.astype(np.uint8) else: frame_np = np.array(frame) n_heads = attn_weights.shape[0] cols = min(n_heads, 8) rows = math.ceil(n_heads / cols) fig, axes = plt.subplots(rows, cols, figsize=(3 * cols, 3 * rows)) if rows == 1 and cols == 1: axes = np.array([[axes]]) elif rows == 1 or cols == 1: axes = axes.reshape(rows, cols) for h in range(n_heads): r, c = divmod(h, cols) head_attn = attn_weights[h] # (patches, patches) scores = head_attn.mean(dim=0) # per-patch importance hmap = attention_to_heatmap(scores, grid_size, image_size) blended = overlay_heatmap(frame_np.copy(), hmap, alpha=0.5) axes[r, c].imshow(blended) axes[r, c].set_title(f"Head {h}", fontsize=9) axes[r, c].axis("off") # Turn off unused subplots for idx in range(n_heads, rows * cols): r, c = divmod(idx, cols) axes[r, c].axis("off") fig.suptitle("SigLIP vision encoder per-head self-attention", fontsize=13, fontweight="bold") plt.savefig(output_path, dpi=120, bbox_inches="tight", facecolor="white") plt.close() print(f" Saved per-head grid: {output_path}") # --------------------------------------------------------------------------- # 5. Model and dataset loading # --------------------------------------------------------------------------- def find_vision_encoder(policy): """ Navigate the SmolVLA model hierarchy to find the SigLIP vision encoder. SmolVLA structure (typical): policy.model .vlm_model (or .model) .vision_model (or .vision_encoder) .encoder .layers[0..N] .self_attn ← this is what we hook into """ model = policy # Traverse known paths search_paths = [ # SmolVLA (LeRobot): policy.model.vlm_with_expert.vlm.model.vision_model "model.vlm_with_expert.vlm.model.vision_model", # SmolVLA / other VLM wrapping "model.vlm.vision_model", "model.vlm.model.vision_model", "model.vlm_model.vision_model", "model.vlm_model.model.vision_model", # Direct transformers-style "model.vision_model", "model.model.vision_model", # Idefics3-style (SmolVLM uses same impl) "model.vlm.model.vision_model.encoder", "model.vlm_model.model.vision_model.encoder", ] for path in search_paths: obj = model parts = path.split(".") try: for part in parts: obj = getattr(obj, part) print(f" Found vision encoder at: policy.{path}") return obj except AttributeError: continue # Fallback: search by module type name (SigLIP, SmolVLM vision, etc.) print(" Searching for vision encoder by module type...") for name, module in model.named_modules(): type_name = type(module).__name__.lower() if "visionencoder" in type_name or "siglip" in type_name: print(f" Found vision encoder: {name} ({type(module).__name__})") return module if "visionmodel" in type_name and "text" not in type_name: print(f" Found vision model: {name} ({type(module).__name__})") return module return None def find_image_keys(dataset): """Find observation image keys in the dataset.""" sample = dataset[0] image_keys = [k for k in sample.keys() if "image" in k.lower()] if not image_keys: image_keys = [k for k in sample.keys() if "pixel" in k.lower()] if not image_keys: image_keys = [k for k in sample.keys() if isinstance(sample[k], torch.Tensor) and sample[k].dim() == 3 and sample[k].shape[0] == 3] return image_keys _image_key_warning_shown = False def _match_image_keys(policy_img_keys, dataset_img_keys): """ Match dataset image keys to policy image keys. Strategy (in priority order): 1. Exact match — dataset key exists in policy keys 2. Suffix match — last segment matches (e.g. both end in 'wrist') 3. Positional fallback — pair by sorted order (with warning) Returns: List of ``(dataset_key, policy_key)`` pairs. """ global _image_key_warning_shown mapping = [] unmatched_pkeys = list(policy_img_keys) unmatched_dkeys = list(dataset_img_keys) # Pass 1: exact match for dkey in list(unmatched_dkeys): if dkey in unmatched_pkeys: mapping.append((dkey, dkey)) unmatched_pkeys.remove(dkey) unmatched_dkeys.remove(dkey) # Pass 2: suffix match (last dotted segment, e.g. "camera1" or "wrist") for dkey in list(unmatched_dkeys): d_suffix = dkey.rsplit(".", 1)[-1] for pkey in list(unmatched_pkeys): p_suffix = pkey.rsplit(".", 1)[-1] if d_suffix == p_suffix: mapping.append((dkey, pkey)) unmatched_pkeys.remove(pkey) unmatched_dkeys.remove(dkey) break # Pass 3: positional fallback if unmatched_dkeys and unmatched_pkeys: positional = list(zip(sorted(unmatched_dkeys), sorted(unmatched_pkeys))) for dkey, pkey in positional: mapping.append((dkey, pkey)) unmatched_pkeys.remove(pkey) unmatched_dkeys.remove(dkey) if not _image_key_warning_shown: pairs = ", ".join(f"'{d}' -> '{p}'" for d, p in positional) print(f" WARNING: No name match for images — mapping by position: {pairs}. " f"Use --image-key if this is wrong.") _image_key_warning_shown = True return mapping def _resolve_task_string(sample, dataset=None): """ Get the task/language instruction for a sample. Priority: 1. ``sample["task"]`` — always present in LeRobot datasets 2. ``dataset.meta.tasks`` — first task in the dataset metadata 3. Generic fallback """ task = sample.get("task") if task is not None: if isinstance(task, list): task = task[0] return task # Try dataset metadata if dataset is not None: try: tasks_df = dataset.meta.tasks if len(tasks_df) > 0: return tasks_df.iloc[0].name except (AttributeError, IndexError): pass return "manipulate object" def build_policy_batch_from_sample(sample, policy, device, batch_size=1, image_key_for_grad=None, dataset=None): """ Build a batch dict that matches the policy's expected keys (e.g. observation.images.camera1), by mapping from the dataset sample keys (e.g. observation.images.up, observation.images.side). Policy expects config.image_features keys; dataset may use different names (up/side vs camera1/2/3). Also tokenizes the task string into ``observation.language.tokens`` and ``observation.language.attention_mask`` which ``select_action()`` requires. """ policy_img_keys = list(getattr(policy.config, "image_features", {})) if not policy_img_keys: # Policy has no image_features config; use sample keys as-is batch = {} for key in sample: val = sample[key] if isinstance(val, torch.Tensor): batch[key] = val.unsqueeze(0).to(device) if batch_size == 1 else val.to(device) elif isinstance(val, str): batch[key] = [val] else: batch[key] = val return batch, None dataset_img_keys = sorted([k for k in sample.keys() if "image" in k.lower() and isinstance(sample.get(k), torch.Tensor)]) if not dataset_img_keys: dataset_img_keys = sorted([k for k in sample.keys() if isinstance(sample.get(k), torch.Tensor) and sample[k].dim() >= 3 and sample[k].shape[0] == 3]) batch = {} for key in sample: if key in dataset_img_keys: continue # Fill with policy keys below val = sample[key] if isinstance(val, torch.Tensor): batch[key] = val.unsqueeze(0).to(device) if batch_size == 1 else val.to(device) elif isinstance(val, str): batch[key] = [val] else: batch[key] = val # Map dataset image keys -> policy image keys (exact > suffix > positional) img_mapping = _match_image_keys(policy_img_keys, dataset_img_keys) grad_pkey = None for dkey, pkey in img_mapping: img = sample[dkey] if batch_size == 1: img = img.unsqueeze(0).to(device) else: img = img.to(device) if image_key_for_grad is not None and dkey == image_key_for_grad: img = img.clone().detach().requires_grad_(True) batch[pkey] = img grad_pkey = pkey else: batch[pkey] = img.clone().detach().requires_grad_(False) # --- Resolve task string --- task_text = _resolve_task_string(sample, dataset) if "task" not in batch: batch["task"] = [task_text] # --- Tokenize the task description into language tokens --- # select_action() expects 'observation.language.tokens' and # 'observation.language.attention_mask' which come from tokenizing # the task string with the model's built-in tokenizer. lang_key = "observation.language.tokens" lang_mask_key = "observation.language.attention_mask" if lang_key not in batch: if isinstance(task_text, list): task_text = task_text[0] try: tokenizer = policy.model.vlm_with_expert.processor.tokenizer tok_out = tokenizer(task_text, return_tensors="pt", padding=True) batch[lang_key] = tok_out["input_ids"].to(device) batch[lang_mask_key] = tok_out["attention_mask"].to(device) except Exception as e: print(f" WARNING: Could not tokenize task string: {e}") return batch, grad_pkey def get_episode_frames(dataset, episode_idx, num_frames, image_key): """ Extract evenly-spaced frames from an episode. Returns list of (frame_index, image_tensor) tuples. """ # Get episode boundaries try: # LeRobot v3 format ep_from = dataset.meta.episodes["dataset_from_index"][episode_idx] ep_to = dataset.meta.episodes["dataset_to_index"][episode_idx] except (AttributeError, KeyError): try: # LeRobot v2 format ep_from = dataset.episode_data_index["from"][episode_idx].item() ep_to = dataset.episode_data_index["to"][episode_idx].item() except (AttributeError, KeyError): # Fallback: assume ~200 frames per episode ep_from = episode_idx * 200 ep_to = min(ep_from + 200, len(dataset)) ep_length = ep_to - ep_from if ep_length <= 0: raise ValueError(f"Episode {episode_idx} is empty (from={ep_from}, to={ep_to})") # Sample evenly spaced frames if num_frames >= ep_length: indices = list(range(ep_from, ep_to)) else: step = ep_length / num_frames indices = [int(ep_from + i * step) for i in range(num_frames)] frames = [] for idx in indices: sample = dataset[idx] img = sample[image_key] frames.append((idx, img)) print(f" Episode {episode_idx}: {ep_length} frames, sampled {len(frames)}") return frames # --------------------------------------------------------------------------- # 6. Main pipeline # --------------------------------------------------------------------------- def extract_attention_maps(policy, dataset, episode_idx=0, num_frames=8, image_key=None, device="cpu", method="last-layer", cross_attention=False, show_heads=False, output_dir="./outputs", raw_attention=False): """ Core function: Run inference and extract attention heatmaps. Args: method: ``"last-layer"`` (default), ``"rollout"``, or ``"all-layers"`` — how to aggregate SigLIP self-attention. cross_attention: If *True* also capture action-expert → vision cross-attention (requires full ``select_action`` forward pass, slower). show_heads: If *True* save a per-head grid for the first frame. output_dir: Where to save per-head grids (only used when *show_heads* is True). raw_attention: If *True* skip positional baseline subtraction (show raw, uncorrected attention). Returns: frames: list of image tensors (C, H, W) heatmaps: list of numpy heatmaps (H, W) in [0, 1] actions: list of predicted actions (or None) cross_attn_heatmaps: list of numpy heatmaps or *None* """ # --- Find vision encoder --- print("\n[1/4] Locating vision encoder...") vision_encoder = find_vision_encoder(policy) if vision_encoder is None: print(" Could not find vision encoder. Dumping model structure:") for name, module in policy.named_modules(): print(f" {name}: {type(module).__name__}") raise RuntimeError("Cannot find SigLIP vision encoder in model") # --- Force eager attention so we get attention weights (SDPA/Flash return None) --- eager_count = 0 for mod in vision_encoder.modules(): if getattr(mod, "config", None) is not None and hasattr(mod.config, "_attn_implementation"): mod.config._attn_implementation = "eager" eager_count += 1 if eager_count: print(" Using eager attention to capture weights") # --- Register attention hooks --- print("\n[2/4] Registering attention hooks...") attn_capture = SigLIPAttentionCapture() attn_capture.register_hooks(vision_encoder) # --- Optionally register cross-attention capture --- cross_capture = None vision_start = vision_end = n_vision_tokens = 0 if cross_attention: try: vlm_with_expert = policy.model.vlm_with_expert cross_capture = ActionVisionAttentionCapture() cross_capture.register(vlm_with_expert) vision_start, vision_end, n_vision_tokens = cross_capture.get_vision_token_range(policy) print(f" Cross-attention capture enabled (vision tokens {vision_start}..{vision_end}, " f"n={n_vision_tokens})") except Exception as e: print(f" WARNING: Could not set up cross-attention capture: {e}") cross_capture = None # --- Compute positional baseline (once) --- baseline_scores = None if not raw_attention: baseline_scores = compute_positional_baseline( vision_encoder, attn_capture, device, method, ) if baseline_scores is not None: print(" Positional baseline computed (subtracting to reveal content-dependent attention)") else: print(" Could not compute positional baseline, using raw attention") # --- Find image key in dataset --- if image_key is None: image_keys = find_image_keys(dataset) if not image_keys: raise ValueError("No image keys found in dataset. Available keys: " + str(list(dataset[0].keys()))) image_key = image_keys[0] print(f" Using image key: {image_key}") if len(image_keys) > 1: print(f" Other available image keys: {image_keys[1:]}") # --- Get frames --- print(f"\n[3/4] Loading {num_frames} frames from episode {episode_idx}...") frame_pairs = get_episode_frames(dataset, episode_idx, num_frames, image_key) # --- Log task string being used --- first_sample = dataset[frame_pairs[0][0]] task_str = _resolve_task_string(first_sample, dataset) print(f" Task string: \"{task_str}\"") # --- Run inference and collect attention --- print(f"\n[4/4] Running forward passes and extracting attention (method={method})...") frames = [] heatmaps = [] cross_attn_heatmaps = [] if cross_capture else None actions = [] policy.eval() # We need a full policy forward pass when cross-attention is requested # (the vision encoder hooks still fire during select_action too) use_full_forward = cross_attention and cross_capture is not None raw_heads_attn = None # populated when show_heads is True (first frame only) n_patches_h = n_patches_w = 0 # set during vision-encoder-only forward for i, (frame_idx, img_tensor) in enumerate(frame_pairs): frames.append(img_tensor.clone()) # Reset captured attention maps attn_capture.reset_maps() if cross_capture: cross_capture.reset_maps() sample = dataset[frame_idx] # ----- forward pass ----- try: with torch.no_grad(): if use_full_forward: # Full policy forward — needed for cross-attention capture. # Reset the action queue so each frame triggers a real # forward pass (otherwise cached actions are returned). policy.reset() batch, _ = build_policy_batch_from_sample( sample, policy, device, batch_size=1, image_key_for_grad=None, dataset=dataset, ) try: policy.select_action(batch) except Exception as e: print(f" select_action error: {e}") else: # Vision-encoder-only forward (faster) img = img_tensor.unsqueeze(0).to(device) if img.max() > 1.0: img = img.float() / 255.0 target_size = getattr( getattr(vision_encoder, "config", None), "image_size", None, ) or getattr(vision_encoder, "image_size", 384) if img.shape[-1] != target_size or img.shape[-2] != target_size: img_resized = F.interpolate(img, size=(target_size, target_size), mode="bilinear", align_corners=False) else: img_resized = img try: enc_dtype = next(vision_encoder.parameters()).dtype img_resized = img_resized.to(enc_dtype) except StopIteration: pass patch_size = getattr(vision_encoder, "patch_size", None) or getattr( getattr(vision_encoder, "config", None), "patch_size", 14 ) n_patches_h = img_resized.size(2) // patch_size n_patches_w = img_resized.size(3) // patch_size patch_mask = torch.ones(1, n_patches_h, n_patches_w, dtype=torch.bool, device=device) try: if hasattr(vision_encoder, 'embeddings') and hasattr(vision_encoder, 'encoder'): embeddings = vision_encoder.embeddings(img_resized, patch_mask) encoder_out = vision_encoder.encoder(embeddings) elif hasattr(vision_encoder, 'forward'): encoder_out = vision_encoder(img_resized) else: encoder_out = vision_encoder(pixel_values=img_resized, patch_attention_mask=patch_mask) except Exception as e: print(f" Direct vision forward failed ({e}), trying full policy...") batch, _ = build_policy_batch_from_sample( sample, policy, device, batch_size=1, image_key_for_grad=None, dataset=dataset, ) try: policy.select_action(batch) except Exception: pass except Exception as e: print(f" Frame {i} forward pass error: {e}") # ----- per-head capture (independent of summary method) ----- if show_heads and i == 0 and raw_heads_attn is None: attn_for_heads = attn_capture.get_last_layer_attention() if attn_for_heads is not None: while attn_for_heads.dim() > 3: attn_for_heads = attn_for_heads[0] if attn_for_heads.dim() == 3: raw_heads_attn = attn_for_heads.clone() # ----- self-attention heatmap ----- if method == "rollout": all_layers = attn_capture.get_all_layer_attentions() rollout_mat = compute_attention_rollout(all_layers) if rollout_mat is not None: patch_scores = rollout_mat.mean(dim=0) # per-patch importance else: patch_scores = None elif method == "all-layers": # For "all-layers" we still produce a single summary heatmap # (averaging last-layer scores) but also save individual layer # grids elsewhere; here fall through to last-layer for the # summary heatmap. attn = attn_capture.get_last_layer_attention() patch_scores = None if attn is not None: while attn.dim() > 3: attn = attn[0] if attn.dim() == 3: attn = attn.mean(dim=0) patch_scores = compute_patch_attention_scores(attn, method="mean") else: # "last-layer" (default) attn = attn_capture.get_last_layer_attention() patch_scores = None if attn is not None: while attn.dim() > 3: attn = attn[0] if attn.dim() == 3: attn = attn.mean(dim=0) patch_scores = compute_patch_attention_scores(attn, method="mean") # Subtract positional baseline to isolate content-dependent signal if patch_scores is not None and baseline_scores is not None: patch_scores = torch.clamp(patch_scores - baseline_scores, min=0) if patch_scores is not None: n_patches = patch_scores.shape[0] grid_side = int(math.sqrt(n_patches)) if grid_side * grid_side != n_patches: grid_h = n_patches_h if n_patches_h > 0 else grid_side grid_w = n_patches_w if n_patches_w > 0 else grid_side else: grid_h = grid_w = grid_side img_h, img_w = img_tensor.shape[1], img_tensor.shape[2] heatmap = attention_to_heatmap(patch_scores, (grid_h, grid_w), (img_h, img_w)) heatmaps.append(heatmap) print(f" Frame {i}: {n_patches} patches → " f"{grid_h}x{grid_w} grid → {img_h}x{img_w} heatmap ({method})") # Per-head grid for first frame if show_heads and i == 0 and raw_heads_attn is not None: head_path = os.path.join(output_dir, f"per_head_ep{episode_idx:03d}.png") create_per_head_grid( img_tensor, raw_heads_attn, (grid_h, grid_w), (img_h, img_w), output_path=head_path, ) raw_heads_attn = None # only once else: print(f" Frame {i}: No attention captured, using uniform heatmap") img_h, img_w = img_tensor.shape[1], img_tensor.shape[2] heatmaps.append(np.ones((img_h, img_w)) * 0.5) # ----- cross-attention heatmap ----- if cross_capture is not None: cross_scores = cross_capture.get_mean_cross_attention(vision_start, vision_end) if cross_scores is not None: # cross_scores: (n_vision_tokens,) # Vision tokens come from SigLIP → connector (pixel shuffle). # After pixel shuffle the spatial grid is halved in each dim. n_vis = cross_scores.shape[0] cs_side = int(math.sqrt(n_vis)) if cs_side * cs_side != n_vis: cs_h = cs_w = cs_side else: cs_h = cs_w = cs_side img_h, img_w = img_tensor.shape[1], img_tensor.shape[2] cross_hm = attention_to_heatmap(cross_scores, (cs_h, cs_w), (img_h, img_w)) cross_attn_heatmaps.append(cross_hm) print(f" Frame {i}: Cross-attention captured ({n_vis} vision tokens)") else: img_h, img_w = img_tensor.shape[1], img_tensor.shape[2] cross_attn_heatmaps.append(np.ones((img_h, img_w)) * 0.5) print(f" Frame {i}: No cross-attention captured, using uniform") # Cleanup attn_capture.clear() if cross_capture: cross_capture.clear() return frames, heatmaps, actions, cross_attn_heatmaps # --------------------------------------------------------------------------- # 7. Alternative: Pure gradient-based visualization (no hooks needed) # --------------------------------------------------------------------------- def gradient_attention_map(policy, dataset, frame_idx, image_key, device="cpu"): """ Compute input-gradient saliency map as a fallback. This doesn't require hooking into attention — it directly computes which input pixels most affect the output actions by backpropagating through the entire model. Interpretation: Bright pixels = changing this pixel would change the predicted action the most. """ sample = dataset[frame_idx] # Build batch using policy-expected image keys (camera1, camera2, ...) # so we don't get "All image features are missing" when dataset uses up/side. batch, grad_pkey = build_policy_batch_from_sample( sample, policy, device, batch_size=1, image_key_for_grad=image_key, dataset=dataset, ) if grad_pkey is None: # No policy image key matched; try legacy: use raw sample keys img = sample[image_key].unsqueeze(0).to(device).float() img.requires_grad_(True) batch = {k: v.unsqueeze(0).to(device) if isinstance(v, torch.Tensor) else ([v] if isinstance(v, str) else v) for k, v in sample.items()} batch[image_key] = img if "task" not in batch: batch["task"] = ["pick and place"] grad_tensor = img else: grad_tensor = batch[grad_pkey] try: policy.train() # Need gradients action = policy.select_action(batch) # Backpropagate from action norm if isinstance(action, dict): action_tensor = list(action.values())[0] elif isinstance(action, torch.Tensor): action_tensor = action else: return None loss = action_tensor.sum() loss.backward() # Saliency = absolute gradient magnitude across channels if grad_tensor.grad is None: return None saliency = grad_tensor.grad.abs().squeeze(0) if saliency.dim() == 3: saliency = saliency.mean(dim=0) # (H, W) saliency = saliency / (saliency.max() + 1e-8) return saliency.detach().cpu().numpy() except Exception as e: print(f" Gradient saliency failed: {e}") return None finally: policy.eval() # --------------------------------------------------------------------------- # 8. Entry point # --------------------------------------------------------------------------- def load_defaults(): """Load defaults from configs/defaults.yaml if it exists.""" config_path = Path(__file__).parent / "configs" / "defaults.yaml" if config_path.exists(): try: import yaml with open(config_path) as f: return yaml.safe_load(f) or {} except ImportError: pass return {} def main(): defaults = load_defaults() parser = argparse.ArgumentParser( description="See what SmolVLA's vision encoder is looking at.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: python inspect_attention.py python inspect_attention.py --episode 3 --num-frames 12 --device cuda python inspect_attention.py --model path/to/finetuned_checkpoint """ ) parser.add_argument("--model", type=str, default=defaults.get("model", "lerobot/smolvla_base")) parser.add_argument("--dataset", type=str, default=defaults.get("dataset", "lerobot/svla_so101_pickplace")) parser.add_argument("--episode", type=int, default=defaults.get("episode", 0)) parser.add_argument("--num-frames", type=int, default=defaults.get("num_frames", 8)) parser.add_argument("--image-key", type=str, default=None) parser.add_argument("--output-dir", type=str, default=defaults.get("output_dir", "./outputs")) parser.add_argument("--device", type=str, default=defaults.get("device", "auto"), choices=["auto", "cpu", "cuda", "mps"]) parser.add_argument("--save-individual", action="store_true", default=defaults.get("save_individual", False)) parser.add_argument("--method", type=str, default=defaults.get("method", "last-layer"), choices=["last-layer", "rollout", "all-layers"], help="Self-attention aggregation method") parser.add_argument("--cross-attention", action="store_true", default=defaults.get("cross_attention", False), help="Capture action-expert → vision cross-attention (slower)") parser.add_argument("--show-heads", action="store_true", default=defaults.get("show_heads", False), help="Save a per-head attention grid for the first frame") parser.add_argument("--raw-attention", action="store_true", default=defaults.get("raw_attention", False), help="Skip positional baseline subtraction (show raw attention)") args = parser.parse_args() os.makedirs(args.output_dir, exist_ok=True) # Auto-detect device if args.device == "auto": if torch.backends.mps.is_available(): args.device = "mps" elif torch.cuda.is_available(): args.device = "cuda" else: args.device = "cpu" device = torch.device(args.device) # ----------------------------------------------------------------------- print("=" * 70) print("SmolVLA Attention Visualizer") print("=" * 70) print(f" Device: {args.device}") # --- Load model --- print(f"\n[Step 1] Loading model: {args.model}") print(" This may download ~1GB on first run...") try: from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy # Suppress noisy warnings from HF/lerobot during model loading: # - "Device 'cuda' is not available. Switching to 'mps'" # - "`torch_dtype` is deprecated! Use `dtype` instead!" # - "Loading ... weights ..." _suppressed_loggers = { name: logging.getLogger(name) for name in ("lerobot.configs.policies", "lerobot", "transformers") } _saved_levels = {name: lg.level for name, lg in _suppressed_loggers.items()} for lg in _suppressed_loggers.values(): lg.setLevel(logging.ERROR) try: with warnings.catch_warnings(): warnings.filterwarnings("ignore", message=".*torch_dtype.*deprecated.*") policy = SmolVLAPolicy.from_pretrained(args.model) finally: for name, lg in _suppressed_loggers.items(): lg.setLevel(_saved_levels[name]) policy.to(device) policy.eval() print(f" Model loaded successfully ({sum(p.numel() for p in policy.parameters()) / 1e6:.1f}M params)") except ImportError: print("\n ERROR: LeRobot not installed. Run:") print(' pip install "lerobot[smolvla]"') sys.exit(1) except Exception as e: print(f"\n ERROR loading model: {e}") print(" Make sure the model ID is correct and you have internet access.") sys.exit(1) # --- Load dataset --- print(f"\n[Step 2] Loading dataset: {args.dataset}") print(" This may download several GB on first run...") try: from lerobot.common.datasets.lerobot_dataset import LeRobotDataset except ImportError: from lerobot.datasets.lerobot_dataset import LeRobotDataset try: dataset = LeRobotDataset(args.dataset) print(f" Dataset loaded: {len(dataset)} frames") # Show dataset info sample = dataset[0] print(f" Available keys: {list(sample.keys())}") image_keys = find_image_keys(dataset) print(f" Image keys found: {image_keys}") except Exception as e: print(f"\n ERROR loading dataset: {e}") sys.exit(1) # --- Extract attention maps --- print(f"\n[Step 3] Extracting attention maps...") cross_attn_heatmaps = None try: frames, heatmaps, actions, cross_attn_heatmaps = extract_attention_maps( policy=policy, dataset=dataset, episode_idx=args.episode, num_frames=args.num_frames, image_key=args.image_key, device=args.device, method=args.method, cross_attention=args.cross_attention, show_heads=args.show_heads, output_dir=args.output_dir, raw_attention=args.raw_attention, ) except Exception as e: print(f"\n Attention extraction failed: {e}") print(" Falling back to input-gradient saliency maps...") # Fallback to gradient-based saliency image_key = args.image_key or find_image_keys(dataset)[0] frame_pairs = get_episode_frames(dataset, args.episode, args.num_frames, image_key) frames = [] heatmaps = [] for frame_idx, img_tensor in frame_pairs: frames.append(img_tensor) saliency = gradient_attention_map(policy, dataset, frame_idx, image_key, args.device) if saliency is not None: heatmaps.append(saliency) else: heatmaps.append(np.ones((img_tensor.shape[1], img_tensor.shape[2])) * 0.5) actions = [] if not frames: print("\nERROR: No frames extracted. Check episode index and dataset.") sys.exit(1) # --- Generate visualizations --- print(f"\n[Step 4] Generating visualizations...") grid_path = os.path.join(args.output_dir, f"attention_grid_ep{args.episode:03d}.png") create_visualization_grid( frames=frames, heatmaps=heatmaps, actions=actions, cross_attn_heatmaps=cross_attn_heatmaps, episode_idx=args.episode, output_path=grid_path, ) if args.save_individual: save_individual_frames( frames=frames, heatmaps=heatmaps, output_dir=os.path.join(args.output_dir, f"episode_{args.episode:03d}"), episode_idx=args.episode, ) # --- Summary --- print(f"\n{'=' * 70}") print("DONE!") print(f"{'=' * 70}") print(f"\nOutputs saved to: {args.output_dir}/") print(f" Grid visualization: {grid_path}") if args.save_individual: print(f" Individual frames: {args.output_dir}/episode_{args.episode:03d}/") print() if __name__ == "__main__": main()