Temporal Disentanglement of Representations for Improved Generalisation in Reinforcement Learning
Authors: Mhairi Dunion
Trevor McInroe
Kevin Sebastian Luck
Josiah P. Hanna
Stefano V Albrecht
Venue: ICLR 2023 poster (2023)
Keywords: Reinforcement Learning
Representation Learning
Disentanglement

Abstract: Reinforcement Learning (RL) agents are often unable to generalise well to environment variations in the state space that were not observed during training. 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. The changed pixels can lead to drastic changes in the agent's latent representation of the image, causing the learned policy to fail. To learn more robust representations, we introduce TEmporal Disentanglement (TED), a self-supervised auxiliary task that leads to disentangled image representations exploiting the sequential nature of RL observations. 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. Since TED enforces a disentangled structure of the representation, our experiments also show 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).

Published as a conference paper at ICLR 2023
TEMPORAL DISENTANGLEMENT OF REPRESENTATIONS FOR IMPROVED GENERALISATION IN REINFORCEMENT LEARNING
Mhairi Dunion
Trevor McInroe University of Edinburgh
Kevin Sebastian Luck
University of Edinburgh
Aalto University kevin.s www.kaaltoeaaaooooacaaooo
mhairi.. dunionfed. ac.uk
wwwncoeeed..ccmmmmmm ac.uk
Josiah P. Hanna
Stefano V. Albrechht University of Edinburgh
University of Wisconsin Madison jphanna@cs, wisc.
s s.albrechh@edd @ed. ac uk
ABSTRACT
Reinforcement Learning (RL) agents are often unable to generalise well to environment variations in the state space that were not observed during training. This issue is especially problematic for image-based RL, where a change in just one variable, such the background colour, change many pixels in the image. The changed pixels can lead to drastic changes the agent's s latent representation the image, causing the learned policy fail. To learn more robust representations, we introduce TEmporal Disentanglement (TED), a self-supervised auxiliary task that leads disentangled image representations exploiting the sequential nature of RL observations. 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- Since TED enforces a disentangled structure of the representation, our experiments also show 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 variables that affect the optimal policy (e.g. goal positions).
representation learning methods.
INTRODUCTION
Real-world environments are often not static and deterministic, but can be subject changes, both incremental or with sudden effects (Luck et al., 2017). Reinforcement Learning (RL) algorithms need to be robust these changes and adapt quickly. Moreover, since many real-world robotics applications rely on images as inputs (Vecerik et al. 2021 Hämäläinen et al., 2019: Chebotar et al.) 2019), RL agents need learn robust representations of images that remain useful after change in the environment. For example, simple change in lighting conditions can change the perceived colour of object, but this should not affect the agent's 's ability to perform a task.
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 al. 2018). The issue is especially problematic for image-based based RL, where change in one environment variable can mean the agent is presented with very different set of pixels for which trained RL policies are often no longer optimal. 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). In practice, this often results in agents needing to adapt their policy after change only one variable. We show experimentally that agents often cannot recover the optimal policy after the environment changes because is too difficult to 'undo' the overfitting.
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). However, in practice, we may be unaware what variations an agent might see in the future. Even if all possible future variations are known, training this full set is often sample inefficient and may

---

Published as a conference paper at ICLR 2023
result in sub-optimal policy as the agent learns to compensate for many different values. It can also be impractical to generate all variations during training due to limitations on laboratory setups or in simulation suites. 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.| 2.22). Some approaches aim to learn representation that is invariant image distractors (Zhang et al. 2020 2021) to improve generalisation to unseen values of only variables irrelevant to the optimal policy. These approaches do not enforce a structure on the remaining task-relevant variables in the representation for generalisation to unseen values of relevant variables.
promising direction towards robust RL is learn a disentangled representation of observations. 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). 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. 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. Higgins et al. (2017b) show that disentangled representation improves generalisation in RL. However, this approach requires a trained or hard-coded policy collect independent and identically distributed (i.i.d.) data pre-train
3-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. data alone (Locatello et al. 2019).
We introduce a self-supervised auxiliary task for learning disentangled representations for the robust encoding of images, which we call TEmporal Disentanglement (TED). Unlike previous work, TED can be implemented with only minimal changes to existing RL algorithms and allows for lifelong learning of disentangled representations. In contrast to Higgins et al. (2017b). our approach uses the non-i.i.d. temporal data from consecutive timesteps in RL to learn the disentangled representation online. Note that TED does not require decoder which lessens computational costs. We provide experimental results across a variety of tasks from the DeepMind Control Suite (Tunyasuvunakool cct al. 2020), Panda Gym (Gallouédee et al. 2021) and Proegen (Cobbe et al. 2020) environments. For our experiments, we train on 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. Our results demonstrate that TED improves generalisation of a variety of base RL algorithms on unseen environment variables that are relevant or irrelevant the task, while state-of-the-ar baselines that achieve equally good training performance still fail to adapt and, in some cases, are unable to recover after overfitting to the training environment. We also evaluate 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.
RELATED WORK
2.1
GENERALISATION IN IMAGE-BASED REINFORCEMENT LEARNING
Image augmentation. Image augmentation artificially increases the size of the dataset by adding
image perturbations improve robustness of representations. [askin al. (2020a) apply a variety
of image augmentation techniques such as translation, cutouts and cropping; Yarats et al. (2021) average over multiple augmentations: [Hansen & Wang (02T) maximise mutual information between the representations of augmented and non non-augmenter augmented images; and Hansen et al. (202T) use both augmented and non-augmented images to stabilise valuee estimation. However, image augmentation approaches to generalisation can still fail when the agent experiences stronger types of variation after training (Kirk et al., 2022). In our experiments, we show that TED can used alongside image augmentation techniques to further improve generalisation while benefiting from the augmentation.
Learning invariant representations. Invariance techniques aim to learn a representation that ignores distractors in the image. Zhang et al. (2020) use causal inference techniques assuming a block MDP structure: Zhang et al. (2021) use a bisimulation metric; and Li et al. (2021) use domain adversarial optimisation to learn representation invariant distractors. These approaches all aim to generalise to unseen values of irrelevant variables, e.g. background colours. In contrast, TED uses disentanglement to enforce a structured representation that applies to both relevant and irrelevant variables.
2

---

Published as a conference paper at ICLR 2023
Encoding inductive biases. Some approaches an auxiliary task to encode inductive biases. Laskin et al. (2220)) learn a representation to maximise similarity between different augmentations of the same observation; Mazoure et al. (2200) maximise similarity between observations at successive timesteps; and Agarwal et al. ) enforce a structured representation for task relevant variables using policy similarity metries. These approaches are based on enforeing similarity constraint between pairs of observations to learn informative features but they do not enforce any structure to the representation, whereas TED encourages a disentangled structure due the form of the classifier without similarity constraint. van den Oord et al. (2018) learn representations that capture information predictive of the future, [Schwarzer et al., (2021) | pretrain an encoder and fine-tune on task specific data, and Jaderberg et al. (2017) uses combination of multiple auxiliary tasks for representation learning, but these approaches also do not require a disentangled structure. To encode the inductive bias of disentanglement, HHggiins et al. (2017)) train B-VAR offline using data from a pre pre-trained or hardcoded agent. In contrast, our approach uses the temporal structure of data available in RL to learn a disentangled representation online with the RL policy.
2.2 DISENTANGLED REPRESENTATIONS
Unsupervised learning. Many variations of the Variational Autoencoder (VAE) (Kingma & Welling 2014) aim to improve disentanglement of the learned representation, such as the V-V (Higgins) ct al. 201 7 (Burgess et al. 2017) and the Factor-VAE (Kim & Mnih 2018 (). Locatello et al. (2019) prove that it is theoretically impossible to learn disentangled representations from i.i.d. data alone. To bypass the impossibility result, many recent approaches extend the 3-VAE to use some form of supervision. (Shu et al. (2020) provide weak supervision using a labelled grouping of images, while Locatello et al. (2020) generate pairs of images with limited factors of variation between pairs. In contrast, our approach exploits non-i.i.d. temporal observations available in RL to learn a disentangled representation without labelling or artificially generating images.
Independent component analysis. Hyvärinen & Morioka (2016) introduce Time-Contrastive Learning to learn a disentangled representation from time-series data with non-stationary factors. Hyvärinen & Morioka (2017) disentangle time-series data with stationary factors, introducing Permutation Contrastive Learning (PCL) to train classifier to discriminate between temporal and non-temporal inputs. We consider RL episodes back-to-back as time-series. 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. 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.
PRELIMINARIES
We assume the environment is a fully-observable Markov Decision Process (MDP), defined as the tuple M = (S,A, P, 7), where S is a of states, A is a set of actions, P S S A [0, 1] is the state-transition function, R S x A -> IR is the reward function, and 7 E [0, 1) is the discount factor. An RL agent chooses an action a, E AA time based on its current state s, E S and its policy ar (S²). The agent then transitions to the next state according to the state-transition probability P(s a,), and receives a reward, r, R(s;, at). The goal of an RL agent is learn a policy ı to maximise the expected discounted cumulative rewards, max, Ep,z Erlol-
In this work, the agent's observation o, O at time image pixels, a high- -igendimensi representation of the true underlying environment state s, E S. We assume the environment has a factored state representation s;, where each factor corresponds to an environment variable. The components of the state vector s; are the unobserved ground truth factors of variation. We consider one observation to be a stack of consecutive frames to ensure all features of 8,, such as velocities, can be extracted from o,, We assume o, is generated by an invertible, non-linear transformation of the state factors S,, h S + O, such that o, = h(s)). The aim of disentanglement is to learn a representation z, E Z that recovers the independent ground truth factors of from the observations o, by approximating the inverse of the transformation, such that z, = f(o,). To simplify notation, we will usually denote 2, f(o:) z(o:). We will use z to denote the i-th component of the vector
3

---

Published as a conference paper at ICLR 2023
Reinforcement Learning
loss
*; [o
Temporal sample (Nf) m).1)
1411
Repley buffer
Non-temporal different opisade sample (#E.aç)
Target encoder of So(of)
Classifier
8o(25, x5)
TED loss
Non-tamporal same episode sample (g))
Figure 1: TED architecture: The classifier is trained to discriminate between temporal and non-tenne tee-ttemrrall samples to encourage the encoder to disentangle the temporal structure in the image observations. The '/'" indicates that the gradient flow is stopped, E {e, e') and f E {t 1, t") depending the observation being processed.
4 TEMPORAL DISENTANGLEMENT
We introduce TEmporal Disentanglement (TED) auxiliary task for representation learning with RL algorithms. Our goal to improve generalisation to both relevant and irrelevant features when their testing variation is unknown priori. TED aims to disentangle the factors of variation that generate an observation o,, such as background colour and the trajectory of an object across the frame stack. We will provide an overview of architecture in Section then discuss the details the TED auxiliary task in Section
4.1 ARCHITECTURE OVERVIEW
The high-level architecture for TED with generic RL algorithm is depicted in Figure 1 TED is designed to encourage the encoder fe O + Z to recover the temporal structure determined by observations consecutive timesteps, O, and O,-1, such that classifier gø discriminate between temporal and non-temporal pairs of observation encodings. We structure the classifier to compare each feature in the representation Z, separately to encourage the encoder to disentangle the temporal structure into the ground truth factors of sy that generated the observation. We simultaneously perform dimensionality reduction such that dim(S) S dim(2) < dim
Z R can
TED requires a batch of N transitions B {(o,, 0t+1): 1 originating from various episodes. The classifier loss is applied as an auxiliary loss to any base RL algorithm. For off off-policy algorithms, the batch can be sampled from the replay buffer B ~ DD following the sampling procedure for the base RL algorithm. For on-policy base algorithms, the batch can created using multiple parallel environments to ensure transitions from different episodes. which is common in practice, for which B D the subsequent explanations. Image augmentations can be applied to the observations where required by the base algorithm, i.c. TED uses the same augmented images as the base algorithm. Both the TED classification loss and the relevant RI. loss are used to learn the encoder parameters 0. For training stability, also for the next observation at a given timestep Oz+1, where Ø' transitions are used for RL, but non-temporal pairs (o,, 0z) are created where f * ttt + to train the classifier, which we describe in more detail in the next section.
target encoder (He et al., 2020 Laskin al. 2020b) før O Z T)0. Only temporal
70' +
4.2 TEMPORAL DISENTANGLEMENT CLASSIFIER
We use the TED auxiliary loss to train classifier to discriminate between temporal and non- temporal pairs of observation encodings. TED encourages the encoder to learn representation that uncovers the temporal structure the data to enable distinguishing whether two observations occurred consecutively not.
Classifier inputs. For each batch of transitions B ~ D, we create three types of observation-pua batches for classifier input: temporal samples X, 2) different episode, non-temporal samples
X', and 3) same episode, non-temporal samples X" A single transition contains (of. of+1) )) E B
where is the episode index corresponding to transition. The temporal sample x; (of, o(41)
4

---

Published as a conference paper at ICLR 2023
Algorithm 1 TED update step
Input: batch of transitions B = (.... (of, of+1);..)>
Create batch Zees of representations for each observation in B: = =f((oi)
Create batch Zuext obs of representations for each next observation in B: zit1 for (0 for each transition in B do
Create temporal sample + (2; X + -XUxx
Sample प.. ~ Zeett such that e' e
Create non-tenn non-temporal different episode sample x t- (zf, z ).x X' X' Sample of. ~ D such that t" {t, + 1), and get representation zi, (0)))
Create non-temporal same episode sample x (zi, zir), X X" Ux! end for
for each sample x & [X,X', do
Classifier prediction y (x) (see Equation D
Calculate binary cross-entropy loss (x, (see Equation 2
end for
Calculate average loss for the batch LIED +- mean (x))
Backpropogate loss to update encoder parameters @ and classifier parameters &
Update target encoder parameters 0' = -0' + -]
Output: Loss LTED and updated parameters Φ, 0, and 0'
consists of two consecutive observations within a given episode e. Due to the episodic nature of RL, we use two types of non-temporal samples. The non-temporal sample x = (of, of consists of non-consecutive timesteps from different episodes, where of B such that e' * e. These samples encourage the encoder, fe. to learn a representation, z(o)), that disentangles episodic features that chosen randomly at the start of an episode, such colours, this is sufficient to distinguish x from the temporal sample X,. The second non-temporal sample x" = (o oi consists of non-consecutive timesteps from the same episode where o D. 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, and x".
Classification objective. Given a batch of samples {X,X', x',X'" X"), a logistic regression classifier is trained discriminate between the corresponding representation of temporal samples, x; X, and non-temporal samples, x, E X' and x"' E X To learn a disentangled representation, we use regression function of the form proposed by HHvärinen Morioka (2017):
y(xx) = Bo (x x2)
Σ (x/=((/) + (x)) + 641
b1)² +c
for all E (X, X', X"), where is the dimensionality of the representation z(of), = (x/. xf). and Φ (k1, k2, b, k, b, e) are the classifier parameters to be trained simultaneously with the encoder parameters 0.
Temporal samples, X, E X, are given a classification label / - 1, and non-temporal samples, x² E X' and x E X", are given a classification label = 0. The classifier is trained using the cross-entropy loss for binary classification for all x, E {X, [X,'X""
(IT 1) (2] log a ()(x,)) (1 1) log(1 (y((x)))
where a is the sigmoid function. Since there are two non-temporal samples, x E X' and x" E X", for each transition in the batch B but only one temporal, x, E X, the positive temporal samples are weighted by 2. The coefficient is hyperparameter to be tuned to the task. It is used to scale up the classification auxiliary loss to ensure not dwarfed by the RL loss and to prioritise representation learning over policy learning at the start of training while the factors of variation, St, are independent. The structure of the classifier = go (x] x²) (Equation !) ensures each feature z' is considered separately to the other features مززز 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 the representation so that the temporal structure of each feature can be determined independently of the other features. The second term of Equation I approximmtes the marginal log -df. Due to this structure, we expect the encoder to learn a disentangled representation and the classifier to approximate the distribution of the samples (Hyvärinen & Morioka 2017). The
5

---

Published as a conference paper at ICLR 2023
pseudocode for an update step is provided in Algorithm [ which is performed for every update of the base RL algorithm. We train the encoder with both CTED and the RL loss. It is important to note that the TED auxiliary loss requires only creating temporal and non-temporal pairs of representations to train the classifier. The classifier is not required for execution so it can be discarded after training.
EXPERIMENTAL RESULTS
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 of the base RI. algorithm after change in environment variables if generalisation is not instant. 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. We use training environment with a subset of some of the environment variables (such as colours), and evaluate generalisation test environment with unseen values of the variables. Images of the train and test environments are provided in Appendix [C] We show results on a variety of different tasks with RAD (Laskin al. 2020a), SVEA (Hansen al. [021] and PPO (Schulman al., (017) 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. Our results show that TED consistently improves the generalisation of the base RL algorithm in all tasks. TED also shows lower variance across seeds compared baselines, making more reliable.
5.1 GENERALISATION TO AAR-IERRELV VARIABLES
To demonstrate generalisation to unseen values of task-irrelevant environment variables. we continuous control tasks from the DeepMind Control Suite (DMC) (Tunyasuvunakool et al. 2020) and Panda Gym (Gallouédec et al.] 2021) simulations of robotics tasks. We adapt DMC wrappers from the Distracting Control Suite Stone et al. 202T) to add colour distractors to the observations as irrelevant factors of variation.
Experiment setup. We show results with RAD and SVEA as two different base algorithms for continuous control. We use the sartpole_swinuup and walker_wald tasks from DMC with SVEA as the base algorithm. For RAD, we use the easier finger-spin .spnn task instead of walker_waid since RAD was unable learn an optimal policy on the walker_wall walk task due the difficulty of the task with the colour distractors. 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. 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. We test on a set of colours of the same size that were not observed during training. Following the original setup of the base RL algorithms, the encoder for SVEA has 11 convolutional layers and RAD has 4 convolutional layers. Full implementation details are described in Appendix B
Baselines. We compare to the base RL algorithm for each task to demonstrate the performance improvement achieved by TED. For further comparison, we also compare our method to representative baselines from each category of RL generalisation method diseussed in Section 2. Our data augmentation baseline is DrQ (Yarats et al.) 2021). We use DBC (Zhang et al. (021) as a baseline method that learns invariant representations. We also compare with CURL (Laskin et al.) 2020b) as a state-of-the-arr contrastive auxiliary task. Finally, we include the base algorithm with domain randomisation, shown as (base algorithm3-DD -DRR in the figures, as privileged baseline that 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 learn the optimal policy. All baselines use the same size encoder as the TED base algorithm on a given task, so results for the same baseline differ slightly depending on the hase algorithm being compared to. For all baselines, we tuned hyperparameters by grid search and report the best performing ones.
Disentanglement metric. To evaluate the learned representations, we measure the disentanglement using the disentanglement metric proposed by Higgins et al. (2017a). 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. For example, the background colour could be held fixed while all other colours and variables are randomly assigned. We use pairs of frame stacks corresponding to o instead of individual images allow the policy to extract velocity information.

---

Published as a conference paper at ICLR 2023
-2.5
800
800
-6.0
Maan eval nours
800
-7.5
oval
-0.00
400
-2.8
15.0
200
200
0
8.
0.5
1.0
155
2.0
-20
8.
0.2
0.4
0.6
0.0
1.0
Triining timeoteps.
1e5
Triiing Timestepe
365
Traiiing timectipe
365
(a) cartpole swingup
(b) finger spin
(c) panda reach
RAD-TED
RAD
RAD-DR
CURL
DBC
DrQ
Figure 2: Generalisation
colours the vertical dotted line with RAD base algorithun. RAD-TED
(ours) recovers more quickly than baselines and achieves higher returns than domain randomisation (RAD-DR). Returns are the average 10 evaluation episodes, averaged over seeds, shaded region shows standard deviation.
800
800
rourts
800
-7.5
Maan eval
-0.00
400
-2.5
200
200
8.0
0.5
1.0
1.5
20.0
0.2
0.4
0.8
0.8
1.0
Training timoonapp.
1eB
Triining timiooap
Truining timesdips
165
(a) cartpole_swingup
(b) walker_walk
(c) panda_reach
SVEA-TED
SVEA
SVEA-DR
CURL
DBC
DrQ
Figure 3: Generalisation to unseen colours at the vertical dotted line. SVEA-TED (ours) recovers more quickly than baselines and achieves similar performance to domain randomisation but with fewer assumptions. Returns are the average of 10 evaluation episodes, averaged over 5 seeds, shaded region shows standard deviation.
One sample for the classifier is calculated for a batch B of observation pairs, given by:
Zdiftt Σ Iz(o, (b) z(0 (0)
where of and (6) have the same fixed factor. We then train linear classifier that takes Zaif as one input and predicts which factor was held fixed across a batch of inputs. The accuracy of this classifier is the disentanglement score. The intuition that there will be low variance in the features corresponding the fixed factor in disentangled representation the classifier will be highly accurate. Although independence of the factors does not strictly hold in practice in many 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. A score of 1.0 is the maximum for fully disentangled representation.
The results for RAD as the base algorithm are shown in Figure 2 and the results for SVEA are shown in Figure 3 The figures show that TED improves the generalisation of both RAD and SVEA across all tasks. In many tasks, TED achieves zero-shot generalisation. In other tasks, c.g. Figure 2a TED experiences some reduction in performance but able to recover more quickly than the base algorithm and other baselines. Surprisingly, the baselines that achieve equally good training performance as TED fail to adapt and, in some cases, are unable to recover after overfitting to the training colours. Even though TED increases the disentanglement score of the base algorithm, shown in Table there can still be an initial drop in performance on the test environment because the RL policy may still unnecessarily put some weight the colour features, but the 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 known priori, because
Results.
7

---

Published as a conference paper at ICLR 2023
caripole swingup
finger spin
panda reach
cartpole swingup
walker walk
panda
reach
RAD-TED
0.99
0.79
0.95
SVEA-THD
0.64
0.79
0.90
RAD
0.88
0.56
0.83
SVEA
0.53
0.71
0.83
RAD-DR
0.67
0.53
0.49
SVEA-DR
0.58
0.62
0.72
CURL
0.87
0.89
0,91
CURL
0.92
0.77
0.65
DBC
0.65
0,46
0.58
DBC
0.61
0.34
DrQ
0.73
0.59
0.91
DrQ
0.86
0.66
0.60
(a) RAD base algorithm (Figure
(b) SVEA base algorithm (Figure
Table I: Disentanglement scores at the end training before changing to the test environment.
it be difficult to learn optimal policy with such large set of colour distractors. The poor training performance of some of the baselines, particularly DBC, is due to the difficulty of learning optimal policy with the colour distractors it increases the size of the state space. DBC aims to maintain performance on difficult distractors resulting in reduced performance on 'easy" distractors.
5.2
GENERALISATION BOTH SESERREEEVNNT AND IRRELEVANT VARIABLES
We the Procgen generalisation benchmark (Cobbe et al. 2020 demonstrate generalisation to unseen values of both task-relevant and irrelevant variables together. Proegen is a set discrete- control tasks that uses procedural generation determine the layout, objects, entities, background, colours and other game details for each level of a game. This produces many factors of variation of which relevant solving the game and others are irrelevant
across the game levels, distractors. Due to the nature of procedural task generation with Procgen, we are unable to fix the factors of variation calculate the disentanglement scores for these tasks.
Experiment setup. We use the coinrun and jumper Procgen environments. We train on 100 levels of the hard difficulty, and test generalisation on 100 unseen hard levels. Following the setup of Cobbe et al. (2220), we use PPO as the base algorithm for the Procgen tasks.
The results are shown in Figure The results show that PPO-TED (ours) recovers more quickly than PPO the unseen levels. PPO is unable to recover optimal performance on both tasks after overfitting the training levels, but PPO-TED converges to higher performance on the unseen levels, similiar to that achieved in training. In the jumper environment, Figure [b] PPO-TED also demonstrates better zero-shot generalisation performance than PPO (at the vertical dotted line).
Results.
Moai wall raurr relrre
PRO-TED
PPO
8.
8.0
Trainin
1eB
(a) coinrun
(b) jumper
Figure Generalisation to unseen levels the vertical dotted line. PPO-TED (ours) recovers more quickly than PPO on the unseen levels. Returns are the average of evaluation episodes, averaged over seeds, shaded region shows standard deviation. The graphs show the 10-poim rolling average for readability.
5.3 ABLATION STUDIES
How does the loss coefficient affect performance? Figure 5a shows how the choice of the loss coefficient a affects performance. We found a 100 to be optimal for this task, and the results show some robustness prioritise disentanglement enough reducing generalisation performance. A higher coefficient reduces performance because the agent is prioritising disentanglement too much over the optimal policy.
200 still achieves good performance. A lower coefficient
50 does not
8

---

Published as a conference paper at ICLR 2023
1000
3330
RAC TED Sameia
800
WW
800
000
Meant aoai raturnok
600
800
Vour 1911
600
eve
400
200
200
200
80
Trairing timesteps
Truning limesteps
105
0.0
05
3.0
1.5
20
2.5
20
25
40
(a) Comparison of different TED
(c) Modified TED using a standard linear classifier. Figure 5: Ablations for RAD-TED on the cartpole swingup task with colour distractors, averaged over 5 seeds.
(b) TED using same episode or different episode samples only.
loss coefficients.
Is one type of non-temporal sample sufficient to improve generalisation? TED requires two types of non-temporal samples: from the same episode, and from different episodes. We explored two simplified versions of TED using non-tew con-temporal samples from the same episode only and non- temporal samples from different episodes only. The results in Figure 5b 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.
How important is the disentangled structure of the TED classifier? We compared TED to a simplified version that uses a standard linear classifier instead of the disentangled structure of the TED classifier. The results in Figure 5c show that while the linear classifier improves the generalisation performance of RAD, it does not generalise as well RAD-TED.
LIMITATIONS AND FUTURE WORK
Our approach has some limitations that could be addressed in future work. Guarantees of disentanglement usually assume the factors of variation are independent and do not generalise to correlated factors (Träuble et al. 202T]. 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. We leave further exploration of how relax this assumption for future work.
We introduced two types of non- non-temporal samples to disentangle features controlled by the agent (which form a non-stationary time-series) and episodic features (which form a stationary time- series). While there exists a theoretical proof that disentanglement is possible with either non- stationary (Hyvärinen & Morioka 2016) or stationary (Hyvärinen & Morioka (017) factors, it is still an open problem to identify approach that guarantees disentanglement in time-series that contains both stationary and non non-stationary factors.
Our approach introduces the TED loss coefficient a as new hyperparameter which must be tuned to the task. Future work could consider automatically tuning a based on the current disentanglement score. Finally, our temporal samples are constructed from the current timestep t and next timestep ++T 1 to allow TED to be easily be added to existing algorithms. Future work could explore extending the horizon of the temporal sample I. + l with k 1.
CONCLUSION
In this work, we introduced TED, auxiliary task for learning disentangled representations in RL. Our approach is the first to consider an auxiliary task based on disentangled representation learning for online RL. TED can be used with existing algorithms and does not require a decoder. We demonstrated experimentally that TED improves generalisation of three different RL base algorithms by adapting with continued learning to previously unseen values of environment variables that are both task-relevant and irrelevant. We also adapted the notion of factors of variation to frame stacking in RI. and used disentanglement metric to show that TED improves the disentanglement of learned representations. TED is a step toward making RI. 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 while reducing catastrophic forgetting.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
This work was supported by the EPSRC Centre for Doctoral Training in Robotics and Autonomous Systems, funded by the UK Engineering and Physical Sciences Research Council and the Edinburgh Centre for Robotics. This work was also supported by the Academy of Finland Flagship programme: Fimnish Center for Artificial Intelligence FCAL The authors wish to acknowledge the generous computational resources provided by the Aalto Science-IT project and the CSC IT Center for Science, Finland.
REPRODUCIBILITY STATEMENT
public and open-source implementation of TED can be found at github. htwadddabob..ooouuuugggggongooo wwtp://oo/oooooen///Eaaa Full details of the architecture and hyperparameter settings for each of our experiments, including implementation details of the disentanglement metric, can found Appendix B
REFERENCES
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization reinforcement learning. In International Conference on Learning Representations, 2021.
Y. Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. In IEEE Transactions on Pattern Analysis and Machine Intelligence, volume 35, pp. 1798-1828, 2013.
Christopher P. Burgess, Irina Higgins, Arka Pal, Loie Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerehner. Understanding disentangling in 8-vae. In 31st Conference on Neural Information Processing Systems (NIPS 2017), 2017.
Yevgen Chebotar, Ankur Handa, Viktor Makoviychuk, Miles Macklin, Jan Issac, Nathan Ratliff, and Dieter Fox. Closing the sim-to-real loop: Adapting simulation randomization with real world experience. In 2019 International Conference Robotics and Automation (ICRA), pp. 8973-8979. IEEE, 2019.
Karl Cobbe, Oleg Klimov, Chris Hesse, Tachoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In Proceedings of the 36th International Conference Machine Learning (ICML 2019), volume 97 of Proceedings of Machine Learning Research, pp. 1282-1289. PMLR, 2019.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning (ICML 2020), volume 119 of Proceedings of Machine Learning Research, pp. PMLR, 2020.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, Shane Legg, and Koray Kavukeuoglu. IMPALA: scalable distributed deep-rl with importance weighted actor-lcarner architectures. In Proceedings of the 35th International Conference on Machine Learning (ICMI. 2018), volume 80 of Proceedings of Machine Learning Research, pp. 1407-1416. PMLR, 2018.
Quentin Gallouédec, Nicolas Cazin, Emmanuel Dellandréa, and Liming Chen. panda-gym: Open- Source Conditoond Environments for Robotic Learning. 4th Robot Learning Workshop: Self-Supervised and Lifelong Learning at NeurlPS, 2021.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with stochastic actor. In Proceedings of the 35th International Conference on Machine Learning (ICML 2018), volume 80 of Proceedings of Machine Learning Research, pp. 1861-1870. PMLR, 2018
10

---

Published as a conference paper at ICLR 2023
Aleksi Hämäläinen, Karol Arndt, Ali Ghadirzadeh, and Ville Kyrki. Affordance learning for end-to- end visuomotor robot control. In 2019 JEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 1781-1788. IEEE, 2019.
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. In International Conference Robotics and Automation, 2021.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. In Conference on Neural Information Processing Systems, 2021.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp- 9726-9735, 2020.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. 3-vae: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017a.
Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. DARLA: Improving zero-shot transfer in reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning (ICML 2017), volume 70 of Proceedings of Machine Learning Research, pp. 1480-1490. PMLR, 2017b.
Aapo Hyvärinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. In Advances in Neural Information Processing Systems (NIPS 2016), 2016.
Aapo Hyvärinen and Hiroshi Morioka. Nonlinear ica of temporally dependent stationary sources. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS), volume 54, 2017.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations, 2017.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In Proceedings of the 35th International Conference on Machine Learning (ICML 2018), volume 80 of Proceedings of Machine Learning Research, pp. 2649-2653 PMLR, 2018.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations, 2014.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. survey of generalisation in deep reinforcement learning. arXiv:2111.09794 2022.
Charline Le Lan, Stephen Tu, Adam Oberman, Rishabh Agarwal, and Mare G.Bellemare. On the generalization of representations in reinforcement learning. In International Conference Artificial Intelligence and Statistics (AISTATS22), 2022.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. In 34th Conference on Neural Information Processing Systems (NeurIPS 2020), 2020a.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representations for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning (ICMI. 2020), volume 1 119 of Proceedings of Machine Learning Research, pp. 5639-5650. PMLR, 2020b.
Bonnie Li, Vincent François Lavet, Thang Doan, and Joelle Pineau. Domain adversarial reinforcement learning. 10X 2021.
11

---

Published as a conference paper at ICLR 2023
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In Proceedings of the 36th International Conference Machine Learning (ICMI. 2019), Proceedings of Machine Learning Research, PMLR, 2019.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michaël Tschannen. Weakly-supervised disentanglement without compromises. In Proceedings of the 37th International Conference on Machine Learning (ICML 2020), volume 119 of Proceedings of Machine Learning Research, pp. 6348-6359) PMLR, 2020.
Kevin Sebastian Luck, Joseph Campbell, Michael Andrew Jansen, Daniel M. Aukes, and Heni Ben Amor. From the lab to the desert: Fast prototyping and learning of robot locomotion. In Robotics: Science and Systems, 2017.
Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. In 34th Conference on Neural Information Processing Systems (NeurIPS 2020), 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms.
2017.
Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R Devon Hjelm, Philip Bachman, and Aaron C Courville. Pretraining representations for data- efficient reinforcement learning. In 35th Conference on Neural Information Processing Systems (NeurIPS 202/). 1), 2021.
Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disentanglement with guarantees. In International Conference on Learning Representations, 2020.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite- challenging benchmark for reinforcement learning from pixels. WxXiv2200.0022:) 2021.
Frederik Träuble, Elliot Creager, Niki Kilbertus. Anirudh Goyal, Francesco Locatello, Bernhard Schölkopf, and Stefan Bauer. Is independence all you need? on the generalization of representations learned from correlated data. In Proceedings of the 38 th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research. PMLR, 2021.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Frez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for continuous control. Software Impacts, 6: 6:100022, 2020.
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv:18800 1807 03748, 2018
Mel Vecerik, lean-Baptiste Regli, Oleg Sushkov, David Barker, Rugile Pevceviciute, Thomas Rothörl, Raia Hadsell, Lourdes Agapito, and Jonathan Scholz. S3k: Self-supervised semantic keypoints for robotic manipulation via multi-view consistency. In Conference on Robot Learning, pp. 449-460. PMLR. 2021.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In Proceedings of the 37th International Conference Machine Learning (ICML 2020), volume 119 of Proceedings of Machine Learning Research, pp. 11214-11224. PMLR, 2020.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. A study on overfitting in deep reinforcement learning. arXiv: 1804 884406683, 2018.
12

---

Published as a conference paper at ICLR 2023
A EXTENDED BACKGROUND
In this section, we provide details of the RL algorithms that we use as the base algorithms for TED in experiments. We use RAD (Laskin al., 2020a) and SVEA (Hansen et al.] 2021) for the experiments in Section 5.1 Both base algorithms are extensions of the Soft Actor-Critic (SAC) algorithm (Haarnoja al. 2018). We use PPO (Schulman al. 2017) as the base algorithm in Section 5.2
SAC. Soft-Actor Critic (SAC) is an off-policy, actor-critic RL. algorithm that learns policy a to maximise the expected discounted future rewards and the entropy of the policy. SAC uses transitions from the replay buffer D to train the critic Q : x A + by minimising the loss:
Lo
(o a.o.a...))...)) [Q(o,,, ar) - Tt 10e--V(oo++i (V(o++)))
The target value function V is estimated by:
((o+1)) = -a min ((00+1,,(00411)) - OSAC log 77 (()+((a++1))) 0
(5)
where Q is the target Q network whose parameters are exponential moving average of the corresponding Q network parameters. We maintain two Q networks, Q1 and Q2, and use the minimum of the two networks for the updates. The actor ı is trained by minimising the loss:
L = EO
Fais asscc log((((a))oo) - min Qi(or, ax) 1=,,2
RAD. RAD adds data augmentations the observations before the SAC network updates. We image padding and random crop augmentations for the observations in each transition sampled from the replay buffer (o,, 0c+1) D.
SVEA. SVEA stabilises Q-learning using combination of augmented and unaugmented images with the updated loss:
SVEA aSvEALo (or, a, 0c+1) + ( "uu at, oc+1) L
(7)
The policy æ is trained using only unaugmented images, using the standard loss in Equation 0
PPO. Proximal Policy Optimisation (PPO) is an on-policy, actor-critic RL algorithm for continuous or discrete action spaces. We use a discrete policy T. for the Procgen environments. PPO learns policy by minimising clipped loss over minibatches of transitions:
CLIP E 1-)()= [min(p:A,, clip(p:(v). cl c)4,)l
(8)
where pe(+) is the ratio of the action probability under the new and old policies, and A₁ is the action advantage given by: A, = Q" (or, 0x) V* (ot).
B IMPLEMENTATION DETAILS
In this section, we provide the implementation details for TED. Our codebase is built on top of the publicly released DrQ PyTorch implementation by Yarats al. (202I), and uses the official implementation of SVEA (Hansen et al. 2021). We adapt the codebase to implement the base RL algorithms well the TED auxiliary task. A public and open-soure source implementation of TED is available at
B. RAD AND SVEA IMPLEMENTATION DETAILS
Encoder architecture. We use the same encoder architecture as Yarats et al. (2021). The encoder weights are shared between the actor and critic Q. For the RAD experiments, the encoder consists of 4 convolutional layers following the original RAD implementation (Laskin et al. 2020a).. For SVEA experiments, convolution layers following the original SVEA paper (Hansen et al.
13

---

Published as a conference paper at ICLR 2023
Baselines follow the same encoder size the base RL algorithm for TED in each experiment. Each convolutional layer has 3 x kernel size and 32 channels. The first layer has a stride of 2, all other layers have stride of 1. There is ReLU activation between each convolutional layer. The convolutional layers are followed by a trunk network with linear layer, layer normalisation, and finally tanh activation.
2021).
Both RAD and SVEA use the critic loss Lo to update the encoder parameters. So our implementation of TED requires both the critic loss Lg and TED loss LTDD update the encoder together. In practice, this can be done by adding the two losses Lenc Lo + LTED and backpropagating the encoder loss CECC update the encoder parameters.
Actor and critic architecture. The actor and critic Q networks are both 2-layer MLPs with hidden dimension of 1024. We apply ReLU activations after each layer except the last layer.
TED architecture. The TED classifier is implemented with the following parameters: k1, k2, k, b, and are vectors the same size as the latent representation; and e is scalar. The output the classifier is defined in Equation
Hyperparameters. Table 2 shows the value of the TED loss coefficient a for each task. All other hyperparameters are provided in Table Unless specified otherwise, use the same hyperparameter settings for both RAD-TED and SVEA-TED algorithms.
Environment cartpole swingup
Base algorithm
Value of a
RAD
100
cartpole_swingup
SVEA
0.5
finger_spin
RAD
25
walker_wall
SVEA
1
panda_reach
RAD
200
randa_reach
SVEA
0.5
Table 2: TED loss coefficient
Hyperparameter name
Value
Replay buffer capacity
100000
Initial steps
1000
Stacked frames
Action repeat
for finger 4 otherwise
Batch size
128
Discount factor
0.99
Optimizer
Adam
Learning (actor, critic and encoder)
1e-3
Target soft-update rate 7
0.01
Actor update frequency
Actor log stddev bounds
[-10, 2]
Latent representation dimension
50 for RAD, 100 for SVEA
Image size
(84, 84)
Image pad
Initial temperature
0.1
Table 3: Hyperparameter values for both RAD-TED and SVEA-TED.
Disentanglement metric. To calculate the disentanglement metric described in Section 5.1 we collect batch of 10,000 samples. To create each sample, we use 32 pairs of observations with the same fixed factor to evaluate Equation 3 with = 32. The batch of samples is split into 8,000 training samples train the classifier, and 2,000 testing samples calculate the accuracy of the classifier as the disentanglement score. We use Scikit-learn logistic regression classifier with L1 regularisation and the saga solver.
14

---

Published as a conference paper at ICLR 2023
B.2 PPO IMPLEMENTATION DETAILS
We follow the PPO architecture and hyperparameters used in the Procgen benchmark (Cobbe et al. 2020 The encoder the IMPALA (Espeholt al., 2018 architecture. PPO augmented with TED by adding the loss terms & = Lpo + LTDD and backpropagating the total loss L with shared optimiser. The hyperparameters are shown Table 4
Hyperparameter name
Value
Image size
(64, 64)
Discount factor 7
0.999
GAE
0.95
Timesteps per rollout
250
Epochs per rollout
# Minibatches rollout
Entropy bonus
0.01
PPO clip range Learning rate # Workers
0.2
5e-4
# Environments per worker LSTM?
No
Frame stack?
No
TED coefficient
1 for coinrun, 0.25 for jumper
Table 4: Hyperparameter values PPO-TED
C ENVIRONMENT IMAGES
In Figure we provide images example observations for each of the DMC and Panda Gym training and testing environments used in our experiments to visualise the generalisation challenge.
In Figure 7 we provide images of example observations for each of the Procgen environments used in our experiments.
15

---

Published as a conference paper at ICLR 2023
(a) cartpole-swingup
7
(b) finger-spin
(c) walker_wall
(d) renda
Figure 6: Example observations for cach task with colour distractors used in our experiments (before image pre-processing to reduce the size). The images the left are examples from the training environment, and the images on the right are examples from the testing environment.
(a) coinrun
(b) jumper
Example observations for each Proegen environment used in experiments. The images the left are examples from the training environment, the images the right are examples from the testing environment.
Figure 7:
16


Reviews:
[{"review_id": "ws_N8_wl_L", "reviewer": "Reviewer_QVh3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Thank you for your submission. I mentioned the main points in each section, and I elaborated on the points in the Detailed Comments.\n\n## Summary\n\nThe paper proposes TED, an auxiliary task for deep RL agents that encourages learning disentangled representations. TED is shown to improve generalization across color variations of training domains, and is evaluated across different task suites and in combination with different \"base\" RL agents.", "summary_of_the_review": "## Concerns\n\nMy main concern is the impact/relevance of the results for the broader community. Modulo any concerns raised by other reviewers, I will raise my score to accept (8) if:\n* the authors improve the clarity of the paper (the three points above)\n* can demonstrate that, in the setting considered, TED's disentanglement component makes a big contribution to generalization.\n\n## Detailed Comments\n\n**Strength 1**. I would like the paper to emphasize this a bit more. It helps make a stronger argument for why we should care about the evaluation setting being considered in the paper. Plus, I think this is quite a surprising fact, because we should care about generalization to tasks that are slightly outside our training set, but we rarely do, and as researchers we may inadvertently discarding/deprioritizing research ideas that fail to outperform existing baselines in training, whereas they could be better than the baselines in ways that training performance does not measure, but that we could also care about (i.e., in this paper, adaptation to new tasks).\n\n**Strength 2**. This could also be investigated further. There's a chance that conclusions about disentangled representations generalizing better carry over to other settings, so the conclusions in this work can be useful for other work too.\n\n**Weakness 1**. This is my only concern about the algorithmic contribution. The method is highly dependent on the negative-example distribution \"making sense\" or making the classification problem challenging enough to drive representation learning. A failure case would be if the time index $t$ can easily be inferred from $o_t$. In that case a representation that discards everything else and retains only information about `t` can classify optimally. This can therefore be an issue with recurrent representations that retain some amount of information about `t` in them. TED may not immediately carry over to settings where we can't design a suitable negative example distribution, but some empirical conclusions of this paper can still be useful more broadly (namely the conclusions from Strengths 1 and 2).\n\n**Weakness 2**. This is related to my request for more emphasis to Strengths 1 and 2, and it's about making explicit some of the conclusions that go beyond the settings where Weakness 2 is not a problem.\n\n**Strength 4 and Weakness 3**. The choice of evaluation (train for N steps, then change tasks and train for N more steps). This is quite different from many previous works on generalization, in particular, previous work that has discussed zero-shot generalization (ZSG). I think this paper made the right decision to look at adaptation because ZSG would not give a full picture of TED's usefulness. However, certain changes could have helped compare the results against the ZSG in previous work. For example, reporting ZSG performances in control tasks at 100k and 500k.\n\n**Comparison to CURL**. I recommend comparing TED to CURL in the presentation because many components are similar. The comparison will help emphasize the main impactful idea in TED, which I believe to be the classification loss that encourages disentanglement. In my opinion the choice of negative example distribution makes sense and is adequate, but is domain dependent (Weakness 1).\n\n**Investigating good ZSG cases**. I noticed that in Figure 3 SVEA-TED is generalizing quite well _zero-shot_. It seems essentially unaffected by the change at the vertical dotted line, and continues to improve the agent in some cases. If this is not an artifact of smoothing, I think it can be quite an interesting finding. For future work, I suggest the authors reflect and investigate what it is about the SVEA+domain combination that is causing this effect, and whether it can be incorporated into the auxiliary task to bear in other domains. Figure 3(b) is a really interesting example of this phenomenon, where the baselines degrade a lot at 1e6 steps, but SVEA-TED and SVEA-DR continue to improve \"almost as though\" there has been no change in the training distribution.\n\n**Relation to theoretical results**. The paper draws some connections to negative results about disentanglement, but only in Sections 1, 2.2 and 6. Without additional context or specific knowledge of the results, the comments in this paper are vague to the reader. Moreover, the \"decomposition\" in terms of stationary and non-stationary factors does not seem to have any bearing in the discussion of the results. It would be good to see this clarified, especially if the authors have intuition on the subject that is guiding how they designed or interpret TED.\n\n**Related work**. It may be worth citing [CURL](https://arxiv.org/pdf/2004.04136.pdf) or [MoCo](https://openaccess.thecvf.com/content_CVPR_2020/papers/He_Momentum_Contrast_for_Unsupervised_Visual_Representation_Learning_CVPR_2020_paper.pdf) in reference to the target encoder.\nThere are some other papers with auxiliary tasks that can be relevant to current and future work, for example: [UNREAL](https://arxiv.org/pdf/1611.05397.pdf), [SimCore](https://proceedings.neurips.cc/paper/2019/file/2c048d74b3410237704eb7f93a10c9d7-Paper.pdf), [PBL](http://proceedings.mlr.press/v119/guo20g/guo20g.pdf) and [SPR](https://arxiv.org/pdf/2007.05929.pdf).\n\n**Figures**. I suggest placing the figures slightly below where they are first mentioned. For example, Figure 2 is a page before the results discussion. Also, should \"Figure 4\" be \"Table 4\"? Finally, the curve colors used may be difficult for colorblind readers to interpret. Please consider whether the color scheme in use is accessible or if it can be improved.\n\n**Hyperparameters** While selecting $\\alpha$ will often require domain-specific experimentation, I have found that a good rule of thumb for auxiliary tasks with RL is to make the scaled auxiliary loss of the same order of magnitude as the RL loss.", "strength_and_weaknesses": "## Strengths\n\n1. To me, one of the most interesting findings in the paper is that TED improves how fast the agent can adapt to new tasks, and how different the ability to adapt can be for auxiliary tasks that look equally good in terms of training performance.\n2. The design that incorporates disentanglement into the predictor is a core algorithmic strength.\n3. The empirical study in the paper has wide coverage in terms of algorithms and task suites.\n4. The evidence provided adequately supports the claims being made.\n\n## Weakness\n\n1. The temporal sampling is a limitation of this particular method. How limiting is a matter of opinion, but in my opinion it is a strong limitation. \n2. In my opinion the focus on demonstrating that TED outperforms previous baselines limits the impact of the paper.\n3. Some choices in the empirical study make the results harder to immediately compare to previous work.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "sctNPwCJCq", "reviewer": "Reviewer_rUnk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a representation learning method for image-based state observations for RL learning algorithms. This is proposed as an auxiliary task built around self-supervision between successive observations, with a distinct focus on disentangling the learned representations to differentiate between relevant and irrelevant aspects of the environment when learning a policy. The proposed temporal disentanglement (TED) augmentation is then shown to improve performance to a handful of selected RL algorithms. ", "summary_of_the_review": "I enjoyed this paper. I am intrigued by the apparent simplicity of the proposed TED auxiliary task. I am however worried about the limited scope of the variations tested on in the construction and execution of the experiments. I am also fairly concerned about whether the experimental evaluation was properly executed in comparison to baselines across the different domains. There are also concerns about the clarity of how the proposed approach is actually implemented. In light of these concerns I have currently rated this paper as a \"borderline reject\". I would be interested in raising my score if the authors were able to adequately clarify my questions and some of the concerns raised in my review. I am however not entirely confident that this can be done without substantial revision of the experiments since it's not clear that they were executed fairly. ", "strength_and_weaknesses": "## Strengths\n\n- The paper is well written and is very clear through the motivation and development of the proposed TED auxiliary task. The clarity falters somewhat in the experimental setup and results sections (specifically addressed in the next section). For the most part, I found the paper easy to understand and could follow the conceptual build-up arguing for the use of SSL techniques when considering representation learning for RL. \n- I appreciate the depth of discussion about the conceptual framework built through related work. This is a great example of good scholarship on the behalf of the authors. I feel confident about where to place the paper's significance as a result of the detailed overview and technical support derived from previous work. There are some missing references however. In particular, there has been a collection of papers looking into SSL within Deep RL in the recent few years that should be cited (more on this below).\n- The construction of TED as an auxiliary task is straightforward and seems like it would be straightforward to incorporate into future research. The major challenge is the engineering aspect of appropriately constructing the batch for SSL. It's promising that the authors intend to release their code to smooth over this barrier. \n- The TED augmented algorithms are quite thoroughly compared to relevant baselines on a variety of environments demonstrating the utility of the approach. It was really nice to see the proposed TED applied to different algorithms to demonstrate that the proposed approach is more flexible rather than taking advantage of unique qualities of a single approach/problem domain. The rigor of the experimental analysis is completed with an ablation study showing that the inclusion of the different components in the SSL batch help contribute to the overall performance.\n\n## Weaknesses\n- The deviations considered as unrelated/related between episodes are quite narrow. It seems that only color is considered where there are potentially more impactful deviations that are not included (e.g. viewing angle, perturbations in the dynamics, occluded aspects of the scene, etc.). Aside from Figure 1 and a limited collection in the Appendix, there are no visual demonstrations of what the environment variations are. This would be a great inclusion in Section 5.1 to solidify what exactly is meant by distractor and relevant/irrelevant variations. So much of the discussion and explanation of these terms (throughout the whole paper) relies perhaps too much on these specifics being completely understood by the reader.\n- The paper doesn't completely outline how $\\mathcal{L}_{TED}$ is integrated into the RL loss. It is mentioned that the representation and policy learning are performed hand-in-hand. In algorithm 1 there's no indication of what is done to accumulate the loss with the different samples $x\\in\\{X', X'', X'''\\}$ for every transition in B. \n- The clarity of the experimental evaluation is much worse than the preceding sections (specific details in the next section) which casts some doubt on the overall utility of the the approach. \n- Figure 1 could probably be better detailed in the text of Section 4.1. There are some missing points of notation (e.g. $z_{\\tilde{t}}^{\\tilde{e}}$ as well as an unclear differentiation between where the representation of the next observation $z_{t+1}^e$ is used in the RL algorithm. Is it assumed that the RL algorithm is multistep? Or is this an implied use of the next observation within the \"target\" of a TD-error update? \n- The disentanglement metric comes across as an afterthought and is not very clear how and where it's used to evaluate the performance. E.g. how significant is a score of 0 vs. a score of 1? Without this insight, the tables in \"Figure\" 4 are relatively meaningless...\n- There is inconsistency between experiments where the same algorithm is used in the same domain (e.g. CURL/DBC/DrQ on cartpole swingup and panda reach). There is no explanation of where the different experiments differ. Are the results of provided in \"Figure\" 4 a single run of these algorithms? As far as I can see, the only change between these experiments is the base algorithm that TED was applied to (RAD --> SVEA). Additionally, it seems that the experiments are set-up much differently since there are different time scales (and performances of these baselines) between Figure 2 and Figure 3. *This is a major concern.*\n- The ablations are only performed on one algorithm-TED combo. How generalizable are the insights seen in Section 5.3?\n- There is a claim that TED enables continual learning without catastrophic forgetting. There are no experiments to test this claim and no other discussion about this term/concept throughout the paper. \n\n### Missing papers \n\n```Mazoure, B., Tachet des Combes, R., Doan, T. L., Bachman, P., & Hjelm, R. D. (2020). Deep reinforcement and infomax learning. Advances in Neural Information Processing Systems, 33, 3686-3698.```\n\n```Schwarzer, M., Rajkumar, N., Noukhovitch, M., Anand, A., Charlin, L., Hjelm, R. D., ... & Courville, A. C. (2021). Pretraining representations for data-efficient reinforcement learning. Advances in Neural Information Processing Systems, 34, 12686-12699.```\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "WbBW_GocKr", "reviewer": "Reviewer_aCmj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes to improve the generalization of RL through a self-supervised auxiliary task which aims to learn disentangled representations. The auxiliary loss is used to train a classifier to discriminate between temporal and non-temporal pairs so that the encoder is encouraged to disentangle temporal structure in the observations. The proposed auxiliary loss can be easily combined to existing RL algorithms and experimental results on three benchmarks are reported.", "summary_of_the_review": "Generalization of RL algorithms is an important topic and the motivation of this paper is good. To solve the generalization problem of RL, the paper proposes a novel disentangled representation learning auxiliary task and experiments on three benchmarks demonstrate its effectiveness. The major concern is the completeness and quality of the experiments.", "strength_and_weaknesses": "Strength:\n1. Generalization of RL algorithms is an important research area, and this paper proposed a self-supervised auxiliary loss to achieve improved generalization in RL.\n2. The idea of combining the TED auxiliary loss with RL is novel.\n3. The paper is clearly presented and well-written.\n4. The proposed method is easy to implement and the experiments are conducted on different benchmarks.\n\nWeakness:\n1. TED requires non-temporal samples from both the same episode and different episodes. However, the necessity of using non-temporal samples from the same episode is not convincing. Why not just use samples from different episodes or mix the two sources during sampling? The ablation study also does not include experiments that use non-temporal samples from different episodes only.\n2. The parameter alpha seems to be a sensitive parameter. However, experiments are only implemented on the cartpole swingup task with three different values. More analysis should be added. Is the optimal value of alpha different in different tasks?\n3. In many experiments, the training does not seem to have converged yet. More training steps are suggested. Also, the number of DMC tasks are quite limited in this paper.\n4. In Figure 4, what if the TED loss is combined with CURL in the cartpole swingup task.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster