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 work proposes using instruction modelling as an auxiliary objective to improve long-term planning. On BabyAI and Crafter, this method demonstrates significant improvement on longer hop levels (e.g. 3 hop BabyAI).