Peeking Inside a Robot Policy

I’m working on tools that save useful signals from vision-language-action models while they control a robot. This page shows one of them: π0‑FAST.

Interpretability Robot Learning JAX π0‑FAST

The short version: the model normally returns an action and throws away the reasoning signals that produced it. My recorder keeps the useful ones so I can study them later.

How one action happens

π0‑FAST takes camera images, a written instruction, and the robot’s current state. It turns them into one long sequence, reasons over that sequence, then writes the next movement as action tokens. The hooks are small observation points along that path.

what the robot gets camera views + instruction + state
“pick up the red mug”
one shared sequence every input gets a position 816 positions
the model information mixes across 18 layers
what comes out one action token at a time then robot motion
observation_input

The exact camera frames, instruction, and robot state that reached the policy after preprocessing.

Click the questions to see which recorded hook helps answer each one.

The basic idea

01

Tap into the forward pass

Each hook sits at a useful point: before the model, inside its layers, or during action decoding.

02

Record without interfering

Extra samples use a separate random key, so turning on the recorder does not change the trajectory the robot executes.

03

Analyze after the rollout

Every inference call becomes one compressed file containing the observation, action, and whichever hooks were enabled.


One detail that matters

A lot of the recorded arrays have an axis of length 816. Those positions are not anonymous: they map back to three camera views, then the instruction and robot state. I record those boundaries too, which is what lets an attention signal become a heatmap over a real image.

base camera
left wrist
empty right slot
text + state

256 + 256 + 256 + 48 = 816 positions

Each camera contributes a 16 × 16 grid of patches. The third camera slot is empty in LIBERO, but π0‑FAST still reserves and attends to those positions.

See all nine hooks and tensor shapes
HookShapeWhat I use it for
observation_input3 × [B,224,224,3]The exact model-ready input
token_spanssmall dictionaryMapping positions back to each camera, text, or state
prefix_embeddings[B,816,2048]How every input is represented before reasoning
prefix_final_hidden_state[B,816,2048]What the model made of the scene
prefix_gradients[B,816,2048]Which inputs could most change the first action
raw_attention_weights[18,B,8,816]Where the first action token looked
value_vectors[B,18,816,1,256]What information was available at those locations
action_chunks[chunks,B,256]Other token rollouts the model considered
insight_metrics6 × [B,256]Confidence and uncertainty during decoding

Why I’m building this

When a robot fails, the final movement only tells part of the story. These recordings make it possible to ask whether the policy missed the object, misunderstood the instruction, focused on the wrong camera region, or simply became uncertain while decoding the action.

The larger goal is to turn those internal signals into useful ways to analyze—and eventually predict—VLA failures.

Last Updated: September 20, 2026