Task-Induced Representation Learning

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 can not 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 Task-Induced . 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 can not 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 can not 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 can not 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

Summary: 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.