Masked Inverse Reinforcement Learning for Language Conditioned Reward Learning

Minyoung Hwang, Alexandra Forsey-Smerek, Andreea Bobu

Keywords: Inverse Reinforcement Learning, Multi-Modal Feedback, Language Conditioning, Reward Learning

Abstract: Natural language provides a flexible interface for specifying robot tasks, but language-conditioned reward learning often assumes that instructions are unambiguous and directly informative. In reality, human language is frequently ambiguous — and may specify not just what to do, but also what matters in the environment. In this work, we propose a method that leverages this duality: we use large language models (LLMs) to extract state feature-level relevance masks from language and demonstrations, and train a reward function that is both conditioned on clarified task language and explicitly invariant to irrelevant parts of the state. We show that this approach improves generalization and sample efficiency in inverse reinforcement learning, particularly in settings with ambiguous instructions, distractor objects, or limited data. Our results highlight that disambiguating language with contextual demonstrations — and using language to guide both goal inference and state abstraction — enables more robust reward learning from natural instructions.

Submission Number: 13