Learning State Representations via Retracing in Reinforcement Learning
Authors: Changmin Yu
Dong Li
Jianye HAO
Jun Wang
Neil Burgess
Venue: ICLR 2022 Poster (2022)
Keywords: Representation learning
model-based reinforcement learning

Abstract: We propose learning via retracing, a novel self-supervised approach for learning the state representation (and the associated dynamics model) for reinforcement learning tasks. In addition to the predictive (reconstruction) supervision in the forward direction, we propose to include "retraced" transitions for representation/model learning, by enforcing the cycle-consistency constraint between the original and retraced states, hence improve upon the sample efficiency of learning. Moreover, learning via retracing explicitly propagates information about future transitions backward for inferring previous states, thus facilitates stronger representation learning 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 efficiency and asymptotic performance, whilst exhibiting behaviours that are indicative of stronger representation learning. 

Published as a conference paper at ICLR 2022
LEARNING STATE REPRESENTATIONS VIA RETRACING IN REINFORCEMENT LEARNING
Changmin Yu' Dong , Jianye Hao Jun Wangh?, Neil Burgess UCLL, London, United Kingdom
Huawei Noah's Ark Lab
College of Intelligence and Computing, Tianjin University
{changmin, yu. 11;; n. sarress]@uucl. ac. uk; (lidong106) 06; w. ) w @.:)eeii. com; j jianye.
cn
ABSTRACT
We propose learning via retracing, a novel self-supervised approach for learning the state representation (and the associated dynamics model) for reinforcement learning tasks. In addition to the predictive (reconstruction) supervision the forward direction, propose to include "retraced" transitions for representa- tion/model learning, by enforcing the cycle-consistency constraint between the original and retraced states, hence improve upon the sample efficiency of learning. Moreover, learning via retracing explicitly propagates information about future transitions backward for inferring previous states, thus facilitates stronger repre- sentation learning for the downstream reinforcement learning tasks. We introduce Sycle-Consistence World Model (CCWM), concrete model-based instantiation of learning via retracing. Additionally we propose a novel adaptive "truncation" mechanism for counteracting the negative impacts brought by "irreversible" tran- sitions such that learning via retracing can be maximally effective. Through extensive empirical studies on aasual-based continuous control benchmarks, we demonstrate that CCWM achieves state-of-thee oott--of-hheaart art performance in terms of sam- ple efficiency and asymptotic performance, whilst exhibiting behaviours that are indicative of stronger representation learning.
INTRODUCTION
Recent developments in deep reinforcement learning (RL) have made great progress in solving complex control tasks (Mnih et al., 2013; Levine al., 2016; Silver al., 2017; Vinyals 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 sigmificant redundancy and might only contain partial information with respect to the associated ground-truth states, hence negatively impacting the policy learning. The field of representation learning offers a wide range of approaches for extracting useful information from figh-dimensional data (with potentially sequential dependency structure) (Bengio 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 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 significantly improve
Please send any enquiry changmin n.
yu. nuyu. 90uc1 ac. uk and burgess@uel .ac.ul

---

Published as a conference paper at ICLR 2022
the sample efficiency of representation learning and the overall RL task, which is Long-standing issue that plagues the practical upplicability of deep RL algorithms. Beyond improved sample efficiency, joint predictive supervision in tem- porally forward and backward directions use in- formation 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.
(a)
(b)
Figure 1: Motivation of "learning via retrac- ing". (a): Retracing in navigation tasks yields faster representation learning and potentially sup- ports stronger generalisation; (b): "Inreversible" transitions (graphical demonstration from the DeepMind Control Suite Tassa et al. (2018)).
As a motivating example, consider rat navigating towards a cheese a cluttered environment (Figure la). Upon first visit the goal state, multiple imaginative retracing trajectories can be randomly simulated. By constraining the temporal cycle-consistency the retracing transitions, the rat quickly builds a state representation that accurately preserves the transitions in the around the actual forward trajectory taken by the rat. Moreover, all retracing simulations pass through the two "doors", allowing the rat 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 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' for some action but no action a' such that s' + s (Figure lb). Under these situations, naively enforcing the similarity between the representations of s and (the retraced state given and a, see Section 3) leads to suboptimal representation space, potentially hindering RI. training. Hence in order to maximally preserve the advantages brought by "learning via retracing", identify "inreversible" transitions and rule them out from representation learning via retracing. To this end, we propose a novel dynamically regulated approach for identifying such "inreversible" states, which we term adaptive truncation (Section 3.3).
essential to
"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- Cice-Connisteeccy World Model (CCWM), 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. Experi- mental results show that CCWM achieves state-of-the-art performance in terms of sample efficiency and asymptotic performance, whilst providing additional benefits such as stronger generalisability and extended planning horizon, indicative of stronger representation learning achieved.
2
PRELIMINARIES
2.1 PROBLEM FORMULATION
We consider reinforcement learning problems in Markov Decision Processes (MDPs). An MDP can be charaterised by the tuple, M (S, A, R, P. 7), where S, are the state and action spaces, respectively; R : S > R is the reward function (we assume determinisite reward functions unless stated otherwise), P S A + [0. is the transition distribution of the task dynamics; 7 R is the discounting factor. The control policy, क S x A + [0, 1]. represents a distribution over actions each state. The goal is to learn the optimal policy, n², such that the expected future reward is maximised across all states, i.c.,
Σ ' ()))s = vs E S, where 8; ~ P(-| for = 1,2,...
n = arg max En xEII
(1)
Here we consider tasks in which the perceivable observation space, O, is high-dimensional, due to either redundant information or simply because only visual inputs are available. Hence it is necessary
2

---

Published as a conference paper at ICLR 2022
to learn embedding function, Ø O Z. such that the embedded observation space, o0, could act as in the MDP, to support efficient learning of the optimal policy using existing RL algorithms.
2.2 GENERATIVE MODELLING OF DYNAMICS
Modelling the transition dynamics using sequential VAE-like structure enables joint learning of the latent representation and the associated latent transitions. Specifically, the dynamics model is defined in terms of the following components (see also top part in Figure 2b).
Observation (context) embedding: e = 96 (Ot),
Latent posterior distribution: p(*++1 |2, ae, Oc+1);
Latent variational transition (prior) distribution: 9v, (2++1 |2, at),
(2)
Latent variational posterior distribution: 94, (#++1 (2, at, c(+1),
Generative distribution: po (0++1/2+111
where = [ih1, 2h2, o) and & represent the parameters associated with the recognition and generative models respectively. Latent variables = E 2 are introduced for more flexible modelling of the distribution of the observed variables. A variational approximation is employed since the true posterior pactiiorrp(z-ee p(z-1/zz 4z, Oc+1) is usually intractable in practice.
At each time step the agent receives an observation Ot, which is then embedded into a context vector é go (O.). After initialisation, the latent space vector is rolled out in a forward fashion given the action at, yielding one-step prediction into the future following the variational transition (prior) distribution qei, (2+-1 4x) (we assume standard first-order Markovian structure in the latent space). The dynamics model 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. are learned by maximising the variational free energy (also known as the ELBO; (Wainwright and Jordan, 2008; Higgins et al., 2016)).
Liorward (0-1) E E -0.
[log po (0++1 +++1 - (DxL. (4pz (ze+ oe, C++1 Ilge (2++1/24, a )]}
(3)
The variational free energy objective consists of the reconstruction error, pa (O++1 [++1). 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 (9 (Ze+1 |2e, de. Ct+1)) serving the main inference engine of the optimal latent representations. Note that the 3 parameter controls the degree of factored structure (disentanglement) of the latent code, which by default is set to 1 (Higgins et al., 2016).
METHOD
We firstly introduce learning via retracing in its most general format, then provide a concrete model- based instantiation based on generative dynamics modelling We finally propose novel adaptive truncation scheme for dealing with the "irreversibility" issue in "learning via retracing".
3.1 LEARNING VIA RETRACING
We always assume the usage of an approximate dynamics model, regardless of the overall RL agent being model-free or model-based (the dynamics model would only be used for representation learning under the model-free setting, hence it is possible to have separate dynamics models for forward and "reversed" transitions). Given a set of observations, O = {O1,..., Or), a dynamics model, M S x A S + [0, 1J. most existing methods of self- supervisdd representation learning involve learning an encoder (Es) and a decoder (D&), trained via minimising the predictive reconstruction.
L(,,0) = d(Or, ( (E A AuT--::M)
(4)
where d is some metric of the observable space (e.g., the L2 distance). The function ((e;M, some function specifying the schedule for predictions, e.g., [(e,, (M,,,,,n a,,a). = MMe,; (t) = ē 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 transition tuple, (s, a, s'), we define the "reversed" transitions being the tuple (s', a'. s). where a' is the "reversed" ac- tion. In situations where the same set of rules apply to forward and backward transitions, the reversed transition given d' could contribute to representation learning via Eq. as an additional training
3

---

Published as a conference paper at ICLR 2022
á
sample (utilising potentially different loss function from the forward supervision, see Fig- ure 2a).
g
Hence by utilising the additional reversed transi- tion for representation learning, improve the sample efficiency of learning without additional interaction with the environment. As mentioned previously, in situations where perfect reversibil- ity is not preserved across all transitions, such "inreversible" transitions could negatively impact the overall learning. Correct identification of such states is hence essential for the successful implementation of "learning via retracing". To this end, we propose novel adaptive truncation scheme in Section 3.3.
d(s,b
8)
(a)
forward
a2
0c+1
0x+2
pass
2++2
O₂
O++1
Os+a
Ö
Despite the intuitive simplicity of learning via retracing, offers a number advantages com- paring to existing representation learning meth- ods. In addition to the improved sample effi- ciency,
PC
PC
Evethane
def1
de42
backward
interpret learning with "reversed" transitions as explicit inference of the current la- tent state given future information. In combina- tion with the forward predictive supervision, the joint learning dynamics is similar the smooth- ing operation in dynamical systems, which is often superior than filtering (corresponds to us- ing 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 benefit from the representations obtain from learning via retracing even in tasks without perfect reversibility, such moving car where the external state feature, such 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.
pass
2x
2++1
2+42
(b)
Figure 2: Graphical illustration of "learning via retracing". (a) "learning via retracing" addition- ally constrains the similarity between the retraced and original states for representation learning; (b) Graphical model of CCWM. The empty circle and filler square nodes represent the stochastic and de- terministic variables, respectively.
We have introduced learning via retracing in its most general form, where 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 free or model-basee based; whether not to use a separate "reversed" dynamics model; loss function for constraining the "retraced" transitions; deterministic or probabilistic dynamics model, just to name few. Below we provide one concrete instantiation of learning via retracing, the Cycle- Cycle-Consistency World Model (CCWM), under the model-based framework based on generative dynamics modelling.
3.2
CYCLE- COLSSSCONSSYYNCCY WORLD MODEL
CCWM is a model-based RL agent that utilises 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 2, ă and é, respectively.
We use the similar dynamics model described in Section 2.2 (top panel in Figure 2b). We additionally define reverse action approximator, Pc 2 x Z A, which takes in tuple of latent states (21+1, 21) and outputs an action /t+1 that approximates the "reversed" action that leads the transition from 2++1 back to 24. Instead of introducing a separate "reversed" dynamics model, we the same dynamics model for both the forward and retracing transitions, which lead to improved sample efficiency of model learning in addition representation learning. The parameters of p. Ç, can be
4

---

Published as a conference paper at ICLR 2022
either learned jointly with the model in end-to-end fashion, trained separately (see Appendix D). The graphical model of CCWM is shown in Figure 2b.
During training, given sample trajectory {O1:T+1, 01.7). we firstly perform one-step forward sweep through all timesteps to compute the variational prior and posterior estimates of latent states.
2 ~ qv1 (z)
(5)
Er+1 Dv. (z)2r, ar, go (o())
for r = o,..., where Eo 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).
2, qoh (1 where 0,+1 pc(()+1, for r =
(6)
The forward and retracing predictive supervision separately contributes 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 based on constraining the deviation of the retraced states from the original states. Intuitively, this utilises the temporal sycls-conss of the transition dynamics: assuming that the dependedd transition mapping is invertible across all timesteps (Dwibedi et al., 2019). The loss function for constraining the coces-conss 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 al., 201 1; Zhang al., 2020).
Cretrssd (Zi, tc) Es, [112: l111
P ( R( (()
W22( (-1z, 7 ((,)), qu, (- (
where R(r|=) represents the learned reward distributions (we assume stochastic rewards), and (2 represents the 2- D-Wasserstein distance (see Eq. 11 in Appendix D). The advantage choosing the bisimulation metric 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 Monte Carlo estimate based on sampled "retracing" states, but empirically we observe that one "retracing" sample is sufficient as we do not observe noticeable improvements for increasing the number of "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 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 linear combination of the forward and "retracing" loss functions.
C(O,0.C) NT
Σ [rooward (07;0, ()) ACyetrace (27, 22, vb ()], m=11=1
(8)
where À is the scalar multiplier for the retrace loss, and N is the batch size. We implement CCWM using a Recurrent State- Spacee Model (RSSM; Hafner et al. (2019)). The complete pseudocode for CCWM training is shown in Algorithm 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.3 REVERSIBILITY AND TRUNCATION
As we noted above, perfect reversibility is not always present across all transitions in many environ- ments. For instance, consider the falling android presented in Figure 1b, it is trivial 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
5

---

Published as a conference paper at ICLR 2022
representation learning (and dynamics model learning in CCWM). Here we propose an approach to deal with such "irreversibility".
Our approach is based adaptive identification of "irreversible" transitions. We propose that the value function of the controller (c.g., an actor-critic agent) possesses some information about the continuity of the latent states (Gelada et al. (2019); see Appendix C for further discussion). Hence we use the value function as an indicator for sudden change in the agent's 's state. Specifically, for each sampled trajectory, we firstly compute the values of each state-action pair using the current value function approximator, [Q(21, a1) ai)..... QZTT, (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 -)-length vector (Q1..... Qr-s]. Any drop/increase in the sliding averages (above some pre-defined threshold) indicates sudden change the value function, hence a sudden change the latent representation given the continuity conveyed by the value function. Given some timestep, T, at which the sudden change occurs, we then remove the transitions {2--S:r;0, S5r from "learning via retracing". Such adaptive scheduling allows us to deal with "irreversibility".
4 RELATED WORKS
Representation learning in RL. Finding useful state representations that could aid RL tasks has long been studied. Early works have investigated representations based on a fixed basis such as tile coding and Fourier basis (Mahadevan, 2005; Sutton and Barto, 2018). With the development of deep learning techniques, recent works explored automatic feature discovery based on neural network training, which can be categorised into three large classes. The first class of methods studies the usage of data augmentation for broadening the data distribution for training more robust feature representation (Laskin al., 2020; Kostrikov et al., 2020; Schwarzer al., 2021; Yarats et al., 2021). The second class explores the role of auxiliary tasks in learning representations, such as weakly-supervised classification and location recognition, for dealing with sparse and delayed supervision (Lee et al., 2020b; Mirowski et al., 2017; Oord et al., 2018). The third class of methods, specifically tailored to model-based RL models, leverages generative modelling of 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 al., 2020a; Schrittwieser et al., 2020; Hafner et al., 2020b).
Cynlesconsistency 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 the validation of matches between cycling through multiple samples. We adopt similar design principles for sequential decision-making tasks: rollouts in temporally forward direction alone yield moner-constr learning of the world model. By additionally incorporating backwards rollouts into model learning in self-supervised fashion, we enforce the inductive bias that the same transition rules govern the dynamics of the task.
Concurrent to work, Yu et al. (2021) proposed PlayVirtual, model-free RL method that integrates a similar cycle-consistency philosophy into training representations with data augmen- tations (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 free and utilising a separate reversed dynamics model, with the latter being the main difference between the premises of Play Virtual and CCWM. By using the same dynamics model for both the for- ward 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 efficiency, but also stronger generalisation, improved predictive rollouts (leads to more accurate policy gradient hence improving policy training).
EXPERIMENTAL STUDIES
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.1 EXPERIMENT SETUP
CCWM can be combined with any value- based or policy- RL algorithms. We implement CCWM with simple actor-critic RI. agent with generalised advantage estimation based standard model-

---

Published as a conference paper at ICLR 2022
Cheetah Run
Finger Spin
Hopper Hap
Hopper Stand
Evaluation Retur
1000
1000
1000
750
200
750
500
500
100
250
250
250
0.2
0.6
1.0
0.5
1.5
2.0
0.2
0.6
1.0
Quadruped Run
Reacher Easy
Walker Run
Walker Walk
Evaluation Return
1000
1000
400
400
200
500
200
500
250
0.5
1.0
1.5
2.0
0.5
1.0
1.5
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Steps
1c6
Steps
1c6
Steps
C3C
Dreamer
ABC (state)
D4PG (state)
SAC (state)
(b)
Figure 3: Evaluation CCWM on DeepMind Control Suite. (a): Graphical demonstration of
selected continuous control task environments, from left to right: hopper stand/hop, walker run/walk, finger spin, reacher easy, cheetah run, quadruped run. (b): Average evaluation returns (+:1 s.d.) during training random seeds). "Learning via retracing" generally improves the performance of learning from pixel inputs in presented tasks comparing to the main baseline Dreamer agent (which could approximately be viewed as CCWM without retracing). CCWM reaches the asymptotic performance of tatte-of-the-aas model-free methods (SAC, D4PG at 10 steps) on several tasks.
based RL framework using model-based rollouts' (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 al. (2018); 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 wate-of-the-art world-model-type model-based RI. agent on visual-based continuous control tasks. We also compare with the following model-free baselines: SAC Haarnoja et al. (2018), D4PG Barth-Maron et al. (2018), A3C Mnih et al. (2016). We implement the SAC agent given the state inputs and directly report the asymptotic performance of the D4PG and A3C algorithms from Tassa et al. (2018). We report the asymptotic scores for the model-free algorithms due to the large gap in sample efficiency comparing to the model- methods.
5.2 EVALUATION ON CONTINUOUS CONTROL TASKS
The performance of CCWM and selected baseline algorithms 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 efficiency of learning (see also Appendix E). CCWM outperforms Dreamer on of the selected tasks, and comparable Dreamer on 2 of the remaining tasks, in terms of both the sample efficiency and final 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, the "Cheetah Run" task, CCWM converges at 900 score with whereas Dreamer, by the time it has received x 10 training steps, is yet reach comparable
10 steps,
Note that for the maximally fair comparison with main baseline, Dreamer (Hafner et al., 2020a), utilise the policy agent Dreamer in order to fully demonstrate the utility of "learning via retracing"

---

Published as a conference paper at ICLR 2022
Context
10
15
20
25
30
True
We Meeer walk
Dreamer COWM
Dreamer CCWM True
cheetah run
Figure 4: Qualitative comparison of long-range predictive reconstruction of CCWM and Dreamer. Predictive rollouts over 30 time-steps given the actions are computed using the rep- resentation models. CCWM consistently generates more accurate predictive reconstructions further into the future than Dreamer, with CCWM becoming noticeably inaccurate by 25 30 timesteps. and Dreamer by 10 15 timesteps. See implementation details and further discussion in Appendix
CHANGED COMPONENTS
CCWM
DREAMER
P-VALUE
SIGNIFICANT?
(MEAN +1 S.D.)
(MEAN #1 S.D.)
(3 S.F.)
(a 0.01)
+ M + R +M F
628.07 +: 36.95
468.82 + 94.73
7.27 x 10 G
YES YES
641.89 = 28.67
562.58 = 93.91
3.97 x 10
Table Evaluation of trained CCWM and Dreamer on ability zero-shot transfer cheetah run tasks with different configurations. (R: Reward; M: Mass; F: Friction; S: Stiffness.)
score. This demonstrates that "learning via retracing" brings more benefits beyond plain sample efficiency, i.e., doubling the training steps does not eliminate the performance gap. This corresponds to our hypothesis that by explicitly conveying future information back 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, provide further training for Dreamer whenever necessary. such that the asymptotic performance comparable with CCWM (see Appendix for details). Figure shows that CCWM consistently yields more accurate predictive reconstructions over longer time span, on both the walker walk and cheetah run tasks. The empirical evidence confirms our hypothesis that by incorporating "learning via retracing" into model learning enables the resulting latent space to support more accurate latent predictions. hence leading to stronger behaviour learning. Increased range of accurate latent prediction additionally enables CCWM to perform better planning. We provide further analysis of predictive reconstruction in Appendix F.
5.3 ZERO-SHOT TRANSFER
Based the motivation that "learning via retracing" will improve the generalisability of the agent (Figure Ia), we empirically test the generalisability of CCWM on the basis of zero-shot transfer tasks. Specifically, modify a number of basic configurations of the cheetah run task, such the mass of the agent and the friction coefficients between the joints. The details of the changes can be found in Appendix D. Despite the increased sample efficiency (Figure 3b), both methods converge at similar values at 2 x 10th 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 (+: s.d.) of both agents over 15 random seeds, as well as the one-sided (-test statistics and significance of the difference between the two sets of evaluations in Table 1. The overall performance zero-shot transfer of our approach is comparable with Dreamer on simpler transfer tasks (full results shown in Appendix 5.3), and significantly outperforms Dreamer more con-trivial modifications to the original task. The introduction of retracing also improves the stability of zero-shot transfer in general (reduced variance in evaluation). These confirm our previous hypothesis that "learning via retracing" improves the ability of within-domain generalisation.
CCWM over Dreamer during training
8

---

Published as a conference paper at ICLR 2022
5.4 ADAPTIVE TRUNCATION OF "LEARNING VIA RETRACING"
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 "inreversibility" of the task comparing to the others such that the naive representation learning by enforcing "learning via retracing" leads to suboptimal training.
From Figure 5, we observe that augmenting CCWM with the proposed flexible truncation schedule yields significant performance increase on the Hopper-Stane task, with consistently bet- ter sample efficiency than both Dreamer and the standard CCWM. For tasks with less degree of
walker walk
hogper stand
Auorage Evsiuation
CCWM
Dreamter
OCWM AT
X4
CEWMAAT
"irreversibility". such as Walker- Walk (Figure right), we do not observe significant improve- ment by the introduction of adaptive truncation since the amount of negative impacts were al- ready minimal in the original settings. Similar patterns observed across the other tasks (no noticeable improvement except for the Hopper- Stand/Hop tasks). We extend our discussion adaptive truncation in Appendix C.
Figure 5: Evaluation of Adaptive Truncation tasks with varying degrees of "inreversiiility." Minimal improvement is observed tasks with low degree of "irreversibility" (walker walk in the left panel); whereas significant improvements are observed on tasks with high degree of "irreversibil- ity" (hopper stand in the right panel).
DISCUSSION
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 the-art art model-based and model-free methods on number of challenging continuous control benchmarks, in terms of both the sample- efficiency 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 critical role in model-based decision maaoonnmaainngg 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 instantia- tion of the reversed hippocampal replay. Similar to intuitions from the computational neuroscience literature (Penny et al., 2013), we also find that "learning via retracing" brings a number of merits comparing its counterparts that only uses forward rollouts. In addition to improved sample effi- 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 "con- trollable" internal states of the agent even when external features (c.g. location) are irreversible. Such representation could be combined with task-specific global contextual informatton/reprreenttiio 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.
9

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
C.Y. is supported by the DeepMind studentship funded through UCL CDT in Foundational Artificial Intelligence. N.B. is supported by the Wellcome Trust. The authors would like thank Maneesh Sahani, Peter Dayan, Jessica Hamricks, and the anonymous reviewers for helpful comments and discussions.
REFERENCES
M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro. G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, Goodfellow, Harp. G. Irving, M. Isard, Y. Jia, Jozefowiez. L. Kaiser, M. Kudlur, J. Levenberg, D. Mané, R. Monga, S. Moore, D. Murray, C. Olah, M. Schuster, Shlens, Steiner, Sutskever, K. Talwar, P. Tucker, V. Vanhoucke, V. Vasudevan, F. Viégas, O. Vinyals, P. Warden, M. Wattenberg, M. Wicke, Y. Yu, and X. Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https: http://ttttoowlwwc..oo/// org// Software available from tensorflow.org.
G. Barth-Maron, M. W. Hoffman, D. Budden, W. Dabney, D. Horgan, Dhruva, A. Muldal, N. Heess, and T. Lillicrap. Distributed distributional deterministic policy gradients. ArXiv, abs/aab4..0 1804088617, 2018.
Y. Bengio, A. Courville, and P. Vincent. Representation learning: A review and perspectives. IEEE transactions on pattern analysis and machine intelligence, 15(8):1998-1825 2013.
G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba. Openai gym. arXiv preprini
2016.
L. Buesing, T. Weber, S. Racanière, Eslami, D. J. Rezende, D. Reichert, F. Viola, F. Besse, K. Gregor, D. Hassabis, and D. Wierstra. Learning and querying fast generative models for reinforcement learning. ArXiv, abs/ 802.03006, 2018.
J. Chung, C. Gulcehre, K. Cho, and Y. Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint 11Xi5:1412.3555 2014.
J. V. Dillon, I. Langmore, D. Tran, E. Brevdo, S. Vasudevan, D. Moore, B. Patton, A. Alemi, M. Hoffman, and R. A. Saurous. Tensorflow distributions. arXiv preprint arXi::1711.1064. 2017.
D. Dwibedi, Y. Aytar. J. Tompson, P. Sermanet, and A. Zisserman. Temporal cycle-consistency learning. 2019 JEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 1801-1810, 2019.
Evans and N. Burgess. Coordinated hippocampal-entorhinall replay as structural inference. In NeurlPS, volume 32. NIPS, 2019.
N. Ferns, P. Panangaden, and Precup. Bisimulation metries for continuous markov decision processes. SIAM J. Comput., 00:1662-1714, 2011.
D. Foster and M. Wilson. Reverse replay of behavioural sequences hippocampal place cells during the awake state. Nature, 440:680-653 68300 2006.
C. Gelada, S. Kumar, J. Buckman, O. Nachum, and M. G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference Machine Learning, pages 21 2179- PMLR, 2019.
D. R. Ha and J. Schmidhuber. World models. ArXiv, abs/ htls/1803. 1002, 2018.
Haamoja, A. Zhou, P. Abbeel, and Levine. Soft actor-critic: Off- off-policy maximum entropy deep reinforcement learning with stochastic actor. In ICMI., 2018.
D. Hafner, T. Lillicrap, Fischer, R. Villegas, D. R. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. ArXiv, abs/1811 0111811.04551.1 2019.
10

---

Published as a conference paper at ICLR 2022
D. Hafner, T. Lillicrap. J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. ArXiv, abs/ abs/1912.01603, 2020a.
D. Hafner, T. Lillicrap, M. Norouzi, and J. Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02/99 2020b.
1. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. 2016.
R. Kalman. A new approach to linear filtering and prediction problems. 1960.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412::1112..980 2014.
V. R. Konda and Tsitsiklis. Actor-critic algorithms. In NIPS, 1999.
I. Kostrikov, D. Yarats, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
M. Laskin, K. Lee, A. Stooke, L. Pinto, P. Abbeel, and A. Srinivas. Reinforcement learning with augmented data. arXiv preprini arXiv:2004. 14990, 2020.
A. X. Lee, A. Nagabandi, Abbeel, and Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. ArXiv, abs/1907.0000 .0055,, 2020a.
L. Lee, Eysenbach, R. Salakhutdinov, S. Gu, and C. Finn. Weakly-supervised reinforcement learning for controllable behavior. ArXiv, abs/2004.02860, 2020b.
S. Levine, C. Finn, T. Darrell, and P. Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 17(1):1334-1373 2016.
Y. Liu, M. G. Mattar, T. E. Behrens, N. D. Daw, and R. J. Dolan. Experience replay is associated with efficient nonlocal learning. Science, 372(6544), 2021.
S. Mahadevan. Proto-value functions: Developmental reinforcement learning. In Proceedings of the 22nd international conference on Machine learning, pages 553-560, 2005.
M. Mattar and N. Daw. Prioritized memory access explains planning and hippocampal replay. Nature neuroscience. 21:1609 1617, 2018
Mirowski, R. Pascanu, Viola, Soyer, A. Ballard, Banino, M. Denil, Goroshin, Sifre, K. Kavukcuoglu, D. Kumaran, and R. Hadsell. Learning to navigate in complex environments. ArXiv, abs/161 1.03673, 2017.
V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. A. Riedmiller. Playing atari with deep reinforcement learning. ArXiv, abs/1312.5602, 2013.
V. Mnih, A. P. Badia, M. Mirza, A. Graves, Т. Lillicrap. T. Harley, D. Silver, and K. Kavukcuoglu. Asynchronous methods for deep reinforcement learning. ArXiv, abs/1602. 2016.
K. Murphy. Machine learning: a probabilistic perspective. 2012.
A. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. ArXiv, abs/1 005/1877.00 80780073374 007403348 2018.
Osband, Blundell, A. Pritzel, and B. Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29, 2016.
W. D. Penny, P. Zeidman, and N. Burgess. Forward and backward inference in spatial cognition. PLoS Comput Biol, (12))e1003388) 2013.
J. Schrittwieser, I. Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lockhart, D. Hassabis, T. Graepel, T. Lillicrap, and D. Silver. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588 7839:604-609, 2020.

---

Published as a conference paper at ICLR 2022
J. Schulman, P. Moritz, S. Levine, M. Jordan, and P. Abbeel. High dishesiimeni continuous control using generalized advantage estimation. arXiv preprint orXiv:1506.02456. 2015.
M. Schwarzer, A. Anand, R. Goel, R. D. Hjelm, A. Courville, and P. Bachman. Data-efficient reinforcement learning with self-predictive representations. In ICLR, 2021.
D. Silver, J. Schritwieser, Simonyan, Antonoglou, Huang, A. Guez, Hubert, L. Baker, M. Lai, A. Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676)) 354 359, 2017.
R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
Y. Tassa, Y. Doron, A. Muldal, T. Erez, Y. Li, D. Casas, D. Budden, A. Abdolmaleki, J. Merel, A. Lefrancq, T. Lillicrap, and M. A. Riedmiller. Deepmind control suite. ArXiv, abs/1 8010006000 2018.
L. Van der Maaten and G. Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11) 1), 2008.
O. Vinyals, T. Ewalds, S. Bartunov, Georgiev, A. S. Vezhnevets, M. Yeo, A. Makhzani, H. Küttler, J. Agapiou, J. Schrittwieser, al. Starcraft ii: A new challenge for reinforcement learning. arXiv preprint arXiv:1008. 7708.0448, 2017.
M. Wainwright and M. Jordan. Graphical models, exponential families, and variational inference. Found. Trends Mach. Learn., 1:1-305, 2008
7. Yang, W. Chen, Wang, and Xu. Improving neural machine translation with conditional sequence generative adversarial nets. arXiv preprint 103.0477, 2017.
D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Mastering visual continuous control: Improved augmenee reinforcement learning. arXiv preprint arXi::2107.09645, 2021.
T. Yu, C. Lan, W. Zeng, M. Feng, and Z. Chen. Playvirtual: Augmenting cycle-consistent virtual trajectories for reinforcement learning. arXiv preprint arXivv2106.04152, 2021.
M. Zeiler, D. Krishnan, G. W. Taylor, and R. Fergus, Deconvolutional networks. In 2010 IEEE Computer Society Conference on computer vision and pattern recognition, pages 2528 2535. IEEE, 2010.
A. Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant representations for reinforcement learning without reconstruction. ArXiv, abs/2006.10744 2020.
T. Zhou, P. Krähenbühl, M. Aubry, Q. Huang, and A. A. Efros. Learning dense correspondence via 3d-guided cycle consistency. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 117-126, 2016.
J.-Y. Zhu, T. Park, Isola, and A. Efros. Unpaired image-to-iman translation using cycle- consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pages 2223-2232, 2017.
12

---

Published as a conference paper at ICLR 2022
A PSEUDOCODE 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.
((0,0,
(0; Cetrace ( 2 (2,,)
(9)
NT
where M. is the entry of the mask vector M with values in {0, 1,, which describe in Algorithm 1.
Algorithm Cycle-Consistenae World Model (CCWM)
I: Given:
policy x (als), encoders (O( qo (z'|z, a). qu (z'|z, a, qo()))) decoder qa (O(2z), reverse action approximator Pç(2, 2'), latent horizon K, average loss Lavg 0, counter adaptive truncation threshold, 17, adaptive truncation sliding window size, S, current estimate of the Q-function, Q,, CCWM parameters: + { 11, 4/2, 8,
2: Input: Batch of N sampled trajectories of length T. {O (0 +n a +
3: for = to N do
4: = (Or
b. t a {a,, 17-1 k
5: for k = 1 to K do
6: Compute prior &2+1144K K and posterior estimates Ek+ 1ik using Eq. Compute the sliding window average, + (7(t+ss
7:
IQI,..., Qk-sl, where Q
Compute the step-wise difference, A -11-000 b.... QK QK- Q1 Compute the adaptive truncation mask, M (0, 1)* where MR. if V j=k-s(A) < <n)) and 0 otherwise
8:
QK
9:
10:
Compute retraced states EdSR K using Eq, 6 Compute the latent model loss Lk using Eq. 9. Update Lavg + c
11:
12:
c+1 Lavg + c+1 Cx. Increment counter, e c
13:
14:
end for
15:
end for
16:
Compute the gradient, Ve
Update the model parameters, O aV
B MODEL-FREE INSTANTIATION OF "LEARNING VIA RETRACING"
As mentioned in Section "learning via retracing" admits many degrees of freedom in its implemen- tation. CCWM provides one such instantiation under the model-hased RL. setting, here we provide an alternative model based "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 Virtual (Yu al., 2021), due 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 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 p in Figure 2b, and we only need to use the ground- truth forward actions for the retracing operations. The learned representation this case would benefit the downstream model-free RL agent since the resulting state representation is efficient for the prediction future states. We note a key difference between our model-free instantiation of learning via retracing and Play-Vitual Virtual (Yu et al., 2021), that have consistently employed probabilistic
13

---

a) Published as a conference paper at ICLR 2022
forward
O₂
Oc+1
Ox+2
pass
21
At++1
ar-1
at
0++1
Öz
Ött1
Ött2
Leethgge
2422 = 2+42
backward
2t
22+1
22+2 pass
Figure 6: Graphical illustration of a model-free instantiation of "learming via retracing". The forward model a state-space model and trained generative fashion under the variational principles. The retracing operations are now performed with separate "reversed" dynamics model (indicated by the red arrows). Given the independent "reversed" dynamics model, we can use the same action as in the forward model for retracing, removing the necessity of using the "reversed" action approximator.
models over deterministic models for modelling the embedding and latent transitions, which naturally provides posterior predictive uncertainty that can be used for various downstream tasks, such as exploration (Osband et al., 2016).
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 of many possible instantiations of "learning via retracing", leave further investigation to future work.
C FURTHER DISCUSSION ON TRUNCATION AND THE DEGREE OF "IRREVERSIBILITY"
In Section 3.3, we introduced the adaptive truncation general-purpon purpose method for improving the performance of CCWM in continuous control tasks. On a high level, we choose to remove the transitions from model/repressntation learning through "learning via retracing" whenever large changes in the Q-values occur, hence the "irreversibility" of the transitions depend on the local continuity the corresponding Q-values. Our argument is largely based on Theorem from Gelada et al. (2019), which we re-iterate below for consistency.
Theorem For an MDP M = (S.A.R,P.Y) and the corresponding DeepMDP M (S, A. R, P, 7) (see Section 2.2 Gelada et (2019) for detailed definition of DeepMDP), let dg be metric over S, o: > S be the embedding function, 1 sup and L suPsES,aeA W(GP(-ls, a), P(: a)) be the global loss function (for training the DeepMDP, where ((..)) is the Wasserstein distance). For any Ky-Lipschita policy II. the representation à guarantees that for all $1, $2 E S, and a A.
I + +KKoL R
1Q* (*2. s Kpdg(((s)) +2
(10)
1- 7
The proof of Theorem can be found in Appendix A.2 in Gelada et al. (2019). The high-level interpretation for Theorem 1 is that the learned representation space should admit the situation where two states with different values collapse onto the same representations. Theorem 1 tells us that given good representation (embedding function), the absolute difference between the Q-values
14

---

Published as a conference paper at ICLR 2022
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 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. Figure 7 shows the TSNE embeddings of a trajectory taken by a successfully trained agent in the Cheetah Run task (we focus on the left figure in both Figure 7(a) and (b), please refer to the full discussion of Figure 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 TSseedmbes of the representation along trajectory in an under-trained CCWM agent in the Hopper Stand task (~~ 2 x 10 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 view as the more "irreversible" task, large proportion of all transitions allows retracing operations, which can contribute as additional samples for the training of representation and the dynamics model through "learning via retracing". The quality of the overall representation learning is dependent on the correct identification of the "inreversible" 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 "inreversibility" of the transitions along the same trajectory, the adaptive truncation is introduced by detecting the changes in the average Q-values over sliding window instead of the single Q-values, which allows 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 I for more details). Hence upon the "delayed" identification of the "irreversible" transition at state 8t, we remove the state as well as certain number of states preceding st (i.e., {st-r, $t- .$ }) from "learning via retracing". We note that the choice of the sliding ............
window size, S, threshold, 17 and the number of preceding states to remove, T, are all chosen based on hyperparameter tuning at this stage. We leave the automatic determination of the parameters associated with adaptive truncation to future work.
We additionally propose an intermediate approach for ruling out the "irretraceable" states from "learning via retracing", based on fixed scheduling of truncating out the final proportions of each collected episode (but still use these samples for representation learning in the forward direction). In many episodic environments, (e.g., LunarLander; Brockman et al. (2016)), each episode terminates upon either successful completion of the task 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/mood 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 IMPLEMENTATION DETAILS
We implement the latent transition dynamics model of CCWM using the RSSM (Hafner et al., 2019), utilising 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, qv (Ze+a /2,, 0t) and qe (#t+1 &e, ag et+1,, from Eq. 2. The context embedding function is given by a convolutional encoder that deterministically embeds the nigh-dimensional observation into context vector. The dimension of the latent space is set to equal 32, there is low sensitivity of
15

---

Published as a conference paper at ICLR 2022
the performance with respect to the dimension of the latent embedding. The generative function qo (Or(2)) is implemented with the combination of an MLP and a deconvolution network (Zeiler et al., 2010), ouiputting the means of the Gaussian distributions for each pixel value of the decoded observation. The reward distribution (r((2) is modelled by an MLP, outputting 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 ML.. We used distributed actors for interacting with the environment for more efficiency 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) 1). The policy network is updated using policy gradient estimated based on generalised advantage estimation (GAE) Osschulman201 Shigh.
We use the Adam optimiser for updating the parameters for the model, value function, and policy networks (Kingma and Ba, 2014).
The specific configurations 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 training steps. The reported evaluation scores are averaged values over random seeds. We adopt the same scheme for setting the action repeats equal to 2 across all tested environments 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 set to 0.99.
The predictive reconstruction (Figure 4. Figure 10) is performed by firstly applying the dynamics model the first 5 frames of the trajectory, and starting from the posterior latent estimates step 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 efficiency, the model of the baseline Dreamer is provided further training, until the training steps is twice the training steps of CCWM and the final performance reaches a comparable score of CCWM. Specifically, for both the walker walk and cheetah run tasks, the Dreamer model used for reconstruction is provided with 2 x 10 training steps.
For the implementation details of the generalisability experiments in Section 5.3, we man- ually changed the corresponding settings of the agent in the task-dependent XML file (see e.g.. https: / //gthbb. com/ deepmind/ 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, to any element of the list: ((.j.k)j E {0.6, 0.8), j E {0.4, 0.5), k E {0.33 0.4)), changing the stiffness constants from 8 to any of the element {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.
_dm_controccoooloooooooooo control/ /
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- 2-Wasserstein distance between two Gaussians has closed-form expression:
2//313 Wo(N(m,, N,, E,) = ||, - M + 1
(11)
where ||. || is the matrix Frobenius norm. This admits analytical computation instead of having sample from the joint distribution hence reducing the variance of loss function.
16

---

Published as a conference paper at ICLR 2022
COMPONENT
ATTRIBUTE
VALUE
RSSM
INTERNAL STATE DIMENSION
256
MLP LAYER UNITS
256
MLP LAYER 2 UNITS
64
EMBEDDING FUNCTION
CONY LAYER NUMBER OF FILTERS
CONV LAYER KERNEL SIZE
CONV LAYER 2 NUMBER OF FILTERS
64
CONY LAYER 2 KERNEL SIZE
4
CONY LAYER NUMBER
128
CONV LAYER 3 KERNEL. SIZE
CONY LAYER NUMBER OF FILTERS
256
CONV LAYER 4 KERNEL SIZE
4
FINAL OUTPUT OPERATION
CONCATENATION
GENERATIVE MODEL
MLP LAYER 1 NUMBER OF UNITS DECONVOLUTION LAYER I NUMBER FILTERS
1024
128
DECONVOLUTION LAYER KERNEL SIZE
DECONVOLUTION LAYER 2 NUMBER OF FILTERS DECONVOLUTION LAYER KERNEL SIZE
64
5
DECONVOLUTION LAYER 3 NUMBER FILTERS
32
DECONVOLUTION LAYER KERNEL
DECONVOLUTION LAYER 4 NUMBER OF FILTERS
DECONVOLUTION LAYER KERNEL SIZE
6
REWARD MODEL
MLP LAYER 1 NUMBER OF UNITS
512
MLP LAYER 2 NUMBER OF UNITS
512
MLP LAYER 3 NUMBER OF UNITS
VALUE MODEL
MLP LAYER 1 NUMBER OF UNITS
512
MLP LAYER 2 NUMBER OF UNITS
512
MLP LAYER 3 NUMBER UNITS
512
MLP LAYER NUMBER OF UNITS
POLICY MODEL
FULLY CONNECTED LAYERS ACTIVATION
ELU
MLP LAYER NUMBER OF UNITS
512
MLP LAYER 2 NUMBER OF UNITS
512
MLP LAYER 3 NUMBER OF UNITS
512
MI.P LAYER 4 NUMBER UNITS
512
ML. LAYER 5 NUMBER OF UNITS
ADAM OPTIMISER
LEARNING RATE FOR MODEL LEARNING
x
LEARNING RATE FOR VALUE MODEL
8 x 10
LEARNING RATE FOR POLICY MODEL
x 10
Table 2: Configurations for the neural network implementations of CCWM.
For our implementation of the adaptive truncation technique (Section 3.3), the associated parameters that require pre-specification are the detection threshold, 7), 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 in Algorithm and Section 3.3), т, and the number of initial steps to be omitted from adaptive truncation (due to under- training the critic, which is set to by default), ç. The default values for the parameters 103.
used for the empirical evaluation shown in Figure 5 are: y 0.10, S = 10, 5, The hyperparameters are determined through standard grid search, and we leave the automatic determination of the parameters associated with adaptive truncation future work.
The python implementation of CCWM can be found at https: : / github.. com/ wwt.c/oooccaiiceeeeee http//hnnngnuuu _ode.
17

---

Published as a conference paper at ICLR 2022
rntraced
-10
(b)
Figure 7: Visualisation similarity between the ground-truth states and retraced states using tSNE. The two-dimensional embeddings for 500-step trajectory (a) and 76-step subtrajectory (b) shown. Color of the seatters indicates the temporal ordering of the states within the trajectories.
30
400
300
100
-30
-40
-30
-20
-10
10
20 30
Figure 8: Visualisation state representations along the same trajectory in Hopper Stand task. We show the TSNE-embeddings of the state representations over a trajectory of length 500 steps in Hopper Stand task, with standard CCWM agent (without adaptive truncation) that has received 2 x 10a 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 efficiency.
75
50-
7.
0-
100
65
25
95
6.
-3
55
-5
-50
25
-100
-50
100
Figure 9: Visualisation of the structure of the state representations with respect to the Q-values. We show the TSNE-embeddings of the state representations of 10000 randomly sampled states (from the cached replay buffer) of CCWM (left) and Dreamer (right) in the Walker Walk task. Visual inspections show that CCWM leads more disentangled state representations with respect to the state-values, hence potentially supports stronger generalisability (Higgins et al., 2016).
18

---

Published as a conference paper at ICLR 2022
Context
10
20
25
45
Thuc
7
Dreamor COMM Thuc Dreamer COWM
Figure 10: Qualitative comparison of long-range predictive reconstruction of CCWM and Dreamer. Predictive rollouts over 45 time-steps given the actions are computed using rep- resentation models. Comparing to Dreamer, CCWM consistently generates more accurate predictive reconstructions over longer time spans. CCWM generates accurate predictions up to steps on Walker task, and 35 steps Cheetah task; Dreamer generates accurate predictions up steps on Walker task, and ~ 15 steps on Cheetah task (evaluated over 5 randomly sampled trajectories).
E FURTHER INVESTIGATION OF THE EFFECTS OF "LEARNING VIA RETRACING" IN THE SAMPLE EFFICIENCY OF LEARNING
In order to demonstrate that "learning via retracing" indeed brings additional "valid" supervision signals for model learning. we need examine whether the retracing operation correctly learned. In Figure 7 we show the t-SNE embeddings of the true and retraced latent states over trajectory of 500 steps of the trained agent on the cheetah run task (Van der Maaten and Hinton, 2008). We observe 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 sub-trajen trajerrjectrry towards the end of the entire episode in Figure 7b, where most states are densely distributed on both sides the ring and more loosely distributed on the paths connecting the 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 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 correctly learned hence provides useful supervision signals for model training.
F LONG-RANGE PREDICTION RECONSTRUCTION
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 able 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 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 trained to acquire more precise control policy such that the "irrecoverable" states that might hinder the overall learning process are minimally encountered, leading to improved efficiency and quality of learning.
19

---

Published as a conference paper at ICLR 2022
Cheetah Run
Finger Spin
Hopper Hop
Hopper Stand
Evaluation Return
1000
200
1000
150
750
500
100
500
250
250
0.2
0.6
1.0
1.5
0.6
1.0
Quadruped Run
Reacher Easy
Walker Run
Walker Walk
Evaluation Retur
1.000
1000
750
200
500
500
200
250
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Steps
Steps
1e5
Stops
166
Steps
12
CCWM-A3C
reconstruction
Figure 11: Full ablation studies of the retracing loss function. CCWM with bisimulation metrics as the retracing loss function consistently outperforms the alternatives using the L2 and reconstruction retracing error.
CHANGED COMPONENTS
CCWM
DREAMER (MEAN S.D.)
P-VALUE (3 S.E.)
SIGNIFICANT? (a 0.01)
(MEAN S.D.)
R
630.40 10.40166 6.49
629.00 1 47.01
5.22 10
No
+
M
635.37 + 40.12
597.43 11.43:87.00 87.68
7.84 x 10
No
R+F
643.58 + 9.29
3.96
9.76 x 10
No
+
634.80 1: 10.92
646.07 -.00777777
9.98 x 10
No
R + M + F
628.07 1 36.95
94.73
7.27 x 10
YES
+M
S
641.89 1 28.67
562.58 1 93.91
3.97 x 10 3
YES
Table 3: Evaluation of trained CCWM and Dreamer on the ability of zero-shot transfer in cheetah run tasks with different configurations. (R: Reward; M: Mass; F: Friction; Stiffness.)
G FULL ABLATION STUDIES ON THE RETRACE LOsS FUNCTION
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 (L2 and reconstruction error). This result conforms with our hypothesis that cycle-consistency supervision given bisimulation metrics poses stronger constraint the learning of the representation, leading to learned representation that better respect the cycle-consistency properties of the environment.
H FULL RESULTS OF ZERO-SHOT TRANSFER
Please refer to Table 3 for full results on the zero-shot transfer experiment. Note that the R-changer only rescale the reward function by constant, hence should not count towards transfer task.
20


Reviews:
[{"review_id": "zZPE4kbGMdz", "reviewer": "Reviewer_dZvG", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The strengths of the paper involve solid experiments and well-structured writing. However, I think some details are not illustrated clearly, so that I have the following questions:\n1. Is the representation learning part jointly trained with RL algorithms or pretrained? \n2. In the $L_{retrace}$, why retracing samples $\\check{z}$ are from the variational predictive distribution but forward samples $\\tilde{z}$ are from the variational posterior distribution? \n3. How the “reversed” action policy $p_{\\eta} $ is learned? by $L_{retrace}$?\n4. The proposed method for dealing with “irreversibility” is too intuitive. I think it might fail in heavily irreversible environments. Why the proposed truncation method is not involved in your full algorithm in the appendix?\n5.  As you mentioned CCWM could be combined with any existing RL algorithms, why not try one or two to show that the new representation learning loss could actually improve the performance. I think the performance of CCWM is comparable with Dreamer given the huge additional computational complexity. \n", "summary_of_the_paper": "This paper proposes a self-supervised approach for learning the state representation of RL tasks. The main contribution apart from prior works is to involve additional 'retracing' trajectory samples in representation training, which are trained by minimizing the propensity between retraced samples and forward posterior samples. They also proposed an intuitive way to mitigate the irreversibility in RL dynamics.", "summary_of_the_review": "The proposed representation learning method for RL is interesting, which involves addtional retracing samples for training. However, more explanations are needed to make it clear. I expect to see more elaborations.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "k4_RFcesn2-", "reviewer": "Reviewer_VSAG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n\n- Interesting technique to retrieve more supervisory signals from the data in RL\n- Mathematically sound approach, with clear figure 2 to visually understand.\n- Extensive experiments to evaluate the idea and interesting ablations and investigations described.\n- Extensive details of algorithm and implementation + code is available.\n \n##########################################################################\n\nCons:\n\n- A few inconsistencies in the paper that make me think that there were extra parts that have been cut out to reach the page limit without reformatting:\n    - Section 5.3: Details of modifications to the environment are said to be found in Appendix A, but Appendix A describes implementation details and not the environment.\n    - Appendix C, last sentence. Model-free instantiation is shown in Figure ?? Figure doesn't exist.\n\n##########################################################################\n\nQuestions during rebuttal period:\n\n- Address/clarify cons above.\n- How do you select the predefined threshold to find the sudden changes in the value function for your adaptive scheduling?\n- I can imagine from Table 1, that the components that you change in the cheetah-run environment are related to Reward, Mass, Friction, Stiffness. I don't understand how changing the reward would affect evaluating an episode, as it doesn't change the dynamic behavior of the physics simulator and so the actual trained model shouldn't have any problem in evaluating.\n- In Section 5.4, you write that the curve for hopper-stand in Figure 3b for CCWM lacks the adaptive truncation of irreversible states. Does this mean that all plots in Figure 3b for CCWM are trained on all transitions, with no truncation? Please clarify, as I imagined that the experimental evaluation of your method included both contributions (CCWM+adaptive truncation).\n- In algorithm 1: I don't understand in the input, why in ${O^n_{t_n:t_n+K}, a^n_{t_n:t_n+K}}$ the $t$ has a subscript $n$, which should represent the n-th sample in the batch. Also, I think there is a typo and the subscript should be to $t+T$ and not $t+K$.\n- In figure 8 (from Appendix F), you explain how CCWM yields less accurate predictions for the irrecoverable states, but to me the first row (true trajectory for Walker), doesn't look like an irrecoverable trajectory.\n\n##########################################################################\n\nTypos/formatting:\n- Section 3.3, first paragraph: \"Here we propose an approaches\" -> \"Here we propose an approach\"\n- None of the citations in Appendix A are in brackets anymore.\n- Appendix A: \"outputting the means the Gaussian distributions\" -> \"outputting the means of the Gaussian distributions\"\n- Appendix D, second paragraph: \"empirical evaluation\" -> \"empirically evaluate\"", "summary_of_the_paper": "This paper proposes \"learning via retracing\" as an approach to learn state representations, through matching a trajectory both in the forward and in the backward direction. This paper then introduces Cycle-Consistency World Model (CCWM) which is a model-based RL algorithm which learns through retracing. This method is sample efficient and provides good state representations for predicting the future and generalization. Furthermore, it proposes a value-based approach to identify \"irreversible\" transitions.\n", "summary_of_the_review": "Overall, I am for accepting. I like the idea and I think it's a novel approach for state representations in RL. My major concern is about the clarity of the paper, as written above. Hopefully the authors can address my concern in the rebuttal period.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "_OHOFELM_As", "reviewer": "Reviewer_HLp2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength: The cycle-consistency constraints could provide extra supervision signal (beyond forward predictive loss) for learning state representations and transition models, without additional interaction with environment. It also improves the representation learning performance (including the zero-shot transfer capability and long-range predictions) by obtaining better latent state inference. \n\nWeakness: The concept of cycle-consistency constraints have also been considered in the PlayVirtual work under model-free setting. Therefore, more thorough discussion about the novel contribution relative to PlayVirtual should be included.\n\nComments:\n- The authors may need to discuss how frequent such cycle-consistency constraints appear in practical applications. For example, it is more common in navigation-type applications? Listing more practical application settings in this regard and discuss them would help justify how widely applicable the proposed approach is.\n- Why the continuity detection on action-value function could detect the irreversible states? It seems that there is not much discussion on this point in Section 3.3. It is crucial to clarify it.\n- VirtualPlay should also be included as a baseline in the experiment section in order to show which approach best exploits cycle-consistency.\n", "summary_of_the_paper": "This paper considers state representation learning problem in deep RL. It exploits the cycle-consistency supervision and develops a “learning via retracing” approach. Such supervision signals are generally inherent in existing data and does not need additional interaction with environment, which leads to better sample efficiency. Learning from predictive supervision from temporally forward and backward directions reveals information from both the future and past to the target state, leading to more accurate latent state inference. In particular, the paper proposes the Cycle-Consistency World Model (CCWM) along with practical considerations (e.g., adaptive truncation to remove irreversible states), under the model-based framework based on generative dynamics modeling (CCWM).", "summary_of_the_review": "The paper proposes CCWM, a learning via retracing method, based on cycle-consistency constraint. It leads to better sample efficiency and final performance by learning better state representations. Need more clarification about its novel contribution relative to a recent work (VirtualPlay), which considers cycle-consistency learning in model-free RL setting.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "KNWF1j1nc0P", "reviewer": "Reviewer_NEVM", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- The motivation for the cycle consistency is not really discussed. Why would enforcing cycle consistency improve representation learning?\n- In equation 7: can you explain the choice of the KL distance between two rewards? How is the 2-wassertstein distance estimated from data?\n- The \"model-free instantiation of learning via retracing\" (appendix C) is not clear. In particular, the paper states that there is a graphical illustration \"shown in Figure\" but there is no reference to any figure.\n- In Figure 3, can you clarify why some of the baselines are straight lines? Can you provide more details than this explanation: '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.\"\n- Section 3.3: The prediction of the irreversibility of a transition is not very clear. How is a \"sudden change in the value function\" a good predictor of the irreversibility of the transition? \n- In the conclusion, it is written \"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 sample efﬁciency and the asymptotic performance.\" However, CCWM is only compared to dreamer and yields to comparable results in most cases? (the other baselines are only given as \"asymptotic scores\" and they have a direct access to the state?)\n\nMinor comments:\n- The paper formalizes an observation space, O that is \"high-dimensional, due to either redundant information or simply because only visual inputs are available\". Can you clarify whether that is the POMDP setting? If it is not the POMDP setting, it might be worth clarifying as this is the usual denomination used in POMDP.\n- The paper is based on the open source implementation of Hafner et al. Is there an open source implementation of the modifications presented in this paper?", "summary_of_the_paper": "The paper investigates a self-supervised approach for learning the state representation in RL tasks. In addition to the predictive (reconstruction) supervision in the forward direction, the authors include a “retraced” transitions for representation/model learning, by enforcing the cycle-consistency constraint between the original and retraced states. The authors claim that this facilitates stronger representation learning and improve upon the sample efﬁciency of learning. As it is not always possible to find such a cycle consistency between two states and for counteracting the negative impacts brought by the “irreversible” transitions, the authors add a novel adaptive “truncation” mechanism.", "summary_of_the_review": "The motivation for the contribution is not very clear and there are a few unclear elements in the formalization (e.g. sudden change in the value function for the prediction of the irreversibility) .", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes augmenting standard forward prediction techniques used for representation learning with backward prediction as well, termed "learning via retracing". The paper implements this idea in a Cycle-Consistency World Model (CCWM) and demonstrates that CCWM improves performance of a Dreamer agent across a number of Control Suite tasks. The paper also proposes a way to detect "irreversible" transitions and exclude them from the backwards prediction step.

This paper generated mixed opinions, and the reviewers did not come to a consensus on whether it should be accepted or rejected. In particular, Reviewer VSAG maintained it should be accepted, while Reviewer NEVM maintained it should be rejected. The other reviewers did not reply; I thought the authors' responses to their questions were reasonable so I assume their concerns were addressed (additionally, I don't believe a comparison to PlayVirtual is a justifiable request per the [reviewing guidelines](https://iclr.cc/Conferences/2022/ReviewerGuide), as it is concurrent work).

The reviewers generally agreed that the cycle-consistency idea proposed by the paper is interesting, well-motivated, and borne out by the experimental results. I agree with these points. The main weakness of the paper, brought up by multiple reviewers, was the justification/motivation for the method for irreversibility detection. The authors clarified in the rebuttal that the motivation comes from the idea that temporally-adjacent states with very different values will tend to be far apart in representation space. While I believe that is true, it's not at all clear to me that this necessarily *entails* irreversible transitions between those states. That said, my feeling is that this is approach is (1) not the main contribution of the paper and (2) empirically seems to work, based on the experiments, even if the motivation is unclear/unjustified. Therefore, I do not think the concern about irreversibility detection is grounds on its own for rejection.

Overall, I find this is a sensible approach to better representation learning in MBRL. I recommend acceptance as a poster.
Decision: Accept (Poster)