Add model health diagnostics and architecture documentation

- Add --model-health mode with weight spectral analysis (WeightWatcher),
attention entropy, and head redundancy checks across all components
- Capture attention from SigLIP vision encoder, VLM+Expert joint
self-attention, and Expert-to-VLM cross-attention via monkey-patched
eager_attention_forward()
- Generate terminal report, markdown report, and 3-panel PNG plot
- Add per-component grouping with vertical separators in plots
- Add descriptive text to each report section explaining the metrics
- Add architecture diagrams (assets/architecture.md) covering data flow,
attention components, execution timeline, and trainable vs frozen parts
- Update README with health diagnostics docs, new CLI flags, and
architecture diagram link

README.md

@@ -9,7 +9,14 @@ 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 **visualizes where the model looks** by extracting attention maps from two places:
+SmolVLA is a **vision-language-action** policy: it takes camera images and a language instruction, then outputs robot actions. This tool has two 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
+
+### Attention visualization
+
+Extracts 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
@@ -19,6 +26,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, and an optional per-head attention grid.
 
+### Model health diagnostics
+
+Runs three diagnostic checks across all model components (SigLIP vision encoder, VLM text model, action expert, connector, and projection heads):
+
+1. **Weight spectral analysis** -- fits a power-law to singular values of each weight matrix using WeightWatcher. The alpha exponent indicates training quality (2-4 = healthy, >6 = severely undertrained).
+2. **Attention entropy** -- measures how focused or diffuse each attention head is across three attention operations: SigLIP self-attention, VLM+Expert joint self-attention, and Expert-to-VLM cross-attention.
+3. **Head redundancy** -- measures pairwise cosine similarity between attention heads within each layer. High similarity means wasted capacity.
+
+**Output:** Terminal report, markdown report (`model_health_report.md`), and a 3-panel plot (`model_health_report.png`).
+
+For a detailed visual walkthrough of the architecture and how it maps to the report, see **[Architecture Diagrams](assets/architecture.md)**.
+
 ---
 
 ## How it works
@@ -34,10 +53,10 @@ That lets you check whether the model attends to task-relevant regions (gripper,
 
 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
+   - `rollout` -- multiplies attention across all layers with residual connections, giving a more complete picture of information flow
    - `all-layers` -- keeps each layer separately
 
-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.
+4. **Capture cross-attention** (`--cross-attention`) -- 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.
 
 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].
 
@@ -99,7 +118,7 @@ python inspect_attention.py
 ### Examples
 
 ```bash
-# Default config: rollout method, cross-attention enabled, per-head grid enabled
+# Default: attention heatmaps with last-layer method
 ./run.sh
 
 # Your fine-tuned model
@@ -108,15 +127,27 @@ python inspect_attention.py
 # More frames, specific episode
 ./run.sh --episode 3 --num-frames 12
 
-# Last-layer only (faster, no rollout)
-./run.sh --method last-layer
+# Rollout aggregation (multiply attention across all layers)
+./run.sh --method rollout
+
+# Enable cross-attention capture (slower, adds rows 4-5)
+./run.sh --cross-attention
 
-# Skip cross-attention (faster, 3-row grid only) — edit configs/defaults.yaml:
-#   cross_attention: false
+# Per-head attention grid for the first frame
+./run.sh --show-heads
 
 # Raw attention without positional baseline subtraction
 ./run.sh --raw-attention
 
+# Model health diagnostics (spectral analysis + entropy + redundancy)
+./run.sh --model-health
+
+# Health check with more sample frames for stable entropy estimates
+./run.sh --model-health --health-frames 10
+
+# Custom thresholds for health warnings
+./run.sh --model-health --entropy-warn 0.85 --redundancy-warn 0.75
+
 # Explicit device override (auto-detected by default: mps > cuda > cpu)
 ./run.sh --device cuda
 ```
@@ -125,6 +156,8 @@ Results land in `outputs/`.
 
 ### CLI flags
 
+**Attention visualization:**
+
 | Flag | Default | Description |
 |------|---------|-------------|
 | `--model` | `lerobot/smolvla_base` | HuggingFace model ID or local path |
@@ -134,12 +167,24 @@ Results land in `outputs/`.
 | `--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 |
+| `--save-individual` | `false` | Save each frame as a separate PNG |
+| `--method` | `last-layer` | `last-layer`, `rollout`, or `all-layers` |
+| `--cross-attention` | `false` | Capture action-expert cross-attention |
+| `--show-heads` | `false` | Save per-head attention grid for first frame |
 | `--raw-attention` | `false` | Skip positional baseline subtraction |
 
+**Model health diagnostics:**
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--model-health` | `false` | Run health diagnostics instead of attention heatmaps |
+| `--health-frames` | `5` | Number of sample frames for entropy/redundancy |
+| `--entropy-warn` | `0.8` | Entropy ratio threshold for "unfocused" warning |
+| `--entropy-critical` | `0.95` | Entropy ratio threshold for "dead" heads |
+| `--entropy-low` | `0.1` | Entropy ratio threshold for "collapsed" heads |
+| `--redundancy-warn` | `0.7` | Cosine similarity threshold for "high redundancy" |
+| `--redundancy-critical` | `0.9` | Cosine similarity threshold for "collapsed" heads |
+
 Defaults can be changed in `configs/defaults.yaml`.
 
 ---
@@ -172,14 +217,33 @@ The cyan overlay highlights regions where **both** the vision encoder and the ac
 
 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.
 
+### Model health report
+
+| Metric | Healthy | Warning | Critical |
+|--------|---------|---------|----------|
+| Spectral alpha | 2-4 | 4-6 (undertrained) | >6 (severely undertrained) or <2 (overcorrelated) |
+| Attention entropy | 0.10-0.80 | >0.80 (unfocused) | >0.95 (dead) or <0.10 (collapsed) |
+| Head redundancy | <0.70 (diverse) | >0.70 (redundant) | >0.90 (collapsed) |
+
+The report covers three attention components mapped to distinct operations in the architecture:
+
+| Report component | Architecture operation | When it runs |
+|-----------------|----------------------|-------------|
+| SigLIP Vision (12L, 12H) | Self-attention inside the vision encoder | Image encoding |
+| VLM+Expert Joint Self-Attn (16L, 15H) | VLM and Expert tokens concatenated, attend to each other | Prefill (initial encoding) |
+| Expert-to-VLM Cross-Attn (16L, 8H) | Expert queries VLM's cached keys/values | Generation (action decoding, 10 steps) |
+
+See **[Architecture Diagrams](assets/architecture.md)** for visual explanations of each component.
+
 ---
 
 ## Project layout
 
 ```
 smolvla-inspect/
-├── inspect_attention.py      # All logic: model loading, hooks, heatmaps, visualization
+├── inspect_attention.py      # All logic: model loading, hooks, heatmaps, health diagnostics
 ├── assets/
+│   ├── architecture.md       # Architecture diagrams and report reference
 │   ├── how_it_works_architecture.png
 │   ├── example_grid.png
 │   └── example_per_head.png
@@ -188,7 +252,9 @@ smolvla-inspect/
 ├── docs/
 │   ├── ELI5.md               # Plain-language explanation of the interpretability approach
 │   └── TESTING.md            # CLI test commands and expected output
-├── outputs/                  # Generated images (gitignored)
+├── outputs/                  # Generated images and reports (gitignored)
+│   ├── model_health_report.md
+│   └── model_health_report.png
 ├── run.sh                    # Wrapper that sets FFmpeg lib path
 ├── requirements.txt
 └── README.md

assets/architecture.md

@@ -0,0 +1,276 @@
+# SmolVLA Architecture Diagrams
+
+Reference diagrams for understanding the SmolVLA model structure and what
+the `--model-health` report measures.
+
+---
+
+## 1. High-level data flow
+
+```
+  Camera Images (512x512)       Task String            Robot State     Noisy Actions
+  +--------+ +--------+         "pick up the           (joint          (from flow
+  |  side  | |   up   |          red block"             positions)       matching)
+  +---+----+ +---+----+              |                     |               |
+      |          |                   |                     |               |
+      v          v                   v                     v               v
+  +------------------+     +---------------+     +---------------+  +---------------+
+  | SigLIP Vision    |     | Tokenizer     |     | state_proj    |  | action_in_proj|
+  | Encoder          |     |               |     | (Linear)      |  | (Linear)      |
+  | (frozen)         |     |               |     |               |  | + time MLP    |
+  | 12L, 12H         |     |               |     | trainable     |  | trainable     |
+  +--------+---------+     +-------+-------+     +-------+-------+  +-------+-------+
+           |                       |                     |                  |
+     1024 patches            token embeds            1 token          10 tokens
+           |                       |                     |                  |
+           v                       |                     |                  |
+  +-------------------+            |                     |                  |
+  | Connector         |            |                     |                  |
+  | (pixel shuffle)   |            |                     |                  |
+  | trainable         |            |                     |                  |
+  +--------+----------+            |                     |                  |
+           |                       |                     |                  |
+     64 tokens                     |                     |                  |
+           |                       |                     |                  |
+           +-----------+-----------+                     |                  |
+                       |                                 |                  |
+                       v                                 |                  |
+           +-----------------------+                     |                  |
+           | PREFIX SEQUENCE       |                     |                  |
+           | [vision] [language]   +<--------------------+                  |
+           | [state]               |     VLM dim (960)                     |
+           +-----------+-----------+                                       |
+                       |                                                   |
+                       |                                 +-----------------+
+                       |                                 |
+                       v                                 v
+           +---------------------------------------------------+
+           |         SmolVLMWithExpert (16 layers)               |
+           |                                                     |
+           |   PREFIX  ------------------->  VLM Text Model      |
+           |   (vision+lang+state)            (frozen)           |
+           |                                  960-dim, 15 heads  |
+           |                                                     |
+           |   SUFFIX  ------------------->  Action Expert       |
+           |   (noisy actions+time)           (trainable)        |
+           |                                  480-dim, 8 heads   |
+           |                                                     |
+           |         (see attention diagrams below)              |
+           +--------------------------+--------------------------|
+                                      |
+                                      |  Expert output (last 10 tokens)
+                                      v
+                          +-------------------+
+                          | action_out_proj   |
+                          | (Linear)          |
+                          | trainable         |
+                          +---------+---------+
+                                    |
+                                    v
+                          Predicted Actions
+                          (10 steps x 6 DOF)
+```
+
+---
+
+## 2. The three attention components
+
+These are the three distinct attention operations that the health report
+measures for entropy and head redundancy.
+
+### A. SigLIP Vision (12 layers, 12 heads)
+
+Standard self-attention inside the vision encoder. Happens before anything
+enters SmolVLMWithExpert. Image patches attend to other image patches.
+
+```
+  patch_1  patch_2  patch_3  ...  patch_1024
+     |        |        |              |
+     v        v        v              v
+  +----------------------------------------------+
+  | SigLIP Self-Attention                         |
+  |                                               |
+  | Each patch can see all patches.               |
+  |                                               |
+  | Q <-- patches        K <-- patches            |
+  | (1024 x 1024 attention matrix)                |
+  +----------------------------------------------+
+
+  Report label: "SigLIP Vision (12L, 12H)"
+```
+
+### B. VLM+Expert Joint Self-Attention (16 layers, 15 heads)
+
+Runs during prefill (initial encoding). VLM and Expert tokens are
+concatenated into one sequence and attend to each other in a single
+attention call.
+
+```
+  VLM tokens (prefix)           Expert tokens (suffix)
+  [vision][lang][state]          [action+time]
+        |                              |
+        v                              v
+  +-------------+               +-------------+
+  | VLM layer   |               | Expert layer|
+  | q/k/v proj  |               | q/k/v proj  |
+  +------+------+               +------+------+
+         |                             |
+         +-------------+---------------+
+                       |
+                       v
+                torch.cat(dim=seq)
+                       |
+                       v
+  +----------------------------------------------+
+  | JOINT Self-Attention                          |
+  |                                               |
+  | All tokens see all tokens:                    |
+  |   vision  <-->  vision                        |
+  |   vision  <-->  language                      |
+  |   vision  <-->  action                        |
+  |   action  <-->  language                      |
+  |   action  <-->  action                        |
+  |   ... etc                                     |
+  |                                               |
+  | Q <-- [VLM+Expert]     K <-- [VLM+Expert]     |
+  | Q_len == K_len  (self-attention)              |
+  +----------------------------------------------+
+
+  Report label: "VLM+Expert Joint Self-Attn (16L, 15H)"
+```
+
+### C. Expert-to-VLM Cross-Attention (16 layers, 8 heads)
+
+Runs during generation (autoregressive action decoding). The Expert
+generates queries from its own hidden states, but keys and values come
+from the VLM's cached representations, re-projected through the Expert's
+k_proj and v_proj. Repeats 10 times (one per action token).
+
+```
+  VLM KV Cache                       Expert hidden states
+  (built during prefill)             (current action token)
+        |                                    |
+        |                                    v
+        |                            +---------------+
+        |                            | Expert q_proj |
+        |                            +-------+-------+
+        |                                    |
+        |     re-projected through           |
+        |     Expert k_proj / v_proj         |
+        v                                    |
+  +-------------+                            |
+  | Expert K, V |<-- VLM cache               |
+  | (from VLM)  |   re-projected             |
+  +------+------+                            |
+         |                                   |
+         +---------------+-------------------+
+                         |
+                         v
+  +----------------------------------------------+
+  | CROSS-Attention                               |
+  |                                               |
+  | Expert reads VLM:                             |
+  |   action --> vision   (what to grab?)         |
+  |   action --> language  (what was asked?)      |
+  |   action --> state     (current pose?)        |
+  |                                               |
+  | VLM does NOT read Expert here.                |
+  |                                               |
+  | Q <-- Expert tokens (few)                     |
+  | K <-- VLM prefix (many)                       |
+  | Q_len != K_len  (cross-attention)             |
+  +----------------------------------------------+
+
+  Report label: "Expert-to-VLM Cross-Attn (16L, 8H)"
+```
+
+---
+
+## 3. Execution timeline
+
+```
+  IMAGE ENCODING            PREFILL                   GENERATION
+  +--------------+    +------------------+    +------+------+     +------+
+  | SigLIP       |    | Joint Self-Attn  |    | XA   | XA   |    | XA   |
+  | 12 layers    +--->| 16 layers        +--->| #1   | #2   |...>| #10  |
+  | (A)          |    | (B)              |    |      |      |    |      |
+  +--------------+    +------------------+    +------+------+    +------+
+                                              <--- 10 action steps --->
+
+  Report:              Report:                 Report:
+  "SigLIP Vision"      "VLM+Expert             "Expert-to-VLM
+                        Joint Self-Attn"        Cross-Attn"
+  12L x 12H            16L x 15H               16L x 8H
+                                                (x10 steps, averaged)
+```
+
+---
+
+## 4. What the health report measures
+
+### Section 1: Weight Spectral Analysis (alpha)
+
+Analyzes weight matrices directly (no data needed). Each transformer layer
+contains multiple weight matrices (q_proj, k_proj, v_proj, o_proj,
+gate_proj, up_proj, down_proj). A power-law is fit to each matrix's
+singular values; the exponent (alpha) indicates training quality.
+
+```
+  +----------------------------------+-------------+---------------------+
+  | Component                        | Wt Matrices | What it is          |
+  +----------------------------------+-------------+---------------------+
+  | Expert (trainable)               | 112 = 16x7  | Action decoder      |
+  | VLM Text Model (frozen)          | 113 = 16x7+ | Language model      |
+  | Vision Encoder (frozen)          |  74 = 12x~6 | SigLIP              |
+  | Connector (trainable)            |   1          | Pixel shuffle       |
+  | Projections (trainable)          | too small    | state/action linear |
+  +----------------------------------+-------------+---------------------+
+
+  Healthy: alpha 2-4     Undertrained: 4-6     Severe: >6
+```
+
+### Section 2: Attention Entropy
+
+Requires a forward pass with real data. Measures how spread out each
+attention head's focus is (normalized by log of sequence length).
+
+```
+  Collapsed: <0.10     Healthy: 0.10-0.80     Unfocused: >0.80     Dead: >0.95
+```
+
+### Section 3: Head Redundancy
+
+Same forward pass data. Measures pairwise cosine similarity between
+flattened head attention patterns within each layer. Heads should learn
+different patterns.
+
+```
+  Diverse: <0.70       High redundancy: >0.70       Collapsed: >0.90
+```
+
+---
+
+## 5. Trainable vs frozen components
+
+```
+  FROZEN (pretrained, not updated during fine-tuning)
+  +----------------------------------------------------+
+  | SigLIP Vision Encoder                               |
+  | 12 layers, 12 heads, 74 weight matrices             |
+  +----------------------------------------------------+
+  | VLM Text Model (SmolLM2)                            |
+  | 16 layers, 15 heads, 113 weight matrices            |
+  +----------------------------------------------------+
+
+  TRAINABLE (learned during fine-tuning)
+  +----------------------------------------------------+
+  | Action Expert                                       |
+  | 16 layers, 8 heads, 112 weight matrices             |
+  +----------------------------------------------------+
+  | Connector (pixel shuffle)                           |
+  | 1 weight matrix                                     |
+  +----------------------------------------------------+
+  | state_proj, action_in_proj, action_out_proj         |
+  | (too small for spectral analysis)                   |
+  +----------------------------------------------------+
+```

configs/defaults.yaml

@@ -13,3 +13,12 @@ method: rollout                         # last-layer | rollout | all-layers
 cross_attention: true                   # true to capture action-expert → vision cross-attention (slower)
 show_heads: true                        # true to save a per-head attention grid for the first frame
 raw_attention: false                    # true to skip positional baseline subtraction
+
+# Model health analysis
+model_health: false
+health_frames: 5
+entropy_warn: 0.8
+entropy_critical: 0.95
+entropy_low: 0.1
+redundancy_warn: 0.7
+redundancy_critical: 0.9

inspect_attention.py

@@ -345,6 +345,129 @@ class ActionVisionAttentionCapture:
         return accum / count
 
 
+# ---------------------------------------------------------------------------
+# 1c. Decoder attention capture — VLM+Expert self-attn & Expert cross-attn
+# ---------------------------------------------------------------------------
+
+class DecoderAttentionCapture:
+    """
+    Captures attention weights from ``SmolVLMWithExpert.eager_attention_forward``.
+
+    SmolVLA bypasses the standard ``LlamaAttention.forward()`` path — instead,
+    ``forward_attn_layer`` / ``forward_cross_attn_layer`` manually project
+    Q/K/V and call ``eager_attention_forward`` directly. So we monkey-patch
+    that method to intercept the attention probabilities.
+
+    Captured maps are categorised by the Q/K sequence lengths:
+      - **Self-attention** (``q_len == k_len``, ``q_len > 1``): prefill phase
+        where VLM + Expert tokens are concatenated.
+      - **Cross-attention** (``q_len != k_len``): Expert queries attending to
+        the VLM prefix KV cache.
+    """
+
+    def __init__(self):
+        self.self_attn_maps = []    # (seq_idx, probs) for prefill self-attn
+        self.cross_attn_maps = []   # (seq_idx, probs) for expert cross-attn
+        self._original_fn = None
+        self._patched_obj = None
+        self._self_idx = 0
+        self._cross_idx = 0
+
+    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
+            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,
+            )
+
+            # ---------- capture attention probs ----------
+            q_len = query_states.shape[1]
+            k_len = key_states.shape[1]
+
+            # Only capture meaningful attention (skip single-token
+            # autoregressive steps which have q_len == 1)
+            if q_len > 1:
+                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)
+
+                if q_len == k_len:
+                    capture.self_attn_maps.append(
+                        (capture._self_idx, probs.detach().cpu())
+                    )
+                    capture._self_idx += 1
+                else:
+                    capture.cross_attn_maps.append(
+                        (capture._cross_idx, probs.detach().cpu())
+                    )
+                    capture._cross_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.self_attn_maps = []
+        self.cross_attn_maps = []
+        self._self_idx = 0
+        self._cross_idx = 0
+
+    def reset_maps(self):
+        """Clear captured maps but keep the monkey-patch active."""
+        self.self_attn_maps = []
+        self.cross_attn_maps = []
+        self._self_idx = 0
+        self._cross_idx = 0
+
+    def get_self_attn_layers(self):
+        """Return prefill self-attention maps sorted by sequential index."""
+        if not self.self_attn_maps:
+            return []
+        return sorted(self.self_attn_maps, key=lambda x: x[0])
+
+    def get_cross_attn_layers(self):
+        """Return expert cross-attention maps sorted by sequential index."""
+        if not self.cross_attn_maps:
+            return []
+        return sorted(self.cross_attn_maps, key=lambda x: x[0])
+
+
 # ---------------------------------------------------------------------------
 # 2. Gradient-based attention fallback (GradCAM-style)
 # ---------------------------------------------------------------------------
@@ -1502,7 +1625,859 @@ def gradient_attention_map(policy, dataset, frame_idx, image_key, device="cpu"):
 
 
 # ---------------------------------------------------------------------------
-# 8. Entry point
+# 8. Model health diagnostics
+# ---------------------------------------------------------------------------
+
+def compute_weightwatcher_alpha(policy):
+    """
+    Run WeightWatcher spectral analysis on each trainable (and frozen
+    reference) component of SmolVLA.
+
+    Returns dict keyed by component name, values are lists of
+    ``{layer, alpha}`` dicts.  Components that are too small for
+    reliable SVD are reported with ``alpha=None``.
+    """
+    try:
+        import weightwatcher as ww
+    except ImportError:
+        print("  ERROR: weightwatcher not installed. Run:  pip install weightwatcher")
+        print("  Skipping spectral alpha analysis.")
+        return None
+
+    components = {}
+
+    # --- helper: run ww on a submodel, return list of {layer, alpha} ---
+    def _analyze(name, submodel):
+        try:
+            watcher = ww.WeightWatcher(model=submodel)
+            details = watcher.analyze(min_evals=50)
+            results = []
+            for idx, row in details.iterrows():
+                results.append({"layer": idx, "alpha": row.get("alpha", None)})
+            components[name] = results
+        except Exception as e:
+            print(f"    WARNING: WeightWatcher failed on {name}: {e}")
+            components[name] = []
+
+    vlm_with_expert = policy.model.vlm_with_expert
+
+    # Trainable components
+    print("  Analyzing expert layers...")
+    _analyze("Expert (trainable)", vlm_with_expert.lm_expert)
+
+    print("  Analyzing connector...")
+    _analyze("Connector (trainable)", vlm_with_expert.get_vlm_model().connector)
+
+    # Projection heads
+    proj_names = [n for n, _ in policy.model.named_children()
+                  if "proj" in n.lower()]
+    if proj_names:
+        print(f"  Analyzing projection heads ({', '.join(proj_names)})...")
+        for pn in proj_names:
+            _analyze(f"Projection/{pn} (trainable)", getattr(policy.model, pn))
+
+    # Frozen reference components
+    print("  Analyzing vision encoder (frozen reference)...")
+    _analyze("Vision Encoder (frozen)", vlm_with_expert.get_vlm_model().vision_model)
+
+    print("  Analyzing VLM text model (frozen reference)...")
+    _analyze("VLM Text Model (frozen)", vlm_with_expert.get_vlm_model().text_model)
+
+    return components
+
+
+def compute_attention_entropy(attn_maps, num_patches):
+    """
+    Per-head, per-layer attention entropy as a fraction of maximum
+    entropy (``log(num_patches)``).
+
+    Args:
+        attn_maps: list of ``(layer_idx, attn_weights)`` tuples.
+            Each ``attn_weights`` has shape ``(batch, heads, patches, patches)``
+            or ``(heads, patches, patches)``.
+        num_patches: total number of patches (for max-entropy normalisation).
+
+    Returns:
+        list of ``{layer, head, entropy, entropy_ratio}`` dicts.
+    """
+    max_entropy = math.log(num_patches) if num_patches > 1 else 1.0
+    eps = 1e-8
+    results = []
+
+    for layer_idx, attn in sorted(attn_maps, key=lambda x: x[0]):
+        # Reduce to (heads, patches, patches)
+        while attn.dim() > 3:
+            attn = attn[0]
+        if attn.dim() == 2:
+            attn = attn.unsqueeze(0)
+
+        n_heads = attn.shape[0]
+        for h in range(n_heads):
+            head_attn = attn[h].float()  # (patches, patches)
+            # Entropy per row, then average across rows
+            ent = -(head_attn * torch.log(head_attn + eps)).sum(dim=-1).mean().item()
+            results.append({
+                "layer": layer_idx,
+                "head": h,
+                "entropy": ent,
+                "entropy_ratio": ent / max_entropy,
+            })
+
+    return results
+
+
+def compute_head_redundancy(attn_maps):
+    """
+    Pairwise cosine similarity between flattened head attention patterns
+    within each layer.
+
+    Args:
+        attn_maps: list of ``(layer_idx, attn_weights)`` tuples.
+
+    Returns:
+        list of ``{layer, mean_redundancy, max_redundancy}`` dicts.
+    """
+    results = []
+
+    for layer_idx, attn in sorted(attn_maps, key=lambda x: x[0]):
+        while attn.dim() > 3:
+            attn = attn[0]
+        if attn.dim() == 2:
+            attn = attn.unsqueeze(0)
+
+        n_heads = attn.shape[0]
+        if n_heads < 2:
+            results.append({
+                "layer": layer_idx,
+                "mean_redundancy": 0.0,
+                "max_redundancy": 0.0,
+            })
+            continue
+
+        # Flatten each head's attention to a vector
+        flat = attn.float().reshape(n_heads, -1)  # (heads, patches*patches)
+        # Normalise
+        flat_norm = F.normalize(flat, dim=1)
+        # Pairwise cosine similarity
+        sim = torch.mm(flat_norm, flat_norm.t())  # (heads, heads)
+
+        # Extract upper triangle (exclude diagonal)
+        mask = torch.triu(torch.ones(n_heads, n_heads, dtype=torch.bool), diagonal=1)
+        pairwise = sim[mask]
+        results.append({
+            "layer": layer_idx,
+            "mean_redundancy": pairwise.mean().item(),
+            "max_redundancy": pairwise.max().item(),
+        })
+
+    return results
+
+
+def classify_health(alpha=None, entropy=None, redundancy=None, thresholds=None):
+    """
+    Classify a single metric value into a status string and ANSI color.
+
+    Alpha thresholds (hardcoded, RMT-derived):
+        <2 overcorrelated, 2-4 healthy, 4-6 undertrained, >6 severely undertrained
+
+    Entropy/redundancy thresholds come from *thresholds* dict.
+
+    Returns:
+        ``(ansi_status_str, ansi_color_code, plain_label)``
+    """
+    GREEN = "\033[92m"
+    YELLOW = "\033[93m"
+    RED = "\033[91m"
+    RESET = "\033[0m"
+
+    if thresholds is None:
+        thresholds = {}
+
+    if alpha is not None:
+        if alpha < 2:
+            return f"{RED}overcorrelated{RESET}", RED, "overcorrelated"
+        elif alpha <= 4:
+            return f"{GREEN}healthy{RESET}", GREEN, "healthy"
+        elif alpha <= 6:
+            return f"{YELLOW}undertrained{RESET}", YELLOW, "undertrained"
+        else:
+            return f"{RED}severely undertrained{RESET}", RED, "severely undertrained"
+
+    if entropy is not None:
+        crit = thresholds.get("entropy_critical", 0.95)
+        warn = thresholds.get("entropy_warn", 0.8)
+        low = thresholds.get("entropy_low", 0.1)
+        if entropy >= crit:
+            return f"{RED}uniform/dead{RESET}", RED, "uniform/dead"
+        elif entropy >= warn:
+            return f"{YELLOW}unfocused{RESET}", YELLOW, "unfocused"
+        elif entropy <= low:
+            return f"{RED}collapsed{RESET}", RED, "collapsed"
+        else:
+            return f"{GREEN}healthy{RESET}", GREEN, "healthy"
+
+    if redundancy is not None:
+        r_crit = thresholds.get("redundancy_critical", 0.9)
+        r_warn = thresholds.get("redundancy_warn", 0.7)
+        if redundancy >= r_crit:
+            return f"{RED}collapsed{RESET}", RED, "collapsed"
+        elif redundancy >= r_warn:
+            return f"{YELLOW}high redundancy{RESET}", YELLOW, "high redundancy"
+        else:
+            return f"{GREEN}diverse{RESET}", GREEN, "diverse"
+
+    return "unknown", "", "unknown"
+
+
+def print_health_report(ww_results, entropy_results, redundancy_results, thresholds):
+    """
+    Print a formatted terminal table with ANSI color-coded status per
+    layer/component.
+    """
+    BOLD = "\033[1m"
+    RESET = "\033[0m"
+    DIM = "\033[2m"
+
+    print(f"\n{'=' * 78}")
+    print(f"{BOLD}MODEL HEALTH REPORT{RESET}")
+    print(f"{'=' * 78}")
+
+    # --- Section 1: Weight Spectral Analysis ---
+    if ww_results is not None:
+        print(f"\n{BOLD}1. Weight Spectral Analysis (alpha){RESET}")
+        print(f"   Fits a power-law to each weight matrix's singular values.")
+        print(f"   Alpha (α) measures how well-trained a layer is — values of 2-4 indicate")
+        print(f"   strong correlation structure learned during training. High alpha means the")
+        print(f"   layer hasn't learned enough structure; low alpha means overcorrelation.")
+        print(f"   {DIM}Healthy: 2-4  |  Undertrained: 4-6  |  Overcorrelated: <2  |  Severe: >6{RESET}")
+        print(f"   {'Component':<35} {'Wt Matrices':>11}  {'Mean α':>8}  {'Min α':>8}  {'Max α':>8}  Status")
+        print(f"   {'-' * 77}")
+        for comp_name, layers in ww_results.items():
+            alphas = [l["alpha"] for l in layers if l.get("alpha") is not None]
+            if not alphas:
+                print(f"   {comp_name:<35} {'—':>11}  {'N/A':>8}  {'N/A':>8}  {'N/A':>8}  {DIM}too small{RESET}")
+                continue
+            mean_a = sum(alphas) / len(alphas)
+            min_a = min(alphas)
+            max_a = max(alphas)
+            status, _, _ = classify_health(alpha=mean_a)
+            print(f"   {comp_name:<35} {len(layers):>11}  {mean_a:>8.2f}  {min_a:>8.2f}  {max_a:>8.2f}  {status}")
+    else:
+        print(f"\n{BOLD}1. Weight Spectral Analysis{RESET}")
+        print(f"   {DIM}Skipped (weightwatcher not available){RESET}")
+
+    # --- Section 2: Attention Entropy ---
+    if entropy_results:
+        print(f"\n{BOLD}2. Attention Entropy (fraction of max){RESET}")
+        print(f"   Measures how spread out each attention head's focus is.")
+        print(f"   Low entropy means the head attends to very few tokens (collapsed/dead).")
+        print(f"   High entropy means the head spreads attention nearly uniformly (unfocused).")
+        print(f"   Healthy heads are selective but not degenerate — attending to a meaningful subset.")
+        print(f"   {DIM}Collapsed: <{thresholds.get('entropy_low', 0.1):.2f}  |  "
+              f"Healthy: {thresholds.get('entropy_low', 0.1):.2f}-{thresholds.get('entropy_warn', 0.8):.2f}  |  "
+              f"Unfocused: >{thresholds.get('entropy_warn', 0.8):.2f}  |  "
+              f"Dead: >{thresholds.get('entropy_critical', 0.95):.2f}{RESET}")
+
+        from collections import defaultdict
+        for comp_name, comp_entries in entropy_results.items():
+            print(f"\n   {BOLD}{comp_name}{RESET}")
+            by_layer = defaultdict(list)
+            for e in comp_entries:
+                by_layer[e["layer"]].append(e["entropy_ratio"])
+
+            print(f"   {'Layer':>6}  {'Mean Ent':>10}  {'Min Ent':>10}  {'Max Ent':>10}  Status")
+            print(f"   {'-' * 56}")
+            for layer in sorted(by_layer.keys()):
+                vals = by_layer[layer]
+                mean_e = sum(vals) / len(vals)
+                min_e = min(vals)
+                max_e = max(vals)
+                status, _, _ = classify_health(entropy=mean_e, thresholds=thresholds)
+                print(f"   {layer:>6}  {mean_e:>10.4f}  {min_e:>10.4f}  {max_e:>10.4f}  {status}")
+    else:
+        print(f"\n{BOLD}2. Attention Entropy{RESET}")
+        print(f"   {DIM}No data{RESET}")
+
+    # --- Section 3: Head Redundancy ---
+    if redundancy_results:
+        print(f"\n{BOLD}3. Head Redundancy (cosine similarity){RESET}")
+        print(f"   Measures how similar the attention heads are to each other within each layer.")
+        print(f"   Each layer has multiple heads that should learn different patterns (e.g., one")
+        print(f"   head for spatial relations, another for color). High similarity means heads are")
+        print(f"   redundant — wasted capacity. Collapsed means nearly identical heads.")
+        print(f"   {DIM}Diverse: <{thresholds.get('redundancy_warn', 0.7):.2f}  |  "
+              f"High: >{thresholds.get('redundancy_warn', 0.7):.2f}  |  "
+              f"Collapsed: >{thresholds.get('redundancy_critical', 0.9):.2f}{RESET}")
+
+        for comp_name, comp_entries in redundancy_results.items():
+            print(f"\n   {BOLD}{comp_name}{RESET}")
+            print(f"   {'Layer':>6}  {'Mean Sim':>10}  {'Max Sim':>10}  Status")
+            print(f"   {'-' * 46}")
+            for r in comp_entries:
+                status, _, _ = classify_health(redundancy=r["mean_redundancy"], thresholds=thresholds)
+                print(f"   {r['layer']:>6}  {r['mean_redundancy']:>10.4f}  {r['max_redundancy']:>10.4f}  {status}")
+    else:
+        print(f"\n{BOLD}3. Head Redundancy{RESET}")
+        print(f"   {DIM}No data{RESET}")
+
+    print(f"\n{'=' * 78}\n")
+
+
+def _status_emoji(label):
+    """Map a plain health label to a markdown-friendly status indicator."""
+    if label in ("healthy", "diverse"):
+        return "OK"
+    elif label in ("undertrained", "unfocused", "high redundancy"):
+        return "WARN"
+    else:
+        return "CRITICAL"
+
+
+def generate_health_markdown(ww_results, entropy_results, redundancy_results, thresholds):
+    """
+    Build a Markdown report string from health diagnostics results.
+    """
+    from collections import defaultdict
+    lines = []
+    w = lines.append
+
+    w("# Model Health Report\n")
+
+    # --- Section 1: Weight Spectral Analysis ---
+    w("## 1. Weight Spectral Analysis (alpha)\n")
+    w("Fits a power-law to each weight matrix's singular values. "
+       "Alpha measures how well-trained a layer is — values of 2-4 indicate "
+       "strong correlation structure learned during training. High alpha means the "
+       "layer hasn't learned enough structure; low alpha means overcorrelation.\n")
+    if ww_results is not None:
+        w("> Healthy: 2-4 | Undertrained: 4-6 | Overcorrelated: <2 | Severe: >6\n")
+        w("| Component | Wt Matrices | Mean alpha | Min alpha | Max alpha | Status |")
+        w("|-----------|------------:|-----------:|----------:|----------:|--------|")
+        for comp_name, layers in ww_results.items():
+            alphas = [l["alpha"] for l in layers if l.get("alpha") is not None]
+            if not alphas:
+                w(f"| {comp_name} | -- | N/A | N/A | N/A | too small |")
+                continue
+            mean_a = sum(alphas) / len(alphas)
+            min_a = min(alphas)
+            max_a = max(alphas)
+            _, _, label = classify_health(alpha=mean_a)
+            badge = _status_emoji(label)
+            w(f"| {comp_name} | {len(layers)} | {mean_a:.2f} | {min_a:.2f} | {max_a:.2f} | {badge} {label} |")
+    else:
+        w("*Skipped (weightwatcher not available)*\n")
+
+    # --- Section 2: Attention Entropy ---
+    w("\n## 2. Attention Entropy (fraction of max)\n")
+    w("Measures how spread out each attention head's focus is. "
+       "Low entropy means the head attends to very few tokens (collapsed/dead). "
+       "High entropy means the head spreads attention nearly uniformly (unfocused). "
+       "Healthy heads are selective but not degenerate — attending to a meaningful subset.\n")
+    if entropy_results:
+        low = thresholds.get("entropy_low", 0.1)
+        warn = thresholds.get("entropy_warn", 0.8)
+        crit = thresholds.get("entropy_critical", 0.95)
+        w(f"> Collapsed: <{low:.2f} | Healthy: {low:.2f}-{warn:.2f} | Unfocused: >{warn:.2f} | Dead: >{crit:.2f}\n")
+
+        for comp_name, comp_entries in entropy_results.items():
+            w(f"\n### {comp_name}\n")
+            by_layer = defaultdict(list)
+            for e in comp_entries:
+                by_layer[e["layer"]].append(e["entropy_ratio"])
+
+            w("| Layer | Mean Ent | Min Ent | Max Ent | Status |")
+            w("|------:|---------:|--------:|--------:|--------|")
+            for layer in sorted(by_layer.keys()):
+                vals = by_layer[layer]
+                mean_e = sum(vals) / len(vals)
+                min_e = min(vals)
+                max_e = max(vals)
+                _, _, label = classify_health(entropy=mean_e, thresholds=thresholds)
+                badge = _status_emoji(label)
+                w(f"| {layer} | {mean_e:.4f} | {min_e:.4f} | {max_e:.4f} | {badge} {label} |")
+    else:
+        w("*No data*\n")
+
+    # --- Section 3: Head Redundancy ---
+    w("\n## 3. Head Redundancy (cosine similarity)\n")
+    w("Measures how similar the attention heads are to each other within each layer. "
+       "Each layer has multiple heads that should learn different patterns (e.g., one "
+       "head for spatial relations, another for color). High similarity means heads are "
+       "redundant — wasted capacity. Collapsed means nearly identical heads.\n")
+    if redundancy_results:
+        r_warn = thresholds.get("redundancy_warn", 0.7)
+        r_crit = thresholds.get("redundancy_critical", 0.9)
+        w(f"> Diverse: <{r_warn:.2f} | High: >{r_warn:.2f} | Collapsed: >{r_crit:.2f}\n")
+
+        for comp_name, comp_entries in redundancy_results.items():
+            w(f"\n### {comp_name}\n")
+            w("| Layer | Mean Sim | Max Sim | Status |")
+            w("|------:|---------:|--------:|--------|")
+            for r in comp_entries:
+                _, _, label = classify_health(redundancy=r["mean_redundancy"], thresholds=thresholds)
+                badge = _status_emoji(label)
+                w(f"| {r['layer']} | {r['mean_redundancy']:.4f} | {r['max_redundancy']:.4f} | {badge} {label} |")
+    else:
+        w("*No data*\n")
+
+    return "\n".join(lines)
+
+
+def plot_health_report(ww_results, entropy_results, redundancy_results, output_path):
+    """
+    3-panel vertical matplotlib figure saved to *output_path*.
+
+    Panel 1: alpha per layer (bars) with reference lines at 2 and 6
+    Panel 2: mean entropy ratio per layer
+    Panel 3: mean head redundancy per layer
+    """
+    from collections import defaultdict
+
+    fig, axes = plt.subplots(3, 1, figsize=(14, 12), constrained_layout=True)
+
+    # --- Panel 1: Spectral Alpha ---
+    ax1 = axes[0]
+    if ww_results is not None:
+        bar_labels = []
+        bar_vals = []
+        bar_colors = []
+        separator_positions = []
+        label_positions = []
+        offset = 0
+
+        for comp_name, layers in ww_results.items():
+            alphas = [l["alpha"] for l in layers if l.get("alpha") is not None]
+            if not alphas:
+                continue
+            comp_start = offset
+            for i, a in enumerate(alphas):
+                bar_labels.append(f"L{layers[i]['layer']}")
+                bar_vals.append(a)
+                if a < 2:
+                    bar_colors.append("#e74c3c")       # red
+                elif a <= 4:
+                    bar_colors.append("#2ecc71")       # green
+                elif a <= 6:
+                    bar_colors.append("#f39c12")       # yellow
+                else:
+                    bar_colors.append("#e74c3c")       # red
+                offset += 1
+            comp_end = offset
+            label_positions.append(((comp_start + comp_end - 1) / 2, comp_name))
+            separator_positions.append(offset - 0.5)
+
+        # Remove last separator
+        if separator_positions:
+            separator_positions.pop()
+
+        if bar_vals:
+            x = range(len(bar_vals))
+            ax1.bar(x, bar_vals, color=bar_colors, edgecolor="white", linewidth=0.5)
+            ax1.axhline(y=2, color="green", linestyle="--", linewidth=1, label="α=2 (lower healthy)")
+            ax1.axhline(y=6, color="red", linestyle="--", linewidth=1, label="α=6 (upper healthy)")
+            ax1.set_ylabel("Alpha (α)")
+            ax1.legend(loc="upper right", fontsize=8)
+
+            for sep_x in separator_positions:
+                ax1.axvline(x=sep_x, color="#888888", linestyle="-", linewidth=0.8, alpha=0.5)
+
+            if len(bar_labels) > 30:
+                step = max(1, len(bar_labels) // 20)
+                ax1.set_xticks(range(0, len(bar_labels), step))
+                ax1.set_xticklabels([bar_labels[i] for i in range(0, len(bar_labels), step)],
+                                    rotation=45, ha="right", fontsize=6)
+            else:
+                ax1.set_xticks(x)
+                ax1.set_xticklabels(bar_labels, rotation=45, ha="right", fontsize=6)
+
+            # Component labels above bars
+            for x_center, comp_label in label_positions:
+                ax1.text(x_center, 1.02, comp_label, ha="center", va="bottom",
+                         fontsize=7, fontweight="bold", transform=ax1.get_xaxis_transform())
+        else:
+            ax1.text(0.5, 0.5, "No alpha data (layers too small)",
+                     ha="center", va="center", transform=ax1.transAxes)
+    else:
+        ax1.text(0.5, 0.5, "WeightWatcher not available",
+                 ha="center", va="center", transform=ax1.transAxes)
+    ax1.set_title("Weight Spectral Analysis — Alpha per Layer", fontweight="bold", pad=20)
+
+    # --- Panel 2: Attention Entropy (grouped by component) ---
+    ax2 = axes[1]
+    if entropy_results:
+        bar_labels = []
+        bar_means = []
+        bar_colors = []
+        separator_positions = []  # x positions for vertical lines between components
+        label_positions = []      # (x_center, comp_name) for component labels
+        offset = 0
+
+        for comp_name, comp_entries in entropy_results.items():
+            by_layer = defaultdict(list)
+            for e in comp_entries:
+                by_layer[e["layer"]].append(e["entropy_ratio"])
+            layers_sorted = sorted(by_layer.keys())
+            comp_start = offset
+
+            for layer in layers_sorted:
+                vals = by_layer[layer]
+                mean_e = sum(vals) / len(vals)
+                bar_labels.append(f"L{layer}")
+                bar_means.append(mean_e)
+                if mean_e >= 0.95:
+                    bar_colors.append("#e74c3c")
+                elif mean_e >= 0.8:
+                    bar_colors.append("#f39c12")
+                elif mean_e <= 0.1:
+                    bar_colors.append("#e74c3c")
+                else:
+                    bar_colors.append("#2ecc71")
+                offset += 1
+
+            comp_end = offset
+            label_positions.append(((comp_start + comp_end - 1) / 2, comp_name))
+            if comp_end < sum(len(v) for v in [defaultdict(list)] * 0) or offset > 0:
+                separator_positions.append(offset - 0.5)
+
+        # Remove last separator (no line after the last component)
+        if separator_positions:
+            separator_positions.pop()
+
+        ax2.bar(range(len(bar_means)), bar_means, color=bar_colors, edgecolor="white", linewidth=0.5)
+        ax2.axhline(y=0.8, color="#f39c12", linestyle="--", linewidth=1, label="warn (0.8)")
+        ax2.axhline(y=0.95, color="#e74c3c", linestyle="--", linewidth=1, label="critical (0.95)")
+        ax2.axhline(y=0.1, color="#e74c3c", linestyle=":", linewidth=1, label="collapsed (0.1)")
+
+        # Vertical separators between components
+        for sep_x in separator_positions:
+            ax2.axvline(x=sep_x, color="#888888", linestyle="-", linewidth=0.8, alpha=0.5)
+
+        # Component labels at top
+        for x_center, comp_label in label_positions:
+            ax2.text(x_center, 1.02, comp_label, ha="center", va="bottom",
+                     fontsize=7, fontweight="bold", transform=ax2.get_xaxis_transform())
+
+        if len(bar_labels) > 40:
+            step = max(1, len(bar_labels) // 30)
+            ax2.set_xticks(range(0, len(bar_labels), step))
+            ax2.set_xticklabels([bar_labels[i] for i in range(0, len(bar_labels), step)],
+                                fontsize=6, rotation=45, ha="right")
+        else:
+            ax2.set_xticks(range(len(bar_labels)))
+            ax2.set_xticklabels(bar_labels, fontsize=6, rotation=45, ha="right")
+        ax2.set_ylabel("Entropy / max entropy")
+        ax2.set_ylim(0, 1.12)
+        ax2.legend(loc="upper right", fontsize=8)
+    else:
+        ax2.text(0.5, 0.5, "No entropy data", ha="center", va="center", transform=ax2.transAxes)
+    ax2.set_title("Attention Entropy per Layer (mean across heads)", fontweight="bold", pad=20)
+
+    # --- Panel 3: Head Redundancy (grouped by component) ---
+    ax3 = axes[2]
+    if redundancy_results:
+        bar_labels = []
+        bar_means = []
+        bar_colors = []
+        separator_positions = []
+        label_positions = []
+        offset = 0
+
+        for comp_name, comp_entries in redundancy_results.items():
+            comp_start = offset
+            for r in comp_entries:
+                bar_labels.append(f"L{r['layer']}")
+                bar_means.append(r["mean_redundancy"])
+                m = r["mean_redundancy"]
+                if m >= 0.9:
+                    bar_colors.append("#e74c3c")
+                elif m >= 0.7:
+                    bar_colors.append("#f39c12")
+                else:
+                    bar_colors.append("#2ecc71")
+                offset += 1
+
+            comp_end = offset
+            label_positions.append(((comp_start + comp_end - 1) / 2, comp_name))
+            separator_positions.append(offset - 0.5)
+
+        # Remove last separator
+        if separator_positions:
+            separator_positions.pop()
+
+        ax3.bar(range(len(bar_means)), bar_means, color=bar_colors, edgecolor="white", linewidth=0.5)
+        ax3.axhline(y=0.7, color="#f39c12", linestyle="--", linewidth=1, label="warn (0.7)")
+        ax3.axhline(y=0.9, color="#e74c3c", linestyle="--", linewidth=1, label="critical (0.9)")
+
+        for sep_x in separator_positions:
+            ax3.axvline(x=sep_x, color="#888888", linestyle="-", linewidth=0.8, alpha=0.5)
+
+        for x_center, comp_label in label_positions:
+            ax3.text(x_center, 1.02, comp_label, ha="center", va="bottom",
+                     fontsize=7, fontweight="bold", transform=ax3.get_xaxis_transform())
+
+        if len(bar_labels) > 40:
+            step = max(1, len(bar_labels) // 30)
+            ax3.set_xticks(range(0, len(bar_labels), step))
+            ax3.set_xticklabels([bar_labels[i] for i in range(0, len(bar_labels), step)],
+                                fontsize=6, rotation=45, ha="right")
+        else:
+            ax3.set_xticks(range(len(bar_labels)))
+            ax3.set_xticklabels(bar_labels, fontsize=6, rotation=45, ha="right")
+        ax3.set_ylabel("Mean cosine similarity")
+        ax3.set_ylim(0, 1.12)
+        ax3.legend(loc="upper right", fontsize=8)
+    else:
+        ax3.text(0.5, 0.5, "No redundancy data", ha="center", va="center", transform=ax3.transAxes)
+    ax3.set_title("Head Redundancy per Layer", fontweight="bold", pad=20)
+
+    plt.savefig(output_path, dpi=150, bbox_inches="tight", facecolor="white")
+    plt.close()
+    print(f"  Saved health report plot: {output_path}")
+
+
+def run_model_health_report(policy, dataset, args):
+    """
+    Orchestrator for ``--model-health`` mode.
+
+    Step 1: Spectral alpha via WeightWatcher (no data needed).
+    Step 2: Sample frames, run forward passes, compute entropy + redundancy.
+    Step 3: Print terminal report and save plot.
+    """
+    device = torch.device(args.device)
+
+    thresholds = {
+        "entropy_warn": args.entropy_warn,
+        "entropy_critical": args.entropy_critical,
+        "entropy_low": args.entropy_low,
+        "redundancy_warn": args.redundancy_warn,
+        "redundancy_critical": args.redundancy_critical,
+    }
+
+    # ------------------------------------------------------------------
+    # Step 1: WeightWatcher spectral analysis
+    # ------------------------------------------------------------------
+    print(f"\n{'=' * 70}")
+    print("MODEL HEALTH DIAGNOSTICS")
+    print(f"{'=' * 70}")
+    print("\n[1/3] Running WeightWatcher spectral analysis...")
+    ww_results = compute_weightwatcher_alpha(policy)
+
+    # ------------------------------------------------------------------
+    # Step 2: Sample frames, capture per-head attention, compute metrics
+    # ------------------------------------------------------------------
+    print(f"\n[2/3] Capturing attention maps over {args.health_frames} frames...")
+
+    # Find vision encoder and set up hooks
+    vision_encoder = find_vision_encoder(policy)
+    if vision_encoder is None:
+        print("  ERROR: Could not find vision encoder. Skipping attention metrics.")
+        entropy_results = {}
+        redundancy_results = {}
+    else:
+        # Force eager attention on SigLIP
+        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"
+
+        # Force per-head weights (disable averaging) on SigLIP MHA
+        mha_modules = []
+        for mod in vision_encoder.modules():
+            if isinstance(mod, torch.nn.MultiheadAttention):
+                mha_modules.append((mod, getattr(mod, "average_attn_weights", True)))
+                mod.average_attn_weights = False
+
+        # Set up capture objects
+        sigclip_capture = SigLIPAttentionCapture()
+        sigclip_capture.register_hooks(vision_encoder)
+
+        vlm_with_expert = policy.model.vlm_with_expert
+
+        # SmolVLA bypasses self_attn.forward() — all attention flows through
+        # eager_attention_forward(). DecoderAttentionCapture patches that
+        # method to capture both self-attention (prefill) and cross-attention
+        # (expert → VLM prefix).
+        decoder_capture = DecoderAttentionCapture()
+        print("  Setting up decoder attention capture (eager_attention_forward)...")
+        decoder_capture.register(vlm_with_expert)
+
+        # Resolve image key
+        image_key = args.image_key
+        if image_key is None:
+            image_keys = find_image_keys(dataset)
+            if image_keys:
+                image_key = image_keys[0]
+            else:
+                print("  ERROR: No image keys found in dataset.")
+                sigclip_capture.clear()
+                decoder_capture.clear()
+                for mod, orig in mha_modules:
+                    mod.average_attn_weights = orig
+                entropy_results = {}
+                redundancy_results = {}
+                image_key = None
+
+        if image_key is not None:
+            # Sample frames
+            frame_pairs = get_episode_frames(
+                dataset, args.episode, args.health_frames, image_key,
+            )
+
+            # Determine SigLIP patch grid for entropy normalisation
+            patch_size = getattr(vision_encoder, "patch_size", None) or getattr(
+                getattr(vision_encoder, "config", None), "patch_size", 14,
+            )
+            img_size = getattr(
+                getattr(vision_encoder, "config", None), "image_size", 384,
+            )
+            n_patches_side = img_size // patch_size
+            sigclip_num_patches = n_patches_side * n_patches_side
+
+            # Component names for the three attention sources
+            COMP_SIGLIP = "SigLIP Vision (12L, 12H)"
+            COMP_DECODER_SA = "VLM+Expert Joint Self-Attn (16L, 15H)"
+            COMP_EXPERT_XA = "Expert-to-VLM Cross-Attn (16L, 8H)"
+
+            all_entropy = {COMP_SIGLIP: [], COMP_DECODER_SA: [], COMP_EXPERT_XA: []}
+            all_redundancy = {COMP_SIGLIP: [], COMP_DECODER_SA: [], COMP_EXPERT_XA: []}
+
+            policy.eval()
+
+            for i, (frame_idx, img_tensor) in enumerate(frame_pairs):
+                sigclip_capture.reset_maps()
+                decoder_capture.reset_maps()
+                sample = dataset[frame_idx]
+
+                with torch.no_grad():
+                    # Use full policy forward so we capture all layers
+                    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"    Frame {i} forward pass error: {e}")
+                        continue
+
+                # --- SigLIP ---
+                counts = []
+                siglip_attns = sigclip_capture.get_all_layer_attentions()
+                if siglip_attns:
+                    frame_ent = compute_attention_entropy(siglip_attns, sigclip_num_patches)
+                    frame_red = compute_head_redundancy(siglip_attns)
+                    all_entropy[COMP_SIGLIP].append(frame_ent)
+                    all_redundancy[COMP_SIGLIP].append(frame_red)
+                    counts.append(f"SigLIP={len(siglip_attns)}")
+                else:
+                    counts.append("SigLIP=0")
+
+                # --- Decoder Self-Attention (prefill: VLM+Expert concatenated) ---
+                sa_attns = decoder_capture.get_self_attn_layers()
+                if sa_attns:
+                    sa_num_tokens = sa_attns[0][1].shape[-1]
+                    frame_ent = compute_attention_entropy(sa_attns, sa_num_tokens)
+                    frame_red = compute_head_redundancy(sa_attns)
+                    all_entropy[COMP_DECODER_SA].append(frame_ent)
+                    all_redundancy[COMP_DECODER_SA].append(frame_red)
+                    counts.append(f"Decoder SA={len(sa_attns)}")
+                else:
+                    counts.append("Decoder SA=0")
+
+                # --- Expert Cross-Attention ---
+                xa_attns_raw = decoder_capture.get_cross_attn_layers()
+                if xa_attns_raw:
+                    # Remap sequential indices → actual layer indices.
+                    # During generation, each autoregressive step runs through
+                    # all decoder layers, so call_idx % num_layers = layer.
+                    num_dec_layers = vlm_with_expert.num_vlm_layers
+                    xa_attns = [(idx % num_dec_layers, probs) for idx, probs in xa_attns_raw]
+                    xa_num_tokens = xa_attns[0][1].shape[-1]
+                    frame_ent = compute_attention_entropy(xa_attns, xa_num_tokens)
+                    frame_red = compute_head_redundancy(xa_attns)
+                    all_entropy[COMP_EXPERT_XA].append(frame_ent)
+                    all_redundancy[COMP_EXPERT_XA].append(frame_red)
+                    counts.append(f"Expert XA={len(xa_attns_raw)}")
+                else:
+                    counts.append("Expert XA=0")
+
+                print(f"    Frame {i}: {', '.join(counts)}")
+
+            # Cleanup
+            sigclip_capture.clear()
+            decoder_capture.clear()
+
+            # Restore MHA averaging setting
+            for mod, orig in mha_modules:
+                mod.average_attn_weights = orig
+
+            # Average metrics across frames, per component
+            from collections import defaultdict
+
+            entropy_results = {}
+            for comp_name, frame_list in all_entropy.items():
+                if not frame_list:
+                    continue
+                ent_accum = defaultdict(lambda: defaultdict(list))
+                for frame_ent in frame_list:
+                    for e in frame_ent:
+                        ent_accum[(e["layer"], e["head"])]["entropy_ratio"].append(e["entropy_ratio"])
+                        ent_accum[(e["layer"], e["head"])]["entropy"].append(e["entropy"])
+
+                comp_results = []
+                for (layer, head), vals in sorted(ent_accum.items()):
+                    comp_results.append({
+                        "layer": layer,
+                        "head": head,
+                        "entropy": sum(vals["entropy"]) / len(vals["entropy"]),
+                        "entropy_ratio": sum(vals["entropy_ratio"]) / len(vals["entropy_ratio"]),
+                    })
+                if comp_results:
+                    entropy_results[comp_name] = comp_results
+
+            redundancy_results = {}
+            for comp_name, frame_list in all_redundancy.items():
+                if not frame_list:
+                    continue
+                red_accum = defaultdict(lambda: {"mean": [], "max": []})
+                for frame_red in frame_list:
+                    for r in frame_red:
+                        red_accum[r["layer"]]["mean"].append(r["mean_redundancy"])
+                        red_accum[r["layer"]]["max"].append(r["max_redundancy"])
+
+                comp_results = []
+                for layer in sorted(red_accum.keys()):
+                    vals = red_accum[layer]
+                    comp_results.append({
+                        "layer": layer,
+                        "mean_redundancy": sum(vals["mean"]) / len(vals["mean"]),
+                        "max_redundancy": sum(vals["max"]) / len(vals["max"]),
+                    })
+                if comp_results:
+                    redundancy_results[comp_name] = comp_results
+
+    # ------------------------------------------------------------------
+    # Step 3: Print report and save plot
+    # ------------------------------------------------------------------
+    print(f"\n[3/3] Generating report...")
+    print_health_report(ww_results, entropy_results, redundancy_results, thresholds)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    md_path = os.path.join(args.output_dir, "model_health_report.md")
+    md_text = generate_health_markdown(ww_results, entropy_results, redundancy_results, thresholds)
+    with open(md_path, "w") as f:
+        f.write(md_text)
+    print(f"  Saved markdown report: {md_path}")
+
+    plot_path = os.path.join(args.output_dir, "model_health_report.png")
+    plot_health_report(ww_results, entropy_results, redundancy_results, plot_path)
+
+    print(f"Done! Reports saved to {args.output_dir}/")
+
+
+# ---------------------------------------------------------------------------
+# 9. Entry point
 # ---------------------------------------------------------------------------
 
 def load_defaults():
@@ -1562,6 +2537,29 @@ Examples:
                         default=defaults.get("raw_attention", False),
                         help="Skip positional baseline subtraction (show raw attention)")
 
+    # Model health diagnostics
+    parser.add_argument("--model-health", action="store_true",
+                        default=defaults.get("model_health", False),
+                        help="Run health diagnostics instead of attention heatmaps")
+    parser.add_argument("--health-frames", type=int,
+                        default=defaults.get("health_frames", 5),
+                        help="Number of sample frames for entropy/redundancy (default: 5)")
+    parser.add_argument("--entropy-warn", type=float,
+                        default=defaults.get("entropy_warn", 0.8),
+                        help="Entropy ratio threshold for 'unfocused' warning (default: 0.8)")
+    parser.add_argument("--entropy-critical", type=float,
+                        default=defaults.get("entropy_critical", 0.95),
+                        help="Entropy ratio threshold for 'uniform/dead' (default: 0.95)")
+    parser.add_argument("--entropy-low", type=float,
+                        default=defaults.get("entropy_low", 0.1),
+                        help="Entropy ratio threshold for 'collapsed' (default: 0.1)")
+    parser.add_argument("--redundancy-warn", type=float,
+                        default=defaults.get("redundancy_warn", 0.7),
+                        help="Cosine similarity threshold for 'high redundancy' (default: 0.7)")
+    parser.add_argument("--redundancy-critical", type=float,
+                        default=defaults.get("redundancy_critical", 0.9),
+                        help="Cosine similarity threshold for 'collapsed' (default: 0.9)")
+
     args = parser.parse_args()
 
     os.makedirs(args.output_dir, exist_ok=True)
@@ -1639,7 +2637,12 @@ Examples:
     except Exception as e:
         print(f"\n  ERROR loading dataset: {e}")
         sys.exit(1)
-    
+
+    # --- Model health mode (early exit) ---
+    if args.model_health:
+        run_model_health_report(policy, dataset, args)
+        return
+
     # --- Extract attention maps ---
     print(f"\n[Step 3] Extracting attention maps...")
     

requirements.txt

@@ -7,3 +7,4 @@ matplotlib>=3.7
 numpy>=1.24
 Pillow>=10.0
 pyyaml>=6.0
+weightwatcher>=0.7