The Diagnostic Agent

How it figures out why your robot policy fails — and what to fix first

The Problem

Your VLA model makes wrong actions. Attention maps and GradCAM show where it looks, but not why it fails or what to change.

The Approach

Automatically analyze the scene, build a multi-signal attribution matrix, form hypotheses, and test them with counterfactual interventions.

The Output

Ranked findings with evidence chains, severity levels, and concrete, actionable fixes — not just pretty heatmaps.

Use arrow keys or the buttons below to navigate

The Pipeline at a Glance

Nine phases, from raw image to actionable diagnosis. Each phase feeds into the next.

AI = LLM decides what to do next

Phase 1Scene
Understanding
Phase 2Dataset
Diversity
Phase 3Signal
ExtractionAI
Phase 4Diagnostic
Matrix
Phase 5Representation
Probes
Phase 6LLM Hypothesis
FormationAI
Phase 7Counterfactual
VerificationAI
Phase 8Spatial vs Object
Diagnosis
Phase 9Report
SynthesisAI

Key Idea: Evidence Chain

Every final recommendation is backed by: signal datasymptom detectionhypothesiscounterfactual test result. No guessing.

Two Run Modes

Integrated: load model + dataset, run everything.
Post-hoc: analyze saved heatmaps from a previous run (faster, no GPU needed for matrix/symptoms).

What Makes It AI-Driven?

The LLM decides which signals to compute, which hypotheses to test, whether to run follow-up experiments, and how to synthesize the narrative. Different evidence leads to different diagnostic paths.

The AI Decision Tree

Each phase's output determines which branch executes next.

START  Sample frame + task string from dataset
├─ Parse task string → noun chunks → object queries
├─ OWL-ViT v2 detection on each query
│   ├─ detections ≥ threshold (0.10)? → use them
│   └─ zero detections? → retry at 0.03 (fallback)
├─ SAM segmentation on each bbox
│   ├─ SAM available? → pixel-precise masks
│   └─ SAM unavailable? → rectangular bbox masks
├─ Always-run cheap signals: attention, saliency, positional baseline, vision-vs-state
├─ DECISION LLM Signal Triage
│   ├─ Reviews cheap signal summary + symptom hints
│   └─ Selects up to 4 expensive signals from registry (different evidence → different signal choices every run)
├─ Build diagnostic matrix + run 13 symptom detectors (fixed thresholds)
├─ Model loaded?
│   ├─ YES → run semantic probe (checks if patches encode target object via cosine similarity) + QK probe (classifies attention heads as semantic vs positional)
│   └─ NO (post-hoc) → skip probes, skip counterfactuals
├─ DECISION LLM Hypothesis Formation
│   ├─ API key present? → LLM generates hypotheses + test params
│   └─ No API key? → rule-based: symptom → template hypothesis
├─ Execute counterfactuals (mandatory 4 + hypothesis-driven)
│   ├─ LOOP Surprising result? → form follow-up hypothesis → re-test
│   └─ Budget exhausted (max 7 total)? → stop iteration
├─ Spatial vs. Object diagnosis
│   ├─ Inconclusive? → run QK decomposition as tiebreaker (decomposes each attention head's Q*K matrix to measure content vs position encoding ratio)
│   └─ Clear verdict? → proceed to report
└─ DECISION LLM Report Synthesis → ranked findings + narrative

4 Decision Points

Signal triage, hypothesis formation, iterative follow-up, and report synthesis.

3 Fallback Branches

Detection confidence, SAM → bbox, LLM → rule-based. Never blocks.

1 Iteration Loop

Unconfirmed hypothesis → new hypotheses → re-test. Budget: 7 total tests.

Key Insight

Two runs on the same model can follow different diagnostic paths.

AI Under the Hood

The diagnostic system orchestrates six distinct AI/ML techniques.

Task String Parsing (NLP)

Regex-based noun chunking: tokenizes instruction, removes stop words, returns noun phrases + implicit "robot gripper" query.

"pick up the white lego block"white lego block robot gripper

Object Detection — OWL-ViT v2

Open-vocabulary ViT. Cross-modal dot products between image patches and text queries → confidence per (query, bbox). Per-class NMS keeps top-3.

No fixed label set — detects any object described in natural language.

Segmentation — SAM

Prompted with each OWL-ViT bbox. Produces pixel-level binary masks (H×W). Background = complement of all object masks.

GradCAM — Gradient Attribution

Which patches causally drive actions? Per-patch importance = attention × gradient magnitude. 1024 patches → 32×32 → upsample to 512×512 heatmap.

Attention alone is representational. GradCAM adds gradient weighting for causality.

Saliency Maps — Input Gradients

Which raw pixels matter? Enable gradients on input, backward pass, extract image.grad → L2 norm across channels → per-pixel map.

Full 512×512 resolution (finer than GradCAM's 32×32 but noisier).

SigLIP Semantic Probe

Do internal features encode the target? Cosine similarity between patch embeddings and text embeddings.

Key: margin = target_sim - best_non_target_sim. Low margin (<0.08) = object confusion.

The SmolVLA Architecture

Before we diagnose, we need to know what we're diagnosing.

Vision EncoderSigLIP
12 layers, 12 heads
1024 patches
ConnectorPixel-Shuffle
1024 → 64 tokens
93.75% compression
Language ModelVLM
16 layers, 15 heads
fuses vision + language
Action ExpertDecoder
16 layers, 8 heads
predicts robot actions

Self-Attention

Representational — shows what the encoder focuses on, but not necessarily what drives the action output. Think of it as "what the model looks at."

Cross-Attention

Closer to causal — the action expert reads from VLM outputs via cross-attention. High cross-attention on a region means the decoder is actively using it.

GradCAM / Saliency

Causal — gradient-based signals that answer: "if I change these pixels, how much does the action change?" This is the strongest evidence.

The diagnostic agent compares these three signal types against each other. When they disagree, that's informative.

Scene Understanding

The agent needs to know what's in the scene before it can measure where the model looks.

Step 1: Parse the Task

Extract object names from the language instruction using rule-based noun chunking.

"pick up the white lego block and place it in the stainless steel cup"

white lego block stainless steel cup robot gripper

Step 2: Detect Objects

OWL-ViT v2 — open-vocabulary detector. Searches the image for each parsed object query. Confidence threshold: 0.10.

Fallback: if any query gets zero detections, re-scans at 0.03 confidence to recover task-critical objects.

Step 3: Segment Objects

SAM (Segment Anything) — takes each bounding box and produces a pixel-level mask.

Fallback: if SAM is unavailable, creates rectangular masks from bounding boxes. Background = everything not covered by any object.

Why This Matters

Without scene segmentation, we can only say "the model looks at the top-left." With it, we can say "the model puts 78% of its causal attention on the background and only 3% on the target cup."

Output: SceneSegmentation

  • Per-object binary masks (H × W)
  • Background mask (complement of all objects)
  • Object labels, bounding boxes, confidence scores

This is used everywhere downstream — the matrix, the counterfactuals, and the probes all need to know which pixels belong to which semantic region.

Signal Extraction

The agent collects multiple views of model behavior. No single signal tells the whole story.

SignalWhat It MeasuresTypeCost
Self-AttentionWhere the SigLIP encoder focuses (last layer or rollout across layers)RepresentationalCheap
GradCAM (SigLIP)Which patches causally affect the action prediction, via gradients at the encoder outputCausalExpensive
Cross-AttentionWhich vision tokens the action decoder reads from during denoisingNear-causalCheap
SaliencyWhich input pixels (not patches) most affect the action, via input gradientsCausalExpensive
GradCAM (Connector)Information flow through the pixel-shuffle bottleneck (1024 → 64 tokens)CausalExpensive
Vision vs. StateFraction of gradient norm from vision pathway vs. proprioceptive stateScalarModerate
Per-Action-DimSeparate GradCAM per action dimension (x, y, z, roll, pitch, yaw, gripper)CausalExpensive
Language DiffGradCAM shift when you change the task instruction to something elseCausalExpensive

Adaptive Budget AI

Cheap signals run first. LLM reads their summary and selects up to 4 expensive signals from the registry based on what symptoms it sees.

Post-Hoc Mode

Loads pre-computed heatmaps from NPZ files. No model needed for matrix — only counterfactuals require loading the model.

The Diagnostic Matrix

The core data structure. It answers: "for each signal, what fraction of its total energy lands on each scene region?"

Attribution Mass

For each (signal, region) pair:

mass(signal, region) = Σ heatmap[region] / Σ heatmap[all]

Each signal's masses sum to ~1.0 across all regions. This normalizes for the fact that different signals have different magnitudes.

Example matrix row (GradCAM):

RegionMass
background0.78
robot gripper0.14
lego block0.03
steel cup0.05

78% of causal signal on background — that's a problem.

Scalar Metrics

Alongside the matrix, the agent computes:

  • Vision share: % of gradient from vision vs. state pathway
  • Positional baseline ratio: cosine similarity to a content-free spatial prior (>0.6 = spatial shortcut)
  • Foreground ratio: total non-background attribution
  • Cross-attention entropy: how diffuse the decoder's attention is (>5.0 = near-uniform)

Multi-Frame Averaging

The matrix is computed per-frame, then averaged. This smooths out frame-specific noise and reveals consistent patterns. Frame-to-frame variance is also tracked (used by the "unstable GradCAM" symptom detector).

Symptom Detection

13 domain-specific detectors run against the matrix. Each checks a known failure pattern.

SymptomSeverityTriggerWhat It Means
High background attributionCriticalGradCAM background >75%Actions driven by background, not task objects
Spatial shortcutCriticalPositional baseline >0.6 + low object attrMemorized positions, not recognizing objects
Low object attributionCriticalTarget object GradCAM <5%Not using the target object
Attention-GradCAM divergenceWarningRegion differs >0.3 attn vs GradCAMLooks at regions it doesn't use
Dead state pathwayWarningVision share >99.5%Proprioceptive state ignored
Language insensitivityWarningLang diff max shift <0.05Instruction change has no effect
Low dataset diversityWarningPosition std <15px, bg div <0.02Training data too uniform
Gripper fixationWarningGripper >40% causal signalOver-reliance on gripper
Cross-attention diffuseWarningCross-attn entropy >5.0Decoder reads all tokens equally
Temporal instabilityWarningCentroid smooth >0.15, tracking <0.3Attention jumps erratically
Action-attention misalignWarningTranslation dims → backgroundSpatial dims use shortcuts
Single region dependencyInfoAll dims attribute to same regionNo per-dimension specialization
Unstable GradCAMInfoGradCAM std/mean >1.0Noisy gradient attribution

Representation Probes

Deeper analysis of what the model represents at the patch level, beyond just heatmap attribution.

Semantic Probe (Patch ↔ Text)

Measures whether patch representations encode the target object's semantics. Cosine similarity between SigLIP patch embeddings and text embeddings.

Key metrics:

  • Semantic peak — max similarity to target across all patches
  • Causal alignment — similarity on GradCAM-selected patches
  • Margin — target sim minus best non-target (low = object confusion)
  • Background gap — target semantics on causal vs background patches

QK Probe (Head Classification)

Classifies each attention head as semantic, positional, or mixed. Computes QK attention map per head, then measures:

  • Semantic correlation: Pearson r with text-similarity map
  • Positional correlation: Pearson r with content-free positional baseline
  • Relocation test: Does attention follow a moved object?

Classification: advantage = semantic_corr - positional_corr + relocation_bonus

>0.08 → semantic   <-0.08 → positional   else → mixed

Positional-dominated model likely uses spatial shortcuts.

LLM Hypothesis Formation AI

The LLM doesn't just list problems — it designs experiments to verify them.

Step 1: Build the Prompt

System assembles: architecture context, task + scene, diagnostic matrix (markdown), symptom list with values, probe results, and available tests (auto-generated from counterfactual registry).

Step 2: LLM Generates JSON

Returns a JSON array of hypotheses. Robust parser handles markdown fences, trailing commas, partial failures.

Fallback: Rule-Based

No API key? Each symptom maps to a HypothesisTemplate via registry lookup. Same pipeline, no LLM reasoning.

Hypothesis Structure

Each of up to 7 hypotheses includes:

descriptionWhat might be wrong
confidence0–1, how sure the LLM is
test_typeWhich counterfactual to run
test_paramsSpecific parameters (LLM-chosen)
confirms_on_changeDirection of confirmation

Step 3: Iterative Follow-Up

Hypothesis Test Result Surprising? New Hypothesis

Budget-capped at 7 total tests. Unconfirmed hypotheses get confidence reduced by 60%.

confirms_on_change

true: Confirmed when action changes.   false: Confirmed when action doesn't change.

Counterfactual Primitives

Eight deterministic primitives the agent digitally applies to the scene. The primitives themselves always do the same thing — but which ones run and with what parameters is decided by the LLM.

1. Background Substitution

Replace background with gray/noise/blur. High delta = model uses background.

2. Object Relocation

Move target by (dx,dy). Tests object tracking vs position memorization.

3. Object Recolor

Shift hue in HSV. Low delta = shape-based, not colour-based recognition.

4. Occlusion (Targeted)

Cover target with gray/noise. Low delta = spatial shortcut (doesn't need to see object).

5. Task String Swap

Replace instruction with "do nothing". Confirmed when actions don't change (language-blind).

6. Distractor Insertion

Insert synthetic ellipse/noise. High delta = model distracted by novel objects.

7. Lighting Perturbation

Adjust brightness/contrast. Tests robustness to lighting changes.

8. Temporal Consistency

Same perturbation across frames. Tests response consistency over time.

Mandatory vs. Hypothesis-Driven Tests

Not all counterfactuals are LLM-chosen. Some always run for comparability across diagnostic runs.

Mandatory (Always Run)

These four tests run on every diagnostic, regardless of what the LLM hypothesizes. This ensures you can compare results across models and runs.

background_substitutiongray fill
object_relocationshift target object
task_string_swap"do nothing"
occlusion_targetedgray fill on target

Hypothesis-Driven AI

The LLM chooses which additional tests to run based on symptoms. The agent also runs iterative follow-ups if initial results are surprising:

  • A high-confidence hypothesis wasn't confirmed → investigate further
  • An unexpectedly large delta appears → probe with a different primitive

Total counterfactual budget is configurable (default: 3 hypothesis-driven + all mandatory).

How Confirmation Works

A counterfactual result is evaluated against the hypothesis:

significant_change = action_delta_L2 > 0.02

If confirms_on_change = true:  confirmed when change IS significant (model depends on the feature).

If confirms_on_change = false:  confirmed when change is NOT significant (model ignores the feature).

Unconfirmed hypotheses have their confidence reduced by 60% in the final report.

Spatial vs. Object Diagnosis

The critical question: is the model recognizing objects, or just memorizing where they usually are?

Four Evidence Sources

  • Positional baseline ratio — high similarity to a content-free spatial pattern suggests spatial bias
  • Semantic margin — can the model distinguish the target from background and other objects in representation space?
  • Relocation follow-through — when the object is moved, does the causal signal follow it (object grounding) or stay at the old spot (spatial shortcut)?
  • QK head classification — what fraction of attention heads encode content vs. position? (decomposes each head's Q*K matrix to measure semantic vs positional encoding ratio)

Verdict

objectModel genuinely recognizes and tracks objects
spatialModel relies on memorized positions
mixedSome evidence of both strategies
inconclusiveNot enough evidence to tell

Why This Matters

A spatial model works in training but fails when anything moves — new object positions, different camera angles, or real-world deployment.

The fix is different depending on the verdict: spatial bias needs data diversity; weak object grounding needs representation improvements.

Report Synthesis AI

The LLM synthesizes all evidence into ranked, actionable findings. It decides severity, writes the narrative, and proposes specific fixes.

Finding Structure

severitycritical warning info
titleShort name
observationCites specific numbers
test_descriptionWhat counterfactual was run
test_resultWhat happened
interpretationWhat it means for the model
fixConcrete, actionable steps
expected_impactQuantified improvement
evidence_refsLinks to symptoms & hypotheses

Evidence Chain

Every finding traces back through the full chain:

Signal Data Symptom Hypothesis Counterfactual Finding

This means you can always ask "why does the report say this?" and trace it back to specific numbers in the diagnostic matrix.

Narrative Summary

The LLM also writes a 2–4 paragraph narrative targeted at robotics engineers (not ML specialists). It answers:

  • What's wrong? — plain language description
  • Why? — root cause analysis
  • What to do? — prioritized action plan
  • What impact to expect? — predicted improvement

Visual Comparison Output

Each counterfactual generates a side-by-side comparison image showing exactly what was changed.

How Comparisons Work

  • Left panel: original image
  • Right panel: modified image
  • Dimming: unaffected pixels are darkened to 40% brightness so the changed region pops out
  • Rectangles: green (original) and red (modified) bounding boxes around the affected region
  • Zoom inset: magnified crop of the affected region (only when the region is small enough to benefit — skipped for whole-image changes like background substitution)

Affected Masks

Each primitive specifies which pixels it modified:

Background sub.background mask (everything except objects)
Object relocationold position OR new position
Object recolorobject mask
Occlusionobject mask
Distractorellipse mask at insertion point
Lightingno mask (whole image changes)
Task string swapper-dimension bar chart (no image change)

Putting It All Together

A pick-and-place model trained on 100 episodes with a plain background and a static cup.

1. Scene Understanding

Detects: robot gripper lego block steel cup

Segments each with SAM (or bbox fallback).

2. Diversity Analysis

Cup position std: 4px (static!)
Background diversity: 0.01 (single environment)

low_dataset_diversity

3. Signal Extraction

Attention: diffuse, mostly on surface
GradCAM: 78% background, 3% lego block

4. Symptoms Fired

high_background
low_object
spatial_shortcut
low_diversity

5. QK Probe

33% semantic, 25% positional, 42% mixed.

Top heads are "semantic" but spread attention across background, not the target object.

6. LLM Hypotheses

H1: "Relies on background" → test: background_substitution
H2: "Spatial shortcut for cup" → test: object_relocation

7. Counterfactual Results

Background sub: delta = 0.30 (confirmed!)
Object relocation: delta = 0.02 (confirmed — model doesn't care where object is)

8. Spatial Verdict

Verdict: spatial

The model memorized a trajectory. It doesn't visually track objects.

9. Final Finding

Critical: Background shortcut.
Fix: Add background augmentation during training. Expected: reduce background attribution from 78% to <30%.

Pluggable Extension Points

Every major component uses a registry pattern. Add a function in one file, and the rest of the pipeline discovers it automatically.

Scene Models

Config

Set scene_model in YAML. Passed as model_id to detect_objects().

Signal Types

agent.py

@register_signal("my_signal", cost_description="~1min", prompt_description="...")
LLM triage prompt auto-generates from registry.

Symptom Detectors

matrix.py

@register_symptom_detector on any function in matrix.py. Matrix, LLM, and report consume symptoms generically.

Counterfactual Primitives

counterfactual.py

@register_primitive("counterfactual.my_test", prompt_description="...", param_schema="...")
LLM hypothesis prompt auto-generates from registry.

Rule-Based Hypotheses

registry.py

register_hypothesis_template("my_symptom", test_type="...", confidence=0.7, ...)
Maps symptom type → counterfactual test automatically.

LLM Providers

registry.py

@register_llm_provider("vertex")
on any async function. Agent checks registry before built-in providers.

Mandatory CF Tests

Config

mandatory_counterfactual_tests: list in diagnostic.yaml. No code changes needed.

Report Generation

Already Generic

Findings are data-driven dataclasses. Any severity, any structure. Nothing to register.

Add a new symptom detector + its hypothesis template + its counterfactual primitive — the LLM prompt, matrix, report, and comparison images all update automatically. Three decorator calls, zero plumbing.

AGENTS.md in the repo root is the onboarding guide for AI coding assistants — it documents the architecture, key invariants, testing patterns, and extension points so that agents (Claude Code, Cursor, etc.) can contribute to the codebase with full context.

MCP Server — AI Agent Integration

48 read-only tools let any MCP-compatible client (Claude Code, Cursor, Windsurf, Continue, Zed, custom agents) query runs, diagnostics, and comparisons directly. No web viewer needed — just ask questions in natural language.

How It Works

MCP Client stdio JSON-RPC MCP Server Run Data

The MCP server reuses the same run_scanner, data_loader, stats_engine, and diagnostic_service modules as the web backend — zero code duplication.

Setup (3 steps)

  1. pip install "mcp>=1.0.0"
  2. Register server in your MCP client config
  3. Restart your client — tools appear automatically

48 Tools in 7 Categories

Discovery 8 list, search, tag, validate runs
Inspection 12 heatmaps, gradients, images
Statistics 3 viz stats, summary, formatted
Diagnostic 11 symptoms, hypotheses, findings
Comparison 5 cross-run diffs, matrix
Registry 7 primitives, signals, templates
Notes 2 get/set run notes

Example Queries (natural language)

"Compare model v2 vs v3 across bridge and aloha datasets" "What symptoms did the diagnostic find in the latest run?" "Show self-attention stats for frame 0"

Cross-Run Comparison

Compare diagnostic runs side-by-side to track training progress, validate fixes, or detect regressions. No GPU or torch required — works from pre-computed reports.

What Gets Compared

  • Attribution deltas — per-signal, per-region mass changes with % shift
  • Counterfactual deltas — action delta L2 changes per test type, confirmation status
  • Symptom tracking — resolved, new, and persistent symptoms across runs
  • Weight spectral alphas — per-component mean alpha deltas with health status

Lightweight by Design

Comparison reads report.json and weightwatcher.json from each run directory. No model loading, no GPU, no torch import.

python regenerate_report.py compare \
  run_2026-03-10/ run_2026-03-15/ \
  --labels "before fix" "after fix" \
  --output-dir ./comparison_output

Auto-Generated Outputs

  • Verdict — concise summary of what improved or regressed
  • Recommendations — actionable next steps based on delta patterns
  • JSON + Markdown — machine-readable and human-readable reports
RunSnapshot A + RunSnapshot B compare_runs() ComparisonReport Verdict + Recommendations

Each RunSnapshot loads diagnostic JSON + weight internals. Deltas are computed per-signal, per-test, and per-component. The verdict and recommendations auto-generate from delta patterns.

Key Takeaways

  • Signals alone don't explain failures. The agent combines representational (attention), near-causal (cross-attention), and causal (GradCAM/saliency) signals into one coherent matrix.
  • Symptom detection is domain-specific. 13 detectors encode robotics failure patterns: background shortcuts, spatial memorization, language blindness, state pathway death.
  • Counterfactuals provide causal evidence. Instead of correlational heatmaps, the agent actually modifies the input and measures action change.
  • Four phases are AI-driven. The LLM decides which signals to extract (Phase 3), which hypotheses and tests to run (Phases 6–7), and how to write the report (Phase 9). Other phases are deterministic for reproducibility.
  • Spatial vs. object is the key diagnostic axis. Most VLA failures come down to: does the model recognize objects, or just memorize positions? QK probe, semantic probe, and relocation counterfactual attack this from different angles.
  • Fixes are concrete. Not "collect more data" but "add background color jitter with brightness=0.3 and re-train for 50 more epochs."

Entity Relationship Diagram

The DiagnosticReport is the single serializable root. Arrows show ownership; dashed lines show cross-references.

1:1 1:1 1:1 1:N 1:N 1:N 1:N 1:1 1:N 1:N 1:1 id ref ref[] ref[] test_type generates 1:N 1:N 1:N DiagnosticReport metadata: dict scene: SceneSegmentation dataset_diversity? matrix: DiagnosticMatrix semantic_probe? qk_probe? spatial_object_diagnosis? symptoms: Symptom[] hypotheses: Hypothesis[] counterfactual_results[] findings: Finding[] llm_synthesis: str SceneSegmentation objects[], background_mask image_shape DetectedObject label, box, score mask: ndarray? DatasetDiversityReport position_stats, bg_diversity lighting_stats, task_diversity DiagnosticMatrix attribution_mass: sig×reg scalars, per_action_dim temporal_trajectories[] occlusion, connector_analysis TemporalTrajectory centroids[], smoothness tracking_corr, signal_type OcclusionMap sensitivity_map, patch_size Symptom type, severity description, evidence Hypothesis id, description, confidence test_type, test_params supporting_symptoms[] CounterfactualResult hypothesis_id, action_delta_l2 confirmed, gradcam_shift Finding severity, title, fix evidence_refs → Symptom, Hypothesis REGISTRIES PrimitiveSpec name, fn, category, cost param_schema, prompt_description SignalSpec name, cost_description, fn HypothesisTemplate symptom_type, test_type test_params_template, confidence PROBES SemanticProbeReport target_object, summary frames[], counterfactuals{} SemanticFrameSummary semantic_peak, margin causal_alignment, bg_gap QKProbeReport dominant_head_type semantic/positional fractions head_index, type, correlations SpatialObjectDiagnosis verdict, confidence spatial/object evidence[] source, score, summary LEGEND owns (1:1 / 1:N) cross-reference child / leaf entity
100%