注意力可视化

相关源文件

生成此 wiki 页面时使用了以下文件作为上下文：

attention_viz 包提供工具，用于在推理期间从 Action Chunking Transformer (ACT) 模型中提取、可视化并交互操作 cross-attention 权重。它用于调试模型关注区域、失败分析和视觉依赖的交互式测试。

概述

在基于 ACT 的策略中，decoder 使用 cross-attention 关注来自多个相机的视觉特征。attention_viz 通过 PyTorch hooks 截获这些权重，并将其映射回原始相机图像，生成 heatmap。

关键能力：

系统架构

可视化流水线连接高性能推理环境（PyTorch）与 ROS 2 消息系统。

数据流图

下图展示 attention 数据如何从模型层流向最终可视化结果。

“Natural Language Space to Code Entity Space: Attention Pipeline”

        graph TD
    subgraph "Inference Node (lerobot_policy_node)"
        A["ACTPolicy"] -->|"Forward Pass"| B["ACTDecoderLayer"]
        B -->|"Intercept via"| C["AttentionWeightHook"]
        C -->|"Monkey-patch"| D["MultiheadAttention.forward"]
        D -->|"Capture"| E["AttentionWeights Msg"]
    end

    subgraph "Visualization Node (attention_visualization_node)"
        E -->|"/attention/weights"| F["AttentionVisualizationNode"]
        G["Camera Drivers"] -->|"/camera/*/image_raw"| F
        F -->|"Process"| H["visualize_attention_single"]
        H -->|"Apply JET Colormap"| I["Heatmap Image"]
    end

    I -->|"Publish"| J["/visualization/heatmap/*"]
    I -->|"Save"| K["save_dir (JPEG Files)"]

来源：src/attention_viz/README.md:21-45，src/attention_viz/attention_viz/attention_hook.py:25-36，src/attention_viz/attention_viz/attention_visualization_node.py:57-61

关键组件

1. AttentionWeightHook

AttentionWeightHook 类负责在不修改 lerobot 库源码的情况下提取权重。它使用 “monkey-patching” 强制 nn.MultiheadAttention 返回权重，即使模型配置为高速推理（通常 need_weights 为 False）src/attention_viz/attention_viz/attention_hook.py:8-10。

2. AttentionVisualizationNode

这是一个 ROS 2 节点，会将 attention 消息与相机帧同步，生成可视化输出。

3. 交互式 Masking

该组件位于 attention_masking.py，提供用于调试视觉依赖的 GUI。

“Natural Language Space to Code Entity Space: Interactive Masking”

        graph LR
    subgraph "GUI Interaction"
        M["_MaskDrawer"] -->|"Mouse Drag"| P["Pixel Mask (uint8)"]
    end

    subgraph "Tensor Conversion"
        P -->|"process_pixel_mask_to_feature_mask"| F["Feature Mask"]
        F -->|"build_transformer_attention_masks"| T["Transformer Mask"]
    end

    T -->|"Apply to"| H["AttentionWeightHook._decoder_attention_mask"]

来源：src/attention_viz/attention_viz/attention_masking.py:53-58，src/attention_viz/attention_viz/attention_hook.py:185-194

配置

该包可通过 visualization_params.yaml 或中心化 robot_config 系统配置。

参数

默认值

说明

visualization_mode

file

file（保存到磁盘）、realtime（GUI）或 interactive src/attention_viz/config/visualization_params.yaml:3-3

queries_to_visualize

[0, 20, 40, 60, 80]

要可视化的 action chunk 中的步骤 src/attention_viz/config/visualization_params.yaml:7-7

blend_alpha

0.4

heatmap overlay 的透明度 src/attention_viz/config/visualization_params.yaml:11-11

save_dir

attention_visualizations

file 模式下 JPEG 输出目录 src/attention_viz/config/visualization_params.yaml:5-5

来源：src/attention_viz/config/visualization_params.yaml:1-14

实现细节

Attention 提取逻辑

系统通过对多头 attention 求平均或选择特定 head 来处理 multi-head attention。extract_attention_per_camera 函数会根据 feature_map_size 和非图像 token（如 [CLS] token 或 goal tokens）数量，对展平的 attention vector 进行切片 src/attention_viz/attention_viz/utils.py:169-195。

# Logic for slicing visual tokens per camera
feature_map_area = feature_map_size[0] * feature_map_size[1]
visual_tokens = attn[:, query_idx, num_non_image_tokens:] # Skip non-image tokens
# ...
for i, cam_key in enumerate(camera_keys):
    start = i * feature_map_area
    end = start + feature_map_area
    result[cam_key] = weights[start:end]

来源：src/attention_viz/attention_viz/utils.py:210-222

ROS 2 消息转换

自定义 ibrobot_msgs/msg/AttentionWeights 消息在 ROS 网络中承载原始 tensor。

来源：src/attention_viz/attention_viz/utils.py:74-167