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: The authors introduce a new type of state abstraction called model-invariance by integrating sparsity with causal invariance in the single-task RL scenarios. In the linear settings, they learn the model-invariance states by directly applying ICP and prove that an optimal policy can be learned over them. In the nonlinear settings, they instead propose a practical method to approximately learn such model-invariance states.