Add gradient-based attribution (saliency + GradCAM) for SmolVLA

Attention maps show where the model looks, but not which image regions
causally influence the predicted action. This adds two gradient methods
behind a new --gradient flag:

- Vanilla saliency (|dA/dpx|) at full pixel resolution
- GradCAM on SigLIP's last encoder layer at patch resolution (32x32)

Key implementation detail: select_action() is decorated with
@torch .no_grad(), blocking all gradient flow. The new gradient module
bypasses this by calling prepare_images/prepare_state/sample_actions
directly under torch.enable_grad(), with a monkey-patch to cast the
Long attention mask to bool (required by torch.where under autograd).

New files: smolvla_inspect/gradient.py
Modified: cli.py (--gradient flag, --gradient-device, --gradient-seed),
viz.py (dynamic row count up to 7, inferno/magma colormaps),
defaults.yaml, README.md (grid row reference table, interpretation guide)

README.md

@@ -9,10 +9,11 @@ See what SmolVLA's vision encoder and action expert are looking at when the mode
 
 ## What this does
 
-SmolVLA is a **vision-language-action** policy: it takes camera images and a language instruction, then outputs robot actions. This tool has two modes:
+SmolVLA is a **vision-language-action** policy: it takes camera images and a language instruction, then outputs robot actions. This tool has three modes:
 
 1. **Attention visualization** (default) -- extracts and visualizes attention heatmaps showing where the model looks
-2. **Model health diagnostics** (`--model-health`) -- runs spectral analysis, attention entropy, and head redundancy checks across all model components
+2. **Gradient-based attribution** (`--gradient`) -- computes saliency maps and GradCAM to show which pixels *causally influence* the predicted action
+3. **Model health diagnostics** (`--model-health`) -- runs spectral analysis, attention entropy, and head redundancy checks across all model components
 
 ### Attention visualization
 
@@ -26,6 +27,18 @@ That lets you check whether the model attends to task-relevant regions (gripper,
 **Input:** A pretrained or fine-tuned SmolVLA policy + a LeRobot dataset (e.g. episodes of pick-and-place).
 **Output:** A multi-row grid PNG per episode, optional per-frame PNGs, an optional per-head attention grid, and a positional baseline diagnostic (`positional_baseline.png`) showing the position-dependent attention pattern that gets subtracted.
 
+### Gradient-based attribution
+
+Attention maps show where the model *allocates compute*, but not whether those regions actually *drive the output*. Gradient attribution answers that question by backpropagating from the predicted action to the input pixels:
+
+1. **Saliency** (`--gradient saliency`) -- computes `|d(action) / d(pixel)|` at full pixel resolution. Highlights the raw input regions whose changes most affect the predicted action.
+2. **GradCAM** (`--gradient gradcam`) -- hooks the last SigLIP encoder layer, weights activations by their gradient, and produces a patch-resolution heatmap. Shows which high-level visual features drive the action.
+3. **Both** (`--gradient` or `--gradient both`) -- runs both methods and adds both rows to the output grid.
+
+Since gradient computation requires `.backward()` through the full model (~176 transformer layer passes), it is slower than attention-only mode. MPS backward support is limited for some ops, so you can run attention on MPS and gradients on CPU with `--gradient-device cpu` (see [Split device execution](#split-device-execution)).
+
+**Output:** Up to 2 additional rows in the grid PNG -- saliency overlay (inferno colormap) and GradCAM overlay (magma colormap).
+
 ### Model health diagnostics
 
 Runs three diagnostic checks across all model components (SigLIP vision encoder, VLM text model, action expert, connector, and projection heads):
@@ -66,17 +79,19 @@ For a detailed visual walkthrough of the architecture and how it maps to the rep
    ![Positional baseline](assets/example_positional_baseline.png)
    *Positional baseline diagnostic: attention pattern from a content-free gray image, showing position-dependent artifacts that get subtracted from real frames.*
 
-6. **Visualize** -- the output grid has up to 5 rows per frame:
+6. **Visualize** -- the output grid has up to 7 rows per frame:
 
-| Row | Content | Colormap |
-|-----|---------|----------|
-| 1 | Original frame | -- |
-| 2 | SigLIP self-attention heatmap | jet (blue-to-red) |
-| 3 | Self-attention overlay on frame | jet |
-| 4 | Action cross-attention heatmap | Greens |
-| 5 | Co-attention overlay (self x cross) | cyan (black-cyan-white) |
+| Row | Content | Colormap | When shown |
+|-----|---------|----------|------------|
+| 1 | Original frame | -- | always |
+| 2 | SigLIP self-attention heatmap | jet (blue-to-red) | always |
+| 3 | Self-attention overlay on frame | jet | always |
+| 4 | Action cross-attention heatmap | Greens | `--cross-attention` |
+| 5 | Co-attention overlay (self x cross) | cyan (black-cyan-white) | `--cross-attention` |
+| 6 | Saliency overlay (\|dA/dpx\|) | inferno | `--gradient saliency` or `both` |
+| 7 | GradCAM overlay (SigLIP last layer) | magma | `--gradient gradcam` or `both` |
 
-Rows 4-5 only appear when cross-attention is enabled. The co-attention overlay multiplies self-attention and cross-attention element-wise, highlighting regions that are **both** visually salient and action-relevant.
+Rows 4-5 only appear when cross-attention is enabled. The co-attention overlay multiplies self-attention and cross-attention element-wise, highlighting regions that are **both** visually salient and action-relevant. Rows 6-7 only appear when gradient attribution is enabled.
 
 ### Per-head grid
 
@@ -133,6 +148,9 @@ python inspect_attention.py
 # More frames, specific episode
 ./run.sh --episode 3 --num-frames 12
 
+# Override the language instruction (useful for multi-task datasets)
+./run.sh --task "pick up the red cube"
+
 # Last-layer method instead of rollout
 ./run.sh --method last-layer
 
@@ -151,6 +169,15 @@ python inspect_attention.py
 # No threshold (show all baseline-subtracted values)
 ./run.sh --attn-threshold 0
 
+# Gradient-based attribution: saliency + GradCAM
+./run.sh --gradient
+
+# Saliency only
+./run.sh --gradient saliency
+
+# Attention on MPS, gradients on CPU (avoids MPS backward issues)
+./run.sh --device mps --gradient both --gradient-device cpu
+
 # Model health diagnostics (spectral analysis + entropy + redundancy)
 ./run.sh --model-health
 
@@ -177,6 +204,7 @@ Results land in `outputs/`.
 | `--episode` | `0` | Episode index to visualize |
 | `--num-frames` | `8` | Number of frames to sample |
 | `--image-key` | auto-detected | Dataset image key override |
+| `--task` | from dataset | Override the language instruction |
 | `--output-dir` | `./outputs` | Output directory |
 | `--device` | `auto` | `auto`, `cpu`, `cuda`, or `mps` |
 | `--save-individual` | `true` | Save each frame as a separate PNG |
@@ -186,6 +214,14 @@ Results land in `outputs/`.
 | `--raw-attention` | `false` | Skip positional baseline subtraction |
 | `--attn-threshold` | `0.5` | Percentile (0-1) below which attention values are zeroed to suppress positional noise |
 
+**Gradient-based attribution:**
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--gradient` | off | `saliency`, `gradcam`, or `both` (bare `--gradient` means `both`) |
+| `--gradient-device` | same as `--device` | Device for gradient computation (`cpu`, `cuda`, `mps`) |
+| `--gradient-seed` | `42` | Fixed noise seed for reproducible gradient attribution |
+
 **Model health diagnostics:**
 
 | Flag | Default | Description |
@@ -226,6 +262,45 @@ Defaults can be changed in `configs/defaults.yaml`.
 
 The cyan overlay highlights regions where **both** the vision encoder and the action expert agree something is important. Bright cyan = high self-attention AND high cross-attention. This is the strongest signal for task-relevant regions.
 
+### Gradient attribution (saliency + GradCAM rows)
+
+| Pattern | Interpretation |
+|---------|----------------|
+| Saliency highlights gripper/object edges | **Healthy** -- action predictions are driven by task-relevant pixel changes |
+| GradCAM highlights same regions as self-attention | Attention and causal influence agree -- strong signal |
+| Attention focused but saliency diffuse | Attention points at the right place, but the action doesn't depend on it (possible shortcut) |
+| Saliency focused on unexpected region (e.g. table edge) | Model may be using a visual shortcut rather than task understanding |
+| GradCAM and saliency disagree | GradCAM operates at patch level while saliency is pixel-level -- some divergence is normal, but large disagreement warrants investigation |
+
+### Grid row reference
+
+| Row | Name | Question it answers | Resolution | Interpretation of hot spots |
+|-----|------|---------------------|------------|----------------------------|
+| SigLIP self-attn | Vision encoder attention | Which patches attend to each other inside the encoder? | 32x32 patches | Encoder's internal processing focus -- structural, not necessarily action-relevant |
+| Action cross-attn | Action-to-vision cross-attention | Which vision tokens does the action decoder query? | 8x8 tokens (post pixel-shuffle) | Visual regions the action decoder pulls information from |
+| Co-attention | Self x cross product | Which regions are both visually salient and action-queried? | 8x8 upsampled | Strongest attention signal for "what the model looks at to decide what to do" |
+| Saliency \|dA/dpx\| | Input-gradient saliency | If I changed this pixel, would the action change? | Full pixel (480x640) | Pixels that causally influence the predicted action -- fine-grained but noisy |
+| GradCAM SigLIP L-1 | Gradient-weighted activations | Which learned feature regions drive the action? | 32x32 patches | Patch regions whose features most influence the action -- coarser but more semantic |
+
+**Attention vs gradient:**
+
+|  | Attention rows | Gradient rows |
+|--|----------------|---------------|
+| Measures | Correlation ("model looked here") | Causation ("changing here changes the action") |
+| High attn + low gradient | Model looks but doesn't use it for action | -- |
+| Low attn + high gradient | -- | Region subtly influences output without dominating attention |
+| Both high | -- | Strong evidence this region genuinely drives behavior |
+
+### Split device execution
+
+Gradient computation requires `.backward()` through the full model, which can fail or be slow on MPS. Use `--gradient-device cpu` to run attention on MPS (fast, forward-only) and gradients on CPU (backward-compatible):
+
+```bash
+./run.sh --device mps --gradient both --gradient-device cpu
+```
+
+The model is moved to the gradient device after attention extraction finishes. Since gradients run last, there is no need to move it back.
+
 ### Per-head patterns
 
 Look for heads that specialize: one head tracking the gripper, another tracking the object, another attending globally. Specialization is a sign of a well-trained encoder. Heads that all look identical suggest the model hasn't learned diverse attention strategies.
@@ -254,22 +329,32 @@ See **[Architecture Diagrams](assets/architecture.md)** for visual explanations
 
 ```
 smolvla-inspect/
-├── inspect_attention.py      # All logic: model loading, hooks, heatmaps, health diagnostics
+├── inspect_attention.py        # Thin entry point (delegates to smolvla_inspect)
+├── smolvla_inspect/            # Main package
+│   ├── __init__.py
+│   ├── cli.py                  # CLI args, orchestration, attention extraction
+│   ├── capture.py              # Attention hook classes (SigLIP, cross-attention)
+│   ├── heatmap.py              # Patch scores, rollout, positional baseline, upsampling
+│   ├── gradient.py             # Gradient attribution (saliency, GradCAM)
+│   ├── data.py                 # Dataset helpers, batch building, image key mapping
+│   ├── viz.py                  # Visualization grid, overlays, per-head grids
+│   ├── health.py               # Model health diagnostics (spectral, entropy, redundancy)
+│   └── _compat.py              # Resize/pad compatibility helpers
 ├── assets/
-│   ├── architecture.md       # Architecture diagrams and report reference
+│   ├── architecture.md         # Architecture diagrams and report reference
 │   ├── how_it_works_architecture.png
 │   ├── example_grid.png
 │   └── example_per_head.png
 ├── configs/
-│   └── defaults.yaml         # Default CLI values (model, dataset, method, flags)
+│   └── defaults.yaml           # Default CLI values (model, dataset, method, flags)
 ├── docs/
-│   ├── ELI5.md               # Plain-language explanation of the interpretability approach
-│   └── TESTING.md            # CLI test commands and expected output
-├── outputs/                  # Generated images and reports (gitignored)
+│   ├── ELI5.md                 # Plain-language explanation of the interpretability approach
+│   └── TESTING.md              # CLI test commands and expected output
+├── outputs/                    # Generated images and reports (gitignored)
 │   ├── positional_baseline.png
 │   ├── model_health_report.md
 │   └── model_health_report.png
-├── run.sh                    # Wrapper that sets FFmpeg lib path
+├── run.sh                      # Wrapper that sets FFmpeg lib path
 ├── requirements.txt
 └── README.md
 ```
@@ -280,7 +365,7 @@ smolvla-inspect/
 
 Attention maps show where the model allocates compute, but not whether those regions actually drive the output. The following interpretability methods would complement the current tooling:
 
-- [ ] **Gradient-based attribution** -- compute `d(action) / d(patch_embedding)` via vanilla saliency, GradCAM, or Integrated Gradients to measure which image patches *causally influence* the predicted action (not just where attention points)
+- [x] **Gradient-based attribution** -- vanilla saliency (`|d(action)/d(pixel)|`) and GradCAM on SigLIP's last encoder layer, with split device support (`--gradient-device`) for running gradients on CPU while attention runs on MPS/CUDA
 - [ ] **Occlusion / perturbation sensitivity** -- mask out image regions or zero out specific prefix tokens (vision, language, state) and measure action MSE change; model-agnostic and directly answers "if I cover the gripper, does the model break?"
 - [ ] **Representation probing** -- train small linear classifiers on intermediate layer representations to test what information is encoded at each stage (e.g., can layer N predict object position? does the Expert encode gripper state?)
 - [ ] **Causal tracing / activation patching** -- replace activations at specific (layer, token) positions with corrupted versions and measure output change; builds a causal map of information flow through the model

configs/defaults.yaml

@@ -16,6 +16,11 @@ show_heads: true                        # true to save a per-head attention grid
 raw_attention: false                    # true to skip positional baseline subtraction
 attn_threshold: 0.5                    # percentile (0-1) below which attention values are zeroed
 
+# Gradient-based attribution
+gradient: null                         # null | saliency | gradcam | both
+gradient_device: cpu                  # null (same as device) | cpu | cuda | mps
+gradient_seed: 42                      # fixed noise seed for reproducible gradient attribution
+
 # Model health analysis
 model_health: false
 health_frames: 5

smolvla_inspect/cli.py

@@ -40,6 +40,7 @@ from .viz import (
     overlay_heatmap,
 )
 from .health import run_model_health_report
+from .gradient import compute_gradient_maps, compute_saliency_map
 
 
 def extract_attention_maps(policy, dataset, episode_idx=0, num_frames=8,
@@ -389,62 +390,15 @@ 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.
+    Delegates to :func:`gradient.compute_saliency_map` which bypasses the
+    ``@torch.no_grad()`` on ``select_action()`` by calling internal model
+    methods directly under ``torch.enable_grad()``.
     """
     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, task_override=task_override,
+    return compute_saliency_map(
+        policy, sample, dataset, image_key, device,
+        task_override=task_override,
     )
-    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()
 
 
 def load_defaults():
@@ -532,6 +486,19 @@ Examples:
                         default=defaults.get("redundancy_critical", 0.9),
                         help="Cosine similarity threshold for 'collapsed' (default: 0.9)")
 
+    # Gradient-based attribution
+    parser.add_argument("--gradient", nargs="?", const="both",
+                        default=defaults.get("gradient", None),
+                        choices=["saliency", "gradcam", "both"],
+                        help="Gradient attribution method (default: off; bare --gradient means 'both')")
+    parser.add_argument("--gradient-device", type=str,
+                        default=defaults.get("gradient_device", None),
+                        choices=["cpu", "cuda", "mps"],
+                        help="Device for gradient attribution (default: same as --device)")
+    parser.add_argument("--gradient-seed", type=int,
+                        default=defaults.get("gradient_seed", 42),
+                        help="Fixed noise seed for reproducible gradient attribution (default: 42)")
+
     args = parser.parse_args()
 
     os.makedirs(args.output_dir, exist_ok=True)
@@ -546,11 +513,17 @@ Examples:
             args.device = "cpu"
     device = torch.device(args.device)
 
+    # Resolve gradient device (defaults to main device)
+    grad_device_str = args.gradient_device or args.device
+    grad_device = torch.device(grad_device_str)
+
     # -----------------------------------------------------------------------
     print("=" * 70)
     print("SmolVLA Attention Visualizer")
     print("=" * 70)
     print(f"  Device: {args.device}")
+    if args.gradient and grad_device_str != args.device:
+        print(f"  Gradient device: {grad_device_str}")
 
     # --- Load model ---
     print(f"\n[Step 1] Loading model: {args.model}")
@@ -660,6 +633,34 @@ Examples:
         print("\nERROR: No frames extracted. Check episode index and dataset.")
         sys.exit(1)
 
+    # --- Gradient-based attribution (after attention hooks are cleaned up) ---
+    saliency_maps = None
+    gradcam_maps = None
+    if args.gradient:
+        image_key_for_grad = args.image_key or find_image_keys(dataset)[0]
+
+        # Move model to gradient device if different from main device
+        if grad_device != device:
+            print(f"\n  Moving model from {device} to {grad_device} for gradient computation...")
+            policy.to(grad_device)
+
+        print(f"\n[Step 3b] Computing gradient attribution (method={args.gradient}, device={grad_device_str})...")
+        saliency_maps, gradcam_maps = compute_gradient_maps(
+            policy=policy,
+            dataset=dataset,
+            episode_idx=args.episode,
+            num_frames=args.num_frames,
+            image_key=image_key_for_grad,
+            device=grad_device_str,
+            method=args.gradient,
+            noise_seed=args.gradient_seed,
+            task_override=args.task,
+        )
+        if saliency_maps:
+            print(f"  Saliency maps: {len(saliency_maps)} frames")
+        if gradcam_maps:
+            print(f"  GradCAM maps: {len(gradcam_maps)} frames")
+
     # --- Generate visualizations ---
     print(f"\n[Step 4] Generating visualizations...")
 
@@ -669,6 +670,8 @@ Examples:
         heatmaps=heatmaps,
         actions=actions,
         cross_attn_heatmaps=cross_attn_heatmaps,
+        saliency_maps=saliency_maps,
+        gradcam_maps=gradcam_maps,
         episode_idx=args.episode,
         output_path=grid_path,
     )

smolvla_inspect/gradient.py

@@ -0,0 +1,314 @@
+"""
+Gradient-based attribution — saliency maps and GradCAM for SmolVLA.
+
+Bypasses ``select_action()``'s ``@torch.no_grad()`` decorator by calling
+internal methods (``prepare_images``, ``prepare_state``, ``sample_actions``)
+directly under ``torch.enable_grad()``.
+"""
+
+import contextlib
+import warnings
+
+import numpy as np
+import torch
+
+from .data import build_policy_batch_from_sample, get_episode_frames, _resolve_task_string
+
+
+# ---------------------------------------------------------------------------
+# Monkey-patch: cast Long attention masks to bool for torch.where under grad
+# ---------------------------------------------------------------------------
+
+@contextlib.contextmanager
+def _patch_eager_attention_bool_mask(policy):
+    """
+    ``eager_attention_forward`` uses ``torch.where(mask, ...)`` where
+    *mask* is Long.  Under ``torch.no_grad()`` PyTorch silently casts it,
+    but with gradients enabled the dtype mismatch raises.  This context
+    manager monkey-patches the method to add ``.bool()`` on the mask.
+    """
+    vlm_expert = getattr(policy.model, "vlm_with_expert", None)
+    if vlm_expert is None:
+        yield
+        return
+
+    orig_fn = vlm_expert.eager_attention_forward
+
+    def _patched(attention_mask, *args, **kwargs):
+        if attention_mask.dtype != torch.bool:
+            attention_mask = attention_mask.bool()
+        return orig_fn(attention_mask, *args, **kwargs)
+
+    vlm_expert.eager_attention_forward = _patched
+    try:
+        yield
+    finally:
+        vlm_expert.eager_attention_forward = orig_fn
+
+
+# ---------------------------------------------------------------------------
+# Internal: forward pass with gradient graph retained
+# ---------------------------------------------------------------------------
+
+def _run_forward_with_grad(policy, batch, device, noise_seed=42):
+    """
+    Replicate ``SmolVLAPolicy._get_action_chunk()`` without the
+    ``@torch.no_grad()`` wrapper so that gradients flow back to inputs.
+
+    Returns:
+        action_scalar: scalar tensor suitable for ``.backward()``
+        (the sum of the first predicted action step)
+    """
+    with torch.enable_grad(), _patch_eager_attention_bool_mask(policy):
+        images, img_masks = policy.prepare_images(batch)
+        state = policy.prepare_state(batch)
+
+        lang_tokens = batch["observation.language.tokens"]
+        lang_masks = batch["observation.language.attention_mask"]
+
+        # Fixed noise for reproducibility across frames
+        bsize = state.shape[0]
+        actions_shape = (
+            bsize,
+            policy.model.config.chunk_size,
+            policy.model.config.max_action_dim,
+        )
+        gen = torch.Generator(device=device)
+        gen.manual_seed(noise_seed)
+        noise = torch.randn(actions_shape, device=device, generator=gen, dtype=torch.float32)
+
+        actions = policy.model.sample_actions(
+            images, img_masks, lang_tokens, lang_masks, state, noise=noise,
+        )
+
+        # Backprop target: sum of the immediate next action (step 0)
+        action_scalar = actions[:, 0, :].sum()
+
+    return action_scalar
+
+
+# ---------------------------------------------------------------------------
+# Saliency map (vanilla input-gradient)
+# ---------------------------------------------------------------------------
+
+def compute_saliency_map(policy, sample, dataset, image_key, device,
+                         noise_seed=42, task_override=None):
+    """
+    Compute ``|d(action)/d(pixel)|`` at pixel resolution.
+
+    Returns:
+        numpy array ``(H, W)`` in ``[0, 1]``, or *None* on failure.
+    """
+    batch, grad_pkey = build_policy_batch_from_sample(
+        sample, policy, device, batch_size=1,
+        image_key_for_grad=image_key, dataset=dataset,
+        task_override=task_override,
+    )
+    if grad_pkey is None:
+        print("    WARNING: Could not identify gradient image tensor")
+        return None
+
+    grad_tensor = batch[grad_pkey]
+
+    try:
+        policy.reset()
+        action_scalar = _run_forward_with_grad(policy, batch, device, noise_seed)
+        action_scalar.backward()
+
+        if grad_tensor.grad is None:
+            print("    WARNING: grad is None — gradient did not flow to input pixels")
+            return None
+
+        saliency = grad_tensor.grad.abs().squeeze(0)  # (C, H, W)
+        if saliency.dim() == 3:
+            saliency = saliency.mean(dim=0)  # (H, W)
+        saliency = saliency / (saliency.max() + 1e-8)
+        return saliency.detach().cpu().numpy()
+
+    except torch.cuda.OutOfMemoryError:
+        torch.cuda.empty_cache()
+        print("    WARNING: CUDA OOM during saliency — skipping this frame")
+        return None
+    except RuntimeError as e:
+        if "MPS" in str(e) or "mps" in str(e):
+            print(f"    WARNING: MPS backward error: {e}")
+            print("    Consider using --gradient-device cpu")
+        else:
+            print(f"    WARNING: Saliency computation failed: {e}")
+        return None
+
+
+# ---------------------------------------------------------------------------
+# GradCAM on SigLIP last encoder layer
+# ---------------------------------------------------------------------------
+
+def compute_gradcam_map(policy, sample, dataset, image_key, device,
+                        noise_seed=42, task_override=None):
+    """
+    Gradient-weighted class activation map at patch resolution (32x32).
+
+    Hooks the last SigLIP encoder layer to capture activations and their
+    gradients, then computes GradCAM weights.
+
+    Returns:
+        numpy array ``(grid_h, grid_w)`` in ``[0, 1]``, or *None* on failure.
+    """
+    from .data import find_vision_encoder
+
+    vision_encoder = find_vision_encoder(policy)
+    if vision_encoder is None:
+        print("    WARNING: Could not find vision encoder for GradCAM")
+        return None
+
+    # Hook last encoder layer
+    try:
+        last_layer = vision_encoder.encoder.layers[-1]
+    except (AttributeError, IndexError):
+        print("    WARNING: Could not access last encoder layer for GradCAM")
+        return None
+
+    activations = {}
+    gradients = {}
+
+    def fwd_hook(module, input, output):
+        # output is typically a tuple; first element is hidden states
+        out = output[0] if isinstance(output, tuple) else output
+        activations["value"] = out
+
+    def bwd_hook(module, grad_input, grad_output):
+        gradients["value"] = grad_output[0]
+
+    fwd_handle = last_layer.register_forward_hook(fwd_hook)
+    bwd_handle = last_layer.register_full_backward_hook(bwd_hook)
+
+    try:
+        batch, _ = build_policy_batch_from_sample(
+            sample, policy, device, batch_size=1,
+            image_key_for_grad=image_key, dataset=dataset,
+            task_override=task_override,
+        )
+
+        policy.reset()
+        action_scalar = _run_forward_with_grad(policy, batch, device, noise_seed)
+        action_scalar.backward()
+
+        if "value" not in activations or "value" not in gradients:
+            print("    WARNING: GradCAM hooks did not fire")
+            return None
+
+        A = activations["value"]   # (B, n_patches, hidden_dim)
+        dA = gradients["value"]    # same shape
+
+        # GAP over patches → per-channel weight
+        alpha = dA.mean(dim=1, keepdim=True)  # (B, 1, hidden_dim)
+
+        # Weighted combination + ReLU
+        cam = (alpha * A).sum(dim=-1)  # (B, n_patches)
+        cam = torch.relu(cam)
+        cam = cam.squeeze(0)  # (n_patches,)
+
+        # Reshape to spatial grid
+        n_patches = cam.shape[0]
+        grid_side = int(n_patches ** 0.5)
+        if grid_side * grid_side != n_patches:
+            # Non-square — try to infer from vision encoder config
+            img_size = getattr(getattr(vision_encoder, "config", None), "image_size", None) or 512
+            patch_size = getattr(getattr(vision_encoder, "config", None), "patch_size", None) or 16
+            grid_h = img_size // patch_size
+            grid_w = grid_h
+        else:
+            grid_h = grid_w = grid_side
+
+        cam_2d = cam[:grid_h * grid_w].reshape(grid_h, grid_w)
+        cam_2d = cam_2d / (cam_2d.max() + 1e-8)
+        return cam_2d.detach().float().cpu().numpy()
+
+    except torch.cuda.OutOfMemoryError:
+        torch.cuda.empty_cache()
+        print("    WARNING: CUDA OOM during GradCAM — skipping this frame")
+        return None
+    except RuntimeError as e:
+        if "MPS" in str(e) or "mps" in str(e):
+            print(f"    WARNING: MPS backward error: {e}")
+            print("    Consider using --gradient-device cpu")
+        else:
+            print(f"    WARNING: GradCAM computation failed: {e}")
+        return None
+    finally:
+        fwd_handle.remove()
+        bwd_handle.remove()
+
+
+# ---------------------------------------------------------------------------
+# Top-level: compute gradient maps for all frames
+# ---------------------------------------------------------------------------
+
+def compute_gradient_maps(policy, dataset, episode_idx, num_frames, image_key,
+                          device, method="both", noise_seed=42,
+                          task_override=None):
+    """
+    Compute saliency and/or GradCAM maps for a set of episode frames.
+
+    Args:
+        method: ``"saliency"``, ``"gradcam"``, or ``"both"``
+
+    Returns:
+        ``(saliency_maps, gradcam_maps)`` — each is a list of numpy arrays
+        or *None* if that method was not requested / all frames failed.
+    """
+    # MPS warning
+    if str(device) == "mps":
+        print("  NOTE: MPS backward support is limited. If gradient computation "
+              "fails, try --gradient-device cpu")
+
+    do_saliency = method in ("saliency", "both")
+    do_gradcam = method in ("gradcam", "both")
+
+    frame_pairs = get_episode_frames(dataset, episode_idx, num_frames, image_key)
+
+    saliency_maps = [] if do_saliency else None
+    gradcam_maps = [] if do_gradcam else None
+
+    for i, (frame_idx, img_tensor) in enumerate(frame_pairs):
+        sample = dataset[frame_idx]
+
+        if do_saliency:
+            policy.zero_grad()
+            smap = compute_saliency_map(
+                policy, sample, dataset, image_key, device,
+                noise_seed=noise_seed, task_override=task_override,
+            )
+            if smap is not None:
+                saliency_maps.append(smap)
+                print(f"    Frame {i}: Saliency computed ({smap.shape})")
+            else:
+                # Uniform fallback
+                h, w = img_tensor.shape[1], img_tensor.shape[2]
+                saliency_maps.append(np.ones((h, w)) * 0.5)
+                print(f"    Frame {i}: Saliency failed, using uniform")
+
+        if do_gradcam:
+            policy.zero_grad()
+            gcam = compute_gradcam_map(
+                policy, sample, dataset, image_key, device,
+                noise_seed=noise_seed, task_override=task_override,
+            )
+            if gcam is not None:
+                gradcam_maps.append(gcam)
+                print(f"    Frame {i}: GradCAM computed ({gcam.shape})")
+            else:
+                h, w = img_tensor.shape[1], img_tensor.shape[2]
+                gradcam_maps.append(np.ones((h, w)) * 0.5)
+                print(f"    Frame {i}: GradCAM failed, using uniform")
+
+    # If all frames failed, return None instead of list of uniforms
+    if saliency_maps is not None and all(
+        (m == 0.5).all() if isinstance(m, np.ndarray) else False for m in saliency_maps
+    ):
+        saliency_maps = None
+    if gradcam_maps is not None and all(
+        (m == 0.5).all() if isinstance(m, np.ndarray) else False for m in gradcam_maps
+    ):
+        gradcam_maps = None
+
+    return saliency_maps, gradcam_maps

smolvla_inspect/viz.py

@@ -72,22 +72,32 @@ def _frame_to_np(frame):
 
 def create_visualization_grid(frames, heatmaps, actions=None,
                                cross_attn_heatmaps=None,
+                               saliency_maps=None, gradcam_maps=None,
                                episode_idx=0, output_path="attention_grid.png"):
     """
     Create a grid visualization showing original frames, heatmaps, and overlays.
 
-    Layout per frame:
+    Layout per frame (dynamic rows):
       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)
+      Row 4-5 (optional): Cross-attention heatmap + co-attention overlay
+      +1 row (optional): Saliency overlay (|dA/dpx|)
+      +1 row (optional): GradCAM overlay (SigLIP last layer)
     """
     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
+    has_saliency = saliency_maps is not None and len(saliency_maps) == n_frames
+    has_gradcam = gradcam_maps is not None and len(gradcam_maps) == n_frames
+
+    # Dynamic row count: base 3 + 2 cross + 1 saliency + 1 gradcam = 7 max
+    n_rows = 3
+    if has_cross:
+        n_rows += 2
+    if has_saliency:
+        n_rows += 1
+    if has_gradcam:
+        n_rows += 1
 
     fig = plt.figure(figsize=(4 * n_frames, 4 * n_rows))
     gs = gridspec.GridSpec(n_rows, n_frames, hspace=0.3, wspace=0.05)
@@ -99,71 +109,106 @@ def create_visualization_grid(frames, heatmaps, actions=None,
         heatmap_resized = _resize_heatmap(heatmap, h, w)
         overlay = overlay_heatmap(frame_np, heatmap_resized, alpha=0.45)
 
+        row = 0
+
         # 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")
+        ax = fig.add_subplot(gs[row, i])
+        ax.imshow(frame_np)
+        ax.set_title(f"Frame {i}", fontsize=10)
+        ax.axis("off")
         if i == 0:
-            ax1.set_ylabel("Original", fontsize=11, rotation=0, labelpad=60, va="center")
+            ax.set_ylabel("Original", fontsize=11, rotation=0, labelpad=60, va="center")
+        row += 1
 
         # 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")
+        ax = fig.add_subplot(gs[row, i])
+        ax.imshow(heatmap_resized, cmap="jet", vmin=0, vmax=1)
+        ax.axis("off")
         if i == 0:
-            ax2.set_ylabel("SigLIP\nself-attn", fontsize=11, rotation=0, labelpad=60, va="center")
+            ax.set_ylabel("SigLIP\nself-attn", fontsize=11, rotation=0, labelpad=60, va="center")
+        row += 1
 
         # Row 3: Self-attention overlay
-        ax3 = fig.add_subplot(gs[2, i])
-        ax3.imshow(overlay)
-        ax3.axis("off")
+        ax = fig.add_subplot(gs[row, i])
+        ax.imshow(overlay)
+        ax.axis("off")
         if i == 0:
-            ax3.set_ylabel("Self-attn\noverlay", fontsize=11, rotation=0, labelpad=60, va="center")
+            ax.set_ylabel("Self-attn\noverlay", fontsize=11, rotation=0, labelpad=60, va="center")
+        row += 1
 
         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")
+            # Cross-attention heatmap
+            ax = fig.add_subplot(gs[row, i])
+            ax.imshow(cross_hm, cmap="Greens", vmin=0, vmax=1)
+            ax.axis("off")
             if i == 0:
-                ax4.set_ylabel("Action\ncross-attn", fontsize=11, rotation=0, labelpad=60, va="center")
+                ax.set_ylabel("Action\ncross-attn", fontsize=11, rotation=0, labelpad=60, va="center")
+            row += 1
 
-            # 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-attention overlay (self-attn × cross-attn)
+            co_attn = heatmap_resized * cross_hm
+            co_attn = co_attn / (co_attn.max() + 1e-8)
             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")
+            ax = fig.add_subplot(gs[row, i])
+            ax.imshow(co_overlay)
+            ax.axis("off")
             if i == 0:
-                ax5.set_ylabel("Co-attention\noverlay", fontsize=11, rotation=0, labelpad=60, va="center")
+                ax.set_ylabel("Co-attention\noverlay", fontsize=11, rotation=0, labelpad=60, va="center")
+            row += 1
+
+        if has_saliency:
+            sal_hm = _resize_heatmap(saliency_maps[i], h, w)
+            sal_overlay = overlay_heatmap(frame_np, sal_hm, alpha=0.45, colormap="inferno")
+
+            ax = fig.add_subplot(gs[row, i])
+            ax.imshow(sal_overlay)
+            ax.axis("off")
+            if i == 0:
+                ax.set_ylabel("Saliency\n|dA/dpx|", fontsize=11, rotation=0, labelpad=60, va="center")
+            row += 1
+
+        if has_gradcam:
+            gc_hm = _resize_heatmap(gradcam_maps[i], h, w)
+            gc_overlay = overlay_heatmap(frame_np, gc_hm, alpha=0.45, colormap="magma")
 
+            ax = fig.add_subplot(gs[row, i])
+            ax.imshow(gc_overlay)
+            ax.axis("off")
+            if i == 0:
+                ax.set_ylabel("GradCAM\nSigLIP L-1", fontsize=11, rotation=0, labelpad=60, va="center")
+            row += 1
+
+    # --- Build legend string dynamically ---
+    legend_parts = [
+        "Original",
+        "SigLIP self-attn heatmap",
+        "Self-attn overlay",
+    ]
     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}"
-        )
+        legend_parts.append("Action cross-attn heatmap")
+        legend_parts.append("Co-attention (self \u00d7 cross)")
+    if has_saliency:
+        legend_parts.append("Saliency |dA/dpx|")
+    if has_gradcam:
+        legend_parts.append("GradCAM SigLIP L-1")
+
+    legend = "  |  ".join(f"Row {j+1}: {lbl}" for j, lbl in enumerate(legend_parts))
+
+    extra_lines = []
+    if has_cross:
+        extra_lines.append("Co-attention: self-attn \u00d7 cross-attn \u2014 bright regions are both visually salient and action-relevant")
+    if has_saliency or has_gradcam:
+        extra_lines.append("Gradient rows show which image regions causally influence the predicted action")
+
+    title = f"SmolVLA Attention \u2014 Episode {episode_idx}\n\n{legend}"
+    if extra_lines:
+        title += "\n\n" + "\n".join(extra_lines)
     fig.suptitle(title, fontsize=14, fontweight="bold", y=0.98)
 
     plt.savefig(output_path, dpi=150, bbox_inches="tight", facecolor="white")