# Coarse-to-fine Q-attention with Tree Expansion

Stephen James and Pieter Abbeel
UC Berkeley
{stepjam<EMAIL_ADDRESS>
###### 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. Videos and code found at:
https://sites.google.com/view/q-attention-qte.

##  I Introduction

Coarse-to-fine Q-attention [1], and its role within the Coarse-to-fine
Attention-driven Robotic Manipulation (C2F-ARM) system, has increased the
feasibility of training vision-based, sparse-rewarded reinforcement learning
agents in the real world. Given a small number of initial demonstrations and
quick exploration phase, C2F-ARM learns to output a series of next-best poses
which when given to a motion planner, leads to successful completion of
manipulation tasks. The algorithm gains its sample efficiency, in part, due to
the discretization of the translation space (voxelization) via the coarse-to-
fine Q-attention. To overcome the need to voxelize a large workspace — which
would be memory intensive and infeasible for a large number of voxels —
coarse-to-fine Q-attention gradually increases the resolution by extracting
the highest valued voxel at each level in the hierarchy, and using it as the
center of the higher-resolution voxelization in the next level. Because this
coarse-to-fine prediction is applied over several (increasingly finer) levels,
it gives a near-lossless prediction of the translation.

Figure 1: At the top-level of the coarse-to-fine Q-attention, it is not
feasible to distinguish similar-looking objects without first inspecting at a
finer resolution. For example, at the coarsest voxelization, the numbers on
the blocks above will not be legible. Q-attention with Tree Expansion (QTE)
expands (dotted black lines) and accumulates value estimates (dashed red
lines) across the top-k voxels at each Q-attention depth.

Although effective in a number of simulation and real-world tasks, Q-attention
suffers from what we call “coarse ambiguity” — when the voxelization is
significantly coarse, it is not feasible to distinguish between similar-
looking objects without first inspecting them at a finer resolution. Take
Figure 1 as a motivating example, where the goal is to grasp the cube with the
number ‘1’; at the top level voxelization, the numbers on the blocks will not
be legible, thus limiting the root Q-attention’s ability to make an informed
decisions about which block to attend to. This “coarse ambiguity” is further
shown in Figure 2. Unfortunately, simply increasing the voxel grid size at the
root of the Q-attention (i.e., directly operating at a finer resolution) would
require significantly memory, and would not scale to larger workspaces;
therefore an alternative solution is required.

In this work, we look to improve the performance of coarse-to-fine Q-attention
by enriching the information available to the coarsest Q-attention layers. We
do this by taking inspiration from Monte Carlo Tree Search (MCTS) — where
Monte Carlo simulation is used to accumulate value estimates to guide towards
highly rewarding trajectories in a search tree. However, rather than searching
across time (i.e., simulating a ‘game’), we instead search and accumulate
value estimates across space, thus allowing values from the fine-layers to
propagate to the coarse layers and overcoming the “coarse ambiguity”.

Our method, Q-attention with Tree Expansion (QTE), is built into the
Attention-driven Robot Manipulation (ARM) system to give our new C2F-ARM+QTE
system. This system is evaluated on 12 simulated RLBench [2] tasks (Figure 3)
and compared against vanilla C2F-ARM. We show that C2F-ARM+QTE can accomplish
a wider range of tasks, and much like C2F-ARM, is capable of learning
sparsely-rewarded real-world tasks.

To summarize, we propose our main contribution, Q-attention with Tree
Expansion (QTE): an extension to coarse-to-fine Q-attention that allows top-
level coarse layers to inspect the lower-level fine layers before action
selection. Along with this contribution, we also propose a new system,
C2F-ARM+QTE, which improves the performance over C2F-ARM on a set of RLBench
tasks.

##  II Related Work

Figure 2: Example of “coarse ambiguity”. At the coarsest levels of the
Q-attention, it is difficult to distinguish between objects. Highlighted red
voxels show the attention.

Robot Manipulation through the use of deep reinforcement learning (RL) has
been used to solve a variety of tasks, including object pushing [3], block
lifting [4], cloth manipulation [5], lego stacking [6], and grasping [7, 8],
to name but a few. Rather than using continuous control actor-critic
reinforcement learning algorithms, a recent trend instead discretizes part, or
all, of the action space [9, 10, 11, 12, 1]. In particular, the recently
proposed C2F-ARM [1] has shown the ability to rapidly learn sparse-reward, and
image-based manipulation tasks by discretizes the action space in a coarse-to-
fine manner via a Q-attention network. The voxel representation used in
C2F-ARM [1] has sporadically been used for manipulation. Notable manipulation
works include MoreFusion [13], which used voxels to perform multi-object
reasoning to improve 6D pose estimation for precise pick-and-place in clutter;
it has also been used for safe object extraction [14], and object
reorientation [15]. Song _et al_. [16] and Breyer _et al_. [17] fed a voxel
grid representation to a neural network to generate 6 DoF grasp actions.

In this work, we aim to improve and extend the coarse-to-fine Q-attention
component of ARM. However, this is not the first work to propose extensions to
ARM. ARM [12], has recently been extended to a Bingham policy parameterization
[18] to improve training stability. Another extension has sought to improve
the control agent within coarse-to-fine ARM to use learned path ranking [19]
to overcome the weaknesses of traditional path planning. We aim to further
improve coarse-to-fine Q-attention and tackle the “coarse ambiguity” issue
that Q-attention suffers from. Our solution shares motivational similarities
with Monte Carlo Tree Search, which we cover below.

Monte Carlo Tree Search (MCTS) has predominantly been used for decision making
within combinatorial games, including Go[20], Poker[21], and Chess[22], to
name but a few. Driven by successes in games, MCTS has increasingly been
applied outside of this area; one such case is within the realm of robot
manipulation. Rearrangement planning — where the goal is to find a sequence of
transit and transfer motions to move a set of objects to a target arrangement
— has recently seen success from MCTS. Labbe et al. [23] combined Monte Carlo
Tree Search with a standard robot motion planning algorithm (e.g., RRT).
Eljuri et al. [24] combined MCTS and a Motion Feasibility Checker (MFC), which
selects a set of feasible poses for the robot from a feasibility database.
Recently, MCTS has been combined with ‘Task and Motion Planning’ (TAMP), by
constructing an extended decision tree for symbolic task planning and high-
dimension motion variable binding [25]. Outside of symbolic planning, Huang et
al. [26] learned to retrieving a target object from clutter performing a Monte
Carlo search over high-level non-prehensile actions. Its search time was
subsequently improved by learning to predict the discounted reward of each
branch in MCTS without the need to roll out [27]. Bai et al. [28] also tackle
non-prehensile manipulation by proposing an MCTS algorithm guided by a policy
network which is trained by imitation and reinforcement.

Another related area that we briefly want to mention is beam search — a search
algorithm that expands the most promising nodes of a graph in a limited set.
Beam search has almost exclusively been used in natural language processing,
and to the best our knowledge, has not been explored for action selection in
robotics.

##  III Background

Figure 3: In simulation, our method is evaluated on a total of 12 RLBench [2]
tasks. Task Set 1 are tasks that have previously been shown to perform well
with C2F-ARM. Task Set 2 are tasks that were chosen due to their ambiguity at
a coarse scale. Note that the position and orientation of objects are placed
randomly at the beginning of each episode.

###  III-A Reinforcement Learning

The reinforcement learning problem consists of and agent and environment. The
agents takes actions $\mathbf{a}\in\mathcal{A}$, and transitions to new states
$\mathbf{s}\in\mathcal{S}$, where $\mathbf{s}_{t}$ and $\mathbf{a}_{t}$ are
the state and action at time step $t$. Upon transitioning to a new state, the
agent receives a reward $r$ from a reward function
$R(\mathbf{s}_{t},\mathbf{a}_{t})$. The goal of reinforcement learning is for
the agent to learn a policy $\pi$ that aims to maximize the expected future
(discounted) reward:
$\operatorname{\mathbb{E}}_{\pi}[\sum_{t}\gamma^{t}R(\mathbf{s}_{t},\mathbf{a}_{t})]$,
where $\gamma\in[0,1)$ is a discount factor. The expected return after taking
action $\mathbf{a}$ in state $\mathbf{s}$ can be estimated via a value
function $Q(s,a)$.

Deep Q-learning [29] approximates the value function $Q_{\psi}$ via a deep
convolutions network with parameters $\psi$. Exploration data is stored into a
replay buffer, and subsequently used for sampling mini-batches of transitions
to minimize:
$\operatorname{\mathbb{E}}_{(\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}_{t+1})\sim\mathcal{D}}[(\mathbf{r}+\gamma\max_{\mathbf{a}^{\prime}}Q_{\psi^{\prime}}(\mathbf{s}_{t+1},\mathbf{a}^{\prime})-Q_{\psi}(\mathbf{s}_{t},\mathbf{a}_{t}))^{2}]$,
where $Q_{\psi^{\prime}}$ is a target network — a periodic copy of $Q_{\psi}$,
which helps prevent Q-value overestimation. The Q-attention module, which we
discuss in the next section, builds upon the work of Deep Q-learning.

###  III-B Coarse-to-fine Q-attention

Algorithm 1 C2F-ARM+QTE. Addition to C2F-ARM in green.

1:Initialize Coarse-to-fine Q-attention $Q$.

2:Initialize buffer $\mathcal{D}$ with demos;

3:apply keyframe discovery and demo augmentation.

4:for each iteration do

5: for each environment step $t$ do

6: $\mathbf{o}_{t}\leftarrow(\mathbf{b}_{t},\mathbf{p}_{t},\mathbf{z}_{t})$

7: $\mathbf{c}^{0}\leftarrow$ Scene centroid

8: $\text{coords}\leftarrow[\text{ }]$

9: for each ($n$ of $N$) Q-attention depths do

10:
$\mathbf{v}^{n}\leftarrow\mathbf{V}(\mathbf{o}_{t},\mathbf{e}^{n},\mathbf{c}^{n})$

11:
$q^{n}\leftarrow\text{QTE}(n,\mathbf{o}_{t},\mathbf{e}^{n},\mathbf{c}^{n})$

12: $(x^{n},y^{n},z^{n})\leftarrow$ retrieve indices of $q^{n}$

13: $\text{coords.append(}(x^{n},y^{n},z^{n})\text{)}$

14: if $n==N$ then

15:
$\alpha,\beta,\gamma,\omega\leftarrow\operatorname*{arg\,max}_{\mathbf{a}^{h}}Q^{h}_{\phi}(\tilde{\mathbf{v}^{N}},\mathbf{a}^{h})$
for $h\in\\{0,1,2,3\\}$

16: $\mathbf{c}^{n+1}\leftarrow(x^{n},y^{n},z^{n})$

17: $\mathbf{a}_{t}\leftarrow(\mathbf{c}^{N},\alpha,\beta,\gamma,\omega)$

18: $\mathbf{o}_{t+1},\mathbf{r}\leftarrow env.step(\mathbf{a}_{t})$

19:
$\mathcal{D}\leftarrow\mathcal{D}\cup\left\\{(\mathbf{o}_{t},\mathbf{a}_{t},\mathbf{r},\mathbf{o}_{t+1},\text{coords})\right\\}$

20: for each gradient step do

21:
$\theta_{n}\leftarrow\theta_{n}-\lambda_{Q}\hat{\nabla}_{\theta_{n}}J_{Q}(\theta_{n})$
for $n\in\\{0,\dots,N\\}$

22: $\theta^{\prime}_{n}\leftarrow\tau\theta_{n}+(1-\tau)\theta^{\prime}_{n}$
for $n\in\\{0,\dots,N\\}$

Coarse-to-fine Q-attention [1] is a reinforcement learning (robot
manipulation) module that accepts a coarse 3D voxelization of the scene, and
learns to attend to salient areas (or objects) in the scene that would lead to
task success. Upon choosing a salient area, it ‘zooms’ into that area and re-
voxelizes the scene at a a higher spacial resolution. By applying this coarse-
to-fine process iteratively, we get a near-lossless discretization of the
translation space. This is also combined with a discretized rotation
prediction to give a next-best pose. This module was built into the C2F-ARM
[1] system, and resulted in the ability to learn sparsely-rewarded, image-
based manipulation tasks in very few environment steps. The C2F-ARM system
takes the next-best pose output from Q-attention, and uses a motion planner to
get the actions required to reach the goal pose. Coarse-to-fine Q-attention is
an extension of the earlier 2D Q-attention [12] which attended to pixels,
rather that voxels, and did not recursively apply the attention in a coarse-
to-fine manner.

In order to overcome exploration in a sparse reward environment, ARM [12] (the
predecessor to C2F-ARM) used a handful of demonstrations along with two
demonstration pre-processing steps: keyframe discovery and demo augmentation.
The Keyframe discovery method aims to assist the the Q-attention to quickly
converge and suggest meaningful points of interest in the initial phase of
training; while demo augmentation is a augmentation strategy that slices the
demonstration trajectories into multiple smaller trajectories, thereby
increases the initial number of demo transitions in the replay buffer.

##  IV Method

Algorithm 2 Q-attention Tree Expansion (QTE).

1:procedure QTE($n,\mathbf{o}_{t},\mathbf{e}^{n},\mathbf{c}^{n}$)

2:
$\mathbf{v}^{n}\leftarrow\mathbf{V}(\mathbf{o}_{t},\mathbf{e}^{n},\mathbf{c}^{n})$

3: if $n==N$ then $\triangleright$ Base case

4: return
$\max_{\mathbf{a}^{\prime}}Q_{\theta_{n}}(\mathbf{v}^{n},\mathbf{a}^{\prime})$

5: else$\triangleright$ Recursive case

6:
$\eta\leftarrow\text{TOPK}_{\mathbf{a}^{\prime}}(Q_{\theta_{n}}(\mathbf{v}^{n},\mathbf{a}^{\prime}))$

7: $\text{values}\leftarrow[\text{ }]$

8: for $k$ in $\\{0,\dots,K\\}$ do

9: $q^{k}\leftarrow\eta[k]$

10: $(x^{k},y^{k},z^{k})\leftarrow$ retrieve indices of $q^{k}$

11: $\mathbf{c}^{n+1}\leftarrow(x^{k},y^{k},z^{k})$

12:
$q^{\prime}=\text{QTE}(n+1,\mathbf{o}_{t},\mathbf{e}^{n+1},\mathbf{c}^{n+1})$

13: $\text{values.append}((q^{k}+q^{\prime})/2)$

14: return $\max\text{values}$

Our addition to Q-attention is simple: we frame 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. This is somewhat akin to (though not the same as)
Monte Carlo Tree Search (MCTS) — where Monte Carlo simulation is used to
accumulate value estimates to guide towards highly rewarding trajectories in a
search tree. Although similar in spirit, our Q-attention with Tree Expansion
(QTE) expands and accumulates value estimates across space, rather than
searching across time — i.e., we do not simulate a rollout in the environment.
Note that an obvious alternative to our approach would be to simply increase
the number of voxels at the root of the Q-attention, however this would
require significantly more memory, and would not be a scalable solution for
larger workspaces. Another obvious alternative to tackle “coarse ambiguity”
would be to simply operate on raw RGB-D images (rather than voxelization),
however these raw-observation policies do not perform as well as C2F-ARM [12].

Figure 4: Learning curves for 8 RLBench tasks. Both C2F-ARM and C2F-ARM+QTE
only receive 10 demos, and are stored in the replay buffer prior to
exploration. Note, these results intend to show that there is no loss in
performance when QTE is not necessary, and we do not expect to outperform
C2F-ARM. In Figure 5 however, the goal is to outperform C2F-ARM. Solid lines
represent the average evaluation over 5 seeds, while the shaded regions
represent the $std$.

###  IV-A C2F-ARM+QTE

Our system assumes we are operating in a partially observable Markov decision
process (POMDP), where an observation $\mathbf{o}$ consists of an RGB image,
$\mathbf{b}$, an organized point cloud, $\mathbf{p}$, and proprioceptive data,
$\mathbf{z}$. Actions consist of a 6D (next-best) pose and gripper action, and
the reward function is sparse, giving $100$ on task completion, and $0$ for
all other transitions.

We first formally define the coarse-to-fine Q-attention. We assume access to a
voxelization function
$\mathbf{v}^{n}=\mathbf{V}(\mathbf{o},\mathbf{e}^{n},\mathbf{c}^{n})$, which
takes the observation $\mathbf{o}$, a voxel resolution $\mathbf{e}^{n}$, and a
voxel grid center $\mathbf{c}^{n}$, and returns a voxel grid
$\mathbf{v}^{n}\in\mathbb{R}^{x×y×z×(3+M+1)}$ at depth $n$, where $n$ is the
depth/iteration of the Q-attention, and where each voxel contains the 3D
coordinates, $M$ features (e.g., RGB values, features, etc), and an occupancy
flag.

Given our Q-attention function $Q_{\theta_{n}}$ at depth $n$, the original
C2F-ARM system would extract the indices of the voxel with the highest value:

$\mathbf{v}_{ijk}^{n}=\operatorname*{\operatorname*{arg\,max}3D}_{\mathbf{a}^{\prime}}Q_{\theta_{n}}(\mathbf{v}^{n},\mathbf{a}^{\prime}),$
(1)

where $\mathbf{v}_{ijk}$ is the extracted voxel index located at $(i,j,k)$,
and $\operatorname*{\operatorname*{arg\,max}3D}$ is an argmax taken across
three dimensions (depth, height, and width). However, we now extract the
indices of the voxel with the highest value from the full tree expansion:

$\mathbf{v}_{ijk}^{n}=\text{extract\\_index}(\text{QTE}(n,\mathbf{o}_{t},\mathbf{e}^{n},\mathbf{c}^{n}))$
(2)

where QTE is the recursive tree expansion procedure discussed later in Section
IV-B. Note that in practice, QTE would also return the indices of the highest
voxels, but these have been omitted for brevity. As laid out in James et al.
[12], we use a separate Q-network for each resolution.

From this point, C2F-ARM proceeds as laid out in James et al. [1]. Using the
extracted indices, $\mathbf{v}_{ijk}^{n}$, we calculate the
$(x^{n},y^{n},z^{n})$ Cartesian coordinates of that voxel using known camera
parameters. Given that the extracted coordinates represent the next-best
coordinate to voxelize at a higher resolution, we can set these coordinates to
be the voxel grid center $\mathbf{c}$ for the next depth:
$\mathbf{c}^{n+1}=(x^{n},y^{n},z^{n})$. However, if this is the last depth of
the Q-attention, then $\mathbf{c}^{N}=\mathbf{c}^{n+1}$ represents the
continuous representation of the translation (i.e., the translation component
of the next-best pose agent).

Figure 5: Learning curves for an additional 4 RLBench tasks. Both methods
receive 10 demos for lift_numbered_block and pick_and_lift_small, 20 demos for
setup_checkers, and 40 demos for stack_blocks. Unlike Figure 4, we do expect
C2F-ARM+QTE to outperform C2F-ARM, as these tasks cause “coarse ambiguity”.
Solid lines represent the average evaluation over 5 seeds, while the shaded
regions represent the $std$.

As done in James et al. [1], the rotation of each axis is discretized in
increments of $5$ degrees, while the gripper is discretized to an open or
close binary. These components are recovered from an MLP branch (with
parameters $\phi$) from the final Q-attention depth:

$\alpha,\beta,\gamma,\omega\leftarrow\operatorname*{arg\,max}_{\mathbf{a}^{h}}Q^{h}_{\phi}(\tilde{\mathbf{v}^{N}},\mathbf{a}^{h})\text{
for }h\in\\{0,1,2,3\\},$ (3)

where $\alpha,\beta,\gamma$ represent the individual rotation axes, $\omega$
is the gripper action, and $\tilde{\mathbf{v}^{N}}$ are bottleneck features
from the final Q-attention depth. The final action (next-best pose) becomes
$\mathbf{a}_{t}=(\mathbf{c}^{N},\alpha,\beta,\gamma,\omega)$. The full system
is laid out in Algorithm 1, with the main contribution highlighted in green.

###  IV-B Tree Expansion

The aim of the tree expansion phase is to allow for accumulation of value
estimates across the top-k voxels at each Q-attention depth. We now walk
through the recursive expansion, which is summarized in Algorithm 2. The QTE
function takes the current Q-attention depth $n$, an observation $\mathbf{o}$,
a voxel resolution $\mathbf{e}^{n}$, and a voxel grid center $\mathbf{c}^{n}$.
The procedure starts by first voxelizing the observations to $\mathbf{v}^{n}$
(line 2). If we are not at the leaf, we run a top-k function
$\eta=\text{TOPK}_{\mathbf{a}^{\prime}}(Q_{\theta_{n}}(\mathbf{v}^{n},\mathbf{a}^{\prime}))$,
which takes the Q-attention output at depth $n$, and returns the top-k highest
Q-values (line 6). Recall that Q-attention builds on top of Deep Q-learning
[29], where each forward pass of the network outputs Q-values for each of the
actions (voxelized locations in our case). This means that the top-k is run
over the output of the network, rather than having to repeatedly run the
network for each action during the top-k search. For each of the $K$ values,
we extract coordinates and set them to be the voxel grid center $\mathbf{c}$
for the next depth: $\mathbf{c}^{n+1}=(x_{k},y_{k},z_{k})$ (line 9-11). We
then expand the tree at the next Q-attention depth $n+1$ by recursively
calling: $\text{QTE}(n+1,\mathbf{o}_{t},\mathbf{e}^{n+1},\mathbf{c}^{n+1})$
(line 12). The output of each of these recursive calls is stored in a list,
and after iterating through all $K$ branches, the max Q-value is returned
(lines 13-14).

If we are at the leaf of the tree expansion (i.e., the base case of
recursion), then we simply return the single highest Q-value (line 4):

$\max_{\mathbf{a}^{\prime}}Q_{\theta_{n}}(\mathbf{v}^{n},\mathbf{a}^{\prime}).$
(4)

To summarize, the key difference between C2F-ARM and C2F-ARM+QTE is the
addition of the top-k tree expansion. Note that when $K=1$, C2F-ARM+QTEbecomes
equivalent to C2F-ARM.

Tree expansion can be integrated at two key areas within C2F-ARM: when
choosing actions (for exploration) or when calculating the Q-value target (for
training). Where to integrate offers different benefits: applying at action
selection allows for better exploration, while using the expansion for
calculating the target Q-value allows for a more accurate estimate of the
value, and thus more stable gradient updates. Note that the latter is more
computationally expensive than the former because of the batched operation of
expanding the tree. The majority of our results applies the tree expansion for
both action selection and target value calculation; we study this further in
Section V.

##  V Results

Our results show the following 4 core takeaways:

  1. 1.

We are able to achieve the original performance of C2F-ARM on tasks where the
tree expansion is not necessary.

  2. 2.

Our method overcomes “coarse ambiguity” across a set of challenging tasks
involving small objects.

  3. 3.

Through a series of ablations, we study how performance varies as we adjust
the tree expansion breadth $K$, and the effect of when to apply the expansion.

  4. 4.

Our real-world qualitative results also show that the advantage of the tree
expansion extends to the real world.

###  V-A Quantitative Simulation Results

For our simulation experiments, we use RLBench [2], which emphasizes vision-
based manipulation and gives access to a wide variety of tasks with expert
demonstrations. Each task has a sparse reward of $1$, which is given only on
task completion, and $0$ otherwise. We follow James et al. [1], where we use a
coarse-to-fine depth of $2$, each with a voxel grid size of $16^{3}$. Unless
otherwise stated, C2F-ARM+QTE uses the tree expansion for both action
selection and target Q-value selection, with $k=10$.

#### No loss in performance when using tree expansion

To compare C2F-ARM with and without tree expansion, we select the same 8 tasks
as in previous work [12, 1]; these are tasks that are achievable using only
the front-facing camera. Additionally, these are tasks where there is no
advantage to having a tree expansion as there is only one object in the scene,
or objects can be distinguished at a coarse level due to their color. Figure 4
shows that there is no loss in performance when using QTE. The results
illustrate that C2F-ARM+QTE achieves either similar or better performance than
C2F-ARM.

#### Gain in performance on more challenging tasks

Unlike the previous experiment, we now purposefully choose 4 RLBench tasks
where there are many objects in the scene that look similar at a coarse level.
We briefly outline the tasks to emphasize why they are hard for vanilla
C2F-ARM.

Figure 6: Ablation of the tree expansion $k$ parameter on 2 RLBench tasks.
Solid lines represent the average evaluation over 5 seeds, while the shaded
regions represent the $std$. Figure 7: Ablation to assess how necessary using
tree expansion is for both action selection and target Q-value calculations.
Solid lines represent the average evaluation over 5 seeds, while the shaded
regions represent the $std$.

  * •

lift_numbered_block: The goal is to lift the white block with the number ‘1’,
among two distractors labeled with ‘2’ and ‘3’. At a coarse level, it is not
possible to read the numbers on the blocks.

  * •

pick_and_lift_small: The goal is to bring the small blue cube to the target
among 4 other small blue distractors with different shapes (e.g., star, moon,
etc). At a coarse level, it is not possible to distinguish the subtle
geometric differences between the small objects.

  * •

setup_checkers: The goal is to complete the checkers board by placing the
checker in the correct board position. At a coarse level, it is not possible
to distinguish the state of the board.

  * •

stack_blocks: The goal is to stack 2 of 4 red blocks on the target, among 4
other distractors. This is challenging because at a coarse level, each of the
4 red blocks could look equally promising to grasp, however at a finer level,
it would be possible to perceive what blocks are easier to grasp first.

Note for these 4 tasks, we use a RGB-D resolution of $256\times 256$ in order
to capture all of the detail required for the task. Figure 5 presents results
on these 4 tasks, and show that C2F-ARM+QTE attains better performance than
C2F-ARM. There is room for improvement, especially on the lift_numbered_block
task; we hypothesize that this is due to the numbers still being difficult to
read in the RGB image, and so further increasing the RGB-D resolution could
lead to performance gains.

#### Ablations

Figure 8: Real world qualitative task. Right: The goal is to reach the
Pentagonal prism, while ignoring the two distractors. Note, this is not the
camera viewpoint. Left: Two rollouts from the learned agent; each column shows
the RGB-D observations (from camera viewpoint) at $t=0$, $t=T/3$, $t=2T/3$ and
$t=T$.

For our ablations, we aim to understand two key design decisions that affect
performance: the effect of varying the tree expansion breadth $K$, and the
effect of when to apply the expansion — either (1) when choosing actions, (2)
when calculating the Q-value target, or (3) applying both. In Figure 6, we
show the effect of varying $K$ across $1,2,5$ and $10$. The results show that
as $K$ increases, performance improves; this makes intuitive sense, as a
higher $K$ allows for more information to be brought up to the coarsest level
of the Q-attention and allows for less ambiguity when choosing the highest
Q-value. In Figure 7, we show where is best to perform the tree expansion. The
results show that using tree expansion for both action selection and target
Q-value selection leads to the best performance. Encouragingly, only using
tree expansion for action selection also performs relatively well; this is
good because the computational burden for the action selection is minimal in
comparison to the target Q-value selection — which is $\sim\times 2$ more
expensive to train due to the batched tree expansion.

###  V-B Qualitative Real-world Results

To highlight the usefulness of the tree expansion in the real world, we design
a real-world task that involves three small yellow objects with different
shapes. The task is to simply reach the Pentagonal prism, while ignoring the
two distractors. We give both C2F-ARM and our C2F-ARM+QTE 3 demonstrations
(through tele-op via HTC Vive), and train for 10 minutes. When evaluated on 6
test episodes, C2F-ARM was only able to reach the target $2/6$ times, while
C2F-ARM+QTE achieved $6/6$. Due to “coarse ambiguity”, C2F-ARM would act as if
randomly selecting one of yellow objects. For the experiments, we use the
(low-cost) UFACTORY xArm 7, and a single RGB-D RealSense camera for all
experiments. These qualitative results are best viewed via the videos on the
project website.

##  VI Conclusion

In this paper, we have presented our tree expansion extension to coarse-to-
fine Q-attention. When the improved Q-attention is integrated into C2F-ARM — a
highly efficient manipulation algorithm — it is able to accomplish a larger
set of tasks; in particular, ones that contain small objects or objects that
look similar at a coarse level.

Below, we highlight weaknesses that would make for exciting future work. The
expansion of each branch can be done independently, however our current
implementation does this sequentially due to memory limits; future work could
improve training efficiency by making the tree expansion parallel. Although
our method has increased the repertoire of solvable tasks, it has
inadvertently introduced another hyperparameter in the from of the expansion
number $K$. A large value may be unnecessary for many tasks, while for other
tasks, a smaller value may cause unstable training; therefore, future work
could investigate how to adapt $K$ depending on task performance.

## Acknowledgments

This work was supported by the Hong Kong Center for Logistics Robotics.

## References

  * [1] S. James, K. Wada, T. Laidlow, and A. J. Davison, “Coarse-to-Fine Q-attention: Efficient Learning for Visual Robotic Manipulation via Discretisation,” _IEEE Conference on Computer Vision and Pattern Recognition_ , 2022.
  * [2] S. James, Z. Ma, D. Rovick Arrojo, and A. J. Davison, “RLBench: The robot learning benchmark & learning environment,” _IEEE Robotics and Automation Letters_ , 2020.
  * [3] L. Pinto, M. Andrychowicz, P. Welinder, W. Zaremba, and P. Abbeel, “Asymmetric actor critic for image-based robot learning,” _Robotics: Science and Systems_ , 2018.
  * [4] S. James and E. Johns, “3D simulation for robot arm control with deep Q-learning,” _Conference on Neural Information Processing Systems Workshop (Deep Learning for Action and Interaction)_ , 2016.
  * [5] J. Matas, S. James, and A. J. Davison, “Sim-to-real reinforcement learning for deformable object manipulation,” _Conference on Robot Learning_ , 2018.
  * [6] T. Haarnoja, V. Pong, A. Zhou, M. Dalal, P. Abbeel, and S. Levine, “Composable deep reinforcement learning for robotic manipulation,” in _2018 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2018, pp. 6244–6251.
  * [7] D. Kalashnikov, A. Irpan, P. Pastor, J. Ibarz, A. Herzog, E. Jang, D. Quillen, E. Holly, M. Kalakrishnan, V. Vanhoucke _et al._ , “Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation,” _Conference on Robot Learning_ , 2018.
  * [8] S. James, P. Wohlhart, M. Kalakrishnan, D. Kalashnikov, A. Irpan, J. Ibarz, S. Levine, R. Hadsell, and K. Bousmalis, “Sim-to-real via sim-to-sim: Data-efficient robotic grasping via randomized-to-canonical adaptation networks,” in _IEEE Conference on Computer Vision and Pattern Recognition_ , 2019, pp. 12 627–12 637.
  * [9] D. Morrison, P. Corke, and J. Leitner, “Closing the loop for robotic grasping: A real-time, generative grasp synthesis approach,” _Robotics: Science and Systems_ , 2018.
  * [10] A. Zeng, S. Song, S. Welker, J. Lee, A. Rodriguez, and T. Funkhouser, “Learning synergies between pushing and grasping with self-supervised deep reinforcement learning,” in _IEEE Intl. Conference on Intelligent Robots and Systems_. IEEE, 2018, pp. 4238–4245.
  * [11] A. Zeng, P. Florence, J. Tompson, S. Welker, J. Chien, M. Attarian, T. Armstrong, I. Krasin, D. Duong, V. Sindhwani _et al._ , “Transporter networks: Rearranging the visual world for robotic manipulation,” _Conference on Robot Learning_ , 2020.
  * [12] S. James and A. J. Davison, “Q-attention: Enabling efficient learning for vision-based robotic manipulation,” _IEEE Robotics and Automation Letters_ , 2022.
  * [13] K. Wada, E. Sucar, S. James, D. Lenton, and A. J. Davison, “MoreFusion: Multi-object reasoning for 6D pose estimation from volumetric fusion,” in _IEEE Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 14 540–14 549.
  * [14] K. Wada, S. James, and A. J. Davison, “SafePicking: Learning safe object extraction via object-level mapping,” _IEEE Intl. Conference on Robotics and Automation_ , 2022.
  * [15] K. Wada, S. James, and A. J. Davison, “ReorientBot: Learning object reorientation for specific-posed placement,” _IEEE Intl. Conference on Robotics and Automation_ , 2022.
  * [16] S. Song, A. Zeng, J. Lee, and T. Funkhouser, “Grasping in the wild: Learning 6dof closed-loop grasping from low-cost demonstrations,” _IEEE Robotics and Automation Letters_ , vol. 5, no. 3, pp. 4978–4985, 2020.
  * [17] M. Breyer, J. J. Chung, L. Ott, R. Siegwart, and J. Nieto, “Volumetric grasping network: real-time 6 dof grasp detection in clutter,” _arxiv preprint arxiv:2101.01132_ , 2021.
  * [18] S. James and P. Abbeel, “Bingham Policy Parameterization for 3D Rotations in Reinforcement Learning,” _arXiv preprint arXiv:2202.03957_ , 2022.
  * [19] S. James and P. Abbeel, “Coarse-to-Fine Q-attention with Learned Path Ranking,” _arXiv preprint arXiv:2204.01571_ , 2022.
  * [20] L. Kocsis and C. Szepesvári, “Bandit based monte-carlo planning,” in _European conference on machine learning_. Springer, 2006, pp. 282–293.
  * [21] G. V. d. Broeck, K. Driessens, and J. Ramon, “Monte-carlo tree search in poker using expected reward distributions,” in _Asian Conference on Machine Learning_. Springer, 2009, pp. 367–381.
  * [22] D. Silver, T. Hubert, J. Schrittwieser, I. Antonoglou, M. Lai, A. Guez, M. Lanctot, L. Sifre, D. Kumaran, T. Graepel _et al._ , “A general reinforcement learning algorithm that masters chess, shogi, and go through self-play,” _Science_ , vol. 362, no. 6419, pp. 1140–1144, 2018.
  * [23] Y. Labbé, S. Zagoruyko, I. Kalevatykh, I. Laptev, J. Carpentier, M. Aubry, and J. Sivic, “Monte-carlo tree search for efficient visually guided rearrangement planning,” _IEEE Robotics and Automation Letters_ , vol. 5, no. 2, pp. 3715–3722, 2020.
  * [24] P. M. U. Eljuri, G. A. G. Ricardez, N. Koganti, J. Takamatsu, and T. Ogasawara, “Combining a monte carlo tree search and a feasibility database to plan and execute rearranging tasks,” _IEEE Access_ , vol. 9, pp. 21 721–21 734, 2021.
  * [25] T. Ren, G. Chalvatzaki, and J. Peters, “Extended tree search for robot task and motion planning,” _arXiv preprint arXiv:2103.05456_ , 2021.
  * [26] B. Huang, S. D. Han, J. Yu, and A. Boularias, “Visual foresight trees for object retrieval from clutter with nonprehensile rearrangement,” _IEEE Robotics and Automation Letters_ , vol. 7, no. 1, pp. 231–238, 2021.
  * [27] B. Huang, T. Guo, A. Boularias, and J. Yu, “Self-supervised monte carlo tree search learning for object retrieval in clutter,” _arXiv preprint arXiv:2202.01426_ , 2022.
  * [28] F. Bai, F. Meng, J. Liu, J. Wang, and M. Q.-H. Meng, “Hierarchical policy for non-prehensile multi-object rearrangement with deep reinforcement learning and monte carlo tree search,” _arXiv preprint arXiv:2109.08973_ , 2021.
  * [29] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski _et al._ , “Human-level control through deep reinforcement learning,” _nature_ , vol. 518, no. 7540, pp. 529–533, 2015.