💻 computer science

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

GuidedAttention is an interpretable imitation learning framework that enhances out-of-distribution robustness in robot manipulation by allowing users to inspect and correct task-relevant visual attention keypoints at rollout initialization, which are then automatically propagated throughout execution to guide a diffusion-based action policy.

Original authors: Masaki Murooka, Ryoichi Nakajo, Keisuke Shirai, Tomohiro Motoda, Hanbit Oh, Ryo Hanai, Yukiyasu Domae

Published 2026-07-27

📖 10 min read🧠 Deep dive

CC BY 4.0

Original authors: Masaki Murooka, Ryoichi Nakajo, Keisuke Shirai, Tomohiro Motoda, Hanbit Oh, Ryo Hanai, Yukiyasu Domae

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine teaching a robot to do a task, like picking up a cup or folding a towel. In the old days, engineers had to write thousands of lines of code to tell the robot exactly how to move its arm, where to look, and what to grab. It was like giving a child a rigid script to follow; if the script didn't match the real world perfectly, the child would freeze. Today, we use a smarter approach called "Imitation Learning." Instead of writing rules, we show the robot a video of a human doing the task, and the robot tries to learn the pattern by watching. It's like a kid learning to ride a bike by mimicking their older sibling.

However, there's a catch. When a robot learns this way, it builds a "black box" brain. It sees the world through a camera, but we have no idea what it is actually looking at. If the robot fails, we can't easily say, "Hey, you were looking at the wrong spot!" because the robot's attention is hidden inside complex math. This becomes a big problem when things change. If you move the cup to a new spot or put a colorful pattern on the table, the robot might get confused and fail because it was trained on a specific setup. It's like a student who memorized the answers to a test but panics when the teacher changes the order of the questions. We need a way to peek inside the robot's mind and say, "No, look at the cup, not the background!"

This is where a new framework called GuidedAttention comes in. Think of it as giving the robot a pair of "magic glasses" that show us exactly where it is focusing its attention. The researchers behind this work, Masaki Murooka and their team, created a system that doesn't just guess what to do; it first points to the important parts of the image, like a student highlighting the key words in a textbook. If the robot highlights the wrong thing, a human can step in, click the correct spot, and say, "Look here!" The robot then uses that corrected spot as a guide for the rest of the task, automatically tracking it as the robot moves.

The paper tests this idea in two ways: inside a computer simulation and in the real world with a real robot arm. They found that when the robot is in a familiar setting, it does pretty well on its own. But when things get weird—like moving the target to a new location or adding distracting colorful blocks to the table—the robot usually gets lost. That's where the "magic glasses" shine. When the human gives a tiny nudge at the very start of the task to correct the robot's focus, the robot's success rate jumps dramatically. In some tricky real-world tests, fixing the attention at the beginning helped the robot succeed about 80% more often than if it had to figure it out alone. It turns out that giving the robot a simple, correctable hint about where to look is the secret sauce for making it robust enough to handle a messy, changing world.

The Problem: The Robot's "Black Box" Brain

Imagine you are teaching a robot to tie its shoes. You show it a video, and the robot learns to copy your movements. But here's the problem: the robot doesn't "see" the shoe the way you do. It sees a jumble of pixels. In modern robot learning, we use something called End-to-End policies. This means the robot takes a picture as input and spits out a movement as output, skipping all the middle steps where a human might explain what is happening.

The issue is that this process is a "black box." We don't know what the robot is paying attention to. If the robot fails to tie the shoe, we can't easily tell if it was looking at the laces, the floor, or a shadow. It's like trying to fix a car engine without being able to see inside the hood. Even worse, if you move the shoe to a different spot on the table (a situation called Out-of-Distribution or OOD), the robot might get completely confused because it was trained on the shoe being in one specific place. It's like a student who memorized the answer "5" for a math problem but fails when the numbers change, because they didn't learn the concept, just the specific example.

The Solution: GuidedAttention (The "Magic Glasses")

The authors propose a solution called GuidedAttention. Instead of letting the robot guess what to look at, they force it to predict a set of keypoints—little dots that mark the important parts of the image. Think of these keypoints like a treasure map. The robot has to draw an "X" on the spot it needs to grab (like the tip of a rope) and another "X" on the target (like the hole it needs to go through).

Here is the cool part: these "X"s are visible to us. They are an interpretable intermediate representation. This means we can see exactly what the robot thinks is important. If the robot draws an "X" on the wrong spot, a human can step in and move the "X" to the right place. This is the correctable part.

Once the human fixes the "X" at the very beginning of the task, the robot doesn't need help anymore. It uses a tracking module (like a smart camera that follows a moving object) to keep those "X"s locked onto the right spots as the robot moves. It's like giving the robot a sticky note that says "Look here!" and then having a friend follow that note with a laser pointer for the rest of the game.

How It Works: The Robot's New Brain

The system uses a type of AI called a Diffusion Policy. You can think of this as a robot that learns by "denoising." Imagine a picture that is covered in static noise. The robot's job is to slowly remove the noise to reveal the correct movement. Usually, the robot tries to guess the movement based on the whole blurry picture.

In this new system, the robot first predicts the keypoints (the "X"s). Then, it uses those "X"s to guide the denoising process. It's like telling the robot, "Ignore the messy background; focus only on the area around these two dots."

The paper introduces a clever trick called the Keypoint Override Mechanism.

Training: The robot learns to predict the dots by watching humans do the task. Humans only mark the dots in the very first frame of the video, and a computer program tracks them for the rest of the video.
Rollout (The Real Task): When the robot tries the task on its own, it predicts the dots. If the robot is doing well, great! But if the robot is confused (maybe because the table looks different), a human can click to fix the dots once at the start.
The Magic: The system has a special math trick that ensures the "fixed" dots still make sense to the robot's brain. It doesn't just swap the dots; it swaps the meaning of the dots so the robot understands the correction perfectly.

The Results: Does It Actually Work?

The team tested this in two places: a computer simulation (a virtual world) and the real world with a physical robot arm (a Universal Robots UR5e).

In the Simulation:
They tested three tricky tasks:

Cable: Guiding a flexible cable through a narrow gap.
Ring: Putting a ring over a pole.
Particle: Scooping small particles into a box.

They made the tasks harder by moving the targets to new spots (Positional OOD) or changing the table's texture with colorful patterns (Appearance OOD).

Without help: The standard robots (baselines) failed a lot when things changed. For example, in the "Appearance OOD" test (colorful patterns), the standard robot succeeded only about 28.9% of the time.
With GuidedAttention (Autonomous): The robot did better, succeeding about 45.6% of the time.
With GuidedAttention (Human Correction): When a human fixed the dots at the start, the success rate jumped to 67.8%. That's a massive improvement!

In the Real World:
They tested with real robots on tasks like putting a cup in another cup, picking up a chain, and folding a towel.

Positional OOD (Moving the target): The standard robots struggled, with the DP baseline achieving only 35.6% success. GuidedAttention, even without help, did better. But with a single human correction at the start, the success rate soared to 71.1%.
Appearance OOD (Messy table): This was the hardest test. A standard robot failed significantly, achieving 0% success on the towel task. GuidedAttention without help also struggled, managing only 13.3% success. But with human correction, the robot succeeded 90% of the time!

Why This Matters

The paper shows that we don't need to build a robot that knows everything. Instead, we can build a robot that knows how to look, and we can give it a little help when it gets confused.

The authors found that the biggest gains came when the robot was in a new or messy environment. In familiar settings, the robot was already okay. But when the world changed, the ability for a human to say, "No, look at the cup, not the background," saved the day.

The paper also rules out some other ideas. They tried just putting a box or a dot on the screen to tell the robot where to look (called "marker overlay prompting"), but that didn't work as well. The robot needs to predict the dots itself so it can learn the pattern, but be able to correct them when it gets it wrong.

The Limits

The authors are honest about what their system can't do yet.

Simple Keypoints: The system assumes that a few dots are enough to describe the whole task. If a task is super complex, like juggling ten balls, a few dots might not be enough.
2D Only: The dots are just on the flat picture. If the robot needs to know how deep something is, or if the dot gets hidden behind an object, the system might get confused.
Tracking Issues: The system relies on a tracker to follow the dots. If the robot moves too fast or the object gets blocked for a long time, the tracker might lose the dot.

Conclusion

GuidedAttention is like giving a robot a pair of glasses that show us what it's thinking. It bridges the gap between a robot's "black box" brain and our human ability to guide it. By letting us correct the robot's focus just once at the start, the robot can handle messy, changing environments much better than before. It's not a magic wand that solves everything, but it's a powerful step toward robots that are smart enough to learn, but flexible enough to listen to us when they get stuck.

Technical Summary: GuidedAttention for OOD-Robust Robot Manipulation

1. Problem Statement

End-to-end visuomotor policies, while capable of learning dexterous behaviors from experience, suffer from a critical limitation: they offer little opportunity for humans to understand or correct the policy's visual attention during execution. In traditional modular pipelines, human operators can intervene based on visual cues, but in end-to-end systems, the policy must discover task-relevant visual cues entirely on its own. Consequently, these policies often fail in out-of-distribution (OOD) scenarios—such as when object positions or appearances shift—simply because they lack a simple, crucial hint that a human could readily provide. Existing methods either rely on continuous human supervision or lack an explicit, interpretable interface for correcting the policy's focus.

2. Methodology: GuidedAttention Framework

The authors propose GuidedAttention, a visuomotor imitation learning framework that introduces interpretable and correctable visual attention as an explicit intermediate representation. The core architecture, termed Diffusion Policy with GuidedAttention (DP-GA), operates as follows:

A. Policy Architecture

Instead of conditioning action generation directly on global visual features (e.g., from CNNs or point cloud encoders), DP-GA inserts an explicit visual attention layer between perception and action:

Input: Multi-view RGB observations and proprioceptive joint states.
Visual Attention Encoder (AttnEnc): A DETR-style architecture (ResNet-18 backbone + Transformer encoder-decoder) predicts a set of $N_k$ task-relevant attention keypoints ( $p_i$ ) and their corresponding attention features ( $f_i$ ) from the camera images.
Action Policy: A CNN-based denoising diffusion policy (U-Net backbone) is conditioned on the concatenation of proprioceptive features and the extracted attention features.

B. Training and Annotation

Data Collection: Users annotate task-relevant keypoints only on the initial frame of each teleoperated demonstration.
Propagation: An off-the-shelf tracking model (Co-Tracker) propagates these keypoints through the remaining frames of the demonstration. Static keypoints remain fixed.
Loss Functions: The model is trained to minimize:
1. Keypoint Loss ( $L_{kp}$ ): Mean Squared Error (MSE) between predicted and ground-truth keypoint coordinates.
2. Feature Alignment Loss ( $L_{feat}$ ): Ensures consistency between features predicted by the encoder and features reconstructed from ground-truth coordinates.

C. Keypoint Override Mechanism (Inference)

A unique feature of GuidedAttention is the ability for users to intervene during rollout without continuous supervision:

Autonomous Execution: The policy predicts keypoints and actions autonomously.
One-Time Correction: If the predicted attention drifts (e.g., due to OOD conditions), the user can click to correct a subset of keypoints only at the rollout initialization.
Automatic Propagation: The corrected coordinates are converted back into attention features via an inverse mapping ( $h^{-1}_\theta$ ) and tracked automatically throughout the entire execution.
Training Robustness: To ensure the diffusion model can handle both predicted and corrected features, the training process employs randomized feature routing, randomly selecting whether each keypoint feature comes from the prediction path or the override path.

3. Key Contributions

Explicit Intermediate Representation: The paper introduces visual attention keypoints as an interpretable, human-readable intermediate layer between perception and action, allowing users to inspect and correct the policy's focus.
DP-GA Architecture: A novel diffusion-based policy that conditions action generation on these attention features, supporting a "one-time correction, automatic tracking" workflow.
Robustness to Distribution Shifts: The framework demonstrates that human-guided correction effectively recovers task-relevant visual focus when autonomous prediction fails under positional and appearance OOD conditions.

4. Experimental Results

The method was evaluated in both simulation (MuJoCo) and real-world settings using a UR5e robot across tasks such as cable threading, ring placement, particle scooping, cup insertion, chain manipulation, and towel folding.

Simulation Results

Performance: DP-GA consistently outperformed baselines (standard Diffusion Policy and ACT) across In-Distribution (ID) and OOD conditions.
OOD Gains: Under positional and appearance OOD, DP-GA with optional human correction achieved 20–30 percentage points (pp) higher success rates than the standard Diffusion Policy baseline.
Ablation: Removing the feature routing mechanism caused a drastic ~60 pp drop in performance when correction was enabled, highlighting the necessity of training the model to accept both predicted and corrected features.

Real-World Results

Generalization: DP-GA variants outperformed baselines in ID conditions by 5–30 pp.
OOD Resilience: In positional OOD, DP-GA maintained higher success rates than baselines even without correction. With correction, it achieved an additional 20–40 pp improvement.
Appearance OOD: In the "Towel Fold" task with severe visual clutter, baselines failed significantly. DP-GA with human correction recovered performance, showing an ~80 pp improvement over its autonomous execution and ~40 pp over the baseline.
Qualitative Analysis: Visualizations confirmed that DP-GA maintained focus on task-relevant structures (e.g., chain tips, cup rims) where baselines lost focus, and that one-time correction successfully guided the robot through complex folds despite visual distractors.

5. Significance and Claims

The paper claims that GuidedAttention provides a practical solution to the "black box" nature of end-to-end visuomotor policies. By treating visual attention as an explicit, correctable interface, the framework bridges the gap between human intuition and autonomous execution. The authors emphasize that this approach allows robots to recover from perception failures in challenging OOD environments without requiring continuous human supervision, thereby enhancing the robustness and reliability of imitation learning in real-world manipulation tasks.

6. Limitations

The authors acknowledge several constraints:

Keypoint Sufficiency: The method assumes a small number of manually defined keypoints can represent all task-relevant structures, which may not hold for highly complex scenarios.
2D Limitations: Annotations rely on 2D coordinates and do not explicitly account for depth or occlusion, potentially degrading performance if keypoints become invisible.
Tracking Dependency: The system relies on off-the-shelf tracking (Co-Tracker); performance may degrade in scenarios with severe deformation or prolonged occlusion, though static keypoints mitigate this partially.