MICo: Improved representations via sampling-based state similarity for Markov decision processes

1 Introduction . The success of reinforcement learning ( RL ) algorithms in large-scale , complex tasks depends on forming useful representations of the environment with which the algorithms interact . Feature selection and feature learning has long been an important subdomain of RL , and with the advent of deep reinforcement learning there has been much recent interest in understanding and improving the representations learnt by RL agents . Much of the work in representation learning has taken place from the perspective of auxiliary tasks [ Jaderberg et al. , 2017 , Bellemare et al. , 2017 , Fedus et al. , 2019 ] ; in addition to the primary reinforcement learning task , the agent may attempt to predict and control additional aspects of the environment . Auxiliary tasks shape the agent ’ s representation of the environment implicitly , typically via gradient descent on the additional learning objectives . As such , while auxiliary tasks continue to play an important role in improving the performance of deep RL algorithms , our understanding of the effects of auxiliary tasks on representations in RL is still in its infancy . In contrast to the implicit representation shaping of auxiliary tasks , a separate line of work on behavioural metrics , such as bisimulation metrics [ Desharnais et al. , 1999 , 2004 , Ferns et al. , 2004 , 2006 ] , aims to capture structure in the environment by learning a metric measuring behavioral similarity between states . Recent works have successfully used behavioural metrics to shape the representations of deep RL agents [ Gelada et al. , 2019 , Zhang et al. , 2021 , Agarwal et al. , 2021a ] . However , in practice behavioural metrics are difficult to estimate from both statistical and computational perspectives , and these works either rely on specific assumptions about transition dynamics to make the estimation tractable , and as such can only be applied to limited classes of environments , or are applied to more general classes of environments not covered by theoretical guarantees . The principal objective of this work is to develop new measures of behavioral similarity that avoid the statistical and computational difficulties described above , and simultaneously capture richer information about the environment . We introduce the MICo ( Matching under Independent Couplings ) distance , and develop the theory around its computation and estimation , making comparisons with existing metrics on the basis of computational and statistical efficiency . We demonstrate the usefulness ∗Equal contribution . Correspondence to Pablo Samuel Castro : psc @ google.com . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of the representations that MICo yields , both through empirical evaluations in small problems ( where we can compute them exactly ) as well as in two large benchmark suites : ( 1 ) the Arcade Learning Environment [ Bellemare et al. , 2013 , Machado et al. , 2018 ] , in which the performance of a wide variety of existing value-based deep RL agents is improved by directly shaping representations via the MICo distance ( see Figure 1 , left ) , and ( 2 ) the DM-Control suite [ Tassa et al. , 2018 ] , in which we demonstrate it can improve the performance of both Soft Actor-Critic [ Haarnoja et al. , 2018 ] and the recently introduced DBC [ Zhang et al. , 2021 ] ( see Figure 1 , right ) . 2 Background . Before describing the details of our contributions , we give a brief overview of the required background in reinforcement learning and bisimulation . We provide more extensive background in Appendix A. Reinforcement learning . We consider a Markov decision process ( X , A , γ , P , r ) defined by a finite state space X , finite action space A , transition kernel P : X ×A →P ( X ) , reward function r : X × A → R , and discount factor γ ∈ [ 0 , 1 ) . For notational convenience we will write P ax and rax for transitions and rewards , respectively . Policies are mappings from states to distributions over actions : π ∈ P ( A ) X and induce a value function V π : X → R defined via the recurrence : V π ( x ) : = Ea∼π ( x ) [ rax + γEx′∼Pax [ V π ( x′ ) ] ] . In RL we are concerned with finding the optimal policy π∗ = arg maxπ∈P ( A ) X V π from interaction with sample trajectories with an MDP , without knowledge of P or r ( and sometimes not even X ) , and the optimal value function V ∗ induced by π∗ . State similarity and bisimulation metrics . Various notions of similarity between states in MDPs have been considered in the RL literature , with applications in policy transfer , state aggregation , and representation learning . The bisimulation metric [ Ferns et al. , 2004 ] is of particular relevance for this paper , and defines state similarity in an MDP by declaring two states x , y ∈ X to be close if their immediate rewards are similar , and the transition dynamics at each state leads to next states which are also judged to be similar . This self-referential notion is mathematically formalised by defining the bisimulation metric d∼ as the unique fixed-point of the operator TK : M ( X ) →M ( X ) , whereM ( X ) = { d ∈ [ 0 , ∞ ) X×X : d symmetric and satisfies the triangle inequality } is the set of pseudometrics on X , given by TK ( d ) ( x , y ) = maxa∈A [ |rax − ray |+ γWd ( P ax , P ay ) ] . Here , Wd is the Kantorovich distance ( also known as the Wasserstein distance ) over the set of distributions P ( X ) with base distance d , defined by Wd ( µ , ν ) = infX∼µ , Y∼ν E [ d ( X , Y ) ] , for all µ , ν ∈P ( X ) , where the infimum is taken over all couplings of ( X , Y ) with the prescribed marginals [ Villani , 2008 ] . The mapping TK is a γ-contraction onM ( X ) under the L∞ norm [ Ferns et al. , 2011 ] , and thus by standard contraction mapping arguments analogous to those used to study value iteration , it has a unique fixed point , the bisimulation metric d∼ . Ferns et al . [ 2004 ] show that this metric bounds differences in the optimal value function , hence its importance in RL : |V ∗ ( x ) − V ∗ ( y ) | ≤ d∼ ( x , y ) ∀x , y ∈ X . ( 1 ) Representation learning in RL . In large-scale environments , it is infeasible to express value functions directly as vectors in RX×A . Instead , RL agents must approximate value functions in a more concise manner , by forming a representation of the environment , that is , a feature embedding φ : X → RM , and predicting state-action values linearly from these features . Representation learning is the problem of finding a useful representation φ . Increasingly , deep RL agents are equipped with additional losses to aid representation learning . A common approach is to require the agent to make additional predictions ( so-called auxilliary tasks ) with its representation , typically with the aid of extra network parameters , with the intuition that an agent is more likely to learn useful features if it is required to solve many related tasks . We refer to such methods as implicit representation shaping , since improved representations are a side-effect of learning to solve auxiliary tasks . Since bisimulation metrics capture additional information about the MDP in addition to that summarised in value functions , bisimulation metrics are a natural candidate for auxiliary tasks in deep reinforcement learning . Gelada et al . [ 2019 ] , Agarwal et al . [ 2021a ] , and Zhang et al . [ 2021 ] introduce auxiliary tasks based on bisimulation metrics , but require additional assumptions on the underlying MDP in order for the metric to be learnt correctly ( Lipschitz continuity , deterministic , and Gaussian transitions , respectively ) . The success of these approaches provides motivation in this paper to introduce a notion of state similarity applicable to arbitrary MDPs , without further restriction . Further , we learn this state similarity explicitly : that is , without the aid of any additional network parameters . 3 Advantages and limitations of the bisimulation metric . The bisimulation metric d∼ is a strong notion of distance on the state space of an MDP ; it is useful in policy transfer through its bound on optimal value functions [ Castro and Precup , 2010 ] and because it is so stringent , it gives good guarantees for state aggregations [ Ferns et al. , 2004 , Li et al. , 2006 ] . However , it has been difficult to use at scale and compute online , for a variety of reasons that we summarize below . ( i ) Computational complexity . The metric can be computed via fixed-point iteration since the operator TK is a contraction mapping . The map TK contracts at rate γ with respect to the L∞ norm onM , and therefore obtaining an ε-approximation of d∼ under this norm requires O ( log ( 1/ε ) / log ( 1/γ ) ) applications of TK to an initial pseudometric d0 . The cost of each application of TK is dominated by the computation of |X |2|A|Wd distances for distributions over X , each costing Õ ( |X |2.5 ) in theory [ Lee and Sidford , 2014 ] , and Õ ( |X |3 ) in practice [ Pele and Werman , 2009 , Guo et al. , 2020a , Peyré and Cuturi , 2019 ] . Thus , the overall practical cost is Õ ( |X |5|A| log ( ε ) / log ( γ ) ) . ( ii ) Bias under sampled transitions . Computing TK requires access to the transition probability distributions P ax for each ( x , a ) ∈ X×Awhich , as mentioned in Section 2 , are typically not available ; instead , stochastic approximations to the operator of interest are employed . Whilst there has been work in studying online , sample-based approximate computation of the bisimulation metric [ Ferns et al. , 2006 , Comanici et al. , 2012 ] , these methods are generally biased , in contrast to sample-based estimation of standard RL operators . ( iii ) Lack of connection to non-optimal policies . One of the principal behavioural characterisations of the bisimulation metric d∼ is the upper bound shown in Equation ( 1 ) . However , in general we do not have |V π ( x ) − V π ( y ) | ≤ d∼ ( x , y ) for arbitrary policies π ∈ Π ; a simple example is illustrated in Figure 2 . More generally , notions of state similarity that the bisimulation metric encodes may not be closely related to behavioural similarity under an arbitrary policy π . Thus , learning about d∼ may not in itself be useful for large-scale reinforcement learning agents . Property ( i ) expresses the intrinsic computational difficulty of computing this metric . Property ( ii ) illustrates the problems associated with attempting to move from operator-based computation to online , sampled-based computation of the metric ( for example , when the environment dynamics are unknown ) . Finally , property ( iii ) shows that even if the metric is computable exactly , the information it yields about the MDP may not be practically useful . Although π-bisimulation ( introduced by Castro [ 2020 ] and extended by Zhang et al . [ 2021 ] ) addresses property ( iii ) , their practical algorithms are limited to MDPs with deterministic transitions [ Castro , 2020 ] or MDPs with Gaussian transition kernels [ Zhang et al. , 2021 ] . Taken together , these three properties motivate the search for a metric without these shortcomings , which can be used in combination with deep reinforcement learning .

Summary: This paper introduces MICo, a new behavioural distance over the state space of an MDP. Compared to previous works on behavioural metrics such as bisimulation metrics, MICo is compatible with sample-based algorithms and therefore can be easily integrated with deep RL methods. Empirical results demonstrate that learning this distance improves performance on the Atari benchmark.