Original: T EMPORAL D ISENTANGLEMENT OF R EPRESENTATIONS FOR I MPROVED G ENERALISATION IN R EINFORCEMENT L EARNING
Revised: T EMPORAL D ISENTANGLEMENT OF R EPRESENTATIONS FOR I MPROVED G ENERALISATION IN R EINFORCEMENT L EARNING

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Reinforcement Learning (RL) agents are often unable to generalise well to environment variations in the state space that were not observed during training.
Revised: Reinforcement Learning (RL) agents are often unable to generalise well to environment variations in the state space that were not observed during training.

Original: This issue is especially problematic for image-based RL, where a change in just one variable, such as the background colour, can change many pixels in the image, which can lead to drastic changes in the agent’s latent representation of the image, causing the learned policy to fail.
Revised: This issue is especially problematic for image-based RL, where a change in just one variable, such as the background colour, can change many pixels in the image, which can lead to drastic changes in the agent’s latent representation of the image, causing the learned policy to fail.

Original: To learn more robust representations, we introduce TE mporal D isentanglement (TED), a self-supervised auxiliary task that leads to disentangled image representations exploiting the sequential nature of RL observations.
Revised: To learn more robust representations, we introduce TE mporal D isentanglement (TED), a self-supervised auxiliary task that leads to disentangled image representations exploiting the sequential nature of RL observations.

Original: We find empirically that RL algorithms utilising TED as an auxiliary task adapt more quickly to changes in environment variables with continued training compared to state-of-the-art representation learning methods.
Revised: We find empirically that RL algorithms utilising TED as an auxiliary task adapt more quickly to changes in environment variables with continued training compared to state-of-the-art representation learning methods.

Original: Since TED enforces a disentangled structure of the representation, we also find that policies trained with TED generalise better to unseen values of variables irrelevant to the task (e.g. background colour) as well as unseen values of variables that affect the optimal policy (e.g. goal positions).
Revised: Since TED enforces a disentangled structure of the representation, we also find that policies trained with TED generalise better to unseen values of variables irrelevant to the task (e.g. background colour) as well as unseen values of variables that affect the optimal policy (e.g. goal positions).

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Real-world environments are often not static and deterministic, but can be subject to changes, both incremental or with sudden effects.
Revised: Real-world environments are often not static and deterministic, but can be subject to changes, both incremental or with sudden effects.

Original: Reinforcement Learning (RL) algorithms need to be robust to these changes and adapt quickly.
Revised: Reinforcement Learning (RL) algorithms need to be robust to these changes and adapt quickly.

Original: Moreover, since many real-world robotics applications rely on images as inputs (Vecerik et al., 2021; H¨am¨al¨ainen et al., 2019; Chebotar et al., 2019), RL agents need to learn robust representations of images that remain useful after a change in the environment.
Revised: Moreover, since many real-world robotics applications rely on images as inputs (Vecerik et al., 2021; H ¨ am ¨ al ¨ ainen et al., 2019; Chebotar et al., 2019), RL agents need to learn robust representations of images that remain useful after a change in the environment.

Original: For example, a simple change in lighting conditions can change the perceived colour of an object, but this should not affect the agent’s ability to perform a task.
Revised: For example, a simple change in lighting conditions can change the perceived colour of an object, but this should not affect the agent’s ability to perform a task.

Original: One of the reasons RL agents fail to generalise to unseen values of environment variables, such as colours and object positions, is that they overfit to variations seen in training (Zhang et al., 2018).
Revised: One of the reasons RL agents fail to generalise to unseen values of environment variables, such as colours and object positions, is that they overfit to variations seen in training (Zhang et al., 2018).

Original: The issue is especially problematic for image-based RL, where a change in one environment variable can mean the agent is presented with a very different set of pixels for which trained RL policies are often no longer optimal.
Revised: The issue is especially problematic for image-based RL, where a change in one environment variable can mean the agent is presented with a very different set of pixels for which trained RL policies are often no longer optimal.

Original: This failure to generalise occurs for both variables that are irrelevant to the optimal policy, such as background colours, and relevant variables, such as goal positions (Kirk et al., 2022).
Revised: This failure to generalise occurs for both variables that are irrelevant to the optimal policy, such as background colours, and relevant variables, such as goal positions (Kirk et al., 2022).

Original: In practice, this often results in agents needing to adapt their policy after a change to only one variable.
Revised: In practice, this often results in agents needing to adapt their policy after a change to only one variable.

Original: We show experimentally that agents often cannot recover the optimal policy after the environment changes because it is too difficult to ‘undo’ the overfitting.
Revised: We show experimentally that agents often cannot recover the optimal policy after the environment changes because it is too difficult to ‘undo’ the overfitting.

Original: One approach to tackle the generalisation issue is to use domain randomisation during training to maximise the environment variations observed (Cobbe et al., 2019; Chebotar et al., 2019).
Revised: One approach to tackle the generalisation issue is to use domain randomisation during training to maximise the environment variations observed (Cobbe et al., 2019; Chebotar et al., 2019).

Original: However, in practice, we may be unaware of what variations an agent might see in the future.
Revised: However, in practice, we may be unaware of what variations an agent might see in the future.

Original: Even if all possible future variations are known, training on this full set is often sample inefficient and may result in a sub-optimal policy as the agent learns to compensate for many different values.
Revised: Even if all possible future variations are known, training on this full set is often sample inefficient and may result in a sub-optimal policy as the agent learns to compensate for many different values.

Original: It can also be impractical to generate all variations during training in the real world or even in simulation due to limitations on laboratory setups or in simulation suites.
Revised: It can also be impractical to generate all variations during training in the real world or even in simulation due to limitations on laboratory setups or in simulation suites.

Original: An alternative to domain randomisation is to learn robust representations that generalise to unseen variations (Kirk et al., 2022), but learning such representations remains an open problem (Lan et al., 2022).
Revised: An alternative to domain randomisation is to learn robust representations that generalise to unseen variations (Kirk et al., 2022), but learning such representations remains an open problem (Lan et al., 2022).

Original: Some approaches aim to learn a representation that is invariant to image distractors (Zhang et al., 2020; 2021) to improve generalisation to unseen values of only variables irrelevant to the optimal policy.
Revised: Some approaches aim to learn a representation that is invariant to image distractors (Zhang et al., 2020; 2021) to improve generalisation to unseen values of only variables irrelevant to the optimal policy.

Original: These approaches do not enforce a structure on the remaining task-relevant variables in the representation for generalisation to unseen values of relevant variables.
Revised: These approaches do not enforce a structure on the remaining task-relevant variables in the representation for generalisation to unseen values of relevant variables.

Original: A promising direction towards robust RL is to learn a disentangled representation of observations.
Revised: A promising direction towards robust RL is to learn a disentangled representation of observations.

Original: Disentanglement techniques aim to learn robust representations for both task relevant and irrelevant variables by separating distinct, informative factors of variation in an image into the unknown ground truth factors that generated the image (Bengio et al., 2013).
Revised: Disentanglement techniques aim to learn robust representations for both task relevant and irrelevant variables by separating distinct, informative factors of variation in an image into the unknown ground truth factors that generated the image (Bengio et al., 2013).

Original: When one factor of variation changes to a previously unseen value, such as a new colour, changing many pixels in the image, only a subset of features in a disentangled representation will change.
Revised: When one factor of variation changes to a previously unseen value, such as a new colour, changing many pixels in the image, only a subset of features in a disentangled representation will change.

Original: The RL agent will still be able to rely on the remaining unchanged features in the representation to adapt quickly, allowing generalisation and performance recovery similar to state-based RL.
Revised: The RL agent will still be able to rely on the remaining unchanged features in the representation to adapt quickly, allowing generalisation and performance recovery similar to state-based RL.

Original: Higgins et al.
Revised: Higgins et al.

Original: (2017b) show that a disentangled representation improves generalisation in RL.
Revised: (2017b) show that a disentangled representation improves generalisation in RL.

Original: However, this approach requires a trained or hard-coded policy to collect independent and identically distributed (i.i.d.) data to pre-train a β -VAE (Higgins et al., 2017a) offline, and it has since been proven that it is theoretically impossible to learn a disentangled representation from i.i.d.
Revised: However, this approach requires a trained or hard-coded policy to collect independent and identically distributed (i.i.d.) data to pre-train a β -VAE (Higgins et al., 2017a) offline, and it has since been proven that it is theoretically impossible to learn a disentangled representation from i.i.d.

Original: data alone (Locatello et al., 2019).
Revised: data alone (Locatello et al., 2019).

Original: We introduce a self-supervised auxiliary task for learning disentangled representations for the robust encoding of images, which we call TE mporal D isentanglement (TED).
Revised: We introduce a self-supervised auxiliary task for learning disentangled representations for the robust encoding of images, which we call TE mporal D isentanglement (TED).

Original: Unlike previous work, TED can be implemented with only minimal changes to existing RL algorithms and allows for lifelong learning of disentangled representations.
Revised: Unlike previous work, TED can be implemented with only minimal changes to existing RL algorithms and allows for lifelong learning of disentangled representations.

Original: In contrast to Higgins et al. (2017b), our approach uses the non-i.i.d.
Revised: In contrast to Higgins et al. (2017b), our approach uses the non-i.i.d.

Original: temporal data from consecutive timesteps in RL to learn the disentangled representation online.
Revised: temporal data from consecutive timesteps in RL to learn the disentangled representation online.

Original: We provide experimental results across a variety of tasks from the DeepMind Control
Revised: We provide experimental results across a variety of tasks from the DeepMind Control

Original: Suite (Tunyasuvunakool et al., 2020), Panda Gym (Gallou´edec et al., 2021) and Procgen (Cobbe et al., 2020) environments.
Revised: Suite (Tunyasuvunakool et al., 2020), Panda Gym (Gallou ´ edec et al., 2021) and Procgen (Cobbe et al., 2020) environments.

Original: For our experiments, we train on a subset of some of the environment variables (such as colours), then evaluate generalisation on a test environment with unseen values of the variables and continue training to demonstrate adaptation to the new environment.
Revised: For our experiments, we train on a subset of some of the environment variables (such as colours), then evaluate generalisation on a test environment with unseen values of the variables and continue training to demonstrate adaptation to the new environment.

Original: Our results demonstrate that TED improves generalisation of a variety of base RL algorithms and adapts more quickly than representation learning baselines on unseen environment variables that are relevant or irrelevant to the task.
Revised: Our results demonstrate that TED improves generalisation of a variety of base RL algorithms and adapts more quickly than representation learning baselines on unseen environment variables that are relevant or irrelevant to the task.

Original: We also evaluate a disentanglement metric (Higgins et al., 2017a) to demonstrate that our approach increases the extent to which the learned representation has disentangled the factors of variation in the image observations compared to baselines.
Revised: We also evaluate a disentanglement metric (Higgins et al., 2017a) to demonstrate that our approach increases the extent to which the learned representation has disentangled the factors of variation in the image observations compared to baselines.

Original:  R ELATED W ORK 2. G ENERALISATIONIN I MAGE - BASED R EINFORCEMENT L EARNING
Revised:  R ELATED W ORK 2. G ENERALISATION

Revised: IN I MAGE - BASED R EINFORCEMENT L EARNING

Original: Image augmentation.
Revised: Image augmentation.

Original: Image augmentation artificially increases the size of the dataset by adding image perturbations to improve robustness of image representations.
Revised: Image augmentation artificially increases the size of the dataset by adding image perturbations to improve robustness of image representations.

Original: Laskin et al.
Revised: Laskin et al.

Original: (2020a) apply a variety of image augmentation techniques such as translation, cutouts and cropping; Yarats et al.
Revised: (2020a) apply a variety of image augmentation techniques such as translation, cutouts and cropping; Yarats et al.

Original: (2021) average over multiple augmentations; Hansen & Wang (2021)
Revised: (2021) average over multiple augmentations; Hansen & Wang (2021)

Original: maximise mutual information between the representations of augmented and non-augmented images; and Hansen et al.
Revised: maximise mutual information between the representations of augmented and non-augmented images; and Hansen et al.

Original: (2021b) use both augmented and non-augmented images to stabilise Q-value estimation.
Revised: (2021b) use both augmented and non-augmented images to stabilise Q-value estimation.

Original: However, image augmentation approaches to generalisation can still fail when the agent experiences stronger types of variation after training (Kirk et al., 2022).
Revised: However, image augmentation approaches to generalisation can still fail when the agent experiences stronger types of variation after training (Kirk et al., 2022).

Original: In our experiments, we show that TED can be used alongside image augmentation techniques to further improve generalisation while benefiting from the augmentation.
Revised: In our experiments, we show that TED can be used alongside image augmentation techniques to further improve generalisation while benefiting from the augmentation.

Original: Learning invariant representations.
Revised: Learning invariant representations.

Original: Invariance techniques aim to learn a representation that ignores distractors in the image.
Revised: Invariance techniques aim to learn a representation that ignores distractors in the image.

Original: Zhang et al.
Revised: Zhang et al.

Original: (2020) use causal inference techniques assuming a block MDP structure, Zhang et al.
Revised: (2020) use causal inference techniques assuming a block MDP structure, Zhang et al.

Original: (2021) use a bisimulation metric.
Revised: (2021) use a bisimulation metric.

Original: Li et al.
Revised: Li et al.

Original: (2021) use domain adversarial optimisation to learn a representation invariant to visual distractors.
Revised: (2021) use domain adversarial optimisation to learn a representation invariant to visual distractors.

Original: These approaches all aim to generalise to unseen values of irrelevant variables, e.g. background colours.
Revised: These approaches all aim to generalise to unseen values of irrelevant variables, e.g. background colours.

Original: In contrast, TED uses disentanglement to enforce a structured representation that applies to both relevant and irrelevant variables.
Revised: In contrast, TED uses disentanglement to enforce a structured representation that applies to both relevant and irrelevant variables.

Original: Encoding inductive biases.
Revised: Encoding inductive biases.

Original: Some approaches use an auxiliary task to encode inductive biases.
Revised: Some approaches use an auxiliary task to encode inductive biases.

Original: Laskin et al.
Revised: Laskin et al.

Original: (2020b) learn a representation to maximise similarity between different augmentations of the same observation, while van den Oord et al.
Revised: (2020b) learn a representation to maximise similarity between different augmentations of the same observation, while van den Oord et al.

Original: (2018) learn representations that capture information predictive of the future.
Revised: (2018) learn representations that capture information predictive of the future.

Original: Agarwal et al.
Revised: Agarwal et al.

Original: (2021) enforce a structured representation for task relevant variables using policy similarity metrics.
Revised: (2021) enforce a structured representation for task relevant variables using policy similarity metrics.

Original: Hansen et al.
Revised: Hansen et al.

Original: (2021a) explore adaptation to new environments with no reward signal using self-supervision to continue training in the new environment.
Revised: (2021a) explore adaptation to new environments with no reward signal using self-supervision to continue training in the new environment.

Original: While these approaches constrain the latent space, they do not aim to disentangle the representation.
Revised: While these approaches constrain the latent space, they do not aim to disentangle the representation.

Original: To encode the inductive bias of disentanglement, Higgins et al.
Revised: To encode the inductive bias of disentanglement, Higgins et al.

Original: (2017b) train a β -VAE offline using data from a pre-trained or hardcoded agent.
Revised: (2017b) train a β -VAE offline using data from a pre-trained or hardcoded agent.

Original: In contrast, our approach uses the temporal structure of data available in RL to learn a disentangled representation online with the RL policy.
Revised: In contrast, our approach uses the temporal structure of data available in RL to learn a disentangled representation online with the RL policy.

Original: 2. D ISENTANGLED R EPRESENTATIONS
Revised: 2. D ISENTANGLED R EPRESENTATIONS

Original: Unsupervised learning.
Revised: Unsupervised learning.

Original: Many variations of the Variational Autoencoder (VAE) (Kingma & Welling, 2014) aim to improve disentanglement of the learned representation, such as the β -VAE (Higgins et al., 2017a)
Revised: Many variations of the Variational Autoencoder (VAE) (Kingma & Welling, 2014) aim to improve disentanglement of the learned representation, such as the β -VAE (Higgins et al., 2017a)

Original: (Burgess et al., 2017) and the Factor-VAE (Kim & Mnih, 2018).
Revised: (Burgess et al., 2017) and the Factor-VAE (Kim & Mnih, 2018).

Original: Locatello et al.
Revised: Locatello et al.

Original: (2019) prove that it is theoretically impossible to learn disentangled representations from i.i.d.
Revised: (2019) prove that it is theoretically impossible to learn disentangled representations from i.i.d.

Original: data alone.
Revised: data alone.

Original: To bypass the impossibility result, many recent approaches extend the β -VAE to use some form of supervision.
Revised: To bypass the impossibility result, many recent approaches extend the β -VAE to use some form of supervision.

Original: Shu et al.
Revised: Shu et al.

Original: (2020) provide weak supervision using a labelled grouping of images, while
Revised: (2020) provide weak supervision using a labelled grouping of images, while

Original: Locatello et al.
Revised: Locatello et al.

Original: (2020) generate pairs of images with limited factors of variation between pairs.
Revised: (2020) generate pairs of images with limited factors of variation between pairs.

Original: In contrast, our approach exploits non-i.i.d.
Revised: In contrast, our approach exploits non-i.i.d.

Original: temporal observations available in RL to learn a disentangled representation without labelling or artificially generating images.
Revised: temporal observations available in RL to learn a disentangled representation without labelling or artificially generating images.

Original: Independent Component Analysis (ICA).
Revised: Independent Component Analysis (ICA).

Original: Hyv ¨ arinen & Morioka (2016) introduce TimeContrastive Learning to learn a disentangled representation from time-series data with non-stationary factors.
Revised: Hyv ¨ arinen & Morioka (2016) introduce TimeContrastive Learning to learn a disentangled representation from time-series data with non-stationary factors.

Original: Hyv¨arinen & Morioka (2017) disentangle time-series data with stationary factors, introducing
Revised: Hyv ¨ arinen & Morioka (2017) disentangle time-series data with stationary factors, introducing

Original: Permutation Contrastive Learning (PCL) to train a classifier to discriminate between temporal and non-temporal inputs.
Revised: Permutation Contrastive Learning (PCL) to train a classifier to discriminate between temporal and non-temporal inputs.

Original: We consider RL episodes back-to-back as a time-series.
Revised: We consider RL episodes back-to-back as a time-series.

Original: Features that are randomised at the start of an episode are stationary, and due to the agent learning and adapting behaviour over time, features controlled by the agent are non-stationary.
Revised: Features that are randomised at the start of an episode are stationary, and due to the agent learning and adapting behaviour over time, features controlled by the agent are non-stationary.

Original: The TED classification objective is based on the PCL classifier structure to encourage disentanglement, but unlike PCL, TED addresses the combination of stationary and non-stationary features in RL.
Revised: The TED classification objective is based on the PCL classifier structure to encourage disentanglement, but unlike PCL, TED addresses the combination of stationary and non-stationary features in RL.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: We assume the environment is a fully-observable Markov Decision Process (MDP), defined as the tuple M = ( S , A , P, R, γ ) , where S is a set of states, A is a set of actions, P : S × S ×
Revised: We assume the environment is a fully-observable Markov Decision Process (MDP), defined as the tuple M = ( S , A , P, R, γ ) , where S is a set of states, A is a set of actions, P : S × S ×

Original: A → [0 , 1] is the state-transition function, R : S ×
Revised: A → [0 , 1] is the state-transition function, R : S ×

Original: A → R is the reward function, and γ ∈ [0 , 1) is the discount factor.
Revised: A → R is the reward function, and γ ∈ [0 , 1) is the discount factor.

Original: An RL agent chooses an action a t ∈ A at time t based on its current state s t ∈ S and its policy a t ∼ π ( s t ) .
Revised: An RL agent chooses an action a t ∈ A at time t based on its current state s t ∈ S and its policy a t ∼ π ( s t ) .

Original: The agent then transitions to the next state according to the state-transition probability
Revised: The agent then transitions to the next state according to the state-transition probability

Original: P ( s t +1 | s t , a t ) , and receives a reward, r t =
Revised: P ( s t +1 | s t , a t ) , and receives a reward, r t =

Original: R ( s t , a t ) .
Revised: R ( s t , a t ) .

Original: The goal of an RL agent is to learn a policy π to maximise the expected discounted cumulative rewards, max π E P [ (cid:80) ∞ t =

Revised: The goal of an RL agent is to learn a policy π to maximise the expected discounted cumulative rewards, max π E P [ (cid:80) ∞ t =

Original: [ γ t r t ]] .
Revised: [ γ t r t ]] .

Original: In this work, the agent’s observation o t ∈ O at time t is image pixels, a high-dimensional representation of the true underlying environment state s t ∈ S .
Revised: In this work, the agent’s observation o t ∈ O at time t is image pixels, a high-dimensional representation of the true underlying environment state s t ∈ S .

Original: We assume the environment has a factored state representation s t , where each factor corresponds to an environment variable.
Revised: We assume the environment has a factored state representation s t , where each factor corresponds to an environment variable.

Original: The components of the state vector s t are the unobserved ground truth factors of variation.
Revised: The components of the state vector s t are the unobserved ground truth factors of variation.

Original: We consider one observation to be a stack of consecutive frames to ensure all features of s t , such as velocities, can be extracted from o t .
Revised: We consider one observation to be a stack of consecutive frames to ensure all features of s t , such as velocities, can be extracted from o t .

Original: We assume o t is generated by an invertible, non-linear transformation of the state factors s t , h : S
Revised: We assume o t is generated by an invertible, non-linear transformation of the state factors s t , h : S

Original: → O , such that o t = h ( s t ) .
Revised: → O , such that o t = h ( s t ) .

Original: The aim of disentanglement is to learn a representation z t ∈
Revised: The aim of disentanglement is to learn a representation z t ∈

Original: Z that recovers the independent ground truth factors of s t from the observations o t by approximating the inverse of the transformation, such that z t = f
Revised: Z that recovers the independent ground truth factors of s t from the observations o t by approximating the inverse of the transformation, such that z t = f

Original: ( o t ) .
Revised: ( o t ) .

Original: To simplify notation, we will usually denote z t = f ( o t ) as z ( o t ) .
Revised: To simplify notation, we will usually denote z t = f ( o t ) as z ( o t ) .

Original: We will use z i to denote the i -th component of the vector z . 
Revised: We will use z i to denote the i -th component of the vector z . [Figure]

Original:  T EMPORAL D ISENTANGLEMENT
Revised:  T EMPORAL D ISENTANGLEMENT

Original: We introduce TEmporal Disentanglement (TED) as an auxiliary task for representation learning with
Revised: We introduce TEmporal Disentanglement (TED) as an auxiliary task for representation learning with

Original: RL algorithms.
Revised: RL algorithms.

Original: Our goal is to improve generalisation to both relevant and irrelevant features when their testing variation is unknown a priori.
Revised: Our goal is to improve generalisation to both relevant and irrelevant features when their testing variation is unknown a priori.

Original: TED aims to disentangle the factors of variation that generate an observation o t , such as background colour and the trajectory of an object across the[Figure] frame stack.
Revised: TED aims to disentangle the factors of variation that generate an observation o t , such as background colour and the trajectory of an object across the frame stack.

Original: We will provide an overview of architecture in Section 4.1 then discuss the details of the
Revised: We will provide an overview of architecture in Section 4.1 then discuss the details of the

Original: TED auxiliary task in Section 4.2.
Revised: TED auxiliary task in Section 4.2.

Original: 4. A RCHITECTURE O VERVIEW
Revised: 4. A RCHITECTURE O VERVIEW

Original: The high-level architecture for TED with a generic RL algorithm is depicted in Figure 1.
Revised: The high-level architecture for TED with a generic RL algorithm is depicted in Figure 1.

Original: TED is designed to encourage the encoder f
Revised: TED is designed to encourage the encoder f

Original: θ : O → Z to recover the temporal structure determined by observations at consecutive timesteps, o t and o t +1 , such that a classifier g ϕ :
Revised: θ : O → Z to recover the temporal structure determined by observations at consecutive timesteps, o t and o t +1 , such that a classifier g ϕ :

Original: Z × Z → R can discriminate between temporal and non-temporal pairs of observation encodings.
Revised: Z × Z → R can discriminate between temporal and non-temporal pairs of observation encodings.

Original: We structure the classifier to compare each feature in the representation z t separately to encourage the encoder to disentangle the temporal structure into the ground truth factors of s t that generated the observation.
Revised: We structure the classifier to compare each feature in the representation z t separately to encourage the encoder to disentangle the temporal structure into the ground truth factors of s t that generated the observation.

Original: We simultaneously perform dimensionality reduction such that dim ( S ) ≤ dim ( Z ) ≪ dim ( O ) .
Revised: We simultaneously perform dimensionality reduction such that dim ( S ) ≤ dim ( Z ) ≪ dim ( O ) .

Original: TED requires a batch of transitions B containing consecutive observations originating from various episodes.
Revised: TED requires a batch of transitions B containing consecutive observations originating from various episodes.

Original: The classifier loss is applied as an auxiliary loss to any base RL algorithm.
Revised: The classifier loss is applied as an auxiliary loss to any base RL algorithm.

Original: For off-policy algorithms, the batch can be sampled from the replay buffer B ∼ D following the sampling procedure for the base RL algorithm.
Revised: For off-policy algorithms, the batch can be sampled from the replay buffer B ∼ D following the sampling procedure for the base RL algorithm.

Original: For on-policy base algorithms, the batch can be created using multiple parallel environments to ensure transitions from different episodes, which is common in practice, for which B
Revised: For on-policy base algorithms, the batch can be created using multiple parallel environments to ensure transitions from different episodes, which is common in practice, for which B

Original: = D in the subsequent explanations.
Revised: = D in the subsequent explanations.

Original: Image augmentations can be applied to the observations where required by the base algorithm, and TED uses the same augmented images as the base algorithm.
Revised: Image augmentations can be applied to the observations where required by the base algorithm, and TED uses the same augmented images as the base algorithm.

Original: Both the TED classification loss and the relevant RL loss are used to learn the encoder parameters θ .
Revised: Both the TED classification loss and the relevant RL loss are used to learn the encoder parameters θ .

Original: For training stability, we also use a target encoder f
Revised: For training stability, we also use a target encoder f

Original: θ ′ :
Revised: θ ′ :

Original: O → Z for the next observation at a given timestep o t +1 , where θ ′ =
Revised: O → Z for the next observation at a given timestep o t +1 , where θ ′ =

Original: τθ ′
Revised: τθ ′

Original: + (1 − τ ) θ .
Revised: + (1 − τ ) θ .

Original: Only temporal transitions are used for RL, but non-temporal pairs are created where ˜ t ̸ = t
Revised: Only temporal transitions are used for RL, but non-temporal pairs are created where ˜ t ̸ = t

Original: + 1 to train the classifier, which we describe in more detail in the next section.
Revised: + 1 to train the classifier, which we describe in more detail in the next section.

Original: 4. T EMPORAL D ISENTANGLEMENT C LASSIFIER
Revised: 4. T EMPORAL D ISENTANGLEMENT C LASSIFIER

Original: We use the TED auxiliary loss to train a classifier to discriminate between temporal and nontemporal pairs of observation encodings.
Revised: We use the TED auxiliary loss to train a classifier to discriminate between temporal and nontemporal pairs of observation encodings.

Original: TED encourages the encoder to learn a representation that uncovers the temporal structure in the data to enable distinguishing whether two observations occurred consecutively or not.
Revised: TED encourages the encoder to learn a representation that uncovers the temporal structure in the data to enable distinguishing whether two observations occurred consecutively or not.

Original: Classifier inputs.
Revised: Classifier inputs.

Original: For each batch of transitions B ∼ D , we create three types of observation-pair batches for classifier input: 1) temporal samples X , 2) different episode, non-temporal samples X ′ , and 3) same episode, non-temporal samples X ′′ .
Revised: For each batch of transitions B ∼ D , we create three types of observation-pair batches for classifier input: 1) temporal samples X , 2) different episode, non-temporal samples X ′ , and 3) same episode, non-temporal samples X ′′ .

Original: A single transition contains ( o et , o et +1 ) ∈ B where e is the episode index corresponding to transition.
Revised: A single transition contains ( o et , o et +1 ) ∈ B where e is the episode index corresponding to transition.

Original: The temporal sample x t =
Revised: The temporal sample x t =

Original: ( o et , o et +1 ) consists of two consecutive observations.
Revised: ( o et , o et +1 ) consists of two consecutive observations.

Original: Due to the episodic nature of RL, we use two types of non-temporalsamples.
Revised: Due to the episodic nature of RL, we use two types of non-temporal

Original: The non-temporal sample x ′ t =

Original: ( o et , o e ′ t ′ ) consists of non-consecutive timesteps from different episodes, where o e ′ t ′ ∼ B

Original: such that e ′ ̸

Original: = e .

Original: These samples encourage the encoder, f θ , to learn a representation, z ( o et ) , that disentangles episodic features that are chosen randomly at the start

Original: Input: batch of transitions B = { ..., ( o et , o et +1 ) , ... } ∼ D Create batch Z obs of representations for each observation in B : z et = f
Revised: Input: batch of transitions B = { ..., ( o et , o et +1 ) , ... } ∼ D Create batch Z obs of representations for each observation in B : z et = f

Original: θ ( o et ) Create batch Z next obs of representations for each next observation in B : z et [-+
-] =
Revised: θ ( o et ) Create batch Z next obs of representations for each next observation in B : z et [-+
-]+1 =

Original: f θ ′
Revised: f θ ′

Original: ( o et +1 ) for each transition in B do
Revised: ( o et +1 ) for each transition in B do

Original: Create temporal sample x t ← ( z et , z et +1 ) , X ← X ∪ x t
Revised: Create temporal sample x t ← ( z et , z et +1 ) , X ← X ∪ x t

Original: Sample z e ′ t ′ ∼ Z next obs such that e ′ ̸
Revised: Sample z e ′ t ′ ∼ Z next obs such that e ′ ̸

Original: = e
Revised: = e

Original: Create non-temporal different episode sample x ′ t ← ( z et , z e ′ t ′ ) , X ′ ← X ′ ∪ x ′ t
Revised: Create non-temporal different episode sample x ′ t ← ( z et , z e ′ t ′ ) , X ′ ← X ′ ∪ x ′ t

Original: Sample o et ′′ ∼ D such that t ′′ / ∈ { t, t + 1 } , and get representation z et ′′ =
Revised: Sample o et ′′ ∼ D such that t ′′ / ∈ { t, t + 1 } , and get representation z et ′′ =

Original: f θ ′ ( o et ′′ )
Revised: f θ ′ ( o et ′′ )

Original: Create non-temporal same episode sample x ′′ t ← ( z et , z et ′′ ) , X ′′ ← X ′′ ∪ x ′′ t end for for each sample x ∈ { X, X ′ , X ′′ } do
Revised: Create non-temporal same episode sample x ′′ t ← ( z et , z et ′′ ) , X ′′ ← X ′′ ∪ x ′′ t end for for each sample x ∈ { X, X ′ , X ′′ } do

Original: Classifier prediction y
Revised: Classifier prediction y

Original: =
Revised: =

Original: g ϕ ( x ) (see Equation 1)
Revised: g ϕ ( x ) (see Equation 1)

Original: Calculate binary cross-entropy loss L TED ( x , l ) (see Equation 2) end for
Revised: Calculate binary cross-entropy loss L TED ( x , l ) (see Equation 2) end for

Original: Calculate average loss for the batch L TED ← mean ( L TED ( x ))
Revised: Calculate average loss for the batch L TED ← mean ( L TED ( x ))

Original: Backpropogate loss to update encoder parameters θ and classifier parameters ϕ
Revised: Backpropogate loss to update encoder parameters θ and classifier parameters ϕ

Original: Update target encoder parameters θ ′ =
Revised: Update target encoder parameters θ ′ =

Original: τθ ′
Revised: τθ ′

Original: + (1 − τ ) θ
Revised: + (1 − τ ) θ

Original: Output: Loss L TED and updated parameters ϕ , θ , and θ ′ of an episode, such as colours, as this is sufficient to distinguish x ′ t  from the temporal sample x t .
Revised: Output: Loss L TED and updated parameters ϕ , θ , and θ ′ samples. The non-temporal sample x ′ t = ( o et , o e ′ t ′ ) consists of non-consecutive timesteps from different episodes, where o e ′ t ′ ∼ B such that e ′ ̸

Revised: = e .

Revised: These samples encourage the encoder, f θ , to learn a representation, z ( o et ) , that disentangles episodic features that are chosen randomly at the start of an episode, such as colours, as this is sufficient to distinguish x ′ t from the temporal sample x t .

Original: The second non-temporal sample x ′′ t =
Revised: The second non-temporal sample x ′′ t =

Original: ( o et , o et ′′ ) consists of non-consecutive timesteps from the same episode where o et ′′ ∼ D .
Revised: ( o et , o et ′′ ) consists of non-consecutive timesteps from the same episode where o et ′′ ∼ D .

Original: These same episode non-temporal samples encourage the encoder to disentangle features that change during the episode, such as agent and object positions, because episodic features will be the same for both x t and x ′′ t .
Revised: These same episode non-temporal samples encourage the encoder to disentangle features that change during the episode, such as agent and object positions, because episodic features will be the same for both x t and x ′′ t .

Original: Classification objective.
Revised: Classification objective.

Original: Given a batch of samples { X, X ′ , X ′′ } , a logistic regression classifier is trained to discriminate between the corresponding representation of temporal samples, x t ∈ X , and non-temporal samples, x ′ t ∈ X ′ and x ′′ t ∈ X ′′ .
Revised: Given a batch of samples { X, X ′ , X ′′ } , a logistic regression classifier is trained to discriminate between the corresponding representation of temporal samples, x t ∈ X , and non-temporal samples, x ′ t ∈ X ′ and x ′′ t ∈ X ′′ .

Original: To learn a disentangled representation, we use a regression function of the form proposed by Hyv ¨ arinen & Morioka (2017):
Revised: To learn a disentangled representation, we use a regression function of the form proposed by Hyv ¨ arinen & Morioka (2017):

Original: [Equation]
Revised: [Equation]

Original: for all x t ∈ { X, X ′ , X ′′ } , where n is the dimensionality of the representation z ( o et ) , x t = ( x 1 t , x 2 t ) , and ϕ = { k 1 , k 2 , b , ¯k , ¯b , c } are the classifier parameters to be trained simultaneously with the encoder parameters θ .
Revised: for all x t ∈ { X, X ′ , X ′′ } , where n is the dimensionality of the representation z ( o et ) , x t = ( x 1 t , x 2 t ) , and ϕ = { k 1 , k 2 , b , ¯k , ¯b , c } are the classifier parameters to be trained simultaneously with the encoder parameters θ .

Original: Temporal samples, x t ∈ X , are given a classification label l = 1 , and non-temporal samples, x ′ t ∈ X ′ and x ′′ t ∈ X ′′ , are given a classification label l = 0 .
Revised: Temporal samples, x t ∈ X , are given a classification label l = 1 , and non-temporal samples, x ′ t ∈ X ′ and x ′′ t ∈ X ′′ , are given a classification label l = 0 .

Original: The classifier is trained using the cross-entropy loss for binary classification for all x t ∈ { X, X ′ , X ′′ } :
Revised: The classifier is trained using the cross-entropy loss for binary classification for all x t ∈ { X, X ′ , X ′′ } :

Original: [Equation]
Revised: [Equation]

Original: where σ is the sigmoid function.
Revised: where σ is the sigmoid function.

Original: Since there are two non-temporal samples, x ′ t ∈ X ′ and x ′′ t ∈ X ′′ , for each transition in the batch B but only one temporal, x t ∈ X , the positive temporal samples are weighted by 2 .
Revised: Since there are two non-temporal samples, x ′ t ∈ X ′ and x ′′ t ∈ X ′′ , for each transition in the batch B but only one temporal, x t ∈ X , the positive temporal samples are weighted by 2 .

Original: The coefficient α is a hyperparameter to be tuned to the task.
Revised: The coefficient α is a hyperparameter to be tuned to the task.

Original: It is used to scale up the classification auxiliary loss to ensure it is not dwarfed by the RL loss and to prioritise representation learning over policy learning at the start of training while the factors of variation, s t , are independent.
Revised: It is used to scale up the classification auxiliary loss to ensure it is not dwarfed by the RL loss and to prioritise representation learning over policy learning at the start of training while the factors of variation, s t , are independent.

Original: The structure of the classifier y =
Revised: The structure of the classifier y =

Original: g ϕ ( x 1 t , x 2 t ) (Equation 1) ensures each feature
Revised: g ϕ ( x 1 t , x 2 t ) (Equation 1) ensures each feature

Original: z i is considered separately to the other features z
Revised: z i is considered separately to the other features z

Original: j ̸ = i .
Revised: j ̸ = i .

Original: This encourages the encoder to not only uncover the temporal structure necessary for classification, but to do so by separating the factors of variation into distinct features in the representation so that the temporal structure of each feature can be determined independently of the other features.
Revised: This encourages the encoder to not only uncover the temporal structure necessary for classification, but to do so by separating the factors of variation into distinct features in the representation so that the temporal structure of each feature can be determined independently of the other features.

Original: The second term of Equation 1 approximates the marginal log-pdf.
Revised: The second term of Equation 1 approximates the marginal log-pdf.

Original: Due to this structure, we expect the encoder to learn a disentangled representation and the classifier to approximate the distribution of the samples (Hyv¨arinen & Morioka, 2017).
Revised: Due to this structure, we expect the encoder to learn a disentangled representation and the classifier to approximate the distribution of the samples (Hyv¨arinen & Morioka, 2017).

Original: The pseudocode for an update step is provided in Algorithm 1, which is performed for every update of the base RL algorithm.
Revised: The pseudocode for an update step is provided in Algorithm 1, which is performed for every update of the base RL algorithm.

Original: We train the encoder with both L TED and the RL loss.
Revised: We train the encoder with both L TED and the RL loss.

Original: It is important to note that the TED auxiliary loss requires only creating temporal and non-temporal pairs of representations to train the classifier.
Revised: It is important to note that the TED auxiliary loss requires only creating temporal and non-temporal pairs of representations to train the classifier.

Original: The classifier is not required for execution so it can be discarded after training.
Revised: The classifier is not required for execution so it can be discarded after training.

Original:  E XPERIMENTAL R ESULTS
Revised:  E XPERIMENTAL R ESULTS

Original: Our experiments are designed to evaluate whether TED allows zero-shot generalisation to unseen values of an environment variable, and whether TED promotes faster adaptation after a change in environment variables if generalisation is not instant.
Revised: Our experiments are designed to evaluate whether TED allows zero-shot generalisation to unseen values of an environment variable, and whether TED promotes faster adaptation after a change in environment variables if generalisation is not instant.

Original: We evaluate our approach on settings where the agent must generalise to unseen values of both task relevant and irrelevant distractor variables with continued learning on the test environment.
Revised: We evaluate our approach on settings where the agent must generalise to unseen values of both task relevant and irrelevant distractor variables with continued learning on the test environment.

Original: We use a training environment with a subset of some of the environment variables (such as colours), and evaluate generalisation on a test environment with unseen values of the variables.
Revised: We use a training environment with a subset of some of the environment variables (such as colours), and evaluate generalisation on a test environment with unseen values of the variables.

Original: We show results on a variety of different tasks with RAD (Laskin et al., 2020a), SVEA (Hansen et al., 2021b) and PPO (Schulman et al., 2017) as different base algorithms utilising the TED loss, which covers on-policy and off-policy, continuous and discrete control, and with and without image augmentations, demonstrating the flexibility of our approach.
Revised: We show results on a variety of different tasks with RAD (Laskin et al., 2020a), SVEA (Hansen et al., 2021b) and PPO (Schulman et al., 2017) as different base algorithms utilising the TED loss, which covers on-policy and off-policy, continuous and discrete control, and with and without image augmentations, demonstrating the flexibility of our approach.

Original: Our results show that TED consistently improves the generalisation of the base RL algorithm in all tasks.
Revised: Our results show that TED consistently improves the generalisation of the base RL algorithm in all tasks.

Original: TED also shows lower variance across seeds compared to baselines, making it more reliable.
Revised: TED also shows lower variance across seeds compared to baselines, making it more reliable.

Original: 5. G ENERALISATION TO T ASK -I RRELEVANT V ARIABLES
Revised: 5. G ENERALISATION TO T ASK -I RRELEVANT V ARIABLES

Original: To demonstrate generalisation to unseen values of task-irrelevant environment variables, we use continuous control tasks from the DeepMind Control Suite (DMC) (Tunyasuvunakool et al., 2020) and
Revised: To demonstrate generalisation to unseen values of task-irrelevant environment variables, we use continuous control tasks from the DeepMind Control Suite (DMC) (Tunyasuvunakool et al., 2020) and

Original: Panda Gym (Gallou´edec et al., 2021) as simulations of robotics tasks.
Revised: Panda Gym (Gallou ´ edec et al., 2021) as simulations of robotics tasks.

Original: We adapt DMC wrappers from the Distracting Control Suite (Stone et al., 2021) to add colour distractors to the observations as irrelevant factors of variation.
Revised: We adapt DMC wrappers from the Distracting Control Suite (Stone et al., 2021) to add colour distractors to the observations as irrelevant factors of variation.

Original: Experiment setup.
Revised: Experiment setup.

Original: We show results with RAD and SVEA as two different base algorithms for continuous control.
Revised: We show results with RAD and SVEA as two different base algorithms for continuous control.

Original: We use the cartpole swingup and walker walk tasks from DMC with SVEA as the base algorithm.
Revised: We use the cartpole swingup and walker walk tasks from DMC with SVEA as the base algorithm.

Original: For RAD, we use the easier finger spin task instead of walker walk since RAD was unable to learn an optimal policy on the walker walk task due to the difficulty of the task with the colour distractors.
Revised: For RAD, we use the easier finger spin task instead of walker walk since RAD was unable to learn an optimal policy on the walker walk task due to the difficulty of the task with the colour distractors.

Original: We also use the Reach task with dense rewards from Panda Gym where the agent receives a reward of the negative of its distance to the goal.
Revised: We also use the Reach task with dense rewards from Panda Gym where the agent receives a reward of the negative of its distance to the goal.

Original: We train each algorithm on a fixed set of colours by varying the RGB colour values within a small bounded region of the original value.
Revised: We train each algorithm on a fixed set of colours by varying the RGB colour values within a small bounded region of the original value.

Original: We test on a set of colours of the same size that were not observed during training.
Revised: We test on a set of colours of the same size that were not observed during training.

Original: Following the original setup of the base RL algorithms, the encoder for SVEA has 11 convolutional layers and RAD hasconvolutional layers.
Revised: Following the original setup of the base RL algorithms, the encoder for SVEA has 11 convolutional layers and RAD hasconvolutional layers.

Original: Full implementation details are described in Appendix B.
Revised: Full implementation details are described in Appendix B.

Original: Baselines.
Revised: Baselines.

Original: We compare to the chosen base algorithm for each task to demonstrate the performance improvement achieved by TED.
Revised: We compare to the chosen base algorithm for each task to demonstrate the performance improvement achieved by TED.

Original: For further comparison, we also compare our method to representative baselines from each category of RL generalisation method discussed in Section 2.1.
Revised: For further comparison, we also compare our method to representative baselines from each category of RL generalisation method discussed in Section 2.1.

Original: Our data augmentation baseline is DrQ (Yarats et al., 2021).
Revised: Our data augmentation baseline is DrQ (Yarats et al., 2021).

Original: We use DBC (Zhang et al., 2021) as a baseline method that learns invariant representations.
Revised: We use DBC (Zhang et al., 2021) as a baseline method that learns invariant representations.

Original: We also compare with CURL (Laskin et al., 2020b) as a state-of-the-art contrastive auxiliary task.
Revised: We also compare with CURL (Laskin et al., 2020b) as a state-of-the-art contrastive auxiliary task.

Original: Finally, we include the base algorithm with domain randomisation, shown as { base algorithm } -DR in the figures, as a privileged baseline that is trained on both the ‘train’ and ‘test’ colours together, demonstrating that when the test variations are known a priori, it is often less sample efficient to use domain randomisation and sometimes unable to learn the optimal policy.
Revised: Finally, we include the base algorithm with domain randomisation, shown as { base algorithm } -DR in the figures, as a privileged baseline that is trained on both the ‘train’ and ‘test’ colours together, demonstrating that when the test variations are known a priori, it is often less sample efficient to use domain randomisation and sometimes unable to learn the optimal policy.

Original: All baselines use the same size encoder as the TED base algorithm on a given task.
Revised: All baselines use the same size encoder as the TED base algorithm on a given task.

Original: For all baselines, we tuned hyperparameters by grid search and report the best performing ones.
Revised: For all baselines, we tuned hyperparameters by grid search and report the best performing ones.

Original: Disentanglement metric.
Revised: Disentanglement metric.

Original: To evaluate the learned representations, we measure the disentanglement using the disentanglement metric proposed by Higgins et al. (2017a).
Revised: To evaluate the learned representations, we measure the disentanglement using the disentanglement metric proposed by Higgins et al. (2017a).

Original: The metric measures disentanglement using pairs of images with one factor of variation fixed to the same value in both images and the other factors randomised.
Revised: The metric measures disentanglement using pairs of images with one factor of variation fixed to the same value in both images and the other factors randomised.

Original: For example, the background colour could be held fixed while all other colours and variables are randomly assigned.
Revised: For example, the background colour could be held fixed while all other colours and variables are randomly assigned.

Original: We use pairs of frame stacks corresponding to o t instead of individual images to allow the policy to extract velocity information. 
Revised: We use pairs of frame stacks corresponding to o t instead of individual images to allow the policy to extract velocity information. [Figure] [Figure] RAD­TED

Original: One sample for the classifier is calculated for a batch B of observation pairs, given by:
Revised: RAD RAD­DR CURL DBC

Revised: DrQ

Revised: [Figure]

Revised: One sample for the classifier is calculated for a batch B of observation pairs, given by:

Original: [Equation]
Revised: [Equation]

Original: where o ( b ) t and o ( b ) ˜ t have the same fixed factor.
Revised: where o ( b ) t and o ( b ) ˜ t have the same fixed factor.

Original: We then train a linear classifier that takes z diff as one input and predicts which factor was held fixed across a batch of inputs.The accuracy of this [Figure] [Figure]
Revised: We then train a linear classifier that takes z diff as one input and predicts which factor was held fixed across a batch of inputs.

Original: [Figure] r e t u r n s [Figure] classifier is the disentanglement score.
Revised: The accuracy of this classifier is the disentanglement score.

Original: The intuition is that a there will be low variance in the features corresponding to the fixed factor in a disentangled representation so the classifier will be highly accurate.
Revised: The intuition is that a there will be low variance in the features corresponding to the fixed factor in a disentangled representation so the classifier will be highly accurate.

Original: Although independence of the factors does not strictly hold in practice in many
Revised: Although independence of the factors does not strictly hold in practice in many

Original: RL environments, it does hold when generating the observations for the metric because the value of each factor is assigned randomly, so we can fairly assess the disentanglement.
Revised: RL environments, it does hold when generating the observations for the metric because the value of each factor is assigned randomly, so we can fairly assess the disentanglement.

Original: A score of 1 . 0 is the maximum for a fully disentangled representation.
Revised: A score of 1 . 0 is the maximum for a fully disentangled representation.

Original: cartpole finger panda swingup spin reach

Original: RAD-TED 0.

Original: RAD 0.

Original: RAD-DR 0.

Original: CURL 0.

Original: DBC 0.

Original: DrQ 0.

Original: (a) Tasks

Original: [Table]

Original: (shown in Figure 2) cartpole walker panda swingup walk reach

Original: SVEA-TED 0.

Original: SVEA 0.

Original: SVEA-DR 0.

Original: CURL 0.

Original: DBC 0.

Original: DrQ 0.

Original: [Figure]

Original: [Figure]

Original: Results.
Revised: Results.

Original: The results for RAD as the base algorithm are shown in Figure 2, and the results for
Revised: The results for RAD as the base algorithm are shown in Figure 2, and the results for

Original: SVEA are shown in Figure 3.
Revised: SVEA are shown in Figure 3.

Original: The figures show that TED improves the generalisation of both
Revised: The figures show that TED improves the generalisation of both

Original: RAD and SVEA across all tasks.
Revised: RAD and SVEA across all tasks.

Original: In many tasks, TED is able to generalise instantly without any reduction in performance at all, while baselines fail to generalise immediately and, in some cases, are unable to recover after overfitting to the training colours.
Revised: In many tasks, TED is able to generalise instantly without any reduction in performance at all, while baselines fail to generalise immediately and, in some cases, are unable to recover after overfitting to the training colours.

Original: In other tasks, e.g. see Figure 2a, TED experiences some reduction in performance but is able to recover more quickly than the base algorithm and other baselines.
Revised: In other tasks, e.g. see Figure 2a, TED experiences some reduction in performance but is able to recover more quickly than the base algorithm and other baselines.

Original: Even though TED increases the disentanglement score of the base algorithm, shown in Figure 4, there can still be an initial drop in performance on the training environment because the RL policy may still unnecessarily put some weight on the colour features, however the disentangled structure allows faster recovery.
Revised: Even though TED increases the disentanglement score of the base algorithm, shown in Figure 4, there can still be an initial drop in performance on the training environment because the RL policy may still unnecessarily put some weight on the colour features, however the [Table] [Figure] [Table]

Revised: [Table]

Original:  TED achieves higher returns than the privileged domain randomisation baseline in many tasks (e.g. Figure 2), which assumes the test colours are known a priori, because it can be difficult to learn an optimal policy withsuch a larger set of colour distractors.
Revised: [Figure] [Table] disentangled structure allows faster recovery. TED achieves higher returns than the privileged domain randomisation baseline in many tasks (e.g. Figure 2), which assumes the test colours are known a priori, because it can be difficult to learn an optimal policy with a larger set of colour distractors.

Original: The poor training performance of some of the baselines, particularly DBC, is due to the difficulty of learningan optimal policy with the colour distractors as it increases the size of the state space.
Revised: The poor training performance of some of the baselines, particularly DBC, is due to the difficulty of learning with the colour distractors as it increases the size of the state space.

Original: DBC aims to maintain performance on difficult distractors resulting in reduced performance on ‘easy’ distractors.
Revised: DBC aims to maintain performance on difficult distractors resulting in reduced performance on ‘easy’ distractors.

Original: 5. G ENERALISATION TO BOTH T ASK -R ELEVANT AND I RRELEVANT V ARIABLES
Revised: 5. G ENERALISATION TO BOTH T ASK -R ELEVANT AND I RRELEVANT V ARIABLES

Original: We use the Procgen generalisation benchmark (Cobbe et al., 2020) to demonstrate generalisation to unseen values of both task-relevant and irrelevant variables together.
Revised: We use the Procgen generalisation benchmark (Cobbe et al., 2020) to demonstrate generalisation to unseen values of both task-relevant and irrelevant variables together.

Original: Procgen is a set of discretecontrol tasks that uses procedural generation to determine the layout, objects, entities, background, colours and other game details for each level of a game.
Revised: Procgen is a set of discretecontrol tasks that uses procedural generation to determine the layout, objects, entities, background, colours and other game details for each level of a game.

Original: This produces many factors of variation across the game levels, some of which are relevant to solving the game and others are irrelevant distractors.
Revised: This produces many factors of variation across the game levels, some of which are relevant to solving the game and others are irrelevant distractors.

Original: Due to the nature of procedural task generation with Procgen, we are unable to fix the factors of variation to calculate the disentanglement scores for these tasks.
Revised: Due to the nature of procedural task generation with Procgen, we are unable to fix the factors of variation to calculate the disentanglement scores for these tasks.

Original: Experiment setup.
Revised: Experiment setup.

Original: We use the coinrun and jumper Procgen environments.
Revised: We use the coinrun and jumper Procgen environments.

Original: We train on 100 levels of the hard difficulty, and test generalisation on 100 unseen hard levels.
Revised: We train on 100 levels of the hard difficulty, and test generalisation on 100 unseen hard levels.

Original: Following the setup of Cobbe et al.
Revised: Following the setup of Cobbe et al.

Original: (2020), we use PPO as the base algorithm for the Procgen tasks.
Revised: (2020), we use PPO as the base algorithm for the Procgen tasks.

Original: Results.
Revised: Results.

Original: The results are shown in Figure 5.
Revised: The results are shown in Figure 5.

Original: The results show that PPO-TED (ours) recovers more quickly than PPO on the unseen levels.
Revised: The results show that PPO-TED (ours) recovers more quickly than PPO on the unseen levels.

Original: PPO is unable to recover optimal performance on both tasks after overfitting to the training levels, but PPO-TED converges to higher performance on the unseen levels, similiar to that achieved in training.
Revised: PPO is unable to recover optimal performance on both tasks after overfitting to the training levels, but PPO-TED converges to higher performance on the unseen levels, similiar to that achieved in training.

Original: In the jumper environment, Figure 5b, PPO-TED also demonstrates better zero-generalisation performance (at the vertical dotted line) as well as faster adaptation to the unseen levels.
Revised: In the jumper environment, Figure 5b, PPO-TED also demonstrates better zero-generalisation performance (at the vertical dotted line) as well as faster adaptation to the unseen levels.

Original: 5. A BLATION S TUDIES
Revised: 5. A BLATION S TUDIES

Original: How does the loss coefficient affect performance?
Revised: How does the loss coefficient affect performance?

Original: Figure 6 shows how the choice of the loss coefficient α affects performance on the cartpole swingup task with colour distractors and RAD[Figure] [Figure] base algorithm.
Revised: Figure 6 shows how the choice of the loss coefficient α affects performance on the cartpole swingup task with colour distractors and RAD base algorithm.

Original: We found α = 100 to be optimal for this task.
Revised: We found α = 100 to be optimal for this task.

Original: A lower coefficient α = 50 does not prioritise the disentanglement enough reducing generalisation performance.
Revised: A lower coefficient α = 50 does not prioritise the disentanglement enough reducing generalisation performance.

Original: A higher coefficient α = 500 reduces the agents ability to recover because the agent is prioritising disentanglement too much over finding the optimal policy.
Revised: A higher coefficient α = 500 reduces the agents ability to recover because the agent is prioritising disentanglement too much over finding the optimal policy.

Original: Are non-temporal samples from different episodes only sufficient to improve generalisation?
Revised: Is one type of non-temporal sample sufficient to improve generalisation?

Original: TED requires two types of non-temporal samples: from the same episode, and from different episodes.
Revised: TED requires two types of non-temporal samples: from the same episode, and from different episodes.

Original: We explored a simplified version of TED using non-temporal samples from the same episode  only.
Revised: We explored two simplified versions of TED using non-temporal samples from the same episode only and nontemporal samples from different episodes only.

Original: This version of TED is more similar to Permutation Contrastive Learning (Hyv¨arinen & Morioka, 2017) adapted to an RL setting, so we name it RAD-PCL.

Original: The results in Figure 7 show that RAD-PCL improves the generalisation of RAD to an extent on the cartpole swingup task with colour distractors, but RAD-TED further improves the generalisation performance. 
Revised: The results in Figure 7 show that using a single type of non-temporal sample improves the generalisation of RAD to an extent, but RAD-TED further improves the generalisation performance. [Figure] [Figure] [Figure]

Original:  L IMITATIONS AND F UTURE W ORK
Revised:  L IMITATIONS AND F UTURE W ORK

Original: Our approach has some limitations that could be addressed in future work.
Revised: Our approach has some limitations that could be addressed in future work.

Original: Guarantees of disentanglement usually assume the factors of variation are independent and do not generalise to correlated factors (Tr ¨ auble et al., 2021).
Revised: Guarantees of disentanglement usually assume the factors of variation are independent and do not generalise to correlated factors (Tr ¨ auble et al., 2021).

Original: While we have shown that TED improves disentanglement in practice, RL observations generated by a learning agent do not have independent factors of variation in general.
Revised: While we have shown that TED improves disentanglement in practice, RL observations generated by a learning agent do not have independent factors of variation in general.

Original: We leave further exploration of how to relax this assumption for future work.
Revised: We leave further exploration of how to relax this assumption for future work.

Original: We introduced two types of non-temporal samples to disentangle features controlled by the agent (which form a non-stationary time-series) and episodic features (which form a stationary timeseries).
Revised: We introduced two types of non-temporal samples to disentangle features controlled by the agent (which form a non-stationary time-series) and episodic features (which form a stationary timeseries).

Original: While there exists a theoretical proof that disentanglement is possible with either nonstationary (Hyv ¨ arinen & Morioka, 2016) or stationary (Hyv ¨ arinen & Morioka, 2017) factors, it is still an open problem to identify an approach that guarantees disentanglement in a time-series that contains both stationary and non-stationary factors.
Revised: While there exists a theoretical proof that disentanglement is possible with either nonstationary (Hyv ¨ arinen & Morioka, 2016) or stationary (Hyv ¨ arinen & Morioka, 2017) factors, it is still an open problem to identify an approach that guarantees disentanglement in a time-series that contains both stationary and non-stationary factors.

Original: Finally, our approach introduces the TED loss coefficient α as a new hyperparameter which must be tuned to the task.
Revised: Finally, our approach introduces the TED loss coefficient α as a new hyperparameter which must be tuned to the task.

Original: Future work could consider automatically tuning α based on the current disentanglement score.
Revised: Future work could consider automatically tuning α based on the current disentanglement score.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we introduced TED, a disentangled representation learning auxiliary task for RL.
Revised: In this work, we introduced TED, a disentangled representation learning auxiliary task for RL.

Original: Our approach is the first to consider an auxiliary task based on disentangled representation learning for online RL.
Revised: Our approach is the first to consider an auxiliary task based on disentangled representation learning for online RL.

Original: TED can be used with existing algorithms and does not require a decoder.
Revised: TED can be used with existing algorithms and does not require a decoder.

Original: We demonstrated experimentally that TED improves generalisation of three different base algorithms by adapting with continued learning to previously unseen values of environment variables that are both task-relevant and irrelevant.
Revised: We demonstrated experimentally that TED improves generalisation of three different base algorithms by adapting with continued learning to previously unseen values of environment variables that are both task-relevant and irrelevant.

Original: We also adapted the notion of factors of variation to frame stacking in
Revised: We also adapted the notion of factors of variation to frame stacking in

Original: RL and used a disentanglement metric to show that TED improves the disentanglement of learned representations.
Revised: RL and used a disentanglement metric to show that TED improves the disentanglement of learned representations.

Original: TED is a step toward making RL algorithms more robust for real-world deployment and life-long learning as the agent is able to quickly recover when presented with previously unseen values of environment variables and continue learning without suffering from catastrophic forgetting. 
Revised: TED is a step toward making RL algorithms more robust for real-world deployment and life-long learning as the agent is able to quickly recover when presented with previously unseen values of environment variables and continue learning without suffering from catastrophic forgetting. A CKNOWLEDGEMENTS

Revised: This work was supported by the EPSRC Centre for Doctoral Training in Robotics and Autonomous

Revised: Systems, funded by the UK Engineering and Physical Sciences Research Council and the Edinburgh

Original: R EPRODUCIBILITY S TATEMENT
Revised: Centre for Robotics. This work was also supported by the Academy of Finland Flagship programme: Finnish Center for Artificial Intelligence FCAI. The authors wish to acknowledge the generous computational resources provided by the Aalto Science-IT project and the CSC – IT Center for

Original:  A public and open-source implementation of TED can be found at github.com/usr/repo [currently anonymised for double blind review].
Revised: Science, Finland. R EPRODUCIBILITY S TATEMENT A public and open-source implementation of TED can be found at github.com/uoe-agents/TED.

Original: Full details of the architecture and hyperparameter settings for each of our experiments, including implementation details of the disentanglement metric, can be found in Appendix B
Revised: Full details of the architecture and hyperparameter settings for each of our experiments, including implementation details of the disentanglement metric, can be found in Appendix B

Original: A E XTENDED B ACKGROUND
Revised: A E XTENDED B ACKGROUND

Original: In this section, we provide details of the RL algorithms that we use as the base algorithms for TED in our experiments.
Revised: In this section, we provide details of the RL algorithms that we use as the base algorithms for TED in our experiments.

Original: We use RAD (Laskin et al., 2020a) and SVEA (Hansen et al., 2021b) for the experiments in Section 5.1.
Revised: We use RAD (Laskin et al., 2020a) and SVEA (Hansen et al., 2021b) for the experiments in Section 5.1.

Original: Both base algorithms are extensions of the Soft Actor-Critic (SAC) algorithm (Haarnoja et al., 2018).
Revised: Both base algorithms are extensions of the Soft Actor-Critic (SAC) algorithm (Haarnoja et al., 2018).

Original: We use PPO (Schulman et al., 2017) as the base algorithm in
Revised: We use PPO (Schulman et al., 2017) as the base algorithm in

Original: Section 5.2.
Revised: Section 5.2.

Original: SAC.
Revised: SAC.

Original: Soft-Actor Critic (SAC) is an off-policy, actor-critic RL algorithm that learns a policy π to maximise the expected discounted future rewards and the entropy of the policy.
Revised: Soft-Actor Critic (SAC) is an off-policy, actor-critic RL algorithm that learns a policy π to maximise the expected discounted future rewards and the entropy of the policy.

Original: SAC uses transitions from the replay buffer D to train the critic Q : O ×
Revised: SAC uses transitions from the replay buffer D to train the critic Q : O ×

Original: A → R by minimising the loss:
Revised: A → R by minimising the loss:

Original: [Equation]
Revised: [Equation]

Original: The target value function ¯ V is estimated by:
Revised: The target value function ¯ V is estimated by:

Original: [Equation]
Revised: [Equation]

Original: where ¯ Q is the target Q network whose parameters are an exponential moving average of the corresponding Q network parameters.
Revised: where ¯ Q is the target Q network whose parameters are an exponential moving average of the corresponding Q network parameters.

Original: We maintain two Q networks, Q 1 and Q 2 , and use the minimum of the two networks for the updates.
Revised: We maintain two Q networks, Q 1 and Q 2 , and use the minimum of the two networks for the updates.

Original: The actor π is trained by minimising the loss:
Revised: The actor π is trained by minimising the loss:

Original: [Equation]
Revised: [Equation]

Original: RAD.
Revised: RAD.

Original: RAD adds data augmentations to the observations before the SAC network updates.
Revised: RAD adds data augmentations to the observations before the SAC network updates.

Original: We use image padding and random crop augmentations for the observations in each transition sampled from the replay buffer ( o t , o t +1 ) ∼ D .
Revised: We use image padding and random crop augmentations for the observations in each transition sampled from the replay buffer ( o t , o t +1 ) ∼ D .

Original: SVEA.
Revised: SVEA.

Original: SVEA stabilises Q-learning using a combination of augmented and unaugmented images with the updated loss:
Revised: SVEA stabilises Q-learning using a combination of augmented and unaugmented images with the updated loss:

Original: [Equation]
Revised: [Equation]

Original: The policy π is trained using only unaugmented images, using the standard loss in Equation 6.
Revised: The policy π is trained using only unaugmented images, using the standard loss in Equation 6.

Original: PPO.
Revised: PPO.

Original: Proximal Policy Optimisation (PPO) is an on-policy, actor-critic RL algorithm for continuous or discrete action spaces.
Revised: Proximal Policy Optimisation (PPO) is an on-policy, actor-critic RL algorithm for continuous or discrete action spaces.

Original: We use a discrete policy π ψ for the Procgen environments.
Revised: We use a discrete policy π ψ for the Procgen environments.

Original: PPO learns a policy by minimising a clipped loss over minibatches transitions:
Revised: PPO learns a policy by minimising a clipped loss over minibatches transitions:

Original: [Equation]
Revised: [Equation]

Original: where ρ t ( ψ ) is the ratio of the action probability under the new and old policies, and A t is the action advantage given by: A t = Q π ( o t , a t ) − V π ( o t ) .
Revised: where ρ t ( ψ ) is the ratio of the action probability under the new and old policies, and A t is the action advantage given by: A t = Q π ( o t , a t ) − V π ( o t ) .

Original: B I MPLEMENTATION D ETAILS
Revised: B I MPLEMENTATION D ETAILS

Original: In this section, we provide the implementation details for TED.
Revised: In this section, we provide the implementation details for TED.

Original: Our codebase is built on top of the publicly released DrQ PyTorch implementation by Yarats et al.
Revised: Our codebase is built on top of the publicly released DrQ PyTorch implementation by Yarats et al.

Original: We adapt the codebase to implement the base RL algorithms as well as the TED auxiliary task.
Revised: We adapt the codebase to implement the base RL algorithms as well as the TED auxiliary task.

Original: A public and open-source implementation of TED can be found in the supplementary material and at github.com/usr/repo [currently anonymised for double blind review].
Revised: A public and opensource implementation of TED can be found in the supplementary material and at github.com/uoeagents/TED.

Original: B. RAD AND SVEA I MPLEMENTATION D ETAILS
Revised: B. RAD AND SVEA I MPLEMENTATION D ETAILS

Original: Encoder architecture.
Revised: Encoder architecture.

Original: We use the same encoder architecture as Yarats et al.
Revised: We use the same encoder architecture as Yarats et al.

Original: The encoder weights are shared between the actor π and critic Q .
Revised: The encoder weights are shared between the actor π and critic Q .

Original: For the RAD experiments, the encoder consists of 4 convolutional layers following the original RAD implementation (Laskin et al., 2020a).
Revised: For the RAD experiments, the encoder consists of 4 convolutional layers following the original RAD implementation (Laskin et al., 2020a).

Original: For
Revised: For

Original: SVEA experiments, we use 11 convolution layers following the original SVEA paper (Hansen et al.,
Revised: SVEA experiments, we use 11 convolution layers following the original SVEA paper (Hansen et al.,

Original: 2021b).
Revised: 2021b).

Original: Baselines follow the same encoder size as the base RL algorithm for TED in each experiment.
Revised: Baselines follow the same encoder size as the base RL algorithm for TED in each experiment.

Original: Each convolutional layer has a 3 × 3 kernel size and 32 channels.
Revised: Each convolutional layer has a 3 × 3 kernel size and 32 channels.

Original: The first layer has a stride of 2, all other layers have a stride of 1.
Revised: The first layer has a stride of 2, all other layers have a stride of 1.

Original: There is a ReLU activation between each convolutional layer.
Revised: There is a ReLU activation between each convolutional layer.

Original: The convolutional layers are followed by a trunk network with a linear layer, layer normalisation, and finally a tanh activation.
Revised: The convolutional layers are followed by a trunk network with a linear layer, layer normalisation, and finally a tanh activation.

Original: Actor and critic architecture.
Revised: Actor and critic architecture.

Original: The actor π and critic Q networks are both 2-layer MLPs with a hidden dimension of 1024.
Revised: The actor π and critic Q networks are both 2-layer MLPs with a hidden dimension of 1024.

Original: We apply ReLU activations after each layer except the last layer.
Revised: We apply ReLU activations after each layer except the last layer.

Original: TED architecture.
Revised: TED architecture.

Original: The TED classifier is implemented with the following parameters: k 1 , k 2 , ¯ k , b , and ¯ b are vectors of length 50; and c is a scalar.
Revised: The TED classifier is implemented with the following parameters: k 1 , k 2 , ¯ k , b , and ¯ b are vectors of length 50; and c is a scalar.

Original: The output of the classifier is defined in Equation 1.
Revised: The output of the classifier is defined in Equation 1.

Original: Hyperparameters.
Revised: Hyperparameters.

Original: Table 1 shows the value of the TED loss coefficient α for each task.
Revised: Table 1 shows the value of the TED loss coefficient α for each task.

Original: All other hyperparameters are provided in Table 2.
Revised: All other hyperparameters are provided in Table 2.

Original: Unless specified otherwise, we use the same hyperparameter settings for both RAD-TED and SVEA-TED algorithms.
Revised: Unless specified otherwise, we use the same hyperparameter settings for both RAD-TED and SVEA-TED algorithms.

Original: [Table]
Revised: [Table]

Original: [Table]

Original: [Table]

Original: [Table]
Revised: [Table]

Original: Disentanglement metric.
Revised: Disentanglement metric.

Original: To calculate the disentanglement metric described in Section 5.1, we collect a batch of 10,000 samples.
Revised: To calculate the disentanglement metric described in Section 5.1, we collect a batch of 10,000 samples.

Original: To create each sample, we use 32 pairs of observations with the same fixed factor to evaluate Equation 3 with B = 32 .
Revised: To create each sample, we use 32 pairs of observations with the same fixed factor to evaluate Equation 3 with B = 32 .

Original: The batch of samples is split into 8, training samples to train the classifier, and 2,000 testing samples to calculate the accuracy of the classifier as the disentanglement score.
Revised: The batch of samples is split into 8, training samples to train the classifier, and 2,000 testing samples to calculate the accuracy of the classifier as the disentanglement score.

Original: We use a Scikit-learn logistic regression classifier with L regularisation and the saga solver.
Revised: We use a Scikit-learn logistic regression classifier with L regularisation and the saga solver.

Original: B. PPO I MPLEMENTATION D ETAILS
Revised: B. PPO I MPLEMENTATION D ETAILS

Original: We follow the PPO architecture and hyperparameters used in the Procgen benchmark (Cobbe et al., 2020).
Revised: We follow the PPO architecture and hyperparameters used in the Procgen benchmark (Cobbe et al., 2020).

Original: The encoder uses the IMPALA (Espeholt et al., 2018) architecture.
Revised: The encoder uses the IMPALA (Espeholt et al., 2018) architecture.

Original: The hyperparameters are shown in Table 3.
Revised: The hyperparameters are shown in Table 3.

Original: [Table]
Revised: [Table]

Revised: Value

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: Entropy bonus

Original: C E NVIRONMENT I MAGES
Revised: [Table] range [Table] Learning rate

Original:  In Figure 8, we provide images of example observations for each of the DMC and Panda Gym training and testing environments used in our experiments to visualise the generalisation challenge.
Revised: [Table] [Table] C E NVIRONMENT I MAGES In Figure 8, we provide images of example observations for each of the DMC and Panda Gym training and testing environments used in our experiments to visualise the generalisation challenge.

Original: A video is also provided in the supplementary materials.
Revised: A video is also provided in the supplementary materials.

Original: [Figure]
Revised: [Figure]

Original: In Figure 9, we provide images of example observations for each of the Procgen environments used in our experiments.
Revised: In Figure 9, we provide images of example observations for each of the Procgen environments used in our experiments.

Original: [Figure]
Revised: [Figure]