Attention-driven Robotic Manipulation

1 INTRODUCTION . Despite their potential , continuous-control reinforcement learning ( RL ) algorithms have many flaws : they are notoriously data hungry , often fail with sparse rewards , and struggle with long-horizon tasks . 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 ) , a privilege that is not feasible for training real-world robotic application across a broad range of tasks . Motivated by the observation that humans focus their gaze close to objects being manipulated ( Land et al. , 1999 ) , we propose an Attention-driven Robotic Manipulation ( ARM ) algorithm that consists of a series of algorithmagnostic components , that when combined , results in a method that is able to perform a range of challenging , sparsely-rewarded manipulation tasks . Our algorithm operates through a pipeline of modules : our novel Q-attention module first extracts interesting pixel locations from RGB and point cloud inputs by treating images as an environment , and pixel locations as actions . Using the pixel locations we crop the RGB and point cloud inputs , significantly reducing input size , and feed this to a next-best-pose continuous-control agent that outputs 6D poses , which is trained with our novel confidence-aware critic . 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 a keyframe state , we use our demo augmentation method which also stores the transition from intermediate points along a 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 methods when evaluated on 10 RLBench ( James et al. , 2020 ) tasks ( Figure 1 ) that range in difficulty . To summarise , we propose the following contributions : ( 1 ) An off-policy hard attention mechanism that is learned via Q-Learning , rather than the on-policy hard attention and soft attention that is commonly seen in the NLP and vision community ; ( 2 ) A confidence-aware Q function that predicts pixel-wise Q values and confidence values , resulting in improved convergence times ; ( 3 ) A keyframe discovery strategy and demo augmentation method that go hand-in-hand to improve the utilisation of demonstrations in RL . 2 RELATED WORK . The use of reinforcement learning ( RL ) is prevalent in many areas of robotics , including legged robots ( Kohl & Stone , 2004 ; Hwangbo et al. , 2019 ) , aerial vehicles ( Sadeghi & Levine , 2017 ) , and manipulation tasks , such as pushing ( Finn & Levine , 2017 ) , peg insertion ( Levine et al. , 2016 ; Zeng et al. , 2018 ; Lee et al. , 2019 ) , throwing ( Ghadirzadeh et al. , 2017 ; Zeng et al. , 2020 ) , ballin-cup ( Kober & Peters , 2009 ) , cloth manipulation ( Matas et al. , 2018 ) , and grasping ( Kalashnikov et al. , 2018 ; James et al. , 2019b ) . Despite the abundance of work in this area , there has yet to be a 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 et al. , 2018 ) , asymmetric actor-critic ( Pinto et al. , 2018 ) , reward shaping ( Rajeswaran et al. , 2018 ) , and auxiliary tasks ( James et al. , 2017 ) ) . Crucial to our method is 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. , 2015 ; Zhang et al. , 2019 ) . Soft attention deterministically multiplies an 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 maximising an approximate variational lower bound or equivalently via ( on-policy ) REINFORCE ( Williams , 1992 ) ) . Given that we perform non-differentiable cropping , our Q-attention is closest to hard attention , but with the distinction that we learn this in an off-policy way . This is key , as ‘ traditional ’ hard attention is 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 4.1 . Our proposed confidence-aware 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 a small amount of work in estimating uncertainty with Q-learning in discrete domains ( Clements et al. , 2019 ; Hoel et al. , 2020 ) ; our work uses a 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 a number of works ( Vecerik et al. , 2017 ; Matas et al. , 2018 ; Kalashnikov et al. , 2018 ; Nair et al. , 2018 ) , but while successful , they make limited use of the demonstrations and still can take many samples to converge . Rather than simply inserting these directly into the replay buffer , we instead make sure of our keyframe discovery and demo augmentation to maximise demonstration utility . 3 BACKGROUND . The reinforcement learning paradigm assumes an agent interacting with an environment consisting of states s ∈ S , actions a ∈ A , and a reward function R ( st , at ) , where st and at are the state and action at time step t respectively . The goal of the agent is then to discover a policy π that results in maximising the expectation of the sum of discounted rewards : Eπ [ ∑ t γ tR ( st , at ) ] , where future rewards are weighted with respect to the discount factor γ ∈ [ 0 , 1 ) . Each policy π has a corresponding value function Q ( s , a ) , which represents the expected return when following the policy after taking action a in state s. Our Q-attention module builds from Deep Q-learning ( Mnih et al. , 2015 ) , a method that approximated the value function Qθ , with a deep convolutional network , whose parameters θ are optimised by sampling mini-batches from a replay buffer D and using stochastic gradient descent to minimise the loss : E ( st , at , st+1 ) ∼D [ ( r + γmaxa′Qθ′ ( st+1 , a′ ) − Qθ ( st , at ) ) 2 ] , where Qθ′ is a target network ; a periodic copy of the online network Qθ which is not directly optimised . Our next-best pose agent builds upon SAC ( Haarnoja et al. , 2018 ) , however , the agent is compatible with any off-policy , continuous-control RL algorithm . SAC is a maximum entropy RL algorithm that , in addition to maximising the sum of rewards , also maximises the entropy of a policy : Eπ [ ∑ t γ t [ R ( st , at ) +αH ( π ( ·|st ) ) ] ] , where α is a 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 : JQ ( θ ) = E ( st , at , st+1 ) ∼D [ ( ( r+ γQθ′ ( st+1 , πφ ( st+1 ) ) − α log πφ ( at|st ) ) −Qθ ( st , at ) ) 2 ] . ( 1 ) The policy is updated towards the Boltzmann policy with temperature α , 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 : πnew = arg min π′∈Π DKL ( π′ ( ·|st ) ‖ 1 α exp ( Q πold ( st , · ) ) Zπold ( st ) ) . ( 2 ) Minimising the expected KL-divergence to learn the policy parameters was shown to be equivalent to maximising the expected value of the soft Qfunction : Jπ ( φ ) = E st∼D [ E a∼πφ [ α log ( πφ ( at|st ) ) −Qπρ ( st , at ) ] ] . ( 3 ) 4 METHOD . Our method can be split into a 3-phase pipeline . Phase 1 ( Section 4.1 ) consists of a high-level pixel agent that selects areas of interest using our novel Q-attention module . Phase 2 ( Section 4.2 ) consists of a next-best pose prediction phase where the pixel location from the previous phase is used to crop the incoming observations and then predict a 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 2 and Algorithm 1 . All experiments are run in RLBench ( James et al. , 2020 ) , a large-scale benchmark and learning environment for vision-guided manipulation built around CoppeliaSim ( Rohmer et al. , 2013 ) 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 a depth image d , along with proprioceptive information z from the arm ( consisting of end-effector pose and gripper open/close state ) . Using known camera intrinsics and extrinsics , we process each depth image to produce a point cloud p ( in world coordinates ) projected from the view of the front-facing camera , producing a ( H ×W × 3 ) ‘ image ’ . 4.1 Q-ATTENTION . Motivated by the role of vision and eye movement in the control of human activities ( Land et 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 organised 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 as the ‘ environment ’ , and pixel locations are treated as the ‘ actions ’ . Given our Q-attention function QAθ , we extract the coordinates of pixels with the highest value : ( xt , yt ) = argmax2D a′ QAθ ( st , a ′ ) . ( 4 ) The parameters of the Q-attention are optimised by using stochastic gradient descent to minimise the loss : JQA ( θ ) = E ( st , at , st+1 ) ∼D [ ( r+ γmax2D a′ QAθ′ ( st+1 , a ′ ) −QAθ ( st , at ) ) 2 + ‖QA‖ ] , ( 5 ) where s = ( b , p ) , QAθ′ is the target Q-function , and ‖QA‖ is an L2 loss on the per-pixel output of the Q function ( which we call Q regularisation ) ; in practice , we found that this leads to increased 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. , 2015 ) , which is summarised in Figure 2 . Example per-pixel outputs of the Q-attention are shown in Figure 3 . With the suggested coordinates from Q-attention , we perform a ( 16 × 16 ) crop on both the ( 128 × 128 ) RGB and organised point cloud data : b′ , p′ = crop ( b , p , ( x , y ) ) . 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 a 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 human-inspired motivation to the attention seen in NLP ( Bahdanau et 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 : ( 1 ) 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 crucial 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 Q-attention .

Summary: Motivated 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.