Source-linked AI summary
Learning Robust Rewards with Adversarial Inverse Reinforcement Learning
Justin Fu, Katie Luo, Sergey Levine
TL;DR
Reward engineering remains a barrier, and existing IRL methods struggle with ambiguity, scale, and changing dynamics. The paper proposes AIRL, an adversarial method for learning portable rewards, and reports effective transfer and stronger performance than prior methods in transfer settings.
Problem
Reward engineering remains difficult, while inverse reinforcement learning has struggled to scale to high-dimensional tasks and distinguish true rewards from dynamics-shaped rewards.
Method
AIRL jointly learns reward and value functions through an adversarial inverse reinforcement learning formulation, with constraints that support disentangled rewards.
Results
AIRL rewards transfer effectively under variation in domain dynamics, while unmodified IRL methods tend to recover brittle rewards and GAIL does not recover reward functions.
Takeaways & Limitations
AIRL provides a practical and scalable IRL algorithm that can learn disentangled rewards and recover ground-truth rewards up to a constant in small unambiguous MDPs.
Takeaways & Limitations
The advantage-based state-action reward is heavily entangled with the training MDP and is not guaranteed to remain robust when environment dynamics change.
Abstract
from arXiv · showhide
Reinforcement learning provides a powerful and general framework for decision making and control, but its application in practice is often hindered by the need for extensive feature and reward engineering. Deep reinforcement learning methods can remove the need for explicit engineering of policy or value features, but still require a manually specified reward function. Inverse reinforcement learning holds the promise of automatic reward acquisition, but has proven exceptionally difficult to apply to large, high-dimensional problems with unknown dynamics. In this work, we propose adverserial inverse reinforcement learning (AIRL), a practical and scalable inverse reinforcement learning algorithm based on an adversarial reward learning formulation. We demonstrate that AIRL is able to recover reward functions that are robust to changes in dynamics, enabling us to learn policies even under significant variation in the environment seen during training. Our experiments show that AIRL greatly outperforms prior methods in these transfer settings.
1 INTRODUCTION
The paper motivates automatic reward acquisition because reward engineering remains a practical barrier, while IRL has struggled with ambiguity, scalability, and transfer across changing dynamics. AIRL addresses these issues through adversarial reward learning and aims to recover rewards that generalize beyond the demonstration environment.
- Reward engineering remains a significant barrier to applying reinforcement learning in practice, despite deep reinforcement learning reducing feature engineering.The difficulty includes specifying behaviors such as socially acceptable conduct.
- IRL infers an expert’s reward function from demonstrations, but prior methods have not scaled to tasks as complex as direct imitation learning.IRL can nevertheless support re-optimizing rewards in novel environments or inferring agent intentions.
- Many rewards and policies can explain demonstrations, so IRL may recover dynamics-shaped rewards that fail to produce optimal behavior when test-time dynamics change.The paper calls rewards invariant to changing dynamics disentangled rewards.
- AIRL jointly learns reward and value functions within an adversarial formulation to recover a generalizable, portable reward function.This design targets more complex settings than trajectory-level methods while retaining an efficient adversarial formulation.
- AIRL outperforms prior IRL methods on continuous, high-dimensional tasks with unknown dynamics and matches GAIL on tasks without transfer.Under considerable environmental variation, GAIL and other IRL methods fail to generalize according to the introduction.
2 RELATED WORK
The related work situates AIRL within maximum causal entropy IRL and adversarial imitation learning, emphasizing its focus on recovering portable rewards rather than only expert policies. It also distinguishes AIRL from trajectory-centric and multi-task generalization approaches.
- Maximum causal entropy IRL removes ambiguity between demonstrations and the expert policy and casts reward learning as maximum likelihood.This connects inverse reinforcement learning to generative model training.
- GAIL recovers the expert policy but does not recover a reward function, making its representation less portable for transfer.The related work notes that GAIL’s discriminator is unsuitable as a reward at optimality.
- AIRL most closely resembles adversarial methods by Uchibe, Ho and Ermon, and Finn, while interleaving policy optimization with reward learning.The paper links this interleaving to improved sampling efficiency in adversarial training.
- Finn et al.’s adversarial IRL formulation was theoretical, and direct implementation was reported as ineffective because full-trajectory operation produced high variance.The paper motivates moving toward single state-action pairs.
- Unlike Amin et al.’s multi-task setting, AIRL studies generalization within the standard inverse reinforcement learning formulation.
3 BACKGROUND
The paper formulates reinforcement learning as entropy-regularized control in an MDP and interprets IRL as inferring rewards from demonstrations generated by an optimal policy. An adversarial formulation alternates reward-function and policy updates so the reward can be recovered from the discriminator.
- The maximum causal entropy IRL framework uses an entropy-regularized MDP with state and action spaces, transition dynamics, rewards, discounting, and an initial-state distribution.In the standard setup, dynamics, initial states, and rewards are unknown.
- Forward reinforcement learning seeks an optimal policy maximizing expected entropy-regularized discounted reward under the dynamics and initial-state distribution.
- IRL infers the reward function from demonstrations assumed to be generated by an optimal policy and frames this as a maximum likelihood problem.
- Trajectory-centric GAN-GCL updates the discriminator as reward learning and the policy as improvement of the sampling distribution used to estimate the partition function.
- At optimality, GAN-GCL can extract an optimal reward from the discriminator up to a constant while recovering the optimal policy.GAIL shares the adversarial structure but lacks discriminator structure that permits reward recovery.
4 ADVERSARIAL INVERSE REINFORCEMENT LEARNING (AIRL)
AIRL converts adversarial IRL from full trajectories to state-action pairs for efficient learning, then analyzes the resulting reward representation. The state-action formulation can recover the optimal-policy advantage, but that reward may remain entangled with training dynamics and fail under environmental changes.
- Using full trajectories in GAN-GCL can produce high-variance estimates and very poor learning compared with using single state-action pairs.
- AIRL proposes converting the trajectory formulation into a discriminator over individual state-action pairs.The discriminator is defined using the learned function and the policy action probability.
- At optimality, the learned function equals log π*(a|s), which is the optimal policy’s advantage function.The paper provides a justification and proof that the algorithm solves the IRL problem.
- The advantage is a valid optimal reward but is heavily entangled because each action is supervised by the optimal policy for the training MDP.
- The entangled reward is not guaranteed to be robust to changed dynamics and can encourage policy mimicry rather than desirable behavior after environmental changes.
5 THE REWARD AMBIGUITY PROBLEM
Reward shaping can preserve optimal behavior in one environment while making inferred rewards non-robust to changed dynamics. The analysis motivates state-only rewards and characterizes when they remain disentangled across dynamics.
- Reward ambiguity: Reward shaping preserves the optimal policy for any function Φ, leaving IRL unable to distinguish rewards within this transformation class from optimal demonstrations.This ambiguity applies unless the learnable reward-function class is restricted.
- Dynamics dependence: Changing environment dynamics can break policy invariance for a shaped reward recovered in the original MDP.The paper studies two MDPs sharing a reward but differing in transition models T and T′.
- Disentangled rewards: Under a decomposability condition, a recovered state-only reward that is optimal in one dynamics model is disentangled with respect to all dynamics.The condition can hold in environments with self-transitions or in the paper’s experimental environments.
- Disentangled rewards: A reward disentangled for all dynamics must be state-only.The result identifies state dependence as necessary for complete disentanglement across dynamics functions.
- Implications: In traditional single-MDP IRL, restricting rewards to state-only functions can recover the true reward up to a constant when the ground-truth reward is also state-only.This conclusion follows under the paper’s stated assumptions about the reward and dynamics.
6 LEARNING DISENTANGLED REWARDS WITH AIRL
AIRL modifies adversarial reward learning to separate a state-only reward approximator from a shaping term. This restriction enables extraction of rewards intended to be disentangled from training-environment dynamics.
- Motivation: The original AIRL method cannot learn a state-only reward and therefore cannot guarantee that learned rewards are unshaped.The proposed modification targets this limitation by decoupling reward from advantage.
- Discriminator design: AIRL decomposes its discriminator score into a reward approximator gθ and a shaping term hφ.The discriminator is modified so fθ,φ is restricted to these two components.
- Training procedure: The shaping term mitigates unwanted shaping effects while AIRL alternates discriminator training on expert and policy data with policy updates.This training procedure resembles GAIL and GAN-GCL.
- Disentangled reward extraction: Parametrizing gθ(s) solely by state allows AIRL to extract rewards disentangled from the environment dynamics used for training.The claim is established for deterministic environments with a state-only ground-truth reward.
- Interpretation: Under deterministic environments with state-only ground truth, the shaping component recovers the optimal value function up to a constant.The paper identifies h* as V* plus a constant, while stochastic settings use a single-sample advantage interpretation.
7 EXPERIMENTS
The experiments are designed to test whether AIRL learns disentangled rewards that remain robust when environment dynamics change.
- Experimental question: The experiments ask whether AIRL can learn disentangled rewards that are robust to changes in environment dynamics.This is stated as the central experimental question.
2. Is AIRL efficient and scalable to high-dimensional continuous control tasks?
AIRL scales to high-dimensional continuous-control transfer settings and learns rewards that remain effective under substantial domain shifts. In tabular MDPs it recovers the ground-truth reward, while in continuous control it supports transfer across changed environments and agents.
- Benchmark imitation: On benchmark imitation tasks without transfer, AIRL performs on par with GAIL and outperforms GAN-GCL.These results use expert demonstrations and report mean scores across five runs.
- Tabular MDPs: AIRL with a state-only reward recovers the ground-truth reward in randomly generated tabular MDPs.The state-action formulation instead recovers a shaped advantage function.
- Tabular MDPs: Under a new transition matrix, optimizing the state-only reward achieves optimal performance, whereas the state-action reward improves only marginally over uniform random policy.The transfer learning curve plots policy returns against value iteration steps.
- Continuous control: In a shifting maze, only AIRL trained with state-only rewards consistently reaches the goal after the test-time wall change.Direct policy transfer and state-action IRL follow the training path, which is blocked in the modified maze.
- Continuous control: On a disabled-ant task, AIRL learns a reward that supports forward movement with a modified gait despite changes to the agent’s front legs.The learned behavior reorients the ant and crawls with its hind legs, unlike the original sideways-running policy.
- Continuous control: Across transfer experiments, AIRL learns disentangled rewards that accommodate significant domain shifts, while GAIL transfers poorly and GAN-GCL performs poorly even on the original task.The continuous-control evaluation retrains policies in test environments using rewards learned in training environments.
8 CONCLUSION
AIRL is presented as a practical, scalable IRL algorithm that learns disentangled rewards and transfers effectively under domain variation. In small MDPs, it can recover ground-truth rewards up to a constant.
- AIRL learns disentangled rewards and greatly outperforms prior imitation-learning and IRL algorithms.
- AIRL rewards transfer effectively when the underlying domain varies, unlike brittle rewards from unmodified IRL methods.
- In small MDPs with unambiguous optimal policies and rewards, AIRL exactly recovers ground-truth rewards up to a constant.
APPENDICES
The appendices connect AIRL’s discriminator and policy objectives to maximum causal entropy IRL and characterize their equilibrium. They also show how mixture sampling and state-action discrimination support the derivation.
- AIRL’s objective matches maximum causal entropy IRL in the undiscounted case.
- A mixture of the policy and demonstration density estimate reduces importance-sampling variance when early policy coverage is poor.
- The discriminator distinguishes expert demonstrations from generated samples using cross-entropy loss.
- The discriminator-gradient derivation matches the earlier objective when the mixture distribution combines dataset and policy samples.
- At equilibrium, the learned policy equals the expert policy, and the discriminator-derived reward corresponds to the expert advantage function.
B STATE-ONLY INVERSE REINFORCEMENT LEARNING
The appendix analyzes when state-only rewards remain disentangled across dynamics and explains AIRL’s recovery guarantee under deterministic, decomposable environments. The proofs identify state-only structure as essential for robustness to dynamics changes.
- Definitions and assumptions: Under the decomposability condition, all MDP states are linked through one-step reachability and transitivity.
- Theoretical results: If IRL recovers a state-only reward producing an optimal policy, that reward is disentangled across all dynamics.
- Theoretical results: The recovered state-only reward must equal the ground-truth reward up to an additive constant.
- Theoretical results: A reward disentangled for every dynamics function must depend only on the state, not the action or next state.
- Counterexample: Changing dynamics can make an action-dependent shaped reward induce a different policy, including infinite negative reward in the altered MDP.
- AIRL recovery guarantee: AIRL’s recovery guarantee assumes deterministic dynamics, a state-only ground-truth reward, and a state-only reward estimator.