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: The paper proposes a new behavioral distance over the state space of an MDP, which can be used to learn representations in deep RL tasks. The proposed distance has three advantages over the existing bisimulation metric: (1) it has better computational complexity; (2) it allows for online approximation; and (3) it has a close connection to the underlying policy. Empirical evidence on toy examples shows the informativeness of the representations learned by using the proposed distance, and large-scale experiments show improvements of deep RL agents when incorporating the proposed distance to learn representations.