FOSSIL: Harnessing Feedback on Suboptimal Samples for Data-Efficient Generalisation with Imitation Learning for Embodied Vision-and-Language Tasks

AI summary

Overview

Research area: Embodied AI — specifically imitation learning (IL) for language-conditioned, pixel-based agents that navigate and manipulate objects in a simulated environment.

Technical level: Intermediate. The paper assumes familiarity with imitation learning, Transformers, and Vision-Language-Action models, but its central idea (using language feedback to contextualise bad demonstrations) is explained in accessible terms.

Scope in one sentence: The paper introduces FOSSIL, a framework that trains Transformer-based IL policies on a mixture of optimal and suboptimal trajectories annotated with language feedback, and evaluates compositional generalisation, robustness, and data efficiency in a custom grid-world environment called BabyAI-XGen.

What This Paper Is About

Imitation learning normally assumes demonstration data is optimal or near-optimal, so models never see recoverable mistakes and may conclude there is only one valid solution to a task. The authors ask whether suboptimal trajectories — mistakes and inefficiencies that a planner has to correct — can be turned into useful training signal when the effect of each action is described in plain language ("positive" feedback for desired interactions with goal objects, "negative" feedback for undesired interactions with distractor objects). Their goal is a data-efficient method that generalises better and is more robust than training on optimal data alone.

Key Contributions

1. FOSSIL, a framework that feeds language feedback embeddings directly into the input sequence of a Transformer-based policy so that suboptimal demonstrations become learning opportunities for imitation learning.
2. BabyAI-XGen, a modified version of the BabyAI environment that supports procedural generation of custom tasks with granular control over environment parameters, enabling controlled compositional generalisation studies.
3. A multi-dimensional evaluation framework covering compositional generalisation (Systematicity and Productivity), robustness (to missing optimal data, external perturbations, and adversarial or absent feedback), solution efficiency, and data efficiency.
4. Empirical evidence that language feedback and shaped scalar rewards, provided at comparable frequency, can be equally effective for language-specified embodied tasks, and that their combination improves robustness.

Main Findings

• Language feedback rivals scalar rewards. On the compositional generalisation metrics, models trained with language feedback perform comparably to models trained with shaped scalar rewards when a feedback-prediction auxiliary objective is used. Averaged across settings, the Combo variant with feedback prediction reaches 65.0 in the "All" column versus 65.8 for Scalar without feedback prediction and 33.1 for the IL baseline with no feedback.
• Feedback unlocks suboptimal data. The paper reports that models trained with suboptimal trajectories plus some form of feedback score more than four times higher than IL baselines on the same settings (for example, 62.8 for Lang with feedback prediction versus 33.1 for None, on the averaged "All" metric).
• Productivity is uneven across dimensions. Language feedback is most effective at extrapolating to unseen goal colours, more complex layouts, and higher numbers of obstacles, and roughly on par with scalar rewards there. It is weakest on object locations never seen in training (13.8 for Lang versus 19.7 for Scalar), which the authors hypothesise reflects difficulty grounding spatial relations in language.
• Robustness to the absence of optimal data. When optimal trajectories are removed entirely, the baseline without step-level feedback drops 19.9% in success rate (15.1 to 12.1), while Combo drops only 6.7% (70.3 to 65.6), Scalar 10.5%, and Lang 10.0%.
• Robustness to external perturbations. Models trained on suboptimal trajectories without feedback fall below 70% of original performance under random action replacements. Lang with feedback prediction retains 75%, matching the scalar-reward models, and Combo with feedback prediction retains 78% — the most robust variant.
• Sensitivity to adversarial or missing feedback. Models trained only on language feedback collapse when given random sentences at random time steps or no feedback at all. Combo with feedback prediction retains 70%–80% of original performance, which the authors attribute specifically to combining scalar rewards as context with the ability to anticipate feedback.
• Solution efficiency. All models, including the one trained only on optimal trajectories, favour solutions less efficient than the oracle on average (mean oracle-normalised path lengths of 220.9% to 253.8%). Training with language feedback plus feedback prediction yields the most efficient solutions learned from suboptimal data (233.6% mean).
• Data efficiency and scaling. Models that contextualise suboptimal trajectories with any feedback form are more data-efficient than either IL from optimal trajectories or a PPO baseline, which the authors report achieves near-random generalisation when trained on equivalent numbers of data points. Doubling data from 25% to 50% raises success by 10.0 points for Scalar, 9.5 for Lang, and 11.1 for Combo with feedback prediction, versus 0.6 for the no-feedback baseline.

Methodology in Plain English

The authors build a single flexible architecture for all experiments: an autoregressive Transformer backbone (Llama2 configuration, but trained from scratch rather than using pretrained weights, and approximately 90M parameters) with an action prediction head and optional auxiliary regression heads. Sentences — both task instructions and feedback — are encoded by a frozen pretrained Sentence-BERT model, and image observations are compressed into a single token by a CNN trained alongside the policy.

Training data is generated procedurally in BabyAI-XGen. For each task, the authors create one dataset of roughly 12K optimal paths produced by a BFS planner, and another containing about 4K optimal and 8K suboptimal paths. Suboptimal paths are made by uniformly sampling from the action space and replacing the planner's action with probability p, after which the planner corrects the mistake if needed. Feedback is generated by rule-based oracles that have privileged access to environment state, so that feedback quality is controlled and not a confounding factor. Feedback comes in two flavours: task feedback about interactions with goal objects (now extended with negative feedback on distractor objects) and affordance feedback about actions with no visible effect, such as walking into walls. Shaped scalar rewards are derived from the same success and failure detectors, so language feedback and rewards are supplied at aligned frequency.

Models are trained with the standard next-action-prediction cross-entropy loss, optionally augmented with self-supervised auxiliary losses that predict the scalar reward or the language feedback embedding at the next timestep using an MSE objective, balanced by learnable weights. Predicted feedback tokens are used only to compute losses and are never fed back in as inputs. Variants are created by masking or unmasking instruction, returns-to-go, and language feedback tokens: the baseline (None) sees only instructions, Scalar sees returns-to-go, Lang sees language feedback, and Combo sees both.

Why This Matters

The work argues that imitation learning does not have to discard suboptimal data, and that language — the same modality used to specify tasks — can serve as a practical substitute or complement for shaped scalar rewards. It also provides a controllable testbed for compositional generalisation research in embodied settings where pixel observations and partial observability make evaluation harder than in symbolic benchmarks.

Real-world applications the paper points to:

• Collaborative robotics, where natural-language communication between humans and robots is already the preferred interaction channel.
• Interactive tutoring systems, where explanatory feedback about why an action was wrong aids learning more than a bare score.
• Multi-agent coordination, where agents need to understand and communicate about errors and corrections.
• Assistive or service robots that must keep functioning when hardware fails or actions go wrong, motivating the robustness-under-perturbation experiments.

Industry relevance: The paper situates itself alongside current Vision-Language-Action models in robotics, arguing that the feedback-conditioning idea is domain-agnostic. The authors note a practical cost caveat: step-level scalar rewards and language feedback will often require human annotation or reward models, and language feedback may be more expensive to collect than scalar scores.

Future Directions

• Scaling to realistic environments. The authors state that BabyAI-XGen lacks the visual realism, continuous action spaces, and physics fidelity of real robotic tasks, and that the approach should be tested where optimal trajectories are even harder to obtain.
• Larger, pretrained backbones. The current ~90M-parameter model is trained from scratch; the paper calls for investigating whether pretrained language models and larger architectures can better reason about multi-step errors and corrections, while noting that training data would need to scale to match.
• Better feedback generation. Feedback here is programmatic and oracle-based. Future work should explore more sophisticated step-level feedback generation, including VLM-based systems, and measure how feedback quality affects learning.
• Cost and collection tradeoffs. The paper suggests systematically evaluating when language feedback is worth its higher annotation cost compared with scalar rewards across different domains.

Target Audience

Researchers and graduate students working on embodied AI, imitation learning, and Vision-Language-Action models will get the most from this paper, particularly those interested in learning from imperfect demonstration data or in alternatives to reward engineering. It is also relevant to practitioners building language-guided robot policies who need to decide what kind of supervision signal to collect, and to researchers studying compositional generalisation who may want to reuse the BabyAI-XGen environment. Readers without a background in sequential decision making will find the experimental tables dense, though the core argument is stated in accessible terms.

Authors’ abstract

Current approaches to embodied AI tend to learn policies from expert demonstrations. However, without a mechanism to evaluate the quality of demonstrated actions, they are limited to learning from optimal behaviour, or they risk replicating errors and inefficiencies. While reinforcement learning offers one alternative, the associated exploration typically results in sacrificing data efficiency. This work explores how agents trained with imitation learning can learn robust representations from both optimal and suboptimal demonstrations when given access to constructive language feedback as a means to contextualise different modes of behaviour. We directly provide language feedback embeddings as part of the input sequence into a Transformer-based policy, and optionally complement the traditional next action prediction objective with auxiliary self-supervised learning objectives for feedback prediction. We test our approach on a range of embodied Vision-and-Language tasks in our custom BabyAI-XGen environment and show significant improvements in agents' compositional generalisation abilities and robustness, suggesting that our data-efficient method allows models to successfully convert suboptimal behaviour into learning opportunities. Overall, our results suggest that language feedback is a competitive and intuitive alternative to intermediate scalar rewards for language-specified embodied tasks.