T ASK -I NDUCED R EPRESENTATION L EARNING
A BSTRACT
A major bottleneck for applying deep reinforcement learning to real-world problems is its sample inefﬁciency, particularly when training policies from highdimensional inputs such as images.
A number of recent works use unsupervised representation learning approaches to improve sample efﬁciency.
Yet, such unsupervised approaches are fundamentally unable to distinguish between task-relevant and irrelevant information.
Thus, in visually complex scenes they learn representations that model lots of task-irrelevant details and hence lead to slower downstream task learning.
Our insight: to determine which parts of the scene are important and should be modeled, we can exploit task information, such as rewards or demonstrations, from previous tasks.
To this end, we formalize the problem of task-induced representation learning (TARP), which aims to leverage such task information in ofﬂine experience from prior tasks for learning compact representations that focus on modelling only task-relevant aspects.
Through a series of experiments in visually complex environments we compare different approaches for leveraging task information within the TARP framework with prior unsupervised representation learning techniques and (1) ﬁnd that task-induced representations allow for more sample efﬁcient learning of unseen tasks and (2) formulate a set of best-practices for task-induced representation learning.
I NTRODUCTION
A central capability of intelligent agents acting in complex environments is to ﬁlter their sensory inputs and retrieve the important bits of information.
Humans are remarkably efﬁcient at this: while hundreds of MB of raw pixel data enter our retina in every second, only roughly 40 bits/sec remain after passing the visual processing system (Zhaoping, 2006).
This highly ﬁltered and compressed representation of reality allows for more efﬁcient learning and acting in the complex environments of our everyday lives.
Similarly, our goal is to efﬁciently train artiﬁcial agents in complex environments by learning which parts of their sensory inputs should be modeled and which can be discarded.
Deep reinforcement learning (RL) replaces manually engineered perception pipelines with policies that are trained directly from raw visual inputs (Mnih et al., 2015; Levine et al., 2016), but requires a prohibitive amount of environment interactions for learning a good representation.
Thus, unsupervised objectives for representation learning have been proposed (Lange and Riedmiller, 2010; Finn et al., 2016; Hafner et al., 2019; Laskin et al., 2020; Stooke et al., 2021).
While substantially more efﬁcient, these approaches are typically only evaluated in clean laboratory or simulated settings (Finn et al., 2016; Lee et al., 2020; Laskin et al., 2020) where most of the perceived information is important for the task at hand and therefore should be modeled.
In contrast, realistic environments feature lots of task-irrelevant detail, like clutter or objects moving in the background.
Trying to capture all the details in the representation will make both representation learning and downstream task learning inefﬁcient or even infeasible.
Thus, we need mechanisms for ﬁltering task-irrelevant details when learning representations in visually complex environments.
Fundamentally however, unsupervised representation learning techniques cannot address the ﬁltering problem in such environments: they do not discriminate between task-relevant and irrelevant information, but instead are trained to model all perceived information.
Without additional supervision, for them, modelling the movement of tree leaves on the road side might be as important as capturing an approaching car in front of the agent.
[Figure]
Our core insight is that we can only decide which information to retain and which to ﬁlter by taking task information into account.
In the example above, humans rely on a lifetime of experience solving roadside tasks to understand that the approaching car is more relevant than the tree leaves.
To this end, we formalize the problem of task-induced representation learning ( TARP , see Figure 1) which aims to leverage task-related information, such as rewards or demonstrations from prior tasks, for learning representations that focus on the task-relevant information in the observations.
We analyze different approaches for leveraging task information within the TARP framework and show that the learned compact, task-driven representations accelerate the training of unseen tasks in domains of real-world visual complexity.
The contributions of our work are threefold: (1) we empirically show that common unsupervised representation learning techniques struggle at learning tasks in visually challenging environments since they cannot distinguish between relevant and irrelevant information, (2)we investigate the role of task information in representation learning by formalizing task-induced representation learning (TARP) that leverage task information from prior tasks and (3) we compare different approaches for leveraging such task information within the TARP framework with prior unsupervised representation learning techniques in three complex visual domains: distracting DMControl, ViZDoom, and CARLA.
We ﬁnd that task-induced representations allow for more sample efﬁcient learning of unseen tasks and formulate a set of best-practices for task-induced representation learning.
R ELATED W ORK
To enable agents to act from high-dimensional observations, classic approaches have relied on engineered perception pipelines that process the input into a manually deﬁned “state” representation.
This has allowed remarkable progress in challenging real-world domains like autonomous driving (Urmson et al., 2008) or humanoid robotics
(Atkeson et al., 2015), but manually deﬁning such perception interfaces for each task is costly and requires substantial human effort.
Instead, deep RL has allowed the training of “end-to-end” policies that directly map raw visual inputs to action commands, without the need for human-engineered perception modules (Mnih et al., 2015; Kalashnikov et al., 2018).
Due to the need for large environment interactions, deep
RL approaches are typically evaluated in clean lab settings with uncluttered input observations that allow for easier learning (Finn et al., 2016; Pathak et al., 2018; Zhang et al., 2019), with a few notable exceptions such as Gupta et al.
(2018) and Kahn et al.
(2021) that rely on large-scale data collection.
In addition to sample inefﬁciency, representations learned through end-to-end task training are speciﬁc to the training task and cannot be reused for solving different downstream tasks.
In contrast, our goal is to learn representations that enable sample efﬁcient learning on unseen tasks.
To improve data efﬁciency in deep RL from high-dimensional observations, a number of recent works have taken inspiration from advances in computer vision and have explored using data augmentation (Yarats et al., 2021;
Laskin et al., 2021) and unsupervised representation learning techniques during RL training.
The latter can be categorized into (1) reconstruction (Lange and
Riedmiller, 2010; Finn et al., 2016), (2) prediction (Hafner et al., 2019; Lee et al., 2020), and (3) contrastive learning (Laskin et al., 2020; Stooke et al., 2021; Zhan et al., 2020; Yang and Nachum, 2021) approaches.
While these works have shown improved sample efﬁciency, fundamentally they are unsupervised and therefore cannot decide which information in the input is relevant or irrelevant to a 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.
Task-induced representation learning addresses this problem by leveraging prior task experience to infer which input information is “interesting”, allowing it to learn representations that enable sample efﬁcient learning of new tasks in scenes of real-world complexity.
Closest to our work are recent approaches that explore reconstruction-free representation learning (Gelada et al., 2019; Zhang et al., 2021).
Speciﬁcally, Zhang et al.
(2021) use a bisimulation objective to learn a reward-induced representation which can ignore distractors.
While Zhang et al.
(2021) focus on single-task, online representation learning, we aim to leverage diverse ofﬂine datasets, collected across multiple tasks, which can enable efﬁcient reuse of previously collected data from a variety of sources.
Further, we compare multiple objectives for task-induced representation learning beyond bisimulation that can leverage diverse forms of task supervision.
P ROBLEM F ORMULATION
Our goal is to efﬁciently learn a policy π that solves a target task T target in an MDP deﬁned by a tuple ( S , A , T , R, ρ, γ ) of states, actions, transition probabilities, rewards, initial state distribution, and discount factor.
The policy is trained to maximize the discounted return E a t ∼ π (cid:2) (cid:80) T − 1 t =
γ t R ( s t , a t ) (cid:3) .
We do not assume access to the underlying state s of the MDP, but instead our policy receives a high-dimensional observation x ∈ X in every step, e.g., an image observation.
To improve training efﬁciency, we aim to learn an encoder φ ( x ) which maps the input observation to a low-dimensional state representation that is input to the policy π ( φ ( x )) .
To learn this representation, we assume access to a dataset of past interactions D
= {D 1 , . . .
, D N } collected across T 1: N ∈ T tasks, with per-task datasets D i
= { x t , a t , ( r t ) , ... } of state-action trajectories and optional reward annotation.
Crucially, the set of training tasks does not include the target task T target / ∈ T 1: N .
We require the training data to include task information for prior tasks, e.g., in the form of step-wise reward annotations r t or by consisting of task demonstrations.
We follow Stooke et al.
and test all representation learning approaches in two phases: we ﬁrst train the representation encoder φ ( x ) from the ofﬂine dataset D , then freeze its parameters and optimize the policy π ( φ ( x )) on the downstream task.
This two-stage experimental protocol has two beneﬁts: (1) by training representations solely from ofﬂine experience we are able to efﬁciently reuse data collected across prior tasks, e.g., from prior training runs (Fu et al., 2020; Çaglar Gülçehre et al., 2020) or human teleoperation (Mandlekar et al., 2018; Cabi et al., 2019), and (2) by freezing the encoder weights after pre-training we can evaluate the quality of the pre-trained representation in isolation.
All tested approaches can be combined with ﬁnetuning of the learned representation on the downstream task to further improve performance, but we leave this investigation for future work.
T ASK -I NDUCED R EPRESENTATION L EARNING
Our core idea is to use task information from prior tasks as a ﬁlter for which aspects of the environment are interesting, and which do not need to be modeled.
This addresses a fundamental problem of the 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 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 i task , S in } .
The superscript i indicates that this division is task-dependent , since some components of the state space are relevant for a particular task T i but irrelevant for others.
ϕ ( x ) h 1 BC ( a | ϕ ( x ))
h 2 BC ( a | ϕ ( x ))
h NBC ( a | ϕ ( x ))
h 1 v ( V | ϕ ( x ))
h 2 v ( V | ϕ ( x ))
h Nv ( V | ϕ ( x ))
h 1 π ( a | ϕ ( x ))
Q 1 ( ϕ ( x ), h 1 π ( ϕ ( x )) )
CQL
N x
TARP-V
TARP-CQL
TARP-BC h 1 bisim ( z | ϕ ( x ))
P ( z t +1 | ¯ z t , a t ) ˜
r 1 ( z )
TARP-Bisim
Bisim
In visually complex environments, we can assume that |S task | (cid:28) |S| , i.e., that the task-relevant part of the state space is much smaller than all state information in the input observations (Zhaoping, 2006).
When choosing how much of this input information to model, end-to-end policy learning and unsupervised representation learning, form two ends of a spectrum.
End-to-end policy learning only models
S
i task for its training task, which is efﬁcient, but only allows transfer from task i to task j if S j task ⊆ S i task , i.e., if the task-relevant components of the target task are a subset of those of the training task.
In contrast, unsupervised representation learning models the full S =
{S task , S n } , which allows for ﬂexible transfer since S j task ⊆ S ∀ T j ∈ T , but training can be inefﬁcient since the learned representation contains many nuisance variables.
In this work, we formalize the problem of task-induced representation learning (TARP), which aims to combine the best of both approaches.
Similar to end-to-end policy learning, task-induced representation learning approaches use task information to learn compact, task-relevant representations.
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 :
S TARP =
S 1 task ∪ · · · ∪ S N task .
Thus, such a representation allows for transfer to a wide range of unseen tasks for which S targettask ⊆ S TARP .
Yet, the representation ﬁlters a large part of the state space which is not relevant for any of the training tasks, allowing sample efﬁcient learning on the target task.
4. A PPROACHES FOR T ASK -I NDUCED R EPRESENTATION L EARNING
We compare multiple approaches for task-induced representation learning which can leverage or even mix different forms of task supervision.
In this work we focus on two forms, 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 by estimating the future discounted return of the data collection policy.
Intuitively, a representation that allows estimation of the value of a state needs to include all task-relevant aspects of this state.
We introduce separate value prediction heads h iv for each task i and train the representation φ ( x ) by minimizing the error in the predicted discounted return:
[Equation]
Ofﬂine 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 ofﬂine training data only, we leverage recently proposed methods for ofﬂine RL (Levine et al., 2020; Kumar et al., 2020) and introduce separate policy heads h iπ and critics Q i for each task.
Following Kumar et al.
(2020), we train the policy to maximize the entropy-augmented future return:
[Equation]
[Figure]
Bisimulation (TARP-Bisim).
We use a bisimulation objective for reward-induced representation learning (Larsen and Skou, 1991; Ferns et al., 2011).
It groups states based on their “behavioral similarity”, measured as their expected future returns under arbitrary action sequences.
Speciﬁcally, we build on the approach of Zhang et al.
(2021) that leverages deep neural networks to estimate the bisimulation distance, and modify it for multi-task learning by adding per-task embedding heads
z i = h
i
[Equation]
φ ( x ) (cid:1) .
The representation learning objective is:
[Equation]
Here P ( ·| z, a ) is a learned latent transition model and W 2 refers to the 2-Wasserstein metric which we can compute in closed form for Gaussian transition models.
Following Zhang et al.
we use a target encoder updated with the moving average of the encoder’s weights for producing ¯ z and we add an auxiliary reward prediction objective with per-task reward predictors ˜
r i bisim ( z ) .
Imitation Learning (TARP-BC).
We can also train task-induced representations from data without reward annotation by directly imitating the data collection policy, thus learning to represent all elements of the state space that were important for the policy’s decision making.
We choose behavioral cloning (BC, Pomerleau (1989)) for imitation learning since it is easily applicable especially in the ofﬂine setting.
We introduce N separate imitation policy heads h i BC and minimize the negative log-likelihood of the training data’s actions:
[Equation]
E XPERIMENTS 5. E XPERIMENTAL S ETUP
We compare the performance of different representation learning approaches in three visually complex environments (see Figure 3).
We will brieﬂy describe each environment; for more details on environment setup and data collection, see appendix, Section A.
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) 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.
Task-induced representations should focus on modeling the agent while ignoring irrelevant information from the background video.
ViZDoom.
A simulator for the ego-shooter game Doom (Wydmuch et al., 2018).
We use pretrained models from Dosovitskiy and Koltun (2017) for data collection and vary their objectives to get diverse behaviors 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).
Autonomous Driving.
We simulate realistic ﬁrst-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 efﬁcient learning, representations need to model relevant driving information such as position and velocity of other cars, while ignoring irrelevant details such as trees, shadows or the color and make of cars in the scene.
We compare the different instantiations of the task-induced representation learning framework to prior representation learning approaches:
• Reconstruction : Trains a beta- VAE (Higgins et al., 2017) or stochastic video prediction model (“ Pred-O ”) on the training data and transfers the encoder.
• Reconstruction + Task-Induced : Combines video prediction with reward prediction for pre-training (“ Pred-R+O ”), similar to Hafner et al.
• Contrastive Learning :
Uses the contrastive learning objective from Stooke et al.
(2021) for pre-training (“ ATC ”).
All approaches are trained on the same ofﬂine dataset D used for training of the TARP approaches.
We also compare to a policy transfer baseline which pre-trains a policy on D using BC and then ﬁnetunes the full policy on the downstream task as an alternative to representation transfer 1 .
In distracting DMControl we further report results for an oracle baseline whose representation is trained with direct supervision from the low-dimensional state representation 2 .
After pre-training of the representation, we transfer the frozen encoder weights to the target task policy, which we train with soft actor-critic (SAC, Haarnoja et al.
(2018a))
on continuous control tasks (distracting DMControl, CARLA), and with PPO (Schulman et al., 2017) on discrete action tasks (ViZDoom).
For more implementation details and hyperparameters, see Appendix E.
5. D OWNSTREAM T ASK L EARNING E FFICIENCY
We report downstream task learning curves for task-induced and unsupervised representation learning methods as 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 signiﬁcantly different behaviors than those observed in the pre-training data, a scenario in which transferring representations can be beneﬁcial over directly transferring behaviors 3 .
All reconstruction-based approaches (green) struggle with the complexity of the tested scenes, especially the method that attempts to predict the scene dynamics ( Pred-O ). We ﬁnd that adding reward prediction to the objective improves performance (red), especially the method that predicts the scene dynamics with the task-induced objective ( Pred-O+R ).
Yet, we ﬁnd that downstream learning is still slow, possibly because the reconstruction objective leads to task-irrelevant agent appearance information being modeled in the learned representation.
The non-reconstructive contrastive approach ATC (brown) achieves stronger results, particularly in the VizDoom environment which features less visual distractors, but its performance deteriorates substantially in environments with more task-irrelevant details, i.e. distracted DMControl and CARLA. [Figure]
Overall, we ﬁnd that task-induced representations enable more sample efﬁcient learning.
Among the
TARP instantiations, we ﬁnd that TARP-V and TARP-Bisim representations can lead to lower transfer performance, since they rely on the expressiveness of the reward function: 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 efﬁcient transfer to downstream policy learning problems.
On the distracting DMControl task we ﬁnd that TARP representations even outperform representations trained with direct supervision through a handcrafted 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. P ROBING T ASK -I NDUCED R EPRESENTATIONS
To better understand the improved learning efﬁciency of the TARP approaches over the unsupervised methods, we visualize what information is captured in the representation in Figure 5: we compare input saliency maps 4 for representations learned with task-induced and unsupervised objectives.
We ﬁnd that task-induced representations can focus only 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 task-irrelevant aspects such as changing background videos, buildings and trees, since they cannot differentiate task-relevant and irrelevant information.
To quantitatively analyze whether task-induced representations effectively represent task-relevant information and discard task-irrelevant features, we train probing networks on 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-irrelevant 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 at classifying the video.
We report the probing networks’ state prediction error in Figure 6a and the background classiﬁcation accuracy in Figure 6b to compare representations learned via TARP to unsupervised objectives (we choose VAE and ATC as comparisons since they achieved the best transfer performance among reconstruction and non-reconstruction baselines respectively).
From this comparison we can see that probing networks trained on top of task-induced representations can more accurately predict the task-relevant state information while they are successfully ﬁltering information about the background video and thus obtain lower background classiﬁcation accuracy.
Therefore, these quantitative results support the qualitative intuition from Figure 5. [Figure] [Figure]
to struggle to achieve and thus cannot conﬁdently classify the background video, while unsupervised approaches fail to ﬁlter the irrelevant information and thus achieve perfect classiﬁcation scores.
(c) Transfer performance for IL in distracting DMControl. Task-induced representations achieve superior sample efﬁciency. 5. T RANSFER R EPRESENTATION FOR I MITATION L EARNING
We show that the representations learned via task-induced representation learning are general and can be used for multiple downstream learning approaches: in addition to the model-free RL results above we demonstrate their effectiveness for imitation learning (IL).
We train policies with the representations learned on the same datasets used in Section5.2 by Soft-Q Imitation Learning (SQIL, Reddy et al.
on the distracting DMControl target running task.
In Figure 6c we show that taskinduced representations also improve downstream performance in visually complex environments for IL and allow for more efﬁcient imitation, since they model only task-relevant information.
Again, TARP-BC and TARP-CQL lead to the best learning efﬁciency.
This shows that the beneﬁt of TARP is not constrained to RL, but the same representations can be used to accelerate IL.
5. D ATA A NALYSIS E XPERIMENTS
Task-induced representation learning approaches are designed to leverage data from a 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 use – our goal in this section is to derive a set of best practices when collecting datasets for task-induced representation learning.
Many tasks vs. lots of data per task.
When collecting large datasets there is a trade-off between collecting a lot of data for few tasks vs. collecting less data each for a larger set of tasks.
To analyze the effects of this trade-off on TARP, we train TARP-BC on data from a varying number of pre-training tasks in distracting DMControl.
When training on data from fewer tasks, we collect more data on each of these tasks, to ensure that the size of the training datasets is constant across all experiments.
The results in Figure 7 show: training on fewer data from a larger number of tasks is beneﬁcial 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]
More data vs. optimal data.
Another trade-off in data collection exists between the amount of data and its optimality: a method that can learn 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 DMControl.
The downstream RL performance comparison in Figureshows that TARP approaches can learn strong representations from low-performance trajectory data and even when trained from random data does not decrease performance over a SAC baseline trained from scratch because, intuitively, task-induced representation learning does not pre-train a model on “how to act” but merely “what to pay attention to”.
We ﬁnd that TARP-CQL’s performance can even increase slightly with the suboptimal data, which we attribute to 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 . [Figure]
Multi-Source Task Supervision.
When collecting large datasets, it can be beneﬁcial to pool data from multiple sources, e.g. data obtained from prior training runs or 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 we can train a single representation using multiple sources of task supervision simultaneously .
In particular, we train “TARP-V+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 propagate gradients from all prediction heads into the same representation and compare downstream learning efﬁciency on distracting DMControl and CARLA in Figure 9.
We ﬁnd that the combined-source model trained from the heterogeneous dataset achieves comparable or superior performance to all single-source models, showing that practicioners should collect diverse datasets, even if they have heterogeneous sources of task-supervision .
C ONCLUSION
In this work, we formalize the problem of task-induced representation learning and analyze multiple approaches for leveraging ofﬂine experience of prior tasks for learning compact representations in visually complex scenes.
Our empirical results show that task-induced representations retain task-relevant information and discard irrelevant features, leading to improved sample efﬁciency on downstream task compared to prior unsupervised representation learning approaches.
Future work should investigate approaches for incorporating other sources of task information such as language commands to learn task-induced representation.
R EPRODUCIBILITY S TATEMENT
We perform all of our experiments on publicly available environments.
We provide a detailed description of the procedures used for the ofﬂine dataset collection of prior tasks in Appendix A.
Furthermore, in Appendix E we list all of the hyperparameters used for the pre-training phase (i e. task-induced and unsupervised representation learning) as well as training RL policies.
We include our codebase with example commands to reproduce our results in the supplementary materials.  A CKNOWLEDGEMENT
We thank our colleagues from the CLVR lab for the valuable discussions that considerably advanced the research.
A E NVIRONMENTS A. D ISTRACTING DMC ONTROL
We increase the visual complexity of DMControl “walker” task (Tassa et al., 2018) by overlaying randomly sampled videos from the DAVIS 2017 dataset (Pont-Tuset et al., 2017), similar to Stone et al.
In our experiments we use expert policies to collect ofﬂine 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., 2018a) in the state space and collect rollouts of the visual observation.
We test our representation on the downstream task of “running”.
Rewards: We use reward functions provided by DMControl “walker” task (Tassa et al., 2018).
For the backward walking task, we invert the sign of walk speed in “forward” task deﬁned in DMControl so that the agent gets higher reward when the agent moves backward instead of forward.
A. V I ZD OOM
Our experiments use “D3 battle”environment provided by Wydmuch et al.
(2018) where the agent’s objective is to defend against enemies while collecting medi-packs and ammunition.
For collection of the prior task dataset, we use the pre-trained models provided by Dosovitskiy and Koltun (2017) and vary the reward weighting parameters (described below) to produce a diverse set of behaviours.
Rewards: A reward function is deﬁned by a linear combination of three measurements (ammunition, health, and frags) with their corresponding coefﬁcients.
[Equation]
where x ammt is a measurement of ammunition, x healtht is health of the agent, and x fragst is a number of frags at timestep t .
The set of coefﬁcients for ammunition, health and frags are represented as ( c ammo , c health , c frags ).
We use the coefﬁcients of (0 , 0 , 1) , (0 , 1 , 0) , and (1 , 1 , − 1) for the prior tasks, and (0 . 5 , 0 .
for the target task.
A. CARLA
In our experiments, we use the map of “Town05” from the CARLA environment (Dosovitskiy et al., 2017).
At the beginning of each episode, we randomly spawn 300 vehicles and 200 pedestrians.
The initial location of the agent is randomly sampled from a task set containing multiple start and goal locations.
We collect the datasets for the tasks of intersection crossing, taking a right turn, and taking a left turn using pre-trained policies.
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 to collect rollouts of visual observations for the datasets.
Rewards: We use the same reward function for all of the tasks.
The reward function consists of terms for speed, centering on a road, angle of the agent, and collision.
[Equation]
(6) where v is velocity of the agent, d center is distance between the center of the road and the agent, r is angle of the agent.
We use constant values of v min = 15 . 0 , v max = 30 . 0 , v target = 25 . 0 , d max = 3 , and r max = 20 .
[Figure]
Our task-induced representation with multiple supervision performs slightly better or as good as the single source representation.
Thus, task-induced representation learning framework is compatible with heterogeneous datasets.
B D ATA C OLLECTION FOR M ULTI -S OURCE T ASK S UPERVISION
We use supervision from heterogeneous data sources to learn task-induced representations.
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 the datasets annotated with rewards for the “right turn” and “left turn” tasks.
For both environments, the datasets are collected with the procedure described in Appendix A.
C F INETUNING LEARNED REPRESENTATIONS
In our experimental evaluation in Section 5 we held the parameters of the pre-trained encoder ﬁxed to cleanly evaluate the quality of the pre-trained representation.
However, in practice pretrained representations are often ﬁnetuned on the target task using target task rewards.
Prior work on representation learning found mixed results when ﬁnetuning the pre-trained representations (Yang and
Nachum, 2021).
Thus, in this section we experimentally compare the different representation learning approaches on the Distracting DMControl environment while ﬁnetuning the learned representations on the target task .
As illustrated in Figure 10(a), we ﬁnd that task-induced representations show improved sample efﬁciency and better performance even in the ﬁne-tuning setting, analogous to the results in Section 5.2.
D T RANSFER P ERFORMANCE FOR DIFFERENT PRIOR AND TARGET TASK SETS
We experimentally validate that we can interchange pre-training vs target task sets.
Figure 10(b) shows the downstream performance on the "walk" task with representations pre-trained on "run", "backward walking", and "standing" tasks.
This quantitative result underlines that task-induced representation lead to more efﬁcient downstream training than unsupervised methods – the same conclusions as in Section 4 apply in this scenario.
[Figure]
Scratch
Pred-O
VAE Reconstruction
ATC Contrastive
Task
(b) ViZDoom (c) CARLA (a)
Distracting DMControl Oracle
Policy Transfer
Reconstruction
[Figure]
Induced
[Figure]
E H YPERPARAMETERS E. H YPERPARAMETERS FOR RL
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
E
H YPERPARAMETERS FOR
PRE
- TRAINING
In TARP-BC, we use recurrent neural networks to predict a sequence of actions over prediction horizon T to learn task-induced representations only in CARLA (see Table 7)
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]