Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

When Demonstrations Meet Generative World Models: A Maximum Likelihood Framework for Offline Inverse Reinforcement Learning

Siliang Zeng, Chenliang Li, Alfredo Garcia, Mingyi Hong

TL;DR

This work addresses offline inverse reinforcement learning by formulating it as a bi-level maximum-likelihood problem that couples a world-model-based lower-level conservative policy optimization with an upper-level likelihood over expert demonstrations. By learning a dynamics model from fixed data and penalizing uncertain, poorly covered regions, the method recovers high-quality rewards that explain observed behavior while mitigating distribution shift. The authors rigorously characterize the relationship between the true likelihood and a surrogate objective, derive sample-complexity guarantees, and provide convergence results for an efficient alternating optimization algorithm. Empirically, Offline ML-IRL outperforms leading offline IRL and imitation learning baselines on MuJoCo and D4RL, and the recovered reward can even transfer across datasets to enable effective offline RL. This framework advances practical reward learning in settings where interaction with the environment is restricted, with implications for safety-critical applications.

Abstract

Offline inverse reinforcement learning (Offline IRL) aims to recover the structure of rewards and environment dynamics that underlie observed actions in a fixed, finite set of demonstrations from an expert agent. Accurate models of expertise in executing a task has applications in safety-sensitive applications such as clinical decision making and autonomous driving. However, the structure of an expert's preferences implicit in observed actions is closely linked to the expert's model of the environment dynamics (i.e. the ``world'' model). Thus, inaccurate models of the world obtained from finite data with limited coverage could compound inaccuracy in estimated rewards. To address this issue, we propose a bi-level optimization formulation of the estimation task wherein the upper level is likelihood maximization based upon a conservative model of the expert's policy (lower level). The policy model is conservative in that it maximizes reward subject to a penalty that is increasing in the uncertainty of the estimated model of the world. We propose a new algorithmic framework to solve the bi-level optimization problem formulation and provide statistical and computational guarantees of performance for the associated optimal reward estimator. Finally, we demonstrate that the proposed algorithm outperforms the state-of-the-art offline IRL and imitation learning benchmarks by a large margin, over the continuous control tasks in MuJoCo and different datasets in the D4RL benchmark.

When Demonstrations Meet Generative World Models: A Maximum Likelihood Framework for Offline Inverse Reinforcement Learning

TL;DR

This work addresses offline inverse reinforcement learning by formulating it as a bi-level maximum-likelihood problem that couples a world-model-based lower-level conservative policy optimization with an upper-level likelihood over expert demonstrations. By learning a dynamics model from fixed data and penalizing uncertain, poorly covered regions, the method recovers high-quality rewards that explain observed behavior while mitigating distribution shift. The authors rigorously characterize the relationship between the true likelihood and a surrogate objective, derive sample-complexity guarantees, and provide convergence results for an efficient alternating optimization algorithm. Empirically, Offline ML-IRL outperforms leading offline IRL and imitation learning baselines on MuJoCo and D4RL, and the recovered reward can even transfer across datasets to enable effective offline RL. This framework advances practical reward learning in settings where interaction with the environment is restricted, with implications for safety-critical applications.

Abstract

Offline inverse reinforcement learning (Offline IRL) aims to recover the structure of rewards and environment dynamics that underlie observed actions in a fixed, finite set of demonstrations from an expert agent. Accurate models of expertise in executing a task has applications in safety-sensitive applications such as clinical decision making and autonomous driving. However, the structure of an expert's preferences implicit in observed actions is closely linked to the expert's model of the environment dynamics (i.e. the ``world'' model). Thus, inaccurate models of the world obtained from finite data with limited coverage could compound inaccuracy in estimated rewards. To address this issue, we propose a bi-level optimization formulation of the estimation task wherein the upper level is likelihood maximization based upon a conservative model of the expert's policy (lower level). The policy model is conservative in that it maximizes reward subject to a penalty that is increasing in the uncertainty of the estimated model of the world. We propose a new algorithmic framework to solve the bi-level optimization problem formulation and provide statistical and computational guarantees of performance for the associated optimal reward estimator. Finally, we demonstrate that the proposed algorithm outperforms the state-of-the-art offline IRL and imitation learning benchmarks by a large margin, over the continuous control tasks in MuJoCo and different datasets in the D4RL benchmark.
Paper Structure (24 sections, 12 theorems, 121 equations, 5 figures, 3 tables, 1 algorithm)

This paper contains 24 sections, 12 theorems, 121 equations, 5 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries and problem formulation
  3. The world model and statistical guarantee
  4. Algorithm design
  5. Convergence analysis
  6. Numerical results
  7. Conclusion
  8. Experiment details
  9. Detailed experiment setting
  10. Additional numerical results
  11. Related work
  12. Auxiliary lemmas
  13. Proof of Lemma \ref{['lemma:obj_decomposition']}
  14. Proof of Lemma \ref{['lemma:generalization_decomposition']}
  15. Proof of Lemma \ref{['lemma:outer_gradient']}
  16. ...and 9 more sections

Key Result

Lemma 1

Under any reward parameter $\theta$, the objective $L(\theta)$ in objective:ML can be decomposed as below: where $\widehat{L}(\theta)$ is a surrogate objective defined as:

Figures (5)

  • Figure 1: Illustration of the modular structure in our algorithmic framework, Offline ML-IRL. In Offline ML-IRL, it first estimates a world model from the dataset of transition samples, and then implements an ML based offline IRL algorithm on the estimated world model to recover the ground-truth reward function from the collected expert trajectories.
  • Figure 2: The performance of Offline ML-IRL given 5,000 expert demonstrations.
  • Figure 3: The performance of Offline ML-IRL in different environments given $1,000$ expert demonstrations
  • Figure 4: The performance of Offline ML-IRL in different environments given $10,000$ expert demonstrations
  • Figure 5: Reward Transfer. The recovered reward by Offline ML-IRL in the medium-expert dataset is transferred to the medium-replay datasets for solving offline RL tasks.

Theorems & Definitions (22)

  • Lemma 1
  • Lemma 2
  • Proposition 1
  • Theorem 1
  • Lemma 3
  • Lemma 4
  • Theorem 2: Convergence Analysis
  • Theorem 3: Optimality Guarantee
  • Lemma 5
  • Lemma 6
  • ...and 12 more