Year: 2021

Abstract:
We propose learning via retracing, a novel self-supervised approach for learning the state representation (and the associated dynamics model) for reinforcement learning tasks. In addition to the predictive (reconstruction) supervision in the forward direction, we propose to include “retraced” transitions for representation/model learning, by enforcing the cycle-consistency constraint between the original and retraced states, hence improve upon the sample efficiency of learning. Moreover, learning via retracing explicitly propagates information about future transitions backward for inferring previous states, thus facilitates stronger representation learning. We introduce Cycle-Consistency World Model (CCWM), a concrete instantiation of learning via retracing implemented under existing model-based reinforcement learning framework. Additionally we propose a novel adaptive “truncation” mechanism for counteracting the negative impacts brought by the “irreversible” transitions such that learning via retracing can be maximally effective. Through extensive empirical studies on continuous control benchmarks, we demonstrates that CCWM achieves state-of-the-art performance in terms of sample efficiency and asymptotic performance.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARIES

heading: 2.1 PROBLEM FORMULATION
text: 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.

heading: 2.2 GENERATIVE MODELLING OF DYNAMICS
text: 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).

heading: 3 METHOD
text: 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”.

heading: 3.1 LEARNING VIA RETRACING
text: 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.

heading: 3.2 CYCLE-CONSISTENCY WORLD MODEL
text: CCWM is a model-based RL agent that utilises a generative world-model, trained given both the predictive reconstruction of future states, and constraining the temporal cycle-consistency of the “retraced” states (i.e., constraining the retraced states to match the original states). For notational convenience, we denote all predictive prior estimates, posterior estimates, and the retraced predictive latent estimates as ẑ, z̃ and ž, respectively.
We use the similar dynamics model described in Section 2.2 (top panel in Figure 2b). We additionally define a reverse action approximator, ρζ : Z × Z → A, which takes in a tuple of latent states (zt+1, zt) and outputs an action ǎt+1 that approximates the “reversed” action that leads the transition from zt+1 back to zt. Instead of introducing a separate “reversed” dynamics model, we use the same dynamics model for both the forward and retracing transitions, which lead to improved sample efficiency of model learning in addition to representation learning. The parameters of ρ, ζ, can be either learned jointly with the model in an end-to-end fashion, or trained separately (see Appendix A). The graphical model of CCWM is shown in Figure 2b.
During training, given a sample trajectory {O1:T+1, a1:T }, we firstly perform a one-step forward sweep through all timesteps to compute the variational prior and posterior estimates of latent states.
ẑτ+1 ∼ qψ1(z|z̃τ , aτ ) z̃τ+1 ∼ qψ2(z|ẑτ , aτ , qφ(Oτ ))
(5)
for τ = 0, . . . , T , where ẑ0 is randomly initialised. Note that we also include reward prediction as part of the dynamics modelling, and we have omitted showing this for simplicity.
Given the predictive estimates in the forward direction, we compute the “retracing” estimates, utilising the same latent transition dynamics (variational predictive prior distribution).
žτ ∼ qψ1(z|z̃τ+1, ǎτ+1), where ǎτ+1 = ρζ(z̃τ+1, z̃τ ), for τ = 1, . . . , T (6) The forward and retracing predictive supervision separately contributes to the model learning of CCWM. For the forward pass, the parameters of the dynamics model are trained to maximise the likelihood of the sampled observations via predictive reconstruction. We follow the variational principle, by maximising the variational free energy (Eq. 3), computed by Monte Carlo estimate given the posterior predictive samples. For the “retracing” operation, model learning is based on constraining the deviation of the retraced states from the original states. Intuitively, this utilises the temporal cycle-consistency of the transition dynamics: assuming that the action-dependent transition mapping is invertible across all timesteps (Dwibedi et al., 2019). The loss function for constraining the cycle-consistency is another degree of freedom of “learning via retracing”. For CCWM, we choose bisimulation metric as the loss function for the retracing operations, which has been shown to yield stronger constraints of the latent states on the MDP level, and also leads to more robust and noise-invariant representation without reconstruction (Ferns et al., 2011; Zhang et al., 2020).
Lretrace(z̃t, žt) = Ežt [ (||z̃t − žt||1 −DKL[R̂(·|z̃t)||R̂(·|žt)]− γW2(qψ1(·|z̃t, π(z̃t)), qψ1(·|žt, π(žt))))2 ] ,
(7)
where R̂(r|z) represents the learned reward distributions (we assume stochastic rewards), andW2(·, ·) represents the 2-Wasserstein distance (see Eq. 9 in Appendix A). The advantage of choosing the bisimulation metric as the retrace loss function is further empirically shown in Appendix G through careful ablation studies (Figure 12). Multiple retracing trajectories can be simulated and the retrace loss is again a Monte Carlo estimate based on sampled “retracing” states, but empirically we observe that one “retracing” sample is sufficient as we do not observe noticeable improvements for increasing the number of sampled “retracing”. We note that here we utilise the same transition dynamics model for both forward and reversed rollouts, which might cause issues in model learning due to the absence of perfect “reversibility” across all valid transitions. Hence we need a method for dynamically assessing the “reversibility” so as to know when to apply learning via retracing (see Section 3.3).
The overall objective for the CCWM dynamics model training is thus a linear combination of the forward and “retracing” loss functions.
L(θ, ψ, ζ) = 1 NT N∑ n=1 T∑ τ=1 [LELBO(Onτ ; θ, ψ) + λLretrace(z̃nτ , žnτ ;ψ, ζ)] , (8)
where λ is the scalar multiplier for the retrace loss, and N is the batch size. We implement CCWM using a Recurrent State-Space Model (RSSM; Hafner et al. (2019)). The complete pseudocode for CCWM training is shown in Algorithm 1 in Appendix B. We note that “learning via retracing” is also applicable under the model-free setting, we describe one such instantiation in Appendix C.

heading: 3.3 REVERSIBILITY AND TRUNCATION
text: As we noted above, perfect reversibility is not always present across all transitions in many environments. For instance, consider the falling android presented in Figure 1b, it is trivial to observe that no valid action is able to transit a falling android to its previous state. Under such situations, naive application of “learning via retracing”, by constraining the temporal cycle-consistency, will corrupt representation learning (and dynamics model learning in CCWM). Here we propose an approach to deal with such “irreversibility”.
Our approach is based on adaptive identification of “irreversible” transitions. We propose that the value function of the controller (e.g., an actor-critic agent) possesses some information about the
continuity of the latent states (Gelada et al. (2019); see Appendix D for further discussion). Hence we use the value function as an indicator for sudden change in the agent’s state. Specifically, for each sampled trajectory, we firstly compute the values of each state-action pair using the current value function approximator, [Q(z1, a1), . . . , Q(zT , aT )]. We then compute the averages of the values over a sliding window of size S through the value vectors of each sampled trajectory, resulting in a (T − S)-length vector [Q̄1, . . . , Q̄T−S ]. Any drop/increase in the sliding averages (above some pre-defined threshold) indicates a sudden change in the value function, hence a sudden change in the latent representation given the continuity conveyed by the value function. Given some timestep, τ , at which the sudden change occurs, We then remove the transitions {zτ−S:τ , aτ−S:τ} from “learning via retracing”. Such adaptive scheduling allows us to deal with “irreversibility” in non-episodic tasks.

heading: 4 RELATED WORKS
text: Representation learning in RL. Finding useful state representations that could aid RL tasks has long been studied. Early works have investigated representations based on a fixed basis such as tile coding and Fourier basis (Mahadevan, 2005; Sutton and Barto, 2018). With the development of deep learning techniques, recent works explored automatic feature discovery based on neural network training, which can be categorised into three large classes. The first class of methods studies the usage of data augmentation for broadening the data distribution for training more robust feature representation (Laskin et al., 2020; Kostrikov et al., 2020; Schwarzer et al., 2021; Yarats et al., 2021). The second class explores the role of auxiliary tasks in learning representations, such as weakly-supervised classification and location recognition, for dealing with the sparse and delayed supervision (Lee et al., 2020b; Mirowski et al., 2017; Oord et al., 2018). The third class of methods, specifically tailored to model-based RL models, leverages generative modelling of the environment dynamics, enabling joint learning of the representations and the dynamics model (Ha and Schmidhuber, 2018; Buesing et al., 2018; Hafner et al., 2019; 2020a; Lee et al., 2020a; Schrittwieser et al., 2020; Hafner et al., 2020b).
Cycle-Consistency. Cycle-consistency is a commonly adopted approach in computer vision and natural language processing (Zhou et al., 2016; Zhu et al., 2017; Yang et al., 2017; Dwibedi et al., 2019), where the core idea is the validation of matches between cycling through multiple samples. We adopt similar design principles for sequential decision-making tasks: rollouts in a temporally forward direction alone yield under-constrained learning of the world model. By additionally incorporating backwards rollouts into model learning in a self-supervised fashion, we enforce the inductive bias that the same transition rules govern the dynamics of the task. Due to the existence of control signals (actions), we need to minimise the distributional shift between the temporally forward and backward rollouts, otherwise naively imposing the cycle-consistency constraints would cause inconsistent model learning. The “inverse” action is approximated with predictive estimates from the forward pass such that the retraced path can be maximally similar to the forward path. We impose the temporal cycle-consistency constraints Dwibedi et al. (2019) by introducing an auxiliary objective that ensures the forward and backward temporal rollouts align with the same MDP structure.
Concurrent to our work, Yu et al. (2021) proposed PlayVirtual, a model-free RL method that integrates a similar cycle-consistency philosophy into training representations with data augmentations (Schwarzer et al., 2021). We note that PlayVirtual falls under the proposed “learning via retracing” framework, but lying on the opposite spectrum comparing to the CCWM agent, being model-free and utilising a separate reversed dynamics model.

heading: 5 EXPERIMENTAL STUDIES
text: The experimental studies aim at examining if “learning via retracing” truly helps with the overall RL training and planning (which implicitly entails better representation learned), improves the generalisability of the learned representation, and whether the truncation schedules proposed in Section 3.3 deals with the irreversibility of some state transitions.

heading: 5.1 EXPERIMENT SETUP
text: CCWM can be combined with any existing value-based or policy-based RL algorithms. We implement CCWM with a simple actor-critic controller with generalised advantage estimation based on standard
model-based RL framework using model-based rollouts1 (Sutton and Barto, 2018; Konda and Tsitsiklis, 1999; Schulman et al., 2015). The details of training and the architecture can be found in Appendix A.
Our experimental studies are based on challenging visual-based continuous control benchmarks for which we choose 8 tasks from the DeepMind Control Suite (Tassa et al. (2018); Figure. 3a).
Baselines: We implement Dreamer2 as our main model-based baseline (Hafner et al., 2020a), which represents the current state-of-the-art world-model-type model-based RL agent on visual-based continuous control tasks.We also compare with the following model-free baselines: SAC Haarnoja et al. (2018), D4PG Barth-Maron et al. (2018), A3C Mnih et al. (2016). We implement the SAC agent given the state inputs and directly report the asymptotic performance of the D4PG and A3C algorithms from Tassa et al. (2018). We report the asymptotic scores for the model-free algorithms due to the large gap in sample efficiency comparing to the model-based methods.

heading: 5.2 EVALUATION ON CONTINUOUS CONTROL TASKS
text: The performance of CCWM and selected baseline algorithms is shown in Figure 3b. The empirical results show that CCWM generally achieves faster behaviour learning comparing to the baselines, which conforms with our hypothesis that utilising backward passes in addition to forward passes provides additional supervision, hence improving the sample efficiency of learning. CCWM outperforms Dreamer on 5 of the selected tasks, and is comparable to Dreamer on 2 of the remaining 3 tasks, in terms of both the sample efficiency and final convergence performance. In Appendix E, we show that “learning via retracing” indeed provides additional “valid” supervision for model learning, hence improving sample efficiency.
1Note that for the maximally fair comparison with our main baseline, Dreamer (Hafner et al., 2020a), we utilise the same policy agent as Dreamer in order to fully demonstrate the utility of “learning via retracing” .
2We use the open source implementation: https://github.com/danijar/dreamer.
Moreover, in the “Cheetah Run” task, CCWM converges at ∼ 900 score with ∼ 5 × 105 steps, whereas Dreamer, by the time it has received 1 × 106 training steps, is yet to reach a comparable score. This demonstrates that “learning via retracing” brings more benefits beyond plain sample efficiency, i.e., doubling the training steps does not eliminate the performance gap. This corresponds to our hypothesis that by explicitly conveying future information back to previous states, “learning via retracing” enables the learning of task-aware representations that support stronger behaviour learning. To test our hypothesis, we empirically evaluate CCWM’s ability of long-term predictive reconstructions and compare with Dreamer. To ensure fair comparison, we provide further training for Dreamer whenever necessary, such that the asymptotic performance is comparable with CCWM (see Appendix A for details). Figure 4 shows that CCWM consistently yields more accurate predictive reconstructions over a longer time span, on both the walker walk and cheetah run tasks. The empirical evidence confirms our hypothesis that by incorporating “learning via retracing” into model learning enables the resulting latent space to support more accurate latent predictions, hence leading to stronger behaviour learning. Increased range of accurate latent prediction additionally enables CCWM to perform better planning. We provide further analysis of predictive reconstruction in Appendix F.

heading: 5.3 ZERO-SHOT TRANSFER
text: Based on the motivation that “learning via retracing” will improve the generalisability of the agent (Figure 1a), we empirically test the generalisability of CCWM on the basis of zero-shot transfer tasks. Specifically, we modify a number of basic configurations of the cheetah run task, such as the mass of the agent and the friction coefficients between the joints. The details of the changes can be found in Appendix A. Despite the increased sample efficiency of CCWM over Dreamer during training (Figure 3b), both methods converge at similar values at 2×106 steps. We directly evaluate the trained agents on the updated cheetah run task without further training to test their abilities on zero-shot transfer. We report the mean evaluation scores (± 1 s.d.) of both agents over 15 random seeds, as well as the one-sided t-test statistics and significance of the difference between the two sets of evaluations in Table 1. We observe that the overall performance on zero-shot transfer of our approach is comparable with Dreamer on simpler transfer tasks, and significantly outperforms Dreamer on more non-trivial modifications to the original task. The introduction of retracing also improves the stability of zero-shot transfer in general (reduced variance in evaluation). These confirm our previous hypothesis that “learning via retracing” improves the ability of within-domain generalisation.

heading: 5.4 ADAPTIVE TRUNCATION OF “LEARNING VIA RETRACING”
text: We wish to examine the effects of the proposed adaptive scheduling of truncation (Section 3.3). From Figure 3b, we observe that the original CCWM is outperformed by Dreamer on the Hopper Stand
task, probably due to the large degree of “irreversibility” of the task comparing to the others such that the naive representation learning by enforcing “learning via retracing” leads to suboptimal training.
From Figure 5, we observe that by augmenting CCWM with the proposed flexible truncation schedule, the performance significantly improves, yielding consistently better sample efficiency than both Dreamer and the original CCWM. For tasks with less degree of “irreversibility”, such as walker-walk, we do not observe significant improvement by the introduction of adaptive truncation since the amount of negative impacts were already minimal in the original settings. These statements are further substantiated by the full evaluations of an intermediate truncation scheme based on fixed scheduling in the Appendix D (Figure 7).

heading: 6 DISCUSSION
text: We proposed “learning via retracing”, a novel representation learning method for RL problems that utilises the temporal cycle-consistency of the transition dynamics in addition to the predictive reconstruction in the temporally forward direction. We introduce CCWM, a concrete model-based instantiation of “learning via retracing” based on generative dynamics modelling. We empirically show that CCWM yields improved performance over state-of-the-art model-based and model-free methods on a number of challenging continuous control benchmarks, in terms of both the sampleefficiency and the asymptotic performance. We also show that “learning via retracing” supports stronger generalisability and more accurate long-range predictions, hence stronger planning, both adhere nicely to our intuition and predictions.
We note that “learning via retracing” is strongly affected by the degree of “irreversibility” of the task. We propose an adaptive truncation scheme for alleviating the negative impacts caused by the “irreversible” transitions, and empirically show the utility of the proposed truncation mechanism in tasks with large degree of “irreversibility”.
Hippocampal replay has long been thought to play a critical role in model-based decision-making in humans and rodents Mattar and Daw (2018); Evans and Burgess (2019). Recently, Liu et al. (2021) showed how reversed hippocampal replays are prioritised due to its utility in non-local (model-based) inference and learning in humans. CCWM can be interpreted as an immediate model-based instantiation of the reversed hippocampal replay. Similar to the intuitions from the neuroscience literature, we also find that “learning via retracing” brings a number of merits comparing to its counterparts that only uses forward rollouts. In addition to improved sample efficiency and asymptotic performance, we show that CCWM also supports stronger generalisability and extends the planning horizon, indicating that stronger model-based inferences are obtained. Based on CCWM, we propose a testable experimental prediction, that suppressing reversed hippocampal replay would negatively impact the subject’s ability of within-domain generalisation.

heading: C MODEL-FREE INSTANTIATION OF “LEARNING VIA RETRACING”
text: As mentioned in Section 3, “learning via retracing” admits many degrees of freedom in its implementation. CCWM provides one such instantiation under the model-based RL setting, here we provide an alternative model based on “learning via retracing” under the model-free RL setting. The graphical illustration of the model-free instantiation is shown in Figure 6.
Visual inspection indicates the high similarity between the graphical models of the model-free version and the model-based version (CCWM), but there are key differences. However, there exists subtle but essential differences. Similar to PlayVirtual (Yu et al., 2021), due to the model-free nature of the model, we no longer requires further supervisory signals obtained from ”learning via retracing” contribute to training of the dynamics model, hence we are free to employ an independent ”reversed” dynamics model (denoted by the red arrows in the reversed direction in Figure 6) for performing the retracing operations. Moreover, given the independent ”reversed” dynamics model, we no longer requires approximation of the ”reversed” actions, hence removing the necessity of using ρ as in Figure 2b, and we only need to use the ground-truth forward actions for the retracing operations. Here we stick with the general architectural choice of using a sequential state-space model for the forward dynamics model as in CCWM, but the ”reversed” dynamics model can be chosen to be deterministic and trained discriminatively jointly with the entire model. Note that Figure 6, like CCWM, only describes one of many possible instantiations of ”learning via retracing”, we leave further investigation to future work.
D FURTHER DISCUSSION ON TRUNCATION AND THE DEGREE OF “IRREVERSIBILITY”
In Section 3.3, we introduced the adaptive truncation as a general-purpose method for improving the performance of CCWM in continuous control tasks. On a high level, we choose to remove the transitions from model/representation learning through ”learning via retracing” whenever large changes in the Q-values occur, hence the ”irreversibility” of the transitions depend on the local continuity of the corresponding Q-values. Our argument is largely based on Theorem 1 from Gelada et al. (2019), which we re-iterate below for consistency.
Theorem 1 For an MDP M = 〈S,A,R,P, γ〉 and the corresponding DeepMDP M̄ = 〈S̄,A, R̄, P̄, γ〉 (see Section 2.2 in Gelada et al. (2019) for detailed definition of DeepMDP), let dS̄ be a metric over S̄ , φ : S → S̄ be the embedding function, L∞R̄ = sups∈S,a∈A |R(s, a)−R̄(φ(s), a)| and L∞P̄ = sups∈S,a∈AW (φP(·|s, a), P̄(·|φ(s), a)) be the global loss function (for training the DeepMDP, where W (·, ·) is the Wasserstein distance). For any KV̄ -Lipschitz policy π̄ ∈ Π̄, the
Algorithm 1 Cycle-Consistency World Model (CCWM) 1: Given: policy π(a|s), encoders qφ(O), qψ1(z′|z, a), qψ2(z′|z, a, qφ(O′)), decoder qθ(O|z),
reverse action approximator ρζ(z, z′), latent horizon K, average loss Lavg = 0, counter c = 0, adaptive truncation threshold, η, adaptive truncation sliding window size, S, current estimate of the Q-function, Qγ , CCWM parameters: Θ = {φ, ψ1, ψ2, θ, ζ}.
2: Input: Batch of N sampled trajectories of length T , {Ontn:tn+T ,antn:tn+T }Nn=1. 3: for n = 1 to N do 4: O = {Ontn+k}Tk=1, a = {antn+k} T−1 k=0 . 5: for k = 1 to T −K do 6: Compute prior ẑk+1:k+K and posterior estimates z̃k+1:k+K using Eq. 5 7: Compute the sliding window average, [Q̄1, . . . , Q̄K−S ], where Q̄i =
1 S ∑S s=1Qγ(z̃i+s, ai+s)
8: Compute the step-wise difference, ∆ = [| Q̄2−Q̄1 Q̄1 |, . . . , | Q̄K−S−Q̄K−S−1 Q̄K−S−1 |] 9: Compute the adaptive truncation mask, M ∈ {0, 1}K , where Mk = 1 if ∨kj=k−S(∆j < η), and 0 otherwise
10: Compute retraced states žk:k+K using Eq. 6 11: Compute the latent model loss Lk using Eq. 10. 12: Update Lavg ← cc+1Lavg + 1c+1Lk. 13: Increment counter, c← c+ 1. 14: end for 15: end for 16: Compute the gradient,∇ΘLavg. 17: Update the model parameters, Θ← Θ + α∇ΘLavg.
representation φ guarantees that for all s1, s2 ∈ S, and a ∈ A,
|Qπ̄(s1, a)−Qπ̄(s2, a)| ≤ KV̄ dS̄(φ(s1), φ(s2)) + 2 L∞R̄ + γKV̄ L ∞ P̄
1− γ (11)
The proof of Theorem 1 can be found in Appendix A.2 in Gelada et al. (2019). The high-level interpretation for Theorem 1 is that the learned representation space should be admit the situation where two states with different values collapsing onto the same representations. Theorem 1 tells us that given a good representation (embedding function), the absolute difference between the Q-
values of two latent states should lower-bound the distance (some metric over the representation space) between the representations of the two states (up to some additive and scalar multiplicative constants). Hence whenever there is a large change in the Q-values of between the adjacent states in a trajectory, there must also be a large jump in the representation space, which we interpret as the ”irreversible” transitions. The latter interpretation can be substantiated with the visualisation of the low-dimensional embeddings of the learned representations at different stages in Figure 8 and Figure 9. Figure 8 shows the TSNE embeddings of a trajectory taken by a successfully trained agent in the Cheetah Run task (we focus on the left figure in both Figure 8(a) and (b), please refer to the full discussion of Figure 8 in Appendix E). We see that the resulting TSNE-embeddings of the representation along a trajectory show that the magnitude of the transitions in the representation space is consistently low, hence respecting the temporal locality of the states along the trajectory. However, in Figure 9, where we show the TSNE-embeddings of the representation along a trajectory in an under-trained CCWM agent in the Hopper Stand task (∼ 2× 105 training steps), we observe clear clusters of the state representations along the same trajectory. Moreover, the clustered structure also shows locality with respect to the temporal ordering (along the trajectory), showing that even in tasks such as Hopper Stand, which we view as the more ”irreversible” task, a large proportion of all transitions allows retracing operations, which can contribute as additional samples for the training of representation and the dynamics model through ”learning via retracing”. The quality of the overall representation learning is dependent on the correct identification of the ”irreversible” transitions (large jumps between adjacent states in Figure 9), which necessitates the usage of the adaptive truncation approach to identify and remove the ”irreversible” transitions from corrupting the learning process. These evidence further substantiate our proposal of the overall framework of ”learning via retracing” and the adaptive truncation technique.
We note that as the agent is undergoing a sequence of consecutive ”irreversible” transitions, such as when the walker agent is falling down shown in Figure 1(b) of the main paper, the transitions the agent goes through are ”irreversible” albeit having similar Q-values (all leading to the inevitable fallen down state). In order to deal with such ”continuity” in the ”irreversibility” of the transitions along the same trajectory, the adaptive truncation is introduced by detecting the changes in the average Q-values over a sliding window instead of the single Q-values, which allows a delayed detection of the sudden (and extended) changes in the state values along the same trajectory (the degree of delay depends on the choice of the threshold of adaptive truncation and the sliding window size, see Section 3.3 and Algorithm 1 for more details). Hence upon the ”delayed” identification of the ”irreversible” transition at state st, we remove the state as well as a certain number of states preceding st (i.e., {st−τ , st−τ+1, . . . , st}) from ”learning via retracing”. We note that the choice of the sliding window size, S, threshold, η and the number of preceding states to remove, τ , are all chosen based on hyperparameter tuning at this stage. We leave the automatic determination of the parameters associated with adaptive truncation to future work.
We additionally propose an intermediate approach for ruling out the “irretraceable” states from “learning via retracing”, based on fixed scheduling of truncating out the final proportions of each collected episode (but still use these samples for representation learning in the forward direction). In many episodic environments, (e.g., LunarLander; Brockman et al. (2016)), each episode terminates upon either successful completion of the task or failing the task (e.g., crashing in LunarLander). For the failure cases, which is more often during the early phase of training, even the transitions that leading to the failure states might be considered “irreversible” (falling android; Figure 1b). Hence early truncation of each episode is able to rule out many such “irreversible” transitions from contributing to representation/model learning. We additionally propose annealing schedules for the truncation proportion as given the training of the controller, successful completion will be more frequently encountered comparing to the failure cases.
In Section 5.4, we argued that the effects of the truncation mechanism is correlated with the degree of “irreversibility” of the task. We augment the original CCWM with the fixed truncation scheme to empirically evaluate the relationship between truncation and the degree of “irreversibility”. The complete evaluation of fixed truncation on all the tasks listed in the main text (Figure 3a) is shown in Figure 7. We observe that in tasks with large intuitive degrees of “irreversibility”, such as walker-run, hopper-stand, quadruped-run, where the agent often goes through abrupt changes in its state (e.g., falling over), truncation generally yields great improvement in performance, whereas the improvement is less significant in other tasks.

heading: E FURTHER INVESTIGATION OF THE EFFECTS OF “LEARNING VIA
text: RETRACING” IN THE SAMPLE EFFICIENCY OF LEARNING
In order to demonstrate that “learning via retracing” indeed brings additional “valid” supervision signals for model learning, we need to examine whether the retracing operation is correctly learned. In Figure 8 we show the t-SNE embeddings of the true and retraced latent states over a trajectory of 500 steps of the trained agent on the cheetah run task (Van der Maaten and Hinton, 2008). We observe a similar dominant ring structure in both embeddings. This matches our intuition, since during running, the simulated cheetah repeats its states periodically to resume running at a high speed. This is more clearly observed from a sub-trajectory towards the end of the entire episode in Figure 8b, where most states are densely distributed on both sides of the ring and more loosely distributed on the paths connecting the two sides for both the original and retraced states. Moreover, during the early stage of the episode (darker dots in Figure 8a), the agent starts to accelerate from stationary and the states of the agent are periodic in nature but differs from the states when the agent is running at full speed. Such early-phase structure is correctly captured by the retracing operation. Collectively these evidence support that the retracing is correctly learned hence provides useful supervision signals for model training.

heading: F LONG-RANGE PREDICTION RECONSTRUCTION
text: In Figure 11, we show the complete 45-step predictive reconstructions for CCWM and Dreamer on “walker walk” and “cheetah run” tasks. We see that CCWM is able to retain accurate predictions over the entire 45 predictive time-steps in the cheetah task, whereas Dreamer fails to do so. By looking more carefully at Figure 11, latent rollouts in the embedding space learned by CCWM yield less accurate predictions for the “irrecoverable” states (e.g., when the agent is falling, see Section 5.4 for further discussion). This property of the learned representation enables “task-aware planning”, i.e., the behaviour learning agent is trained to acquire more precise control policy such that the “irrecoverable” states that might hinder the overall learning process are minimally encountered, leading to improved efficiency and quality of learning.

heading: G FULL ABLATION STUDIES ON THE RETRACE LOSS FUNCTION
text: In Figure 12, we show the full ablation studies of the retracing loss function in CCWM. We observe that the bisimulation metric retracing loss function consistently outperforms the other alternatives (L2 and reconstruction error). This result conforms with our hypothesis that cycle-consistency supervision given bisimulation metrics poses a stronger constraint on the learning of the representation, leading to a learned representation that better respect the cycle-consistency properties of the environment.

Review summaries:
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. 

The paper investigates a self-supervised approach for learning the state representation in RL tasks. In addition to the predictive (reconstruction) supervision in the forward direction, the authors include a “retraced” transitions for representation/model learning, by enforcing the cycle-consistency constraint between the original and retraced states. The authors claim that this facilitates stronger representation learning and improve upon the sample efﬁciency of learning. As it is not always possible to find such a cycle consistency between two states and for counteracting the negative impacts brought by the “irreversible” transitions, the authors add a novel adaptive “truncation” mechanism.

This paper considers state representation learning problem in deep RL. It exploits the cycle-consistency supervision and develops a “learning via retracing” approach. Such supervision signals are generally inherent in existing data and does not need additional interaction with environment, which leads to better sample efficiency. Learning from predictive supervision from temporally forward and backward directions reveals information from both the future and past to the target state, leading to more accurate latent state inference. In particular, the paper proposes the Cycle-Consistency World Model (CCWM) along with practical considerations (e.g., adaptive truncation to remove irreversible states), under the model-based framework based on generative dynamics modeling (CCWM).

This paper proposes a self-supervised approach for learning the state representation of RL tasks. The main contribution apart from prior works is to involve additional 'retracing' trajectory samples in representation training, which are trained by minimizing the propensity between retraced samples and forward posterior samples. They also proposed an intuitive way to mitigate the irreversibility in RL dynamics.