Latest / Embodied AI 101 / Masked IRL: LLM-Guided Reward Disambiguation from Demonstrations and Language
Transcript
- Masked IRL: LLM-Guided Reward Disambiguation from Demonstrations and Language. Imagine you ask a robot to “place the coffee on my desk without disturbing my Zoom call.” You implicitly mean “keep the cup away from me and my laptop,” but if the robot only sees a few example demonstrations (e.g. how you physically moved its arm) it might never learn those critical details. In fact, with limited data a reward model trained from demonstrations often overfits to spurious details in those examples, focusing on irrelevant aspects of the state. Natural language can clarify what really matters (“stay away from the laptop, not from the red mug”), but people’s instructions are often terse or ambiguous. The MIT CSAIL news release introducing Masked IRL captures this dilemma succinctly: “It’s laborious to explain every detail…if you don’t have both [demonstrations and full instructions], the machine is likely to misunderstand what it needs to do”. The new Masked IRL framework addresses exactly this problem by using large language models (LLMs) to jointly leverage demonstrations and language to pinpoint which state features truly matter for a task. Instead of naively treating a human instruction as a flat conditioning signal, the method first has an LLM interpret the instruction (filling in missing detail), then generates a state-relevance mask highlighting which parts of the environment the reward should depend on. During training, the reward model is explicitly encouraged to ignore the irrelevant state components indicated by that mask. In practice this yields huge gains: the authors report solving the task with up to 4.7× fewer demonstrations and achieving about 15% higher success rates than prior language-conditioned IRL baselines. In other words, by “masking out” the unimportant factors (e.g. a bystander leaning on a table) the robot learns much faster and more reliably. In this article we’ll walk through how Masked IRL works, why it matters, and what the experiments show. We’ll see how it uses LLMs in two complementary ways – to clarify vague instructions and to identify relevant environment features – and how these insights plug into an inverse-reinforcement-learning pipeline. We’ll describe the simulation and real-robot evaluations (including one where the robot learned to avoid spilling coffee on a laptop simply from a “stay away” cue) and discuss the implications for robot learning with humans. The challenge: ambiguity in demonstrations and instructions. Robots often learn from humans by inverse reinforcement learning (IRL) or imitation learning: given a few expert demonstrations of a task, the robot infers a reward function (or policy) that explains those examples. But a core issue is that demonstrations alone don’t reveal the intent – that is, “what truly matters” in the task. As Hwang et al. note, “demonstrations show how to do a task but not what matters for that task, causing the model to focus on irrelevant state details”. For example, if someone demonstrates picking up a red block in front of a green wall, a naive reward learner might incorrectly conclude “go near red objects” or “stay by the wall” as important, even though those details are spurious. The model overfits to those unintended correlations when data is scarce, and thereby fails to generalize to new situations. Natural language offers a complementary source of information. The same papers point out that human instructions can “more directly specify what the robot should focus on”. You could say “stay away from the laptop,” which hints that the computer’s location is important to avoid. However, typical language-conditioned IRL methods usually ignore this opportunity. They simply feed the instruction text into the model as an extra input, expecting it to parse it implicitly. As the authors put it, prior methods “treat instructions as simple conditioning signals, without fully exploiting their potential to resolve ambiguity”. And in practice people’s shortcuts make things worse: prompts like “place the cup here” or just “stay close” can leave out the crucial noun (here, what to stay close to). In the coffee-and-Zoom example from the MIT CSAIL press release, you never explicitly said “avoid the person and laptop,” but of course that’s what was wanted. Masked IRL’s key insight is to join the strengths of both modalities: demonstrations show how to move, while language hints at what is important. The framework uses LLMs to extract that “what” from language and demonstrations. In short, an LLM pipeline refines the instructions and predicts a binary mask over state dimensions – marking each feature as either relevant (1) or irrelevant (0) to the user’s preference. The reward learner is then trained to be invariant to the masked-out features. This way, irrelevant details (like a background poster or a bystander’s pose) don’t mislead the learned reward. As summarized in the MIT news release: Masked IRL “automatically clarifies instructions” and “decides which details…to incorporate into a motion plan” using LLM-based reasoning, all while learning from much less data. The Masked IRL method. Masked IRL is an IRL framework where language and demonstration are processed together. The pipeline works in two main LLM-guided steps:. Instruction disambiguation: Given a (possibly vague) instruction and one demonstration, prompt an LLM to resolve missing context. For example, if the command is “stay close” and the demonstration shows moving along the table, the LLM will “elaborate on what might be unclear,” turning “stay close” into something like “stay close to the surface of the table”. In practice, the system provides to the LLM a short reference trajectory (the demonstrated motion or shortest-path action) plus any ambiguous text, and asks the model to clarify what the human meant. This language-refinement step ensures that under-specified cues (like not naming the target object or surface) get filled in using the demonstration as context. The result is a more explicit instruction that better captures the user’s true goal. (In the CSAIL example, the LLM sees that the demonstration hugs the table surface, and infers that “stay close” must mean “stay close to the table.”). State-relevance masking: Next, take the clarified instruction and feed it into a second LLM that has been prompted to consider the environment’s state description (e.g. positions or features of objects around the robot) and produce a binary relevance mask. This LLM effectively answers: “Which parts of the state does this instruction care about?” It outputs a 0/1 label for each state component, marking “1” if that feature should influence the reward, and “0” if it should be ignored. For instance, if the instruction is “stay away from the human” in a table-top task, the LLM-marked state might set the distance-to-human feature to 1 and distances to other irrelevant objects (like the wall or a clutter on the table) to 0. As described in the CSAIL article, this step assigns scores to obstacles and objects: “details the environment, such as the position of obstacles… it ‘masks’… generating 1 for important and 0 for not so much”. In effect, the LLM “pays attention” only to pieces of the scene that the language indicates are important (e.g. the laptop and person, not a detergent bottle). Once we have this binary state mask, Masked IRL uses it to shape the reward learning. In a traditional IRL loss (like maximum-entropy IRL), the model would try to make the demonstrated trajectory have high reward. Masked IRL adds a masking loss that enforces invariance of the reward to any changes in the masked-out dimensions. Concretely, instead of literally zeroing irrelevant inputs, the method perturbs the irrelevant state features with random noise and penalizes changes in the predicted reward. In other words, it trains the reward function so that if you wiggle the “irrelevant” parts, the reward stays the same. Mathematically, they combine this with the standard MaxEnt IRL objective as [ J(\theta) = L_\text{IRL}(\theta) + \lambda, L_\text{mask}(\theta),, ] where (L_\text{mask}) is the expected squared difference in reward under noise on irrelevant dims and (L_\text{IRL}) is the usual likelihood loss on demonstrations. This weighted sum encourages the learned reward to focus (invariance enforced) on just the state parts the LLM marked as relevant. At test time, when the robot receives a new instruction and sees a state, it uses the learned reward model (now with built-in “knowledge” of irrelevance) to plan. The CSAIL release explains that the robot selects a trajectory maximizing the learned reward. Because the reward model was trained to ignore masked-out features, the robot’s planning automatically respects the human’s preference while not getting distracted by irrelevant details. In summary, Masked IRL’s novelty is the integration of LLM reasoning into IRL: an LLM clarifies vague human commands given a demonstration, then another LLM labels environment features as relevant or not, and finally the IRL learner uses these labels to mask the reward training. As the project page puts it, “We clarify ambiguous language instructions using demonstrations and LLM reasoning. We then map disambiguated instructions into state masks, which guide the reward model through a masking loss that enforces invariance to irrelevant state dimensions during training”. Put differently: demonstrations show how to act, language specifies what matters, and LLMs connect the two. Experimental evaluation. The authors validated Masked IRL in a simulated robot task and on a real Panda arm, comparing against baselines of standard IRL with language conditioning (LC-IRL) and a simpler explicit-masking approach. The main simulated task is an object handover scenario with a 7-DoF Franka robot in PyBullet. The robot must move an object (e.g. a snack bag) from one side of a table to another, in the presence of a human (standing or seated), a laptop, and a cluttered table. Crucially, the true reward is defined as a linear combination of five semantic features: e.g. distance of the object to the table surface, to the human, to the laptop, height of the mug, etc. Each training preference is paired with a language instruction that refers to some subset of these features (for example “stay away from the table” or “keep the mug upright”). The baselines are:. Language-conditioned IRL (LC-IRL): a MaxEnt IRL model that receives the same demonstrations and a plain instruction text input, but with no masking logic. Explicit mask IRL: a variant that does use the LLM-generated mask, but instead of a soft loss, literally zeroes out the irrelevant state features during learning. They ran experiments across different reward densities (sparse to dense preferences). The performance metric is win rate: given two trajectories, which one achieves higher true reward. Sample efficiency and accuracy (RQ1). One key result is that Masked IRL learns much faster and more reliably than the LC-IRL baseline. With only a handful of demonstrations, Masked IRL’s learned reward already outperforms LC-IRL trained on many more samples. In the authors’ plot of win-rate vs number of demos, the Masked IRL curve lies steadily above LC-IRL’s. Quantitatively, Masked IRL “can match or exceed LC-RL performance with up to 4.7× fewer demonstrations”. In other words, nearly fivefold data efficiency, consistent with the news article claiming it uses “nearly five times less demonstration data”. At the same time, Masked IRL achieves about 15% higher final success rates than LC-IRL. The MIT release confirms this: video comparisons and statistics show the Masked IRL robot correctly obeyed the (implicit) preferences “up to 15% more often than comparable baselines”. Why is it so much better? The implicit masking loss drives the reward model to ignore irrelevant features, whereas LC-IRL often “stays sensitive to spurious cues” in low-data regimes. One ablation (the “explicit mask” variant) illustrates the danger of naive masking: if you simply delete irrelevant inputs, then any mistakes in the mask can catastrophically mislead learning. Indeed, the authors observe that explicit masking “degrades sharply when the masks are imperfect”, whereas Masked IRL’s soft penalty is robust even if the LLM’s mask is noisy. In short, the masking loss both enforces invariance and provides a smoother way to use the mask information, yielding large gains in sample efficiency. Dealing with ambiguous instructions (RQ2). Another axis of evaluation was instruction ambiguity. The researchers artificially generated under-specified commands (“Stay away” without a referent object, or just a single word like “Table” for a goal) for tasks that logically focus on one feature. Masked IRL’s two-step pipeline is meant to handle this: the first LLM query will disambiguate in context (e.g. given a motion, it might infer “Stay away from the human” or “Stay close to the table” as the true intent). As expected, using this LLM-based clarification dramatically improves performance. The paper reports that learning with the disambiguated instructions (DI) yields higher-quality state masks and better reward predictions than using the raw ambiguous commands (AI). For instance, an instruction “Stay away” could mean “stay away from the table” or “stay away from the human” depending on what the demonstration shows; having the LLM fill that in resulted in a mask that correctly marked the table (or human) as relevant, whereas the ambiguous prompt alone could not. In practice, Masked IRL with this language disambiguation consistently outperforms not using it, demonstrating that the LLM step indeed resolves key ambiguity. Put simply: when instructions become vague, querying the LLM to clarify them in light of the demo lets Masked IRL pick better masks. The MIT report emphasizes this benefit: “robots performed better when an LLM cleared up instructions, instead of having the machine try to follow a vague request”. This is an important advantage in real settings, where people often omit details. When combined with the masking loss, this makes Masked IRL robust to the kinds of underspecified language that would otherwise break a conventional learner. Real-world robot experiments (RQ3). Finally, the authors tested Masked IRL on a real Franka Emika Panda arm doing the object handover task with a person present. They collected kinesthetic demonstrations (humans physically guide the arm) for 50 different language-specified preferences (e.g. “stay away from human,” “stay close to table,” and combinations). Using the same Masked IRL training procedure as in simulation, they then performed zero-shot evaluation: given a new combination of constraints, the system optimized among a set of candidate motions and executed the one with highest learned reward. The real-robot results echo the simulation. Masked IRL achieved higher success rates on these held-out tasks than both LC-IRL and the explicit-mask baseline. It also showed much lower regret (i.e. the chosen trajectory was much closer to the true optimal) and was less sensitive to irrelevant perturbations. The authors report that Masked IRL “obtains higher win rates, lower reward variance under irrelevant perturbations, and substantially lower regret of optimized trajectories” than the baselines. In plain terms, Masked IRL picked safer, more aligned behaviors. To give a concrete flavor: the CSAIL article describes several test scenarios. In one case (“Stay away from human”), the baseline policy would sometimes drift too close to the person, but Masked IRL “maintains safe distance from the human” as required. In another (“Stay close to table surface”), Masked IRL keeps the towel in full contact with the table while wiping, whereas LC-IRL lost contact. In a complex case where the command was effectively “stay away from table and human,” Masked IRL actually avoided both obstacles simultaneously while LC-IRL did not. Perhaps most strikingly, in a scenario with a mug, a co-worker at a laptop, and a cluttered desk, Masked IRL executed a path that “carefully avoids both the human and the laptop” to prevent spills, whereas the language-conditioned baseline took a riskier route. These qualitative outcomes are quantified: after only 50 training demos, the robot obeyed constraints (even those never mentioned) reliably. For example, given a vague “stay away” after training, the robot correctly “avoided colliding with a user’s computer”. In summary, on the real arm Masked IRL zero-shot generalizes in a way the competitors do not. It learned from limited human interaction and elicited the intended behavior (keeping distance, maintaining contact, etc.) even for novel instructions. This confirms that the approach scales beyond simulation: the learned reward function truly captured the underlying human preferences. As the news article notes, after training on 50 demos, the robot could “carefully move a cup toward a human while avoiding colliding with the computer – an obstacle it learned to avoid by elaborating on a more general request to ‘stay away’”. In practice, this meant the Masked IRL robot was seen skillfully maneuvering objects around obstacles (e.g. moving a mug around a laptop) in line with unspoken preferences. Putting the results together, Masked IRL demonstrated three practical gains: improved accuracy of learned reward (15% lower error), much improved sample efficiency (nearly 5× fewer demos), and robustness to unclear language. These translate to more aligned and safe robot behavior with far less human effort – exactly the goals the researchers set out. In their words, this method “minimizes human effort by enabling machines to get to the bottom of what users really want”. Discussion and implications. Masked IRL is a clear example of leveraging pretrained language models as common-sense knowledge providers within robot learning. By explicitly extracting structure from language and feeding it into the reward-learning process, the method overcomes obstacles that purely data-driven IRL struggles with. Some key takeaways and broader implications:. Language as more than context. Many RL and IRL systems have appended language as a conditioning vector. Masked IRL shows that its role can be deeper: language can actively shape the learning objective. In prior “language-conditioned IRL” work (e.g. Fu et al. 2019‘s Language to Goals IRL), instructions were used to define a reward but were fed straightforwardly to the model. Here, by having the LLM hand-annotate which features matter, the instruction directly excludes irrelevant state components. This is a richer use of language that particularly combats one of RL’s long-standing issues: sample inefficiency and overfitting to the training distribution. LLM-robot synergy. The approach exemplifies how an LLM’s broad knowledge can bootstrap small-data learning in a domain like robotics. The first LLM, grounded by the demonstration, effectively supplies commonsense understanding (“if you say ‘stay away’ while showing something near a table, you probably meant the table”). The second LLM encodes spatial/physical reasoning to decide relevance. These are tasks that a small IRL network can’t do on its own. The system integrates this without requiring the model to learn language parsing from scratch – it relies on the LLM’s capability. In the CSAIL story, they point out how the LLM “sees the missing referent” (like table or human) and how it “masks” obstacles the human didn’t mention. Masking vs hard features. The clever idea of a stochastic masking loss (adding noise rather than hard deletion) mitigates the brittleness of imperfect LLM outputs. This design could inspire other methods: whenever we have partial knowledge of what should be invariant, an implicit invariance loss can be effective. The results here suggest these kinds of inductive biases greatly speed up reward inference. Practical skill learning. For real systems, data efficiency matters. Kinesthetic demos are expensive; any method that gets away with “nearly five times less data” as claimed is significant. Masked IRL essentially lets a human “show the robot once how to pour, and just say ‘stay away’,” and the robot will fill in the rest. This is a step toward more natural teaching: a few gestures plus some verbal hints, rather than exhaustive programming or repeated examples. Regret and safety. Not only did Masked IRL improve average performance, but the reductions in regret and variance suggest it learns reward functions that are closer to the true human preference. In the cluttered desk scenario, that meant avoiding spills. In tasks involving humans, this increased caution is crucial. By focusing on relevant factors, the robot avoids unsafe assumptions. As an example, the robot learned to keep distance from a person and expensive electronics even though the demos didn’t explicitly encode those dangers. Of course, some limitations and open questions remain. Masked IRL’s success relies on having a reasonably structured state representation so the LLM can meaningfully name “state dimensions.” In the experiments, features like distances to known objects (table, human, laptop) were used. Scaling this to raw perceptual inputs (like images) would require grounding vision into discrete object/feature labels first. The authors hint at this future direction: one could imagine using perception to identify objects (“table”, “laptop”) and then applying the same pipeline. They also discuss integrating cameras so the robot itself identifies relevant objects before prompting the LLM. Another practical concern is dependency on the LLM’s correctness. The current results show robustness to some LLM noise, but entirely wrong clarifications could still mislead the learner. It will be interesting to see how this approach extends to more complex language (long instructions, multi-step tasks) and what happens if the LLM hallucinates context incorrectly. Finally, using an LLM adds computational cost and reliance on large-model APIs. However, since the LLM is only used during training (off-board), the real-time execution remains a standard IRL/planning process. Overall, Masked IRL points to an exciting future where robots learn human intent more like we do: from a combination of showing and telling, with a smart interpreter in between. As one of the authors says, the system “senses and explains what users leave unsaid”. In concrete terms, this could make robots more data-efficient learners in homes or factories, better at capturing nuanced preferences (e.g. “don’t block the doorway,” “keep a clear workspace”), and safer around people and delicate objects. Conclusion. Contacting the fields of IRL, imitation, and language grounding, Masked IRL is a noteworthy advance in embodied AI. By using LLMs to disambiguate instructions and highlight relevant state features, it addresses a core bottleneck of reward learning from limited demos. The reported experiments – both quantitative and vivid qualitative examples on a real robot – demonstrate that it indeed “gets to the bottom of what users really want” with far fewer examples. For practitioners, the open-source code and project page are available online, making it a compelling approach to try in various human-robot teaching settings. Looking ahead, this work suggests many avenues: applying masking losses to other forms of side information (e.g. preferences, haptic cues), or combining multi-modal inputs (vision tags with language). It also exemplifies a broader paradigm: using the unstructured knowledge in LLMs to shape structured control policies. As the authors envision, soon the robot might directly see the environment and match language to pixel-based perceptions, but even now Masked IRL shows that smart language use can bridge the gap between a few demonstrations and robust, human-aligned robot behavior. Sources: The description above is based on the paper “Masked IRL: LLM-Guided Reward Disambiguation from Demonstrations and Language” (Hwang et al., ICRA 2026) and its associated code and project materials, as well as on a MIT CSAIL news article summarizing the work. The code and project website are publicly available.