Update README to reflect current visualization pipeline

- Add example grid and per-head images at the top
- Document all 5 grid rows with colormaps (jet, Greens, cyan)
- Add CLI flags table with current defaults (rollout, cross-attention on)
- Add co-attention interpretation section
- Update architecture details (12 layers, 12 heads, 32x32 patches)
- Remove outdated references to dual-color overlay and gradient fallback

README.md

@@ -1,56 +1,70 @@
 # smolvla-inspect
 
-See what SmolVLA’s vision encoder is looking at when it predicts robot actions.
+See what SmolVLA's vision encoder and action expert are looking at when the model predicts robot actions.
+
+![Example attention grid](assets/example_grid.png)
+*5-row attention grid for a pick-and-place episode. Row 1: original frames. Row 2: SigLIP self-attention heatmap. Row 3: self-attention overlay. Row 4: action cross-attention heatmap. Row 5: co-attention (self x cross) overlay in cyan.*
 
 ---
 
 ## What this does
 
-SmolVLA is a **vision–language–action** policy: it takes camera images (and optionally language) and outputs robot actions. This repo **visualizes where the vision part “looks”** in each frame by extracting **attention maps** from the vision encoder and turning them into heatmaps over the images.
+SmolVLA is a **vision-language-action** policy: it takes camera images and a language instruction, then outputs robot actions. This tool **visualizes where the model looks** by extracting attention maps from two places:
+
+1. **SigLIP vision encoder** (self-attention) -- which image patches the encoder considers important during feature extraction
+2. **Action expert** (cross-attention) -- which image regions the action decoder actually reads when predicting actions
 
-That lets you check whether the model is attending to task-relevant regions (e.g. gripper, object, goal) or to background (e.g. walls, table texture)—a useful signal for debugging overfitting or distribution shift.
+That lets you check whether the model attends to task-relevant regions (gripper, object, goal) or background (walls, table texture) -- useful for debugging overfitting or distribution shift.
 
-**Input:** A pretrained or fine-tuned SmolVLA policy + a LeRobot dataset (e.g. episodes of pick-and-place).  
-**Output:** For each chosen frame, (1) the original image, (2) an attention heatmap, and (3) an overlay of the heatmap on the image, plus a grid summary PNG.
+**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, and an optional per-head attention grid.
 
 ---
 
-## How it works (for data scientists)
+## How it works
 
-![Architecture and attention-to-heatmap pipeline](assets/how_it_works_architecture.png)  
+![Architecture and attention-to-heatmap pipeline](assets/how_it_works_architecture.png)
 *Left: Where the vision encoder lives and where we hook to capture attention. Right: How attention weights become a spatial heatmap.*
 
-High-level pipeline:
+### Pipeline
+
+1. **Load model and dataset** -- loads a SmolVLA policy (e.g. `lerobot/smolvla_base`) and a LeRobot dataset. The dataset provides image sequences; the script uses images as input and runs inference.
+
+2. **Capture self-attention from SigLIP** -- the vision encoder (SigLIP ViT, 12 layers, 12 heads) splits each image into patches (32x32 grid for 512px images with 16px patches) and runs self-attention. Forward hooks on the attention layers capture the weight matrices.
+
+3. **Aggregate across layers** (`--method`):
+   - `last-layer` -- uses only the final encoder layer
+   - `rollout` (default) -- multiplies attention across all layers with residual connections, giving a more complete picture of information flow
+   - `all-layers` -- keeps each layer separately
 
-1. **Load model and dataset**  
-   The script loads a SmolVLA policy (e.g. `lerobot/smolvla_base`) and a LeRobot dataset. The dataset provides sequences of images and actions; we use the images as input and (optionally) actions only for reference.
+4. **Capture cross-attention** (`--cross-attention`, on by default) -- SmolVLA's VLM builds a KV cache from the prefix (vision + language + state tokens). The action expert queries that cache. The script monkey-patches `eager_attention_forward()` on the expert layers to intercept the softmax attention when expert Q attends to prefix K. Only columns corresponding to vision tokens are kept, giving a heatmap of which image regions the action decoder reads.
 
-2. **Locate the vision encoder**  
-   SmolVLA is built from a **vision encoder** (SmolVLM/SigLIP-style Vision Transformer) plus a language/action head. The script finds the vision encoder inside the policy (e.g. `model.vlm_with_expert.vlm.model.vision_model`) so it can hook into it.
+5. **Turn attention into spatial heatmaps** -- patch-level importance scores are reshaped into a 2D grid, upsampled with bilinear interpolation to image size, and normalized to [0, 1].
 
-3. **Run a forward pass and capture self-attention**
-   The vision encoder is a **Transformer**: it splits the image into **patches**, runs **self-attention** over those patches, and outputs patch-level features. The script:
-   - Registers **forward hooks** on the encoder’s attention layers so that when we run a forward pass, we get the **attention weight matrices** (which patch “looks at” which).
-   - To get weights, the code forces **eager** attention (instead of SDPA/Flash), which returns the full weight tensor.
-   - For each frame, it builds a batch (with the right image keys and dtype), runs the vision encoder (or full policy), and the hooks record the attention.
-   - **Aggregation methods** (`--method`): `last-layer` uses only the final encoder layer, `rollout` multiplies attention across all layers (accounting for residual connections) for a more complete picture of information flow, `all-layers` keeps each layer’s attention separately.
+6. **Visualize** -- the output grid has up to 5 rows per frame:
 
-4. **Optionally capture cross-attention** (`--cross-attention`)
-   SmolVLA doesn’t have explicit cross-attention layers. Instead, the VLM builds a **KV cache** from the prefix sequence (vision tokens + language tokens + state tokens), and the **action expert** queries that cache. The script monkey-patches `eager_attention_forward()` on the expert to intercept the softmax attention probabilities when expert queries attend to prefix keys (detected by Q seq-len != K seq-len). Only the columns corresponding to vision tokens are kept, giving a heatmap of which image regions the action decoder actually reads.
+| 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) |
 
-5. **Turn attention into a spatial heatmap**
-   Attention is in **patch space** (e.g. 32×32 grid for SigLIP with 512px images and 16px patches). For each patch we get an “importance” score (mean attention received). Those scores are reshaped into a 2D grid, **upsampled** with bilinear interpolation to the original image size, and normalized to [0, 1]. That gives a single **heatmap** per frame (bright = high attention).
+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.
 
-6. **Visualize**
-   Default grid (3 rows per frame): **Row 1** original, **Row 2** self-attention heatmap, **Row 3** overlay. With `--cross-attention`, two extra rows are added: **Row 4** cross-attention heatmap (hot colormap), **Row 5** dual-color overlay (self-attention in blue, cross-attention in red). With `--show-heads`, a separate per-head grid is saved for the first frame.
+### Per-head grid
 
-**Fallback:** If the vision encoder can’t be hooked or attention isn’t captured (e.g. wrong architecture), the script can fall back to **input-gradient saliency**: backprop from the action output to the image pixels and use gradient magnitude as a proxy for “what the model uses.” That’s no longer true attention but still highlights influential pixels.
+With `--show-heads`, a separate grid shows each of the 12 SigLIP attention heads individually for the first frame:
+
+![Per-head attention grid](assets/example_per_head.png)
+*Each subplot is one attention head. Look for specialization -- e.g. one head tracking the gripper, another tracking the object.*
 
 ---
 
 ## Setup
 
-**Requirements:** Python 3.10+, and FFmpeg 4–7 for video decoding (LeRobot uses TorchCodec). On macOS, Homebrew’s default `ffmpeg` is often v8; install FFmpeg 6 so the loader can find it:
+**Requirements:** Python 3.10+, and FFmpeg 4-7 for video decoding (LeRobot uses TorchCodec). On macOS, Homebrew's default `ffmpeg` is often v8; install FFmpeg 6:
 
 ```bash
 brew install ffmpeg@6
@@ -60,7 +74,7 @@ Then:
 
 ```bash
 cd smolvla-inspect
-python3.10 -m venv .venv && source .venv/bin/activate   # or python3.11
+python3.11 -m venv .venv && source .venv/bin/activate
 pip install -r requirements.txt
 ```
 
@@ -68,105 +82,120 @@ pip install -r requirements.txt
 
 ## Run
 
-Use `run.sh` so TorchCodec finds FFmpeg 6’s libs (needed for dataset video decoding):
+Use `run.sh` so TorchCodec finds FFmpeg 6's libs:
 
 ```bash
 source .venv/bin/activate
 ./run.sh
 ```
 
-Or set the library path yourself and run Python:
+Or set the library path yourself:
 
 ```bash
 export DYLD_LIBRARY_PATH="/opt/homebrew/opt/ffmpeg@6/lib:$DYLD_LIBRARY_PATH"
 python inspect_attention.py
 ```
 
-Examples:
+### Examples
 
 ```bash
-# Default: pretrained SmolVLA base + real-world SO101 pick-place dataset
+# Default config: rollout method, cross-attention enabled, per-head grid enabled
 ./run.sh
 
 # Your fine-tuned model
 ./run.sh --model path/to/finetuned_checkpoint --dataset path/to/dataset
 
-# More frames, specific episode, explicit GPU
-./run.sh --episode 3 --num-frames 12 --device cuda
-
-# Device is auto-detected by default (mps → cuda → cpu)
-# Override with --device cpu/cuda/mps if needed
+# More frames, specific episode
+./run.sh --episode 3 --num-frames 12
 
-# Save each frame as a separate PNG
-./run.sh --save-individual
+# Last-layer only (faster, no rollout)
+./run.sh --method last-layer
 
-# Attention rollout across all SigLIP layers (instead of last layer only)
-./run.sh --method rollout
+# Skip cross-attention (faster, 3-row grid only) — edit configs/defaults.yaml:
+#   cross_attention: false
 
-# Capture action-expert → vision cross-attention (slower, runs full policy forward)
-./run.sh --cross-attention
+# Raw attention without positional baseline subtraction
+./run.sh --raw-attention
 
-# Show per-head attention patterns for the first frame
-./run.sh --show-heads
-
-# Combine flags
-./run.sh --method rollout --cross-attention --show-heads
+# Explicit device override (auto-detected by default: mps > cuda > cpu)
+./run.sh --device cuda
 ```
 
 Results land in `outputs/`.
 
----
-
-## Project layout
-
-```
-smolvla-inspect/
-├── inspect_attention.py   # Main script: load model/dataset, extract attention, save heatmaps
-├── assets/
-│   └── how_it_works_architecture.png   # Diagram: architecture + hooks + attention→heatmap
-├── configs/
-│   └── defaults.yaml      # Default model/dataset/output paths (optional)
-├── docs/
-│   ├── ELI5.md             # Plain-language explanation of how the interpretability works
-│   └── TESTING.md          # CLI test commands and expected output
-├── outputs/                # Generated grid and per-frame images
-├── run.sh                  # Wrapper that sets FFmpeg lib path
-├── requirements.txt
-└── README.md
-```
+### CLI flags
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--model` | `lerobot/smolvla_base` | HuggingFace model ID or local path |
+| `--dataset` | `lerobot/svla_so101_pickplace` | LeRobot dataset ID or local path |
+| `--episode` | `0` | Episode index to visualize |
+| `--num-frames` | `8` | Number of frames to sample |
+| `--image-key` | auto-detected | Dataset image key override |
+| `--output-dir` | `./outputs` | Output directory |
+| `--device` | `auto` | `auto`, `cpu`, `cuda`, or `mps` |
+| `--save-individual` | `true` | Save each frame as a separate PNG |
+| `--method` | `rollout` | `last-layer`, `rollout`, or `all-layers` |
+| `--cross-attention` | `true` | Capture action-expert cross-attention |
+| `--show-heads` | `true` | Save per-head attention grid for first frame |
+| `--raw-attention` | `false` | Skip positional baseline subtraction |
+
+Defaults can be changed in `configs/defaults.yaml`.
 
 ---
 
 ## What to look for
 
-### Self-attention (vision encoder — rows 2-3)
+### Self-attention (SigLIP vision encoder -- rows 2-3)
 
 | Attention pattern | Interpretation |
 |-------------------|----------------|
-| Bright on gripper + object + goal | **Healthy** — model attends to task-relevant regions |
-| Bright on shelves, cables, table grain | **Background overfitting** — model may be using scene cues |
+| Bright on gripper + object + goal | **Healthy** -- model attends to task-relevant regions |
+| Bright on shelves, cables, table grain | **Background overfitting** -- model may be using scene cues |
 | Uniform / diffuse everywhere | Model may not have learned focused visual features yet |
-| Shifts from background → object across frames | Model is tracking the task over time (good sign) |
+| Shifts from background to object across frames | Model is tracking the task over time (good sign) |
 
-### Cross-attention (action expert → vision — rows 4-5, with `--cross-attention`)
+### Cross-attention (action expert -- rows 4-5)
 
 | Attention pattern | Interpretation |
 |-------------------|----------------|
-| Tight focus on gripper tip + target object | **Healthy** — the action decoder reads exactly the tokens it needs |
-| Diffuse across all vision tokens | Decoder hasn't specialised; may predict generic/averaged actions |
-| Self-attn diffuse but cross-attn focused | Good sign — the decoder learned to select useful tokens despite a noisy encoder |
-| Self-attn focused but cross-attn diffuse | Encoder features are good but the decoder doesn't exploit them well |
+| Tight focus on gripper tip + target object | **Healthy** -- action decoder reads exactly what it needs |
+| Diffuse across all vision tokens | Decoder hasn't specialized; may predict generic actions |
+| Self-attn diffuse but cross-attn focused | Decoder learned to select useful tokens despite a noisy encoder |
+| Self-attn focused but cross-attn diffuse | Encoder features are good but the decoder doesn't exploit them |
+
+### Co-attention (row 5)
 
-### Per-head patterns (with `--show-heads`)
+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.
 
-Look for heads that specialise: one head tracking the gripper, another tracking the object, another attending to the background. Specialisation is a sign of a well-trained encoder.
+### Per-head patterns
 
-### Grid layout
+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.
+
+---
+
+## Project layout
 
-The default grid has three rows per frame: **Row 1 = original**, **Row 2 = self-attention heatmap**, **Row 3 = overlay**. With `--cross-attention`, two extra rows appear: **Row 4 = cross-attention heatmap**, **Row 5 = dual-color overlay** (blue = self-attention, red = cross-attention).
+```
+smolvla-inspect/
+├── inspect_attention.py      # All logic: model loading, hooks, heatmaps, visualization
+├── assets/
+│   ├── how_it_works_architecture.png
+│   ├── example_grid.png
+│   └── example_per_head.png
+├── configs/
+│   └── 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 (gitignored)
+├── run.sh                    # Wrapper that sets FFmpeg lib path
+├── requirements.txt
+└── README.md
+```
 
 ---
 
 ## Note on FFmpeg
 
-If you installed `ffmpeg@6` and linked it (e.g. `brew link --overwrite ffmpeg@6`), your default `ffmpeg` is now 6.x. To switch back to FFmpeg 8 later: `brew unlink ffmpeg@6 && brew link ffmpeg`.
+If you installed `ffmpeg@6` and linked it (`brew link --overwrite ffmpeg@6`), your default `ffmpeg` is now 6.x. To switch back later: `brew unlink ffmpeg@6 && brew link ffmpeg`.