Provably Filtering Exogenous Distractors using Multistep Inverse Dynamics

1 INTRODUCTION . In many real-world applications such as robotics there can be large disparities in the size of agent ’ s observation space ( for example , the image generated by agent ’ s camera ) , and a much smaller latent state space ( for example , the agent ’ s location and orientation ) governing the rewards and dynamics . This size disparity offers an opportunity : how can we construct reinforcement learning ( RL ) algorithms which can learn an optimal policy using samples that scale with the size of the latent state space rather than the size of the observation space ? Several families of approaches have been proposed based on solving various ancillary prediction problems including autoencoding ( Tang et al. , 2017 ; Hafner et al. , 2019 ) , inverse modeling ( Pathak et al. , 2017 ; Burda et al. , 2018 ) , and contrastive learning ( Laskin et al. , 2020 ) based approaches . These works have generated some significant empirical successes , but are there provable ( and hence more reliable ) foundations for their success ? More generally , what are the right principles for learning with latent state spaces ? In real-world applications , a key issue is robustness to noise in the observation space . When noise comes from the observation process itself , such as due to measurement error , several approaches have been recently developed to either explicitly identify ( Du et al. , 2019 ; Misra et al. , 2020 ; Agarwal et al. , 2020a ) or implicitly leverage ( Jiang et al. , 2017 ) the presence of latent state structure for provably sample-efficient RL . However , in many real-world scenarios , the observations consist of many elements ( e.g . weather , lighting conditions , etc . ) with temporally correlated dynamics ( see e.g . Figure 1 and the example below ) that are entirely independent of the agent ’ s actions and rewards . The temporal dynamics of these elements precludes us from treating them as uncorrelated noise , and as such , most previous approaches resort to modeling their dynamics . However , this is clearly wasteful as these elements have no bearing on the RL problem being solved . As an example , consider the setting in Figure 1 . An agent is walking in a park on a lonely sidewalk next to a pond . The agent ’ s observation space is the image generated by its camera , the latent { yefroni , dimisra , akshaykr , jcl } @ microsoft.com , alekhagarwal @ google.com endogenous state is its position on the sidewalk , and the exogenous noise is provided by motion of ducks or people in the background , swaying of trees and changes in lighting conditions , typically unaffected by the agent ’ s actions . While there is a line of recent empirical work that aims to remove causally irrelevant aspects of the observation ( Gelada et al. , 2019 ; Zhang et al. , 2020 ) , theoretical treatment is quite limited ( Dietterich et al. , 2018 ) and no prior works address sample-efficient learning with provable guarantees . Given this , the key question here is : How can we learn using an amount of data scaling with just the size of the endogenous latent state , while ignoring the temporally correlated exogenous observation noise ? We initiate a formal treatment of RL settings where the learner ’ s observations are jointly generated by a latent endogenous state and an uncontrolled exogenous state , which is unaffected by the agent ’ s actions and does not affect the agent ’ s task . We study a subset of such problems called Exogenous Block MDPs ( EX-BMDPs ) , where the endogenous state is discrete and decodable from the observations . We first highlight the challenges in solving EX-BMDPs by illustrating the failures of many prior representation learning approaches ( Pathak et al. , 2017 ; Misra et al. , 2020 ; Jiang et al. , 2017 ; Agarwal et al. , 2020a ; Zhang et al. , 2020 ) . These failure happen either due to creating too many latent states , such as one for each combination of ducks and passers-by in the example above leading to sample inefficiency in exploration , or due to lack of exhaustive exploration . We identify one recent approach developed by Du et al . ( 2019 ) with favorable properties for a class of EX-BMDPs with near-deterministic latent state dynamics . In Section 4 and Section 5 , we develop a variation of their algorithm and analyze its performance for near-deterministic EX-BMDPs . The algorithm , called Path Prediction and Elimination ( PPE ) , learns a form of multi-step inverse dynamics by predicting the identity of the path that generates an observation . For near-deterministic EX-BMDPs , we prove that PPE successfully explores the environment using O ( ( SA ) 2H log ( |F|/δ ) ) samples where S is the size of the latent endogenous state space , A is the number of actions , H is the horizon and F is a function class employed to solve a maximum likelihood problem of inferring the path which generated an observation . Several prior works ( Gregor et al. , 2016 ; Paster et al. , 2020 ) have also considered a multi-step inverse dynamics approach to learn a near optimal policy . However , these works do not consider the EX-BMDP model . Further , it is unknown whether these algorithms have guarantees for EX-BMDP similar to what we present for PPE . Theoretical analysis of the performance of these algorithms in the presence of exogenous noise is an interesting future work direction . Empirically , in Section 6 , we demonstrate the performance of PPE and various prior baselines in a challenging exploration problem with exogenous noise . We show that baselines fail to decode the endogenous state and either end up over-abstracting ( creating an abstract state which aliases essentially-different states ) or under-abstracting ( creating essentially the same clones of states ) . We further , show that PPE is able to recover the latent endogenous model in a visually complex navigation problem , in accordance with the theory . 2 EXOGENOUS BLOCK MDP SETTING . We introduce a novel Exogenous Block Markov Decision Process ( EX-BMDP ) setting to model environments with exogenous noise . We briefly describe notations before presenting a formal treatment of EX-BMDP . Notations . For a given set U , we use ∆ ( U ) to denote the set of all probability distributions over U . For a given natural number N ∈ N , we use the notation [ N ] to denote the set { 1 , 2 , · · · , N } . Lastly , for a probability distribution p ∈ ∆ ( U ) , we define its support as supp p = { u | p ( u ) > 0 , u ∈ U } . We start with describing the Block Markov Decision Process ( BMDP ) Du et al . ( 2019 ) . This process consists of a finite set of observations X , a set of latent states Z with cardinality Z , a finite set of actions A with cardinality A , a transition function T : Z × A → ∆ ( Z ) , an emission function q : Z × A → ∆ ( X ) , a reward function R : X × A → [ 0 , 1 ] , a horizon H ∈ N , and a start state distribution µ ∈ ∆ ( Z ) . The agent interacts with the environment by repeatedly generating H-step trajectories ( z1 , x1 , a1 , r1 , · · · , zH , xH , aH , rH ) where z1 ∼ µ ( · ) and for every h ∈ [ H ] we have xh ∼ q ( · | zh ) , rh = R ( xh , ah ) , and if h < H , then zh+1 ∼ T ( · | zh , ah ) . The agent does not observe the states ( z1 , · · · , zH ) , instead receiving only the observations ( x1 , · · · , xH ) and rewards ( r1 , · · · , rH ) . We assume that the emission distributions of any two latent states are disjoint , usually referred as the block assumption : supp ( q ( · | z1 ) ) ∩ supp ( q ( ·|z2 ) ) = ∅ when z1 6= z2 . The agent chooses actions using a policy π : X → ∆ ( A ) . We also define the set of non-stationary policies ΠNS = ΠH as a H-length tuple , with ( π1 , · · · , πH ) ∈ ΠNS denoting that the action at time step h is taken as ah ∼ πh ( · | xh ) . The value V ( π ) of a policy π is the expected episodic sum of rewards V ( π ) : = Eπ [ ∑H h=1R ( xh , ah ) ] . The optimal policy is given by π ? = arg maxπ∈ΠNS V ( π ) . We denote by Ph ( x|π ) the probability distribution over observations x at time step h when following a policy π. Lastly , we refer to an open loop policy as an element in all AH sequences of actions . An open loop policy follows a pre-determined sequence of actions { a1 , .. , aH } for H time steps , unaffected by state information . Given the aforementioned definitions , we define an EX-BMDP as follows : Definition 1 ( Exogenous Block Markov Decision Processes ) . An EX-BMDP is a BMDP such that the latent state can be decoupled into two parts z = ( s , ξ ) where s ∈ S is the endogenous state and ξ ∈ Ξ is the exogenous state . For z ∈ Z the initial distribution and transition functions are decoupled , that is : µ ( z ) = µ ( s ) µξ ( ξ ) , and T ( z′ | z , a ) = T ( s′ | s , a ) Tξ ( ξ′ | ξ ) . The observation space X can be arbitrarily large to model which could be a high-dimensional real vector denoting an image , sound , or haptic data in an EX-BMDP . The endogenous state s captures the information that can be manipulated by the agent . Figure 1 , center , visualizes the transition dynamics factorization . We assume that the set of all endogenous states S is finite with cardinality S. The exogenous state ξ captures all the other information that the agent can not control and does not affect the information it can manipulate . Again , we make no assumptions on the exogenous dynamics nor on its cardinality |Ξ| which may be arbitrarily large . We note that the block assumption of the EX-BMDP implies the existence of two inverse mappings : φ ? e : X → S to map an observation to its endogenous state , and φ ? g : X → Ξ to map it to its exogenous state . Justification of assumptions . The block assumption has been made by prior work ( e.g. , Du et al . ( 2019 ) , Zhang et al . ( 2020 ) ) to model many real-world settings where the observation is rich , i.e. , it contains enough information to decode the latent state . The decoupled dynamics assumption made in the EX-BMDP setting is a natural way to characterize exogenous noise ; the type of noise that is not affected by our actions nor affects the endogenous state but may have non-trivial dynamic . This decoupling captures the movement of ducks , captured in the visual field of the agent in Figure 1 , and many additional exogenous processes ( e.g. , movement of clouds in a navigation task ) . Goal . Our formal objective is reward-free learning . We wish to find a set of policies , we call a policy cover , that can be used to explore the entire state space . Given a policy cover , and for any reward function , we can find a near optimal policy by applying dynamic programming ( e.g. , Bagnell et al . ( 2004 ) ) , policy optimization ( e.g. , Kakade and Langford ( 2002 ) ; Agarwal et al . ( 2020b ) ; Shani et al . ( 2020 ) ) or value ( e.g. , Antos et al . ( 2008 ) ) based methods . Definition 2 ( α-policy cover ) . Let Ψh be a finite set of non-stationary policies . We say Ψh is an α-policy cover for the hth time step if for all z ∈ Z it holds that maxπ∈Ψh Ph ( z|π ) ≥ maxπ∈ΠNS Ph ( z|π ) − α . If α = 0 we call Ψh a policy cover . For standard BMDPs the policy cover is simply the set of policies that reaches each latent state of the BMDP ( Du et al. , 2019 ; Misra et al. , 2020 ; Agarwal et al. , 2020a ) . Thus , for a BMDP , the cardinality of the policy cover scales with |Z| . The structure of EX-BMDPs allows to reduce the size of the policy cover significantly to |S| |Z| = |S| |Ξ| when the size of the exogenous state space is large . Specifically , we show that the set of policies that reach each endogenous state , and do not depend on the exogenous part of the state is also a policy cover ( see Appendix B , Proposition 4 ) . Further , the proposed algorithm PPE , learns such a policy cover and requires number of samples that depends polynomially on |S| instead of |Z| or |X | .

Summary: This paper introduces a model called the Exogenous Block MDP (EX-BMDP) where the latent state contains both controllable (endogenous) and uncontrollable (exogenous) elements. The paper proposes an algorithm to find a policy cover with sample complexity that depends only on the size of the endogenous state rather than the observation or exogenous state. The algorithm works by training a classifier to predict the actions that were taken to get to a state. Since the exogenous state is not affected by actions, states for which the classifier returns similar results are likely the same endogenous state. The algorithm builds up a set of policies (sequences of actions are sufficient in the near-deterministic setting) which visit unique endogenous states by using the classifier to deduplicate redundant action sequences. Besides proving the sample efficiency of their approach, experiments are provided which show that this algorithm performs better than baseline approaches in terms of performance as well as in ability to decode the state from its representation in a simple combination lock environment as well as a grid world with distractors.