Technical Summary: Multi-Modal, Multi-Environment Machine Teaching for Robust Reward Learning

Problem Statement

As autonomous agents are deployed across diverse operational contexts, aligning their behavior with human intent requires reward functions that generalize robustly across varying dynamics rather than overfitting to a single environment. While Inverse Reinforcement Learning (IRL) offers a principled method for inferring objectives from human feedback, existing machine teaching frameworks for IRL are limited by two primary constraints:

1. Single-Environment Focus: Prior work assumes a fixed Markov Decision Process (MDP). However, demonstrations in one MDP entangle reward information with that environment's specific structure, often leading to rewards that fail to generalize when the agent encounters new dynamics.
2. Single-Modality Limitation: Existing teaching approaches rely almost exclusively on demonstrations, neglecting how heterogeneous feedback types (e.g., comparisons, corrections, emergency stops) jointly constrain reward functions.

The central question addressed is: How can a human use heterogeneous feedback to teach an agent so that its behavior remains aligned with human intent across multiple operating conditions?

Methodology

Theoretical Foundations

The authors analyze the geometric constraints imposed by different feedback modalities on the Generalized Behavioral Equivalence Class (gBEC), defined as the set of reward parameters consistent with a given feedback dataset.

• Feedback Modalities: The paper models four types of feedback under a reward-rational choice framework: Demonstrations, Comparisons (preference between trajectories), Corrections (improved trajectory vs. rollout), and Emergency Stops (halting a rollout).
• Data Regimes:
  • Unlimited-Data Regime: Theoretical analysis shows that comparisons impose strictly stronger global constraints than other modalities because they can order any two trajectories across the entire space. In contrast, demonstrations, corrections, and E-stops are constrained by shared initial states or trajectory prefixes.
  • Limited-Budget Regime: In finite-data settings, demonstrations are more constraint-efficient per query because each demonstration implicitly enforces multiple optimality constraints simultaneously.
• Environment-Dependent Identifiability: The paper proves that reward identifiability is dependent on the environment's dynamics. Even with unlimited feedback in a single MDP, residual reward ambiguity may persist if the environment's feasible trajectories do not span all reward-relevant feature directions. Consequently, resolving ambiguity often requires teaching across multiple environments with complementary constraint spans.

Proposed Algorithm: HSCOT

To address these limitations, the authors introduce Hierarchical Set Cover Optimal Teaching (HSCOT). This framework treats the teaching problem as a two-stage hierarchical optimization:

1. Outer Level (Environment Selection): The algorithm greedily selects a minimal subset of training MDPs whose combined trajectory spaces expose the maximum number of complementary reward constraint directions (maximizing the union of feature-difference spans).
2. Inner Level (Feedback Selection): Within the selected environments, the algorithm greedily selects specific feedback atoms (queries) to cover the remaining uncovered constraint directions in the constraint universe.

The objective is to minimize the teaching cost (lexicographically: number of environments, then number of queries) while ensuring the inferred reward induces low-regret behavior on the training set, serving as a surrogate for generalization to held-out environments.

Key Contributions

1. Analysis of Feedback Modalities: The paper provides a formal analysis showing that while comparisons impose the strongest global constraints in the unlimited-data limit, demonstrations are more efficient per query under tight budgets.
2. Characterization of Identifiability: It formally characterizes reward identifiability as environment-dependent, proving that single-MDP teaching often leaves residual ambiguity that only multi-environment teaching can resolve.
3. HSCOT Framework: The introduction of a hierarchical machine teaching algorithm that jointly optimizes environment selection and feedback query selection to efficiently constrain reward spaces across multiple MDPs.
4. Empirical Validation: Demonstration that HSCOT achieves higher constraint coverage and significantly lower regret on held-out environments compared to uniform teaching baselines under identical feedback budgets.

Experimental Results

The authors evaluated HSCOT on two domains: GridWorld (6x6, 50 generated MDPs) and LavaMiniGrid (4 features, 50 generated MDPs). In both cases, 20% of MDPs were held out for testing.

• Held-Out Regret: HSCOT consistently achieved near-zero or zero regret on held-out environments across all feedback modalities (Demonstration, Comparison, Correction, E-stop). In contrast, uniform teaching baselines incurred substantially higher regret, confirming that strategic environment selection is critical for generalization.
• Constraint Coverage: HSCOT achieved complete coverage of the universal constraint set (the set of all possible linear preference constraints inducible by the training MDPs). Uniform teaching covered only a small fraction of this set, even when using demonstrations.
• Environment Efficiency: HSCOT activated significantly fewer environments than the uniform baseline to achieve the same performance. For instance, in GridWorld with demonstration feedback, HSCOT used an average of 4.3 environments versus 5.2 for uniform teaching.
• Mixed Modalities: Combining demonstrations with other modalities allowed HSCOT to reach near-zero regret in both domains, leveraging complementary constraint directions to tighten the feasible reward region.

Significance and Claims

The paper claims to be the first to study machine teaching for reward learning across multiple environments with heterogeneous feedback modalities. Its primary significance lies in bridging the gap between single-MDP teaching and the reality of multi-context deployment.

The authors argue that:

• Single-MDP teaching is insufficient: Rewards learned in a single environment often fail to generalize because the environment's dynamics obscure certain reward directions.
• Multi-environment teaching is necessary: Variation in environment dynamics is not just noise but a resource that reveals complementary reward constraints.
• Strategic selection matters: Simply collecting more data (even across environments) is less effective than strategically selecting environments and feedback types that maximize constraint coverage.

The work extends the Behavioral Equivalence Class framework to a multi-environment setting and proposes a unified optimization formulation for selecting both environments and feedback. The authors acknowledge limitations, noting that extending HSCOT to continuous domains and handling noisy human feedback remain important directions for future work.