L EARNING S TATE R EPRESENTATIONS VIA R ETRACING IN R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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 efﬁciency 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 efﬁciency and asymptotic performance.
I NTRODUCTION
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 signiﬁcant redundancy, negatively impacting the policy learning.
The ﬁeld 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 efﬁciency 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 efﬁcient policy learning.
Thus we hypothesise that by augmenting representation learning with “learning via retracing” in RL, we can signiﬁcantly improve the sample efﬁciency of learning, which is a long-standing issue that plagues the practical applicability of deep RL algorithms.
Beyond improved sample efﬁciency, 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.
[Figure]
“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 (cid:48) for some action a , but s (cid:48)
(cid:54)→ s for all possible actions a (Figure 1b).
Under these situations, naively forcing the representations for s and ˇ s (the retraced state given s (cid:48) 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 ﬁrst 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 efﬁciency and asymptotic performance, whilst providing additional advantages such as stronger generalisability and extended planning horizon.
P RELIMINARIES 2. P ROBLEM F ORMULATION
We consider reinforcement learning problems in Markov Decision Processes (MDPs).
An MDP can be charaterised by the tuple, M =
(cid:104)S , A , R , P , γ (cid:105) , 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.,
[Equation]
, where s t
[Equation]
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 efﬁcient learning of the optimal policy using existing RL algorithms.
2. G ENERATIVE M ODELLING OF D YNAMICS
Modelling the transition dynamics using a sequential VAE-like structure enables joint learning of the latent representation and the associated latent transitions.
Speciﬁcally, the dynamics model is deﬁned in terms of the following components (see also top part in Figure 2b).
[Equation]
where ψ = { ψ 1 , ψ 2 , φ } and θ represent the parameters associated with the recognition and generative models respectively.
Latent variables z ∈ Z are introduced for more ﬂexible modelling of the distribution of the observed variables.
A variational approximation is employed since the true posterior p ( z t +1 | z t , a t , O t +1 ) is usually intractable in practice.
At each time step t , the agent receives an observation O t , which is then embedded into a context vector e t =
q φ ( O t ) .
After initialisation, the latent space vector is rolled out in a forward fashion given the action a t , yielding one-step prediction into the future following the variational transition distribution q ψ 1 ( z t +1 | z t , a t ) (we assume standard ﬁrst-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).
[Equation]
1 | z t , a t , e t +1 )
||
q ψ 1 ( z t +1 | z t , a t )]] , (3)
The variational free energy loss function consists of the reconstruction error, p θ ( O t +1 | z t +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 ( z t +1 | z t , a t , e t +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).
M ETHOD
We ﬁrstly 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 ﬁnally propose a novel adaptive truncation scheme for dealing with the “irreversibility” issue in “learning via retracing”.
3. L EARNING VIA R ETRACING
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 = { O 1 , . . .
, O T } , 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.
[Equation]
where d is some metric of the observable space (e.g., the L 2 distance).
The function f ( e ; T ) is some function specifying the schedule for predictions, e.g., f ( e t , T )
=
T ( e t ) =
ˆ 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 (cid:48) ) , we deﬁne the “reversed” transitions being the tuple ( s (cid:48) , a (cid:48) , s ) , where a (cid:48) is the “reversed” action.
In situations where the same set of rules apply to forward and backward transitions, the reversed transition given a (cid:48) 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).
s
ˆ
sˇ s a ˇ ad ( s, ˇ s ) (a) ··· z t z t + z t + ··· a t a t + a t +
O t
O t +
O t + ˜ O t ˜ O t + ˜ O t + ρ ζ ˇ a t + ρ ζ ˇ a t + ˇ z t +2 ˇ z t +1 ˇ z t ˇ z t +2 =
ˆ z t +2 L retrace forwardpass backwardpass
[Figure]
Hence by utilising the additional reversed transition for representation learning, we improve the sample efﬁciency 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 identiﬁcation 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 efﬁciency, 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 ﬁltering (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.
3. C YCLE -C ONSISTENCY W ORLD M ODEL
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 deﬁne a reverse action approximator, ρ ζ : Z × Z
→ A , which takes in a tuple of latent states ( z t +1 , z t ) and outputs an action ˇ a t +1 that approximates the “reversed” action that leads the transition from z t +1 back to z t .
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 efﬁciency 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 { O 1: T +1 , a 1: T } , we ﬁrstly perform a one-step forward sweep through all timesteps to compute the variational prior and posterior estimates of latent states.
[Equation]
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).
[Equation]
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
[Equation]
where ˆ R ( r | z ) represents the learned reward distributions (we assume stochastic rewards), and W 2 ( · , · ) 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 11).
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 sufﬁcient 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.
[Equation]
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.
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.
3. R EVERSIBILITY AND T RUNCATION
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 identiﬁcation 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.
; see Appendix D for further discussion).
Hence we use the value function as an indicator for sudden change in the agent’s state.
Speciﬁcally, for each sampled trajectory, we ﬁrstly compute the values of each state-action pair using the current value function approximator, [ Q ( z 1 , a 1 ) , . . .
, Q ( z T , a T )] .
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-deﬁned 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.
R ELATED W ORKS
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 ﬁxed 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 ﬁrst 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 classiﬁcation 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, speciﬁcally 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.
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.
E XPERIMENTAL S TUDIES
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.
5. E XPERIMENT S ETUP
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
[Figure]
model-based RL framework using model-based rollouts 1 (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.
; Figure. 3a).
Baselines : We implement Dreamer 2 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.
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.
We report the asymptotic scores for the model-free algorithms due to the large gap in sample efﬁciency comparing to the model-based methods.
5. E VALUATION ON C ONTINUOUS C ONTROL T ASKS
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 efﬁciency 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 efﬁciency and ﬁnal convergence performance.
In Appendix E, we show that “learning via retracing” indeed provides additional “valid” supervision for model learning, hence improving sample efﬁciency.
T r u e
CC W M
D r e a m e r
Context
T r u e
CC W M
D r e a m e r w a l k e r w a l k c h ee t a h r un
Moreover, in the “Cheetah Run” task, CCWM converges at ∼ 900 score with ∼ 5 × 10 5 steps, whereas Dreamer, by the time it has received 1 × 10 6 training steps, is yet to reach a comparable score.
This demonstrates that “learning via retracing” brings more beneﬁts beyond plain sample efﬁciency, 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 conﬁrms 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.
5. Z ERO -S HOT T RANSFER
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.
Speciﬁcally, we modify a number of basic conﬁgurations of the cheetah run task, such as the mass of the agent and the friction coefﬁcients between the joints.
The details of the changes can be found in Appendix A. Despite the increased sample efﬁciency of CCWM over Dreamer during training (Figure 3b), both methods converge at similar values at 2 × 10 6 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 signiﬁcance 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 signiﬁcantly outperforms Dreamer on more non-trivial modiﬁcations to the original task.
The introduction of retracing also improves the stability of zero-shot transfer in general (reduced variance in evaluation).
These conﬁrm our previous hypothesis that “learning via retracing” improves the ability of within-domain generalisation.
5. A DAPTIVE T RUNCATION OF “L EARNING VIA R ETRACING ”
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
C HANGED COMPONENTS
CCWM
[Table]
[Table]
[Table]
R + M
[Table]
R + F
[Table]
+ S
[Table]
+ M +
[Table]
M + S +
[Table]
probably due
[Figure]
Minimal improvement is observed on tasks with low degree of “irreversibility” (walker walk in (a)); whereas signiﬁcant improvements are observed on tasks with high degree of “irreversibility” (hopper stand in (b)).
From Figure 5, we observe that by augmenting CCWM with the proposed ﬂexible truncation schedule, the performance signiﬁcantly improves, yielding consistently better sample efﬁciency than both Dreamer and the original
CCWM .
For tasks with less degree of “irreversibility”, such as walker-walk, we do not observe signiﬁcant 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 ﬁxed scheduling in the Appendix D (Figure 7).
D ISCUSSION
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 sampleefﬁciency 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 ﬁnd that “learning via retracing” brings a number of merits comparing to its counterparts that only uses forward rollouts.
In addition to improved sample efﬁciency 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.
A I MPLEMENTATION D ETAILS
We implement the latent transition dynamics model of CCWM using the RSSM (Hafner et al., 2019), utilising a Gated Recurrent Unit (GRU) as the core component (Chung et al., 2014) in combination with multi-layer perceptrons (MLP).
The RSSM is used for modelling the predictive inferences, q ψ 1 ( z t +1 | z t , a t ) and q ψ 2 ( z t +1 | z t , a t , e t +1 ) , from Eq. 2.
The context embedding function is given by a convolutional encoder that deterministically embeds the high-dimensional observation into a context vector.
The dimension of the latent space is set to equal 32 , there is low sensitivity of the performance with respect to the dimension of the latent embedding.
The generative function q θ ( O t | z t ) is implemented with the combination of an MLP and a deconvolution network (Zeiler et al., 2010), outputting the means of the Gaussian distributions for each pixel value of the decoded observation.
The reward distribution ˆ R ( r | z ) is modelled by an MLP, outputting a univariate Gaussian distribution.
The supervision of the model learning is based on the ELBO and the bisimulation metric as described in Section 3.
In the implementation of our policy agent, A3C, we implement the critic (value network) using an
MLP, with latent state inputs.
The value network models a Gaussian distribution with means and variances.
The policy network also describes a Gaussian distribution, with means and variances parameterised by an MLP.
We used distributed actors for interacting with the environment for more efﬁciency collection of sampled episodes as in (Mnih et al., 2016).
The value function is optimised by maximising the probability of the “ground-truth” values estimated by the lambda return (Sutton and
Barto, 2018).
The policy network is updated using policy gradient estimated based on generalised advantage estimation (GAE) ()schulman2015high.
We use the Adam optimiser for updating the parameters for the model, value function, and policy networks (Kingma and Ba, 2014).
The speciﬁc conﬁgurations for the neural network implementations is summarised in Table 2.
Note that the activation functions for all non-output layers are assumed to be ReLU activation function unless otherwise stated.
All neural network implementation are carried out using TensorFlow (Abadi et al., 2015) and
TensorFlow Distributions (Dillon et al., 2017).
For the actual training, the batch size is chosen to be 64 , and all sampled trajectories are taken to be 50 timesteps long.
Greedy evaluation is performed every 10 4 training steps.
The reported evaluation scores are averaged values over 5 random seeds.
We adopt the same scheme for setting the action repeats equal to 2 across all tested environments as in (Hafner et al., 2020a).
The parameter λ controlling the weights of the retrace auxiliary loss in Eq.
8 is set to 1 . 0 .
The discounting factor for the expected value function is set to 0 .
The predictive reconstruction (Figure 4, Figure 10) is performed by ﬁrstly applying the dynamics model to the ﬁrst 5 frames of the trajectory, and starting from the posterior latent estimates at step 5 , the latent predictions are performed given solely the action inputs, and the predicted latent estimates are then decoded back into the observable space.
To ensure fair comparison, as well as to demonstrate that CCWM improves model learning beyond solely sample efﬁciency, the model of the baseline
Dreamer is provided further training, until the training steps is twice the training steps of CCWM and the ﬁnal performance reaches a comparable score of CCWM .
Speciﬁcally, for both the walker walk and cheetah run tasks, the Dreamer model used for reconstruction is provided with 2 × 10 6 training steps.
For the implementation details of the generalisability experiments in Section 5.3, we manually changed the corresponding settings of the agent in the task-dependent XML ﬁle (see e.g., https://github.com/deepmind/dm_control/blob/master/dm_control/ suite/cheetah.xml ).
We choose the semantically interpretable physical quantities that we hypothesise to alter the underlying task MDP, including changing mass from 14 to any one of the list: [6 , 8 , 10 , 20] , changing the friction constants from (0 . 4 , 0 . 1 , 0 . 1) to any element of the list: { ( i, j, k ) | i ∈ { 0 . 6 , 0 . 8 } , j ∈ { 0 . 4 , 0 . 5 } , k ∈ { 0 . 3 , 0 . 4 }} , changing the stiffness constants from 8 to any of the element in { 2 , 6 , 10 } .
We directly evaluate (without any further training)
the trained CCWM and Dreamer on the original task structure on the tasks corresponding to each of the above changes and any combination of the individual changes. The change in reward setting is merely a [Table]
constant reward reshaping across all transitions, and should not result in any difference in the task dynamics, hence the evaluations, and we use this as the sanity check and a baseline.
Note that in Eq.
7, we follow (Zhang et al., 2020) to replace the 1-Wasserstein distance with the 2-Wasserstein distance, since the 2-Wasserstein distance between two Gaussians has closed-form expression:
[Equation]
where || · || F is the matrix Frobenius norm.
This admits analytical computation instead of having sample from the joint distribution hence reducing the variance of loss function.
The python implementation of the CCWM can be found in the supplementary materials.
B P SEUDOCODE FOR CCWM
The pseudocode for CCWM training is shown in Algorithm 1.
In the current instance, we show the full algorithm with the augmentation of adaptive truncation (Section 3.3). Given the adaptive
Algorithm
Given: policy π ( a | s ) , encoders q φ ( O ) , q ψ 1 ( z (cid:48) | z, a ) , q ψ 2 ( z (cid:48) | z, a, q φ ( O (cid:48) )) , decoder q θ ( O | z ) , reverse action approximator ρ ζ ( z, z (cid:48) ) , latent horizon K , average loss L avg = 0 , counter c = 0 , adaptive truncation threshold, η , adaptive truncation sliding window size, S , current estimate of the Q-function, Q γ , CCWM parameters:
Θ = { φ, ψ 1 , ψ 2 , θ, ζ } .
Input: Batch of N sampled trajectories of length T , { O nt n : t n + T , a nt n : t n + T } Nn =1 . for n = 1 to N do
O =
{ O nt n + k } Tk =1 , a = { a nt n + k } T − 1 k =0 .
for k = 1 to T − K do Compute prior ˆ z k
+1: k + K and posterior estimates ˜
z k +1:
k + K
using Eq.
Compute the sliding window average, [ ¯ Q 1 , . . , ¯ Q K − S ] , where ¯ Q i = 1 S (cid:80) Ss =1 Q γ ( ˜ z
i + s , a i + s )
Compute the step-wise difference, ∆ =
[ | ¯ Q 2 − ¯ Q 1 ¯ Q 1 | , . .
, | ¯ Q K − S − ¯ Q K − S − 1 ¯ Q K − S −| ]
{+Compute the adaptive truncation mask, M ∈ { 0 , 1 } K , where M k = 
+}
if ∧ kj = k − S (∆ j < η ) , and 0 otherwise
Compute retraced states ˇ z k : k
+ K using Eq.
6 Compute the latent model loss
L k using Eq.
Update
L avg ← cc +1 L avg + 1 c +1 L k .
Increment counter, c ← c + 1 .
end for end for
Compute the gradient, ∇ Θ L avg . Update the model parameters, Θ
[Equation]
truncation, we need to slightly modify the loss function in Eq. 8.
[Equation]
where M nτ is the entry τ of the mask vector M n with values in { 0 , 1 } , which we describe in
Algorithm 1.
C M ODEL -F REE I NSTANTIATION OF “L EARNING VIA R ETRACING ”
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.
· · · z t z t + z t + · · ·
O t
O t +
O t + a t −a t a t + ˜ O t ˜ O t + ˜ O t + ˇ z t +2 ˇ z t +1 ˇ z t ˇ z t +2 =
ˆ z t +2 L retrace forwardpass backwardpass
D F URTHER D ISCUSSION ON T RUNCATION AND THE D EGREE OF “I RREVERSIBILITY ”
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 = (cid:104)S , A , R , P , γ (cid:105) and the corresponding DeepMDP ¯ M = (cid:104) ¯ S , A , ¯ R , ¯ P , γ (cid:105) (see Section 2.2 in Gelada et al. for detailed deﬁnition of DeepMDP), let d ¯ S be a metric over ¯ S , φ : S → ¯ S be the embedding function, L ∞ ¯ R = sup s ∈S ,a ∈A |R ( s, a ) − ¯ R
( φ ( s ) , a ) | and L ∞ ¯ P = sup s ∈S ,a ∈A W ( φ P ( ·| s, a ) , ¯ P ( ·| φ ( s ) , a )) be the global loss function (for training the
DeepMDP, where W ( · , · ) is the Wasserstein distance).
For any K ¯ V -Lipschitz policy ¯ π ∈ ¯Π , the representation φ guarantees that for all s 1 , s 2 ∈ S , and a ∈ A ,
[Equation]
The proof of Theorem 1 can be found in Appendix A.2 in Gelada et al.
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 Qvalues 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 Figureand 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 ﬁgure 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.
[Figure] [Figure] [Figure] 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 × 10 5 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 identiﬁcation 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 additionally propose an intermediate approach for ruling out the “irretraceable” states from “learning via retracing”, based on ﬁxed scheduling of truncating out the ﬁnal 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.
, 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 ﬁxed truncation scheme to empirically evaluate the relationship between truncation and the degree of “irreversibility”.
The complete evaluation of ﬁxed 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 signiﬁcant in other tasks.
[-states         75 74727069 68 676664 63616059 58 57 5452 
-]
E F URTHER I NVESTIGATION OF THE EFFECTS OF “L EARNING VIA R ETRACING ” 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. [Figure] [Figure] [Figure]
[Figure]
[Figure]
retraced {+states [Figure] 50 51 52 53 54 57 58 59 60 61 646566 67 68 69 7072 73true states         75 74727069 68 676664 63616059 58 57 5452 
 50 49 484644 434139 383635 34 33 323026242220 19 18 17 16 151312 11 109  5 4 3 2retraced states+}
[Figure]
[Figure]
We show the TSNE-embeddings of the state representations over a trajectory of length 500 steps in
Hopper Stand task, with a standard CCWM agent (without adaptive truncation) that has received 2 × 10 5 training steps.
Clear clustered structured can be observed in the visualisation, which indicates: (a) the existence of ”irreversible” transitions, and (b) the majority of the transitions are ”reversible”, hence ”learning via retracing” should still apply to facilitate improved sample efﬁciency.
T r u e
CC W M
D r e a m e r
Context
T r u e
CC W M D r e a m e r w a l k e r w a l k c h ee t a h r un 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.
F L ONG -R ANGE P REDICTION R ECONSTRUCTION
In Figure 10, 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 10, 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 efﬁciency and quality of learning.
G F ULL A BLATION S TUDIES ON THE R ETRACE L OSS F UNCTION
In Figure 11, 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 (L 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.
[Figure]