(video) Wrist cam image does not attend to target [PI05 inference, SO-101 arm]

[grahamstelzer]

Hello,

Currently running inference on a fine-tuned pi05 model on the so101 arm through lerobot (with some custom tooling/scripts).

The prompt given is "Pick up the yellow cube", this is the same in both training and inference. For the latest round of training, I took longer to approach the cube and grasp it, with the hope of giving the model more frames with the cube for its trajectory calculating. Previous training was simply moving to the cube and grasping it.

However, regardless of the approach trajectories and camera setups, I cannot seem to get the model to recognize and attend to the cube in the wrist cam. Admittedly, I am not using attention rollout for heatmap overlays of the attention matrices and am instead just averaging the last 4 layers from the denoising loop attention phase BUT regardless, shouldnt there be some heatmap activation over the cube pixels??

And furthermore (but perhaps a different question entirely), why does the model continue to perform pretty poorly at this basic task? The most annoying thing to me is that it will center the gripper on the cube and then jerk to the side.

What am missing?

VIDEO 1: attention heatmap of pi05 trained on 30k steps, dataset had 50k frames, multiple block locations, single task, wrist cam does NOT attend to cube?

pi05_30k.mp4

VIDEO 2: attention heatmap of pi05 trained on 15k steps, dataset had 80k frames, multiple block locations, single task, DESPITE SLOWER APPROACH AND MORE FRAMES OF THE CUBE IN WRIST CAM, block still not attended to??

attn_vis_20260629_121520.mp4

More specs:
prompt: "Pick up the yellow cube"
so101 arm
3090 ti gpu
cameras: 3x 640x480 30fps (different orientations did not work)
pi05 model config:

attempted 224x224 and 512x512 for the siglip image sizes (both did not work)
attempted chunk_sizes of 50 (default), 10, 1 (all did not work)
attempted n_action_steps of 50 (default), 10, 1 (all did not work)
training:
attempted steps of 3k, 6k, 15k, 20k, 26k, 30k, all performed poorly, all did not attend to the cube itself in the wrist cam
all had freeze_vision_encoder=false
all had train_expert_only=false

Thanks for your time