Original: L EARNING S TATE R EPRESENTATIONS VIA T EMPORAL C YCLE -C ONSISTENCY C ONSTRAINT IN MBRL
Revised: L EARNING S TATE R EPRESENTATIONS VIA T EMPORAL C YCLE -C ONSISTENCY C ONSTRAINT IN MBRL

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Representation learning is a popular approach for reinforcement learning (RL) tasks with partially observable Markov decision processes.
Revised: Representation learning is a popular approach for reinforcement learning (RL) tasks with partially observable Markov decision processes.

Original: Existing works on learning representations utilise the dynamics models in model-based RL to perform training through model predictive reconstruction in a temporally forward fashion.
Revised: Existing works on learning representations utilise the dynamics models in model-based RL to perform training through model predictive reconstruction in a temporally forward fashion.

Original: However, temporally backward state predictions also yield useful supervision signals as they convey information about the future states given the action choices.
Revised: However, temporally backward state predictions also yield useful supervision signals as they convey information about the future states given the action choices.

Original: We argue that combining them with forward passes will facilitate stronger representation learning and improve the sample efﬁciency of RL.
Revised: We argue that combining them with forward passes will facilitate stronger representation learning and improve the sample efﬁciency of RL.

Original: Here we propose a general framework for learning state representations for RL tasks, utilising both forward and backward passes by imposing temporal cycle-consistency constraints, which can be integrated with any model-based RL algorithms leveraging a latent dynamics model.
Revised: Here we propose a general framework for learning state representations for RL tasks, utilising both forward and backward passes by imposing temporal cycle-consistency constraints, which can be integrated with any model-based RL algorithms leveraging a latent dynamics model.

Original: We show improved empirical performance in terms of sample-efﬁciency and convergence score over several baselines on continuous control benchmarks.
Revised: We show improved empirical performance in terms of sample-efﬁciency and convergence score over several baselines on continuous control benchmarks

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: yields detection of continuous control tasks helps with identifying catastrophic states and behaviour (graphical demonstration from the DeepMind Control

Original: [Figure]
Revised: Suite (Tassa et al., 2018)).

Original: Reinforcement learning (RL) in Partially Observable Markov
Revised: Reinforcement learning (RL) in Partially Observable Markov

Original: Decision Processes (POMDPs) is usually a non-trivial task due to the curse of dimensionality and curse of history (Kaelbling et al., 1998).
Revised: Decision Processes (POMDPs) is usually a non-trivial task due to the curse of dimensionality and curse of history (Kaelbling et al., 1998).

Original: Representation learning is a commonly adopted approach for extracting abstract features that compactly characterise the high-dimensional raw inputs, such that RL can be performed more efﬁciently.
Revised: Representation learning is a commonly adopted approach for extracting abstract features that compactly characterise the high-dimensional raw inputs, such that RL can be performed more efﬁciently.

Original: Both model-free and modebased RL methods have been studied extensively for solving the POMDPs.
Revised: Both model-free and modebased RL methods have been studied extensively for solving the POMDPs.

Original: Here we focus on learning representations in model-based RL.
Revised: Here we focus on learning representations in model-based RL.

Original: In model-based RL, an environment model is typically approximated and combined with planning methods (Chua et al., 2018) or learning methods (Hafner et al., 2020) in the latent space to facilitate behaviour learning.
Revised: In model-based RL, an environment model is typically approximated and combined with planning methods (Chua et al., 2018) or learning methods (Hafner et al., 2020) in the latent space to facilitate behaviour learning.

Original: In their seminal work, Ha & Schmidhuber (2018) propose to learn world models given trained encoders.
Revised: In their seminal work, Ha & Schmidhuber (2018) propose to learn world models given trained encoders.

Original: The world model can be used either for learning behaviours or for real-time planning.
Revised: The world model can be used either for learning behaviours or for real-time planning.

Original: More recently, a number of works propose joint training of the embedding function and the latent dynamics model (Srinivas et al., 2018;
Revised: More recently, a number of works propose joint training of the embedding function and the latent dynamics model (Srinivas et al., 2018;

Original: Hafner et al., 2019; Schrittwieser et al., 2020b; Lee et al., 2020a).
Revised: Hafner et al., 2019; Schrittwieser et al., 2020b; Lee et al., 2020a).

Original: Such model learning is performed by minimising the deviation of the imaginary predictive states from the real observations.
Revised: Such model learning is performed by minimising the deviation of the imaginary predictive states from the real observations.

Original: However, we argue that the aforementioned model learning is inefﬁcient in terms of the sample usage.
Revised: However, we argue that the aforementioned model learning is inefﬁcient in terms of the sample usage.

Original: More meaningful supervision can be obtained by employing the cycle-consistency constraint between the imaginary state predictions in forward and backward temporal directions.
Revised: More meaningful supervision can be obtained by employing the cycle-consistency constraint between the imaginary state predictions in forward and backward temporal directions.

Original: The key intuition is that if the model can translate the agent from state s 0 to s T by following a trajectory, and the model also has the capability to bring the agent back from s T to s (cid:48) 0 along a similar trajectory, where s 0 and s (cid:48)  are expected to be similar.
Revised: The key intuition is that if the model can translate the agent from state s 0 to s T by following a trajectory, and the model also has the capability to bring the agent back from s T to s (cid:48) 0 along a similar trajectory, where s 0 and s (cid:48) 0 are expected to be similar.

Original: Based on this intuition, backward imaginary predictions can be utilised to facilitate model learning in a self-supervised manner.
Revised: Based on this intuition, backward imaginary predictions can be utilised to facilitate model learning in a self-supervised manner.

Original: As a motivating example, consider a rat trying to navigate towards a goal location represented by a cheese in a cluttered room shown in Figure 1a.
Revised: As a motivating example, consider a rat trying to navigate towards a goal location represented by a cheese in a cluttered room shown in Figure 1a.

Original: By imposing the temporal cycle-consistency constraint, multiple backwards imaginative paths can be simulated.
Revised: By imposing the temporal cycle-consistency constraint, multiple backwards imaginative paths can be simulated.

Original: In order to retrace to the starting position, the rat needs to pass through two “doors” in all paths.
Revised: In order to retrace to the starting position, the rat needs to pass through two “doors” in all paths.

Original: This allows the rat to identify key states (the doorways) that are essential for successful navigation towards the goal.
Revised: This allows the rat to identify key states (the doorways) that are essential for successful navigation towards the goal.

Original: Moreover, backward rollouts can pass information about the future states back to previous states.
Revised: Moreover, backward rollouts can pass information about the future states back to previous states.

Original: This approximately resembles the smoothing operation in state-space models (Kalman, 1960; Sahani, 1999), which potentially supports more accurate inference over latent states.
Revised: This approximately resembles the smoothing operation in state-space models (Kalman, 1960; Sahani, 1999), which potentially supports more accurate inference over latent states.

Original: We also expect that learning representation by retracing will improve the agent’s ability of avoiding catastrophic actions in certain situations.
Revised: We also expect that learning representation by retracing will improve the agent’s ability of avoiding catastrophic actions in certain situations.

Original: Consider a toy example of continuous control in Figure 1b, where no action will take the falling humanoid (right) back to its previous state (left), yielding large discrepancy between the true and retraced states.
Revised: Consider a toy example of continuous control in Figure 1b, where no action will take the falling humanoid (right) back to its previous state (left), yielding large discrepancy between the true and retraced states.

Original: The error will then be propagated back to inform the agent to avoid the behaviours that lead to irrecoverable states.
Revised: The error will then be propagated back to inform the agent to avoid the behaviours that lead to irrecoverable states.

Original: Joining forward and backward inferences in the latent space is also supported by evidence from the neuroscience community.
Revised: Joining forward and backward inferences in the latent space is also supported by evidence from the neuroscience community.

Original: Hippocampal “replay” has been suggested to play an essential role in decision-making and model-based inferences (Mattar & Daw, 2018).
Revised: Hippocampal “replay” has been suggested to play an essential role in decision-making and model-based inferences (Mattar & Daw, 2018).

Original: Connections between the latent state inference and forward and reversed hippocampal “replay” have been proposed and substantiated by experimental evidence (Penny et al., 2013).
Revised: Connections between the latent state inference and forward and reversed hippocampal “replay” have been proposed and substantiated by experimental evidence (Penny et al., 2013).

Original: Following this motivation, we thus expect the retrace operation to play an essential role in learning representations for POMDP tasks.
Revised: Following this motivation, we thus expect the retrace operation to play an essential role in learning representations for POMDP tasks.

Original: The main contributions are summarised as follows.
Revised: The main contributions are summarised as follows.

Original: • We introduce a self-supervised representation learning algorithm via imposing the cycleconsistency constraint to facilitate model learning.
Revised: • We introduce a self-supervised representation learning algorithm via imposing the cycleconsistency constraint to facilitate model learning.

Original: To the best of our knowledge, this is the ﬁrst work that introduce cycle-consistency to learning a world model for RL.
Revised: To the best of our knowledge, this is the ﬁrst work that introduce cycle-consistency to learning a world model for RL.

Original: • We instantiate our method utilising the Recurrent State-Space Model (RSSM (Hafner et al., 2019))
Revised: • We instantiate our method utilising the Recurrent State-Space Model (RSSM (Hafner et al., 2019))

Original: as the world model, resulting in a new model-based RL agent termed as CycleConsistent World Model (CCWM) .
Revised: as the world model, resulting in a new model-based RL agent termed as CycleConsistent World Model (CCWM) .

Original: • We empirically show that CCWM outperforms the baseline algorithms, and is able to reach the asymptotic performance of state-based model-free methods on a number of challenging continuous control benchmarks.
Revised: • We empirically show that CCWM outperforms the baseline algorithms, and is able to reach the asymptotic performance of state-based model-free methods on a number of challenging continuous control benchmarks.

Original:  R ELATED W ORKS
Revised:  R ELATED W ORKS

Original: Representation learning in RL.
Revised: Representation learning in RL.

Original: Finding useful state representations that could aid reinforcement learning has long been studied, from the introduction of manually constructed features such as tile coding and Fourier basis in combination (Sutton & Barto, 2018), to the learned embedding vectors using trained neural networks (Mnih et al., 2013; Schrittwieser et al., 2020a).
Revised: Finding useful state representations that could aid reinforcement learning has long been studied, from the introduction of manually constructed features such as tile coding and Fourier basis in combination (Sutton & Barto, 2018), to the learned embedding vectors using trained neural networks (Mnih et al., 2013; Schrittwieser et al., 2020a).

Original: Even under partial observability, we can train an RL agent in the embedding space using a neural network encoder without the requirements of an explicit model of the POMDP (Kaelbling et al., 1998).
Revised: Even under partial observability, we can train an RL agent in the embedding space using a neural network encoder without the requirements of an explicit model of the POMDP (Kaelbling et al., 1998).

Original: However, it has been noted that naive joint training of RL with representation learning can be inefﬁcient due to the sparse and delayed reward structure of the MDP (Shelhamer et al., 2017).
Revised: However, it has been noted that naive joint training of RL with representation learning can be inefﬁcient due to the sparse and delayed reward structure of the MDP (Shelhamer et al., 2017).

Original: Many auxiliary tasks, e.g. weakly-supervised classiﬁcation and location recognition, have been proposed for resolving such issue (Lee et al., 2020b; Mirowski et al., 2017; Oord et al., 2018).
Revised: Many auxiliary tasks, e.g. weakly-supervised classiﬁcation and location recognition, have been proposed for resolving such issue (Lee et al., 2020b; Mirowski et al., 2017; Oord et al., 2018).

Original: In the current work, we propose an additional cycle-consistency constrained objective for learning the representations.
Revised: In the current work, we propose an additional cycle-consistency constrained objective for learning the representations.

Original: There exists two commonly adopted sources of supervisions for learning representations in RL tasks, the reconstruction of observations (Gelada et al., 2019; Hafner et al., 2020; Lee et al., 2020a), and self-supervised signals such as contrastive losses (Oord et al., 2018).
Revised: There exists two commonly adopted sources of supervisions for learning representations in RL tasks, the reconstruction of observations (Gelada et al., 2019; Hafner et al., 2020; Lee et al., 2020a), and self-supervised signals such as contrastive losses (Oord et al., 2018).

Original: Recently, Zhang et al.
Revised: Recently, Zhang et al.

Original: (2020) and Agarwal et al.
Revised: (2020) and Agarwal et al.

Original: (2021) proposed to learn the embeddings using bisimulation metrics (Ferns et al., 2011), which yields representations that are agnostic to task-irrelevant information, supporting faster learning and stronger generalisation.
Revised: (2021) proposed to learn the embeddings using bisimulation metrics (Ferns et al., 2011), which yields representations that are agnostic to task-irrelevant information, supporting faster learning and stronger generalisation

Original: Here we adopt a combination of the reconstruction and bisimulation metrics supervision signals for model training.
Revised: Here we adopt a combination of the reconstruction and bisimulation metrics supervision signals for model training.

Original: Learning the dynamics model.
Revised: Learning the dynamics model.

Original: In addition to learning state abstractions based solely on the observations, it is also possible to learn the latent representation along with the dynamics model using latent state-space models (LSSM (Murphy, 2012)).
Revised: In addition to learning state abstractions based solely on the observations, it is also possible to learn the latent representation along with the dynamics model using latent state-space models (LSSM (Murphy, 2012)).

Original: Ha & Schmidhuber (2018) introduced the world model, for modelling the dynamics of the latent space based on the pre-trained state embeddings.
Revised: Ha & Schmidhuber (2018) introduced the world model, for modelling the dynamics of the latent space based on the pre-trained state embeddings.

Original: The world model leads to improved performance in POMDP tasks such as Atari games, as well as generates reasonable predictions in the observable space through rollouts in the latent space.
Revised: The world model leads to improved performance in POMDP tasks such as Atari games, as well as generates reasonable predictions in the observable space through rollouts in the latent space.

Original: The state representation can also be learnt jointly with the latent dynamics model.
Revised: The state representation can also be learnt jointly with the latent dynamics model.

Original: A number of works proposed to treat the embedding function as part of the recognition model in a probabilistic graphical model (Buesing et al., 2018; Lee et al., 2020a; Hafner et al., 2020).
Revised: A number of works proposed to treat the embedding function as part of the recognition model in a probabilistic graphical model (Buesing et al., 2018; Lee et al., 2020a; Hafner et al., 2020).

Original: These methods treat the learning of representations and latent dynamics model as a joint task, which can be trained via doing predictive inference over the latent states given the observations and actions.
Revised: These methods treat the learning of representations and latent dynamics model as a joint task, which can be trained via doing predictive inference over the latent states given the observations and actions.

Original: The learned embeddings can then be used as the inputs to the value function approximator in the downstream RL component for learning the optimal policy.
Revised: The learned embeddings can then be used as the inputs to the value function approximator in the downstream RL component for learning the optimal policy.

Original: The learned latent dynamics model can be used for generating imaginary trajectories through latent model-based rollouts for RL training, which improves the sample-efﬁciency of learning.
Revised: The learned latent dynamics model can be used for generating imaginary trajectories through latent model-based rollouts for RL training, which improves the sample-efﬁciency of learning.

Original: Our proposed approach is similar to this line of work since we also leverage latent stochastic sequential modes for joint learning of the representations and the latent dynamics model, and our method can be efﬁciently combined with downstream RL algorithm to learn the optimal policy using model-based rollouts following a similar paradigm.
Revised: Our proposed approach is similar to this line of work since we also leverage latent stochastic sequential modes for joint learning of the representations and the latent dynamics model, and our method can be efﬁciently combined with downstream RL algorithm to learn the optimal policy using model-based rollouts following a similar paradigm.

Original: However, we note a major difference between our method and the existing approaches: all related works use predictive supervision in a temporally forward fashion, whereas our method combines predictive supervision in both the temporally forward and backward directions by imposing a temporal cycle-consistency constraint.
Revised: However, we note a major difference between our method and the existing approaches: all related works use predictive supervision in a temporally forward fashion, whereas our method combines predictive supervision in both the temporally forward and backward directions by imposing a temporal cycle-consistency constraint.

Original: Cycle-Consistency.
Revised: Cycle-Consistency.

Original: Cycle-consistency is a commonly adopted approach in computer vision and natural language processing (Zhou et al., 2016; Zhu et al., 2017;
Revised: Cycle-consistency is a commonly adopted approach in computer vision and natural language processing (Zhou et al., 2016; Zhu et al., 2017;

Original: Yang et al., 2017; Dwibedi et al., 2019), where the fundamental idea is the validation of matches between cycling through multiple samples.
Revised: Yang et al., 2017; Dwibedi et al., 2019), where the fundamental idea is the validation of matches between cycling through multiple samples.

Original: Here we adopt similar design principles for sequential decision-making tasks: rollouts in a temporally forward direction alone yield under-constrained learning of the dynamics model of the task.
Revised: Here we adopt similar design principles for sequential decision-making tasks: rollouts in a temporally forward direction alone yield under-constrained learning of the dynamics model of the task.

Original: We propose to additionally incorporate backwards rollouts to perform model learning in a self-supervised fashion.
Revised: We propose to additionally incorporate backwards rollouts to perform model learning in a self-supervised fashion.

Original: We do this by imposing the temporal cycle-consistency constraint (Dwibedi et al., 2019): we introduce an auxiliary objective based on the bisimulation metric that ensures the forward and backward temporal dynamics yield the same MDP structure.
Revised: We do this by imposing the temporal cycle-consistency constraint (Dwibedi et al., 2019): we introduce an auxiliary objective based on the bisimulation metric that ensures the forward and backward temporal dynamics yield the same MDP structure.

Original:  L EARNING R EPRESENTATION VIA R ETRACING
Revised:  L EARNING R EPRESENTATION VIA R ETRACING

Original: We consider learning in a POMDP, which is represented by the tuple (cid:104)S , O , A , r, T, E, γ (cid:105) , where
Revised: We consider learning in a POMDP, which is represented by the tuple (cid:104)S , O , A , r, T, E, γ (cid:105) , where

Original: S , O , A are the state space, the observation space and the action space, respectively.
Revised: S , O , A are the state space, the observation space and the action space, respectively.

Original: r : S ×
Revised: r : S ×

Original: A → R is the reward function, T : S ×
Revised: A → R is the reward function, T : S ×

Original: A × S → [0 , 1] is the transition distribution between states,
Revised: A × S → [0 , 1] is the transition distribution between states,

Original: E : S ×A → [0 , 1] is the conditional observation distribution, and γ ∈ (0 , 1] is the discounting factor.
Revised: E : S ×A → [0 , 1] is the conditional observation distribution, and γ ∈ (0 , 1] is the discounting factor.

Original: At each time step t , the agent receives an observation O t ∈ O , and chooses its action according to its policy a t ∼ π ( a | O t ) to interact with the environment.
Revised: At each time step t , the agent receives an observation O t ∈ O , and chooses its action according to its policy a t ∼ π ( a | O t ) to interact with the environment.

Original: The goal is to learn the optimal policy that maximises the expected return over all states deﬁned as E π [ (cid:80) ∞ τ = t
Revised: The goal is to learn the optimal policy that maximises the expected return over all states deﬁned as E π [ (cid:80) ∞ τ = t

Original: γ τ − t r τ | s ] .
Revised: γ τ − t r τ | s ] .

Original: In the current work, instead of working directly with O , we wish to learn an embedding function q ψ : O → Z , along with a latent dynamics model, q ψ : Z × A → Z , such that model-based RL can be performed efﬁciently in the learned latent space.
Revised: In the current work, instead of working directly with O , we wish to learn an embedding function q ψ : O → Z , along with a latent dynamics model, q ψ : Z × A → Z , such that model-based RL can be performed efﬁciently in the learned latent space.

Original: 3. L EARNING R EPRESENTATIONS VIA F ORWARD P ASSES
Revised: 3. L EARNING R EPRESENTATIONS VIA F ORWARD P ASSES

Original: Existing works on learning the state representation of a latent dynamics model rely on proposing a stochastic sequential model as the world model, which is trained by performing latent forward rollouts such that the future environmental attributes such as observations or rewards are accurately predicted (Lee et al., 2020a; Hafner et al., 2020). 
Revised: Existing works on learning the state representation of a latent dynamics model rely on proposing a stochastic sequential model as the world model, which is trained by performing latent forward rollouts such that the future environmental attributes such as observations or rewards are accurately predicted (Lee et al., 2020a; Hafner et al., 2020). Speciﬁcally, we are interested in model-based RL algorithms with a dynamics model deﬁned in terms of the following components: embedding [Equation] variational

Original: Speciﬁcally, we are interested in model-based RL algorithms with a dynamics model deﬁned in terms of the following components: embedding function : q ψ ( O t ) , posterior distribution : p ( z t +1 | z t , a t , O t +1 ) , variational posterior distribution : q ψ ( z t +1 | z t , a t , O t +1 ) , variational transition distribution : q ψ ( z t +1 | z t , a t ) , generative distribution : p θ ( O t +1 | z t +1 ) , (1) where ψ and θ represent the parameters associated with the recognition and generative models, respectively.
Revised: [Equation] variational transition [Equation] where ψ and θ represent the parameters associated with the recognition and generative models, respectively.

Original: Note that we use a variational approximation to the posterior distribution since the true posterior p ( z t +1 | z t , a t , O t +1 ) is usually intractable in practice.
Revised: Note that we use a variational approximation to the posterior distribution since the true posterior p ( z t +1 | z t , a t , O t +1 ) is usually intractable in practice.

Original: · · · ˆ z t ˆ z t + ˆ z t + · · · a t a t + a t +
Revised: [Equation] t [Equation] ˆ z t + · · · a t a t + a t +

Original: O t
Revised: O t

Original: O t +
Revised: O t +

Original: O t + ˜ O t ˜ O t + ˜ O t + ˜
Revised: O t + ˜ O t ˜ O t + ˜ O t + ˜

Original: z
Revised: z

Original: t ˜
Revised: t ˜

Original: z
Revised: z

Original: t + ˜
Revised: t + ˜

Original: z
Revised: z

Original: t + ρ ¯ a t + ρ ¯ a t +[-¯ z t +
-]
Revised: t + ρ ¯ a t + ρ [Equation] a t +[-¯ z t +
-]

Original: ¯ z t +1 ¯ z t ¯ z t +2 = ˆ z t +2 L
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: forwardpass backward pass 
Revised: forwardpass backward pass of the model learning process in CCWM .

Original: The forward pass (bounded by the orange rectangle) and retrace phase for learning the representations and latent dynamics model via latent imaginations are shown.
Revised: The forward pass (bounded by the orange rectangle) and retrace phase for learning the representations and latent dynamics model via latent imaginations are shown.

Original: Please refer to Section 3.2 for detailed description.
Revised: Please refer to Section 3.2 for detailed description.

Original: At each time step t , the agent receives an observation O t , which is then embedded into a latent feature vector z t = q ψ ( O t ) .
Revised: At each time step t , the agent receives an observation O t , which is then embedded into a latent feature vector z t = q ψ ( O t ) .

Original: The feature vector is then rolled out in a forward fashion given the action a t , yielding one-step prediction into the future as ˆ z t +1 ∼ q ψ ( z | z ≤ t , a ≤ t ) following the variational transition distribution q ψ ( z (cid:48) | z, a ) .
Revised: The feature vector is then rolled out in a forward fashion given the action a t , yielding one-step prediction into the future as ˆ z t +1 ∼ q ψ ( z | z ≤ t , a ≤ t ) following the variational transition distribution q ψ ( z (cid:48) | z, a ) .

Original: The parameters of the recognition and generative models shown in Eq.
Revised: The parameters of the recognition and generative models shown in Eq.

Original: 1 are jointly learned by maximising the variational lower bound (ELBO, (Wainwright & Jordan, 2008)), which takes the general form as following:
Revised: 1 are jointly learned by maximising the variational lower bound (ELBO, (Wainwright & Jordan, 2008)), which takes the general form as following:

Original: [Equation]
Revised: [Equation]

Original: where the expectation is taken with respect to the variational posterior distribution, q ψ ( z t +1 | z t , a t , O t +1 ) .
Revised: where the expectation is taken with respect to the variational posterior distribution, q ψ ( z t +1 | z t , a t , O t +1 ) .

Original: Note that it is also possible to generate predictions over multiple steps into the future instead of one for model learning, see, e.g. Lee et al.
Revised: Note that it is also possible to generate predictions over multiple steps into the future instead of one for model learning, see, e.g. Lee et al.

Original: (2020a).
Revised: (2020a).

Original: 3. L EARNING R EPRESENTATIONS VIA R ETRACING
Revised: 3. L EARNING R EPRESENTATIONS VIA R ETRACING

Original: The backward rollouts convey information of future states given the actions taken back to the previous states, similar to the smoothing operation in state-space models (Kalman, 1960).
Revised: The backward rollouts convey information of future states given the actions taken back to the previous states, similar to the smoothing operation in state-space models (Kalman, 1960).

Original: This potentially supports more accurate inference of the latent states and stronger representation learning.
Revised: This potentially supports more accurate inference of the latent states and stronger representation learning.

Original: Moreover, the backward rollouts provide additional supervision for training the dynamics model, which are expected to improve the sample efﬁciency of learning.
Revised: Moreover, the backward rollouts provide additional supervision for training the dynamics model, which are expected to improve the sample efﬁciency of learning.

Original: Here we introduce a general framework for learning representations for the latent dynamics model which combines both forward and backward passes through latent imagination while constraining the temporal cycle-consistency between the forward and backward passes.
Revised: Here we introduce a general framework for learning representations for the latent dynamics model which combines both forward and backward passes through latent imagination while constraining the temporal cycle-consistency between the forward and backward passes.

Original: We refer to the proposed approach as Cycle-Consistent World Model ( CCWM ).
Revised: We refer to the proposed approach as Cycle-Consistent World Model ( CCWM ).

Original: The pseudocode for
Revised: The pseudocode for

Original: CCWM is shown in Algorithm 1.
Revised: CCWM is shown in Algorithm 1.

Original: Consider the same dynamics model described in Section 3.1, we additionally deﬁne a reverse action approximator, ρ : Z × Z → A , which takes in a tuple of latent states ( z 1 , z 2 ) and outputs an action ¯ a that approximates the action that leads the transition from z 2 to z 1 .
Revised: Consider the same dynamics model described in Section 3.1, we additionally deﬁne a reverse action approximator, ρ : Z × Z → A , which takes in a tuple of latent states ( z 1 , z 2 ) and outputs an action ¯ a that approximates the action that leads the transition from z 2 to z 1 .

Original: We assume the parameters of ρ are learned jointly with the model learning in an end-to-end fashion.
Revised: We assume the parameters of ρ are learned jointly with the model learning in an end-to-end fashion.

Original: Suppose the imagination horizon for predictive inference is k .
Revised: Suppose the imagination horizon for predictive inference is k .

Original: For l = 1 , . . .
Revised: For l = 1 , . . .

Original: , L − k , the prior estimates for the k -step predictions and posterior estimates given the ground-truth observations are deﬁned similarly as in the one-step case (corresponding to line 7 of Algorithm 1 and the forward dynamics chain in Figure 2).
Revised: , L − k , the prior estimates for the k -step predictions and posterior estimates given the ground-truth observations are deﬁned similarly as in the one-step case (corresponding to line 7 of Algorithm 1 and the forward dynamics chain in Figure 2).

Original: [Equation]
Revised: [Equation]

Original: where the initial condition is taken to be ˆ z l = z l .
Revised: where the initial condition is taken to be ˆ z l = z l .

Original: The retracing steps are similar to the one-step case, which can be computed via backwards imagination using the same latent dynamics model given the predicted prior estimates from the forward passes.¯
Revised: The retracing steps are similar to the one-step case, which can be computed via backwards imagination using the same latent dynamics model given the predicted prior estimates from the forward passes.

Original: [Equation]
Revised: [Equation]

Original: ) for l = 1 , . . 
Revised: for l = 1 , . . .

Original: , L − k , where the initial retraced state is taken to be ¯ z l + k =
Revised: , L − k , where the initial retraced state is taken to be ¯ z l + k =

Original: ˆ z l + k .
Revised: ˆ z l + k .

Original: The retracing computation corresponds to Line 8 in Algorithm 1 and is demonstrated graphically in the reversed dynamics chain in Figure 2. 
Revised: The retracing computation corresponds to Line 8 in Algorithm 1 and is demonstrated graphically in the reversed dynamics chain in Figure 2. [Equation]

Original: 1: Given: policy ˆ π ( a | s ) , encoder q ψ ( O ) and decoder q θ ( O | z ) (Eq. 1), reverse action approximator ρ ( z, z (cid:48) ) , imagination horizon k , running average loss L avg = 0 , counter c = 0 .
Revised: 1: Given: policy ˆ π ( a | s ) , encoder q ψ ( O ) and decoder q θ ( O | z ) (Eq. 1), reverse action approximator ρ ( z, z (cid:48) ) , imagination horizon k , running average loss L avg = 0 , counter c = 0 .

Original: 2: Input: Batchof N sampled trajectories of length L , { O i,t i : t i + L , a i,t i : t i + L } N i =1 . 3: for n = 1 to N do 4:
Revised: Batch N sampled trajectories of length L , { O i,t i : t i + L , a i,t i : t i + L } Ni =1 . 3: for n = 1 to N do 4:

Original: O = { O n,t n + j } kj =1 , a = { a n,t n + j } k − 1 j = 5:
Revised: O = { O n,t n + j } kj =1 , a = { a n,t n + j } k − 1 j = 5:

Original: Compute embedding vectors z = { q ψ ( O j ) } kj =

Revised: Compute embedding vectors z = { q ψ ( O j ) } kj =

Original: (which we only use the ﬁrst L − k for learning representation given k -step imaginations).
Revised: (which we only use the ﬁrst L − k for learning representation given k -step imaginations).

Original: 6: for j = 1 to L − k do 7:
Revised: 6: for j = 1 to L − k do 7:

Original: Compute prior ˆ z j +1: j + k and posterior es- timates ˜
Revised: Compute prior ˆ z j +1: j + k and posterior es- timates ˜

Original: z j +1: j + k using Eq.
Revised: z j +1: j + k using Eq.

Original: Compute retraced states ¯ z j : j + k using
Revised: Compute retraced states ¯ z j : j + k using

Original: Eq.9:
Revised: Eq.9:

Original: Compute the latent model loss L j using
Revised: Compute the latent model loss L j using

Original: Eq. 6.
Revised: Eq. 6.

Original: Update L avg ← cc +1 L avg + 1 c +1 L j . 11:
Revised: Update L avg ← cc +1 L avg + 1 c +1 L j . 11:

Original: Increment counter, c ← c + 1 .
Revised: Increment counter, c ← c + 1 .

Original: 12: end for 13: end for 14: Compute the gradient, ∇ Θ L avg .
Revised: 12: end for 13: end for 14: Compute the gradient, ∇ Θ L avg .

Original: 15: Update the model parameters, Θ ← Θ + α ∇ Θ L avg .
Revised: 15: Update the model parameters, Θ ← Θ + α ∇ Θ L avg .

Original: Model learning supervised by cycleconsistency.
Revised: Model learning supervised by cycleconsistency.

Original: We learn the representation with the auxiliary task for satisfying the temporal cycle-consistency of the dynamics models, i.e., we wish to impose a constraint on the similarity between the original (forwardpredicted) states and the retraced states in addition to the ELBO loss (Eq. 2).
Revised: We learn the representation with the auxiliary task for satisfying the temporal cycle-consistency of the dynamics models, i.e., we wish to impose a constraint on the similarity between the original (forwardpredicted) states and the retraced states in addition to the ELBO loss (Eq. 2).

Original: We utilise the bisimulation metrics as the loss function for the auxiliary retrace consistency task (Zhang et al., 2020).
Revised: We utilise the bisimulation metrics as the loss function for the auxiliary retrace consistency task (Zhang et al., 2020).

Original: [Equation]
Revised: [Equation]

Original: where ˆ R ( r | z ) represents the approximated reward model; ˆ P ( z (cid:48) | z, ˆ π ( z )) denotes the distribution over the predicted state z (cid:48) , parameterised by the latent transition dynamics model, given the current state z and a sampled action (chosen given the current action distribution, ˆ π ( a | z ) ); and W 2 ( · , · ) represents the 2 -Wasserstein distance between probability distributions, which has an analytical expression for Gaussian distributions (Givens & Shortt, 1984).
Revised: where ˆ R ( r | z ) represents the approximated reward model; ˆ P ( z (cid:48) | z, ˆ π ( z )) denotes the distribution over the predicted state z (cid:48) , parameterised by the latent transition dynamics model, given the current state z and a sampled action (chosen given the current action distribution, ˆ π ( a | z ) ); and W 2 ( · , · ) represents the 2 -Wasserstein distance between probability distributions, which has an analytical expression for Gaussian distributions (Givens & Shortt, 1984).

Original: We choose the bisimulation metric as the loss function based on the expectation that the additional supervision on the reward and transition distribution yields the learned representation to be consistent with respect to the task MDP.
Revised: We choose the bisimulation metric as the loss function based on the expectation that the additional supervision on the reward and transition distribution yields the learned representation to be consistent with respect to the task MDP.

Original: Given the imposed temporal cycle-consistency constraint, the task MDPs embedded in the resulting latent space will be identical for both the temporally forward and backward transitions.
Revised: Given the imposed temporal cycle-consistency constraint, the task MDPs embedded in the resulting latent space will be identical for both the temporally forward and backward transitions.

Original: This is an appropriate inductive bias in many tasks such as navigation or continuous control (Figure 1).
Revised: This is an appropriate inductive bias in many tasks such as navigation or continuous control (Figure 1).

Original: Hence, for each l = 1 , . . .
Revised: Hence, for each l = 1 , . . .

Original: , k , the overall objective for training the dynamics model via retracing given imagination over k -steps take the following form (corresponding to line 9 in Algorithm 1).
Revised: , k , the overall objective for training the dynamics model via retracing given imagination over k -steps take the following form (corresponding to line 9 in Algorithm 1).

Original: [Equation]
Revised: [Equation]

Original: where λ is a scalar controlling the strength we wish to impose the temporal cycle-consistency.
Revised: where λ is a scalar controlling the strength we wish to impose the temporal cycle-consistency.

Original: The dynamics model parameters Θ = { θ, ψ } are updated via following the gradients of the variational objective in Eq.
Revised: The dynamics model parameters Θ = { θ, ψ } are updated via following the gradients of the variational objective in Eq.

Original: 6 with respect to Θ , which can be computed using reparameterisation trick followed by backpropagation (Kingma & Welling, 2014; Rezende et al., 2014).
Revised: 6 with respect to Θ , which can be computed using reparameterisation trick followed by backpropagation (Kingma & Welling, 2014; Rezende et al., 2014).

Original: Figure 2 demonstrates the complete probabilistic graphical model for CCWM .
Revised: Figure 2 demonstrates the complete probabilistic graphical model for CCWM .

Original: 3. R EINFORCEMENT L EARNING IN POMDP S
Revised: 3. R EINFORCEMENT L EARNING IN POMDP S

Original: We aim to eventually utilise the learned representation and latent dynamics model for RL in POMDPs.
Revised: We aim to eventually utilise the learned representation and latent dynamics model for RL in POMDPs.

Original: To implement this, we interlace the model learning and the policy learning as in Hafner et al.
Revised: To implement this, we interlace the model learning and the policy learning as in Hafner et al.

Original: For each iteration, we ﬁrstly update the models following the steps detailed in Algorithm 1, then use the current transition and reward models to compute the value estimates for the latent states, which can then be used in conjunction with downstream RL algorithms.
Revised: For each iteration, we ﬁrstly update the models following the steps detailed in Algorithm 1, then use the current transition and reward models to compute the value estimates for the latent states, which can then be used in conjunction with downstream RL algorithms.

Original: Essentially CCWM can be combined with any RL algorithms in a model-based setting.
Revised: Essentially CCWM can be combined with any RL algorithms in a model-based setting.

Original:  E XPERIMENTAL S TUDIES
Revised:  E XPERIMENTAL S TUDIES

Original: We evaluate CCWM on a number of challenging image-based continuous control tasks from the
Revised: We evaluate CCWM on a number of challenging image-based continuous control tasks from the

Original: DeepMind Control Suite (Tassa et al., 2018), as shown in Figure 3a.
Revised: DeepMind Control Suite (Tassa et al., 2018), as shown in Figure 3a.

Original: The experimental studies aim to empirically examine the following aspects of the proposed CCWM framework: a) whether learning representation via retracing indeed aids model-based reinforcement learning in continuous control tasks using solely image inputs; b) the interpretability of the proposed retracing operation; c) the effects of different aspects of CCWM on training.
Revised: The experimental studies aim to empirically examine the following aspects of the proposed CCWM framework: a) whether learning representation via retracing indeed aids model-based reinforcement learning in continuous control tasks using solely image inputs; b) the interpretability of the proposed retracing operation; c) the effects of different aspects of CCWM on training.

Original: The architecture and hyperparameters conﬁgurations of the neural network implementation and training details can be found in Appendix A.
Revised: The architecture and hyperparameters conﬁgurations of the neural network implementation and training details can be found in Appendix A.

Original: 4. E XPERIMENT S ETUP
Revised: E XPERIMENT S ETUP

Original: We instantiate CCWM with the RSSM (Hafner et al., 2019) as the latent dynamics model, and A3C as the downstream RL agent for behaviour learning (Mnih et al., 2016).
Revised: We instantiate CCWM with the RSSM (Hafner et al., 2019) as the latent dynamics model, and A3C as the downstream RL agent for behaviour learning (Mnih et al., 2016).

Original: We refer to the resulting model-based agent as CCWM-A3C .
Revised: We refer to the resulting model-based agent as CCWM-A3C .

Original: Despite introducing the framework of learning representations via retracing in its most general form of multi-step forward and backward rollouts, we commit to one-step case in our implementation, partially based on the empirical ﬁnding that multi-step rollouts have marginal improvements over the one-step rollouts in many model-based planning algorithms (Hamrick et al., 2020).
Revised: Despite introducing the framework of learning representations via retracing in its most general form of multi-step forward and backward rollouts, we commit to one-step case in our implementation, partially based on the empirical ﬁnding that multi-step rollouts have marginal improvements over the one-step rollouts in many model-based planning algorithms (Hamrick et al., 2020).

Original: We discuss additional motivations for sticking with the one-step case in the current implementation in Section 5.
Revised: We discuss additional motivations for sticking with the one-step case in the current implementation in Section 5.

Original: To provide a comprehensive evaluation of the performance of the proposed method, we usetasks from the DeepMind Control Suite (Tassa et al., 2018) as shown in Figure.
Revised: To provide a comprehensive evaluation of the performance of the proposed method, we usetasks from the DeepMind Control Suite (Tassa et al., 2018) as shown in Figure.

Original: 3a.
Revised: 3a.

Original: We compare
Revised: We compare

Original: CCWM-A3C with the following baselines.
Revised: CCWM-A3C with the following baselines.

Original: Dreamer (Hafner et al., 2020):
Revised: Dreamer (Hafner et al., 2020):

Original: We implement Dreamer, which integrates the model-learning and policy learning as approximately independent components for solving long-horizon decision-making problems in POMDPs based on image inputs.
Revised: We implement Dreamer, which integrates the model-learning and policy learning as approximately independent components for solving long-horizon decision-making problems in POMDPs based on image inputs.

Original: Model-Free Baselines : We additionally compare with the following model-free algorithms:
Revised: Model-Free Baselines : We additionally compare with the following model-free algorithms:

Original: SAC (Haarnoja et al., 2018), D4PG (Barth-Maron et al., 2018), A3C (Mnih et al., 2016).
Revised: SAC (Haarnoja et al., 2018), D4PG (Barth-Maron et al., 2018), A3C (Mnih et al., 2016).

Original: Namely, we implement the SAC algorithm using the state inputs and directly report the asymptotic performance of the D4PG and A3C algorithms after 10 8 training steps given in Tassa et al.
Revised: Namely, we implement the SAC algorithm using the state inputs and directly report the asymptotic performance of the D4PG and A3C algorithms after 10 8 training steps given in Tassa et al.

Original: We only report the ﬁnal performance of the model-free algorithms instead of the training curve due to the difference in training steps.
Revised: We only report the ﬁnal performance of the model-free algorithms instead of the training curve due to the difference in training steps.

Original: 4. C ONTINUOUS C ONTROL E VALUATION
Revised: 4. C ONTINUOUS C ONTROL E VALUATION

Original: The performance of CCWM-A3C is shown in Figure 3b, with the comparison with the aforementioned baseline algorithms.
Revised: The performance of CCWM-A3C is shown in Figure 3b, with the comparison with the aforementioned baseline algorithms.

Original: The proposed framework in combination with RSSM achieves similar ﬁnal convergence scores as the state-based model-free agents on 6 out of the 8 presented task environments.
Revised: The proposed framework in combination with RSSM achieves similar ﬁnal convergence scores as the state-based model-free agents on 6 out of the 8 presented task environments.

Original: Since we chose our implementation built upon Dreamer, the comparison with the baseline Dreamer is essential.
Revised: Since we chose our implementation built upon Dreamer, the comparison with the baseline Dreamer is essential.

Original: We observe that CCWM-A3C 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.
Revised: We observe that CCWM-A3C 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.

Original: The only task that CCWM-A3C is weaker than Dreamer is the hopper stand task.
Revised: The only task that CCWM-A3C is weaker than Dreamer is the hopper stand task.

Original: We brieﬂy discuss our intuition and proposed solution for this failure in Section 4.4.
Revised: We brieﬂy discuss our intuition and proposed solution for this failure in Section 4.4.

Original:  The empirical results on the continuous control tasks conform with our hypothesis that utilising backward passes in addition to forward passes provides additional supervision, which improves the sample efﬁciency in RL task.
Revised: [Figure] The empirical results on the continuous control tasks conform with our hypothesis that utilising backward passes in addition to forward passes provides additional supervision, which improves the sample efﬁciency in RL task.

Original: Moreover, imposing the temporal cycle-consistency constraint yields signiﬁcant performance increases on the ﬁnger spin, hopper hop and walker run tasks, demonstrating that retracing can help to learn stronger representation from pixel inputs for decision making.
Revised: Moreover, imposing the temporal cycle-consistency constraint yields signiﬁcant performance increases on the ﬁnger spin, hopper hop and walker run tasks, demonstrating that retracing can help to learn stronger representation from pixel inputs for decision making.

Original: 4. I NTERPRETABILITY OF R ETRACING
Revised: 4. I NTERPRETABILITY OF R ETRACING

Original: One might ask that whether the addition of the retrace stage is truly helpful for learning representations, i.e., whether the latent dynamics model and the reverse action approximator ρ can learn to generate similar predictions during retracing as during forward passes.
Revised: One might ask that whether the addition of the retrace stage is truly helpful for learning representations, i.e., whether the latent dynamics model and the reverse action approximator ρ can learn to generate similar predictions during retracing as during forward passes.

Original: We examine this by checking whether, 
Revised: We examine this by checking whether, after training, retracing yields qualitatively similar predictions over latent states as forward passes.

Original: [Figure] yields qualitatively similar predictions over demonstrate this via showing the dimensional of the  and the  states (encoded given the observations) using tSNE visualisation (Van der Maaten & Hinton, 2008) in Figure 
Revised: We demonstrate this via showing the low-dimensional embeddings of the retraced and the ground-truth latent states (encoded given the observations) using tSNE visualisation (Van der Maaten & Hinton, 2008) in Figure 4.

Original: [Figure]
Revised: [Figure]

Revised: 33 323026242220 19 18 17 16 151312 11 109  5 4 3 2retraced states        (b)

Original: In Figure 4a we show the 2-dimensional tSNE embeddings of the true and retraced states over an episode of 500 steps of the trained agent on the cheetah run task.
Revised: In Figure 4a we show the 2-dimensional tSNE embeddings of the true and retraced states over an episode of 500 steps of the trained agent on the cheetah run task.

Original: We observe a similar dominant ring structure in both embeddings, which is expected, since during running, the simulated cheetah repeats its states periodically to resume running at a high speed.
Revised: We observe a similar dominant ring structure in both embeddings, which is expected, since during running, the simulated cheetah repeats its states periodically to resume running at a high speed.

Original: Such structure is correctly captured by the retracing.
Revised: Such structure is correctly captured by the retracing.

Original: During the early stage of the episode (darker dots), the android starts to accelerate from stationarity and the states of the agent are periodic in nature but differs from the states when the agent is running at full speed.
Revised: During the early stage of the episode (darker dots), the android starts to accelerate from stationarity and the states of the agent are periodic in nature but differs from the states when the agent is running at full speed.

Original: The structures for the states in the early phase of the episode in the embedding space for the true and retraced states have high qualitative similarity, again supporting the idea that the agent has successfully learned the correct retracing actions and states upon training.
Revised: The structures for the states in the early phase of the episode in the embedding space for the true and retraced states have high qualitative similarity, again supporting the idea that the agent has successfully learned the correct retracing actions and states upon training.

Original: We take a closer look at the sampled trajectory, by examining a 76 -step sub-trajectory from it (Figure 4b, where the ﬂoating numbers indicate the timestep).
Revised: We take a closer look at the sampled trajectory, by examining a 76 -step sub-trajectory from it (Figure 4b, where the ﬂoating numbers indicate the timestep).

Original: We can see that the embedding of the true and retraced states share the same circular structure, with most states densely distributed on both sides of the ring and more loosely distributed on the paths connecting the two sides.
Revised: We can see that the embedding of the true and retraced states share the same circular structure, with most states densely distributed on both sides of the ring and more loosely distributed on the paths connecting the two sides.

Original: The dynamics in the forward temporal direction causing clock-wise transitions in both embeddings.
Revised: The dynamics in the forward temporal direction causing clock-wise transitions in both embeddings.

Original: These suggest that our approach yields interpretable retrace actions.
Revised: These suggest that our approach yields interpretable retrace actions.

Original: 4. E ARLY T RUNCATIONS
Revised: 4. E ARLY T RUNCATIONS

Original: [Figure]
Revised: [Figure]

Original: We observe from Figure 3b that our agent behaves relatively poorly on the hopper stand task.
Revised: We observe from Figure 3b that our agent behaves relatively poorly on the hopper stand task.

Original: One possible interpretation is that in the selected continuous control tasks, the simulated androids generally fall a lot, especially during the early stage of training.
Revised: One possible interpretation is that in the selected continuous control tasks, the simulated androids generally fall a lot, especially during the early stage of training.

Original: The retracing operation poses the cycle consistency constraint that the retraced states should be similar to the original states, but consider the situation shown in Figure 1b: no valid action can take the agent back to its previous state.
Revised: The retracing operation poses the cycle consistency constraint that the retraced states should be similar to the original states, but consider the situation shown in Figure 1b: no valid action can take the agent back to its previous state.

Original: This violates the assumption that the same MDP governs the forward and backward temporal dynamics hence might hinder the representation learning.
Revised: This violates the assumption that the same MDP governs the forward and backward temporal dynamics hence might hinder the representation learning.

Original: In tasks such as hopper, where irrecoverable failure states occur more often than in other tasks, the effect is more severe.
Revised: In tasks such as hopper, where irrecoverable failure states occur more often than in other tasks, the effect is more severe.

Original: We propose a potential resolution for this issue, by truncating the sampled episodes such that nonretraceable near-terminal states are left out for learning the representation via retracing.
Revised: We propose a potential resolution for this issue, by truncating the sampled episodes such that nonretraceable near-terminal states are left out for learning the representation via retracing.

Original: The effects of early truncations on the performance is shown in Figure 5.
Revised: The effects of early truncations on the performance is shown in Figure 5.

Original: We generally observe that introduction of appropriate truncation indeed have positive effects on learning, while either over- or under-truncation impedes the overall learning.
Revised: We generally observe that introduction of appropriate truncation indeed have positive effects on learning, while either over- or under-truncation impedes the overall learning.

Original: We expect that the improvement will be most signiﬁcant during the early stage of training, when the agent has not learned good behaviours and often stumbles into non-retraceable failure states.
Revised: We expect that the improvement will be most signiﬁcant during the early stage of training, when the agent has not learned good behaviours and often stumbles into non-retraceable failure states.

Original: In the right panel of Figure 5, we show training curves over the ﬁrst 4 × 10 5 steps for CCWM with different truncation lengths in the hopper stand task.
Revised: In the right panel of Figure 5, we show training curves over the ﬁrst 4 × 10 5 steps for CCWM with different truncation lengths in the hopper stand task.

Original: Comparing to the original CCWM-A3C agent (Figure 3b), truncation consistently yields faster acquisition of good policies.
Revised: Comparing to the original CCWM-A3C agent (Figure 3b), truncation consistently yields faster acquisition of good policies.

Original: With appropriate truncation strength, the resulting agent shows signiﬁcant improvement in terms of sample efﬁciency comparing to the baseline Dreamer.
Revised: With appropriate truncation strength, the resulting agent shows signiﬁcant improvement in terms of sample efﬁciency comparing to the baseline Dreamer.

Original: We suggest that truncation can act as a potential resolution for the issues caused by non-retraceable failure.
Revised: We suggest that truncation can act as a potential resolution for the issues caused by non-retraceable failure.

Original: The truncation length is currently set as a ﬁxed hyperparameter and requires careful tuning for each task.
Revised: The truncation length is currently set as a ﬁxed hyperparameter and requires careful tuning for each task.

Original: Future work could look at the adaptive tuning of the truncation length.
Revised: Future work could look at the adaptive tuning of the truncation length.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: We proposed CCWM , a novel representation learning approach that combines the forward and backward latent rollouts based on the temporal cycle-consistency constraint for model learning.
Revised: We proposed CCWM , a novel representation learning approach that combines the forward and backward latent rollouts based on the temporal cycle-consistency constraint for model learning.

Original: We deﬁne a joint objective for model training, which integrates the predictive reconstruction loss based on the ground-truth observation and the self-supervised loss that imposes temporal cycle-consistency between the original and retraced states.
Revised: We deﬁne a joint objective for model training, which integrates the predictive reconstruction loss based on the ground-truth observation and the self-supervised loss that imposes temporal cycle-consistency between the original and retraced states.

Original: CCWM is a general framework for model learning in modelbased RL, and can be combined with any stochastic sequential model as the latent dynamics model.
Revised: CCWM is a general framework for model learning in modelbased RL, and can be combined with any stochastic sequential model as the latent dynamics model.

Original: We empirically show that the model-based instantiation of the proposed approach, CCWM-A3C , yields improved performance over state-of-the-art model-based and model-free methods on a number of continuous control benchmarks, in terms of sample-efﬁciency and ﬁnal convergence score.
Revised: We empirically show that the model-based instantiation of the proposed approach, CCWM-A3C , yields improved performance over state-of-the-art model-based and model-free methods on a number of continuous control benchmarks, in terms of sample-efﬁciency and ﬁnal convergence score.

Original: We empirically investigated the interpretability and potential improvement schemes of CCWM .
Revised: We empirically investigated the interpretability and potential improvement schemes of CCWM .

Original: We use point estimates from the distribution of the latent variable at the current timestep, instead of the distribution itself, to derive the predictive inference, hence imaginary rollouts can only be performed approximately, and the errors accumulate exponentially with the imagination horizon.
Revised: We use point estimates from the distribution of the latent variable at the current timestep, instead of the distribution itself, to derive the predictive inference, hence imaginary rollouts can only be performed approximately, and the errors accumulate exponentially with the imagination horizon.

Original: Structured VAE deals with probabilistic inferences using deep learning in graphical models (Johnson et al., 2016), but is restricted to conjugate graphical model inferences.
Revised: Structured VAE deals with probabilistic inferences using deep learning in graphical models (Johnson et al., 2016), but is restricted to conjugate graphical model inferences.

Original: Future work could look at neural message passing algorithms for doing predictive inferences in general graphical structures that could support accurate multi-step inferences.
Revised: Future work could look at neural message passing algorithms for doing predictive inferences in general graphical structures that could support accurate multi-step inferences.

Original: Although hippocampal replay is thought to play a role in learning the model for model-based decisionmaking, few computational models have been proposed to account for the occurrence of forward and backward replay (Mattar & Daw, 2018).
Revised: Although hippocampal replay is thought to play a role in learning the model for model-based decisionmaking, few computational models have been proposed to account for the occurrence of forward and backward replay (Mattar & Daw, 2018).

Original: Our model suggests the forward and reversed replay are used interchangeably for learning representations that aid decision-making tasks; the forward replay elicits predictions over future events and the reversed replay supports generalisable behaviours such as identiﬁcation of bottleneck states.
Revised: Our model suggests the forward and reversed replay are used interchangeably for learning representations that aid decision-making tasks; the forward replay elicits predictions over future events and the reversed replay supports generalisable behaviours such as identiﬁcation of bottleneck states.

Original: One testable experimental prediction could look at the effect of suppressing reversed hippocampal replay on the ability of within-domain generalisation.
Revised: One testable experimental prediction could look at the effect of suppressing reversed hippocampal replay on the ability of within-domain generalisation

Original: A I MPLEMENTATION D ETAILS
Revised: A I MPLEMENTATION D ETAILS

Original: All stochastic variables in CCWM are assumed to be Gaussian.
Revised: All stochastic variables in CCWM are assumed to be Gaussian.

Original: We implement the latent transition dynamics model of CCWM using the RSSM Hafner et al.
Revised: We implement the latent transition dynamics model of CCWM using the RSSM Hafner et al.

Original: (2019), utilising a Gated Recurrent
Revised: (2019), utilising a Gated Recurrent

Original: Unit (GRU) as the core component (Chung et al., 2014) in addition to the MLP.
Revised: Unit (GRU) as the core component (Chung et al., 2014) in addition to the MLP.

Original: The RSSM is used for modelling the predictive inferences, q ψ ( z t +1 | z t , a t ) and q ψ ( z t +1 | z t , a t , O t +1 ) , from Eq.
Revised: The RSSM is used for modelling the predictive inferences, q ψ ( z t +1 | z t , a t ) and q ψ ( z t +1 | z t , a t , O t +1 ) , from Eq.

Original: The encoding function is given by a convolutional encoder that deterministically embeds the highdimensional observation into the latent space.
Revised: The encoding function is given by a convolutional encoder that deterministically embeds the highdimensional observation into the latent space.

Original: 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.
Revised: 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.

Original: The generative function q ψ ( O t | z t ) is deﬁned by an MLP followed by a deconvolution network (Zeiler et al., 2010), outputting the means the Gaussian distributions for each pixel location of the decoded observation (note that here we assume all pixel-level Gaussians have variance 1 ).
Revised: The generative function q ψ ( O t | z t ) is deﬁned by an MLP followed by a deconvolution network (Zeiler et al., 2010), outputting the means the Gaussian distributions for each pixel location of the decoded observation (note that here we assume all pixel-level Gaussians have variance 1 ).

Original: The reward distribution ˆ R ( r | z ) is modelled by an MLP, outputting a scalar mean Gaussian distribution (again, assuming unit variance).
Revised: The reward distribution ˆ R ( r | z ) is modelled by an MLP, outputting a scalar mean Gaussian distribution (again, assuming unit variance).

Original: The supervision of the model learning is based on the ELBO and the bisimulation metric as described in Section 3.
Revised: The supervision of the model learning is based on the ELBO and the bisimulation metric as described in Section 3.

Original: In the downstream A3C implementation, we use a Gaussian distribution with means and variances parameterised by an MLP for modelling the value estimates.
Revised: In the downstream A3C implementation, we use a Gaussian distribution with means and variances parameterised by an MLP for modelling the value estimates.

Original: The policy distribution is also modelled by a Gaussian distribution, with means and variances parameterised by an MLP.
Revised: The policy distribution is also modelled by a Gaussian distribution, with means and variances parameterised by an MLP.

Original: We used distributed actors for interacting with the environment for more efﬁciency collection of sampled episodes as in (Mnih et al., 2016).
Revised: We used distributed actors for interacting with the environment for more efﬁciency collection of sampled episodes as in (Mnih et al., 2016).

Original: The value function is optimised by maximising the probability of the “groundtruth” values estimated by the lambda return (Sutton & Barto, 2018).
Revised: The value function is optimised by maximising the probability of the “groundtruth” values estimated by the lambda return (Sutton & Barto, 2018).

Original: The policy distribution is updated using policy gradient estimated based generalised advantage estimation (Schulman et al., 2015).
Revised: The policy distribution is updated using policy gradient estimated based generalised advantage estimation (Schulman et al., 2015).

Original: We use the Adam optimiser for updating the parameters for the model, value function, and policy networks (Kingma & Ba, 2014).
Revised: We use the Adam optimiser for updating the parameters for the model, value function, and policy networks (Kingma & Ba, 2014).

Original: The speciﬁc conﬁgurations for the neural network implementations is summarised in Table 1.
Revised: The speciﬁc conﬁgurations for the neural network implementations is summarised in Table 1.

Original: Note that the activation functions for all non-output layers are assumed to be ReLU activation function unless otherwise stated.
Revised: Note that the activation functions for all non-output layers are assumed to be ReLU activation function unless otherwise stated.

Original: All neural network implementation are carried out using TensorFlow (Abadi et al., 2015) and
Revised: All neural network implementation are carried out using TensorFlow (Abadi et al., 2015) and

Original: TensorFlow Distributions (Dillon et al., 2017).
Revised: TensorFlow Distributions (Dillon et al., 2017).

Original: For the actual training, the batch size is chosen to be 64 , and all sampled trajectories are taken to be 50 timesteps long.
Revised: For the actual training, the batch size is chosen to be 64 , and all sampled trajectories are taken to be 50 timesteps long.

Original: After everything 10 4 training steps, we evaluate the agent over 5 random seeds.
Revised: After everything 10 4 training steps, we evaluate the agent over 5 random seeds.

Original: We adopt the same scheme for setting the action repeats equal to 2 across all tested environments as in (Hafner et al., 2020).
Revised: We adopt the same scheme for setting the action repeats equal to 2 across all tested environments as in (Hafner et al., 2020).

Original: The parameter λ controlling the weights of the retrace auxiliary loss in Eq.is set to 1 . 0 .
Revised: The parameter λ controlling the weights of the retrace auxiliary loss in Eq.is set to 1 . 0 .

Original: The discounting factor for the expected value function is set to 0 .
Revised: The discounting factor for the expected value function is set to 0 .

Original: We implemented the tSNE visualisation shown in Figure 4 by ﬁrstly applying principal component analysis initialisation (taking ﬁrst 50 principal components), followed by standard tSNE embedding.
Revised: We implemented the tSNE visualisation shown in Figure 4 by ﬁrstly applying principal component analysis initialisation (taking ﬁrst 50 principal components), followed by standard tSNE embedding.

Original: The python implementation of the CCWM-A3C can be found in the supplementary materials.
Revised: The python implementation of the CCWM-A3C can be found in the supplementary materials.

Original: B F URTHER D ISCUSSION ON T RUNCATION
Revised: B F URTHER D ISCUSSION ON T RUNCATION

Original: In Section 4.4, we introduced truncation as a general-purpose method for improving the performance of CCWM in continuous control tasks.
Revised: In Section 4.4, we introduced truncation as a general-purpose method for improving the performance of CCWM in continuous control tasks.

Original: The complete ablation studies on the effects of the strength of truncation on the remaining tasks in Figure 3a is shown in Figure 6.
Revised: The complete ablation studies on the effects of the strength of truncation on the remaining tasks in Figure 3a is shown in Figure 6.

Original: We observe that the 0 .truncation generally yields higher performance than the basic CCWM-A3C agent.
Revised: We observe that the 0 .truncation generally yields higher performance than the basic CCWM-A3C agent.

Original: The basic idea of truncation is to cut off the last proportions of the sampled episodes such that the possible non-retraceable states in the episodes will not be involved in training to update the model parameters.
Revised: The basic idea of truncation is to cut off the last proportions of the sampled episodes such that the possible non-retraceable states in the episodes will not be involved in training to update the model parameters.

Original: However, the current scheme for setting a ﬁxed truncation proportion throughout the training process seems cumbersome, and hinders learning when either under- or over-truncation occurs.
Revised: However, the current scheme for setting a ﬁxed truncation proportion throughout the training process seems cumbersome, and hinders learning when either under- or over-truncation occurs.

Original: Speciﬁcally, in Figure 5, we showed that appropriate truncation signiﬁcantly improves the performance during the early stage of learning (ﬁrst 4 × 10 5 training steps) in the hopper stand task, but is quickly surpassed by Dreamer after the early stage as shown in Figure 7.
Revised: Speciﬁcally, in Figure 5, we showed that appropriate truncation signiﬁcantly improves the performance during the early stage of learning (ﬁrst 4 × 10 5 training steps) in the hopper stand task, but is quickly surpassed by Dreamer after the early stage as shown in Figure 6.

Original: Despite the fact that the truncation improves the overall performance comparing to the basic version of CCWM-A3C ,
Revised: Despite the fact that the truncation improves the overall performance comparing to the basic version of CCWM-A3C ,

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table][Figure] [Table] [Figure]
Revised: [Table]

Original: [Table]

Original: [Figure][Table] [Figure] [Table]
Revised: [Figure]

Original: [Figure]

Original: [Table][Figure] [Table] [Figure]
Revised: [Table]

Original: [Table]

Original: [Figure][Table] [Figure] [Table]
Revised: [Figure]

Original: [Figure]

Original: [Table][Figure] [Table] [Figure]
Revised: [Table]

Original: [Table]

Original: [Figure][Table] [Figure] [Table]
Revised: [Figure]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]
Revised: [Table]

Original: Conﬁgurations for the neural network implementations of CCWM-A3C .
Revised: Conﬁgurations for the neural network implementations of CCWM-A3C .

Original: still, it is not comparable with the baseline Dreamer agent in terms of sample efﬁciency beyond the early stage.
Revised: still, it is not comparable with the baseline Dreamer agent in terms of sample efﬁciency beyond the early stage.

Original: We argue that this is due to the fact that the agent quickly learns the ability to control at a relatively high level, which causes signiﬁcant decrease in the frequency of falling of the simulated hopper (note that this is shared between the hopper stand and hopper hop tasks).
Revised: We argue that this is due to the fact that the agent quickly learns the ability to control at a relatively high level, which causes signiﬁcant decrease in the frequency of falling of the simulated hopper (note that this is shared between the hopper stand and hopper hop tasks).

Original: Hence if we stick with a ﬁxed truncation proportion, this will lead to the undesirable consequence that the states containing information that guide the agent to achieve higher performance (in the hopper stand task, this corresponds to the states where the torso is brought to a greater height, leading to higher reward received) will not be used for training the model, hence such information will be embedded into the learnt representation, hence causing slow acquisition of near-optimal behaviours.
Revised: Hence if we stick with a ﬁxed truncation proportion, this will lead to the undesirable consequence that the states containing information that guide the agent to achieve higher performance (in the hopper stand task, this corresponds to the states where the torso is brought to a greater height, leading to higher reward received) will not be used for training the model, hence such information will be embedded into the learnt representation, hence causing slow acquisition of near-optimal behaviours.

Original: We leave the study of the adaptive truncation strength for future work.
Revised: We leave the study of the adaptive truncation strength for future work.

Original: C Z ERO -S HOT T RANSFER
Revised: C Z ERO -S HOT T RANSFER

Original: Based on the motivation that learning with retracing will improve the generalisability of the agent (Figure 1a), we evaluate the generalisability of our approach on the basis of zero-shot transfer tasks. 
Revised: Based on the motivation that learning with retracing will improve the generalisability of the agent (Figure 1a), we evaluate the generalisability of our approach on the basis of zero-shot transfer tasks. [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Revised: C HANGED COMPONENTS

Revised: O UR APPROACH

Revised: D REAMER

Revised: P - VALUE

Original: [Figure]
Revised: S IGNIFICANT ? ( MEAN ± 1 S . D .) ( MEAN ± 1 S .

Original: [Figure]
Revised: D .) (3 S . F .) ( α = 0 .

Original: [Figure]
Revised: R EWARD FUNCTION [Table] R EWARD FUNCTION + M ASS [Table]

Revised: R EWARD FUNCTION + F RICTION

Revised: [Table]

Revised: R EWARD FUNCTION + S TIFFNESS

Original: [Figure]
Revised: [Table]

Revised: R EWARD FUNCTION + M ASS + F RICTION

Revised: [Table]

Revised: Y ES

Revised: R EWARD FUNCTION + M ASS + S

Revised: TIFFNESS

Revised: + F RICTION

Original:  Speciﬁcally, we modify a number of basic conﬁgurations of the cheetah run task, such as the mass of the agent, friction coefﬁcients between the joints, reward function of task, etc.
Revised: [Table] Y [Table] Speciﬁcally, we modify a number of basic conﬁgurations of the cheetah run task, such as the mass of the agent, friction coefﬁcients between the joints, reward function of task, etc.

Original: The details of the changes can be found in the appendix.
Revised: The details of the changes can be found in the appendix.

Original: Judging from the learning curve of the cheetah run task in
Revised: Judging from the learning curve of the cheetah run task in

Original: Figure 3b [Figure] despite the increased sample efﬁciency of our approach during training, both our approach and the baseline agent, Dreamer, converge at a similar value at 2 × 10 6 steps.
Revised: Figure 3b, despite the increased sample efﬁciency of our approach during training, both our approach and the baseline agent, Dreamer, converge at a similar value at 2 × 10 6 steps.

Original: We then evaluate the trained agents directly on the updated cheetah run task without any further training to test their abilities on zero-shot transfer.
Revised: We then evaluate the trained agents directly on the updated cheetah run task without any further training to test their abilities on zero-shot transfer.

Original: We show the mean evaluation scores ( ± 1 s [Figure] d.) of both agents overrandom seeds, as well as the one-sided t-test statistics and signiﬁcance of the difference between the two sets of evaluations in Table [Figure]
Revised: We show the mean evaluation scores ( ± 1 s.d.) of both agents overrandom seeds, as well as the one-sided t-test statistics and signiﬁcance of the difference between the two sets of evaluations in Table 2.

Original: It can be observed from the evaluation results that the overall performance on zero-shot transfer of our approach is comparable with Dreamer on simpler transfer tasks [Figure] and signiﬁcantly outperforms
Revised: It can be observed from the evaluation results that the overall performance on zero-shot transfer of our approach is comparable with Dreamer on simpler transfer tasks, and signiﬁcantly outperforms

Original: Dreamer on more complicated tasks.
Revised: Dreamer on more complicated tasks.

Original: Moreover, we see that the introduction of retracing improves the stability of evaluations on zero-shot transfer tasks in general.
Revised: Moreover, we see that the introduction of retracing improves the stability of evaluations on zero-shot transfer tasks in general.

Original: These conﬁrm our previous hypothesis that retracing poses a positive effect on the generalisability (Figure 1a).[Table] HANGED [Table]
Revised: These conﬁrm our previous hypothesis that retracing poses a positive effect on the generalisability (Figure 1a).

Original: APPROACH

Original: [Table]

Original: FUNCTION

Original: [Table]