Setup & Tutorial

Get started with VLA interpretability experiments and Action Atlas

1

Clone the Repository

Clone with submodules to include LIBERO benchmark and LeRobot dependencies.

git clone --recursive https://github.com/CWRU-AISM/vla_interp.git cd vla_interp
2

Create Environment

Set up a conda environment with Python 3.12 and install PyTorch with CUDA support.

conda create -y -n pi05 python=3.12 conda activate pi05 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
3

Install Dependencies

Install LeRobot (for Pi0.5 support) and the project package.

cd lerobot && pip install -e ".[pi]" && cd .. pip install -e .
4

Download Model Checkpoint

Download the Pi0.5 fine-tuned checkpoint from HuggingFace.

huggingface-cli download lerobot/pi05_libero_finetuned \ --local-dir checkpoints/pi05_libero_finetuned
5

Verify Installation

Run a quick check to make sure everything is configured correctly.

python -c "from lerobot.policies.pi05.modeling_pi05 import PI05Policy; print('OK')"
✅
Requirements: NVIDIA GPU with 24GB+ VRAM, CUDA 12.x, Python 3.12, HuggingFace account with PaliGemma access.

Pi0.5 Experiment Pipeline

Run the full mechanistic interpretability pipeline on Pi0.5 (3B, dual-pathway, flow matching).

1

Train Per-Token SAEs

Train TopK SAEs (k=64, 8x expansion) on action expert residual stream activations.

# Train SAE on layer 17 of the action expert (1024 -> 8192) python experiments/train_pertoken_sae.py \ --suite goal --layer action_expert_layer_17 --expansion 8
⚠️
Critical: Always use per-token SAEs. Mean-pooled SAEs cause 88% task failure despite high reconstruction quality. Always set TORCH_COMPILE_DISABLE=1 for experiment scripts.
2

Contrastive Concept Identification

Identify task-specific features using Cohen's d effect size scoring.

python experiments/pi05_contrastive_concept_id.py \ --suite goal --pathway expert --layers 0-17
3

Concept Ablation & Steering

Zero-ablate or steer specific SAE features during live rollouts to establish causality.

# Zero-ablation python experiments/run_pi05_concept_ablation.py \ --suite goal --layer 8 --concept butter --multiplier 0 # Feature steering (suppress at -3x) python experiments/run_pi05_concept_ablation.py \ --suite goal --layer 8 --concept butter --multiplier -3
4

Linear Probing & Activation Injection

Validate causality with probes and test activation transfer across conditions.

# Multi-layer probing python experiments/pi05_action_probes_rss.py \ --suite goal --layers 0-17 # Pathway comparison (Expert vs PaliGemma) python experiments/pathway_comparison.py \ --suite goal --layers 0-17

OpenVLA-OFT Setup

Additional setup for OpenVLA-OFT (7B, LLaMA-2 backbone, L1 regression action head).

1

Install OpenVLA Dependencies

Install the prismatic-vlms and openvla packages alongside the base environment.

# See setup/SETUP.md for full instructions pip install -e openvla/ pip install -e prismatic_stub/
2

Train OFT SAEs

Train SAEs on all 32 LLaMA-2 layers. Note: apply SAEs to action tokens only (last 7 of 605 total tokens).

# Per-token SAE (4096 -> 32768, TopK k=64) python experiments/openvla/train_sae.py \ --suite goal --layer 16 --expansion 4
ℹ️
Key insight: SAEs on ALL 605 tokens = 0% success. SAEs on action tokens only (last 7) = 99.2% success. The SAE disrupts vision/language tokens that the model depends on.
3

OFT Concept Ablation

Run ablation experiments across all 4 LIBERO suites.

python experiments/openvla/run_oft_concept_ablation.py \ --suite goal --concept put --multiplier 0

Running Action Atlas Locally

Action Atlas is deployed at action-atlas.com, but you can also run it locally with full data access.

1

Start the Backend

The Flask backend serves API endpoints and data from local storage.

python vla_conceptviz/run_backend.py --host 0.0.0.0 --port 6006
2

Start the Frontend

The Next.js frontend provides the interactive visualization interface.

cd vla_conceptviz/frontend npm install # For local dev, set API URL: echo "NEXT_PUBLIC_API_URL=http://localhost:6006" > .env.local npm run dev

Visit http://localhost:3002 to access Action Atlas locally.

3

Or Use Docker

Run both backend and frontend in a single container.

docker compose build && docker compose up -d # Visit http://localhost:3000
💡
Tip: Docker mounts local video directories for full data access. The Fly.io deployment uses Tigris S3 for video storage (48,537+ files).

Exploring the Platform

Feature Explorer

Browse UMAP scatter plots of SAE features, search by semantic query, and inspect individual feature activations across layers and suites.

Video Library

Filter through 49,000+ rollout videos by model (Pi0.5 / OFT), experiment type, suite, and success/failure outcome.

Ablation Studies

Compare baseline vs. ablated behavior side-by-side, view success rate deltas per concept, and explore layer-phase ablation matrices.

Perturbation Testing

Test model robustness to vision perturbations (flip, blur, noise, crop) with real experiment data from 5,691 episodes across Pi0.5 and OFT.