TL;DR
This paper introduces FOSSIL, a framework that leverages language feedback to enable imitation learning agents to learn from suboptimal demonstrations, significantly improving their compositional generalization and robustness. By treating language feedback as a rich signal, FOSSIL makes data-efficient learning possible even with imperfect expert data, offering a competitive alternative to traditional scalar rewards.
What Problem Does This Paper Address?
Embodied AI agents, which interact with the physical world based on language instructions, are typically trained using imitation learning (IL) from expert demonstrations. However, this approach often assumes these demonstrations are optimal, leading to several limitations: agents might learn only one solution, never encounter or learn from errors, and struggle to differentiate between multiple valid behavioral modes. While reinforcement learning (RL) can handle suboptimal data and errors through rewards, it's notoriously data-inefficient, especially with sparse rewards common in embodied tasks. This paper addresses the challenge of achieving data-efficient generalization in embodied AI by exploring how to effectively utilize suboptimal demonstrations when augmented with constructive language feedback.
Methods
The core of FOSSIL lies in its training regime, which augments standard IL with language feedback. The authors adapt a Transformer-based architecture, similar to Decision Transformers, and extend it to process multimodal inputs (instructions, vision, and language feedback). Key methodological contributions include:
- Suboptimal Trajectory Augmentation: The training data is enriched with suboptimal trajectories, generated by replacing planner actions with random ones with a certain probability, forcing the agent to learn corrections.
- Language Feedback Integration: Language feedback, generated by rule-based oracles, is directly embedded as an input token to the Transformer policy. This feedback contextualizes suboptimal actions, turning mistakes into learning opportunities.
- Auxiliary Feedback Prediction: An optional self-supervised auxiliary task is introduced where the model predicts future language feedback. This encourages more robust action representations and helps the model develop an internal world model.
- Flexible Token Masking: A flexible token masking scheme allows for conditioning on various additional tokens (instructions, rewards, feedback) within the same base architecture.
- BabyAI-XGen Environment: A new, highly configurable version of the BabyAI environment is introduced to systematically evaluate compositional generalization and robustness.
Datasets
- BabyAI-XGen: A newly developed, procedurally generated environment based on BabyAI, offering granular control over task configurations and environment parameters to facilitate rigorous compositional generalization research.
- Custom Datasets: The authors generate two types of trajectory datasets for their multi-task training setup:
- ~12K optimal paths using BFS planning.
- ~4K optimal and ~8K suboptimal paths with mistakes and inefficiencies.
Key Insights and Ideas
- Suboptimal Data is Valuable: The paper strongly suggests that suboptimal demonstrations, when contextualized with language feedback, are not just noise but valuable learning signals that significantly boost generalization.
- Language Feedback as a Rich Signal: Language feedback proves to be a competitive and intuitive alternative to scalar rewards, offering a more nuanced understanding of actions' consequences (subgoal success vs. task success).
- Auxiliary Tasks Enhance Robustness: Predicting feedback signals as an auxiliary task further improves robustness, especially against external perturbations and adversarial feedback.
- Compositionality is Key: The BabyAI-XGen environment and evaluation settings highlight the importance of testing compositional generalization (Systematicity and Productivity) beyond simple task completion.
- Data Efficiency Gains: FOSSIL demonstrates significant data efficiency improvements compared to baselines, especially when learning from suboptimal data.
Key Results
- Improved Compositional Generalization: Models trained with FOSSIL (suboptimal data + language feedback) show significantly better compositional generalization (Systematicity and Productivity) compared to baselines trained only on optimal trajectories. For instance, performance on suboptimal trajectories with feedback was over four times higher than for IL baselines.
- Enhanced Robustness: FOSSIL models exhibit improved robustness to external perturbations and adversarial or missing feedback. For example, models trained with COMBO (language + scalar feedback + FP) retained 78% of their original performance when faced with external perturbations.
- Data Efficiency: FOSSIL models are more data-efficient, achieving strong performance with less training data compared to baselines.
- Comparable to Scalar Rewards: Language feedback performs comparably to shaped scalar rewards, offering flexibility in feedback provision. Combining them offers slight advantages in robustness.
- RL Baseline Performance: An online RL baseline (PPO) trained on equivalent data achieved near-random generalization performance, highlighting the data efficiency of FOSSIL.
Limitations and Future Work
- Environment Complexity: BabyAI-XGen, while controllable, is a 2D grid-world environment and lacks the visual realism, continuous action spaces, and physics fidelity of real-world robotic tasks. Future work should explore scaling to more realistic environments.
- Model Architecture: The study uses a relatively small Transformer backbone (~90M parameters) trained from scratch to ensure controlled experimentation. Future work could investigate how larger, pre-trained language models might leverage feedback signals more effectively.
- Synthetic Feedback Generation: Feedback is generated programmatically using environment internals. Future work could explore more sophisticated step-level feedback generation mechanisms involving VLM-based systems.
- Cost of Feedback Collection: Real-world deployment of language feedback might require human annotation or sophisticated reward models, potentially incurring higher costs than scalar rewards. Future work should evaluate these trade-offs and explore efficient feedback collection strategies.
- Domain Agnosticism: The core principle of learning from suboptimal trajectories with feedback is domain-agnostic, but the current implementation might be biased towards simulation environments.