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: Motivated by the difficulty of evaluating RL’s ability to transfer behaviors across environments, the authors propose the CausalWorld benchmark. Unlike prior benchmarks, CausalWorld exposes well-defined casual variables, in the form of task factors, and focuses on robotic manipulation of an open-source robot platform. The authors make CausalWorld easily usable for both training (in defining a learning curriculum) and evaluation (in targeting specific expected generalizations), and make it easy to extend. In their original release, the authors include eight concrete tasks to test generalization, and present baseline results on these tasks.