ELLA: Exploration through Learned Language Abstraction

1 Introduction . A long-standing goal for robotics and embodied agents is to build systems that can perform tasks specified in natural language [ 1 , 8 , 20 , 25 , 38 , 41 , 42 ] . Central to the promise of language is its ability to cleanly specify complex , multi-step instructions . Instructions like make a cup of coffee define long-horizon tasks as abstractions over lower-level components—simple instructions like pick up a cup or turn on the coffee maker . Leveraging these abstractions can help amplify the sample efficiency and generalization potential of our autonomous agents . One way to do this is through the lens of instruction following , which can be framed in several ways . One common framing—and the one we use in this work—is via reinforcement learning ( RL ) : an agent is given a start state , a language instruction , and a corresponding reward function to optimize that usually denotes termination [ 20 , 30 ] . While RL can be a useful framing , such approaches are often not sample efficient [ 11 , 20 ] . Especially in the case of complex , highly compositional language instructions , RL agents can fail to make progress quickly—or at all . There are several reasons for poor performance in these settings ; paramount is that in many environments , these instructions are 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . tied to sparse reward functions that only provide signal upon completion of the high-level task , which drastically hurts sample efficiency . For example , seeing reward for make a cup of coffee would require already having turned on the coffee maker , which could itself be a difficult task . Such “ bottleneck '' states , which have zero intermediate reward , complicate exploration and learning [ 29 , 39 ] . Our goal is to improve sample efficiency for instruction following agents operating in sparse reward settings . Driving our approach is using the principle of abstraction—the fact that complex instructions entail simpler ones—to guide exploration . Consider our coffee example ; we would like to guide the agent to explore in a structured fashion , learning low-level behaviors first ( pick up a cup , turn on the coffee machine ) and building up to solving the high-level task . To do this , we frame our problem via reward shaping , the general technique of supplying auxiliary rewards to guide learning [ 31 , 35 ] . Our approach , Exploration through Learned Language Abstraction ( ELLA ) provides intermediate rewards to an agent for completing relevant low-level behaviors as it tries to solve a complex , sparse reward task . Notably , our approach 1 ) learns to identify the low-level primitives helpful for a high-level language task online , and 2 ) does not require a strict hierarchical decomposition of language instructions to these primitives . Rather , ELLA uses low-level instructions to support agents performing complex tasks , bonusing agents as they complete relevant behaviors . In contrast to prior work [ 3 , 23 ] that assumes strict contracts over instruction decomposition , our contract is simple and general . This prior work assumes each high-level task is comprised of exact series of low-level policies—an assumption that is only possible when the full set of primitives is known ahead of time , and fails when new actions or further exploration are necessary . In order to bonus agents as they complete relevant subtasks , ELLA assumes access to a set of lowlevel instructions and corresponding termination states , similar to the data assumed in prior work [ 6 ] . While collecting such data for complex tasks may require time and effort , annotations for low-level instructions are more tractable . Humans can quickly annotate instructions like pick up the cup or open the cupboard , or even build tools for generating such examples . This is increasingly apparent in prior work [ 5 , 22 , 25 , 43 ] ; low-level instructions are used to simplify the problem of interpreting high-level instructions , through a variety of mechanisms . In our case , these low-level instructions provide the basis for reward shaping . We empirically validate our abstraction-based reward shaping framework on a series of tasks via the BabyAI platform [ 11 ] . We compare against a standard RL baseline as well as to a strong languagebased reward shaping approach [ 17 ] , and find that our method leads to substantial gains in sample efficiency across a variety of instruction following tasks . 2 Related Work . We build on a large body of work that studies language in the context of reinforcement learning ( RL ) [ 27 ] with a focus on methods for instruction following , leveraging hierarchy , and reward shaping . Instruction Following . Instruction following focuses on the automated execution of language commands in simulated or real environments [ 1 , 16 , 26 , 41 , 44 ] . We focus on the RL setting [ 11 , 30 ] , where agents learn language-conditioned policies in a Markov decision process ( MDP ) . Even simple instruction following problems are difficult to learn [ 28 ] , which has made 2D environments important test-beds for new learning methods [ 3 , 11 , 46 ] . Such environments , which are procedurally-generated , are useful because they decouple exploration from the problem of perception , and demand policies that do not overfit to narrow regions in the state space [ 10 , 13 ] . We use BabyAI [ 11 ] as the platform for our tasks . BabyAI ’ s language is synthetic , similar to prior work examining RL for instruction following at scale [ 20 , 21 ] . These environments are scalable starting points for new instruction following approaches , since deep RL remains sample-inefficient but has the potential to flexibly deal with complex environments [ 21 ] . Simplified environments also make investigating concepts such as language abstraction possible , and these concepts can generalize to real human-AI interactions [ 25 , 43 ] . We describe our experimental setup in Section 5 . Hierarchical and Language-Structured Policies . A large body of work in hierarchical RL studies temporal abstractions of MDPs , where policies are designed or learned at multiple levels of hierarchy [ 40 ] . Recent works have made the connection to language : for example , high-level policies can act over the space of low-level instructions [ 23 ] . Other hierarchical language techniques include learning models that predict symbolic representations of reward functions at different granularities [ 5 , 24 ] , and compositional approaches like policy sketches [ 3 ] and modular control [ 15 ] , methods that explicitly decompose high-level instructions into sequences of predefined low-level instructions , or subtasks . In this work , we use language hierarchically , but without the strict decomposition assumption made by prior work : we do not assume high-level tasks are composed of an exact series of low-level tasks . Without this assumption , it is infeasible to use explicitly modular approaches [ 3 , 5 , 15 ] . A further distinguishing factor of our method relative to prior work is that we do not require high-level task decompositions be provided a priori . We expand on these features in Section 3 . Language and Reward Shaping . Reward shaping is a general technique for supplying auxiliary rewards in order to guide an agent ’ s learning process . Certain types of reward transformations , such as potential-based rewards , do not change the optimal policy but can impact sample efficiency [ 31 ] . Prior work has applied reward shaping to language-based instruction following settings . Goyal et al . [ 17 ] train a classification network on ( trajectory , language ) data to predict if a trajectory ( parameterized by action frequencies ) matches a language description , and evaluate the network at every time step to form a potential-based reward shaping function . They evaluate their method on Montezuma ’ s Revenge , which is a complex task , but with a static , deterministic layout that does not speak to the generalization potential of such methods . Relatedly , Waytowich et al . [ 45 ] use a narration-based approach , where high-level tasks are explicitly broken down into low-level tasks as narrations . Finally , Misra et al . [ 30 ] use a shaping function based on spatial distance to goals specified in language , but require the location of the goal to be known . In this work , we use the same high level principle— language can yield useful intermediate rewards—but do not make the restrictive assumptions of prior work about the environment or the way tasks are structured . Other Forms of Guidance . Several other methods exist for improving the sample efficiency of RL in sparse reward settings . These methods reward an aspect of exploration that is orthogonal to our approach , and we see potential for combining these methods with ours . One popular technique is hindsight experience replay ( HER ) , in which failed trajectories are stored in a replay buffer and relabeled with new goals that get nonzero reward [ 4 ] . Jiang et al . [ 23 ] extend the method to hierarchical language settings , assuming a mapping from states to corresponding goals . Cideron et al . [ 12 ] learn this mapping alongside policy training using the environment reward as supervision . Another technique for improving sample efficiency is through intrinsic motivation [ 9 , 32 , 34 , 36 ] . In general , these methods instantiate reward shaping to incentivize accessing novel , diverse , or unpredictable parts of the state space via intrinsic rewards , and they can be extended to language- conditioned settings as well [ 14 ] . In Appendix D , we evaluate the performance of RIDE [ 34 ] , an intrinsic motivation method that rewards actions that cause significant changes to the state , and discuss potential synergies with ELLA . 3 Problem Statement . We consider an augmented MDPM defined by the tuple ( S , A , T , R , G , γ ) where S , A , T , and γ are standard . A consists of primitive actions—in BabyAI , these are navigation primitives ( forward , pick up , etc. ) . G is a set of language instructions , from which a high-level task instruction g is drawn , and R : S × A × G → [ 0 , 1 ] represents the state-action reward given some g. Via RL , we wish to find some policy π : S × G → A that maximizes the expected discounted return . We specifically consider cases where M has a finite horizon H and R is sparse with respect to goal-completion , making exploration difficult . Our aim is to construct a reward transformation R → R′ which is policy invariant with respect to the original MDP ( i.e . an optimal policy for M′ = ( S , A , T , R′ , g , γ ) is also optimal inM and vice versa ) , while providing the agent strong signal conducive to sample-efficient exploration . We assume access to a set of low-level instructions Gℓ such that every g ∈ G is supported by some gℓ ∈ Gℓ . Note that we do not require that high-level tasks fully factorize into low-level tasks . To clarify the distinction , say that Gℓ consists of go to “ x ” instructions : go to the red ball , go to the blue ball , etc . The instruction go to the red ball and then go to the blue ball can be fully factorized into low-level tasks in Gℓ . On the other hand , the instruction put the red ball next to the green key is supported by go to the red ball , but it also requires picking up and putting down the ball—actions not covered by Gℓ . Our framing permits exploration using a mixture of low-level instructions as well as primitive actions in A ( such as pick up and put down ) . We assume that examples of the corresponding termination states of instructions in Gℓ are easy to obtain—via a human expert or automation . This is reasonable for simple low-level tasks like the go to tasks in Figure 1 . However , it is less feasible in domains where data collection is costly and the environment is hard to simulate ; we address this further in Section 7 . The environments in our experiments are partially observable , so we use recurrent networks that ingest sequences of observations ( o1 , o2 , ... , ot ) rather than a single state st [ 19 ] . Though our notation describes the fully observable setting , the extension to the partially observable case is straightforward .

Summary: The paper present ELLA, a three-part solution to instruction following tasks with a high-level instruction and a pre-defined set of low-level instructions:  1. a termination classifier is trained offline to identify the completion of a low-level goal at a state; 2. a relevance classifier is learned online to predict a binary relevance of a low-level goal to a high-level goal. 3. reward shaping is used to encourage low-level goals suggested by the relevance classifier.  ELLA is shown to improve sample complexity on 5/6 BabyAI tasks of diverse environment and instruction complexities, compared to an ablation PPO model with reward shaping, and a baseline model LEARN.