Learning State Representations via Retracing in Reinforcement Learning

1 INTRODUCTION . Recent developments in deep reinforcement learning ( RL ) have made great progress in solving complex control tasks ( Mnih et al. , 2013 ; Levine et al. , 2016 ; Silver et al. , 2017 ; Vinyals et al. , 2017 ; Schrittwieser et al. , 2020 ) . With the increasing capacity of deep RL algorithms , the problems of interests become increasingly complex . An immediate challenge is that the observation space becomes unprecedentedly high-dimensional , and often the perceived observations have significant redundancy , negatively impacting the policy learning . The field of representation learning offers a wide range of approaches for extracting useful information from high-dimensional data ( with potentially sequential dependency structure ) ( Bengio et al. , 2013 ) . Many recent works have explored the application of representation learning in RL ( Ha and Schmidhuber , 2018 ; Hafner et al. , 2019 ; 2020a ; Schrittwieser et al. , 2020 ; Schwarzer et al. , 2021 ; Zhang et al. , 2020 ) , which lead to superior performance comparing to naive embedding . Many of such algorithms rely on predictive ( reconstruction ) supervision for representation learning , i.e. , the state representations are learned such that transitions in the observation space are maximally preserved in the embedding space . Here we argue that existing methods do not fully exploit the supervisory signals inherent in the data . Additionally valid supervision can often be obtained for representation learning by including temporally “ backward ” transitions in situations in which the same set of rules govern both temporally forward and backward transitions . In these cases we are able to use such temporally backward transitions as “ free ” supervision for training the representation/model without having to perform any additional interaction with the environment . Hence , by “ learning via retracing ” , we obtain more training samples for representation learning ( twice as much as existing approaches in tasks that admit perfect reversibility across all transitions ) , hence improve the sample efficiency of representation learning . Moreover , the quality of the representation is strongly correlated with the quality of policy learning ( Gelada et al. , 2019 ) . Hence with faster representation learning , we can subsequently achieve more efficient policy learning . Thus we hypothesise that by augmenting representation learning with “ learning via retracing ” in RL , we can significantly improve the sample efficiency of learning , which is a long-standing issue that plagues the practical applicability of deep RL algorithms . Beyond improved sample efficiency , joint predictive supervision in temporally forward and backward directions conveys information from both the future and past to the target state , similar to the smoothing operation for latent state inference in state-space models ( Kalman , 1960 ; Murphy , 2012 ) , leading to more accurate latent state inference , hence achieving stronger representation learning . “ Learning via retracing ” can be integrated into any representation learning that utilises a transition model , under both the model-free and model-based RL frameworks . One problem that plagues the successful application of “ learning via retracing ” is that the reversibility might not be preserved across all valid transitions , i.e. , there exists transitions such that s→ s′ for some action a , but s′ 6→ s for all possible actions a ( Figure 1b ) . Under these situations , naively forcing the representations for s and š ( the retraced state given s′ and a , see Section 3 ) leads to suboptimal representation space , hence potentially hinders or even completely interrupts the overall RL training . Hence in order to maximally preserve the advantages brought by “ learning via retracing ” , it is essential to identify the “ irreversible ” transitions and rule them out from representation learning via retracing . To this end , we propose a novel dynamically regulated approach for identifying such “ irreversible ” states , which we term as adaptive truncation ( Section 3.3 ) . As a motivating example , consider a rat navigating towards a cheese in a cluttered environment ( Figure 1a ) . Upon first visit to the goal state , multiple imaginative retracing trajectories can be randomly simulated . By constraining the temporal cycle-consistency of the retracing transitions , the rat quickly builds a state representation that accurately preserves the transitions in the area around the actual forward trajectory taken by the rat . Moreover , all retracing simulations pass through the two “ doors ” , allowing the rat to quickly identify the key bottleneck states that are essential for successful navigation towards the goal , hence facilitating stronger transferability across environments ( Section 5.3 ) . We conjecture that such imaginative retracing could be physiologically implemented as the reversed hippocampal “ replay ” that has been observed in both rodents and humans ( Foster and Wilson , 2006 ; Mattar and Daw , 2018 ) ( see Section 6 for further discussion ) . We propose Cycle-Consistency World Model ( CCWM ) , a self-supervised instantiation of “ learning via retracing ” , for joint representation learning and generative model learning under the model-based RL settings . We empirically evaluate the performance of CCWM on challenging visual-based continuous control benchmarks . Experimental results show that CCWM achieves state-of-the-art performance in terms of sample efficiency and asymptotic performance , whilst providing additional advantages such as stronger generalisability and extended planning horizon . 2 PRELIMINARIES . 2.1 PROBLEM FORMULATION . We consider reinforcement learning problems in Markov Decision Processes ( MDPs ) . An MDP can be charaterised by the tuple , M = 〈S , A , R , P , γ〉 , where S , A are the state and action spaces , respectively ; R : S → R is the reward function ( we assume determinisitc reward functions unless stated otherwise ) , P : S ×A×S → [ 0 , 1 ] is the transition distribution of the task dynamics ; γ ∈ R is the discounting factor . The control policy , π : S ×A → [ 0 , 1 ] , represents a distribution over actions at each state . The goal is to learn the optimal policy , π∗ , such that the expected future reward is maximised across all states , i.e. , π∗ = arg max π∈Π Eπ [ ∑ t γtR ( st , at ) |s0 = s ] , ∀s ∈ S , where st ∼ p ( ·|st−1 , at−1 ) for t = 1 , 2 , . . . ( 1 ) Here we consider tasks in which the perceivable observation space , O , is high-dimensional , due to either redundant information or simply because only visual inputs are available . Hence it is necessary to learn an embedding function , φ : O → Z , such that the embedded observation space , φO , could act as S in the MDP , to support efficient learning of the optimal policy using existing RL algorithms . 2.2 GENERATIVE MODELLING OF DYNAMICS . Modelling the transition dynamics using a sequential VAE-like structure enables joint learning of the latent representation and the associated latent transitions . Specifically , the dynamics model is defined in terms of the following components ( see also top part in Figure 2b ) . Observation ( context ) embedding : et = qφ ( Ot ) , Latent posterior distribution : p ( zt+1|zt , at , Ot+1 ) , Latent variational transition ( prior ) distribution : qψ1 ( zt+1|zt , at ) , Latent variational posterior distribution : qψ2 ( zt+1|zt , at , et+1 ) , Generative distribution : pθ ( Ot+1|zt+1 ) , ( 2 ) where ψ = { ψ1 , ψ2 , φ } and θ represent the parameters associated with the recognition and generative models respectively . Latent variables z ∈ Z are introduced for more flexible modelling of the distribution of the observed variables . A variational approximation is employed since the true posterior p ( zt+1|zt , at , Ot+1 ) is usually intractable in practice . At each time step t , the agent receives an observationOt , which is then embedded into a context vector et = qφ ( Ot ) . After initialisation , the latent space vector is rolled out in a forward fashion given the action at , yielding one-step prediction into the future following the variational transition distribution qψ1 ( zt+1|zt , at ) ( we assume standard first-order Markovian structure on the latent variables ) . The dynamics model is trained following the maximum likelihood paradigm , and due to intractability , the parameters of the recognition and generative models in Eq . 2 are learned by maximising the variational free energy ( also known as the ELBO ; Wainwright and Jordan ( 2008 ) . LELBO ( Ot+1 ) = Ezt+1∼qψ2 [ log pθ ( Ot+1|zt+1 ) − βDKL [ qψ2 ( zt+1|zt , at , et+1 ) ||qψ1 ( zt+1|zt , at ) ] ] , ( 3 ) The variational free energy loss function consists of the reconstruction error , pθ ( Ot+1|zt+1 ) , and the KL-divergence between the variational posterior distributions and the predictive prior as regularisation . The intuitive autoencoder-like structure nicely separates the generative process from the inference process , with the variational posterior ( qψ2 ( zt+1|zt , at , et+1 ) ) serving as the main inference engine of the optimal latent representations . Note that the β parameter controls the degree of factored structure ( disentanglement ) of the latent code , which by default is set to 1 ( Higgins et al. , 2016 ) . 3 METHOD . We firstly introduce learning via retracing in its most general format , then provide a concrete instantiation based on generative dynamics modelling under the model-based framework . We finally propose a novel adaptive truncation scheme for dealing with the “ irreversibility ” issue in “ learning via retracing ” . 3.1 LEARNING VIA RETRACING . We always assume the usage of an approximate dynamics model , regardless of the overall RL agent being model-free or model-based ( the dynamics model would only be used for representation learning under the model-free setting , hence it is possible to have separate dynamics model for forward and “ reversed ” transitions ) . Given a set observations , O = { O1 , . . . , OT } , a dynamics model , T : S ×A× S → [ 0 , 1 ] , the classical methods of self-supervised learning of representations involve learning an encoder ( Eφ ) and a decoder ( Dθ ) , trained via minimising the predictive reconstruction . L ( φ , θ ) = d ( O1 : T , D ( f ( E ( O1 : T ) ; T ) ) ( 4 ) where d is some metric of the observable space ( e.g. , the L2 distance ) . The function f ( e ; T ) is some function specifying the schedule for predictions , e.g. , f ( et , T ) = T ( et ) = êt+1 corresponds to learning the representations based on one-step predictive reconstruction . Existing methods explore predictive supervision in a temporally forward direction , however , we argue that the “ reversed ” transitions can also contain useful signals for learning . Consider a transition tuple , ( s , a , s′ ) , we define the “ reversed ” transitions being the tuple ( s′ , a′ , s ) , where a′ is the “ reversed ” action . In situations where the same set of rules apply to forward and backward transitions , the reversed transition given a′ could contribute to representation learning via Eq . 4 as an additional training sample ( utilising a potentially different loss function from the forward supervision , see Figure 2a ) . Hence by utilising the additional reversed transition for representation learning , we improve the sample efficiency of learning without any additional interactions with the environment . As mentioned earlier , in situations where perfect reversibility is not preserved across all transitions , such “ irreversible ” transitions could negatively impact the overall learning . Correct identification of such states is hence essential for the successful implementation of “ learning via retracing ” . To this end , we propose a novel adaptive truncation scheme in Section 3.3 . Despite the intuitive simplicity of learning via retracing , it offers a number of advantages comparing to existing representation learning methods . In addition to the improved sample efficiency , we can interpret learning with “ reversed ” transitions as explicit inference of the current latent state given future information . In combination with the forward predictive supervision , the joint learning dynamics is similar to the smoothing operation in dynamical systems , which is often superior than filtering ( corresponds to using solely the forward predictive supervision ) in terms of inference accuracy ( Kalman , 1960 ; Murphy , 2012 ) . Hence learning via retracing could support stronger representation learning . We have introduced learning via retracing in its most general form , where a large degree of freedom exists such that the method can be tailored and integrated with many of the existing representation learning approaches . There are many free model choices , such as being model-free or model-based ; whether or not using a separate “ reversed ” dynamics model ; loss function for constraining the “ retraced ” transitions ; deterministic or probabilistic dynamics model , just to name a few . Below we provide one possible instantiation of learning via retracing , the Cycle-Consistency World Model ( CCWM ) , under the model-based framework based on generative dynamics modelling .

Summary: This paper proposes "learning via retracing" as an approach to learn state representations, through matching a trajectory both in the forward and in the backward direction. This paper then introduces Cycle-Consistency World Model (CCWM) which is a model-based RL algorithm which learns through retracing. This method is sample efficient and provides good state representations for predicting the future and generalization. Furthermore, it proposes a value-based approach to identify "irreversible" transitions.