Source-linked AI summary
Pre-Trained Language Models for Interactive Decision-Making
Shuang Li, Xavier Puig, Chris Paxton, Yilun Du, Clinton Wang, Linxi Fan, Tao Chen, De-An Huang, Ekin Akyürek, Anima Anandkumar, Jacob Andreas, Igor Mordatch, Antonio Torralba, Yuke Zhu
TL;DR
The paper asks whether pre-trained language models can support general sequential decision-making beyond language tasks. It proposes LID, which represents goals, histories, and observations as embedding sequences for an LM-initialized policy, and finds improved generalization to novel tasks, including a 43.6% increase in task completion rates in VirtualHome.
Problem
The paper investigates whether pre-trained language models can provide a general framework for sequential decision-making tasks that involve no language.
Method
LID encodes observations, goals, and histories as embedding sequences and fine-tunes a policy network initialized with a pre-trained language model to predict actions.
Results
43.6% higher task completion rates on novel VirtualHome tasks were obtained with LM-pretrained policies, while sequential representations and LM weight initialization contributed to generalization.
Takeaways & Limitations
LM pre-training improves combinatorial generalization in policy learning, and its benefits extend to arbitrary sequential input encodings rather than only natural-language strings.
Takeaways & Limitations
The method still fails on challenging tasks because of grounding and policy errors, limiting reliable task completion.
Abstract
from arXiv · showhide
Language model (LM) pre-training is useful in many language processing tasks. But can pre-trained LMs be further leveraged for more general machine learning problems? We propose an approach for using LMs to scaffold learning and generalization in general sequential decision-making problems. In this approach, goals and observations are represented as a sequence of embeddings, and a policy network initialized with a pre-trained LM predicts the next action. We demonstrate that this framework enables effective combinatorial generalization across different environments and supervisory modalities. We begin by assuming access to a set of expert demonstrations, and show that initializing policies with LMs and fine-tuning them via behavior cloning improves task completion rates by 43.6% in the VirtualHome environment. Next, we integrate an active data gathering procedure in which agents iteratively interact with the environment, relabel past "failed" experiences with new goals, and update their policies in a self-supervised loop. Active data gathering further improves combinatorial generalization, outperforming the best baseline by 25.1%. Finally, we explain these results by investigating three possible factors underlying the effectiveness of the LM-based policy. We find that sequential input representations (vs. fixed-dimensional feature vectors) and LM-based weight initialization are both important for generalization. Surprisingly, however, the format of the policy inputs encoding (e.g. as a natural language string vs. an arbitrary sequential encoding) has little influence. Together, these results suggest that language modeling induces representations that are useful for modeling not just language, but also goals and plans; these representations can aid learning and generalization even outside of language processing.
1 Introduction
The paper proposes LID, which uses pre-trained language models to process sequential representations of goals, observations, and history for interactive decision-making. Experiments show improved generalization, including with active data gathering and across arbitrary sequential encodings.
- 1 Introduction: LID converts policy inputs into sequential data and feeds them to a policy initialized with a pre-trained language model.Inputs can include observations, goals, and history represented as language, image patches, or scene graphs.
- 1 Introduction: 43.6% improvement in task completion rates is reported for novel VirtualHome tasks using imitation learning with pre-trained LM policy initialization.The approach also yields 20% more successful policies than other baselines on i.i.d. VirtualHome tasks.
- 1 Introduction: Active data gathering enables policy learning without pre-collected expert data by exploring, relabeling failed trajectories, and updating the policy.The procedure uses useful sub-trajectories that achieve sub-goals and reportedly outperforms reinforcement learning methods.
- 1 Introduction: Different input encoding schemes have negligible performance impact, so natural-language strings are not required for LM-based generalization.The finding extends the benefits of language-model pre-training to arbitrary sequential encodings.
- 1 Introduction: Removing sequential input structure significantly degrades novel-task performance, while removing pre-trained weights reduces success rate by 11.2%.The analyses identify sequential representations and LM-based weight initialization as important factors in generalization.
- 1 Introduction: The results support pre-trained LMs as a general framework for promoting structured generalization in interactive decision-making.The paper presents this as a general scaffold across environments and supervisory settings.
2 Related Work
Related work covers pre-trained language-model representations and their adaptation beyond language modeling, including policy learning for challenging embodied environments. Existing policy-learning methods perform well in standard benchmarks but may struggle with complex embodied tasks.
- 2 Related Work: Pre-trained word and sentence representations have become ubiquitous in natural language processing and related multimodal applications.Examples include instruction following and language-guided image retrieval.
- 2 Related Work: Language representation learning progressed from isolated word forms toward masked or autoregressive language-model training.The passage situates modern pre-trained language models within this broader progression.
- 2 Related Work: Language models can be adapted to diverse language tasks by casting them as word-prediction problems or modifying prediction heads.Applications include parsing, language-to-code translation, machine translation, sentiment classification, and style transfer.
- 2 Related Work: Traditional policy-learning methods perform well on Atari, OpenAI Gym, and MuJoCo but may fail on more challenging embodied environments.Recent work has begun using language models for policy learning in these settings.
3 Decision-Making and Language Modeling
The paper frames partially observed decision-making and language modeling as sequence-prediction problems, then uses pre-trained transformer LMs to initialize policies that predict actions from goals, observations, and history.
- POMDPs: POMDP policies must use current partial observations together with histories of previous observations and actions.The policy is represented as πφ(at|g, ht, ot), conditioning action probabilities on goals, history, and current observation.
- Policy learning: LID converts policy inputs into embedding sequences, processes them with a transformer encoder, and uses a task-specific decoder to predict actions.The framework trains on trajectories containing goals, observations, and actions.
- Language modeling: Autoregressive transformer LMs model sequences by conditioning each next element on preceding elements through learned token embeddings and transformer networks.The paper uses GPT-2 to process policy input sequences rather than predict future text tokens.
- Connection: Because language modeling and POMDP decision-making both predict successive elements from sequences, pre-trained LMs can initialize policies fine-tuned on expert or high-reward trajectories.This shared sequence-prediction structure motivates transferring LM initialization to policy learning.
4 Approach
LID represents goals, observations, and histories as sequences processed by a pre-trained LM, and extends this setup with active data gathering that learns from relabeled failures without pre-collected expert data.
- 4.1 Policy Network: LID encodes observations, goals, and action histories as sequences, pools LM outputs into a context representation, and predicts the next valid action.Training maximizes demonstrated-action probabilities, while inference selects the valid action with the highest probability.
- Environment encodings: VirtualHome goals become templated English predicates, while observations encode visible objects, states, and 3D coordinates.These environment-specific features are incorporated into the policy input sequence.
- Environment encodings: BabyAI combines textual grid descriptions, history actions, and the unchanged language instruction before feeding them to the pre-trained LM.The 7 × 7 observation is converted into text descriptions such as objects, walls, and doors.
- 4.1 Policy Network: The policy network can use arbitrary sequential environment encodings rather than requiring natural-language strings.The framework is also described as potentially supporting pixel observations through discretization.
- 4.1 Policy Network: The same sequential policy framework supports substantially different observation spaces and generalizes to novel tasks in both VirtualHome and BabyAI.The paper also describes the framework as applicable to other embodied environments.
- 4.2.2 Policy Learning with Active Data Gathering: LID-ADG alternates exploration, hindsight relabeling, and policy updates to improve policies without pre-collected expert data.Failed trajectories provide sub-trajectories that can be relabeled with goals corresponding to solved sub-goals.
5 Experiment Setup
The experiments evaluate LID and baselines in VirtualHome and BabyAI across in-distribution performance and generalization to new scenes, layouts, and task combinations.
- Environments: The evaluation uses VirtualHome and BabyAI as the two experimental environments.The paper evaluates the proposed method and baselines in both settings.
- VirtualHome: VirtualHome is a partially observable 3D environment with large action spaces and long time horizons.Its evaluation covers in-distribution tasks, novel scenes, and novel tasks.
- VirtualHome: In-distribution VirtualHome tasks sample goal predicates and counts from the training distribution and use common-sense object layouts.This defines the matched-distribution test condition.
- VirtualHome: Novel-scene VirtualHome tasks place objects randomly without common-sense constraints, such as apples inside dishwashers.This changes initial environment layouts while retaining the environment setting.
- VirtualHome: Novel-task VirtualHome tests combine goal-predicate components that were not seen together during training.For example, Inside(plate, fridge) appears only in the test set although plates and fridges appeared separately during training.
- Evaluation protocol: VirtualHome success requires completing the entire goal within the environment’s maximum allowed steps, using 1,500 evaluation examples per model.Each of three test subsets uses five random seeds and 100 tasks per seed.
- BabyAI: BabyAI evaluates four instruction-following tasks over 500 episodes each, requiring generalization to new layouts and unseen task combinations.Observations are 7 × 7 × 3 partial, local egocentric grids.
6 Experiments
The experiments evaluate LID-Text with expert data across VirtualHome and BabyAI, then assess LID-ADG without pre-collected expert demonstrations. LID-Text outperforms baselines, especially on tasks requiring combinatorial generalization.
- LID-Text converts environment inputs into text descriptions and fine-tunes a pre-trained LM for action prediction conditioned on goals, observations, and histories.
- Results on VirtualHome: 43.6% improvement over all baselines occurs on VirtualHome Novel Tasks, where test goals require combinations never seen during training.The methods are trained on 20K demonstrations; LID-Text also has higher success rates in In-Distribution and Novel Scenes settings.
- Results on BabyAI: LID-Text outperforms BabyAI-Ori with less training data, indicating improved sample efficiency when generalizing to novel tasks.With 10K demonstrations, both methods achieve high success rates.
- Active Data Gathering: LID-ADG is evaluated without pre-collected expert data against random, Goal-Object, PPO, and DQN+HER baselines on VirtualHome.The procedure is designed for settings where agents actively gather their own data.
- Active Data Gathering: LID-ADG is the only successful approach in the no-expert-data VirtualHome comparison and can initialize PPO or provide data for offline learning.Initializing PPO with LID-ADG improves its In-Distribution success rate by 53.7%.
7 Analysis: Understanding the Sources of Generalization
The analysis separates the effects of input encoding, sequential structure, and LM-based initialization on combinatorial generalization. Sequential representations and pre-trained weights matter, while the specific encoding format has little influence given sufficient data.
- Input Encoding Scheme: LID-Index preserves the discrete sequential structure while replacing words with integer indices and training a new embedding layer from scratch.
- Input Encoding Scheme: Different input encodings retain comparable eventual performance, although text encoding is the most sample-efficient and is around 4% higher at 500 demonstrations.LID-Text, LID-Index, and LID-Unnatural all perform below 10% with 100 demonstrations and converge to similar performance with sufficient data.
- Sequential Inputs: Removing sequential structure significantly hurts Novel Tasks performance, while No-Seq remains effective on test tasks closer to training.
- Favorable Weight Initialization: Removing pre-trained weights lowers Novel Tasks success by 11.2% despite fitting In-Distribution data well, showing the weights are important for generalization.
- Summary: The combined results identify sequential input representations and favorable weight initialization as important, whereas the encoding scheme itself has little influence.
8 Qualitative Results
Qualitative examples show LID-Text completing tasks in VirtualHome and BabyAI, while failure analysis identifies grounding and policy errors as the main failure types.
- Successful Examples: LID-Text completes demonstrated In-Distribution and Novel Tasks examples in VirtualHome, plus GoToLocal and PickupLoc tasks in BabyAI.The examples show short trajectories or extracted sub-trajectories.
- Failure Cases: The main failure types are grounding errors, which target the wrong object, and policy errors, which fail to find or interact with target objects.The paper notes that challenging tasks remain unsolved and suggests larger LMs may improve success rates.
9 Conclusion and Broader Impact
LID uses sequential representations and pre-trained LM initialization for decision-making, while active data gathering removes reliance on expert data. The authors report effective combinatorial generalization but note limitations in relabeling rules and LM bias.
- LID converts goals, histories, and observations into sequences processed by a policy initialized with a pre-trained LM.
- Active data gathering enables policy learning without expert data but relies on hand-designed task-relabeling rules.
- Input representation and favorable weight initialization contribute to generalization, whereas the input encoding scheme has little influence.
- Biases in pre-trained LMs may influence behavior, requiring further study before deployment in sensitive downstream applications.
Checklist
The checklist reports that the paper addresses contributions, limitations, potential negative societal impacts, reproducibility materials, training details, error bars, compute resources, and asset attribution.
- The authors report that contributions and scope, limitations, and potential negative societal impacts are discussed.
- The paper reports that code, data, and reproduction instructions are included in the supplemental material.
- Training details, error bars, and compute resources are reported in the main paper or appendices.
- Existing assets are cited, while a license discussion is marked not applicable.
- The checklist marks new assets, consent discussion, personally identifiable information discussion, and human-subject procedures as not applicable or addressed elsewhere.
- The appendix provides pointers to environment, implementation, evaluation, data-gathering, test-subset, and attention-visualization details.
A Convolutional encoding in BabyAI
The BabyAI study compares text and convolutional sequential encodings within LID. Both approaches approach perfect performance with sufficient data, while text encoding is more sample-efficient.
- A Convolutional encoding in BabyAI: LID-Conv converts BabyAI’s 7 × 7 × 3 grid observation into a convolutional feature map, flattens it into a feature sequence, and keeps the rest of LID unchanged.
- A Convolutional encoding in BabyAI: Text and convolutional encodings produce similar results with enough training data, while text encoding is slightly better with fewer training data.
- A Convolutional encoding in BabyAI: Different input encoding schemes have negligible impact on performance, extending pre-training benefits beyond natural-language strings to arbitrary sequential encodings.
- A Convolutional encoding in BabyAI: Table 6 reports success rates for policies trained with text versus convolutional encoding on BabyAI.
- A Convolutional encoding in BabyAI: BabyAI and VirtualHome differ substantially in their goals, states, and action sequences, providing distinct evaluation environments.
- A Convolutional encoding in BabyAI: VirtualHome contains partial observability, large action spaces, long horizons, household objects, and predicate-based goals evaluated within 70 steps.
C More implementation Details of LID in VirtualHome
LID’s VirtualHome implementation embeds goals, histories, and partial observations as sequences for a pre-trained LM policy. Active data gathering relabels failed trajectories, but its rules constrain relabeling scope.
- Policy architecture: The VirtualHome policy has input, pre-trained LM, and output components, with goal, history, and partial observation represented as embedding sequences.
- Observation encoding: Each observed object is encoded from its name, state, and position features before being combined into an object representation.
- Observation encoding: Position features combine world coordinates with displacement to the agent and pass through fully connected layers with an intermediate ReLU layer.
- Training and evaluation: Interactive evaluation samples goals and initial states, repeatedly predicts actions from policy πφ(at|g, ht, ot), and computes success rate as n/Ntest.
- Active data gathering: Hindsight relabeling detects useful tasks in failed trajectories with hand-designed keyword rules, while complex settings are limited to short single-predicate sub-trajectories.