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

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 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.

## introduction embodied ai presents a significant advancement in artificial intelligence, emphasising the importance of physical embodiment and the resulting ability to interact directly with the environment to accomplish goals. by perceiving and reasoning about the effect of their actions on their surroundings, embodied ai systems gain access to rich learning signals, which enable more robust and contextaware models of the world @xcite . for such systems to be useful across real-world deployment scenarios, they must be able to translate language instructions into meaningful sequences of actions. embodied agents that understand language not only unlock an additional learning modality, but also enable richer, situational representations of behaviour by grounding language into additional sensory modalities such as vision. this facilitates the creation of language-guided pixel-based policies for embodied ai (e.g., @xcite . while imitation learning (il) methods are commonly used to train such agents, it is typically assumed that the underlying demonstration data reflects optimal or near-optimal behaviour @xcite . this presents three key limitations. first, models trained exclusively on optimal trajectories may learn that there is exactly one valid solution for a given task. second, such models never encounter recoverable errors and their corrections. third, il lacks mechanisms to evaluate the quality of actions when faced with multiple valid behavioural modes for the same observation. in contrast, reinforcement learning (rl) algorithms explicitly incorporate feedback in the form of scalar rewards to distinguish between varied, more or less optimal behaviours @xcite . however, the associated exploration in rl also means that it is typically less sample-efficient than il. this is exacerbated for sparse rewards, which is typically the case in embodied ai. in this work, we investigate to what extent transformer-based il policies can benefit from the inclusion of suboptimal demonstrations in the training data when the effect of actions is contextualised with constructive language feedback, turning mistakes into learning opportunities. we focus on this family of models due to their ability to better represent multimodal inputs and long input sequences that are common in embodied ai tasks, and to align our work with the foundational architecture and training regimen used for current vision-language-action (vla) models in robotics @xcite @xcite . we test whether our figure 1: our method leverages both optimal and suboptimal trajectories for a given task instance by contextualising modes of behaviour with feedback signals. we leverage different types of feedback and additional self-supervised auxiliary tasks to learn highly generalisable and robust representations of behaviour in a data-efficient manner. we introduce a training regime where we systematically augment the training data with suboptimal variations of the base demonstrations. the use of language feedback is intuitive in the light of possible real-world deployment scenarios for embodied ai, and mirrors the modality of language instructions. in addition, we test if language feedback is a competitive substitute or suitable complement for shaped scalar rewards when conditioning il on feedback. we further test the efficacy of a self-supervised auxiliary task in the form of feedback prediction. the motivation for introducing the feedback prediction objectives is two-fold: 1. to encourage more robust representations of actions, and 2. to give the model the ability to tap into an internal world model of the consequences of actions when no feedback is provided at inference. to allow us to test fine-grained compositional generalisation capabilities systematically, we develop babyai-xgen, a modified version of the babyai @xcite environment, bridging the gap between procedural generation and granular control over task configurations and environment parameters. our findings indicate that transformer-based il policies that are trained from scratch with language feedback-conditioned suboptimal demon-strations generalise significantly better in compositional tasks than a baseline trained only on optimal trajectories, by turning mistakes and inefficient strategies into learning opportunities. we find this effect is consistent across different amounts of training data. notably, we observe that language feedback and scalar rewards provided with similar frequency yield comparable performance, offering practical flexibility depending on which signal is easier to provide. while combining language feedback and scalar rewards does not significantly improve task success rates, it does improve robustness to input perturbations during inference, indicating their complementary strengths. the same is observed for the auxiliary prediction task. lastly, we present further evidence for the sample efficiency of our method by demonstrating that an online rl baseline trained with ppo @xcite and using the same reward function achieves nearrandom generalisation performance when trained on equivalent numbers of data points. in summary, the contributions of this work are as follows: ## related work language-guided embodied ai. we consider our work in the context of language-guided embodied ai, specifically, mobile manipulation tasks. compared with manipulation-only tasks @xcite , and navigationonly tasks @xcite , the combination of navigation and object interaction creates additional degrees of freedom for behaviour to deviate from the optimal solutions, and hence, more meaningful opportunities to provide feedback can arise. existing datasets @xcite @xcite @xcite tend to annotate planner-generated trajectories with language instructions, which means that agents are trained to rely on optimal behaviour, while suboptimal behaviour is seen as something to mitigate, not leverage @xcite . language as feedback. previous work uses llms as judges @xcite , trains models to provide language feedback @xcite or leverages llms to increase the diversity of language feedback @xcite for a range of tasks. in contrast, we investigate how to utilize language feedback systematically to contextualise different modes of behaviour, rather than on the quality or generation mode of the feedback. while definitions of feedback in previous work range widely, it broadly falls into two categories: 1) general language hints @xcite or foresight @xcite )-which often takes the form of granular instructions-and 2) hindsight feedback @xcite . we focus on a form of hindsight feedback. to reduce the number of confounding variables, we procedurally populate a multi-part template which reflects both a judgment of the action and, importantly, an explanation of the judgment. our work is also in juxtaposition to previous work that derives a reward function from language, e.g. by computing the similarity between observations and subgoals described in language @xcite . instead, we use the language feedback embeddings directly as input into a transformer-based @xcite policy, following work by @xcite and @xcite . while both are antecedents to our work, they only compare language feedback with rewards corresponding to the binary failure / success scenarios typically associated with embodied ai tasks, and disregard the interplay of the language feedback with the optimality of the trajectories in the training data, which consists either of demonstration datasets @xcite or randomly generated trajectories @xcite . in contrast, we facilitate a fair comparison of language feedback and rewards by aligning their frequency, and systematically construct datasets composed of both optimal and suboptimal trajectories. another important strength of our work is that, differently from @xcite , we focus on purely pixel-based policies rather than studying environments where agents have access to symbolic state representation foot_0 . self-supervised auxiliary tasks for grounded representations. while predicting the observations and rewards resulting from actions is inherent to world models @xcite @xcite @xcite , where this ability is commonly referred to as latent imagination, the utility of predicting tokens other than actions as self-supervised auxiliary tasks @xcite in il is still under-explored. results from work on other multi-modal tasks, such as multi-modal translation @xcite or guessing games @xcite , suggest that such training objectives are a promising avenue for learning well-grounded multi-modal representations. we adapt this approach for il for embodied instruction following tasks. ## conclusion the continued popularity of language modelling objectives for llms sparked great interest in revisiting il techniques for embodied agents (e.g., @xcite . however, these are limited due to their inability to exploit suboptimal trajectories for learning more generalisable and robust behaviour. in this paper, we demonstrate the potential of language feedback as an efficient and intuitive feedback mechanism for il when the dataset either contains or consists entirely of suboptimal trajectories, and show its potential as a viable alternative to scalar rewards for tasks specified in language. we define an experimental setup based on a highly configurable 2d grid-world environment for learning from pixel-based, partial observations, which we designed expressly to assess performance along different axes of compositional generalisation @xcite . we show that models trained with both optimal and suboptimal data and feedback exhibit superior compositional generalisation abilities and increased robustness. we find that language feedback is not only a suitable alternative for scalar rewards for harnessing the learning potential of suboptimal samples, but also shows how we can increase robustness of models trained with language feedback with auxiliary feedback or by combining the strengths of both approaches. we are optimistic that our findings will inform future work on llmbased policies and that language feedback can play a role in different fine-tuning regimes and for alignment with human preferences.