Year: 2021

Abstract:
A major bottleneck for applying deep reinforcement learning to real-world prob1 lems is its sample inefficiency, particularly when training policies from high2 dimensional inputs such as images. A number of recent works use unsupervised 3 representation learning approaches to improve sample efficiency. Yet, such unsu4 pervised approaches are fundamentally unable to distinguish between task-relevant 5 and irrelevant information. Thus, in visually complex scenes they learn representa6 tions that model lots of task-irrelevant details and hence lead to slower downstream 7 task learning. Our insight: to determine which parts of the scene are important and 8 should be modeled, we can exploit task information, such as rewards or demonstra9 tions, from previous tasks. To this end, we formalize the problem of task-induced 10 representation learning (TARP), which aims to leverage such task information in 11 offline experience from prior tasks for learning compact representations that focus 12 on modelling only task-relevant aspects. Through a series of experiments in visu13 ally complex environments we compare different approaches for leveraging task 14 information within the TARP framework with prior unsupervised representation 15 learning techniques and (1) find that task-induced representations allow for more 16 sample efficient learning of unseen tasks and (2) formulate a set of best-practices 17 for task-induced representation learning. 18

text: A major bottleneck for applying deep reinforcement learning to real-world prob-1 lems is its sample inefficiency, particularly when training policies from high-2 dimensional inputs such as images. A number of recent works use unsupervised3 representation learning approaches to improve sample efficiency. Yet, such unsu-4 pervised approaches are fundamentally unable to distinguish between task-relevant5 and irrelevant information. Thus, in visually complex scenes they learn representa-6 tions that model lots of task-irrelevant details and hence lead to slower downstream7 task learning. Our insight: to determine which parts of the scene are important and8 should be modeled, we can exploit task information, such as rewards or demonstra-9 tions, from previous tasks. To this end, we formalize the problem of task-induced10 representation learning (TARP), which aims to leverage such task information in11 offline experience from prior tasks for learning compact representations that focus12 on modelling only task-relevant aspects. Through a series of experiments in visu-13 ally complex environments we compare different approaches for leveraging task14 information within the TARP framework with prior unsupervised representation15 learning techniques and (1) find that task-induced representations allow for more16 sample efficient learning of unseen tasks and (2) formulate a set of best-practices17 for task-induced representation learning.18
1 INTRODUCTION19
A central capability of intelligent agents acting in complex environments is to filter their sensory20 inputs and retrieve the important bits of information. Humans are remarkably efficient at this: while21 hundreds of MB of raw pixel data enter our retina in every second, only roughly 40 bits/sec remain22 after passing the visual processing system (Zhaoping, 2006). This highly filtered and compressed23 representation of reality allows for more efficient learning and acting in the complex environments of24 our everyday lives. Similarly, our goal is to efficiently train artificial agents in complex environments25 by learning which parts of their sensory inputs should be modeled and which can be discarded.26
Deep reinforcement learning (RL) replaces manually engineered perception pipelines with policies27 that are trained directly from raw visual inputs (Mnih et al., 2015; Levine et al., 2016), but requires a28 prohibitive amount of environment interactions for learning a good representation. Thus, unsupervised29 objectives for representation learning have been proposed (Lange and Riedmiller, 2010; Finn et al.,30 2016; Hafner et al., 2019; Laskin et al., 2020; Stooke et al., 2021). While substantially more efficient,31 these approaches are typically only evaluated in clean laboratory or simulated settings (Finn et al.,32 2016; Lee et al., 2020; Laskin et al., 2020) where most of the perceived information is important for33 the task at hand and therefore should be modeled.34
In contrast, realistic environments feature lots of task-irrelevant detail, like clutter or objects moving35 in the background. Trying to capture all the details in the representation will make both representation36 learning and downstream task learning inefficient or even infeasible. Thus, we need mechanisms for37 filtering task-irrelevant details when learning representations in visually complex environments.38
Fundamentally however, unsupervised representation learning techniques cannot address the filtering39 problem in such environments: they do not discriminate between task-relevant and irrelevant informa-40 tion, but instead are trained to model all perceived information. Without additional supervision, for41 them, modelling the movement of tree leaves on the road side might be as important as capturing an42 approaching car in front of the agent.43
TARP
Visually Complex Input Scene
Models only Taskrelevant Information

heading: Task-Induced
text: Representation Learning
Prior Task Experience w/ Task Information (rewards, demos…)
Downstream Policy
π
Visualize
{s, a, r1} {s, ademo} {s, a, r2} {s, ademo} Task 1 Task 2 Task 3 Task 4
The contributions of our work are threefold: (1) we empirically show that common unsupervised53 representation learning techniques struggle at learning tasks in visually challenging environments54 since they cannot distinguish between relevant and irrelevant information, (2)we investigate the role55 of task information in representation learning by formalizing task-induced representation learning56 (TARP) that leverage task information from prior tasks and (3) we compare different approaches for57 leveraging such task information within the TARP framework with prior unsupervised representation58 learning techniques in three complex visual domains: distracting DMControl, ViZDoom, and CARLA.59 We find that task-induced representations allow for more sample efficient learning of unseen tasks60 and formulate a set of best-practices for task-induced representation learning.61
2 RELATED WORK62
To enable agents to act from high-dimensional observations, classic approaches have relied on63 engineered perception pipelines that process the input into a manually defined “state” represen-64 tation. This has allowed remarkable progress in challenging real-world domains like autonomous65 driving (Urmson et al., 2008) or humanoid robotics (Atkeson et al., 2015), but manually defining66 such perception interfaces for each task is costly and requires substantial human effort.67
Instead, deep RL has allowed the training of “end-to-end” policies that directly map raw visual68 inputs to action commands, without the need for human-engineered perception modules (Mnih69 et al., 2015; Kalashnikov et al., 2018). Due to the need for large environment interactions, deep70 RL approaches are typically evaluated in clean lab settings with uncluttered input observations that71 allow for easier learning (Finn et al., 2016; Pathak et al., 2018; Zhang et al., 2019), with a few72 notable exceptions such as Gupta et al. (2018) and Kahn et al. (2021) that rely on large-scale data73 collection. In addition to sample inefficiency, representations learned through end-to-end task training74 are specific to the training task and cannot be reused for solving different downstream tasks. In75 contrast, our goal is to learn representations that enable sample efficient learning on unseen tasks.76
To improve data efficiency in deep RL from high-dimensional observations, a number of recent77 works have taken inspiration from advances in computer vision and have explored using data78 augmentation (Yarats et al., 2021; Laskin et al., 2021) and unsupervised representation learning79 techniques during RL training. The latter can be categorized into (1) reconstruction (Lange and80 Riedmiller, 2010; Finn et al., 2016), (2) prediction (Hafner et al., 2019; Lee et al., 2020), and81 (3) contrastive learning (Laskin et al., 2020; Stooke et al., 2021; Zhan et al., 2020; Yang and Nachum,82 2021) approaches. While these works have shown improved sample efficiency, fundamentally they83 are unsupervised and therefore cannot decide which information in the input is relevant or irrelevant to84 a certain task. We show that this leads to deteriorating performance in visually complex environments85 such as autonomous driving scenarios with lots of non-relevant information in the input observations.86 Task-induced representation learning addresses this problem by leveraging prior task experience to87 infer which input information is “interesting”, allowing it to learn representations that enable sample88 efficient learning of new tasks in scenes of real-world complexity.89
Closest to our work are recent approaches that explore reconstruction-free representation learn-90 ing (Gelada et al., 2019; Zhang et al., 2021). Specifically, Zhang et al. (2021) use a bisimulation91 objective to learn a reward-induced representation which can ignore distractors. While Zhang et al.92 (2021) focus on single-task, online representation learning, we aim to leverage diverse offline datasets,93 collected across multiple tasks, which can enable efficient reuse of previously collected data from a94 variety of sources. Further, we compare multiple objectives for task-induced representation learning95 beyond bisimulation that can leverage diverse forms of task supervision.96
3 PROBLEM FORMULATION97
Our goal is to efficiently learn a policy π that solves a target task Ttarget in an MDP defined by a tuple98 (S,A, T , R, ρ, γ) of states, actions, transition probabilities, rewards, initial state distribution, and99 discount factor. The policy is trained to maximize the discounted return Eat∼π [∑T−1 t=0 γ tR(st, at) ] .100
We do not assume access to the underlying state s of the MDP, but instead our policy receives a101 high-dimensional observation x ∈ X in every step, e.g., an image observation. To improve training102 efficiency, we aim to learn an encoder φ(x) which maps the input observation to a low-dimensional103 state representation that is input to the policy π(φ(x)). To learn this representation, we assume access104 to a dataset of past interactions D = {D1, . . . ,DN} collected across T1:N ∈ T tasks, with per-task105 datasets Di = {xt, at, (rt), ...} of state-action trajectories and optional reward annotation. Crucially,106 the set of training tasks does not include the target task Ttarget /∈ T1:N . We require the training data to107 include task information for prior tasks, e.g., in the form of step-wise reward annotations rt or by108 consisting of task demonstrations.109
We follow Stooke et al. (2021) and test all representation learning approaches in two phases: we110 first train the representation encoder φ(x) from the offline dataset D, then freeze its parameters and111 optimize the policy π(φ(x)) on the downstream task. This two-stage experimental protocol has two112 benefits: (1) by training representations solely from offline experience we are able to efficiently reuse113 data collected across prior tasks, e.g., from prior training runs (Fu et al., 2020; Çaglar Gülçehre et al.,114 2020) or human teleoperation (Mandlekar et al., 2018; Cabi et al., 2019), and (2) by freezing the115 encoder weights after pre-training we can evaluate the quality of the pre-trained representation in116 isolation. All tested approaches can be combined with finetuning of the learned representation on the117 downstream task to further improve performance, but we leave this investigation for future work.118
4 TASK-INDUCED REPRESENTATION LEARNING119
Our core idea is to use task information from prior tasks as a filter for which aspects of the environment120 are interesting, and which do not need to be modeled. This addresses a fundamental problem of the121 unsupervised representation learning approaches: by maximizing the mutual information between122 observations and representation they are trained to model every bit of information equally and thus123 struggle in visually complex environments with lots of irrelevant details.124
Formally, the state of an environment S can be divided into task-relevant and task-irrelevant or125 nuisance components S = {Sitask,Sin}. The superscript i indicates that this division is task-dependent,126 since some components of the state space are relevant for a particular task Ti but irrelevant for others.127
In visually complex environments, we can assume that |Stask| |S|, i.e., that the task-relevant part of128 the state space is much smaller than all state information in the input observations (Zhaoping, 2006).129
When choosing how much of this input information to model, end-to-end policy learning and130 unsupervised representation learning, form two ends of a spectrum. End-to-end policy learning only131 models Sitask for its training task, which is efficient, but only allows transfer from task i to task j132 if Sjtask ⊆ Sitask, i.e., if the task-relevant components of the target task are a subset of those of the133 training task. In contrast, unsupervised representation learning models the full S = {Stask,Sn},134 which allows for flexible transfer since Sjtask ⊆ S ∀ Tj ∈ T , but training can be inefficient since the135 learned representation contains many nuisance variables.136
In this work, we formalize the problem of task-induced representation learning (TARP), which137 aims to combine the best of both approaches. Similar to end-to-end policy learning, task-induced138 representation learning approaches use task information to learn compact, task-relevant represen-139 tations. Yet, by combining the task information from a wide range of prior tasks they learn a140 representation that combines the task-relevant components of all tasks in the training dataset D:141 STARP = S1task ∪ · · · ∪ SNtask. Thus, such a representation allows for transfer to a wide range of unseen142 tasks for which S targettask ⊆ STARP. Yet, the representation filters a large part of the state space which is143 not relevant for any of the training tasks, allowing sample efficient learning on the target task.144
4.1 APPROACHES FOR TASK-INDUCED REPRESENTATION LEARNING145
We compare multiple approaches for task-induced representation learning which can leverage or even146 mix different forms of task supervision. In this work we focus on two forms, reward annotations147 and task demonstrations; but, future work can extend the TARP framework to other forms of task148 supervision like language commands or human preferences. All instantiations of TARP train a shared149 encoder network φ(x) with separate task-supervision heads (see Figure 2).150
Value Prediction (TARP-V). We can leverage reward annotations as task supervision by estimating151 the future discounted return of the data collection policy. Intuitively, a representation that allows152 estimation of the value of a state needs to include all task-relevant aspects of this state. We introduce153 separate value prediction heads hiv for each task i and train the representation φ(x) by minimizing154 the error in the predicted discounted return:155
LTARP-V = N∑ i=1 E(xt,rt)∼Di ( hiv(φ(xt))− T∑ t′=t γt ′−trt′ )2 . (1)
Offline RL (TARP-CQL). Alternatively, we can learn a representation by training a policy to156 directly maximize the discounted reward of the training tasks. Since we aim to use offline training157 data only, we leverage recently proposed methods for offline RL (Levine et al., 2020; Kumar et al.,158 2020) and introduce separate policy heads hiπ and critics Q
i for each task. Following Kumar et al.159 (2020), we train the policy to maximize the entropy-augmented future return:160
LTARP-CQL = − N∑ i=1 Ex∼Di ( Qi(x, hiπ(φ(x))) + αH ( hiπ(φ(x)) )) . (2)
Here H(·) denotes the entropy of the policy’s output distribution and α is a learned weighting161 factor (Kumar et al., 2020; Haarnoja et al., 2018b).162
Bisimulation (TARP-Bisim). We use a bisimulation objective for reward-induced representation163 learning (Larsen and Skou, 1991; Ferns et al., 2011). It groups states based on their “behavioral164 similarity”, measured as their expected future returns under arbitrary action sequences. Specifically,165 we build on the approach of Zhang et al. (2021) that leverages deep neural networks to estimate the166 bisimulation distance, and modify it for multi-task learning by adding per-task embedding heads167 zi = hibisim ( φ(x) ) . The representation learning objective is:168
LTARP-Bisim = N∑ i=1 E (xj ,aj ,rj) (xk,ak,rk) ∼Di ( |zij−zik|−|rj−rk|−γW2 ( P (·|z̄ij,t, aj,t), P (·|z̄ik,t, ak,t) )) (3)
Here P (·|z, a) is a learned latent transition model andW2 refers to the 2-Wasserstein metric which169 we can compute in closed form for Gaussian transition models. Following Zhang et al. (2021) we use170 a target encoder updated with the moving average of the encoder’s weights for producing z̄ and we171 add an auxiliary reward prediction objective with per-task reward predictors r̃ibisim(z).172
Imitation Learning (TARP-BC). We can also train task-induced representations from data without173 reward annotation by directly imitating the data collection policy, thus learning to represent all174 elements of the state space that were important for the policy’s decision making. We choose behavioral175 cloning (BC, Pomerleau (1989)) for imitation learning since it is easily applicable especially in the176 offline setting. We introduce N separate imitation policy heads hiBC and minimize the negative177 log-likelihood of the training data’s actions:178
LTARP-BC = − N∑ i=1 E(x,a)∼Di ( log hiBC(a|φ(x)) ) (4)
5 EXPERIMENTS179
5.1 EXPERIMENTAL SETUP180
We compare the performance of different representation learning approaches in three visually complex181 environments (see Figure 3). We will briefly describe each environment; for more details on182 environment setup and data collection, see appendix, Section A.183
Distracting DMControl. We use the environment of Stone et al. (2021), in which the visual184 complexity of the standard DMControl “Walker” task (Tassa et al., 2018) is increased by overlaying185 randomly sampled natural videos from the DAVIS 2017 dataset (Pont-Tuset et al., 2017). We train186 on data collected from standing, forward and backward walking policies and test on a downstream187 running task. Task-induced representations should focus on modeling the agent while ignoring188 irrelevant information from the background video.189
ViZDoom. A simulator for the ego-shooter game Doom (Wydmuch et al., 2018). We use pre-190 trained models from Dosovitskiy and Koltun (2017) for data collection and vary their objectives to191 get diverse behaviors such as maximizing the number of collected medi-packs or minimizing the loss192 of health points. Learned representations should focus on important aspects such as the location of193 enemies or medi-packs, while ignoring irrelevant details such as the texture of walls or appearance194 features of medi-packs etc. We test on the full “battle” task from Dosovitskiy and Koltun (2017).195
Autonomous Driving. We simulate realistic first-person driving scenarios using the CARLA196 simulator (Dosovitskiy et al., 2017). We collect training data from intersection crossings as well as197 right and left turns using pre-trained policies and test on a long-range point-to-point driving task that198 requires navigating an unseen part of the environment. For efficient learning, representations need199 to model relevant driving information such as position and velocity of other cars, while ignoring200 irrelevant details such as trees, shadows or the color and make of cars in the scene.201
We compare the different instantiations of the task-induced representation learning framework to202 prior representation learning approaches:203
• Reconstruction: Trains a beta-VAE (Higgins et al., 2017) or stochastic video prediction204 model (“Pred-O”) on the training data and transfers the encoder.205
• Reconstruction + Task-Induced: Combines video prediction with reward prediction for206 pre-training (“Pred-R+O”), similar to Hafner et al. (2019) or uses the approach of Fu207 et al. (2021) (“TIA“) that additionally aims to disentangle task-relevant and irrelevant208 representations.209
• Contrastive Learning: Uses the contrastive learning objective from Stooke et al. (2021)210 for pre-training (“ATC”).211
All approaches are trained on the same offline dataset D used for training of the TARP approaches.212 We also compare to a policy transfer baseline which pre-trains a policy on D using BC and then213 finetunes the full policy on the downstream task as an alternative to representation transfer1. In214 distracting DMControl we further report results for an oracle baseline whose representation is trained215 with direct supervision from the low-dimensional state representation2.216
After pre-training of the representation, we transfer the frozen encoder weights to the target task217 policy, which we train with soft actor-critic (SAC, Haarnoja et al. (2018a)) on continuous control218 tasks (distracting DMControl, CARLA), and with PPO (Schulman et al., 2017) on discrete action219 tasks (ViZDoom). For more implementation details and hyperparameters, see Appendix E.220
5.2 DOWNSTREAM TASK LEARNING EFFICIENCY221
We report downstream task learning curves for task-induced and unsupervised representation learning222 methods as well as direct policy transfer in Figure 4. The low performance of the policy transfer223 baseline (purple) shows that in most tested environments the downstream task requires significantly224 different behaviors than those observed in the pre-training data, a scenario in which transferring225 representations can be beneficial over directly transferring behaviors3. All reconstruction-based226 approaches (green) struggle with the complexity of the tested scenes, especially the method that227 attempts to predict the scene dynamics (Pred-O). We find that adding reward prediction to the228 objective improves performance (red), especially for TIA which learns to successfully disentangle229 foreground and background in the Distracting DMControl environment (see Figure 10). Yet, we230 find that downstream learning is still slow, possibly because the reconstruction objective leads to231 task-irrelevant agent appearance information being modeled in the learned representation.232
The non-reconstructive contrastive approach ATC (brown) achieves stronger results, particularly in233 the VizDoom environment which features less visual distractors, but its performance deteriorates234
1We also tried pre-training policies on each of the individual task datasets D1, . . . ,DN but found the finetuning performance of the policies trained on the full dataset D to be superior.
2The other environments do not provide a low-dimensional state representation. 3The policy transfer baseline achieves good performance in the distracting DMControl environment since it
can reuse behaviors from the walk-forward training task for the target run-forward task.
Scratch
TARP-V TARP-CQL TARP-BC Task-Induced Representation TARP-Bisim
(a) Finetuning learned representations (b) Performance of downstream walk task
Pred-O VAE Reconstruction

heading: ATC
text: Contrastive Pred-R+O Reconstruction+Task-Induced
substantially in environments with more task-irrelevant details, i.e. distracted DMControl and235 CARLA.236
Overall, we find that task-induced representations enable more sample efficient learning. Among the237 TARP instantiations, we find that TARP-V and TARP-Bisim representations can lead to lower transfer238 performance, since they rely on the expressiveness of the reward function: if future rewards can be239 predicted without modeling all task-relevant features, the representations will not capture them. In240 contrast, TARP-CQL and TARP-BC learn representations via policy learning, which can enable more241 efficient transfer to downstream policy learning problems. On the distracting DMControl task we find242 that TARP representations even outperform representations trained with direct supervision through243 a handcrafted oracle state representation, since they can learn to represent concepts like joint body244 parts of the walker during pre-training, while the oracle needs to learn these during downstream RL.245
5.3 PROBING TASK-INDUCED REPRESENTATIONS246
To better understand the improved learning efficiency of the TARP approaches over the unsupervised247 methods, we visualize what information is captured in the representation in Figure 5: we compare248 input saliency maps4 for representations learned with task-induced and unsupervised objectives. We249 find that task-induced representations can focus only on the important aspects of the scene, such as250 the walker agent in distracting DMControl and other cars in CARLA. In contrast, the unsupervised251 approaches have high saliency values for scattered parts of the input and often represent task-irrelevant252 aspects such as changing background videos, buildings and trees, since they cannot differentiate253 task-relevant and irrelevant information.254
To quantitatively analyze whether task-induced representations effectively represent task-relevant255 information and discard task-irrelevant features, we train probing networks on top of the learned256 representations in distracting DMControl. We test whether (1) task-relevant information is modeled257 by predicting oracle joint states and (2) whether task-irrelevant information is ignored by classifying258 the ID of the used background video. The more irrelevant background information is captured in the259 representation, the better the probing network will be at classifying the video.260
We report the probing networks’ state prediction error in Figure 6a and the background classification261 accuracy in Figure 6b to compare representations learned via TARP to unsupervised objectives (we262 choose VAE and ATC as comparisons since they achieved the best transfer performance among263 reconstruction and non-reconstruction baselines respectively). From this comparison we can see264 that probing networks trained on top of task-induced representations can more accurately predict the265 task-relevant state information while they are successfully filtering information about the background266
4Saliency 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.
video and thus obtain lower background classification accuracy. Therefore, these quantitative results267 support the qualitative intuition from Figure 5.268
5.4 TRANSFER REPRESENTATION FOR IMITATION LEARNING269
We show that the representations learned via task-induced representation learning are general and270 can be used for multiple downstream learning approaches: in addition to the model-free RL results271 above we demonstrate their effectiveness for imitation learning (IL). We train policies with the272 representations learned on the same datasets used in Section5.2 by Soft-Q Imitation Learning (SQIL,273 Reddy et al. (2020)) on the distracting DMControl target running task. In Figure 6c we show that task-274 induced representations also improve downstream performance in visually complex environments for275 IL and allow for more efficient imitation, since they model only task-relevant information. Again,276 TARP-BC and TARP-CQL lead to the best learning efficiency. This shows that the benefit of TARP277 is not constrained to RL, but the same representations can be used to accelerate IL.278
5.5 DATA ANALYSIS EXPERIMENTS279
Task-induced representation learning approaches are designed to leverage data from a variety of280 sources, such as prior training runs or human teleoperation. Thus, analyzing what characteristics of281 this training data lead to successful transfer is key to their practical use – our goal in this section is to282 derive a set of best practices when collecting datasets for task-induced representation learning.283
Many tasks vs. lots of data per task. When collecting large datasets there is a trade-off between284 collecting a lot of data for few tasks vs. collecting less data each for a larger set of tasks. To285 analyze the effects of this trade-off on TARP, we train TARP-BC on data from a varying number286 of pre-training tasks in distracting DMControl. When training on data from fewer tasks, we collect287 more data on each of these tasks, to ensure that the size of the training datasets is constant across all288 experiments. The results in Figure 7 show: training on fewer data from a larger number of tasks289 is beneficial over training on lots of data from few tasks. Intuitively, since downstream tasks can290 leverage the union of task-relevant components of all training tasks, having a diverse set of tasks is291 more important for transfer than having lots of data for few tasks.292
More data vs. optimal data. Another trade-off in data col-293 lection exists between the amount of data and its optimality: a294 method that can learn from sub-optimal trajectory data is able to295 leverage much larger and more diverse datasets. Since training296 on diverse data is important for successful transfer (see above)297 robustness to sub-optimal training data is an important feature298 of any representation learning approach. We test the robustness299 of task-induced representation learning approaches by training300 on sub-optimal trajectory data, collected from only partially301 trained and completely random policies, in distracting DMCon-302 trol. The downstream RL performance comparison in Figure 8303 shows that TARP approaches can learn strong representations304 from low-performance trajectory data and even when trained305 from random data does not decrease performance over a SAC306 baseline trained from scratch because, intuitively, task-induced307 representation learning does not pre-train a model on “how308
to act” but merely “what to pay attention to”. We find that TARP-CQL’s performance can even309 increase slightly with the suboptimal data, which we attribute to its increased state coverage. Thus310 we conclude that task-induced representation learning is robust to the optimality of the pre-training311 data and collecting larger, diverse datasets is more important than collecting optimal data.312
Multi-Source Task Supervision. When collecting large313 datasets, it can be beneficial to pool data from multiple sources,314 e.g. data obtained from prior training runs or via human teleop-315 eration. In Section 4 we introduced different instantiations of316 the TARP framework that are able to leverage different forms317 of task supervision. Here, we test whether we can train a single318 representation using multiple sources of task supervision simul-319 taneously. In particular, we train “TARP-V+BC” models that320 assume reward annotations for only some of the tasks in the321 pre-training dataset and demonstration data for the remaining322 tasks (for more details on data collection, see Section B). We323 propagate gradients from all prediction heads into the same324 representation and compare downstream learning efficiency on325 distracting DMControl and CARLA in Figure 9. We find that326
the combined-source model trained from the heterogeneous dataset achieves comparable or superior327 performance to all single-source models, showing that practicioners should collect diverse datasets,328 even if they have heterogeneous sources of task-supervision.329
6 CONCLUSION330
In this work, we formalize the problem of task-induced representation learning and analyze multiple331 approaches for leveraging offline experience of prior tasks for learning compact representations332 in visually complex scenes. Our empirical results show that task-induced representations retain333 task-relevant information and discard irrelevant features, leading to improved sample efficiency on334 downstream task compared to prior unsupervised representation learning approaches. Future work335 should investigate approaches for incorporating other sources of task information such as language336 commands to learn task-induced representation.337
7 REPRODUCIBILITY STATEMENT338
We perform all of our experiments on publicly available environments. We provide a detailed339 description of the procedures used for the offline dataset collection of prior tasks in Appendix A.340 Furthermore, in Appendix E we list all of the hyperparameters used for the pre-training phase (i.e.341 task-induced and unsupervised representation learning) as well as training RL policies. We include342 our codebase with example commands to reproduce our results in the supplementary materials.343
REFERENCES344
Christopher G Atkeson, Benzun P Wisely Babu, Nandan Banerjee, Dmitry Berenson, Christoper P345 Bove, Xiongyi Cui, Mathew DeDonato, Ruixiang Du, Siyuan Feng, Perry Franklin, et al. No falls,346 no resets: Reliable humanoid behavior in the darpa robotics challenge. In 2015 IEEE-RAS 15th347 International Conference on Humanoid Robots (Humanoids), pages 623–630. IEEE, 2015.348
Serkan Cabi, Sergio Gomez Colmenarejo, Alexander Novikov, Ksenia Konyushkova, Scott Reed,349 Rae Jeong, Konrad Zolna, Yusuf Aytar, David Budden, Mel Vecerik, Oleg Sushkov, David Barker,350 Jonathan Scholz, Misha Denil, Nando de Freitas, and Ziyu Wang. Scaling data-driven robotics351 with reward sketching and batch reinforcement learning. RSS, 2019.352
Alexey Dosovitskiy and Vladlen Koltun. Learning to act by predicting the future. ICLR, 2017.353
Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. CARLA:354 An open urban driving simulator. In Proceedings of the 1st Annual Conference on Robot Learning,355 pages 1–16, 2017.356
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov357 decision processes. SIAM Journal on Computing, 40(6):1662–1714, 2011.358
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep359 spatial autoencoders for visuomotor learning. In Proceedings of IEEE International Conference on360 Robotics and Automation, 2016.361
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep362 data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.363
Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions.364 In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Ma-365 chine Learning Research. PMLR, 2021. URL http://proceedings.mlr.press/v139/366 fu21b.html.367
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp:368 Learning continuous latent space models for representation learning. In International Conference369 on Machine Learning, pages 2170–2179. PMLR, 2019.370
Abhinav Gupta, Adithyavairavan Murali, Dhiraj Gandhi, and Lerrel Pinto. Robot learning in homes:371 Improving generalization and reducing dataset bias. NeurIPS, 2018.372
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy373 maximum entropy deep reinforcement learning with a stochastic actor. ICML, 2018a.374
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash375 Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and376 applications. arXiv preprint arXiv:1812.05905, 2018b.377
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James378 Davidson. Learning latent dynamics for planning from pixels. ICML, 2019.379
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick,380 Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with a381 constrained variational framework. In ICLR, 2017.382
Gregory Kahn, Pieter Abbeel, and Sergey Levine. Badgr: An autonomous self-supervised learning-383 based navigation system. IEEE Robotics and Automation Letters, 6(2):1312–1319, 2021.384
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre385 Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Qt-opt: Scalable deep386 reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning,387 pages 651–673, 2018.388
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline389 reinforcement learning. NeurIPS, 2020.390
Sascha Lange and Martin Riedmiller. Deep auto-encoder neural networks in reinforcement learning.391 In The 2010 International Joint Conference on Neural Networks (IJCNN), 2010.392
Kim G Larsen and Arne Skou. Bisimulation through probabilistic testing. Information and computa-393 tion, 94(1):1–28, 1991.394
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations395 for reinforcement learning. In International Conference on Machine Learning, 2020.396
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas.397 Reinforcement learning with augmented data. In Proceedings of Neural Information Processing398 Systems (NeurIPS), 2021.399
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic:400 Deep reinforcement learning with a latent variable model. NeurIPS, 2020.401
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep402 visuomotor policies. The Journal of Machine Learning Research, 17(1):1334–1373, 2016.403
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial,404 review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.405
Ajay Mandlekar, Yuke Zhu, Animesh Garg, Jonathan Booher, Max Spero, Albert Tung, Julian406 Gao, John Emmons, Anchit Gupta, Emre Orbay, Silvio Savarese, and Li Fei-Fei. Roboturk: A407 crowdsourcing platform for robotic skill learning through imitation. In CoRL, 2018.408
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle-409 mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen,410 Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra,411 Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning.412 Nature, 518:529–533, 02 2015.413
Deepak Pathak, Parsa Mahmoudieh, Guanghao Luo, Pulkit Agrawal, Dian Chen, Yide Shentu, Evan414 Shelhamer, Jitendra Malik, Alexei A Efros, and Trevor Darrell. Zero-shot visual imitation. In415 Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops,416 pages 2050–2053, 2018.417
Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. In Proceedings of418 Neural Information Processing Systems (NeurIPS), pages 305–313, 1989.419
Jordi Pont-Tuset, Federico Perazzi, Sergi Caelles, Pablo Arbeláez, Alex Sorkine-Hornung, and420 Luc Van Gool. The 2017 davis challenge on video object segmentation. arXiv preprint421 arXiv:1704.00675, 2017.422
Siddharth Reddy, Anca D. Dragan, and Sergey Levine. {SQIL}: Imitation learning via reinforcement423 learning with sparse rewards. In International Conference on Learning Representations, 2020.424 URL https://openreview.net/forum?id=S1xKd24twB.425
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy426 optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.427
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite –428 a challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722,429 2021.430
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning431 from reinforcement learning. ICML, 2021.432
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden,433 Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint434 arXiv:1801.00690, 2018.435
Chris Urmson, Joshua Anhalt, Drew Bagnell, Christopher Baker, Robert Bittner, MN Clark, John436 Dolan, Dave Duggins, Tugrul Galatali, Chris Geyer, et al. Autonomous driving in urban environ-437 ments: Boss and the urban challenge. Journal of Field Robotics, 25(8):425–466, 2008.438
Marek Wydmuch, Michał Kempka, and Wojciech Jaśkowski. Vizdoom competitions: Playing doom439 from pixels. IEEE Transactions on Games, 2018.440
Mengjiao Yang and Ofir Nachum. Representation matters: Offline pretraining for sequential decision441 making. arXiv preprint arXiv:2102.05815, 2021.442
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing443 deep reinforcement learning from pixels. In ICLR, 2021.444
Albert Zhan, Ruihan Zhao, Lerrel Pinto, Pieter Abbeel, and Michael Laskin. A framework for445 efficient robotic manipulation. arXiv preprint arXiv:2012.07975, 2020.446
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant447 representations for reinforcement learning without reconstruction. ICLR, 2021.448
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew J Johnson, and Sergey Levine.449 Solar: deep structured representations for model-based reinforcement learning. ICLR 2019, 2019.450
Li Zhaoping. Theoretical understanding of the early visual processes by data compression and data451 selection. Network: Computation in Neural Systems, 2006.452
Çaglar Gülçehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez Colmenarejo,453 Konrad Zolna, Rishabh Agarwal, Josh Merel, Daniel J. Mankowitz, Cosmin Paduraru, Gabriel454 Dulac-Arnold, Jerry Li, Mohammad Norouzi, Matthew Hoffman, Nicolas Heess, and Nando455 de Freitas. Rl unplugged: A suite of benchmarks for offline reinforcement learning. In456 NeurIPS, 2020. URL https://proceedings.neurips.cc/paper/2020/hash/457 51200d29d1fc15f5a71c1dab4bb54f7c-Abstract.html.458
A ENVIRONMENTS459
A.1 DISTRACTING DMCONTROL460
We increase the visual complexity of DMControl “walker” task (Tassa et al., 2018) by overlaying461 randomly sampled videos from the DAVIS 2017 dataset (Pont-Tuset et al., 2017), similar to Stone462 et al. (2021). In our experiments we use expert policies to collect offline datasets for the pre-training463 tasks of standing, forward walking, and backward walking respectively. To collect these datasets, we464 pre-train polices with SAC(Haarnoja et al., 2018a) in the state space and collect rollouts of the visual465 observation. We test our representation on the downstream task of “running”.466
Rewards: We use reward functions provided by DMControl “walker” task (Tassa et al., 2018). For467 the backward walking task, we invert the sign of walk speed in “forward” task defined in DMControl468 so that the agent gets higher reward when the agent moves backward instead of forward.469
A.2 VIZDOOM470
Our experiments use “D3 battle”environment provided by Wydmuch et al. (2018) where the agent’s471 objective is to defend against enemies while collecting medi-packs and ammunition. For collection472 of the prior task dataset, we use the pre-trained models provided by Dosovitskiy and Koltun (2017)473 and vary the reward weighting parameters (described below) to produce a diverse set of behaviours.474
Rewards: A reward function is defined by a linear combination of three measurements (ammunition,475 health, and frags) with their corresponding coefficients.476
RViZDoom = cammo · xammt − xammt−1
7.5 + chealth · xhealtht − xhealtht−1 30.0 + cfrags · xfragst − x frags t−1 1.0 .
(5)
where xammt is a measurement of ammunition, x health t is health of the agent, and x frags t is a number477 of frags at timestep t. The set of coefficients for ammunition, health and frags are represented as478 (cammo, chealth, cfrags). We use the coefficients of (0, 0, 1), (0, 1, 0), and (1, 1,−1) for the prior479 tasks, and (0.5, 0.5, 1.0) for the target task.480
A.3 CARLA481
In our experiments, we use the map of “Town05” from the CARLA environment (Dosovitskiy et al.,482 2017). At the beginning of each episode, we randomly spawn 300 vehicles and 200 pedestrians. The483 initial location of the agent is randomly sampled from a task set containing multiple start and goal484 locations. We collect the datasets for the tasks of intersection crossing, taking a right turn, and taking485 a left turn using pre-trained policies. We pre-train the policies with SAC (Haarnoja et al., 2018a)486 using segmentation masks of the environment as the input, and then use the learnt policies to collect487 rollouts of visual observations for the datasets.488
Rewards: We use the same reward function for all of the tasks. The reward function consists of489 terms for speed, centering on a road, angle of the agent, and collision.490
Rspeed = v
vmin · 1v≤vmin + (1.0− v − vtarget vmax − vtarget ) · 1v≥vmax + 1.0 · 1vmin≤v≤vmax .
Rcentering = max(1.0− dcenter dmax , 0).
Rangle = max(1.0− | r
(rmax · π180 ) |, 0).
RCARLA = Rspeed +Rcentering +Rangle − 10−4 · collision_intensity.
(6)
where v is velocity of the agent, dcenter is distance between the center of the road and the agent, r is491 angle of the agent. We use constant values of vmin = 15.0, vmax = 30.0, vtarget = 25.0, dmax = 3,492 and rmax = 20.493
B DATA COLLECTION FOR MULTI-SOURCE TASK SUPERVISION494
We use supervision from heterogeneous data sources to learn task-induced representations. For the495 distracting DMControl environment, the dataset is composed of demonstrations for the “forward496 walking“ and “backward walking“ tasks, and a dataset with reward annotations for the “stand“ task.497 For the CARLA environment, the task-induced representations are trained on a dataset composed of498 demonstrations for the “intersection crossing” task, and the datasets annotated with rewards for the499 “right turn” and “left turn” tasks. For both environments, the datasets are collected with the procedure500 described in Appendix A.501
C FINETUNING LEARNED REPRESENTATIONS502
In our experimental evaluation in Section 5 we held the parameters of the pre-trained encoder503 fixed to cleanly evaluate the quality of the pre-trained representation. However, in practice pre-504 trained representations are often finetuned on the target task using target task rewards. Prior work on505 representation learning found mixed results when finetuning the pre-trained representations (Yang and506 Nachum, 2021). Thus, in this section we experimentally compare the different representation learning507 approaches on the Distracting DMControl environment while finetuning the learned representations508 on the target task. As illustrated in Figure 11(a), we find that task-induced representations show509 improved sample efficiency and better performance even in the fine-tuning setting, analogous to the510 results in Section 5.2.511
D TRANSFER PERFORMANCE FOR DIFFERENT PRIOR AND TARGET TASK SETS512
We experimentally validate that we can interchange pre-training vs target task sets. Figure 11(b)513 shows the downstream performance on the "walk" task with representations pre-trained on "run",514 "backward walking", and "standing" tasks. This quantitative result underlines that task-induced515 representation lead to more efficient downstream training than unsupervised methods – the same516 conclusions as in Section 4 apply in this scenario.517
E HYPERPARAMETERS518
E.1 HYPERPARAMETERS FOR RL519
Under review as a conference paper at ICLR 2022
Scratch Pred-O VAE Reconstruction
ATC Contrastive
TARP-V TARP-CQL TARP-BC Task-Induced Representation
(b) ViZDoom (c) CARLA(a) Distracting DMControl Oracle
Pred-R+O

heading: Policy Transfer
text: TARP-BisimTIA R construction+Task-Induced
E.2 HYPERPARAMETERS FOR PRE-TRAINING520
In TARP-BC, we use recurrent neural networks to predict a sequence of actions over prediction521 horizon T to learn task-induced representations only in CARLA (see Table 7).522

Review summaries:
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.  Fine-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.   

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.   The 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. 

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.  

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.