Q-based Variational Inverse Reinforcement Learning

Inferring human preferences from expert demonstrations is a central challenge in building safe and beneficial AI. Inverse reinforcement learning (IRL) addresses this by treating observed behavior as optimal under an unknown reward function. While Bayesian IRL provides a principled framework by recovering a full posterior distribution over rewards—essential for uncertainty quantification and active learning—prior methods are largely limited to small discrete state spaces or require computationally expensive Markov chain Monte Carlo (MCMC) sampling.

This paper introduces Q-based Variational IRL (QVIRL), which combines the scalability of variational inference with the rigor of Bayesian uncertainty quantification. Instead of repeatedly solving the forward reinforcement learning problem to map rewards to policies, QVIRL models a variational posterior directly over optimal Q-values. Using the inverse Bellman operator and efficient approximations like Clark's method, the Q-value posterior is mapped to a closed-form reward posterior. This design allows QVIRL to scale to high-dimensional problems, including continuous control tasks and raw pixel observations.

Evaluated across randomized gridworlds, continuous control tasks like Lunar Lander and Highway Environment, and ATARI games, QVIRL demonstrates strong performance in apprenticeship learning and active learning. In tabular gridworlds, QVIRL closely approximates the true Bayesian posterior generated by MCMC reference methods while requiring a fraction of the computational time. In more complex environments, QVIRL successfully utilizes its learned uncertainty to guide active learning loops—selecting informative demonstration states to rapidly improve policy performance. Notably, QVIRL is the first Bayesian IRL method capable of effectively training directly from raw pixel observations.

By bridging the gap between scalable deep reinforcement learning and full Bayesian uncertainty estimation, QVIRL paves the way for deploying preference-learning algorithms in safety-critical, high-dimensional domains. Its ability to quantify uncertainty over rewards and policies enables risk-sensitive decision-making and active data collection, minimizing the risks associated with misspecified reward point estimates in autonomous systems.