GuidedAttention: Interpretable and Correctable Visual Attention for OOD-Robust Robot Manipulation via Imitation Learning

GuidedAttention: Interpretable and Correctable Visual Attention for OOD-Robust Robot Manipulation via Imitation Learning

Paper: GuidedAttention: Interpretable and Correctable Visual Attention for OOD-Robust Robot Manipulation via Imitation Learning
Authors: Masaki Murooka, Ryoichi Nakajo, Keisuke Shirai, et al. (AIST / CNRS-AIST JRL)
Link: arXiv:2607.21049　|　Project: mmurooka.github.io/guided-attention-project-page　|　Framework: RoboManipBaselines

One-line summary: An interpretable, correctable attention keypoint layer is inserted between perception and action in a diffusion policy—clicking once at rollout start fixes the robot's visual drift under OOD conditions, boosting success rates by 20-30 percentage points.

Background and Motivation

End-to-end visuomotor policies generate actions directly from high-dimensional sensory inputs, eliminating manual feature engineering, but provide virtually no opportunity for humans to understand or correct the policy's visual attention during execution. Even when a human can easily identify the task-relevant object or location, the policy must discover these visual cues on its own. When the policy fails for lack of a simple but crucial hint, the human can only watch—this is especially acute in out-of-distribution (OOD) environments.

Existing visual prompting approaches have limitations at different levels. RT-Sketch requires drawing a goal-scene sketch for every manipulation command, a non-negligible user burden. Robotic Visual Instruction (RoVI) lets users write semantic instructions on images, but the visual prompt represents the task intent itself rather than auxiliary attention. Directly overlaying marker boxes on input images (Marker Overlay Prompting) provides similar spatial cues but lacks explicit learning of the marker-task structure association, performing poorly under OOD. The common problem: visual prompts are either persistent inputs needed every execution, or cannot deeply couple with the policy's internal attention mechanism.

GuidedAttention's core insight: make visual attention an explicit intermediate representation between perception and action, enabling humans to inspect and correct it when needed, without requiring continuous supervision. Attention takes the form of task-relevant keypoints—annotated only in the first frame of each demonstration during training, with remaining frames auto-propagated by a tracking module. At rollout, the policy autonomously predicts keypoints and actions; when autonomous prediction becomes unreliable under OOD, the user can correct selected keypoints once at rollout initialization, after which the correction is auto-tracked throughout execution.

Method

1. Overall Policy Architecture (DP-GA)

DP-GA (Diffusion Policy with GuidedAttention) takes multi-view RGB observations and proprioceptive joint states as input. For each camera image, a visual attention encoder (AttnEnc) predicts task-relevant attention keypoints and extracts corresponding attention features. These features are concatenated with proprioceptive features to condition a CNN-based denoising diffusion policy built on a U-Net backbone. Unlike conventional diffusion policies that condition action generation directly on CNN global visual features, DP-GA introduces an interpretable, correctable attention representation between perception and action.

Figure 2: Policy architecture overview. (A) Denoising diffusion conditioned on attention features and proprioceptive states. (B) ResNet+Transformer encoder-decoder predicts pixel-level keypoints. (C) Users can optionally correct keypoints during rollout.

2. Visual Attention Encoder

AttnEnc adopts a DETR-style architecture with a Transformer encoder and decoder, but instead of predicting bounding boxes, it predicts 2D coordinates of attention keypoints. RGB images are processed by a ResNet-18 backbone for patch features, augmented with positional embeddings, and fed to the Transformer encoder. The decoder receives learnable queries (corresponding to a pre-defined number of keypoints) and produces attention features. A coordinate head implemented as a linear layer maps each attention feature to a 2D coordinate:

$$p_{i}=h_{\theta}(f_{i})=A\,f_{i}+b$$

where $p_{i}\in\mathbb{R}^{2}$ is the coordinate of the $i$-th keypoint and $f_{i}\in\mathbb{R}^{d_f}$ is its attention feature. Each image outputs $N_k$ attention features $\{f_{i}\}_{i=1}^{N_k}$. Training optimizes both action prediction and keypoint prediction. The keypoint loss is the MSE between predicted and annotated coordinates:

$$\mathcal{L}_{kp}=\frac{1}{N_{k}}\sum_{i=1}^{N_{k}}\lVert p_{i}-p_{i}^{gt}\rVert_{2}^{2}$$

By conditioning the diffusion process through keypoint supervision, the policy learns to focus on spatially relevant locations essential for manipulation.

3. Keypoint Override Mechanism

When predicted attention is unreliable under OOD, the user can correct a subset of keypoint coordinates. Corrected coordinates are converted to attention features via the inverse of $h_{\theta}(\cdot)$ and replace the corresponding predicted features in the final conditioning set:

$$z_{i}=\begin{cases}f_{i} & \text{(prediction path)}\\[4pt]h_{\theta}^{-1}(p_{i}^{gt}) & \text{(override path)}\end{cases}$$

Three techniques ensure consistency and reliability of the override mechanism:

Forward-inverse weight sharing: Since $h_{\theta}$ is a linear mapping from $d_f$ dimensions to 2D, its analytical inverse requires no additional parameters:

$$h_{\theta}^{-1}(p_{i})=A^{+}(p_{i}-b)$$

where $A^{+}$ is the Moore-Penrose pseudoinverse. Weight sharing encourages mutual consistency between forward prediction and inverse recovery.

Feature-space alignment loss: When $d_f>2$, the inverse mapping is underdetermined—multiple feature vectors may correspond to the same coordinate. To align forward-predicted and inverse-restored features, a consistency loss is introduced:

$$\mathcal{L}_{feat}=\frac{1}{N_{k}}\sum_{i=1}^{N_{k}}\left\lVert f_{i}-h_{\theta}^{-1}(p_{i}^{gt})\right\rVert_{2}^{2}$$

Randomized feature routing: During training, each keypoint feature $z_i$ is randomly selected from either the prediction path or the override path, preventing the diffusion model from depending on a single pathway and stabilizing early training when predicted keypoints are still inaccurate.

graph LR
  A[RGB Image] --> B[ResNet-18 Backbone]
  B --> C[Transformer Enc-Dec]
  C --> D[Attention Feature f_i]
  C --> E[Coordinate Head h_theta]
  E --> F[Keypoint Coord p_i]
  D --> G{Random Routing}
  H[User Correction] --> I[Inverse Map h_theta^-1]
  I --> G
  G --> J[Conditioning Feature z_i]
  J --> K[Diffusion Policy U-Net]
  K --> L[Robot Action]
  style G fill:#f5a623,stroke:#b97316,color:#fff
  style H fill:#e74c3c,stroke:#a93226,color:#fff
  style K fill:#4a90d9,stroke:#2c5f8a,color:#fff

4. Keypoint Annotation and Tracking

During training, a human clicks keypoint locations on the first frame of each demonstration; subsequent ground-truth keypoints are obtained via tracking:

$$p_{i}^{gt}[t]=G(p_{i}^{gt}[t-1],\,I[t],\,I[t-1])$$

where $G(\cdot)$ is the Co-Tracker tracking module and $I[t]$ is the RGB image at time $t$. Known static keypoints (e.g., a pole base on the table) remain fixed throughout. This design keeps annotation cost minimal—only one frame per demonstration.

Experimental Results

Simulation Experiments

Evaluation in MuJoCo uses a UR5e arm on three tasks: Cable (threading a flexible cable between poles), Ring (placing a ring onto a pole), and Particle (scooping particles into a box). Each task collects 30 teleoperated demonstrations, evaluated under four conditions: ID, Positional OOD, Appearance OOD, and Position+Appearance OOD.

Figure 5: Simulation manipulation tasks. Each uses two attention keypoints: a dynamic keypoint (on the manipulated object) and a static keypoint (on the target or environment).

Policy	Correction	ID	Pos-OOD	App-OOD	Pos+App-OOD
DP-GA (default)	—	83.3%	55.6%	45.6%	34.4%
DP-GA (default)	Enabled	84.4%	70.0%	76.7%	67.8%
DP (diffusion baseline)	—	67.8%	45.6%	54.4%	28.9%
ACT	—	53.3%	28.9%	6.7%	8.9%

Fully autonomous, DP-GA improves over DP by ~15pp under ID and ~10pp in most OOD conditions. With attention correction enabled, DP-GA further improves OOD performance by 15-30pp over autonomous execution, achieving 20-30pp higher success rates than the DP baseline.

Ablation Studies

Variant	Correction	ID (Δ from default)
No feature routing	—	81.1% (-2.2pp)
No feature routing	Enabled	25.5% (-58.9pp)
No attention supervision	—	52.2% (-31.1pp)
No attention supervision	Enabled	17.8% (-66.7pp)

Disabling feature routing causes a ~60pp crash when correction is enabled, showing that proper integration of override-path features during training is critical for leveraging human correction. Removing attention supervision drops performance by ~65pp and ~30pp with and without correction respectively, confirming that robustness requires attention features aligned with task-relevant structures.

Comparison with Marker Overlay Prompting

Compared to directly overlaying marker boxes/points on input images (MOP-Box/MOP-Point), DP-GA outperforms by 20+pp across ID and OOD conditions. Simply overlaying markers on input images is insufficient for reliably exploiting task-relevant information; explicitly predicting attention points with auxiliary supervision is far more effective.

Figure 1: GuidedAttention framework overview. Demonstrations are annotated with attention keypoints only in the initial frame. During rollout, users may optionally correct keypoints once, after which they are auto-tracked throughout execution.

Limitations

Author-stated: Tracking failures were observed in tasks involving severe deformation, which may limit applicability. While static keypoints are unaffected by the tracking module and the framework can integrate improved trackers, tracking quality remains a partial performance bottleneck.

Analysis: Human correction is limited to the rollout's first frame—mid-execution intervention is impossible if attention drifts later (e.g., after object occlusion and reappearance). The number of keypoints is pre-defined as fixed, not adapting to tasks where the number of attention foci dynamically changes. All simulation tasks use a single overhead camera; multi-view interactive correction is not thoroughly evaluated. The linear coordinate head's inverse mapping is underdetermined when $d_f>2$; although mitigated by pseudoinverse and alignment loss, it may theoretically lose information in high-dimensional features.

Conclusion and Future Work

GuidedAttention's core contribution elevates visual attention from an implicit internal state of end-to-end policies to an explicit, inspectable, correctable intermediate representation. Through three techniques—DETR-style keypoint prediction, forward-inverse mapping consistency of the linear coordinate head, and randomized feature routing—humans can correct OOD attention drift by clicking once at the first frame. Ablation studies prove that robustness does not arise from the Transformer architecture itself but from correctly integrating human-corrected attention into policy conditioning during training. Comparison with marker overlay prompting shows that explicit attention prediction with auxiliary supervision far outperforms simple pixel-level overlays.

The best interpretability isn't letting humans see what the black box thinks—it's letting them gently tug its attention back when it drifts.