Technical Summary: Supervised Reward Inference (SRI)

Problem Statement

Traditional approaches to reward inference, such as Inverse Reinforcement Learning (IRL), typically rely on the assumption that human demonstrations are generated according to a specific, known behavior model (e.g., noisy optimality, bounded rationality). However, real-world human behavior is often suboptimal due to poor planning, execution errors, or intentional non-optimal gestures used solely to communicate goals. Existing methods struggle when the assumed behavior model is misspecified, leading to incorrect reward inference. While some prior work attempts to learn behavior models from data, these approaches often face limitations, such as being restricted to tabular settings or failing to handle arbitrarily suboptimal behaviors effectively.

The core problem addressed is how to infer reward functions from a wide range of human behaviors—including those that are arbitrarily suboptimal or intended as communication rather than optimization—without relying on strong, potentially incorrect assumptions about the underlying behavior model.

Methodology: Supervised Reward Inference (SRI)

The authors propose Supervised Reward Inference (SRI), a framework that reframes reward inference as a supervised learning problem. Instead of learning a behavior model that maps rewards to actions and then inverting it, SRI directly learns a mapping from behavior trajectories to reward functions using a dataset of paired behaviors and known ground-truth rewards.

Core Framework

Given a dataset of tasks where each task $k$ consists of:

1. A set of behavior trajectories $\{\tau_{k,i}\}$ (which may be arbitrarily suboptimal).
2. A set of state-reward pairs $\{(s_{k,j}, r_{k,j})\}$ sampled from the same task.

SRI trains a neural network to predict the reward $r(s)$ for a given state $s$, conditioned on the behavior trajectories. The architecture consists of two main components:

• Task Encoder ($f_{\theta_f}$): Encodes a set of behavior trajectories into a fixed-dimensional task embedding $\psi$. In the experiments, this utilizes a Transformer for trajectory encoding and a Set Transformer for pooling multiple trajectories.
• Reward Model ($g_{\theta_g}$): Takes the task embedding $\psi$ and a state $s$ as input to predict the reward value.

The training objective is to minimize the mean squared error between the predicted rewards and the ground-truth rewards over the dataset:
$$\arg \min_{\theta_f, \theta_g} \mathbb{E} \left[ L \left( g_{\theta_g}(S, f_{\theta_f}(\{\tau_n\})), R(S) \right) \right]$$

Theoretical Guarantees

The paper provides a theoretical proof that SRI is asymptotically Bayes-optimal. Under standard assumptions (compact state/action spaces, bounded returns, and sufficient model capacity), as the dataset size approaches infinity, the inferred reward function converges uniformly to the posterior mean of the reward given the behavior. Consequently, the optimal policy derived from this inferred reward converges to the policy that maximizes the expected return under the true posterior distribution of rewards. This holds for any fixed class of behavior, provided the behavior distribution at deployment matches the training distribution.

Extensions

The authors demonstrate that this supervised inference template generalizes to other intent inference tasks:

• Supervised Goal Inference (SGI): Predicting latent goal representations directly from behavior.
• Supervised Action Inference (SAI): Predicting optimal actions directly from behavior and state.

Key Contributions

1. Novel Framework: Introduction of SRI, which eliminates the need for explicit behavior modeling by treating reward inference as a direct supervised learning task.
2. Theoretical Optimality: Proof that SRI achieves asymptotic Bayes-optimality for reward inference and apprenticeship learning under standard assumptions, without requiring knowledge of the behavior model.
3. Robustness to Suboptimality: Demonstration that SRI can infer accurate rewards from highly suboptimal, noisy, and even non-demonstration behaviors (e.g., gestures that never complete the task) where traditional IRL methods fail.
4. Generalization: Extension of the framework to Supervised Goal Inference (SGI) and Supervised Action Inference (SAI), showing the versatility of the supervised intent inference approach.

Experimental Results

The authors evaluated SRI on two distinct benchmarks:

1. Meta-World Robotics Tasks (Continuous Control)

• Setup: Tasks included "Reach" and "Pick-Place" with randomly distributed goals.
• Behavior Classes: The system was tested on five classes of demonstrations, including noisy actions, systematically wrong offsets, and "Gesture" demonstrations where the robot pointed to the goal but never picked or placed the object.
• Findings:
  • SRI achieved near-ceiling performance (0.822 normalized goal proximity) on Pick-Place tasks using only "Gesture" demonstrations, whereas all imitation learning baselines (BC, GAIL, AIRL) failed (performance $\le 0$).
  • SRI demonstrated robustness to arbitrary suboptimalities, inferring near-optimal policies even when demonstrations were highly noisy or systematically biased.
  • Performance scaled smoothly with data quantity, showing strong data efficiency.

2. Shah et al. (2019) Gridworld Benchmark (Tabular)

• Setup: A tabular benchmark comparing SRI against prior methods for arbitrary suboptimal behavior, including the "Inverse Planning" (IP) algorithm from Shah et al. (2019).
• Findings:
  • SRI (adapted to take policy inputs) achieved near-ceiling performance across all demonstrator classes (optimal, naive, sophisticated, myopic, etc.) and their noisy variants.
  • SRI outperformed or matched the state-of-the-art IP methods, suggesting that directly learning the inverse map (policy-to-reward) is more effective than learning a planner and inverting it, even in tabular settings.

Significance and Claims

The paper claims that SRI represents a shift in paradigm from behavior-model engineering to data-driven supervised learning. By viewing behavior as an indication of intent rather than an optimization attempt, SRI bypasses the brittleness associated with model misspecification.

The authors emphasize that:

• SRI is a "parsimonious yet powerful solution" that leverages standard supervised learning techniques.
• It allows for the use of advances in representation learning and pretrained models (e.g., Transformers) without needing to design specific RL inductive biases.
• The method is universally applicable to any fixed behavior distribution, provided sufficient task-indexed reward supervision is available during training.
• While the current experiments assume access to simulator rewards for training, the framework is designed to be compatible with future extensions where rewards are estimated from sparse annotations or preferences.

The paper concludes that less-biased learning algorithms, when provided with sufficient data, can outperform more-biased, model-based approaches in reward inference, supporting the "bitter lesson" that scaling data and computation often yields better results than hand-crafted assumptions.