CausalDyna: Improving Generalization of Dyna-style Reinforcement Learning via Counterfactual-Based Data Augmentation

1 INTRODUCTION . Classical model-free reinforcement learning approaches require a massive amount of data collected in the environment to work , which slows down its success in tasks where data collection is timeconsuming or costly , like robot manipulation . Model-based reinforcement learning ( MBRL ) methods alleviate this issue by maintaining a world model that simulates the real environment . The world model can serve as a surrogate of the real environment for the agent to interact with to reduce the amount of the required time-consuming interaction in the real environment . MBRL methods ( Kaelbling et al. , 1996 ; Wang et al. , 2019 ; Janner et al. , 2019 ) learn from model rollouts of previously observed states . Recently , CTRL ( Lu et al. , 2020 ) takes a structural causal model ( SCM ) approach that can generate samples counterfactually had a different action had been taken for a state previously observed . However , these methods are limited for robotic manipulation tasks since the environment is often the key limiting factor . In this paper , we perform counterfactual reasoning on the object properties . For example , when the task manipulates objects with different masses , the real environment may not have a uniform distribution of object masses . Furthermore , to avoid damaging the robot , certain exploration of the gripper torque may be limited during training . To this end , we propose a Dyna-style MBRL method , CausalDyna in robotics that improves the policy performance by counterfactual reasoning of physics properties of objects and enriching the diversity of the generated rollouts . We leverage the structural causal model ( SCM ) to model the state dynamics . CausalDyna can be applied to generate episodes with unseen or rarely seen objects to improve the sample efficiency and generalization of the policy . Our contributions are summarized as follows . • We introduce a novel Dyna-style causal reinforcement learning algorithm , dubbed as CausalDyna that learns from counterfactually generated episodes with intervened object property values . • We compare with state-of-the-art model based reinforcement learning algorithm , MBPO and model free algorithm , SAC on the CausalWorld environment . Experimental results show that CausalDyna outperforms MBPO and SAC on sample efficiency by up to 17 % and generalization by up to 30 % when manipulating objects with unseen or rarely seen properties . 2 RELATED WORK . Causal Inference in Reinforcement Learning There is an increasing interest in causal inference in the field of reinforcement learning . Counterfactually-Guided Policy Search ( CF-GPS ) ( Buesing et al. , 2018 ) assumes that the real transition , observation , and reward functions are all known . They show that any partially observable Markov decision process ( POMDP ) can be represented as a structural causal model ( SCM ) . Therefore , counterfactual inference can be applied to improve the offpolicy evaluation and policy-guided search . CounTerfactual Reinforcement Learning ( CTRL ) ( Lu et al. , 2020 ) leverages bidirectional conditional GAN to model the environment dynamic for data augmentation . The model takes a noise vector as input besides the state and action to model the randomness of the environment . Before generating counterfactual data given alternative actions , they first infer the value of this noise vector . Then , the inferred noise is used to generate predictions with new actions . Causal Partial Models ( CPM ) ( Rezende et al. , 2020 ) studies the causal incorrectness of world models that don ’ t condition on the full observation . To fix this issue , CPM introduces a backdoor variable that helps the rollout of the model to be causally correct . We propose an SCM framework to model the physics properties of objects across the temporal dimension . In addition , we show that generating episodes with counterfactual object properties helps improve the generalization of the learned policy . Model-Based Reinforcement Learning Model-based Reinforcement Learning ( MBRL ) approaches have shown a potential to improve the sample efficiency by a large margin compared to classical model-free approaches ( Kaelbling et al. , 1996 ; Wang et al. , 2019 ) . Autoencoder-based algorithms like World Models ( Ha & Schmidhuber , 2018 ) and Dreamer ( Hafner et al. , 2019 ; 2020 ) use the world model to better represent the visual observation and faster the policy training . Policy Search with Backpropagation algorithms like PILCO ( Deisenroth & Rasmussen , 2011 ; Deisenroth et al. , 2013 ; Kamthe & Deisenroth , 2018 ) and GPS ( Levine & Koltun , 2013 ; Levine & Abbeel , 2014 ; Montgomery & Levine , 2016 ) train the policy by maximizing the simulated return of the policy in the world model . Because the world model is differentiable , the policy can be directly trained by gradient descent . Shooting algorithms like PETS-RS ( Chua et al. , 2018 ) and MB-MF ( Nagabandi et al. , 2018 ) alleviate the receding horizon problem in model predictive control ( MPC ) . Recent works include Ross & Bagnell ( 2012 ) , MOPO , ( Yu et al. , 2020 ) and Morel ( Kidambi et al. , 2020 ) show that MBRL can work well in the offline RL setting . Unlike the traditional MBRL that approximates the local transition function , L3P ( Zhang et al. , 2021 ) builds the world model as a graph of states for better reasoning ability . Dyna-style algorithms ( Sutton , 1990 ; 1991a ; b ) use the learned world model to roll out simulated episodes to reduce the demand for real data for policy training . As a recent development of Dyna-style algorithms , ME-TRPO ( Kurutach et al. , 2018 ) uses an ensemble of world models to catch the epistemic uncertainty ; MB-MPO ( Clavera et al. , 2018 ) viewed each model in the ensemble as a task and meta-learn a policy that adapts quickly to handle the model-bias issue ; MBPO ( Janner et al. , 2019 ) rolls out short episodes branched from real data to improve the generation quality . Our method follows the Dyna-style framework and targets designing and using a causal world model to generate better and more diverse rollouts in robotic environments . 3 BACKGROUND . 3.1 STRUCTURAL CAUSAL MODEL . Structural Causal Model ( SCM ) is a widely used framework to describe the causal mechanism of a system . Let ’ s denote X = { x1 , ... , xN } as the set of N variables in a system . Knowing their causal relationships allows us to build a directed acyclic causal graph to describe this system . Each node represents a variable , which is directly caused by its parent nodes . In this way , a node xn can be modeled as the following function : xn = fi ( Paobs ( xn ) , un ) ( 1 ) Here , Paobs ( xn ) denotes the observed parent nodes of xn . un is a noise that represents the effect of omitted factors . This function is also called a causal mechanism . SCM is the set of these causal mechanisms that describes the whole system . SCM defines a joint distribution of the variables p ( x1 , ... , xN ) following the causal Markov assumption : given its direct causes , each variable xn is independent of other indirect causal variables . 3.2 DYNA-STYLE MODEL-BASED REINFORCEMENT LEARNING . Dyna-style model-based reinforcement learning uses the world model to roll out simulated episodes , which can be viewed as data augmentation . The training of Dyna-style MBRL is composed of three steps : First , the agent interacts with the real environment and collects real data to train the world model . Then , this world model is used as a simulator of the real environment for the agent to interact and collect simulated data . After that , the agent can be trained together with the real and the simulated data using classical model-free reinforcement learning algorithms . These three steps are executed repeatedly until the training converges . In case we apply Dyna-Style algorithm on RL algorithms with experience-replay buffers and would like to collect whole simulated episodes , as the world model is trained to only approximate the transition of the environment p ( st+1|st , at ) , we need an initial state to start the simulated episodes . A usual way to solve it is using the first state or a randomly sampled state st from the collected real episodes as the start point of the simulated episodes . As the real episode already contains the future of st under the original action sequence { at , at+1 , ... } executed in this episode , generating new simulated episodes starting from st under different action sequences can be viewed as answering a counterfactual “ what if ” question : What would happen if the agent behave differently this time instead of doing { at , at+1 , ... } ? The world model gives the agent a chance to figure out the answer without interacting in the real environment , and helps the agent learn faster . Algorithm 1 : Counterfactual Property Generation Data : Rollout length K , Real experience buffer Dr , Policy pπ , World model pWM , Counterfactual property space M , Empty episode buffer B Result : B 1 Sample a state s = [ s−m ; m ] from the real experience buffer Dr , B.append ( s ) 2 Sample a counterfactual property value mCF from M , set s̃ = [ s−m ; mCF ] 3 for K steps do 4 ã ∼ pπ ( a|s̃ ) , s̃′ ∼ pWM ( s′|s̃ , ã ) 5 B.append ( ã , s̃′ ) , s̃← s̃′ 6 end 4 METHOD . 4.1 STRUCTURE CAUSAL MODEL OF A ROBOT ENVIRONMENT . Let ’ s consider an environment where a robot needs to manipulate an object . We can describe this environment using different states . Many of these states are changing over time , including the object position and the end-effect position . It is important to model them as they directly contain the dynamic information of the environment . Some other states are time-invariant , like the object mass or the floor friction coefficient . Although their values are fixed , they determine the environment dynamics and affect how other time-variant states change . Let ’ s denote the total state at step t as st. st = [ s−m , t ; m ] is the concatenation of the time-variant state s−m , t at step t and the object timeinvariant property m. The motor torque to execute at step t is denoted as at . As shown in Fig.2 , we can build a structural causal model ( SCM ) to describe this environment . The time-invariant property m is modeled as a fixed node across the temporal dimension , which affects all the causal mechanisms . 4.2 COUNTERFACTUAL PROPERTY GENERATION . Policy generalization ability is essential as the testing environment of the policy is not always the same as the training environment . For example , when learning to lift an object , the robot might only interact with objects whose masses are in a suitable range . Lifting frequently a too-heavy object might reduce its service life , and most reinforcement learning algorithms need a large amount of interaction data to work . However , knowing how to lift a heavy object is still desirable when deploying the robot . A typical Dyna-style method generates simulated rollouts branching from a starting state seen in previous real episodes . If the world model takes physics properties as input , it is possible to go a step further and intervene in these properties . For example , we could modify the mass of an object in the world model to make it heavier . So the agent can learn to manipulate them in the world model as much as we want without harming its service life . Inspired by this , we design a simple generation strategy to enrich the simulated rollouts by modifying the original object ’ s property to improve the policy generalization . Concretely , instead of taking a starting state st = [ s−m , t ; m ] sampled from real episodes as it is like most of the Dyna-style methods , we replace the object property m by a desired counterfactual value mCF sampled from a predefined counterfactual property space M before rolling out the simulated episodes . We name this type of episodes generation as counterfactual property generation , illustrate it in Fig.1 and show the process in Alg.1 .

Summary: This paper considers the problem of generalization of a dynamics model to different versions of the same environment, indexed by parameter m. For example, this parameter $m$ might be the mass of an object to be picked up. The problem that the authors tackle is to generalize the agent's policy to environments with either different parameters $m$ or a different distribution of parameters $m$. The solution the authors propose is to randomly sample different values of $m$ during training of the model, as opposed to using the $m$'s observed from the agent's rollouts. Empirically we observe that the proposed method performs similar to MBPO in generalization to out-of-distribution $m$ (Section 5.2) and that generalizes better on heavier masses than MBPO in the unbalanced distribution setting (Section 5.3).