Add fast single-counterfactual test runner

- run_test.sh / run_single_counterfactual.py: run one counterfactual
test on an existing run dir without the full diagnostic pipeline.
Reloads model, dataset, and scene from saved state.
- Supports --param overrides, --viz-only for chart regeneration, --list
- Save baseline/modified actions in task_string_swap metrics for
offline visualization regeneration

run_single_counterfactual.py

@@ -0,0 +1,426 @@
+#!/usr/bin/env python3
+"""Run a single counterfactual test on an existing run directory.
+
+Loads model, dataset, and scene data from a previous run, then executes
+one counterfactual test and saves the result + comparison image.
+Much faster than re-running the full diagnostic pipeline.
+
+Usage:
+    python run_single_counterfactual.py <run_dir> <test_name> [--param key=value ...]
+    python run_single_counterfactual.py --list
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run a single counterfactual test on an existing run directory",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  %(prog)s outputs/mthirumalai/finetuned_model background_substitution
+  %(prog)s outputs/mthirumalai/finetuned_model distractor_insertion --param position='[200,200]'
+  %(prog)s outputs/mthirumalai/finetuned_model task_string_swap --viz-only
+  %(prog)s --list
+""",
+    )
+    parser.add_argument("run_dir", nargs="?", help="Path to existing run directory")
+    parser.add_argument("test_name", nargs="?", help="Counterfactual test name")
+    parser.add_argument("--list", action="store_true", help="List available tests")
+    parser.add_argument("--param", action="append", default=[],
+                        help="Test parameter as key=value (repeatable)")
+    parser.add_argument("--viz-only", action="store_true",
+                        help="Only regenerate visualization from existing result.json")
+    parser.add_argument("--device", default="auto", help="Device (cuda/cpu/auto)")
+    parser.add_argument("--episode", type=int, default=None,
+                        help="Episode index (default: from manifest)")
+
+    args = parser.parse_args()
+
+    # Import here so --list/--help are fast
+    _ensure_imports()
+
+    if args.list:
+        _list_tests()
+        return
+
+    if not args.run_dir or not args.test_name:
+        parser.error("run_dir and test_name are required (or use --list)")
+
+    run_dir = args.run_dir
+    test_name = args.test_name
+
+    # Validate run directory
+    if not os.path.isdir(run_dir):
+        print(f"ERROR: Run directory not found: {run_dir}")
+        sys.exit(1)
+
+    manifest_path = os.path.join(run_dir, "run_manifest.json")
+    if not os.path.exists(manifest_path):
+        print(f"ERROR: No run_manifest.json in {run_dir}")
+        sys.exit(1)
+
+    # Parse test params
+    test_params = _parse_params(args.param)
+
+    if args.viz_only:
+        _regenerate_viz(run_dir, test_name)
+        return
+
+    # Load manifest for model/dataset info
+    with open(manifest_path) as f:
+        manifest = json.load(f)
+
+    cli_args = manifest.get("cli_args", {})
+    model_id = cli_args.get("model")
+    dataset_id = manifest.get("dataset_info", {}).get("dataset_id") or cli_args.get("dataset")
+    image_key = cli_args.get("image_key")
+    image_map_str = cli_args.get("image_map")
+    episode_idx = args.episode if args.episode is not None else cli_args.get("episode", 0)
+
+    if not model_id:
+        print("ERROR: Cannot determine model from manifest. Specify --model?")
+        sys.exit(1)
+    if not dataset_id:
+        print("ERROR: Cannot determine dataset from manifest.")
+        sys.exit(1)
+
+    # Device
+    import torch
+    device = args.device
+    if device == "auto":
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            device = "mps"
+        else:
+            device = "cpu"
+
+    print(f"\n{'=' * 50}")
+    print(f"  Quick Counterfactual Test")
+    print(f"{'=' * 50}")
+    print(f"  Run dir:  {run_dir}")
+    print(f"  Test:     {test_name}")
+    print(f"  Model:    {model_id}")
+    print(f"  Dataset:  {dataset_id}")
+    print(f"  Device:   {device}")
+    if test_params:
+        print(f"  Params:   {test_params}")
+    print(f"{'=' * 50}\n")
+
+    # ── Load model ──
+    t0 = time.time()
+    print("  Loading model...", end="", flush=True)
+    from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
+    policy = SmolVLAPolicy.from_pretrained(model_id)
+    policy.to(device)
+    policy.eval()
+    print(f" done ({time.time() - t0:.1f}s)")
+
+    # ── Load dataset ──
+    t0 = time.time()
+    print("  Loading dataset...", end="", flush=True)
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
+    dataset = LeRobotDataset(dataset_id)
+    print(f" done ({time.time() - t0:.1f}s)")
+
+    # Resolve image key
+    if image_key is None:
+        from smolvla_inspect.data import find_image_keys
+        image_keys = find_image_keys(dataset)
+        image_key = image_keys[0] if image_keys else "observation.images.top"
+
+    # Parse image map
+    image_map = None
+    if image_map_str:
+        from smolvla_inspect.data import parse_image_map
+        image_map = parse_image_map(image_map_str)
+
+    # ── Load scene data ──
+    print("  Loading scene data...", end="", flush=True)
+    scene = _load_scene(run_dir)
+    if scene is None:
+        print("\n  WARNING: No scene data found — running scene detection...")
+        scene = _detect_scene(dataset, episode_idx, image_key, device)
+    else:
+        print(" done")
+
+    # ── Get sample ──
+    print("  Loading sample...", end="", flush=True)
+    first_frame_idx = _get_first_frame_idx(dataset, episode_idx)
+    sample = dataset[first_frame_idx]
+    print(f" done (frame {first_frame_idx})")
+
+    # ── Validate test name ──
+    from smolvla_inspect.diagnostic.registry import REGISTRY
+    primitive_name = f"counterfactual.{test_name}"
+    if primitive_name not in REGISTRY:
+        print(f"\n  ERROR: Unknown test '{test_name}'")
+        print(f"  Available: {', '.join(n.removeprefix('counterfactual.') for n in REGISTRY if n.startswith('counterfactual.'))}")
+        sys.exit(1)
+
+    # ── Apply default params if not provided ──
+    test_params = _apply_defaults(test_name, test_params, scene)
+
+    # ── Run the test ──
+    spec = REGISTRY[primitive_name]
+    params = dict(test_params)
+    params.update({
+        "policy": policy,
+        "sample": sample,
+        "dataset": dataset,
+        "image_key": image_key,
+        "device": device,
+        "image_map": image_map,
+    })
+    if "segmentation" in spec.fn.__code__.co_varnames:
+        params["segmentation"] = scene
+    if "episode_idx" in spec.fn.__code__.co_varnames:
+        params["episode_idx"] = episode_idx
+
+    print(f"\n  Running {test_name}...", flush=True)
+    t0 = time.time()
+    result = spec.fn(**params)
+    elapsed = time.time() - t0
+
+    # ── Save result ──
+    import numpy as np
+
+    cf_dir = os.path.join(run_dir, "diagnostic", "counterfactuals", test_name)
+    os.makedirs(cf_dir, exist_ok=True)
+
+    # Save result.json
+    result_dict = {
+        "hypothesis_id": result.hypothesis_id,
+        "test_type": result.test_type,
+        "action_delta_l2": float(result.action_delta_l2),
+        "action_delta_per_dim": [float(x) for x in result.action_delta_per_dim],
+        "gradcam_shift": float(result.gradcam_shift) if result.gradcam_shift else 0.0,
+        "attribution_shift_per_region": result.attribution_shift_per_region or {},
+        "confirmed": result.confirmed,
+        "metrics": result.metrics or {},
+    }
+    result_path = os.path.join(cf_dir, "result.json")
+    with open(result_path, "w") as f:
+        json.dump(result_dict, f, indent=2)
+
+    # Save comparison image
+    if result.visual_comparison is not None:
+        from PIL import Image
+        comp_path = os.path.join(cf_dir, "comparison.png")
+        Image.fromarray(result.visual_comparison).save(comp_path)
+        print(f"  Saved: {comp_path}")
+
+    print(f"\n  Result ({elapsed:.1f}s):")
+    print(f"    Action delta (L2): {result.action_delta_l2:.4f}")
+    print(f"    Confirmed: {result.confirmed}")
+    print(f"    Saved to: {result_path}")
+    print(f"{'=' * 50}\n")
+
+
+def _ensure_imports():
+    """Check that the package is importable."""
+    try:
+        import smolvla_inspect  # noqa: F401
+    except ImportError:
+        # Try adding the project root to sys.path
+        root = os.path.dirname(os.path.abspath(__file__))
+        sys.path.insert(0, root)
+
+
+def _list_tests():
+    """Print available counterfactual tests."""
+    _ensure_imports()
+    # Force registry population by importing the counterfactual module
+    import smolvla_inspect.diagnostic.counterfactual  # noqa: F401
+    from smolvla_inspect.diagnostic.registry import list_primitives
+
+    print("\nAvailable counterfactual tests:\n")
+    for spec in list_primitives(category="counterfactual"):
+        name = spec.name.removeprefix("counterfactual.")
+        print(f"  {name}")
+        print(f"    {spec.description}")
+        if spec.param_schema:
+            print(f"    Params: {spec.param_schema}")
+        print()
+
+
+def _parse_params(param_list: list[str]) -> dict:
+    """Parse --param key=value arguments into a dict."""
+    params = {}
+    for p in param_list:
+        if "=" not in p:
+            print(f"ERROR: Invalid param '{p}' — expected key=value")
+            sys.exit(1)
+        key, val = p.split("=", 1)
+        # Try JSON parsing for lists, numbers, bools
+        try:
+            params[key] = json.loads(val)
+        except (json.JSONDecodeError, ValueError):
+            params[key] = val
+    return params
+
+
+def _load_scene(run_dir: str):
+    """Reconstruct SceneSegmentation from saved files."""
+    import numpy as np
+    from smolvla_inspect.diagnostic.models import SceneSegmentation, DetectedObject
+
+    scene_dir = os.path.join(run_dir, "diagnostic", "scene")
+    det_path = os.path.join(scene_dir, "detections.json")
+    seg_path = os.path.join(scene_dir, "segmentation.npz")
+
+    if not os.path.exists(det_path):
+        return None
+
+    with open(det_path) as f:
+        det_data = json.load(f)
+
+    seg_data = {}
+    if os.path.exists(seg_path):
+        seg_data = dict(np.load(seg_path))
+
+    objects = []
+    for det in det_data["objects"]:
+        mask_key = det["label"].replace(" ", "_")
+        mask = seg_data.get(mask_key)
+        if mask is not None:
+            mask = mask.astype(bool)
+        objects.append(DetectedObject(
+            label=det["label"],
+            box=tuple(det["box"]),
+            score=det["score"],
+            mask=mask,
+        ))
+
+    bg_mask = seg_data.get("background")
+    if bg_mask is not None:
+        bg_mask = bg_mask.astype(bool)
+
+    h, w = det_data["image_shape"]
+    return SceneSegmentation(
+        objects=objects,
+        background_mask=bg_mask,
+        image_shape=(h, w),
+    )
+
+
+def _get_first_frame_idx(dataset, episode_idx: int) -> int:
+    """Get the dataset index of the first frame in an episode."""
+    try:
+        return dataset.meta.episodes["dataset_from_index"][episode_idx]
+    except (AttributeError, KeyError):
+        try:
+            return dataset.episode_data_index["from"][episode_idx].item()
+        except (AttributeError, KeyError):
+            return episode_idx * 200
+
+
+def _detect_scene(dataset, episode_idx: int, image_key: str, device: str):
+    """Run scene detection from scratch (fallback when no saved scene)."""
+    from smolvla_inspect.diagnostic.scene import detect_scene
+
+    first_idx = _get_first_frame_idx(dataset, episode_idx)
+    sample = dataset[first_idx]
+    return detect_scene(sample, image_key, device)
+
+
+def _apply_defaults(test_name: str, params: dict, scene) -> dict:
+    """Fill in sensible defaults for test params that weren't provided."""
+    if test_name == "background_substitution":
+        params.setdefault("replacement", "gray")
+    elif test_name == "object_relocation":
+        if "target_object" not in params and scene:
+            params["target_object"] = _pick_target(scene)
+        params.setdefault("shift_pixels", [100, -80])
+    elif test_name == "object_recolor":
+        if "target_object" not in params and scene:
+            params["target_object"] = _pick_target(scene)
+        params.setdefault("hue_shift", 0.5)
+    elif test_name == "occlusion_targeted":
+        if "target_object" not in params and scene:
+            params["target_object"] = _pick_target(scene)
+        params.setdefault("fill", "gray")
+    elif test_name == "distractor_insertion":
+        params.setdefault("position", [256, 256])
+        params.setdefault("distractor_size", 80)
+    elif test_name == "task_string_swap":
+        params.setdefault("replacement_task", "do nothing")
+    elif test_name == "lighting_perturbation":
+        params.setdefault("brightness_delta", 0.3)
+        params.setdefault("contrast_delta", 0.3)
+    elif test_name == "temporal_consistency":
+        params.setdefault("perturbation_type", "background_substitution")
+        params.setdefault("num_frames", 5)
+    return params
+
+
+def _pick_target(scene) -> str:
+    """Pick the most likely manipulation target from scene objects."""
+    skip = {"robot gripper", "robot arm", "gripper", "arm"}
+    for obj in scene.objects:
+        if obj.label.lower() not in skip and obj.mask is not None:
+            return obj.label
+    # Fallback to first object with a mask
+    for obj in scene.objects:
+        if obj.mask is not None:
+            return obj.label
+    return scene.objects[0].label if scene.objects else "object"
+
+
+def _regenerate_viz(run_dir: str, test_name: str):
+    """Regenerate only the visualization from an existing result.json."""
+    import numpy as np
+
+    cf_dir = os.path.join(run_dir, "diagnostic", "counterfactuals", test_name)
+    result_path = os.path.join(cf_dir, "result.json")
+
+    if not os.path.exists(result_path):
+        print(f"ERROR: No result.json at {result_path}")
+        print(f"  Run the test first (without --viz-only)")
+        sys.exit(1)
+
+    with open(result_path) as f:
+        result_data = json.load(f)
+
+    # For task_string_swap, regenerate the action delta chart
+    if test_name == "task_string_swap":
+        from smolvla_inspect.diagnostic.counterfactual import _make_action_delta_chart
+        metrics = result_data.get("metrics", {})
+        baseline = metrics.get("baseline_actions")
+        modified = metrics.get("modified_actions")
+
+        if baseline is None or modified is None:
+            # Reconstruct from deltas (approximate — modified = baseline + delta)
+            # but we don't have absolute values, so the grouped-bar top panel
+            # won't render. Re-run the test without --viz-only instead.
+            print("  WARNING: result.json does not contain baseline/modified actions.")
+            print("  Re-run the test without --viz-only to get the two-panel chart.")
+            print("  (Older results lack this data; only delta bars will be shown.)")
+            return
+
+        baseline = np.array(baseline)
+        modified = np.array(modified)
+        original_task = metrics.get("original_task", "original task")
+        replacement_task = metrics.get("replacement_task", "replacement task")
+
+        chart = _make_action_delta_chart(baseline, modified, original_task, replacement_task)
+        from PIL import Image
+        comp_path = os.path.join(cf_dir, "comparison.png")
+        Image.fromarray(chart).save(comp_path)
+        print(f"  Regenerated chart: {comp_path}")
+    else:
+        print(f"  --viz-only currently supports: task_string_swap")
+        print(f"  For image-based tests, re-run the test (model needed for comparison).")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

run_test.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# Fast test runner — re-run individual counterfactual tests or regenerate
+# visualizations on an existing run directory WITHOUT the full pipeline.
+#
+# Usage:
+#   ./run_test.sh <run_dir> <test_name> [--param key=value ...] [--viz-only]
+#
+# Examples:
+#   # Re-run background_substitution on the finetuned model
+#   ./run_test.sh outputs/mthirumalai/finetuned_model background_substitution
+#
+#   # Re-run distractor with custom params
+#   ./run_test.sh outputs/mthirumalai/finetuned_model distractor_insertion \
+#       --param position='[200,200]' --param distractor_size=60
+#
+#   # Re-run task_string_swap with a different replacement task
+#   ./run_test.sh outputs/mthirumalai/finetuned_model task_string_swap \
+#       --param replacement_task='pick up the red cube'
+#
+#   # Just regenerate the visualization from existing result (no model needed)
+#   ./run_test.sh outputs/mthirumalai/finetuned_model task_string_swap --viz-only
+#
+#   # List available counterfactual tests
+#   ./run_test.sh --list
+
+set -e
+cd "$(dirname "$0")"
+
+# macOS: Homebrew ffmpeg@6 for TorchCodec compatibility
+FFMPEG6_LIB="/opt/homebrew/opt/ffmpeg@6/lib"
+if [[ -d "$FFMPEG6_LIB" && -f "$FFMPEG6_LIB/libavutil.58.dylib" ]]; then
+  export DYLD_LIBRARY_PATH="${FFMPEG6_LIB}${DYLD_LIBRARY_PATH:+:$DYLD_LIBRARY_PATH}"
+fi
+
+# Linux: Ensure system FFmpeg 4.x libs load first
+SYS_FFMPEG="/lib/x86_64-linux-gnu"
+if [[ -f "$SYS_FFMPEG/libavutil.so.56" ]]; then
+  export LD_LIBRARY_PATH="${SYS_FFMPEG}${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
+fi
+
+exec python3 run_single_counterfactual.py "$@"

smolvla_inspect/diagnostic/counterfactual.py

@@ -1122,6 +1122,12 @@ def task_string_swap(
         attribution_shift_per_region={},
         confirmed=delta_l2 > 0.01,
         visual_comparison=comparison,
+        metrics={
+            "original_task": original_task,
+            "replacement_task": replacement_task,
+            "baseline_actions": baseline_actions.tolist(),
+            "modified_actions": modified_actions.tolist(),
+        },
     )