Transferable Visual Control Policies Through Robot-Awareness

1 INTRODUCTION . Raw visual observations provide a versatile , high-bandwidth , and low-cost information stream for robot control policies . However , despite the huge strides in machine learning for computer vision tasks in the last decade , extracting actionable information from images remains challenging . As a result , even simple robotic tasks , such as pushing objects using image observations , commonly require data collected over many hours of robot interaction to synthesize effective visual policies . If such policies could transfer reliably and easily to new target robots , then this data collection cost would be amortized . For example , a hospital adding a new robot to its robot fleet could simply plug in their existing policies and start using it immediately . Going further , other hospitals looking to automate the same tasks could purchase a robot of their choice and download the same policy models . However , such transferable policies are difficult to achieve in practice . Even when the task setting , such as the hospital , remains unchanged , the changed visual appearance of the robot itself leads to out-of-distribution inputs for visual policies pre-trained on other robots . This issue particularly affects manipulation tasks : manipulation involves operating in intimate proximity with the environment , and any cameras set up to observe the environment can not avoid also observing the robot . There is a way out of this bind : most robots are capable of highly precise proprioception and kinesthesis to sense body poses and movements through internal sensors . We propose to develop “ robot-aware ” policies that can benefit from distinguishing pixels corresponding to the robot agent from those corresponding to the rest of the “ world ” in image observations . Our robot-aware policies treat the robot and the world differently , to their advantage . While this general principle applies broadly to all visual controllers , we demonstrate the advantages of robot awareness in model-based reinforcement learned policies that work by planning through visual dynamics models that are trained to predict the consequences of agent actions . First , we inject robot awareness into the dynamics model , by factorizing it into robot-specific “ robot ” dynamics and robot-invariant “ world ” dynamics . Figure 1 shows a schematic . Our experiments show that composing these two modules permits reliably transferring visual dynamics models even across robots that look and move very differently . Next , we design a robot-aware planning cost over the separated robot and world pixels . We show that this not only allows visual task reward specifications to transfer from a source to a target robot , it even leads to gains on the source robot itself by allowing the policy to reason separately about the robot and its environment . We evaluate our method , robot-aware control ( RAC ) , on simulated and real-world tabletop object pushing and pick-and-place tasks , demonstrating the importance of each robot-aware component . To our knowledge , RAC achieves two firsts for zero-shot transfer of visuomotor skills across robots : ( 1 ) successful transfer between real-world robots , and ( 2 ) successful transfer after training on just a single source robot . In particular , we show that a robot-aware policy trained on a single , low-cost 5-DoF WidowX arm can transfer entirely “ zero-shot ” to a very different 7-DoF Franka Panda arm in the real world . 2 BECOMING AWARE OF THE ROBOT AGENT IN THE SCENE . We start with the question : what does it mean to be “ aware of the robot ” ? Consider , for example , a manipulation task involving a standard robot arm , as in Fig 1 . At each discrete timestep t , the input to the agent ’ s policy is an RGB image observation o ( t ) from a static camera observing the workspace of the robot . Since the robot must operate in close proximity with objects in the workspace , these images contain parts of the robot ’ s body as well as the rest of the workspace . We define a visual projection function I ( · ) that maps from true robot state r and “ world ” ( i.e. , the rest of the workspace ) state w to the image observation o ( t ) = I ( r ( t ) , w ( t ) ) . Robot states r could be robot joint angles , and world states w could be object poses in the workspace . A robot-aware agent must first locate the robot in the visual observation o ( t ) from the camera . Note that the robot ’ s state r ( t ) is readily available in high precision through internal joint encoders . Further , a full geometric specification of the robot shape and joint degrees of freedom may be assumed available ( usually from the manufacturer ) . Finally , techniques for camera-to-robot calibration are increasingly robust and efficient , to the point of permitting calibration from even a single image when the robot shape is known as above ( Labbé et al. , 2021 ; Lee et al. , 2020 ) . Combining these , we may simply project the articulated robot shape through the camera parameters to obtain the visual projection of the robot in the scene . When parts of this projection may be occluded , we may further use RGB-D observations to easily identify and handle them , see Supp . A.11 for details . We represent this as an image segmentation mask Mr ( t ) of the same height and width as image observations o ( t ) : Mr is 1 on robot pixels and 0 elsewhere . We show examples of such masks in Fig 2 . This simple process effectively spatially disentangles robot pixels from world pixels in image observations o ( t ) . This spatial disentanglement is a key building block for our robot-aware visual policies , since it permits treating robot and world pixels differently.1 Secondarily , with common high-level robot action spaces used in reinforcement learning ( such as end-effector displacement actions ) , it may even be possible to analytically predict future robot states r′ given the current state r and the action a : Pr ( r , a ) = r′ . This is not a strict requirement for robot-aware control , and we may instead simply learn the robot-specific dynamics model Pr from a small amount of experience . However , as we will show in our experiments , analytical Pr models work well for our manipulation tasks , and further , they permit fully zero-shot transfer to a new robot . 1Alternative techniques permit automatically “ self-recognizing ” the robot in a scene ( Michel et al. , 2004 ; Natale et al. , 2007 ; Edsinger & Kemp , 2006 ; Yang et al. , 2020 ) even without prior knowledge of robot embodiment . 3 TRANSFERABLE ROBOT-AWARE MODEL-BASED RL POLICIES . Above , we have described how a policy might be “ robot-aware ” by spatially disentangling the robot in its observations , and utilizing knowledge of its dynamics . We now show how such robotawareness enables effective visual control policies that can be readily transferred between robots . In particular , we instantiate a robot-aware model-based reinforcement learning agent . Continuing in the manipulation setting from Sec 2 , assume that tasks are specified through a goal image that exhibits the target configuration of the workspace . Given the goal image og = I ( rg , wg ) , the agent must execute a sequence of actions aT0 = ( a ( 0 ) , a ( t+ 1 ) , ... a ( T ) ) to reach that goal , where T is a time limit . Success is typically measured by how close the final world state is to the goal image . For example , the goal image might show a specific goal orientation of objects on a table , in which case success might be measured in terms of how close the objects are to their goals . To solve such tasks , we inject robot-awareness into a widely used model-based reinforcement learning paradigm for robots : visual foresight ( VF ) ( Finn & Levine , 2017 ; Ebert et al. , 2018 ) . Standard VF involves two key steps : • Visual dynamics modeling : The first step is to perform exploratory data collection on the robot to generate a dataset D of transitions ( o ( t ) , a ( t ) , o ( t+ 1 ) ) . Dropping time indices , we sometimes use the shorthand ( o , a , o′ ) to avoid clutter . Then , a visual dynamics model P is trained on this dataset D to predict o′ given o and a as inputs , i.e . o′ ≈ P ( o , a ) . When the robot state r is available , it is sometimes included as a third input to P to assist in dynamics modeling . • Visual MPC : Given the trained dynamics model P , VF approaches search over action sequences aT0 to find the sequence whose outcome will be closest to the goal specification og , as predicted by P . For outcome prediction , they apply P recursively , as ô ( o ( 0 ) , a ) = P ( P ( . . . P ( o ( 0 ) , a ( 0 ) ) , a ( 1 ) ) , . . . a ( T ) ) . Then , they pick the action sequence a∗ ( o ( 0 ) , og ) = argmina C ( ô ( o ( 0 ) , a ) , og ) , where the cost function C is commonly the mean pixel-wise ` 2 error between the predicted image and the goal image . Sometimes , rather than measure the error of only the final image , the cost function may sum the errors of all intermediate predictions . For closed-loop control , only the first action a∗ ( 0 ) is executed ; then , a new image o ( 1 ) is observed , and a new optimal action sequence aT1 ∗ ( o ( 1 ) , og ) is computed , and the process repeats . For action sequence optimization , we found the cross-entropy method ( CEM ) ( De Boer et al. , 2005 ) to be sufficient , although more sophisticated optimization methods ( Zhang et al. , 2019 ; Rybkin et al. , 2021 ) could also work well . To train reliable dynamics models P , visual foresight approaches commonly require many hours of exploratory robot interaction ( Finn & Levine , 2017 ; Ebert et al. , 2018 ) even for simple tasks . Data requirements may be reduced a little by interleaving model training with data collection : in this case , data is collected by selecting goal images and then running visual MPC using the most recently trained dynamics model to reach those goals . However , such goal selection must be designed to match the target task ( s ) , which may trade off task generality for sample complexity . 3.1 IMPROVING VISUAL FORESIGHT THROUGH ROBOT-AWARENESS . Having trained one data-hungry visual dynamics model P on one robot , do we still need to repeat this process from scratch on a new robot aiming to perform the same tasks ? In standard VF , the answer is unfortunately yes , for two main reasons : ( 1 ) With the new robot , all observations o ( t ) = I ( rnew ( t ) , w ( t ) ) are out-of-domain for P . Further , the visual dynamics of the new robot may be very different , as between a green 3-DOF robot arm , and a red 5-DOF robot arm . ( 2 ) As described above , standard VF approaches operate by aiming to match the goal image . However , since the observations contain the robot , they commonly require the task specification image og = I ( rg , wg ) itself to contain the robot in a plausible goal-reaching pose . This makes task specification more difficult , and also robot-specific : it is impossible to plan to reach og using a different robot . We show how VF policies may overcome these obstacles through robot-awareness .

Summary: The paper presents an extension of previous works on training a visual policy to control a robot arm in a manipulation context. The policy is pre-trained on a large robot experience dataset (e.g. different robots interacting with different objects in a table-top setting). The dataset is used to train a model for the visual dynamics (i.e. next visual & robot state given current visual & robot state and applied action). The visual dynamics are used to reach a desired visual state using an existing technique: visual foresight (VF) by Finn & Levine, 2017 and Ebert et al., 2018.   The main novelties of the paper of the paper are two. First, authors propose a technique to learn a decoupled visual dynamic model: the world model (i.e. a model to propagate forward the world-observations consisting of an image of all objects in the scene excluding the robot) is decoupled from the robot model (i.e. a model to propagate forward the end-effector pose and a mask which represents the robot body in the visual space). Second, the visual foresight planning relies on a newly proposed cost which combines and decouples two costs: a desired robot-state cost and a desired world-state cost.