Learning Robust State Abstractions for Hidden-Parameter Block MDPs

1 INTRODUCTION A key open challenge in AI research that remains is how to train agents that can learn behaviors that generalize across tasks and environments . When there is common structure underlying the tasks , we have seen that multi-task reinforcement learning ( MTRL ) , where the agent learns a set of tasks simultaneously , has definite advantages ( in terms of robustness and sample efficiency ) over the singletask setting , where the agent independently learns each task . There are two ways in which learning multiple tasks can accelerate learning : the agent can learn a common representation of observations , and the agent can learn a common way to behave . Prior work in MTRL has also leveraged the idea by sharing representations across tasks ( D ’ Eramo et al. , 2020 ) or providing pertask sample complexity results that show improved sample efficiency from transfer ( Brunskill & Li , 2013 ) . However , explicit exploitation of the shared structure across tasks via a unified dynamics has been lacking . Prior works that make use of shared representations use a naive unification approach that posits all tasks lie in a shared domain ( Figure 1 , left ) . On the other hand , in the single-task setting , research on state abstractions has a much richer history , with several works on improved generalization through the aggregation of behaviorally similar states ( Ferns et al. , 2004 ; Li et al. , 2006 ; Luo et al. , 2019 ; Zhang et al. , 2020b ) . In this work , we propose to leverage rich state abstraction models from the single-task setting , and explore their potential for the more general multi-task setting . We frame the problem as a structured super-MDP with a shared state space and universal dynamics model conditioned on a task-specific hidden parameter ( Figure 1 , right ) . This additional structure gives us better sample efficiency , both ∗Corresponding author : amy.x.zhang @ mail.mcgill.ca theoretically , compared to related bounds ( Brunskill & Li , 2013 ; Tirinzoni et al. , 2020 ) and empirically against relevant baselines ( Yu et al. , 2020 ; Rakelly et al. , 2019 ; Chen et al. , 2018 ; Teh et al. , 2017 ) . We learn a latent representation with smoothness properties for better few-shot generalization to other unseen tasks within this family . This allows us to derive new value loss bounds and sample complexity bounds that depend on how far away a new task is from the ones already seen . We focus on multi-task settings where dynamics can vary across tasks , but the reward function is shared . We show that this setting can be formalized as a hidden-parameter MDP ( HiP-MDP ) ( DoshiVelez & Konidaris , 2013 ) , where the changes in dynamics can be defined by a latent variable , unifying dynamics across tasks as a single global function . This setting assumes a global latent structure over all tasks ( or MDPs ) . Many real-world scenarios fall under this framework , such as autonomous driving under different weather and road conditions , or even different vehicles , which change the dynamics of driving . Another example is warehouse robots , where the same tasks are performed in different conditions and warehouse layouts . The setting is also applicable to some cases of RL for medical treatment optimization , where different patient groups have different responses to treatment , yet the desired outcome is the same . With this assumed structure , we can provide concrete zero-shot generalization bounds to unseen tasks within this family . Further , we explore the setting where the state space is latent and we have access to only high-dimensional observations , and we show how to recover robust state abstractions in this setting . This is , again , a highly realistic setting in robotics when we do not always have an amenable , Lipschitz low-dimensional state space . Cameras are a convenient and inexpensive way to acquire state information , and handling pixel observations is key to approaching these problems . A block MDP ( Du et al. , 2019 ) provides a concrete way to formalize this observation-based setting . Leveraging this property of the block MDP framework , in combination with the assumption of a unified dynamical structure of HiP-MDPs , we introduce the hidden-parameter block MDP ( HiP-BMDP ) to handle settings with high-dimensional observations and structured , changing dynamics . Key contributions of this work are a new viewpoint of the multi-task setting with same reward function as a universal MDP under the HiP-BMDP setting , which naturally leads to a gradientbased representation learning algorithm . Further , this framework allows us to compute theoretical generalization results with the incorporation of a learned state representation . Finally , empirical results show that our method outperforms other multi-task and meta-learning baselines in both fast adaptation and zero-shot transfer settings . 2 BACKGROUND . In this section , we introduce the base environment as well as notation and additional assumptions about the latent structure of the environments and multi-task setup considered in this work . A finite1 , discrete-time Markov Decision Process ( MDP ) ( Bellman , 1957 ; Puterman , 1995 ) is a tuple 〈S , A , R , T , γ〉 , where S is the set of states , A is the set of actions , R : S × A → R is the reward function , T : S × A → Dist ( S ) is the environment transition probability function , and γ ∈ [ 0 , 1 ) is the discount factor . At each time step , the learning agent perceives a state st ∈ S , takes an action at ∈ A drawn from a policy π : S ×A → [ 0 , 1 ] , and with probability T ( st+1|st , at ) enters next state st+1 , receiving a numerical reward Rt+1 from the environment . The value function of policy π is defined as : Vπ ( s ) = Eπ [ ∑∞ t=0 γ tRt+1|S0 = s ] . The optimal value function V ∗ is the maximum value function over the class of stationary policies . Hidden-Parameter MDPs ( HiP-MDPs ) ( Doshi-Velez & Konidaris , 2013 ) can be defined by a tuple M : 〈S , A , Θ , Tθ , R , γ , PΘ〉 where S is a finite state space , A a finite action space , Tθ describes the transition distribution for a specific task described by task parameter θ ∼ PΘ , R is the reward function , γ is the discount factor , and PΘ the distribution over task parameters . This defines a family of MDPs , where each MDP is described by the parameter θ ∼ PΘ . We assume that this parameter θ is fixed for an episode and indicated by an environment id given at the start of the episode . Block MDPs ( Du et al. , 2019 ) are described by a tuple 〈S , A , X , p , q , R〉 with an unobservable state space S , action spaceA , and observable space X . p denotes the latent transition distribution p ( s′|s , a ) for s , s′ ∈ S , a ∈ A , q is the ( possibly stochastic ) emission mapping that emits the observations q ( x|s ) for x ∈ X , s ∈ S , and R the reward function . We are interested in the setting where this 1We use this assumption only for theoretical results , but our method can be applied to continuous domains . mapping q is one-to-many . This is common in many real world problems with many tasks where the underlying states and dynamics are the same , but the observation space that the agent perceives can be quite different , e.g . navigating a house of the same layout but different decorations and furnishings . Assumption 1 ( Block structure ( Du et al. , 2019 ) ) . Each observation x uniquely determines its generating state s. That is , the observation space X can be partitioned into disjoint blocks Xs , each containing the support of the conditional distribution q ( ·|s ) . Assumption 1 gives the Markov property inX , a key difference from partially observable MDPs ( Kaelbling et al. , 1998 ; Zhang et al. , 2019 ) , which has no guarantee of determining the generating state from the history of observations . This assumption allows us to compute reasonable bounds for our algorithm in k-order MDPs2 ( which describes many real world problems ) and avoiding the intractability of true POMDPs , which have no guarantees on providing enough information to sufficiently predict future rewards . A relaxation of this assumption entails providing less information in the observation for predicting future reward , which will degrade performance . We show empirically that our method is still more robust to a relaxation of this assumption compared to other MTRL methods . Bisimulation is a strict form of state abstraction , where two states are bisimilar if they are behaviorally equivalent . Bisimulation metrics ( Ferns et al. , 2011 ) define a distance between states as follows : Definition 1 ( Bisimulation Metric ( Theorem 2.6 in Ferns et al . ( 2011 ) ) ) . Let ( S , A , P , r ) be a finite MDP and met the space of bounded pseudometrics on S equipped with the metric induced by the uniform norm . Define F : met 7→ met by F ( d ) ( s , s′ ) = max a∈A ( |ras − ras′ |+ γW ( d ) ( P as , P as′ ) ) , where W ( d ) is the Wasserstein distance between transition probability distributions . Then F has a unique fixed point d̃ which is the bisimulation metric . A nice property of this metric d̃ is that difference in optimal value between two states is bounded by their distance as defined by this metric . Theorem 1 ( V ∗ is Lipschitz with respect to d̃ ( Ferns et al. , 2004 ) ) . Let V ∗ be the optimal value function for a given discount factor γ . Then V ∗ is Lipschitz continuous with respect to d̃ with Lipschitz constant 11−γ , |V ∗ ( s ) − V ∗ ( s′ ) | ≤ 1 1− γ d̃ ( s , s′ ) . Therefore , we see that bisimulation metrics give us a Lipschitz value function with respect to d̃ . For downstream evaluation of the representations we learn , we use Soft Actor Critic ( SAC ) ( Haarnoja et al. , 2018 ) , an off-policy actor-critic method that uses the maximum entropy framework for soft policy iteration . At each iteration , SAC performs soft policy evaluation and improvement steps . The policy evaluation step fits a parametric soft Q-function Q ( st , at ) using transitions sampled from the replay buffer D by minimizing the soft Bellman residual , J ( Q ) = E ( st , st , rt , st+1 ) ∼D [ ( Q ( st , at ) − rt − γV̄ ( xt+1 ) ) 2 ] . The target value function V̄ is approximated via a Monte-Carlo estimate of the following expectation , V̄ ( xt+1 ) = Eat+1∼π [ Q̄ ( xt+1 , at+1 ) − α log π ( at+1|st+1 ) ] , where Q̄ is the target soft Q-function parameterized by a weight vector obtained from an exponentially moving average of the Q-function weights to stabilize training . The policy improvement step then attempts to project a parametric policy π ( at|st ) by minimizing KL divergence between the policy and a Boltzmann distribution induced by the Q-function , producing the following objective , J ( π ) = Est∼D [ Eat∼π [ α log ( π ( at|st ) ) −Q ( st , at ) ] ] . 2Any k-order MDP can be made Markov by stacking the previous k observations and actions together .

Summary: This paper studies a family of Markov Decision Process (MDP) models with a low-dimensional unobserved state, called the block MDP. The authors assume that system dynamics of the underlying MDP is sufficiently summarized by a parameter $\theta$. This learning setting could be seen as a combination of Block MDP and the Hidden Parameter MDP; hence the name HiP-BMDP.