STITCH-OPE: Trajectory Stitching with Guided Diffusion for Off-Policy Evaluation

TL;DR

STITCH-OPE introduces a diffusion-based model framework for accurate off-policy evaluation in high-dimensional, long-horizon problems, significantly reducing variance and improving estimation accuracy.

Contribution

It presents a novel diffusion-guided trajectory stitching method that enhances long-horizon off-policy evaluation with theoretical variance reduction guarantees.

Findings

Outperforms existing methods on D4RL and OpenAI Gym benchmarks.

Achieves lower mean squared error and higher correlation in estimates.

Demonstrates exponential variance reduction in theoretical analysis.

Abstract

Off-policy evaluation (OPE) estimates the performance of a target policy using offline data collected from a behavior policy, and is crucial in domains such as robotics or healthcare where direct interaction with the environment is costly or unsafe. Existing OPE methods are ineffective for high-dimensional, long-horizon problems, due to exponential blow-ups in variance from importance weighting or compounding errors from learned dynamics models. To address these challenges, we propose STITCH-OPE, a model-based generative framework that leverages denoising diffusion for long-horizon OPE in high-dimensional state and action spaces. Starting with a diffusion model pre-trained on the behavior data, STITCH-OPE generates synthetic trajectories from the target policy by guiding the denoising process using the score function of the target policy. STITCH-OPE proposes two technical innovations that make it advantageous for OPE: (1) prevents over-regularization by subtracting the score of the behavior policy during guidance, and (2) generates long-horizon trajectories by stitching partial trajectories together end-to-end. We provide a theoretical guarantee that under mild assumptions, these modifications result in an exponential reduction in variance versus long-horizon trajectory diffusion. Experiments on the D4RL and OpenAI Gym benchmarks show substantial improvement in mean squared error, correlation, and regret metrics compared to state-of-the-art OPE methods.