smolvla_inspect/cli.py

"""
CLI entry point — extract_attention_maps, gradient_attention_map, load_defaults, main.

Imports from all other package modules.
"""

import argparse
import logging
import math
import os
import sys
import warnings
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt

try:
    from tqdm import tqdm
except ImportError:
    def tqdm(iterable, **kwargs):
        return iterable

from ._compat import resize_with_pad
from .capture import SigLIPAttentionCapture, ActionVisionAttentionCapture
from .heatmap import (
    compute_padding_patches,
    attention_to_heatmap,
    compute_patch_attention_scores,
    compute_attention_rollout,
    compute_positional_baseline,
)
from .data import (
    find_vision_encoder,
    find_image_keys,
    _resolve_task_string,
    build_policy_batch_from_sample,
    get_episode_frames,
    parse_image_map,
)
from .viz import (
    create_visualization_grid,
    save_individual_frames,
    create_per_head_grid,
    create_per_step_cross_attn_grid,
    create_vlm_layer_grid,
    create_per_action_dim_grid,
    create_language_diff_grid,
    create_vision_vs_state_chart,
    overlay_heatmap,
)
from .internals import run_model_internals_report
from .gradient import (
    compute_gradient_maps,
    compute_saliency_map,
    compute_gradcam_connector_maps,
    compute_gradcam_vlm_layers_maps,
    compute_vision_vs_state_maps,
    compute_per_action_dim_maps,
    compute_language_conditional_maps,
)
from .export import (
    create_run_dir,
    save_frames,
    save_self_attention,
    save_cross_attention,
    save_gradient_data,
    save_model_internals_data,
    build_manifest,
    collect_model_info,
    collect_dataset_info,
    collect_image_paths,
    build_available_viz,
)


def extract_attention_maps(policy, dataset, episode_idx=0, num_frames=8,
                           image_key=None, device="cpu",
                           method="last-layer", cross_attention=False,
                           show_heads=False, output_dir="./outputs",
                           raw_attention=False, attn_threshold=0.5,
                           task_override=None, per_step_cross_attention=False,
                           image_map=None):
    """
    Core function: Run inference and extract attention heatmaps.

    Args:
        method: ``"last-layer"`` (default), ``"rollout"``, or
                ``"all-layers"`` — how to aggregate SigLIP self-attention.
        cross_attention: If *True* also capture action-expert → vision
                         cross-attention (requires full ``select_action``
                         forward pass, slower).
        show_heads: If *True* save a per-head grid for the first frame.
        output_dir: Where to save per-head grids (only used when
                    *show_heads* is True).
        raw_attention: If *True* skip positional baseline subtraction
                       (show raw, uncorrected attention).
        attn_threshold: Percentile (0–1) below which attention values are
                        zeroed to suppress residual positional noise.
        task_override: If not *None*, use this string as the language
                       instruction instead of the one from the dataset.
        per_step_cross_attention: If *True*, also return per-denoising-step
                                  cross-attention maps.

    Returns:
        frames: list of image tensors (C, H, W)
        heatmaps: list of numpy heatmaps (H, W) in [0, 1]
        actions: list of predicted actions (or None)
        cross_attn_heatmaps: list of numpy heatmaps or *None*
        per_step_cross_attn: list of (list of Tensors per step) per frame, or *None*
    """

    # --- Find vision encoder ---
    print("\n[1/4] Locating vision encoder...")
    vision_encoder = find_vision_encoder(policy)

    if vision_encoder is None:
        print("  Could not find vision encoder. Dumping model structure:")
        for name, module in policy.named_modules():
            print(f"    {name}: {type(module).__name__}")
        raise RuntimeError("Cannot find SigLIP vision encoder in model")

    # --- Force eager attention so we get attention weights (SDPA/Flash return None) ---
    eager_count = 0
    for mod in vision_encoder.modules():
        if getattr(mod, "config", None) is not None and hasattr(mod.config, "_attn_implementation"):
            mod.config._attn_implementation = "eager"
            eager_count += 1
    if eager_count:
        print("  Using eager attention to capture weights")

    # --- Register attention hooks ---
    print("\n[2/4] Registering attention hooks...")
    attn_capture = SigLIPAttentionCapture()
    attn_capture.register_hooks(vision_encoder)

    # --- Optionally register cross-attention capture ---
    cross_capture = None
    vision_start = vision_end = n_vision_tokens = 0
    if cross_attention:
        try:
            vlm_with_expert = policy.model.vlm_with_expert
            cross_capture = ActionVisionAttentionCapture()
            cross_capture.register(vlm_with_expert)
            vision_start, vision_end, n_vision_tokens = cross_capture.get_vision_token_range(policy)
            print(f"  Cross-attention capture enabled (vision tokens {vision_start}..{vision_end}, "
                  f"n={n_vision_tokens})")
        except Exception as e:
            print(f"  WARNING: Could not set up cross-attention capture: {e}")
            cross_capture = None

    # --- Find image key in dataset ---
    if image_key is None:
        image_keys = find_image_keys(dataset)
        if not image_keys:
            raise ValueError("No image keys found in dataset. Available keys: " +
                           str(list(dataset[0].keys())))
        image_key = image_keys[0]
        print(f"  Using image key: {image_key}")
        if len(image_keys) > 1:
            print(f"  Other available image keys: {image_keys[1:]}")

    # --- Get frames ---
    print(f"\n[3/4] Loading {num_frames} frames from episode {episode_idx}...")
    frame_pairs = get_episode_frames(dataset, episode_idx, num_frames, image_key)

    # --- Log task string being used ---
    first_sample = dataset[frame_pairs[0][0]]
    task_str = _resolve_task_string(first_sample, dataset, task_override=task_override)
    print(f"  Task string: \"{task_str}\"")

    # --- Compute positional baseline (once, after frames are loaded so we
    #     know the input aspect ratio for a properly padded baseline) ---
    first_img = frame_pairs[0][1]  # (C, H, W)
    input_hw = (first_img.shape[1], first_img.shape[2])

    baseline_scores = None
    per_head_baseline = None
    if not raw_attention:
        baseline_scores, per_head_baseline = compute_positional_baseline(
            vision_encoder, attn_capture, device, method, input_hw=input_hw,
        )
        if baseline_scores is not None:
            print("  Positional baseline computed (subtracting to reveal content-dependent attention)")
        else:
            print("  Could not compute positional baseline, using raw attention")

    # --- Compute padding-patch crop so heatmaps exclude pad regions ---
    target_size = getattr(getattr(vision_encoder, "config", None), "image_size", None) or 384
    patch_size_cfg = getattr(vision_encoder, "patch_size", None) or getattr(
        getattr(vision_encoder, "config", None), "patch_size", 14)
    content_crop = compute_padding_patches(input_hw, target_size, patch_size_cfg)
    if content_crop != (0, 0):
        print(f"  Padding crop: {content_crop[0]} top rows, {content_crop[1]} left cols of patches")

    # --- Save positional baseline diagnostic heatmap ---
    if baseline_scores is not None:
        n_bl = baseline_scores.shape[0]
        bl_side = int(math.sqrt(n_bl))
        grid_h_bl = target_size // patch_size_cfg if target_size and patch_size_cfg else bl_side
        grid_w_bl = grid_h_bl
        img_h_bl, img_w_bl = first_img.shape[1], first_img.shape[2]
        bl_heatmap = attention_to_heatmap(
            baseline_scores, (grid_h_bl, grid_w_bl), (img_h_bl, img_w_bl),
            content_crop=content_crop,
        )
        bl_path = os.path.join(output_dir, "positional_baseline.png")
        fig_bl, ax_bl = plt.subplots(figsize=(6, 6))
        ax_bl.imshow(bl_heatmap, cmap="jet")
        ax_bl.set_title("Positional baseline (gray-image attention)", fontsize=11)
        ax_bl.axis("off")
        fig_bl.savefig(bl_path, dpi=100, bbox_inches="tight")
        plt.close(fig_bl)
        print(f"  Saved positional baseline diagnostic: {bl_path}")

    # --- Run inference and collect attention ---
    print(f"\n[4/4] Running forward passes and extracting attention (method={method})...")
    frames = []
    heatmaps = []
    cross_attn_heatmaps = [] if cross_capture else None
    per_step_cross_attn = [] if (per_step_cross_attention and cross_capture) else None
    actions = []

    policy.eval()

    # We need a full policy forward pass when cross-attention is requested
    # (the vision encoder hooks still fire during select_action too)
    use_full_forward = cross_attention and cross_capture is not None

    raw_heads_attn = None   # populated when show_heads is True (first frame only)
    n_patches_h = n_patches_w = 0  # set during vision-encoder-only forward

    for i, (frame_idx, img_tensor) in enumerate(tqdm(frame_pairs, desc="Attention extraction", unit="frame")):
        frames.append(img_tensor.clone())

        # Reset captured attention maps
        attn_capture.reset_maps()
        if cross_capture:
            cross_capture.reset_maps()

        sample = dataset[frame_idx]

        # ----- forward pass -----
        try:
            with torch.no_grad():
                if use_full_forward:
                    # Full policy forward — needed for cross-attention capture.
                    # Reset the action queue so each frame triggers a real
                    # forward pass (otherwise cached actions are returned).
                    policy.reset()
                    batch, _ = build_policy_batch_from_sample(
                        sample, policy, device, batch_size=1,
                        image_key_for_grad=None, dataset=dataset,
                        task_override=task_override, image_map=image_map,
                    )
                    try:
                        policy.select_action(batch)
                    except Exception as e:
                        print(f"    select_action error: {e}")
                else:
                    # Vision-encoder-only forward (faster)
                    img = img_tensor.unsqueeze(0).to(device)
                    if img.max() > 1.0:
                        img = img.float() / 255.0

                    target_size = getattr(
                        getattr(vision_encoder, "config", None), "image_size", None,
                    ) or getattr(vision_encoder, "image_size", 384)
                    if img.shape[-1] != target_size or img.shape[-2] != target_size:
                        img_resized = resize_with_pad(img, target_size, target_size, pad_value=0)
                    else:
                        img_resized = img
                    img_resized = img_resized * 2.0 - 1.0  # normalize to [-1, 1] matching SigLIP

                    try:
                        enc_dtype = next(vision_encoder.parameters()).dtype
                        img_resized = img_resized.to(enc_dtype)
                    except StopIteration:
                        pass

                    patch_size = getattr(vision_encoder, "patch_size", None) or getattr(
                        getattr(vision_encoder, "config", None), "patch_size", 14
                    )
                    n_patches_h = img_resized.size(2) // patch_size
                    n_patches_w = img_resized.size(3) // patch_size
                    patch_mask = torch.ones(1, n_patches_h, n_patches_w, dtype=torch.bool, device=device)

                    try:
                        if hasattr(vision_encoder, 'embeddings') and hasattr(vision_encoder, 'encoder'):
                            embeddings = vision_encoder.embeddings(img_resized, patch_mask)
                            encoder_out = vision_encoder.encoder(embeddings)
                        elif hasattr(vision_encoder, 'forward'):
                            encoder_out = vision_encoder(img_resized)
                        else:
                            encoder_out = vision_encoder(pixel_values=img_resized,
                                                          patch_attention_mask=patch_mask)
                    except Exception as e:
                        print(f"    Direct vision forward failed ({e}), trying full policy...")
                        batch, _ = build_policy_batch_from_sample(
                            sample, policy, device, batch_size=1,
                            image_key_for_grad=None, dataset=dataset,
                            task_override=task_override, image_map=image_map,
                        )
                        try:
                            policy.select_action(batch)
                        except Exception:
                            pass

        except Exception as e:
            print(f"    Frame {i} forward pass error: {e}")

        # ----- per-head capture (independent of summary method) -----
        if show_heads and i == 0 and raw_heads_attn is None:
            attn_for_heads = attn_capture.get_last_layer_attention()
            if attn_for_heads is not None:
                while attn_for_heads.dim() > 3:
                    attn_for_heads = attn_for_heads[0]
                if attn_for_heads.dim() == 3:
                    raw_heads_attn = attn_for_heads.clone()

        # ----- self-attention heatmap -----
        if method == "rollout":
            all_layers = attn_capture.get_all_layer_attentions()
            rollout_mat = compute_attention_rollout(all_layers)
            if rollout_mat is not None:
                patch_scores = rollout_mat.mean(dim=0)  # per-patch importance
            else:
                patch_scores = None
        elif method == "all-layers":
            # For "all-layers" we still produce a single summary heatmap
            # (averaging last-layer scores) but also save individual layer
            # grids elsewhere; here fall through to last-layer for the
            # summary heatmap.
            attn = attn_capture.get_last_layer_attention()
            patch_scores = None
            if attn is not None:
                while attn.dim() > 3:
                    attn = attn[0]
                if attn.dim() == 3:
                    attn = attn.mean(dim=0)
                patch_scores = compute_patch_attention_scores(attn, method="mean")
        else:
            # "last-layer" (default)
            attn = attn_capture.get_last_layer_attention()
            patch_scores = None
            if attn is not None:
                while attn.dim() > 3:
                    attn = attn[0]
                if attn.dim() == 3:
                    attn = attn.mean(dim=0)
                patch_scores = compute_patch_attention_scores(attn, method="mean")

        # Subtract positional baseline to isolate content-dependent signal
        if patch_scores is not None and baseline_scores is not None:
            patch_scores = torch.clamp(patch_scores - baseline_scores, min=0)

        if patch_scores is not None:
            n_patches = patch_scores.shape[0]
            grid_side = int(math.sqrt(n_patches))
            if grid_side * grid_side != n_patches:
                grid_h = n_patches_h if n_patches_h > 0 else grid_side
                grid_w = n_patches_w if n_patches_w > 0 else grid_side
            else:
                grid_h = grid_w = grid_side

            img_h, img_w = img_tensor.shape[1], img_tensor.shape[2]
            heatmap = attention_to_heatmap(patch_scores, (grid_h, grid_w), (img_h, img_w),
                                         content_crop=content_crop,
                                         threshold_pct=attn_threshold)
            heatmaps.append(heatmap)

            print(f"    Frame {i}: {n_patches} patches \u2192 "
                  f"{grid_h}x{grid_w} grid \u2192 {img_h}x{img_w} heatmap ({method})")

            # Per-head grid for first frame
            if show_heads and i == 0 and raw_heads_attn is not None:
                head_path = os.path.join(output_dir, f"per_head_ep{episode_idx:03d}.png")
                create_per_head_grid(
                    img_tensor, raw_heads_attn,
                    (grid_h, grid_w), (img_h, img_w),
                    output_path=head_path,
                    content_crop=content_crop,
                    baseline_per_head=per_head_baseline if not raw_attention else None,
                    threshold_pct=attn_threshold,
                )
                raw_heads_attn = None  # only once
        else:
            print(f"    Frame {i}: No attention captured, using uniform heatmap")
            img_h, img_w = img_tensor.shape[1], img_tensor.shape[2]
            heatmaps.append(np.ones((img_h, img_w)) * 0.5)

        # ----- cross-attention heatmap -----
        if cross_capture is not None:
            cross_scores = cross_capture.get_mean_cross_attention(vision_start, vision_end)
            if cross_scores is not None:
                # cross_scores: (n_vision_tokens,)
                # Vision tokens come from SigLIP → connector (pixel shuffle).
                # After pixel shuffle the spatial grid is halved in each dim.
                n_vis = cross_scores.shape[0]
                cs_side = int(math.sqrt(n_vis))
                if cs_side * cs_side != n_vis:
                    cs_h = cs_w = cs_side
                else:
                    cs_h = cs_w = cs_side

                img_h, img_w = img_tensor.shape[1], img_tensor.shape[2]
                cross_hm = attention_to_heatmap(cross_scores, (cs_h, cs_w), (img_h, img_w),
                                               threshold_pct=attn_threshold)
                cross_attn_heatmaps.append(cross_hm)
                print(f"    Frame {i}: Cross-attention captured ({n_vis} vision tokens)")
            else:
                img_h, img_w = img_tensor.shape[1], img_tensor.shape[2]
                cross_attn_heatmaps.append(np.ones((img_h, img_w)) * 0.5)
                print(f"    Frame {i}: No cross-attention captured, using uniform")

        # ----- per-step cross-attention -----
        if per_step_cross_attn is not None and cross_capture is not None:
            step_maps = cross_capture.get_per_step_cross_attention(
                vision_start, vision_end,
            )
            if step_maps is not None:
                per_step_cross_attn.append(step_maps)
                print(f"    Frame {i}: Per-step cross-attention: {len(step_maps)} steps")
            else:
                per_step_cross_attn.append([])

    # Cleanup
    attn_capture.clear()
    if cross_capture:
        cross_capture.clear()

    return frames, heatmaps, actions, cross_attn_heatmaps, per_step_cross_attn


def gradient_attention_map(policy, dataset, frame_idx, image_key, device="cpu",
                           task_override=None, image_map=None):
    """
    Compute input-gradient saliency map as a fallback.

    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]
    return compute_saliency_map(
        policy, sample, dataset, image_key, device,
        task_override=task_override, image_map=image_map,
    )


def load_defaults(config_path=None):
    """Load defaults from a YAML config file.

    Resolution order:
      1. Explicit *config_path* argument (from ``--config``)
      2. ``configs/defaults.yaml`` in the project root
    """
    if config_path is None:
        config_path = Path(__file__).resolve().parent.parent / "configs" / "defaults.yaml"
    else:
        config_path = Path(config_path)
    if config_path.exists():
        try:
            import yaml
            with open(config_path) as f:
                return yaml.safe_load(f) or {}
        except ImportError:
            pass
    elif config_path != Path(__file__).resolve().parent.parent / "configs" / "defaults.yaml":
        # Only error if the user explicitly asked for a config that doesn't exist
        print(f"ERROR: Config file not found: {config_path}")
        sys.exit(1)
    return {}


def _default_value(defaults, *keys, fallback=None):
    """Return the first configured key present in *defaults*."""
    for key in keys:
        if key in defaults:
            return defaults[key]
    return fallback


def add_pipeline_args(parser, defaults, exclude=None):
    """Add inspection pipeline arguments to a parser.

    These are the flags controlling attention extraction, gradient attribution,
    model internals, and output settings.  Shared between the main CLI and the
    ``diagnose`` subcommand (integrated mode).

    Parameters
    ----------
    parser : argparse.ArgumentParser
    defaults : dict
        Loaded YAML defaults (from ``load_defaults``).
    exclude : set | None
        Long option names (e.g. ``"--device"``) to skip because the caller
        already registered them.
    """
    exclude = set(exclude or [])

    def _add(name, **kwargs):
        if name in exclude:
            return
        parser.add_argument(name, **kwargs)

    _add("--num-frames", type=int,
         default=defaults.get("num_frames", 8))
    _add("--image-key", type=str, default=None)
    _add("--image-map", type=str, default=None,
         help="Explicit dataset\u2192policy image key mapping. "
              "Comma-separated pairs using = delimiter. "
              "Accepts full keys or short suffixes. "
              "Example: front=camera2,side=camera3")
    _add("--task", type=str, default=defaults.get("task", None),
         help="Override the task/language instruction (default: from dataset)")
    _add("--output-dir", type=str,
         default=defaults.get("output_dir", "./outputs"))
    _add("--device", type=str,
         default=defaults.get("device", "auto"),
         choices=["auto", "cpu", "cuda", "mps"])
    _add("--save-individual", action="store_true",
         default=defaults.get("save_individual", False))
    _add("--skip-attention", action="store_true",
         default=False,
         help="Skip hook-based attention extraction (Step 3); only run gradient features")
    _add("--export-data", action="store_true",
         default=defaults.get("export_data", True),
         help="Save structured .npz + JSON alongside PNGs (default: true)")
    if "--export-data" not in exclude:
        _add("--no-export-data", action="store_false", dest="export_data",
             help="Disable structured data export")
    _add("--run-name", type=str,
         default=defaults.get("run_name", None),
         help="Run folder name (default: timestamped run_YYYY-MM-DD_HH-MM-SS)")
    _add("--method", type=str,
         default=defaults.get("method", "last-layer"),
         choices=["last-layer", "rollout", "all-layers"],
         help="Self-attention aggregation method")
    _add("--cross-attention", action="store_true",
         default=defaults.get("cross_attention", False),
         help="Capture action-expert \u2192 vision cross-attention (slower)")
    _add("--show-heads", action="store_true",
         default=defaults.get("show_heads", False),
         help="Save a per-head attention grid for the first frame")
    _add("--raw-attention", action="store_true",
         default=defaults.get("raw_attention", False),
         help="Skip positional baseline subtraction (show raw attention)")
    _add("--attn-threshold", type=float,
         default=defaults.get("attn_threshold", 0.5),
         help="Percentile threshold (0-1) below which attention values are zeroed")

    # Model internals report
    _add("--internals-only", action="store_true",
         default=_default_value(defaults, "internals_only", "model_health", fallback=False),
         help="Run the model internals report instead of attention or gradient analysis")
    _add("--with-internals", action="store_true",
         default=defaults.get("with_internals", False),
         help="Also run the model internals report after the standard analysis")
    _add("--internals-frames", type=int,
         default=_default_value(defaults, "internals_frames", "health_frames", fallback=5),
         help="Number of sample frames for entropy/redundancy in the internals report (default: 5)")
    if "--internals-only" not in exclude:
        _add("--model-health", action="store_true", dest="internals_only",
             default=argparse.SUPPRESS, help=argparse.SUPPRESS)
    if "--internals-frames" not in exclude:
        _add("--health-frames", type=int, dest="internals_frames",
             default=argparse.SUPPRESS, help=argparse.SUPPRESS)
    _add("--entropy-warn", type=float,
         default=defaults.get("entropy_warn", 0.8),
         help="Entropy ratio threshold for 'unfocused' warning (default: 0.8)")
    _add("--entropy-critical", type=float,
         default=defaults.get("entropy_critical", 0.95),
         help="Entropy ratio threshold for 'uniform/dead' (default: 0.95)")
    _add("--entropy-low", type=float,
         default=defaults.get("entropy_low", 0.1),
         help="Entropy ratio threshold for 'collapsed' (default: 0.1)")
    _add("--redundancy-warn", type=float,
         default=defaults.get("redundancy_warn", 0.7),
         help="Cosine similarity threshold for 'high redundancy' (default: 0.7)")
    _add("--redundancy-critical", type=float,
         default=defaults.get("redundancy_critical", 0.9),
         help="Cosine similarity threshold for 'collapsed' (default: 0.9)")

    # Gradient-based attribution
    _add("--gradient", nargs="?", const="both",
         default=defaults.get("gradient", None),
         choices=["saliency", "gradcam", "both"],
         help="Gradient attribution method (default: off; bare --gradient means 'both')")
    _add("--gradient-device", type=str,
         default=defaults.get("gradient_device", None),
         choices=["cpu", "cuda", "mps"],
         help="Device for gradient attribution (default: same as --device)")
    _add("--gradient-seed", type=int,
         default=defaults.get("gradient_seed", 42),
         help="Fixed noise seed for reproducible gradient attribution (default: 42)")
    _add("--smooth-grad", type=int,
         default=defaults.get("smooth_grad", 1),
         help="SmoothGrad samples for saliency (1 = vanilla, >1 = averaged over N noisy inputs)")
    _add("--smooth-grad-sigma", type=float,
         default=defaults.get("smooth_grad_sigma", 0.15),
         help="Gaussian noise std for SmoothGrad (default: 0.15)")

    # Extended attribution features
    _add("--per-step-cross-attention", action="store_true",
         default=defaults.get("per_step_cross_attention", False),
         help="Visualize cross-attention at each denoising step")
    _add("--gradcam-connector", action="store_true",
         default=defaults.get("gradcam_connector", False),
         help="GradCAM on VLM connector output (post-pixel-shuffle)")
    _add("--gradcam-vlm-layers", nargs="?", const="4,8,12,16",
         default=defaults.get("gradcam_vlm_layers", None),
         help="GradCAM on VLM intermediate layers (comma-separated 1-indexed, default: 4,8,12,16)")
    _add("--vision-vs-state", action="store_true",
         default=defaults.get("vision_vs_state", False),
         help="Compare gradient attribution between vision and state inputs")
    _add("--per-action-dim", action="store_true",
         default=defaults.get("per_action_dim", False),
         help="Per-action-dimension GradCAM (uses retain_graph \u2014 GPU recommended)")
    _add("--language-diff", nargs="?", const="auto",
         default=defaults.get("language_diff", None),
         help="Language-conditional comparison (auto-select or provide alt task string)")


def run_inspection_pipeline(args, policy, dataset, device, grad_device, run_dir, image_map=None):
    """Run the full inspection pipeline (attention, gradients, internals, viz, export).

    Parameters
    ----------
    args : argparse.Namespace
        Parsed args containing all pipeline flags (from ``add_pipeline_args``).
    policy : SmolVLAPolicy
        Loaded and eval-mode policy on *device*.
    dataset : LeRobotDataset
    device : torch.device
    grad_device : torch.device
    run_dir : str | None
        Path to the run directory (already created), or ``None`` if export is off.
    image_map : dict | None

    Returns
    -------
    run_dir : str | None
    """
    internals_results = None
    model_internals_ran = False
    frames = None
    heatmaps = None
    actions = None
    cross_attn_heatmaps = None
    per_step_cross_attn = None
    saliency_maps = None
    gradcam_maps = None
    connector_gradcam_maps = None
    vlm_layer_results = None
    vlm_layer_indices = None
    vision_vs_state_results = None
    per_action_dim_maps = None
    per_action_dim_mags = None
    language_diff_results = None
    action_dim_names = None
    grid_path = None
    model_on_gradient_device = False

    grad_device_str = str(grad_device)

    if args.internals_only:
        internals_results = run_model_internals_report(policy, dataset, args)
        model_internals_ran = True
    else:
        # --- Extract attention maps ---
        if args.skip_attention:
            print(f"\n[Step 3] Skipping attention extraction (--skip-attention)")
            image_key = args.image_key or find_image_keys(dataset)[0]
            frame_pairs = get_episode_frames(dataset, args.episode, args.num_frames, image_key)
            frames = [img for _, img in frame_pairs]
            heatmaps = [np.ones((f.shape[1], f.shape[2])) * 0.5 for f in frames]
            actions = []
            print(f"  Loaded {len(frames)} frames (attention maps placeholder only)")
        else:
            print(f"\n[Step 3] Extracting attention maps...")
            try:
                frames, heatmaps, actions, cross_attn_heatmaps, per_step_cross_attn = extract_attention_maps(
                    policy=policy,
                    dataset=dataset,
                    episode_idx=args.episode,
                    num_frames=args.num_frames,
                    image_key=args.image_key,
                    device=str(device),
                    method=args.method,
                    cross_attention=args.cross_attention,
                    show_heads=args.show_heads,
                    output_dir=args.output_dir,
                    raw_attention=args.raw_attention,
                    attn_threshold=args.attn_threshold,
                    task_override=args.task,
                    per_step_cross_attention=args.per_step_cross_attention,
                    image_map=image_map,
                )
            except Exception as e:
                print(f"\n  Attention extraction failed: {e}")
                print("  Falling back to input-gradient saliency maps...")

                image_key = args.image_key or find_image_keys(dataset)[0]
                frame_pairs = get_episode_frames(dataset, args.episode, args.num_frames, image_key)

                frames = []
                heatmaps = []
                for frame_idx, img_tensor in frame_pairs:
                    frames.append(img_tensor)
                    saliency = gradient_attention_map(
                        policy, dataset, frame_idx, image_key, str(device),
                        task_override=args.task, image_map=image_map,
                    )
                    if saliency is not None:
                        heatmaps.append(saliency)
                    else:
                        heatmaps.append(np.ones((img_tensor.shape[1], img_tensor.shape[2])) * 0.5)

                actions = []

        if not frames:
            print("\nERROR: No frames extracted. Check episode index and dataset.")
            sys.exit(1)

        # --- Export attention data ---
        if run_dir:
            save_frames(run_dir, frames)
            save_self_attention(run_dir, heatmaps)
            if cross_attn_heatmaps:
                save_cross_attention(
                    run_dir,
                    cross_attn_heatmaps,
                    per_step_data=per_step_cross_attn,
                )

        # --- Gradient-based attribution ---
        if args.gradient:
            image_key_for_grad = args.image_key or find_image_keys(dataset)[0]

            if grad_device != device:
                print(f"\n  Moving model from {device} to {grad_device} for gradient computation...")
                policy.to(grad_device)
                model_on_gradient_device = True

            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,
                smooth_n=args.smooth_grad,
                smooth_sigma=args.smooth_grad_sigma,
                image_map=image_map,
            )
            sal_label = f"SmoothGrad (N={args.smooth_grad})" if args.smooth_grad > 1 else "Saliency"
            if saliency_maps:
                print(f"  {sal_label} maps: {len(saliency_maps)} frames")
            if gradcam_maps:
                print(f"  GradCAM maps: {len(gradcam_maps)} frames")

        # --- Extended attribution features ---
        has_extended = any([
            args.gradcam_connector,
            args.gradcam_vlm_layers,
            args.vision_vs_state,
            args.per_action_dim,
            args.language_diff,
        ])

        if has_extended:
            image_key_for_grad = args.image_key or find_image_keys(dataset)[0]

            if grad_device != device and not args.gradient:
                print(f"\n  Moving model from {device} to {grad_device} for gradient computation...")
                policy.to(grad_device)
                model_on_gradient_device = True

            enabled_features = []
            if args.gradcam_connector:
                enabled_features.append("Connector GradCAM")
            if args.gradcam_vlm_layers:
                enabled_features.append("VLM layer GradCAM")
            if args.vision_vs_state:
                enabled_features.append("Vision vs State")
            if args.per_action_dim:
                enabled_features.append("Per-action-dim")
            if args.language_diff:
                enabled_features.append("Language diff")
            print(f"\n[Step 3c] Extended attribution ({len(enabled_features)} features: "
                  f"{', '.join(enabled_features)})...")

            feat_idx = 0
            if args.gradcam_connector:
                feat_idx += 1
                print(f"\n  [{feat_idx}/{len(enabled_features)}] Connector GradCAM...")
                connector_gradcam_maps = compute_gradcam_connector_maps(
                    policy=policy,
                    dataset=dataset,
                    episode_idx=args.episode,
                    num_frames=args.num_frames,
                    image_key=image_key_for_grad,
                    device=grad_device_str,
                    noise_seed=args.gradient_seed,
                    task_override=args.task,
                    image_map=image_map,
                )
                if connector_gradcam_maps:
                    print(f"  Connector GradCAM: {len(connector_gradcam_maps)} frames")

            if args.gradcam_vlm_layers:
                feat_idx += 1
                layer_str = args.gradcam_vlm_layers
                vlm_layer_indices = [int(x.strip()) - 1 for x in layer_str.split(",")]

                try:
                    text_model = policy.model.vlm_with_expert.get_vlm_model().text_model
                    num_layers = len(text_model.layers)
                    invalid = [i + 1 for i in vlm_layer_indices if i >= num_layers]
                    if invalid:
                        print(f"  WARNING: Model has {num_layers} VLM layers — "
                              f"skipping out-of-range layers {invalid} (1-indexed)")
                        vlm_layer_indices = [i for i in vlm_layer_indices if i < num_layers]
                    if not vlm_layer_indices:
                        print("  WARNING: No valid VLM layers to compute, skipping")
                        vlm_layer_results = None
                except AttributeError:
                    pass

                if vlm_layer_indices:
                    valid_str = ",".join(str(i + 1) for i in vlm_layer_indices)
                    print(f"\n  [{feat_idx}/{len(enabled_features)}] VLM layer GradCAM (layers {valid_str})...")
                    vlm_layer_results = compute_gradcam_vlm_layers_maps(
                        policy=policy,
                        dataset=dataset,
                        episode_idx=args.episode,
                        num_frames=args.num_frames,
                        image_key=image_key_for_grad,
                        device=grad_device_str,
                        layer_indices=vlm_layer_indices,
                        noise_seed=args.gradient_seed,
                        task_override=args.task,
                        image_map=image_map,
                    )

            if args.vision_vs_state:
                feat_idx += 1
                print(f"\n  [{feat_idx}/{len(enabled_features)}] Vision vs. state attribution...")
                vision_vs_state_results = compute_vision_vs_state_maps(
                    policy=policy,
                    dataset=dataset,
                    episode_idx=args.episode,
                    num_frames=args.num_frames,
                    image_key=image_key_for_grad,
                    device=grad_device_str,
                    noise_seed=args.gradient_seed,
                    task_override=args.task,
                    image_map=image_map,
                )

            if args.per_action_dim:
                feat_idx += 1
                print(f"\n  [{feat_idx}/{len(enabled_features)}] Per-action-dim GradCAM (retain_graph — GPU recommended)...")
                try:
                    action_dim_names = list(dataset.meta.names.get("action", []))
                except (AttributeError, TypeError):
                    action_dim_names = None
                print("  Computing per-action-dim GradCAM...")
                per_action_dim_maps, per_action_dim_mags = compute_per_action_dim_maps(
                    policy=policy,
                    dataset=dataset,
                    episode_idx=args.episode,
                    num_frames=args.num_frames,
                    image_key=image_key_for_grad,
                    device=grad_device_str,
                    noise_seed=args.gradient_seed,
                    task_override=args.task,
                    action_dim_names=action_dim_names,
                    image_map=image_map,
                )

            if args.language_diff:
                feat_idx += 1
                ld = args.language_diff
                alt_task = None if ld is True or ld == "auto" else ld
                print(f"\n  [{feat_idx}/{len(enabled_features)}] Language-conditional comparison...")
                language_diff_results = compute_language_conditional_maps(
                    policy=policy,
                    dataset=dataset,
                    episode_idx=args.episode,
                    num_frames=args.num_frames,
                    image_key=image_key_for_grad,
                    device=grad_device_str,
                    noise_seed=args.gradient_seed,
                    task_override=args.task,
                    alt_task=alt_task,
                    image_map=image_map,
                )

        # --- Export gradient data ---
        if run_dir and (
            saliency_maps or gradcam_maps or connector_gradcam_maps
            or vlm_layer_results or per_action_dim_maps
            or language_diff_results or vision_vs_state_results
        ):
            save_gradient_data(
                run_dir,
                saliency_maps=saliency_maps,
                gradcam_maps=gradcam_maps,
                connector_maps=connector_gradcam_maps,
                vlm_layer_results=vlm_layer_results,
                per_action_dim_maps=per_action_dim_maps,
                per_action_dim_mags=per_action_dim_mags,
                language_diff_results=language_diff_results,
                vision_vs_state_results=vision_vs_state_results,
            )

        # --- Generate visualizations ---
        print(f"\n[Step 4] Generating visualizations...")

        grid_path = os.path.join(args.output_dir, f"episode_dashboard_ep{args.episode:03d}.png")
        create_visualization_grid(
            frames=frames,
            heatmaps=heatmaps,
            actions=actions,
            cross_attn_heatmaps=cross_attn_heatmaps,
            saliency_maps=saliency_maps,
            gradcam_maps=gradcam_maps,
            connector_gradcam_maps=connector_gradcam_maps,
            language_diff_maps=language_diff_results,
            episode_idx=args.episode,
            output_path=grid_path,
            smooth_n=args.smooth_grad,
        )

        if per_step_cross_attn and any(len(s) > 0 for s in per_step_cross_attn):
            psc_path = os.path.join(args.output_dir, f"per_step_cross_attn_ep{args.episode:03d}.png")
            create_per_step_cross_attn_grid(
                frames=frames,
                per_step_maps=per_step_cross_attn,
                episode_idx=args.episode,
                output_path=psc_path,
            )

        if vlm_layer_results and vlm_layer_indices:
            vlm_path = os.path.join(args.output_dir, f"vlm_layers_ep{args.episode:03d}.png")
            vlm_lang_tokens = None
            try:
                tokenizer = policy.model.vlm_with_expert.processor.tokenizer
                _img_key = args.image_key or find_image_keys(dataset)[0]
                first_sample = dataset[get_episode_frames(dataset, args.episode, 1, _img_key)[0][0]]
                task_str = _resolve_task_string(first_sample, dataset, task_override=args.task)
                token_ids = tokenizer.encode(task_str, add_special_tokens=False)
                vlm_lang_tokens = [tokenizer.decode([tid]) for tid in token_ids]
            except Exception:
                pass
            create_vlm_layer_grid(
                frames=frames,
                vlm_layer_results=vlm_layer_results,
                layer_indices=vlm_layer_indices,
                episode_idx=args.episode,
                output_path=vlm_path,
                lang_tokens=vlm_lang_tokens,
            )

        if per_action_dim_maps and any(m is not None for m in per_action_dim_maps):
            pad_path = os.path.join(args.output_dir, f"per_action_dim_ep{args.episode:03d}.png")
            create_per_action_dim_grid(
                frames=frames,
                per_dim_maps=per_action_dim_maps,
                action_dim_names=action_dim_names,
                episode_idx=args.episode,
                output_path=pad_path,
                magnitudes=per_action_dim_mags,
            )

        if language_diff_results and any(r is not None for r in language_diff_results):
            ld_path = os.path.join(args.output_dir, f"language_diff_ep{args.episode:03d}.png")
            create_language_diff_grid(
                frames=frames,
                lang_diff_results=language_diff_results,
                episode_idx=args.episode,
                output_path=ld_path,
            )

        if vision_vs_state_results:
            vs_path = os.path.join(args.output_dir, f"vision_vs_state_ep{args.episode:03d}.txt")
            lines = []
            v_total = 0
            n_valid = 0
            for fi, r in enumerate(vision_vs_state_results):
                if r is not None:
                    lines.append(f"Frame {fi}: vision={r['vision_norm']:.2f} "
                                 f"({r['vision_share']:.0%}), "
                                 f"state={r['state_norm']:.2f} "
                                 f"({1-r['vision_share']:.0%})")
                    v_total += r["vision_share"]
                    n_valid += 1
                else:
                    lines.append(f"Frame {fi}: failed")
            if n_valid > 0:
                avg_v = v_total / n_valid
                lines.append(f"Average: {avg_v:.0%} vision / {1 - avg_v:.0%} state")
            report = "\n".join(lines)
            with open(vs_path, "w") as f:
                f.write(report + "\n")
            print(f"\n  Vision vs. State report:\n    " + "\n    ".join(lines))
            print(f"  Saved: {vs_path}")

            vs_chart_path = os.path.join(args.output_dir, f"vision_vs_state_ep{args.episode:03d}.png")
            create_vision_vs_state_chart(
                vision_vs_state_results=vision_vs_state_results,
                episode_idx=args.episode,
                output_path=vs_chart_path,
            )

        if args.save_individual:
            save_individual_frames(
                frames=frames,
                heatmaps=heatmaps,
                output_dir=os.path.join(args.output_dir, f"episode_{args.episode:03d}"),
                episode_idx=args.episode,
            )

        if args.with_internals:
            if model_on_gradient_device:
                print(f"\n  Moving model from {grad_device} back to {device} for model internals...")
                policy.to(device)
                model_on_gradient_device = False
            print("\n[Step 5] Running model internals report...")
            internals_results = run_model_internals_report(policy, dataset, args)
            model_internals_ran = True

    if run_dir and model_internals_ran and internals_results is not None:
        save_model_internals_data(run_dir, internals_results)

    # --- Build and save manifest ---
    if run_dir:
        _manifest_task = args.task
        if _manifest_task is None:
            try:
                _img_key = args.image_key or find_image_keys(dataset)[0]
                _first = dataset[get_episode_frames(dataset, args.episode, 1, _img_key)[0][0]]
                _manifest_task = _resolve_task_string(_first, dataset)
            except Exception:
                _manifest_task = ""
        try:
            _action_names = list(dataset.meta.names.get("action", []))
        except (AttributeError, TypeError):
            _action_names = None
        _image_keys = find_image_keys(dataset) if dataset else None

        model_info = collect_model_info(policy)
        manifest_frame_count = args.internals_frames if args.internals_only else args.num_frames
        dataset_info_dict = collect_dataset_info(
            dataset, args.episode, manifest_frame_count, _manifest_task,
            action_dim_names=_action_names, image_keys=_image_keys,
        )
        available_viz = build_available_viz(
            args,
            heatmaps=heatmaps,
            cross_attn_heatmaps=cross_attn_heatmaps,
            saliency_maps=saliency_maps,
            gradcam_maps=gradcam_maps,
            connector_maps=connector_gradcam_maps,
            vlm_layer_results=vlm_layer_results,
            per_step_cross_attn=per_step_cross_attn,
            per_action_dim_maps=per_action_dim_maps,
            language_diff_results=language_diff_results,
            vision_vs_state_results=vision_vs_state_results,
            model_internals_ran=model_internals_ran,
        )
        image_paths = collect_image_paths(run_dir)
        build_manifest(run_dir, args, model_info=model_info,
                       dataset_info=dataset_info_dict,
                       available_viz=available_viz,
                       image_paths=image_paths)
        print(f"\n  Run manifest saved: {os.path.join(run_dir, 'run_manifest.json')}")

    # --- Summary ---
    output_label = run_dir if run_dir else args.output_dir
    print(f"\n{'=' * 70}")
    print("DONE!")
    print(f"{'=' * 70}")
    print(f"\nOutputs saved to: {output_label}/")
    if grid_path:
        print(f"  Episode dashboard:  {grid_path}")
    if args.save_individual and grid_path:
        print(f"  Individual frames:  {args.output_dir}/episode_{args.episode:03d}/")
    if per_step_cross_attn and any(len(s) > 0 for s in per_step_cross_attn):
        print(f"  Per-step cross-attn: per_step_cross_attn_ep{args.episode:03d}.png")
        print(f"  Centroid trajectory: per_step_cross_attn_ep{args.episode:03d}_trajectory.png")
    if vlm_layer_results and vlm_layer_indices:
        print(f"  VLM layer GradCAM:   vlm_layers_ep{args.episode:03d}.png")
        print(f"  Lang token attrib:   vlm_layers_ep{args.episode:03d}_lang_tokens.png")
    if per_action_dim_maps and any(m is not None for m in per_action_dim_maps):
        print(f"  Per-action-dim:      per_action_dim_ep{args.episode:03d}.png")
    if language_diff_results and any(r is not None for r in language_diff_results):
        print(f"  Language diff:       language_diff_ep{args.episode:03d}.png")
    if vision_vs_state_results:
        print(f"  Vision vs state:     vision_vs_state_ep{args.episode:03d}.txt")
        print(f"  Vision vs state chart: vision_vs_state_ep{args.episode:03d}.png")
    if model_internals_ran:
        print("  Model internals report: model_internals_report.md")
        print("  Model internals plot:   model_internals_report.png")
    if run_dir:
        print(f"\n  Run directory: {run_dir}")
        print(f"  Export data:   {os.path.join(run_dir, 'data')}/")

    print()
    return run_dir


def main():
    # Pre-parse --config so we can load defaults before building the full parser
    pre_parser = argparse.ArgumentParser(add_help=False)
    pre_parser.add_argument("--config", type=str, default=None,
                            help="Path to YAML config file (default: configs/defaults.yaml)")
    pre_args, _ = pre_parser.parse_known_args()
    defaults = load_defaults(pre_args.config)

    parser = argparse.ArgumentParser(
        description="Inspect SmolVLA attention, attribution, and model internals.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python inspect_attention.py
  python inspect_attention.py --episode 3 --num-frames 12 --device cuda
  python inspect_attention.py --model path/to/finetuned_checkpoint
  python inspect_attention.py --internals-only
  python inspect_attention.py --with-internals
  python inspect_attention.py --config configs/gpu.yaml
        """
    )

    parser.add_argument("--config", type=str, default=None,
                        help="Path to YAML config file (default: configs/defaults.yaml)")
    parser.add_argument("--model", type=str,
                        default=defaults.get("model", "lerobot/smolvla_base"))
    parser.add_argument("--dataset", type=str,
                        default=defaults.get("dataset", "lerobot/svla_so101_pickplace"))
    parser.add_argument("--episode", type=int,
                        default=defaults.get("episode", 0))
    add_pipeline_args(parser, defaults)

    args = parser.parse_args()

    if args.internals_only and args.with_internals:
        parser.error("--internals-only and --with-internals are mutually exclusive")

    if "--model-health" in sys.argv[1:]:
        print("  NOTE: --model-health is deprecated; use --internals-only")
    if "--health-frames" in sys.argv[1:]:
        print("  NOTE: --health-frames is deprecated; use --internals-frames")

    # Auto-enable cross-attention if per-step is requested
    if args.per_step_cross_attention and not args.cross_attention:
        args.cross_attention = True
        print("  NOTE: --per-step-cross-attention auto-enables --cross-attention")

    # Validate language-diff requires gradient gradcam
    if args.language_diff and not args.gradient:
        args.gradient = "gradcam"
        print("  NOTE: --language-diff auto-enables --gradient gradcam")

    # --- Set up run directory ---
    base_output_dir = args.output_dir
    os.makedirs(base_output_dir, exist_ok=True)

    if args.export_data:
        run_dir = create_run_dir(base_output_dir, args.run_name)
        images_dir = os.path.join(run_dir, "images")
        os.makedirs(images_dir, exist_ok=True)
        # Redirect output_dir so all PNGs go into images/ under the run folder
        args.output_dir = images_dir
    else:
        run_dir = None

    # Auto-detect device
    if args.device == "auto":
        if torch.backends.mps.is_available():
            args.device = "mps"
        elif torch.cuda.is_available():
            args.device = "cuda"
        else:
            args.device = "cpu"
    device = torch.device(args.device)

    # 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 Inspector")
    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}")
    print("  This may download ~1GB on first run...")

    try:
        from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
        # Suppress noisy warnings from HF/lerobot during model loading:
        #  - "Device 'cuda' is not available. Switching to 'mps'"
        #  - "`torch_dtype` is deprecated! Use `dtype` instead!"
        #  - "Loading ... weights ..."
        _suppressed_loggers = {
            name: logging.getLogger(name)
            for name in ("lerobot.configs.policies", "lerobot", "transformers")
        }
        _saved_levels = {name: lg.level for name, lg in _suppressed_loggers.items()}
        for lg in _suppressed_loggers.values():
            lg.setLevel(logging.ERROR)
        try:
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", message=".*torch_dtype.*deprecated.*")
                policy = SmolVLAPolicy.from_pretrained(args.model)
        finally:
            for name, lg in _suppressed_loggers.items():
                lg.setLevel(_saved_levels[name])
        policy.to(device)
        policy.eval()
        print(f"  Model loaded successfully ({sum(p.numel() for p in policy.parameters()) / 1e6:.1f}M params)")
    except ImportError:
        print("\n  ERROR: LeRobot not installed. Run:")
        print('    pip install "lerobot[smolvla]"')
        sys.exit(1)
    except Exception as e:
        print(f"\n  ERROR loading model: {e}")
        print("  Make sure the model ID is correct and you have internet access.")
        sys.exit(1)

    # --- Load dataset ---
    print(f"\n[Step 2] Loading dataset: {args.dataset}")
    print("  This may download several GB on first run...")

    try:
        from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
    except ImportError:
        from lerobot.datasets.lerobot_dataset import LeRobotDataset

    try:
        dataset = LeRobotDataset(args.dataset)
        print(f"  Dataset loaded: {len(dataset)} frames")

        # Show dataset info
        sample = dataset[0]
        print(f"  Available keys: {list(sample.keys())}")
        image_keys = find_image_keys(dataset)
        print(f"  Image keys found: {image_keys}")
    except Exception as e:
        print(f"\n  ERROR loading dataset: {e}")
        sys.exit(1)

    # --- Parse image map ---
    image_map = parse_image_map(getattr(args, "image_map", None))
    if image_map:
        print(f"\n  Image map overrides: {image_map}")

    run_inspection_pipeline(args, policy, dataset, device, grad_device, run_dir, image_map)