Attention-driven Robotic Manipulation
Authors: Stephen James
Andrew Davison
Venue:  (2021)
Keywords: Robotics
Robot Manipulation
Reinforcement Learning

Abstract: Despite the success of reinforcement learning methods, they have yet to have their breakthrough moment when applied to a broad range of robotic manipulation tasks. This is partly due to the fact that reinforcement learning algorithms are notoriously difficult and time consuming to train, which is exacerbated when training from images rather than full-state inputs. As humans perform manipulation tasks, our eyes closely monitor every step of the process with our gaze focusing sequentially on the objects being manipulated. With this in mind, we present our Attention-driven Robotic Manipulation (ARM) algorithm, which is a general manipulation algorithm that can be applied to a range of real-world sparse-rewarded tasks without any prior task knowledge. ARM splits the complex task of manipulation into a 3 stage pipeline: (1) a Q-attention agent extracts interesting pixel locations from RGB and point cloud inputs, (2) a next-best pose agent that accepts crops from the Q-attention agent and outputs poses, and (3) a control agent that takes the goal pose and outputs joint actions. We show that current state-of-the-art reinforcement learning algorithms catastrophically fail on a range of RLBench tasks, whilst ARM is successful within a few hours.

Under review as a conference paper at ICLR 2021
ATTENTION-DRIVEN ROBOTIC MANIPULATION
Anonymous authors
Paper under double-blind review
ABSTRACT
Despite the success of reinforeement learning methods, they have yet to have their breakthrough moment when applied lation tasks. This is partly due to the fact that reinforcement learning algorithms are notoriously difficult and time consuming to train, which exacerbated when training from images rather than full-state inputs. As humans perform manip- ulation tasks, our eyes closely monitor every step of the process with our gaze focusing sequentially on the objects being manipulated. With this in mind, we present our Attention-driven Robotic Manipulation (ARM) algorithm, which a general manipulation algorithm that can be applied to range of real-world sparse- rewarded tasks without any prior task knowledge. ARM splits the complex task of manipulation into 3 stage pipeline: (I) Q-attention agent extracts interest- ing pixel locations from RGB and point cloud inputs, (2) next-best pose agent that accepts crops from the Q-attention agent and outputs poses, and (3) control agent that takes the goal pose and outputs joint actions. We show that current state-of-the-arr reinforcement learning algorithms catastrophically fail on a range of RLBench tasks, whilst ARM is successful within a few hours.
broad range of robotic manipu-
INTRODUCTION
Despite their potential, continuous-conte reinforcement learning (RL) algorithms have many flaws: they are notoriously data hungry, often fail with sparse rewards, and struggle with The algorithms for both discrete and continuous RL are almost always evaluated on benchmarks that give shaped rewards (Brockman et al.) 2016 Tassa et al. 2018). privilege that is not feasible for training real-world robotic application across broad range of tasks. Motivated by the observation that humans focus their gaze close to objects being manipulated (Land al., 1999), we propose an Attention-driven Robotic Manipulation (ARM) algorithm that consists of a series of algorithm- agnostic components, that when combined, results in a method that is able to perform range of challenging, sparsely-rewarded manipulation tasks.
tasks.
Our algorithm operates through a pipeline of modules: novel Q-attention module first extracts interesting pixel locations from RGB and point cloud inputs by treating images as an environment, and pixel locations actions. Using the pixel locations crop the RGB and point cloud inputs, comtinuous-control agent that
significantly reducing input size, and feed this to a next-best- outputs 6D poses, which is trained with our novel contidence-aware eritic. These goal poses are then used by a control algorithm that continuously outputs motor velocities.
As is common with sparsely-rewarded tasks, we improve initial exploration through the use of demonstrations. However, rather than simply inserting these directly into the replay buffer, we use a keyframe discovery strategy that chooses interesting keyframes along demonstration trajectories that is fundamental to training our Q-attention module. Rather than storing the transition from an initial state to keyframe state, we use our demo augmentation method which also stores the tran- sition from intermediate points along trajectories to the keyframe states; thus greatly increasing the proportion of initial demo transitions in the replay buffer.
All of these improvements result in an algorithm that starkly outperforms other state-of-the-art meth- ods when evaluated on 10 RLBench (James et al. 2020) tasks (Figure that range in difficulty. To summarise, propose the following contributions: (1) An off off-policy hard attention mechanism that learned via Qeaarnnn rather than the on policy hard attention and soft attention that is commonly seen in the NL. and vision community; (2) A confidence-aware Q function that predicts

---

Under review as a conference paper at ICLR 2021
best the busz
phone on base
pic, and in!
pick up.
PuC rubbashhhhhhhhm tm
reach, target
Xack sine
Sake Jrd olf soucopan
sake, tollt, roll stand
lake umbroll out, of umbrolla stand
Figure I: The RLBench tasks used for evaluation. Current state-of-the-an aht reinforcement learning algorithms catastrophically fail on all tasks, whilst our method succeeds within modest number of steps. Note that the positions of objects are placed randomly at the beginning of each episode.
pixel-wise values and confidence values, resulting in improved convergence times; (3) keyframe discovery strategy and demo augmentation method that go hand-in-hand to improve the utilisation of demonstrations in RL.
RELATED WORK
The of reinforcement learning (RL.) is prevalent in many areas of robotics, including legged
robots (Kohl & Stone 2004 Hwannho et al.] (019), aerial vehicles (Sadeghi & Levine 2017).
and manipulation tasks, such as pushing (Finn& Levine 2017). peg insertion (Levine
2016
Zeng et al. 2018 Lee et al. 2019), throwing (Ghadirzadeh et al. 2017 Zeng et al. 2020). ball-
in-cup (Kober Peters 2009), cloth manipulation (Matas al. 2018). and grasping KKllshhikkov et al. 2018 James et al. 2019b). Despite the abundance of work in this area, there has yet to
be general manipulation method that can tackle a range of challenging, sparsely-rewarded tasks without needing access to privileged simulation-only abilities (e.g. reset to demonstrations (Nair jet al., 2018), asymmetric actor-critic (Pinto al., 2018). reward shaping (Rajeswaran et al. 2018), and auxiliary tasks (James et al. 2017)).
Crucial to our method the proposed Q-attention. Soft and hard attention are prominent methods in both natural language processing (NLP) (Bahdanau et al., 2015 Vaswani et al. 2017 Devlin et al.) 2018) and computer vision (Xu et al. multiplies attention map over the image feature map, whilst hard attention uses the attention map stochastically to sample one or a few features on the feature map (which is optimised by maximis- ing approximate variational lower bound equivalently via (on-policy) REINFORCE (Williams) 1992)). Given that we perform non non-differentiable cropping, our Q-attention is closest to hard atten- tion. but with the distinction that we learn this in an off-policy way. This is key, as 'traditional' hard attention unable to be used in an off-policy setting. We therefore see Q-attention as an off-policy hard attention. We elaborate further on these differences in Section
et al.) 2019). Soft attention deterministically
Our proposed confidence-awand critic (used to train the next-best pose agent) takes its inspiration from the pose estimation community (Wang et al., 2019 Wada et al. 2020). There exists small amount of work in estimating uncertainty with Q-learning in discrete domains (Clements et al.) 2019 Hoel et al. (2000;; our work uses continuous Q-function to predict both Q and confidence values for each pixel, which lead to improved stability when training, and is not used during action selection.
Our approach makes use of demonstrations, which has been applied in number of works (Vecerik et al. 2017 Matas et al. 2018 Kalashnikov et al., 2018 Nair et al.] 20188). but while successful, they make limited use of the demonstrations and still take many samples converge. Rather than simply inserting these directly into the replay buffer, we instead make sure our keyframe discovery and demo augmentation maximise demonstration utility.
2

---

Under review as a conference paper at ICLR 2021
Ounr S4h 31
KINA
Dow
Corvili 125 H
Cenei 26,
Ned Cest Pose
(128xx17))
Cendt sob
Seviih 20c
(16-30)
DaNvOOL S0, 3
Srt 11
8 Up
Careit
Com
Pen Placl Qovllues
(x,y)
propeSecrpete
(16x1662)
Figure 2: Summary and architecture of our method. RGB and organised point cloud crops made by extracting pixel locations from our Q-attention module. These crops are then fed to a continuous control RL algorithm that suggests next-best poses that is trained with confidence-awe The next best pose is given to goal-condition control agent that outputs joint velocities. Conv block represented Conv(#channels, filter size, strides).
critic.
BACKGROUND
The reinforcement learning paradigm assumes agent interacting with environment consisting of states $ E S, actions a E A, and a reward function R(s,, a,), where S, and a, are the state and action time step respectively. The goal of the agent is then discover policy that results in maximising the expectation of the sum of discounted rewards: Ex |Σ, 7 R(s,, at)]. where future rewards are weighted with respect the discount factor ? E [0. 1). Each policy ı has a corresponding value function ((s, a), which represents the expected return when following the policy after taking action in state s.
Our Q-attention module builds from Deep Q-learning (Mnih et al. 2015). a method that approx- imated the value function Qo. with a deep convolutional network, whose parameters 0 are op- timised by sampling mini-batches from replay buffer D and using stochastic gradient descent to minimise the loss: a++t))) + yymar, (8+ a') Qo(s,,
where Q" is target network; periodic copy of the online network Qa which is not directly optimised. Our next-best pose agent builds upon SAC (Haarnoja et al. 2018), however, the agent compati- ble with any off-policy, continuous-control RL algorithm. SAC is maximum entropy RL algo- rithm that, in addition to maximising the sum rewards, also maximises the entropy of a policy: E,, E,, (R(s,, a (-(((((s:))) [+]]] where is temperature parameter that determines the relative importance between the entropy and reward. The goal then becomes to maximise a soft Q-function by minimising the following Bellman residual:
(Q)
E (s. (s.
((r + (((( a log logss(ay)).))
Qo(s,, a1))
(I)
The policy is updated towards the Boltzmann policy with temperature a, with the Q-function taking the role of (negative) energy. Specifically, the goal is to minimise the Kullback-Leibler divergence between the policy and the Boltzman policy:
arg min DKL ( (-(ss) a exp TQT (st, a (st)
||
T new
E'Ell
Minimising the expected KL-divergence to learn the policy parameters was shown to be equivalent to maximising the expected value of the soft Qfunction:
(+)
E Beru awas
(s, at)]]-
(st,
4 METHOD
Our method can be split into a 3-phase pipeline. Phase (Section 4. consists of a high-level pixel agent that selects areas of interest using novel Q-attention module. Phase (Section 4.2 consists
3

---

Under review as a conference paper at ICLR 2021
pick up cup
take toilet roll off stand
beat
put rubbish
take lid off saucepan
stack wine
Figure 3: Visualising the Q values across 4 different points in time on 6 tasks. At each step, RGB and organised point cloud crops are made by extracting pixel locations that have the highest Q-value. We can see that as time progresses, the attention strength shifts depending progress in the task; e.g. 'stack wine' starts with high attention on the bottle, but after grasping. attention shifts to the wine rack.
of next-best pose prediction phase where the pixel location from the previous phase is used crop the incoming observations and then predict 6D pose. Finally, phase 3 (Section 4.3) is a low-level control agent that accepts the predicted next-best pose and executes a series of actions to reach the given goal. Before training, we fill the replay buffer with demonstrations using our keyframe discovery and demo augmentation strategy (Section (4.4) that significantly improves training speed. The full pipeline is summarised in Figure and Algorithm
experiments are run RLBench (Jaees et al. 2020), large-scale benchmark and learning envi- ronment for vision-guided manipulation built around CoppeliaSim (Rohmer et al. (013) and PyRep (James et al. 2019a). At each time step, we extract an observation from the front-facing camera that consists of an RGB image b and depth image d, along with proprioceptive information from end-effector pose and gripper open/close state). Using known camera intrin-
the arm (consisting sics and extrinsics, process each depth image produce point cloud p (in world coordinates) projected from the view of the front-facing camera, producing a (H W x 3) 'image'.
4.1 Q-ATTENTION
Motivated by the role of vision and eye movement in the control of human activities LLand al. 1999). we propose a Q-attention module that, given RGB and organised point cloud inputs, outputs 2D pixel locations of the next area of interest. With these pixel locations, we crop the RGB and ganised point cloud inputs and thus drastically reduce the input size to the next stage of the pipeline. Our Q-attention is explicitly learned via Q-learning, where images are treated the 'environment', and pixel locations are treated as the 'actions'.
Given our Q-attention function Q Ag, we extract the coordinates of pixels with the highest value:
(x,, yr) argmax2D (s,
The parameters of the Q-attention are optimised by using stochastic gradient descent to minimise the loss:
(Q) at.
E (r max2D (Ag QAg(s,, az))² 2 + [[QA]]]
(5)
where s = (b, p), QAg the target fuunction, and [|QA[] an L2 loss on the per- pixel output of the function (which we call Q regularisation); in practice, found that this leads to increased
4

---

Under review as a conference paper at ICLR 2021
robustness against the common problem of overestimation of Q values. The Q-attention network follows a light-weight U-Net style architecture (Ronneberger et al. 2013). which is summarised in Figure 2 Example per-pixel outputs of the Q-attention are shown Figure 3 With the suggested coordinates from Q-attention, we perform a (16 x 16) crop on both the (128 x 128) RGB and organised point cloud data: b', b',p' = crop(b, p. (x)
Notably, there is no explicit reward for choosing a pixel, but instead an implicit reward that comes from the output of the method pipeline as whole (i.e. the same reward signal is used to train both the Q-attention and the next-best pose agent). This leads to a cyclic dependency between the two agents: the lower-level next-best pose agent relies on receiving good crops from the Q-attention agent, whilst the Q-attention agent needs the next-best pose agent to perform well in order to get its implicit reward. This is where delicate handling of demonstrations is key, which we discuss in Section 4.4
The module shares similar humann-inspiree motivation to the attention seen in NLP (Bahdanau al. 2015 Vaswani et al. 2017: Devlin et al.) 2018) and computer vision (Xu et al. 2015 Zhang et al., 2019). but differs in its formulation. Soft attention multiplies an attention map over the image feature map, whilst hard attention uses the attention map to sample one or a few features on the feature maps or inputs. Given that we perform non-differentiable cropping, we categorise our Q-attention as hard attention, but with 2 core differences: (I) most importantly 'traditional' hard attention is optimised (on-policy) by maximising an approximate variational lower bound or equivalently via REINFORCE (Williams 1992). whereas our Q-attention is trained off-policy; this is erucial because our demonstration data, by definition, is off-policy, and therefore renders existing hard attention approaches unusable for demonstration-driven RL. (2) The output of 'traditional' hard attention carry different semantics: a score function in the case of REINFORCE hard attention, and Q-value (expected cumulative reward of choosing that crop) in the case of D-attention.
4.2 NEXT-BEST POSE AGENT
Our next-best pose agent accepts cropped RGB b' and organised point cloud p' inputs, and outputs 6D pose. This next-best best pose agent is run every time the robot reaches the previously selected pose. We represent the 6D pose via translation output of q to positive number, therefore restricting the network to output unique unit quaternions. The gripper action h E 'R1 lies between 0 and 1, which is then discretised to binary open/close value. The combined action therefore is a = {e. q, h).
R2 and unit quaternion q R*, and restrict the
To train this next-best pose agent, we use modified version of SAC (Haarnoja al. 2018) where we modify the soft Q-function (Equation] I to be confidence-aware soft Q-function. Recent work in 6D pose estimation (Wang et al. 2019, Wada al. 2020) has seen the inclusion of a confidence score e with the pose prediction output for each dense- jense-pixel. Inspired by this, we augment our function with a per-pixel confidence Cij. where we output a confidence score for each Q-value prediction (resulting in a (16 x 16 x 2) output). To achieve this, we weight the per-pixel Bellman loss with the per-pixel confidence, and add a confidence regularisation term:
JQz (p) =
(e [(I (s41) (541) - log (ails:)) - (, an)) w log(
E
where s = (b',p',z) With this, low confidence will result in low Bellman error but would incur high penalty from the second term, and vice versa. We use the value that has the highest confidence when training the actor. As an aside, we also tried applying this confidence-aware method to the policy, though empirically we found no improvement. In practice we make use of the clipped double-Q trick (Fujimoto et al. 2018), which takes the minimum Q-value between two Q networks, but have omitted in the equations for brevity. Finally, the actor's policy parameters can optimised by minimising the loss as defined in Equation 3
4.3 CONTROL AGENT
Given the best-best pose suggestion from the previous stage, we give this to a goal-conditioned control function f(s;, gt), which given state s, and goal gi, outputs motor velocities that drives the
5

---

Under review as a conference paper at ICLR 2021
Algorithm 1 ARM
Initialise Q-attention networks QA QAoe, critic networks 2% O72 and actor network T with random parameters 8, 02. P1. P2. Øb Initialise target networks 02 6- 01. 022 O2 p/1 + Pl. P/2 +- P2
Initialise replay buffer with demos and apply keyframe selection and demo augmentation for each iteration do for each environment step do
(b,, Pt, 2,) St (01, m1) < argmax2D, n' Q p(), by, P/ crop(b,, Pt. (Izz v6))
bi Use Q-attention to get pixel coords
a (a:aa(bb, . P/, 2r))
D Sample pose from the policy
while target not reached do
V t f()
Get joint velocities from control agent
St+1, env.step(v)
D DUDU ((s,, a. S (1x, 9c))
>> Store the transition in the replay pool
for each gradient step do
0g 0; AQa Vo, (QA( for i E (1,2)
D Update Q-attention parameters
Pi - Pt Ag= V Jor (p:) for E {1,2)
Di Update critie parameters
@ - Φ d V.Va (
> Update policy weights
0 r0,
7)
T 3 for E {1,2)
Update Q-attention target network weights >> Update critic target network weights
P t TP: + (1 T)p for i {1,2}
end-effector towards the goal. This function can take many forms, but two noteworthy solutions would be either motion planning combination with a (eedback-contie or a learnable policy trained with imitation/reinforcmmenn learning. Given that the environmental dynamics limited in the benchmark, we opted for the motion planning solution.
Given the target pose, we perform path planning using the SBL (Sánchez & Latombe 2003 planner within OMPL (Sucan et al. 2012). and use Reflexxes Motion Library for on-line trajec- tory generation. If the target pose is out of reach, terminate the episode and supply ward -1. This path planning and trajectory generation is conveniently encapsulated by the 'ABS EE POSE_PLAN WWWLD FRAME' action mode in RLBench (James et al. 20220).
4.4 KEYFRAME SELECTION & DEMO AUGMENTATION
In this section, we outline how we maximise the utility of given demonstrations in order to com- plete sparsely reward tasks. We assume to have teacher policy x² (c.g. motion planners hu- man telcoperatives) that can generate trajectories consisting of series of states and actions: [(s,, ..)....,(s7) ar), In this case, we assume that the demonstrations come from RLBench (James) et al. 2020).
The keyframe selection process iterates over each of the demo trajectories and runs each of the state- action pairs as through a function K E which outputs a boolean deciding if the given tra- jectory point should be treated as a keyframe. The keyframe function K could include number of con- straints. In practice we found that performing a dis- junction over two simple conditions worked well; these included (1) change in gripper state (a com- mon occurrence when something is grasped or released), and (2) velocities approaching near (a common occurrence when entering pre-grasp poses or entering a new phase of a task). is likely that tasks get more complex, K will inevitably need become more sophisticated via learning
Figure 4: Keyframe selection and demo aug- mentation, where the black line represents trajectory, '!' represents keyframes, and dashed blue lines represent the augmented transitions to the keyframes.

---

Under review as a conference paper at ICLR 2021
take lid_ off saucepan
put_rubbish in_ bin
pick_up upuup___u
take Like__follet__eee _ool_ooffooffft __stand
h=1
I=2
MR3
1m4 ko5
Figure 5:
Visualising RGB observations of keyframes from the keyframe selection process on tasks. Here k is the keyframe number.
or simply through more conditions, e.g. sudden changes in direction or joint velocity, large changes in pixel values, etc. Figure shows RGB observations from the keyframe selection process from tasks.
At each keyframe, we use the known camera intrinsics and extrinsics to project the end-effector pose at state S2+1 into the image plane of state st, giving us pixel locations of the end-effector at the next keyframe. This stage crucial to breaking the cyclic dependency (mentioned Section 4.1) between the Q-attention and next-best pose agent, as these projected pixel coordinates act optimal actions for the Q-attention agent.
transi- Using this keyframe selection method, each trajectory results in N length tions being stored into the replay buffer. To further increase the utility of demonstrations, we apply demo augmentation which stores the transition from intermediate point along the trajectories to the keyframe states. Formally, for each point (s,, az) along the trajectory starting from keyframe k3, we calculate the transformation of the end-effector pose (taken from s₂) at time step t to the end-effector pose at the time step associated with keyframe ki+i- This transformation can then be used the action for the next-best pose agent. We repeat this process for every Mth point along the trajectory (which we set to M = 5). The keyframe selection and demo augmentation is visualised in Figure
RESULTS
In this section, we aim to answer the following questions: (1) Are we able to successfully learn across a range of sparsely-rewarded manipulation tasks? (2) Which of our proposed components contribute the most to our success? (3) How sensitive is our method to the number of demonstrations and to the crop size? To answer these, we benchmark our approach using RLBench (James et al.) 2020). Of the 100 tasks, we select 10 (shown in Figure I that believe to be achievable from using only the front-facing camera. We leave tasks that require multiple cameras to future work. RLBench was chosen due to its emphasis it gives access to a wide variety of tasks with demonstrations. Each task has a completely sparse reward of +1 which is given only on task completion, and 0 otherwise.
vision-based manipulation benchmarking and because
The first of our questions can be answered by attending to Figure All baseline algorithms (SAC, TD3 and QT-Opt) in their 'vanilla' form, and do not contain any of our proposed contributions: Q-attention, confidence-aware critic, keyframe selection, and demo augmentation. All methods receive the exact same 200 demonstration sequences, which are loaded into the replay buffer prior to training. The baseline agents are architecturally similar to the next-best pose agent, but with few differences to account for missing Q-attention (and so receives the full, uncropped RGB and organised point cloud data) and missing confidence-aware critic (and so outputs single Q-values rather than per-pixel values). Specifically, the architecture uses the same RGB and point cloud fusion as shown in Figure 2 Feature maps from the shared representation are concatenated with the reshaped proprioceptive input and fed to both the actor and critic. The baseline actor uses convolution layers (64 channels, 3 x 3 filter size, 2 stride), who's 's output feature maps maxpooled and sent through 2 dense layers (64 nodes) and results in an action distribution output. The critic baseline uses 3 residual convolution blocks (128 channels, 3 x filter size, 2 stride), who's output
7

---

Under review as a conference paper at ICLR 2021
beat, the
phone on, Dase
peck
PICK WP. Sup
put, www.buubbbh.....
prigee
marth W W
oll Mor M
M Mp yorr
Sapp
Saepe
Steps
Dlrs
SAC 1D3
Q1-Cpt
Figure 6: Learning curves for 10 RLBench tasks. Methods include Ours (ARM), SAC (Haarnoja al. 2018). TD3 (Fujimoto et al., 2018), and QT-Opt (Kalashnikov al. 2018). ARM uses the 3-stage pipeline (Q-attention, next-best pose, and control agent), while baselines use the 2-stage pipeline (next-best pose and control agent). All methods receive 200 demos which are stored in the replay buffer prior to training. Solid lines represent the average evaluation over 5 seeds, where the shaded regions represent the min and max values across those trials.
feature maps are maxpooled and sent through 2 dense layers (64 nodes) and results in single Q- value output. All methods use the LeakyReI.U activation, layer normalisation in the convolution layers, learning rate of 3 x 10-3 soft target update of 5 5-4, and reward scaling of 100. Training and exploration were done asynchronously with a single agent (to emulate a real-world robot training scenario) that would continuously load checkpoints every 100 training steps.
The results in Figure 6 show that baseline state-of-the-art methods are unable to accomplish any RLBench tasks, whilst our method is able accomplish the tasks in small number of environ- ment steps; 5, 000 environment steps equating to about an hour of robot interaction time (meaning "take lid off saucepan' being solved in about two hours). We suggest that the reason why our results starkly outperform other ofate-of-the-art methods is because of two key reasons that go hand-in- hand: (1) Reducing the input dimensionality through Q-attention that immensely reduces the burden on the (often difficult and unstable train) continuous control algorithm; (2) Combining this with our keyframe selection method that enables the Q-attention network to quickly converge and sug- gest meaningful points of interest to the bext-best pose agent. We wish to stress that perhaps given enough training time some of these baseline methods may eventually start to succeed, however we found evidence of this.
In Figure 7a we perform an ablation study to evaluate which of the proposed components con- tribute the most the success of our method. To perform this ablation, we chose 2 tasks of varying difficulty: 'take loike_lddofff_see 'iddceffsaucepan" and 'put-rubbish. iin bnn'.'. The ablation clearly shows that the Q-attention (combined with keyframe selection) is crucial achieving the tasks, whilst the demo augmentation, confidence -wwae aware critic, and Q regularisation aid in overall stability and increase final performance. When swapping the Q-attention module with a soft attention Xu et al. (2015) module, we found that performance was similar to that of the 'vanilla' baselines. This result is unsurprising, as soft attention is implicitly learned (i.e. without an explicit loss), where as our Q-attention ex- plicitly learned via (off- (off-policy) Q-learning, and it make greater use of the -ighly-ifformative keyframes given from the keyframe selection process. Note that we cannot compare to 'traditional' hard attention because it requires on-policy training, explained in Section 4.1
Figure 7b shows how robust our method when varying the number of demonstrations given. The results show that our method performs robustly, when given 50% fewer demos, however the task difficulty increase (from 'take lid off saucepan' to 'put rubbish in bin )) the harmful effect of having less demonstrations is more severe. Our final set of experiments in Figure shows how
8

---

Under review as a conference paper at ICLR 2021
take sescupen
takeL Jud, off anacepan
caocepan
(a) Effect of removing components (b) Effect of number of demos (c) Effect of crop size perfor- from our method.
performance.
mance.
Figure 7: rput.rubbish.n task.
Ablation study
across the easier 'take_ Loff saucepan" task and harder
method performs across varying crop sizes. As the task difficulty increases, the harmful effect of larger crop size becomes more prominent; suggesting that one of the key benefits of the Q-attention drastically reduce the input size the next-best pose agent, making the RL optimisation much easier. It's clear that trade-off must be made between choosing smaller crops to increase training size, and choosing larger crops to incorporate more of the surrounding area. We found that setting the crops at 16 x 16 across all tasks performed well.
6 CONCLUSION
We have presented our Attention-driven Robotic Manipulation (ARM) algorithm, which is general manipulation algorithm that be applied to range of real-world sparsely-rewarded tasks. We validate our method on 10 RLBench tasks of varying difficulty, and show that many commonly used -tate-of-the-art methods catastrophically fail. We show that Q-attention (along with the keyframe selection) is key to our success, whilst the confidence-awarn critic and demo augmentation contribute achieving high final performance. Despite our strong experimental results, there are undoubtedly areas of weakness. The control agent (final agent in the pipeline) uses path planning and on-line trajectory generation, which for these tasks are adequate; however, this would need be replaced with an alternative agent for tasks that have dynamic environments (e.g. moving target objects, moving obstacles, etc) or complex contact dynamics (e.g. peg-in-hole). We look future work for swapping this with goal-conditioned reinforcement learning policy, or similar. Another weakness that we only evaluate on tasks that can be done with the front-facing camera; however we are keen to explore many of the other tasks RLBench has to offer by adapting the method to accommodate multiple camera inputs in future work.
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. Intl, Conference on Learning Representations, 2015.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv:1606.011400. 2016.
William R Clements, Benoît-Marie Robaglia, Bastien Van Delft, Reda Bahi Slaoui, and Sébastien Toth. Estimating risk and uncertainty in deep reinforcement learning. arXiv preprint /oXiv:9905.09638, 2019.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 18 2018.

---

Under review as a conference paper at ICLR 2021
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In IEEE Intl. Conference on Robotics and Automation, pp.
IEEE, 2017.
Scott Fujimoto, Herke Van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. Inil. Conference on Machine Learning. 2018.
Ali Ghadirzadeh, Atsuto Maki, Danica Kragic, and Marten Björkman. Deep predictive policy train- ing using reinforcement learning. In IEEE Intl. Conference Intelligent Robots and Systems, pp. 2351-2358 IEEE, 2017.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and appli- cations. arXiv preprint arXiv:1812.0500 2018.
Carl-Johan Hoel, Krister Wolff, and Leo Laine. Tactical decision-making autonomous driving by reinforcement learning with uncertainty estimation. arXiv preprint arXiv:2004. 10439, 2020.
Jemin Hwangbo, Joonho Lee, Alexey Dosovitskiy, Dario Bellicoso, Vassilios Tsounis, Vladlen Koltun, and Marco Hutter. Learning agile and dynamic motor skills for legged robots. Science Roboties, 4(26), 2019.
Stephen James, Andrew Davison, and Edward Johns. Transferring end-to-end visuomotor control from simulation real world for multi-stage task. Conference Robot Learning, 2017.
Stephen James, Mare Freese, and Andrew J Davison. Pyrep: Bringing v-rep to deep robot learning. arXiv preprint arXiv:1906. 11176, 2019a.
Stephen James, Paul Wohlhart, Mrinal Kalakrishnan, Dmitry Kalashnikov. Alex Irpan, Julian Ibarz, Sergey Levine, Raia Hadsell, and Konstantinos Bousmalis. Sim-to-real via sim-to-sim: Data- efficient robotic grasping via Computer Vision and Pattern Recognition, pp. 12627-12637, 2019b.
adaptation networks. In IEEE Conference
Stephen James, Zicong Ma, David Rovick Arrojo, and Andrew Davison. RLBench: The robot learning benchmark & learning environment. IEEE Robotics and Automation Letters, 2020.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation. Conference Robot Learning, 2018.
Jens Kober and Jan R Peters. Policy search for motor primitives in robotics. In Advances in Neural Information Processing Systems, pp. 849-856, 2009.
Nate Kohl and Peter Stone. Machine learning for fast quadrupedal locomotion. In Association for the Advancement of Artificial Intelligence, volume pp. 61 611-616, 2004.
Michael Land, Neil Mennie, and Jennifer Rusted. The roles of vision and eye movements in the control of activities of daily living. Perception, 28(1 28(11):1311) 11:1311-13222 1999.
Michelle Lee, Yuke Zhu, Krishman Srinivasan, Parth Shah, Silvio Savarese, Li Fei-Fei, Ani- mesh Garg, and Jeannette Bohg. Making sense of vision and touch: Self-supervised learning of multimodal representations for contact-rich tasks. In IEEE Intl. Conference on Robotics and Automation, pp. 19943-8950 IEEE, 2019.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. Journal of Machine Learning Research, 17(1): 2016.
Jan Matas, Stephen James, and Andrew J Davison. Sim-to-real reinforcement learning for de- formable object manipulation. Conference on Robot Learning, 2018
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 518(7500):59-5333 2015.
10

---

Under review as a conference paper at ICLR 2021
Ashvin Nair, Bob McGrew, Marcin Andrychowicz. Wojeiech Zaremba, and Pieter Abbeel. Over- coming exploration in reinforcement learning with demonstrations. In IEEE Intl. Conference on Robotics and Automation, pp. 2292-6299. IEEE, 2018.
Lerrel Pinto, Marcin Andrychowicz, Peter Welinder, Wojciech Zaremba, and Pieter Abbeel. Asym- metric actor critic for image-based robot learning. Robotics: Science and Systems, 2018.
Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. Robotics: Science and Systems, 2018.
Eric Rohmer, Surya PN Singh, and Mare Freese. V-rep: A versatile and scalable robot simulation framework. In IEEE Intl. Conference on Intelligent Robots and Systems. IEEE, 2013.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. Springer, 2015.
Fereshteh Sadeghi and Sergey Levine. Cad2rl: Real single-image flight without a single real image. Robotics: Science and Systems, 2017.
Gildardo Sánchez and Jean-Claude Latombe. A single-query bi-directional probabilistic roadmap planner with lazy collision checking. In Robotics research, pp. 403-417. Springer, 2003.
loan A. Sucan, Mark Moll, and Lydia E. Kavraki. The Open Motion Planning Library. IEEE doi: 10.1 0..100//.222012
Robotics & Automation Magazine, 19(4):22 82, December 2012.
2205651. https::/ompl, www.:mppllaaaaalalaammamm.m waa.kiilabboo.. org
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, al. Deepmind control suite. arXiv preprint arXiv:1801.0060) 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, pp. 5998-6008, 2017.
Mel Vecerik, Todd Hester, Jonathan Scholz, Fumin Wang. Olivier Pietquin, Bilal Piot, Nico- las Heess, Thomas Rothörl, Thomas Lampe, and Martin Riedmiller. Leveraging demonstra- tions for deep reinforcement learning on robotics problems with sparse rewards. arXiv preprint arXiv:1707.08817, 2017.
Kentaro Wada, Edgar Sucar, Stephen James, Daniel Lenton, and Andrew J Davison. Morefusion: Multi-object reasoning for 6d pose estimation from volumetric fusion. In IEEE Conference on Computer Vision and Pattern Recognition, pp. 14540-14549, 2020.
Chen Wang, Danfei Xu, Yuke Zhu, Roberto Martín-Martin Cewu Lu, Li Fei-Fei, and Silvio Savarese. Densefusion: 6d object pose estimation by iterative dense fusion. In IEEE Confer- ence on Computer Vision and Pattern Recognition, pp. 3343-3352, 2019..
Ronald J Williams. Simple statistical gralient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229-22 256, 1992.
Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In Intl. Conference on Machine Learning, pp. 2048-2057, 2015.
Andy Zeng, Shuran Song, Stefan Welker, Johnny Lee, Alberto Rodriguez, and Thomas Funkhouser. Learning synergies between pushing and grasping with self-supervised deep reinforcement learn- ing. In IEEE Intl. Conference on Intelligent Robots and Systems, pp. 4238 4245. IEEE, 2018.
Andy Zeng, Shuran Song, Johnny Lee, Alberto Rodriguez, and Thomas Funkhouser. Tossinghot: Learning to throw arbitrary objects with residual physics. IEEE Transactions on Robotics, 2020.
Han Zhang, Ian Goodfellow, Dimitris Metaxas, and Augustus Odena. Self-attention generative adversarial networks. In Inil. Conference on Machine Learning, pp. 7354-7363, 2019.


Reviews:
[{"review_id": "yGu6FXTgVkP", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \n\nThis work focuses on sparse-reward robotic manipulation tasks from image and point cloud inputs, given a few demonstrations, and proposes an algorithm that consists of a Q-attention module and a confidence-aware critic. The Q-attention module is an RL agent, which takes image and point cloud inputs with pixel positions as actions. The image and point cloud observations are then cropped, centered at this pixel position output, and passed down the pipeline to the critic. To make the best use of the demonstrations, the replay buffer is initialized with the demonstrations and further augmented with transitions from demo states, chosen through a keyframe selection process. \n\n\n\nStrengths: \n\n- The results demonstrate how, in the difficult setting of learning from sparse rewards and image inputs, only the proposed approach can learn meaningful behavior while the other RL approaches achieve basically zero return. Though, it would be good to include the returns of the expert agent in the plot as well to show the upper bound performance.\n- Through an ablation study, it’s clear that the Q-attention agent is important for its ability to reduce the input dimension to the continuous control agent. The demo augmentation process is also quite critical to ensure that the Q-attention and continuous control agents can correctly learn together, since the success of one agent relies on the success of the other agent. \n- The paper is overall well-organized and written, and the figures and illustrations are helpful for the reader’s understanding.\n\n\n\nWeaknesses:\n\n- It is not completely clear what advantages the Q-attention agent brings over standard attention modules. While a discussion is included, a plot or some other form of quantitative comparison would be more illustrative. \n- It’s also not clear why an image crop is the best choice of input for the control agent. For instance, an alternative might be the pixel position itself. This is also highly related to keypoint representations, and there isn’t a discussion of or comparison to them.\n- There’s also a key assumption made here that only a fixed region of the image input is relevant to the task at any given timestep. I’m not sure if this is entirely reasonable for environments with different camera viewpoints and different sized objects to manipulate.\n- The keyframe selection is based on heuristics: it looks at frames where (1) the gripper state changes as it corresponds to the agent grasping or releasing an object and (2) the agent velocity is near zero as it often corresponds to a new stage of the task.\n- The choice of baselines also seems quite arbitrary.\n\n\n\nRecommendation:\n\nI am recommending to reject this paper. There are a number of choices in this approach that could be further justified, especially with regards to the Q-attention agent.\n\n\n\nQuestions:\n\n- Is the pixel position outputted by the Q-attention agent also passed into the continuous control agent? If not, could doing so speed up learning since the cropped image doesn’t provide information about where objects are?\n\n- Instead of a cropped image, is it sufficient to just pass the pixel position, or perhaps the top k pixel positions, given by the Q-attention agent to the continuous control agent? On that note, this is highly related to keypoint representations and it would be good to include a discussion of this in the related work section. \n\n- How well does the Q-attention agent scale to higher-dimensional image inputs? Additionally, how does the crop size affect performance and efficiency of the continuous control agent? That is, with larger objects and a different viewpoint, a much larger crop may be necessary to capture enough context to solve the task. \n\n\nMinor Comments:\n\n-The labels of Fig. 6 need to be bigger.\n\n\n\nUpdate after Reading Authors' Response:\n\nI appreciate the authors’ thorough response and the new results for different crop sizes. It’s not too surprising that the performance varies with the crop size and that the best crop size depends on the task and environment (among other variables such as camera resolution, camera viewpoint, and object sizes). And so, I share the concern of the other reviewers: it’s unclear whether this approach would be similarly successful in other manipulation settings without prior knowledge of the task outside of RLBench.", "extra_scores": {}, "extra_text": {}}, {"review_id": "ptOr0bqsw23", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper tackles the problem of training a robotic manipulation model, starting with a set of examples of successful expert trajectories. Training starts with a select set of samples from the expert episodes as the seed of an experience buffer and proceeds to generate new episodes to train on through an RL procedure.\n\nThe robotic manipulation model model first applies an attention mechanism in the form of choosing an area of the input images to crop around (though use of a cropping Q function). It then selects the next desired robotic pose from the cropped input using a version of SAC that uses a confidence-augmented Q-function loss and finally applies a motion control mechanism to translate the desired end effector pose into robotic actuation in the simulated setting. \n\nThe method makes use of some assumptions regarding what would be the important states and makes use of a selection mechanism to populate an experience buffer. Specifically, it assumes that points in time where the gripper is actuated or velocities approach zero represent necessary steps (\"keyframes\") in the completion of a task. It then generates examples where the desired robotic action for a state along an example trajectory is the end effector pose of the next keyframe and the desired crop is around the projected location of the desired end-effector pose in the camera image.\n\nResults show that whereas vanilla implementations of SAC, TD3 and QT-Opt with a standard replay buffer can not learn the tasks, the proposed method succeeds. A further ablation study shows the crucialness of the attention-driven cropping step in this setting and the importance of the sample selection strategy and the confidence layer of the SAC Q function.\n\nI find the combination of these applied characteristics of the method to be sensible and powerful:\n - Starting with expert-provided trajectories (and breaking them down into steps / example-poses using the proposed example sampling method)\n - Focusing of the input signal towards the region of interest (bootstrapping with end effector locations) as a better state representation.\n - Using high level control as the action space (next desired pose as opposed to immediate joint actuations)\n\nThe method is presented as a general manipulation algorithm. While it does seem to generalize to a rich set of examples in RLBench, it makes use of some properties that are limited to a subset of manipulation tasks. Specifically, tasks where the important activity is limited to a narrow region of the observed field of view and is comprised of steps that can be treated as series of intermediate goals (poses in which velocities are small or the gripper is actuated). It would have been interesting to see whether there are RLBench tasks that should be solvable using a single camera view on which this method fails and to understand the failure reasons.\n\nThe methods compared against feel like straw-men. Surely there must be some off-the-shelf methods targeted at learning from expert examples that do not completely fail on these tasks?\n\nThere are multiple methods these days that start by transforming the observations into task-oriented state representations (certainly a focus of interest for ICLR). It would have been interesting to see comparisons with alternatives to image crop as a representation (https://arxiv.org/abs/2007.07170 presents one such example).\n\nThe paper does not provide code to replicate the experiments. Seeing as it aims to address a known application, this is unfortunate as it would help clarify implementation details and support further research.\n\nNits:\n\nI found the representation of the point cloud unclear. It seems like it is treated as another RGB from a different view (e.g. cropped to 16x16 pixels?) but if that is true, why call it a point cloud?\n\nThe use of the per-pixel Q values (the downstream Q function that has confidence values) is a bit unclear. It might be a no-brainer to someone more familiar with SAC but I found its use to be unclear just from reading this paper. Same is true for the use of 2x Q-functions (referred to as \"the double Q trick\").\n\nIt was unclear to me whether during use, the next iteration of pose selection is performed once the robot has reached the previously selected pose (though the low-level motion control agent) or whether the pose selection occurs at every simulation step regardless of distance from the previously generated pose.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Fj_dDx8nPAb", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nMotivated by the fact that RL algorithms are notoriously sample inefficient from pixels and that humans use attention, the authors propose Attention-driven Robotic Manipulation (ARM), which they claim can be applied to several robotics tasks without prior task knowledge. Compared to current methods which fail to train, ARM trains in only a few hours and is successful at the end manipulation tasks. The authors’ main algorithmic improvements with ARM include 1) a Q-attention agent that extracts interesting pixel locations with an explicit attention module, 2) a confidence-aware critic that leads to increased stability of training, and 3) a data augmentation strategy around expert demonstrations. \n\n---\n\n### Strengths\n\n- The paper is well motivated, and the authors recognize a salient problem with most RL systems failing with sparse rewards and being sample inefficient.\n- The authors propose several potentially novel contributions, including their Q-attention agent with an explicit attention module, the confidence-aware critic that improves performance, the keyframe discovery strategy, and the demo augmentation method.\n- The authors’ demo augmentation strategy is clever in that the required pose transformations are easy to collect and calculate, and this strategy clearly helps performance.\n- The authors’ experiments clearly answer the questions they posed around how well this technique learns and how sensitive it is to the number of expert demonstrations. In addition their ablation experiments around which components contribute to the overall pipeline’s success are useful.\n- The experimental results show a clear improvement over prior techniques, working with fewer demonstrations and fewer environment steps. The authors’ intuitive explanations for why their method works when others fail are also well-reasoned.\n\n---\n\n### Weaknesses\n\n- The authors claim that their technique can be applied to several tasks without prior task knowledge. This claim only seems weakly supported in the paper. Specifically, to make this claim I would expect to see results across multiple robot platforms. The keyframe selection and demo augmentation strategy does seem specific to the current platform and set of tasks and may not work as well in stochastic environments.\n- The authors replace the hand-engineering of reward functions with several hand-engineered modules, as such the theoretical motivation for these modules is less strong. Although I expected to see each module ablated and tested in isolation, these results were lacking in the experiments.\n- For the Q-attention agent, the authors claim that a L2 loss prevents overestimation of the Q-values. Because the authors claim that this results was achieved experimentally, such an ablation or analysis should be present in the experiments or supplementary figures.\n- The motivation for the Q-attention agent is not made immediately clear. Specifically, the Q-attention agent seems rather similar to performing hard attention with REINFORCE or other policy gradient methods, except here the hard attention is performed with Q-learning. As such, I’m unsure of the actual novelty of this technique.\n- The next-best pose agent is particularly confusing since it mentions per-pixel Q-values and per-pixel Bellman losses. This formulation makes sense in the Q-attention, since the action is the pixel being chosen, but makes less sense in the next-best pose agent, since presumably the state is the entire 16x16 window of pixels, and the action is the pose. An ablation on this component’s design would be extremely helpful. It’s not clear to me that the performance gain comes from using 16x16 q-values as opposed to 1x1 q-values, or whether it comes from using confidence values attached to the state action pairs.\n- The sparse reward structure being used is missing from the discussion.\n- The authors point out that the keyframe selection method likely needs to become more complex with more sophisticated tasks. This again reduces the motivation and impact of the algorithm, if the hand-engineering of the reward function must now be replaced with hand-engineering of the keyframe selection method.\n- I would like to see some qualitative analyses of the keyframe selection method and how successful it is. Specifically, what do the sample demonstrations look like, and what percent of frames are keyframes? Intuitively, are there some failure modes of the keyframe selection method where some interesting keyframes are not included?\n- I have some confusion around the replay buffers. The pseudocode in algorithm 1 makes it seem that there is a single buffer that stores the final action (pose) taken, but the description claims that the projected pixel coordinates act as optimal actions for the Q-attention agent. Are these projected pixel coordinates also used in a separate replay buffer to warmstart the Q-attention agent?\n- It sounds like the proposed methods around data augmentations and discovery strategies are only necessary because the method has multiple agents as opposed to a single end-to-end agent. \n- Although none of the compared baselines work on the RLBench tasks, do the authors have some intuition on what would be required to make them successful? Is this due to using a sparse reward structure instead of a dense one? Or not having enough demonstrations? Or not having enough samples to interact with the environment? How does the authors’ method’s performance compare to these baselines in these other regimes? Including these would make the comparison much stronger.\n- The ablations present in Figure 6 study whether specific components are necessary (for example whether Q-attention is useful), but not on whether the formulation of that component is the correct formulation (for example, whether Q-attention is better than hard attention). Because the proposed modules are relatively complex and unintuitive, these ablations are critical.\n- I would recommend including ablations across all 10 tasks, potentially in the supplementary pages. The current evidence is less convincing that the augmentation strategy or confidence methods are necessary for success, due to the limited number of tasks profiled.\n\n---\n\n### Recommendation\n\nOverall, I vote for rejecting. I think the pipeline proposed by the authors is overly complex, and I am not yet convinced about the motivation for each individual module the authors include. The experimental results are currently lacking in terms of providing this motivation, and missing crucial ablations.\n\nIf the authors respond to my comments above with stronger motivations, clearer explanations for how the modules work, and more detailed ablative experiments, their submission will be stronger and many of my concerns would be resolved.\n\n---\n\n### Minor Feedback\n\n- Figure 2 and 3 are unclear since multiple pixels are highlighted, but earlier it is implied that only a single pixel is selected by the Q-attention agent.\n- Including videos of demonstrations would be potentially helpful as a qualitative comparison of the behaviors learned between the authors’ proposal and the baselines.\n- A qualitative analysis of the failure modes of the authors’ method would also be helpful, but not required.\n- Is there a way to visualize the behavior of the next-best pose agent? For instance, how does the pose vary depending on the crop passed in to the agent?\n- I don’t understand the reasoning in the appendix (section 7). Specifically, the authors claim that “as in hard-attention style approaches, we are unable to propagate gradients [...] by instead cropping the observation, we alleviate this problem.” The authors’ method still performs non-differentiable cropping, so I don’t see how they are able to propagate gradients. Likewise, cropping observations is commonly done with hard attention, so I don’t see how this method is particularly novel.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes  a general manipulation algorithm for tasks that have sparse rewards. The algorithm uses a Q-attention to extract interesting pixel locations with an explicit attention module. A data augmentation method is also proposed to generalize expert demonstrations. 
While the proposed method and experiments seem interesting, two out of the three reviewers had several issues regarding the clarity of the paper. The main issue is that a better ablation study needs to be performed to assess the proposed system. For instance, the reviewers question the advantages the Q-attention agent brings over a standard attention module. The authors provided detailed explanations and a video showing examples of the expert demonstrations. It is not clear however if the clarity issues are completely resolved. 
Decision: Reject