[Paper Review] Coarse-to-fine Q-attention with Tree Expansion

Stephen James, Pieter Abbeel|arXiv (Cornell University)|Apr 26, 2022

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Q-attention with Tree Expansion (QTE), an extension to coarse-to-fine Q-attention that alleviates 'coarse ambiguity' by propagating value estimates from fine-resolution layers back to coarse layers via tree expansion. By accumulating top-k value estimates across hierarchical levels, QTE enables better decision-making in tasks with visually similar or small objects, significantly improving sample efficiency and success rates in both simulation and real-world robot manipulation tasks compared to standard Q-attention.

ABSTRACT

Coarse-to-fine Q-attention enables sample-efficient robot manipulation by discretizing the translation space in a coarse-to-fine manner, where the resolution gradually increases at each layer in the hierarchy. Although effective, Q-attention suffers from "coarse ambiguity" - when voxelization is significantly coarse, it is not feasible to distinguish similar-looking objects without first inspecting at a finer resolution. To combat this, we propose to envision Q-attention as a tree that can be expanded and used to accumulate value estimates across the top-k voxels at each Q-attention depth. When our extension, Q-attention with Tree Expansion (QTE), replaces standard Q-attention in the Attention-driven Robot Manipulation (ARM) system, we are able to accomplish a larger set of tasks; especially on those that suffer from "coarse ambiguity". In addition to evaluating our approach across 12 RLBench tasks, we also show that the improved performance is visible in a real-world task involving small objects.

Motivation & Objective

• To address 'coarse ambiguity' in coarse-to-fine Q-attention, where similar-looking or small objects cannot be distinguished at low-resolution voxelizations.
• To improve sample efficiency and decision accuracy in vision-based, sparse-reward reinforcement learning for robot manipulation.
• To extend the C2F-ARM framework by enabling top-level attention to access information from finer-resolution layers before action selection.
• To evaluate the method on challenging RLBench tasks involving small or visually similar objects, particularly those where standard Q-attention fails.

Proposed method

• QTE introduces a tree structure where each node represents a voxel at a given resolution level, with children representing higher-resolution voxels centered on the top-k highest-value voxels from the parent layer.
• At each level, the top-k value estimates from the current layer are used to expand the tree and propagate value estimates upward, enriching coarse-level decisions with fine-level information.
• Tree expansion is applied during both action selection and Q-value target computation, allowing the agent to make informed decisions based on accumulated fine-grained value estimates.
• The method integrates into the C2F-ARM system as C2F-ARM+QTE, maintaining the original architecture while enhancing attention with hierarchical value accumulation.
• The tree expansion breadth K controls how many top candidates are considered at each level, with higher K improving ambiguity resolution at the cost of increased computation.
• The approach is inspired by Monte Carlo Tree Search but adapted for spatial, rather than temporal, search, enabling value propagation across resolution levels.

Experimental results

Research questions

• RQ1 Can hierarchical value propagation through tree expansion reduce 'coarse ambiguity' in vision-based robotic manipulation?
• RQ2 How does the number of expanded top-k candidates (K) affect performance in tasks with visually ambiguous or small objects?
• RQ3 Does applying tree expansion during both action selection and Q-value target computation yield better performance than applying it in isolation?
• RQ4 Can QTE improve generalization and success rates on real-world tasks involving small or similar-looking objects compared to standard Q-attention?
• RQ5 To what extent does increasing RGB-D resolution improve performance when combined with QTE on tasks with fine visual distinctions?

Key findings

• C2F-ARM+QTE achieved higher success rates than C2F-ARM across 12 RLBench tasks, particularly excelling in tasks with small or visually similar objects.
• On the 'lift_numbered_block' task, QTE significantly improved performance, though some failure cases remained due to poor number readability even at 256×256 resolution.
• In real-world testing with the uFactory xArm 7, C2F-ARM+QTE achieved 6/6 success in a small-object reaching task, compared to only 2/6 for standard C2F-ARM.
• Ablation studies showed that using tree expansion during both action selection and Q-value target computation yielded the best performance, with increasing K improving success rates.
• Tree expansion applied only during action selection still performed well, suggesting a viable trade-off between performance and computational cost.