Masked IRL: LLM-Guided Reward Disambiguation from Demonstrations and Language

1. Problem Statement

The paper addresses the fundamental challenge in Inverse Reinforcement Learning (IRL): learning a reward function from human demonstrations is an ill-posed problem. A single demonstration trajectory can be explained by infinitely many reward functions.

• The Core Issue: Demonstrations show how to act but do not explicitly convey what matters (the user's intent). Consequently, reward models often overfit to spurious correlations in irrelevant state components (e.g., a robot avoiding a table simply because the human happened to be near it, rather than because the table is an obstacle).
• The Role of Language: Natural language instructions (e.g., "Stay away from the laptop") can theoretically disambiguate these preferences. However, existing language-conditioned methods treat instructions merely as conditioning signals without fully exploiting them to filter irrelevant states.
• The Ambiguity Gap: Real-world instructions are often underspecified or ambiguous (e.g., "Stay away" without specifying what to stay away from). Naive conditioning on such ambiguous text leads to unreliable reward learning.

2. Methodology: Masked IRL

The authors propose Masked Inverse Reinforcement Learning (Masked IRL), a framework that leverages Large Language Models (LLMs) to bridge the gap between demonstrations and language. The method operates on the insight that demonstrations and language are complementary: demonstrations show the behavior, while language specifies relevance.

The framework consists of three key components:

A. LLM-Generated State Relevance Masks

Instead of treating the entire state vector as relevant, the system uses an LLM to infer which parts of the state vector are relevant to a specific instruction.

• Process: Given a demonstration trajectory ($\tau$) and an instruction ($\ell$), the LLM outputs a binary mask $m \in \{0, 1\}^d$, where $d$ is the state dimension. $m(j)=1$ if state component $j$ is relevant, and $0$ otherwise.
• Example: For the instruction "Stay away from the laptop," the LLM identifies the laptop's position as relevant ($1$) and the human's position as irrelevant ($0$).

B. Implicit Masking Loss

Rather than explicitly zeroing out irrelevant state dimensions (which makes the model brittle to mask errors), Masked IRL enforces invariance through a novel loss function.

• Mechanism: The model is trained to ensure that perturbing irrelevant state components (those where $m(j)=0$) does not change the predicted reward.
• Loss Function:
  $$\mathcal{L}_{mask}(\theta) = \mathbb{E} \left[ \sum_{j=1}^{d} (1 - m(j)) | r_\theta(s^{(j)} | \ell) - r_\theta(s | \ell) | \right]$$
  Where $s^{(j)}$ is the state with noise added to dimension $j$. This penalizes the model if it relies on irrelevant features, forcing it to learn robust, task-specific features.

C. LLM-Based Disambiguation

To handle ambiguous instructions (e.g., "Stay away"), the system uses the demonstration context to clarify the intent.

• Process: The LLM is prompted with the environment description, the ambiguous instruction, the demonstration trajectory, and a "reference trajectory" (the shortest path).
• Reasoning: The LLM reasons about the difference between the demonstration and the reference path to hypothesize the missing referent (e.g., inferring "Stay away" implies "Stay away from the table" because the robot deviated from the shortest path to avoid the table).
• Output: The system generates a disambiguated instruction (e.g., "Stay away from the table"), which is then used to generate the state mask and train the reward model.

D. Reward Model Architecture

The core reward model is a language-conditioned neural network using Feature-wise Linear Modulation (FiLM).

• The language instruction is encoded into an embedding $h_{lang}$ via a frozen T5 transformer.
• $h_{lang}$ generates scaling ($\gamma$) and shifting ($\beta$) parameters that modulate the state input $s$ before it is passed to an MLP.
• The total training objective combines the standard Maximum Entropy IRL loss with the masking loss:
  $$J(\theta) = \mathcal{L}_{IRL}(\theta) + \lambda \mathcal{L}_{mask}(\theta)$$

3. Key Contributions

1. Language-Guided State Relevance Masks: A novel mechanism to use LLMs to identify and filter irrelevant state dimensions, preventing overfitting to spurious correlations.
2. Implicit Masking Loss: A training objective that enforces invariance to irrelevant state perturbations, offering robustness against imperfect LLM-generated masks compared to naive explicit masking.
3. Contextual Disambiguation: An LLM-based pipeline that resolves underspecified language instructions by reasoning jointly with demonstration data.
4. Unified Framework: A single language-conditioned reward model that generalizes across diverse preferences without needing to retrain for each new user.

4. Experimental Results

The authors evaluated Masked IRL on a 7-DoF Franka Emika robot arm in both simulation (PyBullet) and the real world, focusing on an object handover task involving a table, laptop, and human.

• Sample Efficiency: Masked IRL achieved comparable or superior performance to baseline language-conditioned IRL (LC-RL) using up to 4.7 times fewer demonstrations. It could generalize well with as few as 5 demonstrations per preference.
• Robustness to Ambiguity: In experiments with ambiguous instructions (e.g., "Stay away"), the disambiguation pipeline improved the F1 score of state mask prediction by 16.9% compared to using ambiguous text directly. The final reward model showed a 21.4% higher average win rate than baselines on test preferences.
• Generalization: Masked IRL outperformed baselines across sparse, medium, and dense reward densities.
• Real-World Performance: On the physical robot, Masked IRL demonstrated:
  • Higher average win rates.
  • Significantly lower reward variance when irrelevant state dimensions were perturbed (proving invariance).
  • 44.8% to 59.4% lower regret compared to baselines, indicating it learned reward functions that led to better optimized trajectories.

5. Significance

This work represents a significant step forward in human-robot interaction and sample-efficient learning.

• Solving the "What Matters" Problem: It moves beyond simply using language as a label and uses it as a structural guide to filter state space, addressing the core cause of overfitting in IRL.
• Handling Real-World Noise: By combining LLM reasoning with demonstration context, the system becomes robust to the natural ambiguity and underspecification found in human communication.
• Practical Deployment: The ability to learn robust reward functions with very few demonstrations makes the approach viable for real-world robotics, where collecting large datasets of human demonstrations is costly and time-consuming.

In summary, Masked IRL effectively leverages the complementary strengths of demonstrations (behavioral data) and language (semantic constraints) via LLMs to create reward models that are more data-efficient, generalizable, and robust to ambiguity.