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: This 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.