Explicit attention for action generation
The policy predicts spatial keypoints from RGB observations and uses them as an explicit intermediate representation for diffusion-based action generation.
End-to-end visuomotor policies provide little opportunity for humans to understand or correct the policy's visual attention. We propose GuidedAttention, a visuomotor imitation learning framework that introduces interpretable and correctable visual attention as an explicit intermediate representation. Task-relevant attention keypoints are predicted from camera images and condition a diffusion-based action policy. Users can inspect and optionally correct selected keypoints once at rollout initialization, after which the corrected attention is automatically propagated throughout execution by a tracking module. Experiments in simulation and the real world demonstrate that GuidedAttention consistently improves robot manipulation performance, particularly under positional and appearance out-of-distribution (OOD) conditions.
Demonstrations are annotated with task-relevant attention keypoints in only their initial frames. An off-the-shelf tracker propagates the annotations, and the resulting keypoints supervise an attention encoder that conditions action generation.
The policy predicts spatial keypoints from RGB observations and uses them as an explicit intermediate representation for diffusion-based action generation.
Task-relevant locations are represented as keypoints rather than remaining hidden inside an end-to-end visual representation, making the policy's focus directly inspectable.
At rollout initialization, a user may override only the unreliable keypoints. The corrected attention is then tracked automatically throughout execution, without continuous supervision.
For each demonstration, a user annotates task-relevant keypoints by clicking them only in the initial frame. Dynamic keypoints are automatically propagated through the remaining frames using an off-the-shelf tracker, while static keypoints remain fixed. This procedure provides dense keypoint supervision without requiring frame-by-frame manual annotation.
We evaluate GuidedAttention on dexterous simulation and real-world manipulation tasks under in-distribution, positional OOD, and appearance OOD conditions.
GuidedAttention improves average success rates over the action-only baselines across both real-world and simulation tasks. The gains are most pronounced under positional and appearance shifts, where correcting inaccurate attention keypoints substantially improves robustness.
The following examples illustrate the complementary benefits of explicit visual attention and one-time human correction: explicit attention improves task-relevant visual grounding, while human correction helps recover from inaccurate attention under distribution shifts.
@misc{murooka2026guidedattention,
title = {GuidedAttention: Interpretable and Correctable Visual Attention for OOD-Robust Robot Manipulation via Imitation Learning},
author = {Masaki Murooka and Ryoichi Nakajo and Keisuke Shirai and Tomohiro Motoda and Hanbit Oh and Ryo Hanai and Yukiyasu Domae},
year = {2026},
eprint = {2607.21049},
archivePrefix = {arXiv},
primaryClass = {cs.RO},
url = {https://arxiv.org/abs/2607.21049}
}