Improving Policy Learning via Language Dynamics Distillation
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Recent work has shown that augmenting environments with language descriptionsimproves policy learning.
However, for environments with complex languageabstractions, learning how to ground language to observations is difﬁcult dueto sparse, delayed rewards.
We propose Language Dynamics Distillation ( LDD ),which pretrains a model to predict environment dynamics given demonstrationswith language descriptions, and then ﬁne-tunes these language-aware pretrainedrepresentations via reinforcement learning (RL).
In this way, the model is trained toboth maximize expected reward and retain knowledge about how language relatesto environment dynamics.
On SILG, a benchmark of ﬁve tasks with languagedescriptions that evaluate distinct generalization challenges on unseen environments (NetHack, ALFWorld, RTFM, Messenger, and Touchdown), LDD outperforms tabula-rasa RL, as well as methods that learn from unlabeled demonstrationsin inverse RL and reward shaping with pretrained experts.
In our analyses, weshow that language descriptions in demonstrations improve sample-efﬁciencyand generalization across environments, and that dynamics modeling with expertdemonstrations is more effective than with non-experts.
Introduction
Language is a powerful medium that humans use to reason about abstractions—its compositionalityallows efﬁcient descriptions that generalize across environments and tasks.
Consider an agent thatfollows instructions to clean the house (e g. ﬁnd the dirty dishes and wash them).
In tabula-rasareinforcement learning (RL), the agent observes raw perceptual features of the environment, thengrounds these visual features to language cues to learn how to behave through trial and error.
Incontrast, we can provide the agent with language descriptions that describe abstractions which arepresent in the environment (e.g. there is a sink to your left and dishes on a table to your right ), therebysimplifying the grounding challenge.
Language descriptions of observations occur naturally in manyenvironments such as text prompts in graphical user interfaces [Liu et al., 2018], dialogue [He et al.
,2018], and interactive games [Küttler et al., 2020].
Recent work has also shown improvements invisual manipulation [Shridhar et al., 2021] and navigation [Zhong et al., 2021, Tam et al., 2022] bycaptioning the observations with language descriptions.
Despite these gains, learning how to interpret language descriptions is difﬁcult through RL, especially on environments with complex languageabstractions [Zhong et al., 2021].
We present Language Dynamics Distillation ( LDD ), a method that improves RL by learning a dynamics model on cheaply obtained unlabeled (i.e. no action labels) demonstrations with language
Observations from expert
[Figure] [Figure]
descriptions.
When learning how to use language descriptions effectively, one central challengeis how to disentangle language understanding from policy performance from sparse, delayed rewards.
Our motivation is to learn initial language grounding via dynamics modeling from an ofﬂinedataset, away from the credit assignment and non-stationarity challenges posed by RL.
While labeleddemonstrations that tell the agent how to to act in each situation are expensive to collect, for manyenvironments one can cheaply obtain unlabeled demonstrations (e.g. videos of experts performing thetask)
[Yang et al., 2019, Stadie et al., 2017].
Intuitively, LDD exploits these unlabeled demonstrationsto learn how to associate language descriptions with abstractions in the environment.
This knowledgeis then used to bootstrap and more quickly learn policies that generalize to instructions and manualsin new environments.
Given unlabeled demonstrations with language descriptions (e.g. captions ofscene content), we ﬁrst pretrain the model to predict the next observation given prior observations,similar to language modeling.
A copy of this model is stored as a ﬁxed teacher that grounds languagedescriptions to predict environment dynamics.
We then train a model with RL, while distilling intermediate representations from the teacher to avoid catastrophic forgetting of how to interpret languagedescriptions for dynamics modeling.
In this way, the model learns to both maximize expected rewardwhile retaining knowledge about how language descriptions relate to environment dynamics.
We evaluate LDD on the recent SILG benchmark [Zhong et al., 2021], which consists of ﬁve diverse environments with language descriptions including NetHack [Küttler et al., 2020], ALFWorld [Shridharet al., 2021], RTFM [Zhong et al., 2020], Messenger [Hanjie et al., 2021], and Touchdown [Chen et al.
These environments present unique challenges in language-grounded policy-learning acrosscomplexity of instructions, visual observations, action space, reasoning procedure, and generalization.
By learning a dynamics model from cheaply obtained unlabeled demonstrations, LDD consistentlyoutperforms reinforcement learning with language descriptions both in terms of sample efﬁciencyand generalization performance.
Moreover, we compare LDD to other techniques that inject priorknowledge in inverse reinforcement learning [Hanna and Stone, 2017, Torabi et al., 2018, Guo et al.,
2019] and reward shaping with a pretrained expert [Merel et al., 2017].
LDD achieves top performanceon all environments in terms of task completion and reward.
In addition to comparing LDD to othermethods, we ablate LDD to quantify the effect of language observations in dynamics modeling, and theimportance of dynamics modeling with expert demonstrations.
On two environments where we cancontrol for the presence of language descriptions (NetHack game messages and Touchdown panoramacaptions), we show that language descriptions improve sample-efﬁciency and generalization.
Finally,across all environments, we ﬁnd that dynamics modeling with expert demonstrations is more effectivethan with non-expert rollouts.
Related Work
Learning by observing language.
Recent work studies generalization to language instructionsand manuals that specify new tasks and environments.
These settings range from photorealistic/3Dnavigation [Anderson et al., 2018, Chen et al., 2018, Ku et al., 2020, Shridhar et al., 2020] to multihop reference games [Narasimhan et al., 2015, Zhong et al., 2020, Hanjie et al., 2021].
We use acollection of these tasks to evaluate LDD .
There is also work where understanding language is notnecessary to achieve the task, however its inclusion (e.g. via captions, scene descriptions) makeslearning more efﬁcient.
Shridhar et al.
[2021] show that one can quickly learn policies in a simulatedkitchen environment described in text, then transfer this policy to the 3D visual environment.
Zhonget al.
[2021] similarly transform photorealistic navigation to a symbolic form via image segmentation,then learn a policy that transfers to the original photorealistic setting.
In work concurrent to ours, Tamet al.
[2022] generate oracle captions of observations for simulated robotic control and city navigation,which improve policy learning.
LDD is complementary to these—in addition to incorporating languagedescriptions as features, we show that learning a dynamics model from unlabeled demonstrationswith language descriptions improves sample efﬁciency and results in better policies.
Imitation learning from observations.
There is prior work on model-free as well as model-basedimitation learning from observations.
Model-free methods encourage the imitator to produce statedistributions similar to those produced by the demonstrator, for example via generative adversariallearning [Merel et al., 2017] and reward shaping [Kimura et al., 2018].
In contrast, LDD only requiresintermediate representations extracted from an expert dynamics model on states encountered by thelearner, which are cheaper to compute than rollouts from an expert policy.
Model-based approacheslearn dynamics models that predict state-transitions given the current state and an action.
Hanna and
Stone [2017] learn an inverse model to map state-transitions to actions, which is then used to annotateunlabeled trajectories for imitation learning.
Edwards et al.
[2019] learn a forward dynamics modelthat predicts future states given state and latent action pairs.
In contrast, LDD does not assume priorsover the action space distribution.
For instance, on ALFWorld, our method works even though it isimpossible to enumerate the action space.
In our experiments, we extend model-free reward shapingand model-based inverse dynamics modeling to account for language descriptions and compare LDD tothese methods.
Auxiliary feedback in reinforcement learning.
One common use of auxiliary reward in reinforcement learning is to encourage (intrinsic) exploration as judged by the difference between the stateand prior states [Strehl and Littman, 2008], raw visual observations [Jaderberg et al., 2017], learnedagent representations [Raileanu and Rocktäschel, 2020], and random network observations [Burdaet al., 2019].
In contrast, the distillation objective in Language Dynamics Distillation encouragessimilarity to expert behaviour, as opposed to dissimilarity to prior agent experience.
Moreover, thedynamics model is not trained on trajectories obtained from the agent policy but an expert policy.
Inthis way, Language Dynamics Distillation is complementary to intrinsic exploration methods.
Recent work improves policy learning by augmenting environment observations with language descriptions [Shridhar et al., 2021, Zhong et al., 2021, Tam et al., 2022].
For environments with complexlanguage abstractions, however, learning how to associate language to environment observations isdifﬁcult through RL due to sparse, delayed rewards.
In Language Dynamics Distillation ( LDD ), wepretrain the model on unlabeled demonstrations (i.e. no annotated actions) with language descriptions to predict the dynamics of the environment, then ﬁne-tune the language-aware model via RL.
LDD consists of two phases.
In the ﬁrst dynamics modeling phase, we pretrain the model to predictfuture observations given unannotated demonstrations.
We store a copy of the model as a ﬁxedteacher that has learned grounded representations useful for predicting how the environment behavesunder an expert policy.
In the second reinforcement learning phase, we ﬁne-tune the model throughpolicy learning, while distilling representations from the teacher.
This way, the model is trainedto both maximize expected reward and retain knowledge about the dynamics of the environment.
Fig 1 illustrates the components of LDD .
3.1 Background
Markov decision process.
Consider a MDP M = {S , A , P, r,  } .
Here, S and A respectively arethe discrete state and action spaces of the problem.
P ( s t +1 | s t , a t ) is the transition probability oftransitioning into state s t +1 by taking action a t from state s t .
r ( s, a ) is the reward function givensome state and action pair.
is a discount factor to prioritize short-term rewards.
Actor-critic methods for policy learning.
In RL, we learn a policy ⇡ ( s ; ✓ ) that maps fromobservations to actions ⇡ : S ! A .
Let R ( ⌧ ) denote the total discounted reward over the trajectory⌧ .
The objective is to maximize the expected reward J ⇡ ( ✓ ) =
E ⇡
[ R ( ⌧ )] following the policy ⇡ byoptimizing its parameters ✓ .
For trajectory length T , the policy gradient is
[Equation]
[Equation]
1) where G t =
{+P 1 k =
+}
k r t + k [-+
-]+1 is the return or discounted future reward at time t .
We considerthe actor-critic family of policy gradient methods, where a critic is learned to reduce variance inthe gradient estimate.
Let V ( s ) =
E ⇡ [ G t | s t = s ] denote the state value, which corresponds to theexpected returns by following the policy ⇡ from a state s .
Actor critic methods estimate the state valuefunction by learning another parametrized function V to bootstrap the estimation of the discountedreturn G t .
For instance, with one-step bootstrapping, we have G t ⇡ r t {++
+} +  V
( s t +1 ;  ) .
The criticobjective is then J V (  ) =
12 ( r t +1 +  V
( s t +1 ;  )  V ( s t ;  )) {+
+}
We minimize a weighted sum of thepolicy objective and the critic objective J ac ( ✓ ,  ) =
J ⇡ ( ✓ ) + ↵ V J V (  ) .
3.2 Dynamics modeling during pretraining
In addition to policy learning, Language Dynamics Distillation learns a dynamics model from unlabeled demonstrations to initialize and distill into the policy learner.
Consider a set of demonstrationswithout labeled actions T  = { ⌧ 1 , ⌧ 2 , . . .
⌧ n } obtained by rolling out some policy  ( a t , s t ) , whereeach demonstration ⌧ =
[ s 1 , s 2 , . . .
s T ] consists of a sequence of observations.
We learn a dynamicsmodel  ( s 1 . . .
s t ; ⇣ ) to predict the next observation s t +1 given the previous observations.
[Equation]
where sim is a differentiable similarity function between the predicted state  ( s 1 , . . .
s t ) and theobserved state s t
+1 .
In the environments we consider, sim is the cross-entropy loss across a grid ofsymbols denoting entities present in the scene.
3.3 Dynamics distillation during policy learning
[Figure]
Fig 1 shows the decomposition of the model into a representation network f rep , a policy head f ⇡ , avalue head f V , and a dynamics head f  .
The three heads share parameters because their inputs areformed by the same representation network.
[Equation]
In the ﬁrst phase of dynamics modeling, we pretrain the model to predict future observationsgiven demonstrations by optimizing J  .
We then store a copy of the model as a ﬁxed teacher˜  ( s t ) =
˜ f  ⇣ ˜
[Equation] ⇣ s
[Equation] ˜ ✓ [Equation] ⌘⌘ .
Let | X  Y | denote the L2 distance between X and Y .
During thesecond phase, in addition to policy learning, we optimize a distillation objective J d ( ✓ rep , ✓ ⇡ , ✓ V ) =|
˜ f
rep ⇣ s t ; ˜ ✓ rep ⌘  f rep ( s t ; ✓ rep ) | to avoid catastrophic forgetting of how to interpret language descriptions for dynamics modeling.
This quantity is the similarity between the feature representationproduced by the ﬁxed teacher (e.g. the pretrained dynamics model) and the feature representationproduced by the model.
Because ˜  is frozen, the parameters ˜ ✓ rep are not included in the objectivefunction J d .
The joint loss for Language Dynamics Distillation is then
[Equation]
To summarize, using unlabeled demonstrations with language descriptions, LDD learns a dynamicsmodel of the environment that grounds language descriptions to environment observations (Section 3.2).
This prior knowledge is then injected into reinforcement learning via initialization anddistillation (Section 3.3).
Experiments
We evaluate Language Dynamics Distillation on the Situated Interactive Language Grounding benchmark (SILG)
[Zhong et al., 2021].
SILG consists of ﬁve different language grounding environmentswith diverse challenges in term of complexity of observation space, action space, language, andreasoning procedure.
The ﬁve different environments are as follows.
4.1 Environments
NetHack [Küttler et al., 2020]: The agent must descend a procedurally generated dungeon.
Itsprimarily challenge is in large state space and partial observability, as the map remains obscured until exploration.
Observations include a symbolic grid of entity IDs, in-game message, and descriptionof character stats.
The agent chooses among a ﬁxed set of actions such as movement, pickingup/buying/selling items, and attacking.
We evaluate not on SILGNethack but the full NetHackchallenge, a difﬁcult game for humans with ⇠ 15% expert win rate [Street, 2013].
SymTouchdown [Zhong et al., 2021]: A symbolic version of Touchdown [Chen et al., 2018]where the agent navigates segmentation maps of Street View panoramas following long instructions.
The primary challenge is reading long, natural instructions that describe photorealistic images.
Evaluation is on new navigation instructions.
Observations include a grid of segmentation class IDscorresponding to a discretized Google Street View panorama, synthetic captions of where objects arerelative to the agent (e.g. to your right you see a lot of road and some cars ), and language navigationinstructions.
The agent selects from a list of radial directions to proceed to the next panorama.
ALFWorld [Shridhar et al., 2021]: The agent navigates and manipulates objects inside a kitchenwhich is described via textual descriptions.
ALFWorld is challenging due to its large (>50) text actionspace that vary across scenes.
Evaluation is on unseen instructions.
Observations include a textualdescription of the scene and language goals (e.g. put a clean sponge on the metal rack ).
The agentchooses among a variable set of language actions (e.g. open drawer 1 ).
RTFM [Zhong et al., 2020]: The agent interprets a game manual and instruction to acquire thecorrect items to ﬁght the correct monsters.
Its main challenge is in multi-step reasoning that combinesworld observations with texts describing multiple entities.
Evaluation is on a set of manuals distinctfrom those in training; hence the agent cannot memorize training manuals and must learn to readcorrectly in order to generalize.
Observations include a symbolic grid containing names of entitiespresent, manual describing high level game rules, agent inventory, and the language instruction.
Theagent chooses among a ﬁxed set of movements.
We train and evaluate on the ﬁrst curriculum stage.
Messenger [Hanjie et al., 2021]: The agent delivers a message from a source entity to a targetwhile avoiding an enemy.
The entities are referred to in text by many names, which have no lexicaloverlap with their symbol ID, hence the core challenge is in mapping language entity referencesin text to observed symbolic entity IDs.
Evaluations are on new entity-role assignments (e.g. whocarries the message).
Observations include a symbolic grid containing symbol IDs of entities present,and a manual of entities and roles.
The agent chooses among a ﬁxed set of movements.
We train andevaluate on the second curriculum stage.
[Figure]
Reinforcement learning with language descriptions from scratch.
We train a base tabula-rasapolicy learner from random initialization.
For NetHack, we train the base policy learner from [Küttleret al., 2020] using moolib [Mella et al., 2022].
For RTFM, ALFWorld, and Touchdown, we train the
SIR model from [Zhong et al., 2021] using Torchbeast [Küttler et al., 2019, Espeholt et al., 2018].
For
Messenger, we train the EMMA model from Hanjie et al.
[2021] using PPO [Schulman et al., 2017].
Language Dynamics Distillation ( LDD )
We train LDD variants of the baseline policy learners foreach environment, where we pretrain the model to perform dynamics modeling on unannotateddemonstrations.
For NetHack, we use 100k screen-recordings (where actions are not annotated andcannot be trivially reverse engineered due to ambiguity in observations) of human-playthroughsfrom the alt.org NetHack public server.
For Touchdown, we use unannotated demonstrations byhuman players.
For ALFWorld, we use trajectories obtained from an A* planner with full state andgoal knowledge, with actions removed.
For RTFM and Messenger, we train expert trajectories untilconvergence, and sample 10k rollouts from the experts from which we remove action labels.
Reward shaping with expert.
Methods such as Merel et al.
[2017] reward shape with an expert byencouraging the agent to produce states similar to a demonstrator.
To compare LDD to this idea, we usethe dynamics model to predict the next observation under expert policy.
The difference (e.g. accuracyin symbol prediction across grid) between the predicted observation and the actual observation after taking the action proposed by the agent is used as a penalty (e.g. negative auxiliary reward).
Thismethod is listing as reward in experiment ﬁgures and uses the same unlabeled demonstrations as LDD .
Inverse reinforcement learning.
Another class of methods learn a inverse dynamics model withwhich infer actions in unlabeled demonstrations, then learn to imitate the pseudo-labeled demonstrations
[Hanna and Stone, 2017, Torabi et al., 2018].
To compare to this method, we ﬁrst train the basepolicy learner for 10k episodes, then collect 10k rollouts to train an inverse dynamics model thatpredicts current action given current and future observations.
This inverse model is used to annotatethe original unlabeled demonstrations for imitation learning [Torabi et al., 2018].
Because theseimitation policies do not generalize to novel environments and goals found in SILG evaluation, weadditionally ﬁne-tune them via RL similar to Guo et al.
This method is listing as inversein experiment ﬁgures.
In addition to the data used by reward and LDD , inverse uses additionalrollouts to train the inverse dynamics model.
4.3 Results and ablations
LDD consistently improves performance across environments.
We evaluate on held-out environments across 4 random seeds for NetHack in Fig 2, Touchdown in Fig 3, ALFWorld in Fig 4,
RTFM in Fig 5, and Messenger in Fig 6.
LDD obtains top performance compared to tabula-rasa policylearning with language descriptions, reward shaping using the dynamics model, and inverse reinforcement learning.
This is consistent across challenges in multi-step reasoning (RTFM), language-entitygeneralization (Messenger), large language action spaces (ALFWorld), large procedurally generated states (NetHack), and long natural language instructions with complex visual scenes (Touchdown).
We also ablate LDD by removing the initialization step ( LDD-init ) or the distillation step( LDD-distill ).
On Messenger, methods that pretrain to initialize ( LDD , inverse , LDD-distill )outperform those that do not ( reward , LDD-init ).
On Touchdown, methods that distill or rewardshape ( LDD , reward , LDD-init ) outperform those that do not.
Learning curves for each methodacross environments are shown in Appendix section G. LDD converges faster and to a higher win ratethan other methods, with the exception of Touchdown, where it achieves lower training but higherevaluation win-rate.
Language descriptions improve initialization and distillation gains in policy learning.
Weuse NetHack (Fig 2) and SymTouchdown (Fig 3) to study the importance of language descriptionsin dynamics modeling (others are not solvable without descriptions).
no-msg removes NetHackmessages describing events near the agent (e.g. kitten attacks the bat! ).
no-msg,stats additionallyremoves character state descriptions (e.g. health, achievements, dungeon level).
no-desc removes
Touchdown captions that describe objects locations scene relative to the agent (e.g. on your left, there are many building, some people, and few cars ).
These variants differ only in the observationspace used in dynamics modeling.
In ﬁne-tuning, they receive the same observation space withlanguage descriptions.
For both NetHack and Touchdown, removal of in-game messages, characterstate descriptions, and synthetic captions degrade performance.
Appendix Fig 12 further shows thatadding language descriptions results in dynamics models with higher accuracy.
This suggests thatmodeling language descriptions in the observation space result in better initialization and distillation,which improve subsequent policy learning.
Expert demonstrations cover late-stage strategy that result in asymptotic gains.
In complexenvironments, experts demonstrate late stage strategies difﬁcult to explore via random sampling.
Because these stages are rarely reached, accurate dynamics models are especially helpful in providingsignals when rewards are sparse.
Take NetHack as an example: unlike experts, non-expert policiesnever proceed to the deeper dungeons of the game.
A dynamics model trained on non-expert rolloutstherefore struggles to generalize to unseen deeper dungeons.
As the agent learns and descends deeperin the dungeon, a dynamics model trained on non-expert demonstrations results in less distillationgains.
Fig 7 compares dynamics modeling from observations using expert demonstrations vs. usingrollouts from a random policy.
The evaluation is on a held-out set of expert demonstrations.
Acrossenvironments, training on expert demonstrations outperforms training on non-expert demonstrations.
This effect is less apparent on environments where random policies can (eventually) discover most ofthe state space (e.g, RTFM, Messenger) and more apparent on partially observed environments whereonly strategic expert policies can encounter rare states indicative of success (e.g. long-term planningin NetHack and Touchdown, choosing from large language action space in ALFWorld).
Conclusion
While recent work showed that augmentation with language descriptions result in better policies,learning how to ground language descriptions to observations is difﬁcult through naive RL with sparse,delayed rewards.
We proposed Language Dynamics Distillation, which pretrains a dynamics modelusing cheaply-obtained unlabeled demonstrations with language descriptions to initialize and distillinto the policy learner.
On ﬁve tasks with language descriptions, LDD improved sample efﬁciency andresulted in better policies than RL from scratch, inverse RL, and expert reward shaping.
In addition,the beneﬁt from initialization and distillation differ on an environment basis, but are complementaryacross environments.
Moreover, language descriptions improved initialization and distillation gains inpolicy learning.
Finally, learning to model dynamics with expert demonstrations was more effectivethan with non-expert rollouts.
A promising direction for future research is studying whether dynamicsmodeling with language descriptions is similarly effective in robotic control where naive RL can beprohibitively expensive, but unlabeled demonstrations with synthetic captions are cheap to obtain.