MAVRL: Learning Reward Functions from Multiple Feedback Types with Amortized Variational Inference
Authors: Raphaël Baur, Yannick Metz, Maria Gkoulta, Mennatallah El-Assady, Giorgia Ramponi, Thomas Kleine Buening
Organizations: ETH AI Center, ETH Zurich, Zurich, Switzerland · Department of Computer Science, ETH Zurich, Zurich, Switzerland · Department of Informatics, University of Zurich, Zurich, Switzerland
Reward learning typically relies on a single feedback type or combines multiple feedback types using manually weighted loss terms. Currently, it remains unclear how to jointly learn reward functions from heterogeneous feedback types such as demonstrations, comparisons, ratings, and stops that provide qualitatively different signals. We address this challenge by formulating reward learning from multiple feedback types as Bayesian inference over a shared latent reward function, where each feedback type contributes information through an explicit likelihood. We introduce a scalable amortized variational inference approach that learns a shared reward encoder and feedback-specific likelihood decoders and is trained by optimizing a single evidence lower bound. Our approach avoids reducing feedback to a common intermediate representation and eliminates the need for manual loss balancing. Across discrete and continuous-control benchmarks, we show that jointly inferred reward posteriors outperform single-type baselines, exploit complementary information across feedback types, and yield policies that are more robust to environment perturbations. The inferred reward uncertainty further provides interpretable signals for analyzing model confidence and consistency across feedback types.
Figures & tables
Figure 1: Visualizations of inferred reward functions from 2 demonstrations, 256 pairwise comparisons, or 128 ratings on a 10×10 grid_trap environment. The final column shows the result obtained when combining all feedback modalities.
Table 1: Normalized mean returns and mean EPIC distances ( n=10 ) for different combinations of feedback modalities across environments. Returns (top, higher is better) are normalized such that 100 corresponds to the performance of an approximately optimal policy trained on the ground-truth reward and 0 corresponds to the performance of a uniformly random policy. EPIC distances (bottom, lower is better) measure the divergence between learned and ground-truth rewards. Cells shaded in dark gray indicate the best result for the corresponding environment, and cells shaded in light gray indicate the next best result outside this margin.
Post-Hoc Avg
MCMC
MAVRL
Environment
Perf.
Perf.
W.-Time
Perf.
W.-Time
grid_sparse
94.1
100.0
154.4 s
100.0
4.5 s
grid_cliff
54.9
100.0
156.0 s
100.0
4.9 s
grid_trap
47.3
93.0
155.0 s
95.3
4.6 s
LunarLander-v3
59.9
–
–
93.5
46.5 s
Table 2: Comparison of MAVRL with non-variational baselines. Performance (Perf.) is normalized return averaged over n=10 seeds (higher is better); wall-time (W.-Time) is per reward-model training run on identical hardware. Best per row in dark gray , second-best in light gray .
Figure 2: Normalized mean returns ( n=10 ) of policies trained on rewards inferred by each method under three dynamics perturbation scenarios. Reward models and baselines are trained in the unperturbed setting and remain fixed throughout the variations. Error bars denote standard error. (left) Increasing environmental stochasticity in grid_cliff and grid_trap. (middle) Increasing ratio between pendulum handle lengths in Acrobot-v1. (right) Increasing gravity and wind-power in LunarLander-v3.
Table 3: Misspecification robustness ( n=10 ; grid_trap in App. D ). Top: single-modality models; bottom: MAVRL with full feedback set. Each row corrupts the data-side parameter of the named channel while the model assumes the well-specified value. Strike-through denotes the corrupted modality; e.g., corresponds to misspecified preferences. Base. = well-specified normalized return; Misspec. = mean ± SEM; Rat. = Misspec./Base. Best per column in dark gray , second-best in light gray .
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Table 4: Normalized mean returns and mean EPIC distances ( n=10 ) for different combinations of feedback modalities across environments under equal feedback budget . Returns (top, higher is better) are normalized such that 100 corresponds to the performance of an approximately optimal policy trained on the ground-truth reward and 0 corresponds to the performance of a uniformly random policy. EPIC distances (bottom, lower is better) measure the divergence between learned and ground-truth rewards. Cells shaded in dark gray indicate the best result for the corresponding environment (within 1% for returns), and cells shaded in light gray indicate the next best result outside this margin.
Table 5: Misspecification robustness ( n=10 ). Top: single-modality models; bottom: MAVRL with full feedback set. Each row corrupts the data-side parameter of the named channel while the model assumes the well-specified value. Strike-through denotes the corrupted modality; e.g., corresponds to misspecified preferences. Base. = well-specified normalized return; Misspec. = mean ± SEM; Rat. = Misspec./Base. Best per column in dark gray , second-best in light gray .
Figure 3: Normalized mean returns ( n=10 ) of policies trained on rewards inferred by each method under three dynamics perturbation scenarios. Reward models and baselines are trained in the unperturbed setting and remain fixed throughout the variations. Returns are normalized such that 100.0 corresponds to the performance of an approximately optimal policy trained on the ground-truth reward in the unperturbed setting and, analogously, 0.0 corresponds to the performance of a uniformly random policy. Error bars denote standard error. (a) and (b) Increasing environmental stochasticity in grid_cliff and grid_trap. (c) Increasing ratio between pendulum handle lengths in Acrobot-v1. (d) Increasing gravity and wind-power in LunarLander-v3.
Figure 4: grid_sparse : Visualizations of inferred reward functions from 2 demonstrations, 256 pairwise comparisons, or 128 ratings on a 10×10 grid_sparse environment. The final column shows the result obtained when combining all feedback modalities. (a) Joint visual encoding of mean and variance. (b) Separate visual encoding of mean and variance for the same data.
Figure 5: grid_cliff : Visualizations of inferred reward functions from 2 demonstrations, 256 pairwise comparisons, or 128 ratings on a 10×10 grid_cliff environment. The final column shows the result obtained when combining all feedback modalities. (a) Joint visual encoding of mean and variance. (b) Separate visual encoding of mean and variance for the same data.
Figure 6: grid_trap : Visualizations of inferred reward functions from 2 demonstrations, 256 pairwise comparisons, or 128 ratings on a 10×10 grid_trap environment. The final column shows the result obtained when combining all feedback modalities. (a) Joint visual encoding of mean and variance. (b) Separate visual encoding of mean and variance for the same data.
Existing approaches to reward inference typically assume that humans provide demonstrations according to specific behavior models. However, humans often indicate their goals through a wide range of behaviors, from actions that are suboptimal due to poor planning or execution to behaviors intended to communicate goals rather than achieve them. One existing solution for inferring rewards from such behavior \unicodex2013 provided it is drawn from the same distribution at training and deployment \unicodex2013 is to construct a dataset of behavior paired with known rewards, and to learn the mapping from behavior to rewards; however, prior methods in this family face notable limitations, such as restrictions to tabular settings. Given such a dataset, we propose instead that supervised learning offers a parsimonious yet powerful solution, which we term Supervised Reward Inference (SRI). Theoretically, we prove that SRI is asymptotically Bayes-optimal under standard assumptions. Empirically, SRI achieves near-ceiling performance on a prior benchmark for reward inference from suboptimal behavior, while on Meta-World robotics tasks, it infers rewards from even arbitrarily suboptimal demonstrations as accurately as those demonstrations allow. Finally, we demonstrate our framework's universality with straightforward generalizations to action- and goal-prediction.
Will Schwarzer, Jordan Schneider, Philip S. Thomas +1
University of Massachusetts · Anthropic · Work done while at the University of Texas.
As autonomous agents are increasingly deployed across diverse operational contexts, aligning their behavior with human intent demands reward functions that remain robust to such changes rather than overfitting to any single environment. Inverse reinforcement learning (IRL) provides a principled way to infer such objectives from human feedback. However, existing analyses of optimal teaching approaches for IRL focus on single-environment, demonstration-only settings, leaving underexplored how heterogeneous feedback modalities and environment dynamics jointly constrain reward functions that generalize across multiple environments. Because demonstrations in one MDP entangle reward information with that environments specific structure, the resulting rewards frequently fail to generalize when the agent is deployed in a new setting. We first analyze how different feedback modalities constrain rewards, showing that, in the unlimited-data regime, comparisons impose strictly stronger global constraints than other modalities. Beyond this theoretical analysis, we introduce a hierarchical machine teaching algorithm for reward learning that operates across multiple MDPs. The algorithm first greedily selects informative environments that expose complementary reward constraints, then strategically queries low-cost feedback within those environments. Empirically, our method achieves substantially lower regret and stronger generalization to held-out environments than uniform teaching baselines under identical feedback budgets, demonstrating the importance of multi-environment, multi-modal teaching for learning dynamics-robust reward functions.
Reinforcement learning from human feedback (RLHF) has emerged as a powerful yet sample-inefficient approach for learning reward models from human preferences, making active learning a critical component in synthesizing informative preference queries. However, effective uncertainty quantification required for active learning remains a key challenge for large neural network reward models. In this paper, we introduce PreferenceEKF, a sample-efficient approach that tracks reward model uncertainty by framing active preference learning as a sequential Bayesian filtering problem. Instead of relying on computationally prohibitive posterior inference over the full neural network parameter space, our method performs sequential inference via an extended Kalman filter within a low-dimensional parameter subspace, continuously updating the reward model posterior as new preference queries arrive. Our approach enables scalable sampling of neural network parameters to efficiently compute acquisition functions for active reward learning. Experiments on the D4RL and V-D4RL benchmarks demonstrate that our approach achieves better sample efficiency, runtime, scalability, and calibration compared to other Bayesian deep learning approaches, and the learned reward models lead to competitive offline reinforcement learning policy performance. This highlights the potential of scalable Bayesian methods for preference-based reward modeling in RLHF. Our code is available at https://github.com/yutaizhou/bnn_pref.
Yutai Zhou, Erdem Bıyık
Thomas Lord Department of Computer Science University of Southern California, Los Angeles, USA