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: The paper presents a method to improve the generalization of model-based RL by means of interventional data augmentation. The key idea is to intervene the value of a particular variable (e.g., object property) in the learned dynamic model for episode simulations. Experimental results show that it improves (i) the generalization ablity in the OoD scenarios with respect to the intervened variable, (ii) sample efficiency in the presence of unbalanced training distribution.