Task-Induced Representation Learning
Authors: Jun Yamada
Karl Pertsch
Anisha Gunjal
Joseph J Lim
Venue: ICLR 2022 Poster (2022)
Keywords: representation learning
reinforcement learning
transfer learning
visually complex observations

Abstract: In this work, we evaluate the effectiveness of representation learning approaches for decision making in visually complex environments. Representation learning is essential for effective reinforcement learning (RL) from high-dimensional in- puts. Unsupervised representation learning approaches based on reconstruction, prediction or contrastive learning have shown substantial learning efficiency gains. Yet, they have mostly been evaluated in clean laboratory or simulated settings. In contrast, real environments are visually complex and contain substantial amounts of clutter and distractors. Unsupervised representations will learn to model such distractors, potentially impairing the agent’s learning efficiency. In contrast, an alternative class of approaches, which we call task-induced representation learning, leverages task information such as rewards or demonstrations from prior tasks to focus on task-relevant parts of the scene and ignore distractors. We investi- gate the effectiveness of unsupervised and task-induced representation learning approaches on four visually complex environments, from Distracting DMControl to the CARLA driving simulator. For both, RL and imitation learning, we find that representation learning generally improves sample efficiency on unseen tasks even in visually complex scenes and that task-induced representations can double learning efficiency compared to unsupervised alternatives.

Published as a conference paper at ICLR 2022
TASK-INDUCED REPRESENTATION LEARNING
Jun Yamada': Karl Pertsch", Anisha Gunjal?, Joseph J. Lim University of Oxford, University of Southern California,
3 Korea Advanced Institute of Science and Technology. Naver Al Lab
ABSTRACT
In this work, we evaluate the effectiveness of representation learning approaches for decision making visually complex environments. Representation learning is essential for effective reinforcement learning (RL. from high-dimensional in- puts. Unsupervised representation learning approaches based on reconstruction, prediction or contrastive learning have shown substantial learning efficiency gains. Yet, they have mostly been evaluated in clean laboratory or simulated settings. In contrast, real environments are visually complex and contain substantial amounts of clutter and distractors. Unsupervised representations will learn to model such distractors, potentially impairing the agent's learning efficiency. In contrast, an alternative class of approaches, which we call fask-induced representation learning, leverages task information such as rewards or demonstrations from prior tasks to focus on task-relevant parts of the scene and ignore distractors. We investi- gate the effectiveness of unsupervised and task-induced representation learning approaches on four visually complex environments, from Distracting DMControl to the CARLA driving simulator. For both, RL and imitation learning, we find that representation learning generally improves sample efficiency on unseen tasks even in visually complex scenes and that task induced representations can double learning efficiency compared to unsupervised alternatives.
INTRODUCTION
The ability to compress sensory inputs into a compact representation is crucial for effective decision making. Deep reinforcement learning (RL) has enabled the learning of such representations via end-to-end policy training (Mnih et al., 2015; Levine al., 2016). However, learning representations from reward feedback requires many environment interactions. Instead, unsupervised objectives for representation learning (Lange and Riedmiller, 2010; Finn et al., 2016; Hafner et al., 2019; Oord et al., 2018; Chen et al., 2020) directly maximize lower bounds the mutual information between the sensory inputs and the learned representation. Empirically, such unsupervised objectives can accelerate the training of deep RL agents substantially (Laskin et al., 2020; Stooke et al., 2021; Yang and Nachum, 2021). However, prior works have evaluated such representation learning approaches in clean laboratory simulated settings (Finn al., 2016; Lee et al., 2020; Laskin al., 2020) where most of the perceived information is important for the task at hand. In contrast, realistic environments feature lots of task-irrelevant detail, such as clutter objects moving in the background. This has the potential to impair the performance of unsupervised representations for RL since they aim to model all information in the input equally.
The goal of this paper is to evaluate the effectiveness of representation learning approaches for decision making in such more realistic, visually complex environments with distractors and clutter. We evaluate two classes of representation learning approaches: (1) unsupervised representation learning based on reconstruction, prediction or contrastive objectives, and (2) approaches that leverage task information from prior tasks to shape the representation, which we will refer to as task-induce representation learning approaches. Such task information can often be available at no extra cost, e.g., in the form of prior task rewards if the data has been collected from RL training runs. Task-induced representations can leverage this task information to determine which features important in visually complex scenes, and which are distractors that can be ignored. We focus
Work done while at USC f AT Advisor at NAVER Lab Project website: elvrai com/tarp

---

Published as a conference paper at ICLR 2022
Prior Task Experience w/ Task Information (rewards, demos...)
Target Task
Task
Task 2
Task 3
Task 4
T
(ss r
S. a demo
(s, a, 2
(S. a demo
Visually Complex Scene
Unsupervised Representation Learning
π
Task-Induced
Pre-trained Representation
Downstream Policy
Representation Learning
Overview of our pipeline for evaluating representation learning in visually complex scenes: given a multi-task dataset of prior experience we pre-train representations using unsupervised objectives, such as prediction and contrastive learning, or task-induced approaches, which leverage task- inskiinformation from prior tasks to learn to focus on task-relevant aspects of the scene. We then evaluate the efficiency of the pre-trained representations for learning unseen tasks.
Figure 1:
our evaluation on a practical representation learning setting, shown in Figure I: given large offline dataset of experience collected across multiple prior tasks, we perform offline representation learning with different approaches. Then, we transfer the pre-trained representations for training a policy on a new task and compare efficiency training the policy from scratch.
To evaluate the effectiveness of different representation learning objectives, we compare common unsupervised approaches based reconstruction, prediction and contrastive learning four task- induced representation learning methods. We perform experiments four visually complex envi- ronments, ranging from scenes with distracting background videos the Distracting DMControl bench (Stone et al., 2021) to realistic distractors such as clouds, weather and urban scenery in the CARLA driving simulator (Dosovitskiy al., 2017). Across RL and imitation learning experiments we find that pre-trained representations accelerate downstream task learning, even in visually complex environments. Additionally, in our comparisons task-induced representation learning approaches achieve up to double the learning efficiency of unsupervised approaches by learning to ignore distractors and task-irrelevant visual details.
In summary, the contributions of our work are threefold: (1) we formalize the class of task-induced representation learning approaches which leverage task-information from prior tasks for shaping the leamned representations, (2) we empirically compare widely used unsupervised representation learning approaches and task-induced representation learning methods across four visually complex environments with numerous distractors, showing that task-induced representations more effective learning and (3) through analysis experiments we develop a set of best practices for task-induced representation learning.
lead to
RELATED WORK
Deep RL can train "end-to-end" policies that directly map raw visual inputs to action com- mands (Mnih et al., 2015; Kalashnikov al., 2018). To improve data efficiency in deep RL from high- dimensional observations, number of recent works have explored using data augmentation (Yarats al., 2021: Laskin al., 2021) and unsupervised representation learning techniques during RL training. The latter can be categorized into (1) reconstruction (Lange and Riedmiller, 2010; Finn al., 2016). (2) prediction (Hafner et al., 2019; Lee et al., 2020), and (3) contrastive learning (Laskin al., 2020: Stooke et al., 2021; Zhan et al., 2020; Yang and Nachum, 2021) approaches. While these works have shown improved sample efficiency, fundamentally they are unsupervised and therefore cannot decide which information in the input is relevant irrelevant certain task. We show that this leads to deteriorating performance in visually complex environments such as autonomous driving scenarios with lots of non-relevant information in the input observations.
Another line of work has used task information to shape representations that learn to ignore distractors. Specifically, Fu al. (2021) predict future rewards learn representations that focus only on task-
2

---

Published as a conference paper at ICLR 2022
relevant aspects of a scene. Zhang et al. (2021) leverage rewards as part of a bisimulation objective to learn representations without reconstruction. In this work, we provide a unifying framework for such task-induced representation learning approaches that encompasses the objectives of Fu et al. (2021) and Zhang et al. (2021) among others, and perform a systematic comparison to unsupervised representations in visually complex environments.
Closest to our work is the work Yang and Nachum (2021), which evaluates unsupervised and task- induced representation learning objectives for pre-training in multiple environments from the D4RL benchmark (Fu et al., 2020). Their work shows that representation pre-training can substantially improve downstream learning efficiency. While Yang and Nachum (2021)'s evaluation is performed in visually clean OpenAl gym environments (Brockman et al., 2016) without distractors, we focus our evaluation on more realistic, visually complex environments with distractors, which can pose additional challenges, especially for unsupervised representation learning approaches.
PROBLEM FORMULATION
Our goal is efficiently learn a policy ı that solves target task Tunget in MDP defined by tuple (S,A, T,RR R. p, 7) of states, actions, transition probabilities, rewards, initial state distribution, and discount factor. The policy is trained to maximize the discounted return Bar-st [E 7 R[s;, a:)]
We do not assume access to the underlying state of the MDP, but instead our policy receives a nigh-dimennsonal observation efficiency, we aim to learn an encoder d(x) which maps the input observation to a low-dimensional state representation that is input to the policy (()()). To learn this representation, we assume access to a dataset of past interactions D = {D....., Dx} collected across Ti:N E T tasks, with per-task datasets D, {ara; ae, (r+),-..) of state-action trajectories and optional reward annotation. The set of training tasks does not include the target task Tiageet E TI:N- We assume that the training data contains task information for prior tasks, e.g., in the form of step-wise reward annotations r; or by consisting of task demonstrations. Such task information can often be available at no extra cost, e.g., if the data was collected from prior training runs (Fu al., 2020; Çaglar Gülçehre al., 2020) or human teleoperation (Mandlekar et al., 2018; Cabi et al., 2019).
X in every step, e.g., an image observation. To improve training
We follow Stooke et al. (2021) and test all representation learning approaches in two phases: we first train the representation encoder Φ(x) from the offline dataset D, then transfer its parameters to the policy x (((x)) and optimize it on the downstream task, either freezing or finetuning the encoder parameters. This allows us to efficiently reuse data collected across prior tasks.
3.1 UNSUPERVISED REPRESENTATION LEARNING
Unsupervised representation learning approaches aim to learn the low-dimensional representation (xx), by maximizing lower bounds on the mutual information max; [xx (x)))
Generative Representation Learning. Maximizes the mutual information via generative model- ing, either of the current state (reconstruction) or of future states (prediction). A decoder (()(z)) trained by maximizing a lower bound on the data likelihood p(z) (Kingma and Welling, 2014; Rezende et al., 2014; Denton and Fergus, 2018; Hafner et al., 2019).
Contrastive Representation Learning. Avoids training complex generative models by using con- trastive objectives for maximizing the mutual information (Oord et al., 2018; 8; Laskin et al., 2020; Stooke et al., 2021). During training, muni-batches of positive and negative data pairs are assembled, where positive pairs can for example constitute consecutive frames of the same trajectory while negative pairs mix frames across trajectories. The representation is then trained with the InfoNCE objective (Gutmann and Hyvärinen, 2010; Oord et al., 2018), which is lower bound on I (2z p(x)).
4 TASK-INDUCED REPRESENTATION LEARNING
An alternative class of representation learning approaches uses task information from prior tasks as filter for which aspects of the environment interesting, and which others do not need to be
3

---

Published as a conference paper at ICLR 2022
TARP-V
TARP-CQL
TARP-Bisim
TARP-BC
Moune (al (((x))
Aactal (bxx)
e'60,, (((60))))))0)))) h((00))
(10)
CQL
P(a3
Bisim
Aadaa (b(x))
Alta (400)
Miaes 1
Kactal ((0))
Figure 2: Instantiations of our task-induced representation learning framework. Left to right: Representation learning via multi-task value prediction (TARP-V), via multi-task offline RL (TARP. CQL), via bisimulation (TARP-Bisim) and via multi-task imitation learning (right, TARP-BC).
modeled. This addresses fundamental problem of unsupervised representation learning approaches: by maximizing the mutual information between observations and representation they are trained to model every bit of information equally and thus can struggle in visually complex environments with lots of irrelevant details.
Formally, the state of an environment S can be divided into task-relevant and task-irrelevant or nuisance components S {S m. Sy The superscript indicates that this division is lask-dependent, since some components the state space are relevant for particular task T; but irrelevant for others. In visually complex environments, the state space is much smaller than all state information in the input observations (Zhaoping, 2006).
assume that [Susk] < |S] i... that the task-relevant part of
When choosing how much of this input information to model, end-to-end policy learning and unsupervised representation learning, form two ends of spectrum. End-to-end policy learning only models Sam.k for its training task, which is efficient, but only allows transfer from task task S lask C S i.e., if the task-relevant components of the target task are subset of those of the training task. In contrast, unsupervised representation learning models the full S = {Stakk S,, which allows for flexible transfer since S S S Y T T, but training be inefficient since the learned representation contains many nuisance variables.
In this work, we formalize the family of task-induced representation learning (TARP) approaches, which aim to combine the best of both worlds. Similar to end-to-end policy learning, task-induced representation learning approaches tations. Yet, by combining the task information from a wide range of prior tasks they learn a representation that combines the task-relevant components of all tasks in the training dataset D: STARP = Slisk U U.N. Thus, such representation allows for transfer to a wide range of unseen tasks for which Singe C STARP. Yet, the representation filters large part of the state space which is not relevant for any of the training tasks, allowing sample efficient learning on the target task.
task information learn compact, task represen-
4.1 APPROACHES FOR
REPRESENTATION LEARNING
We compare multiple approaches for laskeinduced representation learning which can leverage different forms task supervision. We focus on reward annotations and task demonstrations; but, future work can extend the TARP framework to other forms of task supervision like language commands or human preferences. All instantiations of TARP train a shared encoder network Φ(x) with separate task-supervision heads (see Figure 2).
Value Prediction (TARP-V). We can leverage reward annotations as task supervision, similar to Fu et al. (2021). by estimating the future discounted return of the data collection policy. Intuitively, representation that allows estimation of the value of state needs to include all task-relevant aspects of this state. We introduce separate value prediction heads N for each task and train the representation o(z) by minimizing the error in the predicted discounted return:
N
T
LTARP B(urn)DD
Σ :'=
Offline RL (TARP-CQL). Alternatively, we can learn a representation by training a policy to directly maximize the discounted reward of the training tasks. Since we aim to use offline training
4

---

Published as a conference paper at ICLR 2022
(a) Distracting Control
(b) ViZDoom
(c) Distracting MetaWorid
(d) CARLA
Figure 3: Environments with high visual complexity and substantial task irkeeerell detail for testing the learned representations. (a) Distracting DMControl (Stone et al., 2021) with randomly overlayed videos, (b) ViZDoom (Wydmuch et al., 2018) with diverse textures and enemy appearances, (c) Dis- tracting Meta World (Yu et al., 2020) with randomly overlayed videos, and (d) CARLA (Dosovitskiy et al., 2017) with realistic outdoor driving scenes and weather simulation.
data only, we leverage recently proposed methods for offline RL (Levine et al.. 2020; Kumar et al., 2020) and introduce separate policy heads h and critics Q' for each task. Following Kumar et al. (2020), train the policy maximize the entropy augmented future return:
entropy-augmented
LARRP-CL Σ B.D.D (0't,, (((a,M
+H((
Here H()) denotes the entropy of the policy's output distribution and a is learned weighting factor (Kumar al., 2020; Haarnoja al., 2018b).
Bisimulation (TARP-Bisim). We bisimulation objective for reward-induced representation learning (Larsen and Skou, 1991; Ferns et al., 2011). It groups states based on their "behavioral similarity", measured their expected future returns under arbitrary action sequences, Specifically, we build on the approach of Zhang et al. (2021) that leverages deep neural networks to estimate the bisimulation distance, and modify for multi-task learning by adding per-task embedding heads bisim (())). The representation learning objective is:
LTARP.Bist
- (P(1)), Az,),, P(. P( axe))
(3)
Here P(-|z, a) is learned latent transition model and W2 refers to the 2-Sassrrneiin metric which we can compute in closed form Gaussian transition models. Following Zhang et al. (2021) we use target encoder updated with the moving average of the encoder's weights for producing and add an auxiliary reward prediction objective with per-task reward predictors Pisim (z).
can also train task-induced representations from data without reward annotation by directly imitating the data collection policy, thus learning represent all elements of the state space that were important for the policy's 's decision making. We choose behavioral cloning (BC, Pomerleau (1989)) for imitation lcarning since it is easily applicable especially in the offline setting. We introduce N separate imitation policy heads héc and minimize the negative eigllikelihodd of the training data's 's actions:
Imitation Learning (TARP-BC).
N Σ
log (og((((()))(()))
LTARP-BC
BieayDD
(4)
EXPERIMENTS
5.1 EXPERIMENTAL SETUP
We compare the performance of different representation learning approaches in four visually complex environments (see Figure 3). We will briefly describe each environment; for more details on environment setup and data collection, see appendix, Section A.
5

---

Published as a conference paper at ICLR 2022
Distracting DMControl. We use the environment of Stone et al. (2021), in which the visual complexity of the standard DMControl "Walker" task (Tassa et al., 2018) 8) is increased by overlaying randomly sampled natural videos from the DAVIS 2017 dataset (Pont- Tuset et al., 2017). We train on data collected from standing, forward and backward walking policies and test on a downstream running task. An efficient representation should focus modeling the agent while ignoring irrlevant information from the background video.
ViZDoom. A simulator for the ego- shooter game Doom (Wydmuch al., 2018). We use pre- trained models from Dosovitskiy and Koltun (2017) for data collection and vary their objectives to get diverse hehaviors such as maximizing the number of collected medi-packs or minimizing the loss of health points. Learned representations should focus on important aspects such as the location of enemies or medi-packs, while ignoring irrelevant details such as the texture of walls or appearance features of medi-packs etc. We test on the full "battle" task from Dosovitskiy and Koltun (2017).
Distracting MetaWorld. Based on the MetaWorks multi-task robotic manipulation environ- ment (Yu al., 2020). We increase the visual complexity by overlaying the background with the same natural videos used in the distracting DMControl, but define a much larger set of 15 training and 6 testing tasks, which require manipulation of diverse objects with Sawyer robot arm. All objects manipulated in the testing tasks are present in at least one of the training tasks. During evaluation, we report the average performance across the six testing tasks and normalize each task's performance with the score of the best- best-performing policy.
Autonomous Driving. We simulate realistic first-person driving scenarios using the CARLA simulator (Dosovitskiy et al., 2017). We collect training data from intersection crossings as well as right and left turns using pre-trained policies and test on a long-range point-to-point driving task that requires navigating an unseen part of the environment. For efficient learning, representations need to model relevant driving information such as position and velocity of other cars, while ignoring irrelevant details such trees, shadows or the color and make of cars in the scene.
We compare the different instantiations of the task-induced representation learning framework from Section 4.1 1 to common unsupervised representation learning approaches:
Reconstruction: We train a beta- VAE (Higgins et al., 2017) or a stochastic video prediction model ("Pred-O") on the training data and transfers the encoder.
Contrastive Learning: We use the contrastive learning objective from Stooke et al. (2021) for pre-training ("ATC").
We also compare approach that combines video prediction with reward prediction for pre- training ("Pred-R+O"), similar to Hafner et al. (2019), thus combining elements form unsupervised and task-induced representation learning. Additionally, we report performance of a policy transfer baseline, which pre-trains a policy on D using BC and then finetunes the full policy on the downstream task as an alternative representation transfer For environments that provide a din-dimensional state (DMControl, MetaWorld), we further report results for an oracle baseline.
After pre-training, we transfer the frozen encoder weights the target task policy, which we train with soft actor-critic (SAC, Haarnoja et al. (201 8a)) on continuous control tasks (distracting DMControl, distracting Meta World, CARLA), and with PPO (Schulman et al., 2017) on discrete action tasks (ViZDoom).² For more implementation details and hyperparameters, see Appendix D.
5.2 DOWNSTREAM TASK LEARNING EFFICIENCY
We report downstream task learning curves for task-induced and unsupervised representation learning methods well as direct policy transfer in Figure 4. The low performance of the policy transfer baseline (purple) shows that in most tested environments the downstream task requires significantly different behaviors than those observed in the pre-training data, a scenario in which transferring
We also tried pre-training policies on each of the individual task datasets D1,..., but found the finetuning performance the policies trained on the full dataset D to be superior. We report results with finetuning of the pre-trained encoder in appendix, Section C, and find no substantial difference to the setting with frozen encoder.

---

Published as a conference paper at ICLR 2022
Aummalan Epiicodaa Rusward
Epioodo Romoadd
Eprande Rewoord
Epicodo Roward
230
1.30
130
12
130
Ensironmmnt slaps (1M)
Envirommen steps (1M)
10
06
Enviroment stepa (1M)
611
0.2
53
0.4
25
Envrrnment (1M)
(4) Distracting DMControl
(b) Vi2Doom
(6) Distracting MataWorrd
(d) CARLA
Scratch
Oracle
Policy Transfer
Pred-O
VAE
ATC
Reconstruction
Contrastive
Pred-Rt+O
TARP-V
TARP-COL
TARP-BC ... TARP-BIsm miniw
Reconstruction Task-Induced
Tack-induced Representation
Figure 4: Performance of transferred representations on unseen target tasks. The task-induced representations (blue) lead to higher learning efficiency than the fully unsupervised representations or direct policy transfer and achieve comparable performance to the Oracle baseline. All results averaged across three seeds. See Figure 9 for per-task MetaWorld performances.
representations be beneficial over directly transferring behaviors*. All reconstruction-baat approaches (green) struggle with the complexity of the tested scenes, especially the method that attempts predict the scene dynamics (Pred-O). We find that adding reward prediction the ob- jective improves performance (red, Pred-O+R). Yet, downstream learning is still slow, since the reconstruction objective leads task-irrelevant information being modeled in the learned representa- tion. The non-reconstructive contrastive approach ATC (brown) achieves stronger results, particularly in VizDoom which features less visual distractors, but its performance deteriorates substantially in environments with more task-irrelevant details, i.e. distracted DMControl, Meta World and CARLA.
Overall, we find that task ind-ind representations enable more sample efficient learning. Among the TARP instantiations, we find that TARP-V and TARP-Bisim representations can lead to lower transfer performance, since they rely on the expressiveness of the reward funetion: if future rewards can be predicted without modeling all task-relevant features, the representations will not capture them. In contrast, TARP-CQL and TARP-BC learn representations via policy learning. which can enable more efficient transfer to downstream policy learning problems. On the distracting DMControl task we find that TARP representations even outperform representations trained with direct supervision through handerafted oracle state representation, since they can learn to represent concepts like joint body parts of the walker during pre-training, while the oracle needs to learn these during downstream RL.
5.3 PROBING TASK-INDUCED REPRESENTATIONS
To better understand TARP's improved learning efficiency, we visualize what information captured in the representation in Figure 5: we compare input saliency maps for representations learned with task-induced and unsupervised objectives. Saliency maps visualize the average gradient magnitude for each input pixel with respect to the output Φ(x: and thus capture the contribution of each part of the input to the representation. We find that task-induced representations focus on the important aspects of the scene, such as the walker agent in distracting DMControl and other cars in CARLA. In contrast, the unsupervised approaches have high saliency values for scattered parts of the input and often represent aask-ireeevvant aspects such changing background videos, buildings and trees, since they cannot differentiate task-relevant and irrelevant information.
For quantitative analysis, train probing networks top of the learned representations in distracting DMControl. We test whether (1) task-relevant information is modeled by predicting oracle joint states and (2) whether task-inreeevant information is ignored by classifying the ID of the used background video. The more irrelevant background information is captured in the representation, the better the probing network will be classifying the video. The results show that probing networks trained with task-induced representations more accurately predict the task-relevant state information (Figure 6a) while successfully filtering information about the background video and thus obtaining lower background classification accuracy (Figure 6b).
'The policy transfer baseline achieves good performance in the distracting DMControl environment since it reuse behaviors from the walk-forwaat training task for the target run-forward task.
7

---

Published as a conference paper at ICLR 2022
Distracting DMControl
CARLA
Original
TARP
VAE
ATC
Original
TARP
VAE
ATC
Figure 5: Visualization of the learned representations. Left to right: saliency maps for represen- tations learned with task-induced representation learning (TARP-BC) and the comparisons for reconstruction- based (VAE) and reconstruction-fre (ATC) representation learn- ing. Left: Distracting DMControl environment. Right: CARLA environment. Only task-induced representations can ignore distracting information and focus on the important aspects of the scene.
Scrtthh
Episode Reward
Bardigroudd Dlesiltaatio COCARRO
200
Prodecton
2-
150
TARP-COL
TARP-B
100
TARP-Bsk
0.00
TARP TARPAY
WIE
ATC
TARP TARP-V
WAE
ATC
Environment steps (1M)
(a) State prediction
(b) Background classification
(c) Transfer imitation learning
Figure 6: (a) Task-induced representations (TARP-BC/V) allow for more accurate prediction of the task-relevant joint states. Unsupervised approaches aim to model all information in the input thus probing networks can still learn to predict state information, but struggle achieve comparably low error rates. (b) Task lesk-induced representations successfully filter task irrelevant background information and thus cannot confidently classify the background video, while unsupervised approaches fail filter the irrelevant information and thus achieve perfect classification scores. (c) Transfer performance for IL distracting DMControl. Task-induced representations achieve superior sample efficiency.
5.4 TRANSFERRING REPRESENTATIONS FOR IMITATION LEARNING
We test whether the conclusions from the model-fre free RL experiments above equally apply when instead performing imitation learning (II.) on the downstream task. We train policies with the pre- trained representations from Section 5.2 by Soft-Q Imitation Learning (SQIL,, Reddy the distracting DMControl target running task. In Figure 6c show that task-induced representations also improve downstream performance visually complex environments for IL and allow for more efficient imitation, since they model only task-relevant information. Again, TARP-BE and TARP. CQL lead to the best learning efficiency. This shows that the benefit of modeling task-relevant aspects constrained RL, but the same representations be used accelerate IL.
(2020)) on
5.5 DATA ANALYSIS EXPERIMENTS
In the previous sections, we found that task-induced representations improve learning efficiency over common unsupervised representation learning approaches in visually complex environments. These task-induced representation learning approaches designed leverage data from variety of sources, such as prior training runs or human teleoperation. Thus, analyzing what characteristics of this training data lead to successful transfer is key to their practical to derive a set of best practices when collecting datasets for task-induced representation learning.
goal in this section is
8

---

Published as a conference paper at ICLR 2022
Nermalized Episode Rewarrd
Many tasks vs. lots of data per task. When collecting large datasets there is a trade-off between collecting a lot data for few tasks vs. collecting less data each for larger set of tasks. To analyze the effects of this trade-off on TARP, we train TARP- BC data from varying number of pre-training tasks in distracting Meta World. When training on data from fewer tasks,
Taea (Somiit))
Tasks
5 Taala
lasa
0,4
we collect more data on each of these tasks, ensure that the
Environment sleps (1M)
size of the training datasets is constant across all experiments. The results in Figure show: training on fewer data from larger number of tasks is beneficial over training on lots of data from few tasks. Intuitively, since downstream tasks can leverage the union of task-relevant components of all training tasks, having a diverse set of tasks is more important for transfer than having lots of data for few tasks.
Figure 7: Transfer performance vs. number of pre-training tasks. Col- lecting training data from a larger number of tasks is more important for transfer performance than hav- ing lots of data for few tasks.
More data vs. optimal data. Another trade-off in data col- lection exists between the amount of data and its optimality: method that can Iearn from sub-optimal trajectory data is able to leverage much larger and more diverse datasets. Since training on diverse data is important for successful transfer (see above) robustness to sub-optimal training data is an important feature of any representation learning approach. We test the robustness of task-induced representation learning approaches by training on sub-optimal trajectory data, collected from only partially trained and completely random policies, in distracting DMCon- trol. The downstream RL performance comparison in Figure 8 shows that TARP approaches can learn strong representations from low- wwwpprformance trajectory data and even when trained from random data, performance does not decrease compared SAC baseline trained from scratch. Intuitively, task-induced representation learning does not pre-train a model on "how to act" but merely "what to pay attention to".. We find that TARP-CQL's performance can even increase slightly with the suboptimal data, which we attribute its increased state coverage. Thus we conclude that task-induced representation learning is robust to the optimality of the pre-training data and collecting larger, diverse datasets is more important than collecting optimal data.
DARR-OOA
Eptoods Roward
TARK-
Expert
Subrotim
Randon
Robustness to pre- training data optimality. TARP ap- proaches can learn good representa- tions that allow for effective trans- fer even from sub-optimal data.
Figure 8:
Multi-Source Task Supervision. When collecting large datasets, it can be beneficial to pool data from multiple sources, e.g. data obtained from prior training runs via human teleoperation. In Section 4 we introduced different instantiations of the TARP framework that are able to leverage different forms of task supervision. Here, we test whether using multiple sources of task supervision simultancously. In particular, we train "TARP-++BC" models that assume reward annotations for only some of the tasks in the pre-training dataset and demonstration data for the remaining tasks (for more details on data collection, see Section B). We compare downstream learning efficiency on distracting DMControl and CARLA in Figure 10 We find that the combined-source model trained from the heterogeneous dataset achieves comparable or superior performance to all single-source models, showing that practitioners should collect diverse datasets, even if they have heterogeneous sources of task-supervision.
can train a single representation
6 CONCLUSION
In this work, we investigate the effectiveness of representation learning approaches for transfer in visually complex scenes. We define the family of busk-induced representation learning approaches, that leverage fusk-information from prior tasks to only model task-relevant aspects of a scene while ignoring distractors. We compare task-induced representations to common unsupervised representation learning approaches across four visually complex environments with substantial distractors and find that they lead to improved sample efficiency on downstream tasks compared to unsupervised approaches. Future work should investigate approaches for incorporating other sources of task information such as language commands to learn task-induced representation.

---

Published as a conference paper at ICLR 2022
7 REPRODUCIBILITY STATEMENT
We provide a detailed description of all used environments, the procedures for offline dataset collec- tion, and descriptions of the downstream evaluation tasks in appendix, Section A. Furthermore, in appendix, Section D we list all hyperparameters used for the pre-training phase (i.e. task-induced and unsupervised representation learning) as well for training of the RL policies. We open- source our codebase with example commands to reproduce our results on the project website: clvrai. com/tarp.
ACKNOWLEDGEMENT
We thank colleagues from the CLVR lab at USC for the valuable discussions that considerably helped to shape this work.
REFERENCES
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv: 015400000000000...0.000000000000 2016.
Serkan Cabi, Sergio Gomez Colmenarejo, Alexander Novikov, Ksenia Konyushkova, Scott Reed, Rae Jeong, Konrad Zolna, Yusuf Aytar, David Budden, Mel Vecerik, Oleg Sushkov, David Barker, Jonathan Scholz, Misha Denil, Nando de Freitas, and Ziyu Wang. Scaling data-driven robotics with reward sketching and batch reinforcement learning. RSS, 2019.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations, In International conference machine learning. 2020.
E. Denton and R. Fergus. Stochastic video generation with a learned prior. In ICML, 2018.
Alexey Dosovitskiy and Vladlen Koltun. Learning to act by predicting the future. ICLR, 2017.
Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. CARLA: An open urban driving simulator. In Proceedings of the 1st Annual Conference Robot Learning, pages 1-16, 2017.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 4066):1662-1114 2011.
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In Proceedings of IEEE International Conference on Robotics and Automation, 2016.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 2020.
Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Ma- chine Learning Research. PMLR, 2021. URL http: http::///ooceeeeeeeeeeoee www.m.r....... www.mm.rrr........ fu2 Fu.lb. html.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the T'iiteenth International Conference on Artificial Intelligence and Statistics, pages 297-304. JMLR Workshop and Conference Proceedings, 2010.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. ICML, 2018a.
Tuomas Haarojja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv: 18/2.0590. 2018b.
10

---

Published as a conference paper at ICLR 2022
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. ICML, 2019.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with constrained variational framework. In ICLR, 2017.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Qt-opt: Scalable deep reinforcement learning for vision-basee robotic manipulation. In Conference on Robot Learning, pages 651-673, 2018.
Diederik Kingma and Max Welling. Auto-encoding variational Bayes. In ICLR, 2014.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. NeurlPS, 2020.
Sascha Lange and Martin Riedmiller. Deep auto-encoder neural networks in reinforcement learning. In The 2010 International Joint Conference on Neural Networks (IJCNN) 2010.
Kim G Larsen and Ame Skou. Bisimulation through probabilistic testing. Information and computa- tion, 94(1):1-28, 1991.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning. 2020.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Proceedings of Neural Information Processing Systems (NeurIPS), 2021.
Alex X Lee, Anusha Nagabandi, Picter Abbeel. and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. NeurlPS, 2020.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 17(1):1334-1377 2016.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint wrXiv:2005.01643 2020.
Ajay Mandlekar, Yuke Zhu, Animesh Garg, Jonathan Booher, Max Spero, Albert Tung, Julian Gao, John Emmons, Anchit Gupta, Emre Orbay, Silvio Savarese, and Li Fei-Fei. Roboturk: A crowdsourcing platform for robotic skill learning through imitation. In CoRL, 2018.
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518:529-552 533,,9993,,..0 02 2015.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: /807.03777 2018.
Dean A Pomerleau. Alvinn: An autonomous land vehicle in neural network. In Proceedings of Neural Information Processing Systems (NeurIPS), pages 305-313, 1989.
Jordi Pont-Tuset, Federico Perazzi, Sergi Caelles, Pablo Arbeláez, Alex Sorkine- HokkineHHrnuug, and Luc Van Gool. The 2017 davis challenge on video object segmentation. arXiv preprint
17Xiv:1704.00675, 2017.
Siddharth Reddy, Anca D. Dragan, and Sergey Levine. (SQIL); Imitation learning via reinforcement learning with sparse rewards. In International Conference Learning Representations, 2020. URL https: /ttp::////ppeeeiieeeee net wtt./orrrtttttttttt//mm
forum? dd=$1xkd24tws
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In ICMI,,
11

---

Published as a conference paper at ICLR 2022
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101. 02722. 2021.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. ICML. 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 2018.
Marcus Loo Vergara. Accelerating training of deep reinforcement bearning-based autonomous driving agents through comparative study of agent and environment designs. Master's thesis, Oct 2019, URL wwtph:tt/////ccttm /tce//nntee wtw.upunnnnnnn ntnu. no nunu- xmlui: waii/haaddee e /1111/2222222111111
Marek Wydmuch, Michal Kempka, and Wojciech Jaskowski. Vizdoom competitions: Playing doom from pixels. IEFE Transactions on Games, 2018.
Mengjiao Yang and Ofir Nachum. Representation matters: Offline pretraining for sequential decision making. arXiv preprint arXiv:2102.05815, 2021.
Denis Yarats, Nya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In ICLR, 2021.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: benchmark and evaluation for multi-task and meta reinforcement learning. In Conference Robot Learning (CoRL), 2020.
Albert Zhan, Ruihan Zhao, Lerrel Pinto, Pieter Abbeel, and Michael Laskin. A framework for efficient robotic manipulation. arXiv preprint arXiv:2012.07975, 2020.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. ICLR, 2021.
Li Zhaoping. Theoretical understanding the early visual processes by data compression and data selection. Network: Computation in Neural Systems, 2006.
Çaglar Gülçehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez Colmenarejo, Konrad Zolna, Rishabh Agarwal, Josh Merel, Daniel J. Mankowitz, Cosmin Paduraru, Gabriel Dulac-Arolld, Jerry Li, Mohammad Norouzi, Matthew Hoffman, Nicolas Heess, and Nando de Freitas. RI unplugged: A suite of benchmarks for offline reinforcement learning. In NeurlPS, 2020. URL. https::/ppoccedd ngs neurips. httpscc//paae pwtpppppp////0.0......... 2020////////0//hh/. //aah// 112002200tccdSFfaaacccaaobbbodd Abstract html.
12

---

Published as a conference paper at ICLR 2022
A ENVIRONMENTS
A.1 DISTRACTING DMCONTROL.
Following Stone et al. (2021), increase the visual complexity of the DMControl "walker" task (Tassa et al., 201 by overlaying randomly sampled videos from the DAVIS 2017 dataset (Pont- Tuset al., 2017) in the background. In experiments datasets for the pre-training tasks of standing, forward walking, and backward walking respectively, To collect these datasets, we pre-train polices with SAC(Haarnoja et al., 2018 8a) in the state space and collect rollouts of the visual observation. We test our representation on the downstream task of "running".
expert policies to collect offline
We use reward functions provided by DMControl "walker" task (Tassa et al., 2018). For the backward walking task, we invert the sign of the walking speed in the "forward" task defined in DMControl, so that the agent gets higher reward when it moves backwards instead of forwards.
Rewards:
A.2 VIZDOOM
Our experiments use the "D33 environenvirn provided by Wydmuch et al. (2018) where the agent's 's objective is to defend against enemies while collecting medi-packs and ammunition. For collection of the pre-training task dataset, we use the pre-trained models provided by Dosovitskiy and Koltun (2017) and vary the reward weighting parameters (described below) to produce diverse set of behaviors.
Rewards: A reward function is defined by linear combination of three measurements (ammunition, health, and frags) with their corresponding coefficients.
momm
amm
bealth
heolth It 1t-1
It
frags
21-1
Frags
RvziDoom
Camma
Chealth
+ Cfrags
7.5
30.0
1.0
(5)
where n is measurement of ammunition, x zheatth is the health of the agent, and a fragg the number of frags timestep t. The set of coefficients for ammunition, health and frags are represented as (Cammos Chealth, Cfrage). We use the coefficients (0, 0, 1). (0, 10,1,0) and (1,1, -1) for collecting the pre-training data, and (0.5, 0.5, 1.0) for the target task.
A.3 DISTRACTING METAWORLD
We increase the visual complexity of MetaWorld environment (Yu et al., 2020) by overlaying randomly sampled videos from the DAVIS 2017 dataset (Pont-Tuset et al., 2017), similar Stone et al. (2021). For data collection, we pre-train polices with SAC(Haarnoja et al., 2018a) using low-dimensional state information and collect rollouts of the visual observations. All objects manipulated in the downstream tasks are present in at least one of the training tasks. We define 15 tasks for data collection:
button press plate slide
button press topdown plate slide back handle pull side door lock
button press topdown wall plate slide side handle press coffee button
handle pull
door open
coffee push
push
push back
We evaluated transfer performance 6 downstream tasks:
button press wall door unlock
plate slide back side coffee pull
handle press side
push wall
We report the average performance across the six downstream tasks and normalize each task's performance with the episode reward of the best bert-performing policy. Furthermore, show the performance of each downstream task in Figure 9.
Rewards: We reward functions provided by Metaworld (Yu et al., 2020).
13

---

Published as a conference paper at ICLR 2022
A.4 CARLA
We use the map "Town05" from the CARLA environment (Dosovitskiy et al., 2017) for our experiments. At the beginning of each episode, randomly spawn 300 vehicles and 200 pedestrians. The initial location the agent randomly sampled from task set containing multiple start and goal locations. We collect pre-training datasets for the tasks of intersection crossing, taking right turn, and taking a left turn using pre-trained policies. Then, we test on a long-range curvy road point-to-point driving task that requires navigating an unseen part of the environment. We pre-train the policies with SAC (Haarnoja et al., 2018a) using segmentation masks of the environment as the input, and then use the learnt policies collect rollouts of visual observations for the datasets.
Rewards: We modify reward function used in Vergara (2019) for all of the tasks. The reward function consists of terms for speed, centering on road, angle the agent, and collision.
V
V Utarget
Rapeed
I Stmin + (1.0
1 22 Fmaz
+ 1.0 I miin Se Dimus"
Umin
Umna
Ptarget
dcenter 0). dmax
Reentering
max(1.0 1.0
(6)
Rangle
max(1.0- (Tmox
x 1.11).
180
RCARLA= = Ropeee + Reentering + Rangle 10 10-4, collisiooninteesiiii
where is velocity of the agent, deenter is distance between the center of the road and the agent, is angle the agent. We use constant values of Urain 15.0, Umax 30.0, Utarget 25.0, dmax and Imaz 20.
B DATA COLLECTION FOR MULTI-SOURCE TASK SUPERVISION
We detail the composition of the pre-training dataset for the experiment in which train TARP from heterogeneous data sources. For the distracting DMControl environment, the dataset is composed of demonstrations for the "forward walking" and "backward walking" tasks, and a dataset with reward annotations for the "stand" task. For the CARLA environment, the task-induced representations are trained on a dataset composed of demonstrations for the "intersection crossing" task, and data annotated with rewards for the "right turn" and "left turn" tasks. For both environments, the datasets are collected with the procedures described in appendix, Scction A.
C FINETUNING LEARNED REPRESENTATIONS
In our experimental evaluation in Section we held the parameters of the pre-trained encoder fixed to cleanly evaluate the quality of the pre-trained representation. However, practice pre-trained representations are often finetuned on the target task using target task rewards. Prior work on representation learning found mixed results when finetuning the pre-trained representations (Yang and Nachum, 2021). Thus, in this section experimentally compare the different representation learning approaches on the Distracting DMControl, ViZDoom, and CARLA environment while finetuning the learned representations task-induced representations show improved sample efficiency and better performance even in the fine-tuning setting, analogous the results in Section 5.2.
the target task. As illustrated in Figure 1 11 we find that
14

---

Published as a conference paper at ICLR 2022
1000
Episode Reward
Episode Reward
1000
Episode Reward
1000
800
200
200
200
0.0
0.1
0.3
0.0
0.1
0.2
0.3
0.4
00
0.1
0.2
0.4
Environment steps (1M)
Environment steps (1M)
Environment steps (1M)
(a) Button press wall
(b) Plate slide back side
(c) Handle press side
1000
1000
Episode Reward
800
Episode Reward
Episode Reward
800
200
400
200
0.0
0.1
0.3
0.4
0.0
0.1
02
0.3
0.4
Environment steps (1M)
Environment steps (1M)
Environment steps (1M)
(d) Door unlock
(e) Coffee pull
(1) Push wall
Scratch
Oracle
Pred-O
Reconstruction
TARP-V
TARP-CQL
TARP-BC
TARP-Blaim
Contrastive
Task-Induced Representation
Figure 9: Transfer performance for each downstream task in the distracting MetaWorld environment. The task-induced representations generally show better performance than unsupervised representa- tions and achieve comparable performance the oracle baseline.
300
TARP-V
TARP-V
TARP-BC
TARP-BC
Episode Reward
250
TARP-V+BG
Episode Reward
200
TARP-V+BC
200
150
150
100
100
50
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Environment steps (1M)
Environment steps (1M)
(a) Distracting DMControl
(b) CARLA
Figure 10: Transfer performance for single source and multiple source task-induced representation. The task-induced representation with multiple supervision sources performs slightly better or as good the single source representation. This shows that task-induced representations can be effectively trained from datasets with heterogeneous forms of task supervision.
15

---

Published as a conference paper at ICLR 2022
300
17.5
250
Episods Reward
Episode Reward
Episode Reward
100
160
60
100
2.5
20
Environment steps (1M)
Environment steps (1M)
Environment steps (1M)
(a) Distracting DMControl
(b) ViZDoom
CARLA
Scratchh
Pred-O
VAE
ATC
Pred-R+O
Reconstruction
Recomstrudionn-Taskknddoo
TARP-Y
TARP-OOL
TARP-BC
TARP-Sreim
Task Induced Representation
Figure 11: Performance of transferred representations with finetuning on unseen target tasks on distracting DMControl, ViZDoom, and CARLA environments. The task-induced representations (blue) show better sample efficiency and performance compared to the representations learned with unsupervised objectives.
D HYPERPARAMETERS
D.1 HYPERPARAMETERS FOR RL
Table 1: Common SAC hyperparameter
Parameter
Value
Stacked frames Optimizer
Adam
Learning rate
3e-4
Discount factor (v)
0.99
Latent dimension Convolution filters
256
[8, 16, 32, 64] [2, 2]
Convolution strides
Convolution filter size
Hidden Units (MLP)
[1024]
Nonlinearity
ReLU
Target smoothing coefficient (r) Target entropy
0.005
-dim(A)
Table Distracting DMControl and MetaWorld SAC hyperparameter
Parameter
Value
Observation Rendering
(64, 64), RGB
Initial steps
5 x 10
Action repeat
2
Replay buffer size
105
Minibatch size
256
Target update interval
I
Actor update interval
I
Initial temperature
1
16

---

Published as a conference paper at ICLR 2022
Table CARLA SAC hyperparameter
Parameter
Value
Observation Rendering
(128, 128), RGB
Initial steps
3 x 10
Action repeat
I
Replay buffer size
10
Minibatch size
128
Target update interval
2
Actor update interval
2
Initial temperature
0.1
Table 4: ViZDoom PPO hyperparameter
Parameter
Value
Observation rendering
(64, 64) Grey 4
Stacked frames
Action repeat Optimizer
1
Adam
Learning rate
3e-4
PPO epoch
10
Buffer size
2048
Convolution filters
[8, 16, 32, 64]
Convolution filter sizes
[2, 2, 2, 2] [256]
Hidden units (MLP)
Generalized advantage estimation
0.95
Entropy bonus coefficient
4e-3
Discount factor (?)
0.99
Minibatch size
256
Nonlinearity
ReLU
Table 5: Hyperparameters for CQL
Environment
Trade-off factor for Q-values
Number of action samples
Distracting DMControl
ViZDoom
1.
CARLA
3.
D.2 HYPERPARAMETERS FOR PRE-TRAINING
In TARP-BC, we use recurrent neural networks to predict sequence of actions over prediction horizon T' learn task-induced representations only in CARLA (see Table 7).
Table 6: Common model parameters
Parameter
Value
Batch size
128
Hidden units (MLP)
[256]
Learning rate
le44
17

---

Published as a conference paper at ICLR 2022
Table 7: Hyperparameters for TARP-BC
Environment
LSTM hidden units
Prediction horizon
Distracting DMControl
ViZDoom
Distracting MetaWorld
CARLA
512
Table 8: Hyperparameters for TARP-V
Environment
Discount rate
All environments
0.4
Table Hyperparameters for TARP-Bisim
Environment
Reward predictive loss weight
Bisimulation loss weight T
Distracting DMControl
0.1
ViZDoom
1.
10.
Distracting World
1.
0.1
CARLA
1.
0.01
Table 10: Hyperparameters for VAE
Environment
8 constraint
Distracting DMControl
100.
ViZDoom
100.
Distracting MetaWorld
100.
CARLA
10.
Table 11: Hyperparameters for Pred-S and Pred-R+S
Environment
3 constraint
Prediction horizon
Distracting DMControl
50.
6
ViZDoom
10.
Distracting Meta World
50.
6
CARLA
5.
8
Table 12: Hyperparameters for ATC
Environment
Random shift probability
Temporal shift
All environments
1.
18


Reviews:
[{"review_id": "hUhKrh4pcdF", "reviewer": "Reviewer_s9b6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "Strengths of the paper:\n\n- The paper is very well written.\n- The authors evaluate the proposed way of learning task-induced representations to various other \"task\" agnostic ways of learning representations, and show that the proposed method is able to capture task-relevant information. \n- The ablations done in section 5.5 are very informative i.e., set of best practices while collecting datasets for learning task-induced representations \n\nWeakness:\n-  For transferring information between different tasks, there are normally three different ways (a) by learning representations of the raw data (b) by learning policies/skills (c) by learning world models. In this work, paper compares to the proposed framework with the scenario when the information is transferred by learning representations of the raw data. It would be also useful to compared to methods/benchmarks which learn the representation of skills/dynamic models. \n- It may be useful to compare to methods which explicitly capture task relevant information like information bottleneck approaches [1, 2, 3].\n- One other drawback is, that the paper only evaluates by modifying the visual perception by adding distractors. Since the paper uses dynamics model, it would probably be also meaningful to see as to what happens when some \"irrelevant\" dynamics are added during testing.\n\n[1] Information asymmetry in KL-regularized RL, https://arxiv.org/abs/1905.01240\n[2] InfoBot, https://arxiv.org/abs/1901.10902\n[3] Learning invariant representations for reinforcement learning without reconstruction, https://arxiv.org/abs/2006.10742\n\nAfter rebuttal: \n\nI've read the rebuttal. The paper makes it clear they are comparing to methods  that directly learn parametric model (either dynamics model or directly transfer policy)  dynamics models from the offline data (Pred-O and Pred-O+R in Figure 4) and others that directly transfer learned behaviors from the offline data (Policy transfer in Figure 4). I keep my original score.\n", "summary_of_the_paper": "In this work, authors formalized the problem of task-induced representation learning. It basically involves investigating various different ways for learning representations from the raw data (for ex. offline dataset) containing \"privileged\" information (task relevant) from many different tasks. The aim is to leverage task relevant information for learning representations that capture task-relevant information. \n\nThe authors evaluate their idea on three different suite of environments (a) Distracting DMControl (b) ViZDoom (c) Autonomous Driving. The authors compare the proposed way of learning task-induced representation to various other methods for representation learning like reconstruction based approaches, mutual information based methods, prediction based methods etc. The results show that the proposed way of learning task-induced representations capture task-relevant information. ", "summary_of_the_review": "The reviewer  found the paper very well written, and tackling an important problem. I like the ablations done in the paper, and the focus on \"constructing datasets\" for learning task-induced representations.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "W9_MklmegiU", "reviewer": "Reviewer_uzb1", "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": "strengths:\n* experiments comparing with reasonable baselines\n* paper is clear\n\nweaknesses:\n* i am not sure this paper addresses a real problem. In the limit of high empowerment e.g. human, an agent could achieve a very large number of goals and all potential parts of the screen should be relevant. In that sense this approach only makes sense when a precise goal is targeted and the task space in pre-training can be used to cover it.\n* salience maps are not convincing to me. I am not sure what they are meant to convey and in which sense the TARP ones are better than other ones. Background / foreground prediction is also not very relevant in this context.\n\nquestions:\n* for unsupervised representation learning methods are the unsupervised losses still optimized during fine-tuning. For a fair comparison they should be because it is possible and would make the claim of transfer stronger if they still don't get better results than TARP.", "summary_of_the_paper": "The paper proposes a representation learning approach for multi-task reinforcement learning. The idea is to pre-train a model using multi-task data either with rewards or behavior cloning. Then fine-tune it on a new task hopefully more efficiently. They contrast it to unsupervised representation learning.\n\nFine-tuning experimentation is performed with reasonable baselines on 3 different visual domains. The authors probe both what the representation looks like and what pretraining tasks give better performance.\n\n\n", "summary_of_the_review": "The contribution and novelty are both very limited.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "QVLU679K54q", "reviewer": "Reviewer_k7xA", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "Advantages:\n\n- I found the analysis combining saliency maps and state predictions vs background video classification interesting. It was nice to see the authors attempting to gain intuition by combining different techniques.\n- The authors attempt to cover a number of interesting questions, from what should we predict, to how homogeneous or optimal the dataset should be. I found this interesting and informative.\n\nDisadvantages:\n\n- The authors only chose a single training-testing transfer pair of tasks from each environment type. Given how different the rankings of the performance of the methods are between each environment, I am curious whether there’s a high variance between different training-testing transfer pairs from the same environment. Though the results and narrative make sense, I am wondering whether the conclusions would hold in additional training-testing pairs within these three environments.\n- I’m not clear what Section 5.4 and figure 6c show. What are the exact training-testing settings here? How do they differ from the experiments whose results are shown in figure 4?\n- Though the authors attempt to answer many questions, they only provide analysis results on some of the methods. For example, how do saliency maps and state vs background predictions of every other method compare? Wasn’t it possible to undertake a similar analysis for non-TARP methods too?\n- I recommend rewriting the contributions paragraph. Points 1 and 3 seem to be very similar and I don’t think the grouping provided by point 2 is necessary. Recognising that task relevant information is important for representation learning is a worthy contribution, but in my humble opinion “formalize task induced representation learning (TARP) as an alternative family of representation learning approaches that leverage task information from prior tasks“ is not really a separate contribution. \n", "summary_of_the_paper": "The paper attempts to answer the question of how to pre-train encoders in order to get task-relevant representations that are useful for a downstream task in each of three different environments (ViZDoom,  Distracting DMContro and Autonomous Driving). Three main approaches are compared (Behavioral Cloning, Value prediction only or another commonly used offline RL loss), along with a number of additional baselines (VAEs, predicting combinations of observations and rewards, contrastive methods and others), in a setting where an encoder is pre-trained offline on data coming from a certain task and its weights are frozen and used to provide representations to a policy head trained to solve an unseen task, that’s held out in the original dataset. The authors also provide an analysis using saliency maps to explore which parts of the input image are important for decreasing the loss of each method (task relevant vs background) and also show the performance of the representation on a task relevant and a task irrelevant task. Finally, the authors explore how the quality of the data (coming from many tasks vs coming from a single task) during the representation learning phase, and whether they come from an optimal, random or suboptimal policy, affects downstream performance. The main contribution of the paper, in my opinion, is this investigation of the different methods. In particular the comparison of performance in each of the three environments, the saliency analysis and predictability of state vs background information and finally the exploration of the effect of data on performance.", "summary_of_the_review": "In summary, I think the paper is interesting, but importantly, I did not find many of the findings surprising enough to recommend for acceptance. There are additionally two main issues with the paper in its current form. First, the contributions section in the introduction should, in my opinion, be re-written (see disadvantages section above). Second, experimentally there are results on only a single training-testing transfer pair for each of three environments. I would expect a more thorough investigation of multiple of these pairs to assess the robustness of methods that make use of task information. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "IZmolZu3j_N", "reviewer": "Reviewer_TaF6", "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": "Strengths:\n1.\tThe paper presents an interesting discussion on task relevant supervision vs. purely unsupervised techniques. The empirical results are mostly in support of the proposed argument and contain diverse domains of tasks that the technique was tested on. \n2.\tThe paper follows up the main results with with an additional analysis and discussion of why TARP techniques perform better than the baselines, specifically, the role data plays in this performance boost. This analysis adds value by showing a closer look at the features and the kind of data that would result in better performance; and is valuable for the representation learning community in general.\n\nWeaknesses:\n1.\tFirstly, the crux of the paper deals with a comparison between using task-based supervision to learn representations, and unsupervised representation learning techniques. While the argument that task supervision is more valuable compared to unsupervised learning is interesting, there is an obvious difference in the sense that task supervision directly allows for extracting task-relevant features in a representation. When unsupervised learning techniques do not account for any relevance, or there are no particular constraints as in techniques like beta VAE, it is obvious that those representations would not fare as well as a task-induced one in the absence of any additional supervision, constraints, hand crafted loss functions etc. In that way, it is not exactly an apples-to-apples comparison. This is made evident from the saliency map comparison, in the absence of any finetuning / pre-defined constraints, there is no way for an unsupervised learning method to know which features are important for an arbitrary task. \n\nThis is why I feel that the discussion of novelty and the comparison with existing techniques is somewhat insufficient. For instance, it would have been more interesting and valuable to compare the proposed technique with algorithms such as DREAMER [1,2].  DREAMER also contains similar subgoals as TARP such as a value learning network. Although it is true that TARP proposes training value prediction or reward maximization on a ‘set’ of tasks as opposed to a single task, how is that better than, for instance, applying DREAMER itself to multiple tasks and expecting it to build a generalizable representation? Similarly, another relevant algorithm would be TIA (Task induced abstractions) [3] that proposes targeted reconstructions to solve issues with generalization as well as distractors and outperforms DREAMER. Comparing with other techniques in similar classes would help gauge the novelty and efficacy of the current method and localize it properly within the current literature. \n\n2.\tSecondly, there are some parts of the evaluation I am not clear about. What is the main intuition for why the baseline representations are unable to focus on task relevance during the policy learning process? During the pretraining phase, an unsupervised model likely encodes both task-relevant and task irrelevant features in each scene – which gives rise to two possibilities. Either the model is failing to encode features in a disentangled way, which ruins performance during finetuning because it is not able to isolate task-relevant features from the others. If that is not the case, then it should be possible to isolate which of the encoded features are useful during finetuning for a task. If the baseline models can do the latter, then they should be applicable for diverse tasks similar to the way TARP is claimed to be. Hence, I think it would be very valuable to have some experiment that includes finetuning the representations as well while learning the policy. Speaking just conceptually, if we consider only a single task, it is not clear why TARP based learning is fundamentally different from finetuning a representation learnt in an unsupervised way for a given task, assuming that the representation is a well structured one. \n\n\n3.\tI think the paper would also benefit from a discussion of how the proposed techniques could compare against the current increasing push towards unsupervised pretraining. Unsupervised pretraining through large “foundation models”, as well as reward-free pretraining are gaining traction in the world of representation learning and reinforcement learning, where specific tasks or rewards are not considered, and in cases were found to benefit downstream reinforcement learning or imitation learning [4]. In the current paper, the authors mention through empirical findings that it is more important to obtain (although less) data from a large number of tasks, as opposed to lots of data from few tasks. But this also means a lot of effort and handcrafting in setting up the tasks, the associated rewards and so on. \n\nMinor comments/clarification questions:\n\n1. What is the downstream task corresponding to Figure 7? \n2. It is seen that the contrastive learning approach (ATC) performs really badly on distracting DMControl. How much of that is due to a lack of data augmentation? Would a contrastive learning approach perform significantly better if it was trained on a sufficiently large dataset with multiple distractors? \n3. In Figure 8, is there a significant difference in the size of datasets between expert and suboptimal?\n\n[1] Hafner, D., Lillicrap, T., Ba, J., & Norouzi, M. (2019). Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603.  \n[2] Hafner, D., Lillicrap, T., Norouzi, M., & Ba, J. (2020). Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193.  \n[3] Fu, X., Yang, G., Agrawal, P., & Jaakkola, T. (2021, July). Learning task informed abstractions. In International Conference on Machine Learning (pp. 3480-3491). PMLR.  \n[4] Yang, M., & Nachum, O. (2021). Representation matters: Offline pretraining for sequential decision making. arXiv preprint arXiv:2102.05815.", "summary_of_the_paper": "This paper discusses the problem of learning meaningful representations that can efficiently focus only on the features that are relevant for downstream tasks. As opposed to unsupervised representation learning approaches that do not differentiate between task-specific and other information in a given dataset, the paper proposes a framework named “Task Induced Representation Learning” which leverages task information to guide the representation learning. Through this framework, the paper suggests three possible ways to induce task-relevance in representations: by predicting the value of a state, maximizing the discounted reward for a group of tasks, or by imitating an expert policy again for a group of tasks. Through a set of experiments in three domains: distracting DMControl, VizDoom and CARLA, the authors show higher task learning efficiency using the proposed framework as compared to standard approaches such as reconstruction, contrastive learning, state-reward prediction etc.  ", "summary_of_the_review": "This paper proposes a framework called Task-induced representation learning (TARP), with the claim that having task supervision during representation learning is a lot more beneficial than unsupervised, task-free representation learning. Moreover, TARP is able to use supervision from a set of multiple, perhaps diverse tasks so that the representations can benefit from different kinds of guidance and thereby encode features that are useful for downstream tasks. Through a set of experiments on DMC, VizDoom, CARLA etc. the authors show that TARP representations outperform usual unsupervised learning approaches. While the paper is well written and the discussion is interesting, it is well known that task induced supervision will induce more relevance compared to unconstrained, unsupervised learning, and I feel the paper is not strong enough because it does not compare to more relevant baselines such as DREAMER, Task Informed Abstractions.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper is an interesting take on representation learning, using (prior) tasks to determine which information is important. The problem setting is somewhat difficult to pin down, so that that finding the correct comparisons is not obvious and opinions differ on many details of the setup. However, this is not a fault of the paper; it is a general problem the further one moves away from clean settings like classical supervised learning.

There was a lengthy and detailed back-and-forth between the authors and reviewers, where the authors clarified most of the points raised, extended their results, resulting in one reviewer switching from reject to accept.
Decision: Accept (Poster)