Towards Robust Bisimulation Metric Learning

1 Introduction . Complex reinforcement learning ( RL ) problems require the agent to infer a useful representation of the world state from observations . The utility of this representation can be measured by how readily it can be used to learn and enact a policy that solves a specific task . As an example , consider a robot using visual perception to pour a cup of coffee . The clutter on the counter and the colour of the walls have little effect on the correct action ; even more pertinent aspects , such as two mugs with slightly different patterns and shapes , should be treated nearly the same by the policy . In other words , many states are equivalent for a given task , with their differences being task-irrelevant distractors . Thus , a natural approach to generalization across environmental changes is constructing a representation that is invariant to such nuisance variables – effectively “ grouping ” such equivalent states together . One method of obtaining such a representation uses the notion of a bisimulation metric ( BSM ) [ 13 , 14 ] . The goal is to abstract the states into a new metric space , which groups “ behaviourally ” similar states in a task-dependent manner . In particular , bisimilar states ( i.e. , in this case , those close under the BSM ) should yield similar stochastic reward sequences , given the same action sequence . In a recursive sense , this requires bisimilar states to have similar ( i ) immediate rewards and ( ii ) transition distributions in the BSM space ( e.g. , see Eq . 1 ) . Thus , ideally , the resulting representation space contains the information necessary to help the policy maximize return and little else . ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Recently , the Deep Bisimulation for Control ( DBC ) algorithm [ 53 ] showed how to map observations into a learned space that follows a policy-dependent ( or on-policy ) BSM ( PBSM ) for a given task [ 10 ] , resulting in a powerful ability to ignore distractors within high-dimensional observation spaces . Using a learned dynamics model , DBC trains an encoder to abstract states to follow the estimated PBSM , using the aforementioned recursive formulation for metric learning . However , this method can suffer from issues of robustness under certain circumstances . Conceptually , differentiating state encodings requires trajectories with different rewards . In other words , two states that only differ in rewards far in the future should still not be bisimilar . Yet , in practice , the encoder and policy are learned simultaneously as the agent explores ; hence , in the case of uninformative ( e.g. , sparse or near-constant ) rewards , it may incorrectly surmise that two states are bisimilar , as most trajectories will look very similar . This leads to a premature collapse of the embedding space , grouping states that should not be grouped . On the other hand , we can show that the formulation of the metric learning loss is susceptible to embedding explosion if the representation space is left unconstrained2 . In our work , we build upon the DBC model in an attempt to tackle both problems : ( 1 ) we address embedding explosion by stabilizing the state representation space via a norm constraint and ( 2 ) we prevent embedding collapse by altering the encoder training method . Contributions We ( i ) generalize theoretical value function approximation bounds to a broader family of BSMs , with learned dynamics ; ( ii ) analyze the BSM loss formulation and identify sources of potential embedding pathologies ( explosion and collapse ) ; ( iii ) show how to constrain representation space to obtain better embedding quality guarantees ; ( iv ) devise theoretically-motivated training modifications based on intrinsic motivation and inverse dynamics regularization ; and ( v ) show on a variety of tasks that our method is more robust to both sparsity and distractors in the environment . 2 Related Work . Our work builds on two major aspects of DRL : representation learning , particularly task-specific encodings following the BSM , and methods for sparse reward environments , ( e.g. , intrinsic motivation ) . Representation Learning for Control In the context of RL , there is a spectrum of representations from task-agnostic to task-specialized . One task-agnostic approach to learning latent representation spaces uses a reconstruction loss . In particular , some methods utilize this approach for model-based planning [ 27 , 54 , 25 ] , while others do so for model-free RL [ 33 , 17 , 30 ] . Such approaches often utilize generative models , which have shown use for learning via simulation , in addition to representation learning [ 22 , 51 ] . Instead of fully reconstructing an input , other approaches use self-supervised representation learning methods , especially contrastive learning algorithms [ 34 , 28 , 44 ] . While task-agnostic reconstruction is reusable and transferable , it can store unnecessary or distracting details . In contrast , state abstraction methods [ 31 ] can obtain compact state representations that ignore irrelevant information , by finding ( approximately ) equivalent aggregated MDPs that have ( nearly ) identical optimal policies . To learn such MDP homomorphisms [ 38 , 39 ] , previous efforts have exploited various structural regularities ( e.g. , equal rewards and transition probabilities [ 19 , 46 ] , state-action symmetries and equivariances [ 46 , 48 , 49 ] ) and assumptions [ 6 ] . State similarity metric learning [ 29 ] is closely related , as one can aggregate -close states to produce a new MDP . For example , inter-state distances can be assigned based on the difference between policy distributions conditioned on respective states [ 2 , 18 ] . Inverse dynamics prediction can also help encourage holding only information relevant to what the agent can control [ 3 , 36 ] . For model-based RL , prior work constructed equivalence classes of models [ 21 ] , as well as other abstractions [ 12 , 16 ] , to save computation and memory . In our work , we build on the use of task-specific bisimulation relations [ 19 ] , which define equivalence classes of states based on rewards and transition probabilities . However , obtaining equivalence classes is brittle , as our estimates of these quantities may be inexact . Hence , Ferns et al . [ 13 , 14 , 15 ] instead turn to a bisimulation metric between states , which can vary smoothly as the rewards and transition probabilities change . More recently , Zhang et al . [ 53 ] devised Deep Bisimulation for Control ( DBC ) , which applies a metric learning approach to enforce its representation to approximately follow bisimulation-derived state aggregation . While DeepMDP [ 17 ] proves that its representation provides an upper bound on the bisimulation distance , DBC [ 53 ] directly imposes the bisimulation metric 2Embedding collapse and explosion issues also appear in the computer vision literature [ 26 , 42 , 52 ] . structure on the latent space . Herein , we extend prior analysis and methods for representation learning using PBSMs , with the goal of improving robustness and general performance . In particular , we focus on preventing embedding pathologies , in the context of uninformative reward signals and distraction . Reinforcement Learning under Uninformative Rewards Many real-world tasks suffer from uninformative rewards , meaning that the agent receives sparse ( mostly zero ) or static ( largely constant ) rewards throughout its trajectories . For example , games like Montezuma ’ s Revenge only provide reward signals after a large number of steps – and the agent is likely not to receive one at all , in most cases . Thus , while common , this situation significantly increases the difficulty of the RL task , as the agent receives little signal about its progress . To deal with such situations , a common tactic is to encourage the agent to explore in the absence of an extrinsic reward signal . This is often done by providing intrinsic motivation via a self-derived reward , resulting in “ curiosity-driven ” behaviour [ 8 , 4 ] . Such approaches include surprise [ 1 , 5 ] ( where experiencing unexpected dynamics is rewarded ) , and empowerment [ 40 , 20 ] ( where the agent prefers states in which it has more control ) . In this work , we specifically target the case of uninformative rewards , showing that the DBC model is especially susceptible to instability or collapse for such tasks , and consider ways to ameliorate this issue . In particular , we utilize the forward prediction error as an intrinsic reward [ 36 , 8 , 43 ] , thus augmenting the sparse extrinsic rewards and encouraging exploration . We also regularize the latent space by learning an inverse dynamics model on top of it , which does not rely on the extrinsic reward signal and has previously been used for distraction-robust representation learning in DRL [ 3 , 36 ] . 3 Technical Approach . In this section , we first describe our notation and problem setting , and review prior definitions of bisimulation metrics and relevant metric learning objectives . Then , we consider theoretical connections to value functions and convergence properties , and show that these formulations may be susceptible to ( i ) instabilities in optimization due to embedding explosion ( i.e. , large norms ) , and ( ii ) convergence to trivial solutions where all states are represented as a single point ( embedding collapse ) . Based on our analysis , we propose an extension of deep bisimulation metric learning , with theoretically-motivated constraints on the optimization objective , including alterations to the forward dynamics model , intrinsic motivation , and inverse dynamics regularization . 3.1 Preliminaries . In this work , we consider a discounted Markov Decision Process ( MDP ) given by a tuple , 〈S , A , P , R , ρ0〉 . At the beginning of each episode , an initial state , s0 ∈ S , is sampled from the initialstate distribution ρ0 over the state space S . Then , at each discrete time-step t ≥ 0 , an agent takes an action , at ∈ A , according to a policy π ( at|st ) . As a result , the MDP transitions to the next state according to a transition distribution P ( st+1|st , at ) . The agent collects a scalar reward , rt = R ( st , at ) , from the environment according to a bounded reward function , R : S × A → [ Rmin , Rmax ] . In infinite- and long-horizon settings , a discount factor , γ ∈ [ 0 , 1 ) , is used to calculate the agent ’ s discounted return in a given episode , G = ∑ t≥0 γ trt . RL algorithms aim to find an optimal policy , π∗ : = argmaxπ∈Π E [ G ] , for a class of stationary policies Π . In high-dimensional , continuous state ( or observation ) spaces , this learning problem is rendered tractable via a state encoder ( e.g. , a neural network ) , φ : S → Rn , which is used to learn a policy of the form π ( a|φ ( s ) ) . The following ( pseudo3 ) -metric , based on the Wasserstein metric ( see Appendix A for a review ) , is of particular relevance to this work . Distances are assigned to state pairs , ( si , sj ) ∈ S × S , according to a pessimistic measure [ 9 ] of how much the rewards collected in each state and the respective transition distributions differ . A distance of zero for a pair implies state aggregation , or bisimilarity . Definition 1 ( Bisimulation metric for continuous MDPs , Thm . 3.12 of [ 14 ] ) . The following metric exists and is unique , given R : S ×A → [ 0 , 1 ] and c ∈ ( 0 , 1 ) for continuous MDPs : d ( si , sj ) = max a∈A ( 1− c ) |R ( si , a ) −R ( sj , a ) | + cW1 ( d ) ( P ( ·|si , a ) , P ( ·|sj , a ) ) . ( 1 ) An earlier version of this metric for finite MDPs used separate weighting constants cR , cT ≥ 0 for the first and second terms respectively , and required that cR + cT ≤ 1 [ 13 ] . Here , when the weighting 3Pseudo-metrics are a generalization of metrics that allow distinct points to have zero distance . For simplicity , we broadly use the term `` metric '' in the remainder of this paper at the cost of imprecision . constant cT of the W1 ( d ) term is in [ 0 , 1 ) , the RHS is a contraction mapping , F ( d ) : met→ met , in the space of metrics . Then , the Banach fixed-point theorem can be applied to ensure the existence of a unique metric , which also ensures convergence via fixed-point iteration for finite MDPs . For more details , we refer the reader to [ 13 , 14 ] and the proof of Remark 1 in Appendix B . Notice that c ( or cT ) determines a timescale for the bisimulation metric , weighting the importance of current versus future rewards , analogously to the discount factor γ . More recently , an on-policy bisimulation metric ( also called π-bisimulation ) was proposed to circumvent the intractibility introduced by taking the max operation over high-dimensional action spaces ( e.g. , continuous control ) , as well as the inherent pessimism of the policy-independent form [ 10 ] . Definition 2 ( On-policy bisimulation metric [ 10 ] ) . Given a fixed policy π , the following on-policy bisimulation metric exists and is unique : dπ ( si , sj ) : = |rπi − rπj | + γW1 ( dπ ) ( Pπ ( ·|si ) , Pπ ( ·|sj ) ) , ( 2 ) where rπi : = Ea∼π [ R ( si , a ) ] and Pπ ( ·|si ) : = Ea∼π [ P ( ·|si , a ) ] . Zhang et al . [ 53 ] proposed to learn a similar on-policy bisimulation metric directly in the embedding space via an MSE objective . They proposed an algorithm for jointly learning a policy π ( a|φ ( s ) ) with an on-policy bisimulation metric . Below , we define a generalized variant of their objective : J ( φ ) : = 1 2 E [ ( d̂π , φ ( si , sj ) − |rπi − rπj | − γW2 ( ‖·‖q1 ) ( P̂π ( ·|φ ( si ) ) , P̂π ( ·|φ ( sj ) ) ) ) 2 ] , ( 3 ) where d̂π , φ ( si , sj ) : = ‖φ ( si ) − φ ( sj ) ‖q2 and they used ( q1 = 2 , q2 = 1 ) . Notice that the recursion induced by this objective is different from prior metrics in three ways ; ( i ) a W2 metric was used instead of aW1 metric sinceW2 has a convenient closed form for Gaussian distributions when q1 = 2 , ( ii ) the distance function used for the bisimulation metric and the Wasserstein metric are different ( L1 and L2 respectively ) , and ( iii ) a forward dynamics model is used instead of the ground truth dynamics . While they may introduce practical benefits , these differences violate the conditions under which the existence of a unique bisimulation metric has been proven [ 10 , 13 , 14 , 15 ] . Thus , we will ( i ) assume a unique metric exists for all Wasserstein metrics Wp when necessary ( see Assumption 1 ) , ( ii ) study losses that use a matching metric , i.e. , q1 = q2 = q , and ( iii ) introduce a constraint on forward models in Sec . 3.2.2 . Next , we note that they recommend the use of stop gradients for the W2 term . The resulting gradient updates may be considered as approximate fixed-point iteration in the space of metrics : F̂ ( d̂π ( φn , P̂ ) ) ( si , sj ) : = d̂π ( si , sj ; φn + αn∇φJ ( φn ) , P̂ ) , ( 4 ) where αn is a learning rate . However , in practice , π and P̂ may also be updated in training , which is of particular relevance when they have form π ( a|φ ( s ) ) and P̂ ( φ ( s′ ) |φ ( s ) , a ) as in [ 53 ] . In the next section , we will discuss conditions under which joint updates to a policy π and a metric defined by state encodings φ may or may not converge in practical settings : lim n→∞ F ( n ) ( π , d̂π , φ ) ? = dπ∗ , ( 5 ) where dπ∗ is the on-policy bisimulation metric for the optimal policy π∗ .

Summary: In this paper, the author proposed a robust bisimulation metric that could combat the collapse and explosion issues. The authors analyzed the cause of these issues theoretically and further proposed  a set of remedies including a norm constraint on the representation space, and intrinsic rewards and latent space regularization. The authors did experiments in the gym environment to verify the proposed methods.