L EARNING S TATE R EPRESENTATIONS VIA R ETRACING IN R EINFORCEMENT L EARNING
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 for the downstream reinforcement learning tasks.
We introduce Cycle-Consistency World Model ( CCWM ), a concrete model-based instantiation of learning via retracing .
Additionally we propose a novel adaptive “truncation” mechanism for counteracting the negative impacts brought by “irreversible” transitions such that learning via retracing can be maximally effective.
Through extensive empirical studies on visual-based continuous control benchmarks, we demonstrate that CCWM achieves state-of-the-art performance in terms of sample efﬁciency and asymptotic performance, whilst exhibiting behaviours that are indicative of stronger representation learning.
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 and might only contain partial information with respect to the associated ground-truth states, hence 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 such algorithms rely on predictive (reconstruction) supervision for representation learning, such that the effects of actions in the observable space are maximally preserved in the learned representation space.
Here we argue that existing methods do not fully exploit the supervisory signals inherent in the data.
Additional 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.
Hence, with “ learning via retracing ”, we obtain more training samples for representation learning without additional interaction with the environment (twice as much as existing approaches in tasks that admit perfect reversibility across all transitions).
Therefore, we hypothesise that by augmenting representation learning with “learning via retracing”, we can signiﬁcantly improve
[Figure]
the sample efﬁciency of representation learning and the overall RL task, 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 use information from both the future and the past for the inference of states, 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.
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 and generalisation to other task topologies (Section 5.3).
We conjecture that such imaginative retracing could be neurally implemented by the reversed hippocampal “replay” that has been observed in both rodents and humans (Foster and Wilson, 2006; Penny et al., 2013; Liu et al., 2021) (see Section 6 for further discussion). One problem that hinders the successful application of “learning via retracing” is that 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 no action a (cid:48) such that s (cid:48) → s (Figure 1b). Under these situations, naively enforcing the similarity between the representations of s and ˇ s (the retraced state given s (cid:48) and a , see
Section 3) leads to a suboptimal representation space, potentially hindering RL training. Hence in order to maximally preserve the advantages brought by “learning via retracing”, it is essential to identify “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 adaptive truncation (Section 3.3). “Learning via retracing” can be integrated into any representation learning method that utilises a transition model, under both the model-free and model-based RL frameworks.
Here 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 setting.
We empirically evaluate 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 beneﬁts such as stronger generalisability and extended planning horizon, indicative of stronger representation learning is achieved.
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 ∼ p
[Equation]
·| s t − 1 , a t − 1 ) for t = 1 , 2 , . . .
[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 (prior) distribution q ψ 1 ( z t +1 | z t , a t ) (we assume standard ﬁrst-order Markovian structure in the latent space).
The dynamics model is trained via maximum likelihood learning, and due to intractability, we adopt standard amortised variational inference, and 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; Higgins et al., 2016)).
[Equation]
The variational free energy objective consists of the reconstruction error, p θ ( O t +1 | z t +1 ) , and the KLdivergence 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 modelbased instantiation based on generative dynamics modelling 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 models for forward and “reversed” transitions).
Given a set of observations, O = { O 1 , . . .
, O T } , a dynamics model,
M : S ×
A × S → [0 , 1] , most existing methods of self-supervised representation learning 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 ; M , a ) is some function specifying the schedule for predictions, e.g., f ( e t , M , a t ) =
M ( e t , a t ) =
ˆ e t [Equation] 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 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]
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 efﬁciency of learning without additional interaction with the environment.
As mentioned previously, 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 for the downstream RL task.
We note that the overall RL agent could still beneﬁt from the representations obtain from learning via retracing even in tasks without perfect reversibility, such as a moving car where the external state feature, such as the absolute location, and controllable internal state features, such as the velocity and acceleration, are not jointly reversible.
In such cases, we expect learning via retracing would dissect out the controllable internal features from the external features and prioritise the training with respect to such features.
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 to use 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 concrete 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 D).
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., 2011; Zhang et al., 2020).
[Equation]
where ˆ R ( r | z ) represents the learned reward distributions (we assume stochastic rewards), and W 2 ( · , · ) represents the 2 -Wasserstein distance (see Eq.
11 in Appendix D).
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 “retraced” samples.
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 A.
We note that “learning via retracing” is also applicable under the model-free setting, we describe one such instantiation in Appendix B.
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 C 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”.
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 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 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.
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, with the latter being the main difference between the premises of PlayVirtual and CCWM.
By using the same dynamics model for both the forward and reversed transitions, we hope to exploit and embed the context prior of reversible transitions into the learnt representations, hence the induced advantages extend beyond the explicit advantage of improved sample efﬁciency, but also stronger generalisation, improved predictive rollouts (leads to more accurate policy gradient hence improving policy training).
E XPERIMENTAL S TUDIES
The experimental studies aim at examining if “learning via retracing” truly helps with the overall RL training and planning, improves the generalisability of the learned representation, and whether the truncation schedule proposed in Section 3.3 deals with the irreversibility of some state transitions.
5. E XPERIMENT S ETUP
CCWM can be combined with any value-based or policy-based RL algorithms.
We implement CCWM with a simple actor-critic RL agent with generalised advantage estimation based on standard model-
[Figure]
based RL framework using model-based rollouts 1 (Sutton and Barto, 2018; Konda and Tsitsiklis, 1999; Schulman et al., 2015).
We base our experimental studies on the challenging visual-based continuous control benchmarks for which we choose 8 tasks from the DeepMind Control Suite (Tassa et al.
; Figure. 3a).
The details of training and the architecture can be found in Appendix D.
Baselines : We implement Dreamer 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 (without adaptive truncation introduced in Section 3.3) 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 (see also Appendix E).
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.
We note the relative poor performance of CCWM on the Hopper Stand task, which might be accredited to the inherent large degree of irreversibility of the task, we shall examine how to deal with such irreversible tasks in Section 5.4. 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
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
[Figure]
VALUE
S IGNIFICANT ?
( MEAN ± [Figure] D .) ( MEAN ± 1 S .
D .)
(3 S .
F .)
( α = 0 .
[Figure]
07 ± 36 .468 .
82 ± 94 .7 .
[Figure]
[Figure]
Y ES
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] 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 D 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 D. 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.
The overall performance on zero-shot transfer of our approach is comparable with Dreamer on simpler transfer tasks (full results shown in Appendix 5.3), 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 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.
[Figure]
From Figure 5, we observe that augmenting CCWM with the proposed ﬂexible truncation schedule yields signiﬁcant performance increase on the Hopper-Stand task, with consistently better sample efﬁciency than both Dreamer and the standard CCWM .
For tasks with less degree of “irreversibility”, such as Walker-Walk (Figureright), 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.
Similar patterns are observed across the other tasks (no noticeable improvement except for the HopperStand/Hop tasks). We extend our discussion on adaptive truncation in Appendix C.
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 a 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 intuitions from the computational neuroscience literature (Penny et al., 2013), 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.
Moreover, as discussed previously, “learning via retracing” can enable the learning of reversible “controllable” internal states of the agent even when external features (e.g. location) are irreversible.
Such representation could be combined with task-speciﬁc global contextual information/representation provided by an independent system to facilitate stronger policy learning.
The independent system could be supported by the hippocampal area, where reversed hippocampal “replay” of the allocentric “cognitive map” could in turn drive reverse replay of egocentric sensory or motoric representations.
A CKNOWLEDGEMENT
C.Y. is supported by the DeepMind studentship funded through UCL CDT in Foundational Artiﬁcial
Intelligence.
N.B. is supported by the Wellcome Trust.
The authors would like to thank Maneesh
Sahani, Peter Dayan, Jessica Hamricks, and the anonymous reviewers for helpful comments and discussions.
A 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 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.
Algorithm
1: 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 , θ, ζ } .
2: Input: Batch of N sampled trajectories of length T , { O nt n : t n + T , a nt n : t n + T } Nn =1 . 3: for n
[Table]
1 to N do 4:
O =
{ O nt n + k } Tk =1 , a = { a nt n + 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.
[Table]
Compute the sliding window average, [ ¯ Q 1 , . .
, ¯ Q K − S ] , where ¯ Q i = [Equation] [Equation]
z i +
[Equation] , a i + s ) 8:
Compute the step-wise difference, ∆ =
[ | ¯ Q [Equation] Q [Equation]
[Equation] | ] 9:
Compute the adaptive truncation mask, M ∈ { 0 , 1 } K , where M k = 1 if ∨ kj = k − S (∆ j < η ) , and 0 otherwise 10:
Compute retraced states ˇ z k : k + K using Eq.11:
Compute the latent model loss L k using Eq. 9.
Update L avg ← cc +1 L avg + 1 c +1 L k . 13:
Increment counter, c ← c + 1 .
14: end for 15: end for 16: Compute the gradient, ∇ Θ L avg .
17: Update the model parameters, Θ ←
[Equation]
B 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 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” to 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 groundtruth forward actions for the retracing operations.
The learned representation in this case would beneﬁt the downstream model-free RL agent since the resulting state representation is efﬁcient for the prediction of future states.
We note a key difference between our model-free instantiation of learning via retracing and PlayVirtual (Yu et al., 2021), that we have consistently employed probabilistic
· · · z
[Figure]
z
[Figure]
z t
[Figure]
O
[Figure]
O t +
[Figure]
O t +
[Figure]
forwardpass backwardpass
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.
C 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 admit the situation where two states with different values collapse 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 an ”irreversible” transition.
The latter interpretation can be substantiated with the visualisation of the low-dimensional embeddings of the learned representations at different stages in Figure 7 and Figure 8.
Figureshows 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 7(a) and (b), please refer to the full discussion of Figure 7 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 8, 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 8), 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” identiﬁcation of the ”irreversible” transition at state s t , we remove the state as well as a certain number of states preceding s t (i.e., { s t − τ , s t − τ +1 , . . .
, s t } ) 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 ﬁ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 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.
D 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 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.
[Table]
Conﬁgurations for the neural network implementations of CCWM .
For our implementation of the adaptive truncation technique (Section 3.3), the associated parameters that require pre-speciﬁcation are the detection threshold, η , adaptive truncation sliding window size,
S . There are also the optional parameters, the number of states preceding the detected changes to be omitted from ”learning via retracing” (which is set to equal to S by default, see Line 9 in Algorithmand Section 3.3), τ , and the number of initial steps to be omitted from adaptive truncation (due to under-training of the critic, which is set to 0 by default), ξ . The default values for the parameters we used for the empirical evaluation shown in Figure 5 are: η = 0 . 10 , S = 10 , τ = 5 , ξ = 1 × 10 5 .
The hyperparameters are determined through standard grid search, and we leave the automatic determination of the parameters associated with adaptive truncation to future work.
The python implementation of the CCWM can be found in the supplementary materials.
[Figure] [Figure] [Figure] [Figure]
retraced states
[Figure]
[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. [Figure] leads 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 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 7 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 7b, 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 7a), 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.
[Figure]
C HANGED COMPONENTS
CCWM D REAMER P - VALUE S IGNIFICANT ?
( MEAN ± 1 S . D .) ( MEAN ± 1 S . D .)
(3 S .
F .)
( α = 0 .
R 630 . 40 ± 6 .629 . 00 ± 47 .5 . 22 × 10 −
N O
R + M 635 . 37 ± 40 .597 .
43 ± 87 .7 .
84 × 10 − N O R + F 643 . 58 ± 9 .649 .
27 ± 3 .9 .
76 × 10 −
N O
R + S 634 . 80 ± 10 .646 . [Table]
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.
H F ULL R ESULTS OF Z ERO -S HOT T RANSFER
Please refer to Table 3 for full results on the zero-shot transfer experiment.
Note that the R -changes only rescale the reward function by a constant, hence should not count towards a transfer task.