CausalWorld: A Robotic Manipulation Benchmark for Causal Structure and Transfer Learning

1 INTRODUCTION Benchmarks have played a crucial role in advancing entire research fields , for instance computer vision with the introduction of CIFAR-10 and ImageNet ( Krizhevsky et al. , 2009 ; 2012 ) . When it comes to the field of reinforcement learning ( RL ) , similar breakthroughs have been achieved in domains such as game playing ( Mnih et al. , 2013 ; Silver et al. , 2017 ) , learning motor control for high-dimensional simulated robots ( Akkaya et al. , 2019 ) , multi-agent settings ( Baker et al. , 2019 ; Berner et al. , 2019 ) and for studying transfer in the context of meta-learning ( Yu et al. , 2019 ) . Nevertheless , trained agents often fail to transfer the knowledge about the learned skills from a training environment to a different but related environment sharing part of the underlying structure . This can be attributed to the fact that it is quite common to evaluate an agent on the training environments themselves , which leads to over- fitting on these narrowly defined environments ( Whiteson et al. , 2011 ) , or that algorithms are com∗Equal contribution . Correspondence to : < ossama.ahmed @ mail.mcgill.ca > , < frederik.traeuble @ tuebingen.mpg.de > †Equal Advising 1 https : //sites.google.com/view/causal-world/home pared using highly engineered and biased reward functions which may result in learning suboptimal policies with respect to the desired behaviour ; this is particularly evident in robotics . In existing benchmarks ( Yu et al. , 2019 ; Goyal et al. , 2019a ; Cobbe et al. , 2018 ; Bellemare et al. , 2013 ; James et al. , 2020 ) the amount of shared causal structure between the different environments is mostly unknown . For instance , in the Atari Arcade Learning environments , it is unclear how to quantify the underlying similarities between different Atari games and we generally do not know to which degree an agent can be expected to generalize . To overcome these limitations , we introduce a novel benchmark in a robotic manipulation environment that we call CausalWorld . It features a diverse set of environments that , in contrast to previous designs , share a large set of parameters and parts of the causal structure . Being able to intervene on these parameters ( individually or collectively ) permits the experimenter to evaluate agents ’ generalization abilities with respect to different types and magnitudes of changes in the environment . These parameters can be varied gradually , which yields a continuum of similar environments . This allows for fine-grained control of training and test distributions and the design of learning curricula . A remarkable skill that humans learn to master early on in their life is building complex structures using spatial-reasoning and dexterous manipulation abilities ( Casey et al. , 2008 ; Caldera et al. , 1999 ; Kamii et al. , 2004 ) . Playing with toy blocks constitutes a natural environment for children to develop important visual-spatial skills , helping them ‘ generalize ’ in building complex composition designs from presented or imagined goal structures ( Verdine et al. , 2017 ; Nath & Szücs , 2014 ; Dewar , 2018 ; Richardson et al. , 2014 ) . Inspired by this , CausalWorld is designed to aid in learning and investigating these skills in a simulated robotic manipulation environment corresponding to the open-source TriFinger robot platform from Wüthrich et al . ( 2020 ) , which can be built in the real world . Tasks are formulated as building 3D goal shapes using a set of available blocks by manipulating them - as seen in Fig . 1 . This yields a diverse familiy of tasks , ranging from relatively simple ( e.g . pushing a single object ) to extremely hard ( e.g . building a complex structure from a large number of objects ) . CausalWorld improves upon previous benchmarks by exposing a large set of parameters in the causal generative model of the environments , such as weight , shape and appearance of the building blocks and the robot itself . The possibility of intervening on any of these properties at any point in time allows one to set up training curricula or to evaluate an agent ’ s generalization capability with respect to different parameters . Furthermore , in contrast to previous benchmarks ( ChevalierBoisvert et al. , 2018 ; Cobbe et al. , 2018 ) , researchers may build their own real-world platform of this simulator at low cost , as detailed in Wüthrich et al . ( 2020 ) , and transfer their trained policies to the real world . Finally , by releasing this benchmark we hope to facilitate research in causal structure learning , i.e . learning the causal graph ( or certain aspects of it ) as we operate in a complex real-world environment whose dynamics follow the laws of physics , which induce causal relations between the variables . Changes to the variables we expose can be considered do-interventions on the underlying structural causal model ( SCM ) . Consequently , we believe that this benchmark offers an exciting opportunity to investigate causality and its connection to RL and robotics . Our main contributions can be summarized as follows : • We propose CausalWorld , a new benchmark comprising a parametrized family of robotic manipulation environments for advancing out-of-distribution generalization and causal structure learning in RL . • We provide a systematic way of defining curricula and disentangling generalization abilities of RL agents with respect to different changes in the environment , since we allow for dointerventions to be performed on different environment variables ( parameters and states ) individually . • We establish baseline results for some of the available tasks under different learning algorithms , thus verifying the feasibility of the tasks . • We show how different learning curricula affect generalization across different axes by reporting some of the in-distribution and out-of-distribution generalization capabilities of the trained agents . 2 CAUSALWORLD BENCHMARK . Here we make the desiderata outlined in the introduction more precise : 1 . The set of environments should be sufficiently diverse to allow for the design of challenging transfer tasks . 2 . We need to be able to intervene on different properties ( e.g . masses , colors ) individually , such that we can investigate different types of generalization . 3 . It should be possible to convert any environment to any other environment by gradually changing its properties through interventions ; this requirement is important for evaluating different levels of transfer and for defining curricula . 4 . The environments should share some causal structure to allow algorithms to transfer the learned causal knowledge from one environment to another . 5 . There should be a unified measure of success , such that an objective comparison can be made between different learning algorithms . 6 . The benchmark should make it easy for users to define meaningful distributions of environments for training and evaluation . In particular , it should facilitate evaluation of indistribution and out-of-distribution performance . 7 . The simulated benchmark should have a real-world counterpart to allow for sim2real . In light of these desiderata , we propose a setup in which a robot must build goal shapes using a set of available objects . It is worth noting that similar setups were proposed previously in a less realistic setting , e.g . in ( Janner et al. , 2018 ; Bapst et al. , 2019 ; McCarthy et al . ; Akkaya et al. , 2019 ; Fahlman , 1974 ; Winston , 1970 ; Winograd , 1972 ) . Specifically , a task is formulated as follows : given a set of available objects , the agent needs to build a specific goal structure , see Fig . 1 for an example . The vast amount of possible target shapes and environment properties ( e.g . mass , shape and appearance of objects and the robot itself ) makes this a diverse and challenging setting to evaluate different generalization aspects . CausalWorld is a simulated version ( using the Bullet physics engine ( Coumans et al. , 2013 ) ) of the open-source TriFinger robot platform from Wüthrich et al . ( 2020 ) . Each environment is defined by a set of variables such , as gravity , floor friction , stage color , floor color , joint positions , various block parameters ( e.g . size , color , mass , position , orientation ) , link colors , link masses and the goal shape . See Table 3 in the Appendix for a more extensive list of these variables . Desideratum 1 is satisfied since different environment properties and goal shapes give rise to very different tasks , ranging from relatively easy ( e.g . re-positioning a single cube ) to extremely hard ( e.g . building a complex structure ) . Desideratum 2 is satisfied because we allow for arbitrary interventions on these properties , hence users or agents may change parameters individually or jointly . Desideratum 3 is satisfied because the parameters can be changed gradually . Desideratum 4 is satisfied because all the environments share the causal structure of the robot , and one may also use subsets of environments which share even more causal structure . We satisfy desideratum 5 by defining the measure of success for all environments as the volumetric overlap of the goal shape with available objects . Further , by splitting the set of parameters into a set A , intended for training and in-distribution evaluation , and a set B , intended for out-of-distribution evaluation , we satisfy desideratum 6 . Finally , since the TriFinger robot ( Wüthrich et al. , 2020 ) can be built in the real-world , we satisfy desideratum 7 . Desideratum 7 and 2 are in partial conflict since sim2real is only possible for the tasks which are constrained to the variables on which the robot can physically act upon . Task generators : To generate meaningful families of similar goal shapes , CausalWorld allows for defining task generators which can generate a variety of different goal shapes in an environment . For instance , one task generator may generate pushing tasks , while another one may generate towerbuilding tasks ( see Fig . 2 ) . Each task generator is initialized with a default goal shape from its corresponding family and comes with a sampler to sample new goal shapes from the same family . Additionally , upon construction , one can specify the environments ’ initial state and initial goal shape structure when deviating from the default . The maximum episode time to build a given shape is number of blocks×10 seconds . CausalWorld comes with eight pre-defined task generators ( see Fig . 2 ) . • Three generators create goal shapes with a single block : Pushing with the goal shape on the floor , Picking having the goal shape defined above the floor and Pick and Place where a fixed obstacle is placed between the initial block and goal pose . • Stacking2 involves a goal shape of two stacked blocks , which can also be considered one instance of the Towers generator . • The remaining generators use a variable number of blocks to generate much more complex and challenging target shapes , as detailed in the appendix : Towers , Stacked Blocks , Creative Stacked Blocks and General . Given that building new environments using current physics simulators is often tedious , we provide a simple API for users who wish to create task generators for new challenging shape families , which may be added to CausalWorld ’ s task generators repository . Action and Observation Spaces : The robot can be chosen to operate in either joint position control mode , joint torque control mode , end-effector position control mode , or the delta of each . In any of these cases , the action is 9-dimensional ( one per joint ) . We provide two observation modes : structured and pixel . In the structured mode , the observation vector is constructed using a rule for the ordering of the relevant variables , such as joint positions , joint velocities , block positions , etc . Thus , the size of the observation space depends on the number of blocks , which could potentially change with every new goal sampled , e.g . in Towers , ( Creative ) Stacked Blocks and General . In contrast , in the pixel mode , the agent receives six RGB images ( hence the dimension of the observation is 6×3×128×128 ) , the first three images are rendered from the three cameras mounted around the TriFinger robot , and the last three images specify the goal image of the target shape rendered from the same cameras . Additionally , CausalWorld allows users to set up a fully customized observation space . Rewards : The reward function r is defined as the fractional volumetric overlap of the blocks with the goal shape , which ranges between 0 ( no overlap ) and 1 ( complete overlap ) . Since this reward function is shared across all tasks , an agent that learned r from some training tasks could in principle use it to solve unseen goal structures . There is also the possibility of modifying the reward function to 1 ) sparsify the reward further by returning a binary reward signal instead , or 2 ) add a dense reward function in order to introduce inductive biases via domain knowledge and solution guidance . We hope that the considerable complexity and diversity of goal shapes motivate and accelerate the development of algorithms that are not dependent on hand-tuned reward functions . ` Training and evaluation spaces : In this benchmark , a learning setting consists of an allowed training space ( ATS ) and an evaluation space ( ES ) , both of which are subspaces of the full parameter space . During training , in the simplest setting , parameters are sampled iid from the ATS . However , unlike existing benchmarks , CausalWorld allows in addition for curricula within the ATS as well as settings where the agent itself intervenes on the parameters within an episode ( see Fig . 3 ) . Similarly , during evaluation , parameters may be sampled iid from the evaluation space at each episode reset , or there can be interventions within an episode . Moreover , in order to retrieve the setting considered in most RL benchmarks , we could set the ATS and the ES to be identical and intervene only on object and robot states ( and keep other environment properties constant ) at each episode reset . However , to evaluate out-of-distribution generalization , one should set the two spaces ( ATS and ES ) to be different ; possibly even disjoint . Additionally , to evaluate robustness with respect to a specific parameter ( e.g . object mass ) , one may define the training and evaluation spaces which only differ in that particular parameter . In order to facilitate the definition of appropriate training and evaluation settings , we pre-define two disjoint sets , Ai and Bi , for each parameter i . Through this , one can for instance define the training space to be A1 ×A2 × ... and the evaluation space to be B1 ×B2 × ... to assess generalization with respect to all parameters simultaneously . Alternatively , the evaluation space could be defined as A1 ×A2 × ... ×Bi ×Ai+1 × ... to assess generalization with respect to parameter i only . Lastly , users may also define their own spaces which could then be integrated into the benchmark to give rise to new learning settings . Intervention actors : To provide a convenient way of specifying learning curricula , we introduce intervention actors . At each time step , such an actor takes all the exposed variables of the environment as inputs and may intervene on them . To encourage modularity , one may combine multiple actors in a learning curriculum . This actor is defined by the episode number to start intervening , the episode number to stop intervening , the timestep within the episode it should intervene and the episode periodicity of interventions . We provide a set of predefined intervention actors , including an actor which samples parameters randomly at each episode reset , which corresponds to domainrandomization . It is also easy to define custom intervention actors , we hope that this facilitates investigation into optimal learning curricula ( see Fig . 3 ) . Probing the Causal Structure in RL : The problem setting in RL is usually formulated using the language of Markov Decision Processes ( MDPs ) or Partially Observable Markov Decision Processes ( POMDPs ) ( Sutton & Barto , 1999 ) , but can be also represented by Structural Causal Models ( SCMs ) , as shown in ( Buesing et al. , 2018 ) , refer to section E in the Appendix for a detailed explanation . This is achieved by formulating all conditional probability distributions as deterministic functions that take independent noise variables as inputs . These independent noise variables can specify different scenarios while the deterministic functions reflect the causal mechanisms of the system ( Schölkopf et al. , 2021 ) . Changes in the environment can stem from two different sources : 1 . The agent may alter the state of the environment ( e.g . the position of a block ) indirectly , through its actions ( e.g . pushing the block by applying appropriate torques at the motors ) . 2 . During the execution of a learning curriculum or an evaluation protocol , we may directly intervene on any variable of the SCM , including all the latent variables of the causal model that are not accessible to the RL agent ( such as gravity , object mass or color ) . ( 1 ) is the default type of admissible interventions in RL benchmarks , whereas CausalWorld allows for interventions of type ( 2 ) in addition . The idea is that interventions on these latent variables , e.g . during a learning curriculum , will allow the agent to distinguish between spurious correlations that are only present in a particular setting and true causal relations that will hold across all settings ( i.e . interventions ) . If the agent is able to learn such a representation of the underlying SCM structure , we would expect it to perform well even in out-of-distribution scenarios ( Schölkopf et al. , 2021 ; Dittadi et al. , 2021 ) because the causal structure remains the same , even when the functional form of certain relations may vary ( e.g . when transferring to the real robot ) . Moreover , we hope that by having access to a broad range of interventions in CausalWorld it will aid the inference of the underlying SCM structure through the different causal discovery methods ( see Figure 4 for a subset of the expected SCM to be learned ) , which in turn addresses the lack of causal discovery benchmarks for real world challenges . 𝑓1 𝑓2 𝑓3 𝑓4

Summary: This paper proposed a new benchmark for studying reinforcement learning and its generalization in the context of the robotic manipulation problem. To study the generalization of a learned policy, the proposed benchmark is equipped with an interface that makes intervention easy. This interface helps to define a training space and an evaluation space so that one can systemically study both in-distribution and out-of-distribution generalization of a learned policy. At the same time, the proposed benchmark simulates an open-source robot platform, which makes sim2real transfer experiments easier.