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 proposes a new class of decision problems, the exogenous block-MDP (EX-BMDP), and an algorithm for acquiring minimal state representations for EX-BMDPs. EX-BMDPs are an extension of block MDPs that allows for additional "exogenous" state components that may have arbitrary Markovian dynamics, but cannot be influenced by actions, including through the actions' effects on "endogenous" state. The proposed algorithm, PPE, is able to recover a mapping from observed states to endogenous states (i.e. ignoring the latent distractor variables) in the setting where the endogenous state dynamics and initial state distribution are near-deterministic. It does this by comparing the set of policies that can reach each pair of states: informally, if $p(\pi \mid s_1) = p(\pi \mid s_2)$ for all policies $\pi$ (assuming a uniform $p(\pi)$), then $s_1$ and $s_2$ are treated as equivalent. In the deterministic case, policies reduce to action sequences ("paths"). Illustrative experiments on simple domains show that the proposed approach is able to quickly recover a state representation that depends mostly on the endogeneous component.