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: This paper considers the problem of learning good representations for reinforcement learning problems. It approaches this from the perspective of bisimulation metrics, extending earlier work on on-policy bisimulation for control (DBC). The paper identifies issues with DBC that also appear in other types of metric learning, and includes a theoretical analysis of why this happens. To remedy the issues the paper proposes normalizing the representations and including an extension to DBC by including an inverse dynamics model and intrinsic rewards.