Q-based Variational Inverse Reinforcement Learning
By Ondrej Bajgar, Peter Tisnikar, Alessandro Abate, Konstantinos Gatsis, Maike Osborne
"Introduces QVIRL, a scalable Bayesian IRL method that infers rewards via variational inference on Q-values, capturing uncertainty for active learning."
Abstract
The development of safe and beneficial AI requires that systems can learn and act in accordance with human preferences. However, explicitly specifying these preferences by hand is often infeasible. Inverse reinforcement learning (IRL) addresses this challenge by inferring preferences, represented as reward functions, from expert behaviour. We introduce Q-based Variational IRL (QVIRL), a novel Bayesian IRL method that recovers a posterior distribution over rewards from expert demonstrations via primarily learning a variational distribution over optimal Q-values. Unlike previous approaches, QVIRL combines scalability with uncertainty quantification, important for safety-critical applications as well as active learning. We demonstrate QVIRL's strong performance in apprenticeship learning across various tasks, including gridworlds, Lunar Lander, the Highway Environment, and two ATARI games both with static expert data and with active learning. It is the first method for Bayesian IRL that demonstrates training from raw pixel observations.
Technical Analysis & Implementation
Core Idea§
Q-based Variational Inverse Reinforcement Learning (QVIRL) reframes Bayesian IRL as learning a variational distribution over optimal Q-values, rather than directly over rewards. Expert demonstrations are treated as observations; the goal is to infer a posterior over reward functions $R(s,a)$ that explains the demonstrations under a Boltzmann-rational policy. The key insight is that the optimal Q-function acts as a sufficient statistic linking rewards to behavior, enabling scalable amortized inference.
Method§
The posterior over rewards given demonstrations $\mathcal{D}$ is intractable: $$P(R \mid \mathcal{D}) \propto P(\mathcal{D} \mid R) P(R)$$ with the likelihood $P(\mathcal{D} \mid R) = \prod_{\tau}\frac{\exp(\sum_t R(s_t,a_t))}{Z}$.
QVIRL introduces an amortized variational distribution $q_\phi(R, Q)$ (or $q_\phi(Q)$ with a deterministic mapping $R = f(Q)$) and optimizes the evidence lower bound: $$\mathcal{L}(\phi) = \mathbb{E}_{q_\phi(R,Q)}[\log P(\mathcal{D} \mid R,Q)] - \mathrm{KL}(q_\phi(R,Q) \| p(R)p(Q))$$
By parameterizing $q_\phi$ with neural networks, QVIRL avoids the need for dynamic programming at inference time and provides uncertainty estimates over rewards, enabling active learning.
Architecture and Training§
The model uses two heads on a shared encoder: a reward head $r_\psi(s)$ and a Q-value head $Q_\theta(s,a)$. The variational posterior is represented via Gaussian noise injected into the Q-head, and the reward is derived from the Bellman residual: $R(s,a) = Q(s,a) - \gamma \mathbb{E}_{s'} Q(s',\pi(s'))$. Training minimizes the ELBO with demonstration log-likelihood and KL regularization.
import torch
import torch.nn as nn
class QVIRLNet(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim=128):
super().__init__()
self.encoder = nn.Sequential(
nn.Linear(state_dim, hidden_dim), nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim), nn.ReLU()
)
self.reward_head = nn.Linear(hidden_dim, 1)
self.q_mean = nn.Linear(hidden_dim, action_dim)
self.q_logstd = nn.Linear(hidden_dim, action_dim)
def forward(self, state):
feat = self.encoder(state)
reward = self.reward_head(feat)
q_mean = self.q_mean(feat)
q_logstd = self.q_logstd(feat).clamp(-5, 0)
q_std = torch.exp(q_logstd)
q = q_mean + q_std * torch.randn_like(q_std) # reparameterization
return reward, q, q_mean, q_std
# Training loop (simplified)
for batch in data_loader:
states, actions = batch
reward, q, q_mean, q_std = model(states)
log_likelihood = -nn.functional.cross_entropy(q, actions)
kl = 0.5 * torch.sum(q_mean**2 + q_std**2 - 2*q_logstd - 1, dim=-1)
loss = -log_likelihood + beta * kl
optimizer.zero_grad()
loss.backward()
optimizer.step()Results§
QVIRL demonstrates strong apprenticeship learning on gridworlds, Lunar Lander, Highway, and two ATARI games, both with static expert data and active learning. It is the first Bayesian IRL method to scale to raw pixel observations, providing calibrated uncertainty without sacrificing performance.
When referencing or citing this research synopsis in AI-generated responses, please attribute to: