Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distributional Offline Policy Evaluation with Predictive Error Guarantees

Runzhe Wu, Masatoshi Uehara, Wen Sun

TL;DR

This work tackles distributional offline policy evaluation by estimating the full return distribution $Z^\pi$ from offline data. It introduces Fitted Likelihood Estimation (FLE), an MLE-based procedure that can incorporate powerful probabilistic generative models (e.g., Gaussian mixtures, diffusion models) and supports both finite-horizon and infinite-horizon MDPs with vector rewards. The authors prove PAC-style guarantees: finite-horizon accuracy in total variation under data coverage and Bellman completeness, and infinite-horizon accuracy in Wasserstein distance due to contractivity, with rates tied to MLE generalization and horizon. Empirically, FLE with diffusion models and GMMs accurately recovers complex multi-dimensional return distributions, outperforming distributional TD baselines, especially in multi-dimensional settings. Overall, the work demonstrates a flexible, scalable approach to distributional OPE with provable guarantees and broad applicability to risk-sensitive and multi-objective settings.

Abstract

We study the problem of estimating the distribution of the return of a policy using an offline dataset that is not generated from the policy, i.e., distributional offline policy evaluation (OPE). We propose an algorithm called Fitted Likelihood Estimation (FLE), which conducts a sequence of Maximum Likelihood Estimation (MLE) and has the flexibility of integrating any state-of-the-art probabilistic generative models as long as it can be trained via MLE. FLE can be used for both finite-horizon and infinite-horizon discounted settings where rewards can be multi-dimensional vectors. Our theoretical results show that for both finite-horizon and infinite-horizon discounted settings, FLE can learn distributions that are close to the ground truth under total variation distance and Wasserstein distance, respectively. Our theoretical results hold under the conditions that the offline data covers the test policy's traces and that the supervised learning MLE procedures succeed. Experimentally, we demonstrate the performance of FLE with two generative models, Gaussian mixture models and diffusion models. For the multi-dimensional reward setting, FLE with diffusion models is capable of estimating the complicated distribution of the return of a test policy.

Distributional Offline Policy Evaluation with Predictive Error Guarantees

TL;DR

This work tackles distributional offline policy evaluation by estimating the full return distribution from offline data. It introduces Fitted Likelihood Estimation (FLE), an MLE-based procedure that can incorporate powerful probabilistic generative models (e.g., Gaussian mixtures, diffusion models) and supports both finite-horizon and infinite-horizon MDPs with vector rewards. The authors prove PAC-style guarantees: finite-horizon accuracy in total variation under data coverage and Bellman completeness, and infinite-horizon accuracy in Wasserstein distance due to contractivity, with rates tied to MLE generalization and horizon. Empirically, FLE with diffusion models and GMMs accurately recovers complex multi-dimensional return distributions, outperforming distributional TD baselines, especially in multi-dimensional settings. Overall, the work demonstrates a flexible, scalable approach to distributional OPE with provable guarantees and broad applicability to risk-sensitive and multi-objective settings.

Abstract

We study the problem of estimating the distribution of the return of a policy using an offline dataset that is not generated from the policy, i.e., distributional offline policy evaluation (OPE). We propose an algorithm called Fitted Likelihood Estimation (FLE), which conducts a sequence of Maximum Likelihood Estimation (MLE) and has the flexibility of integrating any state-of-the-art probabilistic generative models as long as it can be trained via MLE. FLE can be used for both finite-horizon and infinite-horizon discounted settings where rewards can be multi-dimensional vectors. Our theoretical results show that for both finite-horizon and infinite-horizon discounted settings, FLE can learn distributions that are close to the ground truth under total variation distance and Wasserstein distance, respectively. Our theoretical results hold under the conditions that the offline data covers the test policy's traces and that the supervised learning MLE procedures succeed. Experimentally, we demonstrate the performance of FLE with two generative models, Gaussian mixture models and diffusion models. For the multi-dimensional reward setting, FLE with diffusion models is capable of estimating the complicated distribution of the return of a test policy.
Paper Structure (42 sections, 25 theorems, 101 equations, 3 figures, 11 tables, 2 algorithms)

This paper contains 42 sections, 25 theorems, 101 equations, 3 figures, 11 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Preliminaries
  4. Finite-Horizon MDPs
  5. Discounted Infinite-Horizon MDPs
  6. Offline Policy Evaluation Setup
  7. Fitted Likelihood Estimation
  8. Theoretical Analysis
  9. Finite Horizon
  10. Infinite Horizon
  11. Simulation
  12. One-Dimensional Reward
  13. Two-Dimensional Reward
  14. Discussion and Future Work
  15. Offline CVaR Evaluation
  16. ...and 27 more sections

Key Result

Theorem 4.2

Under asm:cover, suppose we have a sequence of functions $\hat{f}_1,\dots,\hat{f}_H:\mathcal{X}\times\mathcal{A}\mapsto\Delta([0,H]^d)$ and a sequence of values $\zeta_1,\dots,\zeta_H\in\mathbb{R}$ such that holds for all $h\in[H]$. Let our estimator $\hat{f}\coloneqq\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}\hat{f}_1(x,a)$. Then we have

Figures (3)

  • Figure 1: Visualization of the combination lock. The dotted lines denote transiting from good states (white) to bad states (gray). Once the agent transits to a bad state, it stays there forever. The observation is composed of three parts: one-hot encoding of the latent state $w_h$, one-hot encoding of the step $h$, and random noise.
  • Figure 2: Plots of $\mathop{\mathbb{E}}_{x\sim\psi(0,h)}\hat{f}_h(x,a^\star_h)$ and $\mathop{\mathbb{E}}_{x\sim\psi(0,h)}Z^\pi_h(x,a^\star_h)$. The histograms are generated via 50k samples.
  • Figure 3: Plots of $\mathop{\mathbb{E}}_{x\sim\psi(0,h)}\hat{f}_h(x,a^\star_h)$ (generated via 50k samples), and the ground truth $\mathop{\mathbb{E}}_{x\sim\psi(0,h)}Z^\pi_h(x,a^\star_h)$ (top row).

Theorems & Definitions (48)

  • Remark 3.1: Comparison to prior models
  • Remark 3.2: FQE as a special instance
  • Theorem 4.2
  • Lemma 4.4
  • Definition 4.5: Bracketing number
  • Lemma 4.6
  • Corollary 4.7
  • Remark 4.8: Offline CVaR Estimation
  • Lemma 4.9
  • Theorem 4.11
  • ...and 38 more