Model-Invariant State Abstractions for Model-Based Reinforcement Learning

Accuracy and generalization of dynamics models is key to the success of modelbased reinforcement learning ( MBRL ) . As the complexity of tasks increases , learning accurate dynamics models becomes increasingly sample inefficient . However , many complex tasks also exhibit sparsity in dynamics , i.e. , actions only have a local effect on the system dynamics . In this paper , we exploit this property with a causal invariance perspective in the single-task setting , introducing a new type of state abstraction called model-invariance . Unlike previous forms of state abstractions , model-invariant state abstraction leverages causal sparsity over state variables . This allows for compositional generalization to unseen states , something that nonfactored forms of state abstractions can not do . We prove that an optimal policy can be learned over exact model-invariance state abstraction and show improved generalization in a simple toy domain . Next , we propose a practical method to approximately learn a model-invariant representation for complex domains and validate our approach by showing improved modelling performance over standard maximum likelihood approaches on challenging tasks , such as the MuJoCo-based Humanoid . Finally , within the MBRL setting we show strong performance gains with respect to sample efficiency across a host of continuous control tasks . 1 INTRODUCTION . Model-based reinforcement learning or MBRL [ 4 , 15 ] is a popular framework for data-efficient learning of control policies . At the core of MBRL is learning an environmental dynamics model and using it to : 1 ) fully plan [ 14 , 11 ] , 2 ) augment the data used by a model-free solver [ 51 , 54 ] , or 3 ) use as an auxiliary task while training [ 35 , 58 ] . However , learning a dynamics model — similar to other supervised learning problems — suffers from the issue of generalization , since the data we train on is not necessarily the data we test on . This is a persistent issue that is worsened in MBRL as even a small inaccuracy in the dynamics model or changes in the control policy can result in visiting completely unexplored parts of the state space [ 1 ] . This advocates for the need to learn models capable of generalizing well . Various workarounds for this issue have been explored in the past ; for example , combining local but simple models with global , more expressive models [ 36 , 22 ] , using an ensemble of models to handle uncertainty in estimates [ 11 , 33 ] and coupling the model and policy learning processes [ 34 ] so that the model is always accurate to a certain threshold . However , these approaches do not stem from a representation learning viewpoint and thus fail to leverage special structure in dynamics for better generalization . This paper studies how to improve generalization capabilities through careful state abstraction . In particular , we leverage two existing concepts to define a new kind of state abstraction . The first concept is that many real world problems exhibit sparsity in the local dynamics — given a set of state variables , each variable at timestep t+ 1 only depends on a small subset ( i.e. , local ) of all variables in the previous timestep t ( see Figure 1 ) . The second concept is the principle of causal invariance , which dictates that given a set of candidate features , we should aim to build representations that comprise only those features that are consistently necessary for predicting the target variable of interest across different interventions [ 46 ] . In the MBRL context , we can cast model learning with causal invariance as a supervised objective where the target variables are the next state variables and input features are the current state and action variables ( the probable set of causal predictors of the target ) . Intuitively , since we learn a predictor that only takes into account invariant features that consistently predict the target variable well , it is likely to contain the true causal features , and therefore , will generalize well to all possible shifts in the data distribution . The two concepts of sparsity and causality are intertwined in that they both are forms of inductive biases that surround the agent dynamics [ 23 ] . This paper shows how causal invariance tools can be effectively used to exploit sparsity in dynamics , leading to improved model generalization . Given basic exploratory assumptions , we analyze this question theoretically and show empirically that we can learn a model that generalizes well on state distributions induced by policies distinct from the ones used while learning it . To do this , we introduce a new state abstraction , model-invariance , which leverages sparsity over state variables . We connect this abstraction viewpoint to a concrete problem of generalization in model-based RL , that of arising spurious correlations in the dynamics model , even when all state variables are task relevant . Having built enough intuition , we then introduce a practical method that approximates learning a model-invariant representation for more complex domains that use function approximation ( i.e. , neural networks ( NNs ) ) . We empirically observe that model invariance leads to better model generalization for domains such as the MuJoCo-based Humanoid . Our method is simple to implement and shows strong performance on multiple MuJoCo tasks , outperforming state of the art model-based methods in the low-data regime . 2 PRELIMINARIES . RL Setup . We consider the agent ’ s interaction with the environment as a discrete time -discounted Markov Decision Process ( MDP ) [ 47 ] M = X , A , P , R , , µ0 , where X ✓ Rd is a finite but arbitrarily large state space and A is the and action space ; P ⌘ P x 0|x , a is the transition kernel ; R ⌘ r ( x , a ) is the reward function with the maximum value of Rmax ; 2 [ 0 , 1 ) is the discount factor ; and µ0 is the initial state distribution . Let ⇡ : X ! A be a stationary memoryless policy , where A is the set of probability distributions on A . The value function of a policy ⇡ at a state x 2 X is defined as V ⇡ ( x ) ⌘ E [ P t 0 t r ( xt , at ) |x0 = x , ⇡ ] . Similarly , the action-value function of ⇡ is defined as Q⇡ ( x , a ) = E [ P t 0 t r ( xt , at ) |x0 = x , a0 = a , ⇡ ] . The Bellman optimality operator T : R|X⇥A| ! R|X⇥A| is defined as T Q ( x , a ) = r ( x , a ) + ⌦ P ( ·|x , a ) , maxa0 Q ( · , a0 ) ↵ . In this work , we assume that all d state variables x1 , x2 , ... , xd are useful for the task in hand , i.e. , we are given the full state . Furthermore , we assume that the transition dynamics over the full state are factorized . More formally : Assumption 1 . ( Transition Factorization ) For given full state vectors xt , xt+1 2 X , action a 2 A , and x i denoting the i th dimension of state x we have P xt+1|xt , a = Q i P x i t+1|xt , a . Note that this is a weaker assumption than factored MDPs [ 31 , 24 ] as we do not assume a corresponding factorization of the reward function . Invariant Causal Prediction . Invariant causal prediction ( ICP ) [ 46 ] considers the problem of learning an invariant representation w.r.t . spurious correlations that arise due to noise in the underlying causal model ( unknown ) describing a given system . The key observation is that if one considers the direct causal parents of a response/target variable of interest ( Y ) , then the conditional distribution of Y given these direct causes PA ( Y ) does not change across interventions on any variable except Y . Therefore , ICP suggests collecting data into different environments ( corresponding to different interventions ) , and to output the set of variables Xi for which a learned predictor of Y remains the same given Xi across the multiple environments with high probability . State Abstractions . State abstractions allow us to map behaviorally-equivalent states into a single abstract state , thus simplifying the learning problem , which then makes use of the ( potentially much smaller set of ) abstract states instead of the original states [ 5 ] . In principle , any function approximation architecture can act as an abstraction , since it attempts to group similar states together . Therefore , exploring the properties of a representation learning scheme as a state abstraction help develop stronger intuition for building practical algorithms . In the next section , we build our theory based on this connection . A well known state abstraction is bisimulation [ 19 , 48 , 37 ] . Formally , an abstraction : X 7 ! S is a bisimulation if for any two states x1 , x2 and next state x 2 X , abstract state s 2 S , a 2 A where ( x1 ) = ( x2 ) , we have R ( x1 , a ) = R ( x2 , a ) , andP x2 1 ( s ) P x|x1 , a = P x2 1 ( s ) P x|x2 , a . Since an exact equivalence is not practical , prior work deals with approximate variants through the notion of ✏-closeness [ 29 ] . 3 MODEL INVARIANCE AND ABSTRACTIONS . Current state abstractions such as bisimulation do not provide support for sparse structures . Instead , closeness in abstract states is defined based on probabilities of all state variables together , i.e. , x . To remedy this , we introduce a new state abstraction , called model-invariance , which is specific to each state variable xi . Formally , an invariant abstraction i is one which has the same transition probability of the ith next state variable xi for any two given states x1 and x2 , i.e. , P x i|x1 , a = P x i|x2 , a . Note that if we assume factored rewards , we can define a corresponding reward-based invariant abstraction that parallels the bisimulation abstraction more closely , but we focus here on the rewardfree setting . Since it is impractical to ensure this equivalence exactly , we can use an approximate definition which ensures an ✏-closeness . Definition 1 . ( Approximate Model Invariance ) is ✏i , P -model-invariant if for each index i , sup a2A , x , x1 , x22X , ( x1 ) : = ( x2 ) P x i|x1 , a P x i|x2 , a ✏i , P . is ✏R-model-invariant if ✏R = supa2A , x1 , x22X , ( x1 ) = ( x2 ) R ( x1 , a ) R ( x2 , a ) . Based on the definition , it intuitively seems like applying the abstraction over the original transition distribution should be close to the transition distribution over abstract states . This can be precisely written in the following lemma , which is a necessary result in ensuring that an optimal policy can still be learnt over a model-invariant abstraction : Lemma 1 . ( Model Error Bound ) Let be an ✏i , P -approximate model-invariant abstraction on MDP M . Given any distributions psi : si 2 i ( X ) where psi is supported on 1 ( si ) and ps = Q d i=1 psi , we define M = ( ( X ) , A , P , R , ) where P ( s , a ) = Ex⇠ps ⇥ P ( ·|x , a ) ⇤ and R = Ex⇠ps⇥ R ( x , a ) ⇤ . Then for any x 2 X , a 2 A , kP ( s , a ) P ( x , a ) k dX i=1 ✏i , P , where P denotes the lifted version of P , where we take the next-step transition distribution from observation space X and lift it to latent space S ( Proof in Appendix 2 ) . P refers to the transition probability of a MDP that acts on the states ( X ) , rather than the original MDP . Finally , note that we are particularly concerned with the case where each xi is atomic in nature , i.e. , it is not further divisible . Such a property ensures that model-invariance does not collapse to bisimulation1 .

Summary: This work presents a method for learning representations for model based reinforcement learning by leveraging assumed sparsity in the causal graph to improve the generalization performance of the learned model in the low-data regime. The authors propose a method which, at a high level, aims to utilize data collected under different conditions (e.g. policies) to extract a representation under which the transition model is approximately invariant. Furthermore, they provide a bound on difference between the optimal value function on the learned representation and the true optimal value function as a function of the approximation quality of the invariant model. They show how invariant causal prediction can be used in the the linear tabular setting to yield improved performance relative to learning a model directly in the original state space. Next, they present an approach to achieve a similar objective for the nonlinear case, presenting a neural network architecture based on that proposed in ReLIC, and show that the proposed architecture improves the low-data performance in learning dynamics models, and leads to improved performance relative to baselines.