Editing the Goal: Per-Token SAE Steering and Its Depth Limits in π0.5
Rocio Perales Valdes, Animesh Garg Georgia Institute of Technology
LIBERO-Spatial, scene 7: “pick up the black bowl on the stove and place it on the plate”
π0.5-LIBERO unchanged
Picks up the bowl on the stove (10/10 rollouts)
Swap top-16 “stove” SAE features for top-16 “cabinet”
Picks up the bowl on the cabinet instead (15/20 rollouts) Swap for 16 random features: ≤ 2/20 rollouts
Figure 1: Swapping one word’s SAE features at layer 5 changes the robot’s goal in 15/20 rollouts. Both runs see scene 7 and read the same instruction, but they pick different bowls.
At layer 0, one BatchTopK feature set makes the robot complete prompt B’s task in 20/20 rollouts while instructed with prompt A.
(a) LIBERO-Spatial (of 20)
BatchTopK
Matryoshka
Scene
16, k
rand.
16, k
rand.
Layer 0, word-selected
7: stove → cabinet
,
,
9: cabinet → stove
,
,
Layer 0, contrastive
7: stove → cabinet
,
,
9: cabinet → stove
,
,
Layer 5, word-selected
7: stove → cabinet
,
,
9: cabinet → stove
,
,
Layer 5, contrastive
7: stove → cabinet
,
,
9: cabinet → stove
,
,
(b) LIBERO-Object (of 20)
BatchTopK
Matryoshka
Scene
16, k
rand.
16, k
rand.
Layer 0, word-selected
0: soup → dressing
,
,
1: cream cheese → soup
,
,
4: ketchup → BBQ sauce
,
,
6: butter → tomato sauce
,
,
Layer 0, contrastive
0: soup → dressing
,
,
1: cream cheese → soup
,
,
4: ketchup → BBQ sauce
,
,
6: butter → tomato sauce
,
,
Layer 5, word-selected
highest over scenes
,
,
Layer 5, contrastive
highest over scenes
,
,
completes prompt B’s task
Table 1: Rollouts (of 20) that complete prompt B’s task after an SAE edit at A’s goal word at layer L (scene: A’s goal → B’s). 16, k: latents per word; k: all eligible (up to 64) for word selection, 64 for contrastive. rand.: range over three draws of 16 random latents per word. Unedited policy with prompt B completes task B: 20/20 (a), 19–20/20 (b). Bold: ≥ 75%.
Single features
The exception, a BatchTopK “stove”-related feature as active on the camera rows as on the word, drops scene 7 to < 20% but leaves scene 9 at 90%; ablating it only at the camera tokens gives 6/10 and only at the word 9/10.
keep the original task’s success
SAEs are one of many viable tools
Replacing the goal word’s state with its state under prompt B transfers the task in 10/10 rollouts at layers 0 and 5 on both LIBERO-Spatial scenes. Adding a fixed difference vector, which needs no forward pass with prompt B at test time, does nearly as well (5–9/10 at 1×, 8–10/10 at 2×), while random vectors of equal norm never do (0/10).
scene 7
scene 9
layer 0
layer 5
layer 0
layer 5
Activation patching
Difference vector, 1×
Difference vector, 2×
random vectors of equal norm
completes prompt B’s task
Depth of editability
Target wordlayer L only
cameras pick up the black bowl on the stove and place it on the plate
Target wordlayers L–17
cameras pick up the black bowl on the stove and place it on the plate
Whole instructionlayers L–17
cameraspick up the black bowl on the stove and place it on the plate
Instruction + cameraslayers L–17
cameras pick up the black bowl on the stove and place it on the plate
Location word · LIBERO-Spatial, scenes 7 and 9
Object word · LIBERO-Object, 4 scenes
Figure 2: Share of rollouts completing B’s task after patching at layer L (10 per scene, pooled; 20 for dashed blue at L ≤ 6, left, and L ∈ {0, 2} in two scenes, right). Bands: scene range.