A Relational Intervention Approach for Unsupervised Dynamics Generalization in Model-Based Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) has shown great success in solving sequential decision-making problems , such as board games ( Silver et al. , 2016 ; 2017 ; Schrittwieser et al. , 2020 ) , computer games ( e.g . Atari , StarCraft II ) ( Mnih et al. , 2013 ; Silver et al. , 2018 ; Vinyals et al. , 2019 ) , and robotics ( Levine & Abbeel , 2014 ; Bousmalis et al. , 2018 ) . However , solving real-world problems with RL is still a challenging problem because the sample efficiency of RL is low while the data in many applications is limited or expensive to obtain ( Gottesman et al. , 2018 ; Lu et al. , 2018 ; 2020 ; Kiran et al. , 2020 ) . Therefore , model-based reinforcement learning ( MBRL ) ( Janner et al. , 2019 ; Kaiser et al. , 2019 ; Schrittwieser et al. , 2020 ; Zhang et al. , 2019b ; van Hasselt et al. , 2019 ; Hafner et al. , 2019b ; a ; Lenz et al. , 2015 ) , which explicitly builds a predictive model to generate samples for learning RL policy , has been widely applied to a variety of limited data sequential decision-making problems . However , the performance of MBRL methods highly relies on the prediction accuracy of the learned environmental model ( Janner et al. , 2019 ) . Therefore , a slight change of environmental dynamics may cause a significant performance decline of MBRL methods ( Lee et al. , 2020 ; Nagabandi et al. , 2018a ; Seo et al. , 2020 ) . The vulnerability of MBRL to the change of environmental dynamics makes them unreliable in real world applications . Taking the robotic control as an example ( Nagabandi et al. , 2018a ; Yang et al. , 2020 ; Rakelly et al. , 2019 ; Gu et al. , 2017 ; Bousmalis et al. , 2018 ; Raileanu et al. , 2020 ; Yang et al. , 2019 ) , dynamics change caused by parts damages could easily lead to the failure of MBRL algorithms . This problem is called the dynamics generalization problem in MBRL , where the training environments and test environments share the same state S and action space A but the transition dynamics between states p ( st+1|st , at ) varies across different environments . Following previous works ( Petersen et al. , 2018 ; Gottesman et al. , 2018 ) , we focus on the unsupervised dynamics generalization setting , i.e . the id or label information of dynamics function in training MDPs is not available . This setting appears in a wide range of applications where the information of dynamics function is difficult to obtain . For example , in healthcare , patients may respond differently to the same treatment , i.e. , p ( st+1|st , at ) varies across patients . However , it is difficult to label which patients share similar dynamics . To build a generalized dynamics prediction function that can generalize to different transition dynamics , the shift of the transition dynamics can be modelled as the change of unobserved factors across different environments , i.e . there are hidden environment-specified factors Z ∈ Z which can affect the environmental dynamics . This is analogous to the human intuition to understand the change of dynamics , e.g . patients may show different responses to identical treatments because the differences of their gene sequences can affect how well they absorb drugs ( Wilke et al. , 2007 ) . It is natural to assume that Z is the same in a single environment but varies across different environments . As such , these unobserved environment-specified factors do not violate the nature of MDP in a single environment , but their changes can affect the dynamics functions across environments . Therefore , the dynamics function f : S × A → S can naturally be augmented by incorporating Z to be f : S ×A×Z → S ( Rakelly et al. , 2019 ; Zhou et al. , 2019 ; Lee et al. , 2020 ; Seo et al. , 2020 ) . Learning the augmented dynamics function is difficult because the environment-specified factor Z is unobservable . Previous methods ( Rakelly et al. , 2019 ; Zhou et al. , 2019 ; Lee et al. , 2020 ) try to extract information from historical transition segments and use it as a surrogate for Z ( Figure 1a ) . However , in the unsupervised dynamics generalization setting , the extracted information from historical transition segments inevitably contains redundant information unrelated to the dynamics in transition segments . The redundant information would cause the surrogate for Z to lose a crucial property that characterizes Z : Z should be similar in the same environment and dissimilar in different environments . As shown in Figure 1b , the learned environment-specified information of a state-of-theart method CaDM ( Lee et al. , 2020 ) does not form clear clusters for different environments . Because the learned Ẑ fails to represent environmental information , the learned dynamics prediction function will deviate from the true one , which undermines the generalization ability . To alleviate this problem , TMCL ( Seo et al. , 2020 ) directly clusters the environments by introducing multiple prediction heads , i.e . multiple prediction functions . However , TMCL needs to choose the proper prediction head for each new environment , making it hard to be deployed into the scenario with consistently changing environments , e.g . robots walk in the terrain which is constantly changing . To avoid adaptation at the deployment time , we thus need to learn a single generalized prediction function f̂ . To ensure that f̂ can learn modals of transition dynamics in different environments , we need to cluster Z according to their belonging environment . In this paper , we provide an explicit and interpretable description to learn Z as a vector Ẑ ( i.e . the estimation of Z ) from the history transition segments . To cluster Ẑ from the same environment , we introduce a relational head module as a learnable function to enforce the similarity between Ẑs learned from the same environments . However , because environment label is not available , we can only cluster the Ẑs from the same trajectory , so we then propose an interventional prediction module to identify the probability of a pair of ẑi , ẑj belonging to the same environment through estimating Ẑ ’ s direct causal effect on next states prediction by do-calculus ( Pearl , 2000 ) . Because Z from the same environment surely have the same causal effect , we can directly maximize the similarity of Ẑs with the similar causal effect using the relational head , and thus can cluster Ẑ according to the estimated environmental similarity and alleviate the redundant information that varies in an environment , e.g . historical states and actions . In the experiments , we evaluate our method on a range of tasks in OpenAI gym ( Brockman et al. , 2016 ) and Mujoco ( Todorov et al. , 2012 ) , and empirically show that Ẑ estimated by our method enjoy less redundant information than baselines , and the experimental results show that our method significantly reduces the model prediction errors and outperforms the state-of-art model-based RL methods without any adaptation step on a new environment , and even achieve comparable results with the method directly cluster Ẑ with the true environment label . 2 RELATED WORK . Dynamics Generalization in MBRL Several meta-learning-based MBRL methods are proposed Nagabandi et al . ( 2018a ; b ) ; Sæmundsson et al . ( 2018 ) to adapt the MBRL into environments with unseen dynamics by updating model parameters via a small number of gradient updates Finn et al . ( 2017 ) or updating hidden representations of a recurrent model Doshi-Velez & Konidaris ( 2016 ) , and then Wang & van Hoof ( 2021 ) proposes the graph structured model to model the dynamics accurately to improve dynamics forecasting . Ball et al . ( 2021 ) focus on the offline setting and proposes an augmented model method to achieve zero-shot generalization . Lee et al . ( 2020 ) ; Seo et al . ( 2020 ) try to learn a generalized dynamics model by incorporating context information or clustering dynamics implicitly using multi-choice learning , aiming to adapt any dynamics without training . However , how to explicitly learn the meaningful dynamics change information remains a big challenge . Relational Learning Reasoning relations between different entities is a way to build knowledge of this world for intelligent agents Kemp & Tenenbaum ( 2008 ) . In the past decades , relational paradigm have been applied to a wide range of deep learning-based application , e.g. , reinforcement learning Zambaldi et al . ( 2018 ) , question answer Santoro et al . ( 2017 ) ; Raposo et al . ( 2017 ) , graph neural network Battaglia et al . ( 2018 ) , sequential streams Santoro et al . ( 2018 ) , few-shot learning Sung et al . ( 2018 ) , object detection Hu et al . ( 2018 ) and self-supervised learning Patacchiola & Storkey ( 2020 ) . Different from previous methods that perform binary relational reasoning on entities , our method can also perform multiplies relations between entities through the learned similarity of entities , and thus can learn more compact and meaningful entity representation . Causality in Reinforcement Learning Many works focus on the intersection area of reinforcement learning and causal inference . For example , some works aims to alleviate the causal confusion problem in the imitation learning ( de Haan et al. , 2019 ; Zhang et al. , 2020c ; Kumor et al. , 2021 ) , batch learning ( Bannon et al. , 2020 ) , and partial observability settings ( Forney et al. , 2017 ; Kallus & Zhou , 2018 ; Zhang et al. , 2019a ; Lu et al. , 2018 ) in the online environment ( Lyle et al. , 2021 ; Zhang & Bareinboim , 2018 ) , ( Wang et al. , 2020 ) also try to apply causal inference in the offline setting , where the observational data is always confounded . Lee & Bareinboim ( 2018 ) ; Bareinboim et al . ( 2015 ) ; Lattimore et al . ( 2016 ) ; Mozifian et al . ( 2020 ) ; Volodin et al . ( 2020 ) also explore how to design an optimal intervention policy in bandits or RL settings . In addition , ( Zhang et al. , 2020a ; b ) improve the generalization ability of state abstraction . Different from these methods , we focus on the setting of unsupervised dynamics generalization , and measure the direct causal effect ( Pearl , 2013 ) between Ẑ and the next state to estimate the probability of them belonging to the same environment . 3 METHODS . In the section , we first introduce the formulation of the dynamic generalization problem . Then we present the details on how relational intervention approach learns the environment-specified factors . 3.1 PROBLEM SETUP . The standard reinforcement learning task can be formalized as a Markov decision process ( MDP ) M = ( S , A , r , p , γ , ρ0 ) over discrete time ( Puterman , 2014 ; Sutton & Barto , 2018 ) , where S , A , γ ∈ ( 0 , 1 ] , ρ0 are state space , action space , the reward discount factor , and the initial state distribution , respectively . The reward function r : S ×A → R specifies the reward at each timestep t given st and at , and transition dynamics p ( st+1|st , at ) gives the next state distribution conditioned on the current state st and action at . The goal of RL is to learn a policy π ( ·|s ) mapping from state s ∈ S over the action distribution to maximize the cumulative expected return over timesteps Est∈S , at∈A [ ∑∞ t=0 γ t r ( st , at ) ] . In model-based RL , a model f is used to approximate the transition dynamics p , and then f can provide training data to train policy π or predict the future sequences for planning . Benefiting from data provided by learned dynamics model f , model-based RL has higher data efficiency and better planing ability compared with model-free RL . Here we consider the unsupervised dynamics generalization problem in model-based RL , where we are given K training MDPs { Mtri } Ki=1 and L test MDPs { Mtej } Lj=1 that have the same state and action space but disjoint dynamics functions , and we randomly sample several MDPs from training MDPs in each training iteration . We assume that all these MDPs have a finite number of dynamics functions , meaning that the MDPs can be categorized into a finite number of environments and the MDPs in each environment share the same dynamics function but the environment id of MDPs is unavailable in the training process . In the context of model-based RL , how to learn a generalized dynamics model f is the key challenge to solve unsupervised dynamics generalization problem .

Summary: The paper studies multi-task generalization in model-based reinforcement learning (MBRL). One of the standard ways to approach the problem is given by inferring a latent variable Z encoding each task and then conditioning the dynamics model on it. The authors propose to encode segments of state-action trajectories into Z vectors and maximize similarity between Zs from the same trajectory. Using the mechanism, the paper then studies the ability of the agents to generalize to unseen variations of the training environments. For example, if the agent was trained on MuJoCo HalfCheetah with body mass 1, it is asked to generalize to body masses 0.5 and 1.5. The paper compares the method against the baselines across 3 axes: in terms of dynamics prediction error, in terms of returns on testing environments, and in terms of separability of the inferred Z for different environments.