Minyoung Hwang, Alexandra Forsey-Smerek, Nathaniel Dennler, Andreea Bobu
7 min
Abstract
Robots can adapt to user preferences by learning reward functions from demonstrations, but with limited data, reward models often overfit to spurious correlations and fail to generalize. This happens because demonstrations show robots how to do a task but not what matters for that task, causing the model to focus on irrelevant state details. Natural language can more directly specify what the robot should focus on, and, in principle, disambiguate between many reward functions consistent with the demonstrations. However, existing language-conditioned reward learning methods typically treat instructions as simple conditioning signals, without fully exploiting their potential to resolve ambiguity. Moreover, real instructions are often ambiguous themselves, so naive conditioning is unreliable. Our key insight is that these two input types carry complementary information: demonstrations show how to act, while language specifies what is important. We propose Masked Inverse Reinforcement Learning (Masked IRL), a framework that uses large language models (LLMs) to combine the strengths of both input types. Masked IRL infers state-relevance masks from language instructions and enforces invariance to irrelevant state components. When instructions are ambiguous, it uses LLM reasoning to clarify them in the context of the demonstrations. In simulation and on a real robot, Masked IRL outperforms prior language-conditioned IRL methods by up to 15% while using up to 4.7 times less data, demonstrating improved sample-efficiency, generalization, and robustness to ambiguous language. Project page: https://MIT-CLEAR-Lab.github.io/Masked-IRL and Code: https://github.com/MIT-CLEAR-Lab/Masked-IRL
Alex: Right—like tagging what's on-topic in a video game level. But how does the reward model actually use those tags without messing up if the LLM guesses wrong?
Sam: They avoid hard cuts by adding small random wiggles—noise—to the ignored parts of the scene during training. The reward shouldn't change much if those wiggles happen, because they're distractions; but it should shift if a key part like the laptop moves closer. The system learns this through a masking loss, which adds a penalty whenever the reward reacts to irrelevant noise, forcing it to stay steady on distractions while sensitive to what counts.
Alex: Huh. So it's like training a basketball player to focus on the hoop, not the crowd noise—jiggle the background, but the shot value only changes if the hoop shifts.
Sam: Precisely. This invariance weeds out fake patterns from sparse demos, like fixating on a weird arm twist instead of obstacle avoidance.
Alex: Yeah, but to make that invariance stick during training, there has to be some overall goal the model chases, right?
Sam: Yes. The training combines two parts: one pulls the reward toward matching the demonstrations, like standard practice. The other—the masking loss—adds a penalty if the reward shifts when irrelevant scene bits get nudged with noise. They balance these with a tunable weight, so the model fits the demos without chasing distractions.
Alex: Okay, so it's not just ignoring stuff upfront—it's learning to stay flat on the noise. How does the language actually plug into computing the reward itself?
Sam: They start by turning the words into a compact summary using a pre-trained language processor. That summary then tweaks the scene data directly: it creates numbers to scale up or down each part—like dimming background lights in a photo while brightening the main subject—and shifts them too. The adjusted scene feeds into a simple network that spits out a single reward score.
Alex: Huh—like sliders in a video editor, but for robot senses. Does that handle the vague instructions better than just sticking words next to the data?
Sam: It helps, but they go further for vagueness. The LLM doesn't just look at the instruction and demo—it also gets a "reference" path: the straight, shortest route from start to end, ignoring obstacles. By comparing the human's curvy demo to that plain path, it spots what extra care happened—like swerving from the laptop—and clarifies the words to match.
Alex: That seems solid for simulations. But the paper tests on a real arm too—what's the setup there?
Sam: They use a Franka Emika robot arm handing a mug in a setup with table, laptop, human. Preferences vary, like avoid laptop or stay near human. They test win rate: for random trajectory pairs, how often the learned reward picks the truly better one per human preference.
Alex: Win rate makes sense—like judging which play scores higher. With few demos, how does the masking stack up?
Sam: The study compares their soft masking to hard zeroing-out and plain language conditioning without masks. Masked IRL matches or beats the baseline across sparse to dense preferences, where sparse means just one or two key features matter. This holds on test sets unseen in training. On the real robot, it replicates with better success matching user prefs.
Alex: So fewer demos, steadier learning... a practical edge without magic. But with real human demos, which aren't perfect, does it still pick up user preferences reliably across different densities?
Sam: The simulations covered sparse, medium, and dense cases—and Masked IRL stayed robust after pretraining and fine-tuning, often matching oracle masks while beating language-conditioned baselines. On ambiguous instructions, like "Stay close" without naming the object, their disambiguation step lifted mask quality.
Alex: Right, so the LLM isn't flawless on masks. Where does that leave room for issues?
Sam: The paper notes reliance on LLM mask accuracy as a limit—imperfect guesses, especially on nuanced language, could weaken robustness. Evaluations stick to structured manipulation, like handover tasks, so scaling to dynamic or multi-object scenes needs work.
Alex: Makes sense—it's a targeted advance, not universal yet. Still, pulling demos and vague words together to learn rewards from just a handful of examples feels like a step toward household robots that adapt without endless practice.
Sam: Precisely. Demos give the how, language the why, and this invariance bridges them efficiently. The real-robot gains underscore potential for quick preference learning in everyday settings. Overall, a meaningful contribution to sample-efficient reward modeling.
Alex: Yeah, grounded progress on a stubborn robotics challenge. Thanks, Sam—that's our look at Masked IRL for today. Thanks for listening to ResearchPod.