Technical Summary: PERRY – Policy Evaluation with Confidence Intervals using Auxiliary Data

Problem Statement

Off-policy evaluation (OPE) is critical for estimating the value of a new reinforcement learning (RL) policy before deployment, particularly in high-stakes domains like healthcare where direct experimentation is risky. Standard OPE methods often struggle when the target policy differs significantly from the behavior policy due to limited dataset coverage. Recent advances suggest leveraging auxiliary datasets, such as synthetic trajectories generated by generative models, can improve OPE accuracy. However, these auxiliary datasets may introduce bias, and existing methods for using data augmentation in OPE lack principled uncertainty quantification. In high-stakes settings, reliable confidence intervals (CIs) are essential for safe policy selection. Currently, no existing methods provide formal guarantees for CIs in OPE settings that combine offline (real) data with synthetic trajectories.

Methodology

The paper proposes two methods, CP-Gen and DR-PPI, to construct valid confidence intervals for OPE in Markov Decision Processes (MDPs) utilizing both real and synthetic data.

1. CP-Gen: Conformal Prediction for State-Conditioned Values

CP-Gen addresses the task of estimating the policy value conditioned on a specific initial state, $V^\pi(s)$. This is crucial in scenarios where stakeholders must select policies for individuals starting in the same state (e.g., patients with identical disease stages).

• Core Concept: The method decomposes the policy value into a simulator estimate (using synthetic trajectories) and a "model bias" term (the discrepancy between real and synthetic returns).
• Mechanism: It employs a novel conformal prediction framework tailored for MDPs with continuous state spaces.
  • It constructs a prediction interval for the return difference between a real trajectory and a corresponding synthetic trajectory starting from the same state.
  • To handle the distribution shift between the behavior policy ($\pi_b$) and the target policy ($\pi_e$), and the continuous nature of the state space, the authors introduce an $\epsilon$-approximation technique. This approximates the importance weights by conditioning on local neighborhoods (balls of radius $\epsilon_s$ around the state and $\epsilon_r$ around the return difference) rather than exact points, making the weight estimation tractable.
  • The final confidence interval for $V^\pi(s)$ is formed by adding the conformal band of the return discrepancy to the synthetic estimate.

2. DR-PPI: Doubly Robust Estimation for Unconditional Values

DR-PPI targets the more common task of estimating the average policy performance over the distribution of initial states, $V^\pi = \mathbb{E}_{s_0}[V^\pi(s_0)]$.

• Core Concept: This method draws inspiration from Doubly Robust (DR) estimation and Prediction-Powered Inference (PPI).
• Mechanism:
  • The behavior dataset is split into two parts ($D_1, D_2$) for cross-fitting.
  • A generative model is trained on one split to produce synthetic trajectories (rollouts) under the target policy.
  • The estimator combines a model-based term (averaging returns from synthetic trajectories) with a correction term derived from the real data. The correction term uses importance sampling (IS) on real trajectories to adjust for the bias introduced by the generative model.
  • Specifically, for each real trajectory, multiple synthetic trajectories are generated from the same initial state to estimate the expected synthetic return, which is then subtracted from the re-weighted real return.
  • The final estimate is the average of two cross-fitted estimators, and the variance is calculated to construct an asymptotic confidence interval.

Key Contributions

1. Formalization: The authors formalize the problem of uncertainty quantification for OPE in MDPs that leverage synthetically generated trajectories.
2. Novel Algorithms: They introduce CP-Gen and DR-PPI, two methods designed for state-conditioned and unconditional policy value estimation, respectively.
3. Theoretical Guarantees:
   • CP-Gen: Proven to produce valid conformal prediction intervals with coverage guarantees that hold asymptotically or within a small margin of error for finite samples, even in continuous state spaces. The margin of error is explicitly bounded as a function of the approximation radii ($\epsilon_s, \epsilon_r$).
   • DR-PPI: Proven to yield asymptotically valid confidence intervals, even if the generative model is misspecified, provided the importance-sampling correction is consistent.
4. Empirical Validation: Extensive experiments across four domains (Inventory Control, Sepsis treatment, Robotic control, and MIMIC-IV EHR data) demonstrate that these methods effectively leverage auxiliary data to produce CIs that cover ground truth values, often with tighter intervals than baselines that rely solely on real data.

Experimental Results

The authors evaluated their methods against several baselines, including Importance Sampling (IS), Direct Method (DM), Doubly Robust (DR), and their augmented variants (AugIS, AugDR), as well as conformal baselines from prior work (Foffano et al., 2023).

• Coverage: Both CP-Gen and DR-PPI consistently produced confidence intervals that covered the ground truth policy values across all tested domains.
• Interval Width: In many settings, particularly those with continuous state spaces (Inventory, D4RL HalfCheetah, MIMIC-IV), the proposed methods produced significantly tighter intervals than baselines. For instance, in the Inventory domain, CP-Gen produced an interval length of 5531.60 compared to 8550.00 for the Foffano et al. baseline, while still covering the ground truth.
• Robustness: The methods were found to be robust to moderate misspecification of the behavior policy and variations in the quality (noise) of the synthetic trajectories.
• Baseline Failures: Baselines that naively augmented data without principled uncertainty correction (e.g., AugIS, AugDR) often produced biased intervals that failed to cover the ground truth or resulted in exponentially large intervals due to variance issues in Q-learning.

Significance and Claims

The paper claims to take a significant step toward enabling uncertainty-aware OPE in regimes where both real and synthetic trajectories are used. The authors emphasize that their work is foundational for high-stakes domains where safe deployment of RL policies requires rigorous uncertainty estimates.

• Bridging the Gap: The work addresses a critical gap where existing OPE methods with data augmentation lack formal assurances on estimate quality.
• Scalability: By introducing the $\epsilon$-approximation, CP-Gen extends conformal prediction to MDPs with continuous and higher-dimensional state spaces, a setting where prior conformal OPE methods struggled.
• Practical Utility: The empirical results suggest that leveraging synthetic data, when paired with these specific uncertainty quantification techniques, can improve the precision of OPE estimates without sacrificing validity.

The authors remain modest regarding limitations, noting that the choice of hyperparameters ($\epsilon_s, \epsilon_r$) becomes increasingly consequential in higher-dimensional settings and that future work is needed to develop more principled procedures for their selection and to explore alternative generative models like diffusion models. They also acknowledge that the theoretical guarantees for CP-Gen involve a small margin of error dependent on these parameters.