Improving Long-Horizon Imitation Through Language Prediction

1 INTRODUCTION . Intelligent agents ought to be able to complete complex , long horizon tasks and generalize to new scenarios . Unfortunately , policies learned by modern deep-learning techniques often struggle to acquire either of these abilities . This is particularly true in planning regimes where multiple , complex , steps must be completed correctly in sequence to complete a task . Realistic constraints , such as partial observability , the underspecification of goals , or the sparse reward nature of many planning problems make learning even harder . Reinforcement learning approaches often struggle to effectively learn policies and require billions of environment interactions to produce effective solutions ( Wijmans et al. , 2019 ; Parisotto et al. , 2020 ) . Imitation learning is an alternative approach based on learning from expert data , but can still require millions of demonstrations to learn effective planners ( Chevalier-Boisvert et al. , 2019 ) . Such high data constraints make learning difficult and expensive . Unfortunately the aforementioned issues with behavior learning are only exacerbated in the low data regime . First , with limited training data agents are less likely to act perfectly at each environment step , leading to small errors that compound overtime in the offline setting . Ultimately , this leads to sub-par performance over long horizons that can usually only be improved by carefully collecting additional expert data ( Ross et al. , 2011 ) . Second , deep-learning based policies are more likely to overfit small training datasets , making them unable to generalize to new test-time scenarios . On the other hand , humans have the remarkable ability to interpolate previous knowledge and solve unseen long-horizon tasks . After observing an environment , we might deduce plan or sequence of the steps to follow to complete our objective . However , imitation learning agents are not required to construct plans by default – they are trained to only output the direct next action given seen observations . This begs the question : how can we make agents reason better in long-horizon tasks ? An attractive solution lies in language instructions , the same medium humans use for mental planning ( Gleitman & Papafragou , 2005 ) . Several prior works directly provide agents with language instructions to follow ( Anderson et al. , 2018 ; Shridhar et al. , 2020 ; Chen et al. , 2019 ) . Unfortunately , such approaches require the specification of exhaustive instructions at test time for systems to function . A truly intelligent agent ought to be able to devise its own plan and execute it , with only a handful of demonstrations . We propose improving policy learning in the low-data regime 1Code is available at : < github will be made public after reviewing period > by having agents predict high-level planning instructions in addition to their immediate next action . As we do not input instructions to the policy , we can plan without their specification at test time . Though prior works have used hierarchical structures that generate their own instructions to condition on ( Chen et al. , 2021c ; Hu et al. , 2019 ; Jiang et al. , 2019 ) , we surprisingly find that just predicting language instructions is in itself a powerful objective to learn good representations for planning . Teaching agents to output language instructions for completing tasks has two concrete benefits . First , it forces them to learn at a higher level of abstraction where generalization is easier . Second , by outputting multi-step instructions agents explicitly consider the future . Practically , we teach agents to output instructions by adding an auxiliary instruction prediction network to transformer-based policy networks , as in seq2seq translation ( Vaswani et al. , 2017 ) . Our approach can be interpreted as translating observations or trajectories into instructions . We test our representation learning method in limited data settings and combinatorially complex enviornments . In the BabyAI benchmark , we attain near the same performance as Chevalier-Boisvert et al . ( 2019 ) with only 5 % of the data . To our knowledge , we achieve the highest success rate in the Crafting environment from ( Chen et al. , 2021c ) , while using fewer demonstrations and no reinforcement learning . We find that in many settings higher performance can be attained by relabeling existing demonstrations with language instructions instead of collecting new ones , creating a new , scalable type of data collection for practitioners . Furthermore , our method is conceptually simple and easy to implement . This work is the first to show that direct representation learning with language can accelerate imitation learning . To summarize , our contributions are as follows . First , we introduce a method for training transformer based planning networks on paired demonstration and instruction data via an auxiliary language prediction loss . Second , we test our objective in long-horizon planning based environments with limited data and find that it substantially outperforms contemporary approaches . Finally , we analyze the scenarios in which predicting instructions provides fruitful training signal , concluding that instruction modeling is a valuable objective when tasks are sufficiently complex . 2 RELATED WORK . Language in the context of policy learning has been heavily studied ( Luketina et al. , 2019 ) , usually to communicate a task objective . Uniquely , we use natural language instructions to aid in learning via an auxiliary objective . Here we survey the most relevant works to our approach . Language Goals . Language offers a natural medium to communicate goals to intelligent agents . As such , several prior work have focused on learning language goal conditioned policies , particularly for robotics ( Nair et al. , 2021 ; Stepputtis et al. , 2020 ; Kanu et al. , 2020 ; Hill et al. , 2020 ; Akakzia et al. , 2021 ; Goyal et al. , 2021 ; Shridhar et al. , 2021a ) , or for games ( Chevalier-Boisvert et al. , 2019 ; Chaplot et al. , 2018 ; Hermann et al. , 2017 ) . Others in the area of inverse reinforcement learning use language to specify reward functions ( Fu et al. , 2019 ; Bahdanau et al. , 2018 ; Williams et al. , 2018 ) or shape them ( Mirchandani et al. , 2021 ; Goyal et al. , 2019 ) . Unlike these works , we use language instructions that dictate how an agent should complete a task instead of language goals that specify what the task itself is . Other works , particularly in the visual navigation space , provide agents with instructions similar to those we use , sometimes in addition to language goals . Anderson et al . ( 2018 ) ; Fried et al . ( 2018 ) ; Chen et al . ( 2019 ) ; Krantz et al . ( 2020 ) ; Chen et al . ( 2021a ; 2019 ) use instructions for visual navigation or drones ( Blukis et al. , 2019 ) , while Shridhar et al . ( 2020 ) ; Pashevich et al . ( 2021 ) ; Shridhar et al . ( 2021b ) use instructions for household tasks . Critically unlike our method , these approaches use both language goals and instructions as input , and consequently require the creation of instructions at test-time for every desired task . Other proposed environments ( Zhong et al. , 2019 ; Wang & Narasimhan , 2021 ) assess understanding by prompting agents with necessary information about task dynamics , precluding the removal of text-prompting at test time . Language and Hierarchical Learning . Instead of directly using instructions as policy inputs , other works use language instructions as an intermediary representation for hierarchical policies . Usually , a high-level planner outputs language instructions for a low-level executor to follow . Andreas et al . ( 2017 ) and Oh et al . ( 2017 ) provide agents with fixed , hand-designed high-level language instructions or policy “ sketches ” . Such approaches require new instruction labels at test-time for every new task unlike our method . Jiang et al . ( 2019 ) and Shu et al . ( 2017 ) provide interactive language labels to agents to train hierarchical policies with reinforcement learning . In the imitation learning setting , Hu et al . ( 2019 ) learn a hierarchical policy using behavior cloning for a strategy game . Unlike the planning problems we consider , their environment has no oracle solution and does not consider generalization to unseen tasks . Most related to our work , Chen et al . ( 2021c ) use latent representations from a learned high-level instruction predictor to aid a low-level policy . However , unlike Chen et al . ( 2021c ) , we learn latent representations that can predict instructions , but do not explicitly condition on them at test-time . While hierarchical approaches have shown great promise , the quality of learned policies is inherently limited by the amount of language data available for training . Even with a perfect low-level policy , inaccurate high-level languages commands will yield poor overall performance . This is not an issue for our loss-based approach , as our instruction prediction network can be completely detached from the policy . Additionally , this allows our method to work on a mix of instruction annotated and unannotated data , letting it more easily scale than hierarchical approaches particularly in data-limited scenarios . Auxiliary Objectives . The learning community has extensively studied the use of auxiliary objectives in policy learning . Though to our knowledge no prior works use instructions as an auxiliary objective , auxiliary objectives in general have been found to aid policy learning ( Jaderberg et al. , 2017 ) . Laskin et al . ( 2020 ) and Stooke et al . ( 2021 ) demonstrated the success of contrastive auxiliary objectives in robotic reinforcement learning domains . Schwarzer et al . ( 2020 ) and Anand et al . ( 2019 ) did the same in the Atari game-playing environments . We were inspired by their effectiveness . Additionally , works like Andreas et al . ( 2018 ) have previously used language question and answering for representation leaning in visual domains . Transformers . Our approach is based on several innovations involving transformer networks . Vaswani et al . ( 2017 ) previously showed state of the art results in machine translation using transformers . While the application of transformers has extended to behavior learning ( Zambaldi et al. , 2018 ; Parisotto et al. , 2020 ; Chen et al. , 2021b ) , prior works in the area have not leveraged the transformer decoder . Closest to our domain , Lin et al . ( 2021 ) generate captions from video . The architecture of our policy networks take inspiration from recent works adapting transformers to mediums beyond text , namely in vision ( Kolesnikov et al. , 2021 ) and offline reinforcement learning ( Chen et al. , 2021b ) . 3 METHOD . In this section we formally describe the problem of imitation learning with instruction prediction , then describe our implementation for both Markovian and non-Markovian environments . 3.1 PROBLEM SETUP . The standard learning from demonstrations setup assumes access to a dataset of expert trajectory sequences containing paired observations and actions o1 , a1 , o2 , a2 , ... , oT , aT . The goal of imitation learning is to learn a policy π ( at|· ) that predicts the correct actions an agent should take . In our work we consider both fully observed ( MDP ) and partially observed ( POMDP ) settings . In the partially observed case , policies are given access to previous observations in order to infer state ( Kaelbling et al. , 1998 ) , and we denote the policy as π ( at|o1 , ... ot ) . In the fully-observed setting this is unnecessary , and the policy is simply π ( at|ot ) . It is common for policies to be goal conditioned , or even conditioned on language goals as is the case in our experiments . This means they take an encoding of the desired task or goal g as input . As our approach works with or without goal conditioning we omit it from the rest of this section for brevity . A standard imitation learning technique is behavior cloning , which in discrete domains maximizes the likelihood of the actions in the dataset using a negative log likelihood objective , Laction = − ∑ t log π ( at|· ) . In this work , we assume access to oracle language instructions that tell an agent how it should complete a task to provide useful training signal . As mentioned in Section 2 , for the purposes of our method we distinguish goals from instructions . Language goals tell an agent what to do , whereas language instructions communicate how an agent should complete a task . Each trajectory may have several language instructions x ( 1 ) , x ( 2 ) , ... , x ( n ) corresponding to different steps in the task . For example , a language instruction like “ open the door ” only applies to the part of the demonstration before the agent opens the door and after it completes the last instruction . The i-th instruction x ( i ) thus corresponds to an interval [ Ti , Ti+1 ) where Ti marks the time the instruction was given and Ti+1 denotes the start of the next instruction . A depiction of an example instruction sequence can be found in Figure 1 . While language instructions are an additional data requirement , they can be cheap to obtain , particularly in scenarios where demonstrations are expensive to collect . If one has to collect demonstrations in the real world , providing instructions while they are collected likely only constitutes a marginal increase in effort . Moreover , humans can easily re-label existing demonstrations with instructions . Video data could easily be captioned with voice-over . Similar statements can be made for simulators – if one can code an oracle policy , instructions are likely easy to generate along the way . These modifications can easily be done on simulators with planning stacks , as we do with BabyAI . Moreover , we focus on the data limited regime , where the cost of setting up an environment and collecting more demonstrations is likely be higher than annotating an existing small set of demonstrations . Next , we describe how we train agents to predict language instructions to aid in imitation learning .

Summary: This paper investigates the usage of language prediction for imitation learning. The authors distinguish between a goal (a short description of the task) and an instruction (a detailed description of subtasks that span a sequence of observations). They show that predicting instructions at each step as an auxiliary loss improves the performance using lesser number of demonstrations. They experiment with minigrid and crafting environments to show the efficacy of the proposed approach.