High-Dimensional Off-Policy Evaluation without Sample Splitting using an Onsager Correction
Abstract
Reward models used for off-policy evaluation (OPE) are commonly selected using held-out data. When sample size and context dimension grow proportionally, the difference in test risk between models fitted on the full sample and on a fixed fraction of it need not vanish. Held-out evaluation can therefore consistently estimate the test risk of the model trained on the smaller sample, but because the relative ranking of candidate reward models can change after full-data refitting, held-out selection can be suboptimal. We develop an Onsager-corrected, target-policy-weighted criterion that estimates the full-data reward model's out-of-sample risk under the target policy using only its training observations. This enables early stopping along the full-data gradient-descent trajectory without sample splitting. For neural networks with fixed depth and hidden width under Gaussian contexts, single-index rewards, context-dependent multi-index target policies, and known context-independent logging, we prove uniform consistency of our criterion over a fixed finite set of gradient-descent iterates and a near-oracle guarantee for the selected reward model. The same correction yields a doubly robust estimator that consistently estimates the true policy value without sample splitting, even when reward-model prediction error does not vanish. Experiments show decreasing risk-estimation error as dimension grows at a fixed sample-to-dimension ratio, and improved reward-model selection relative to held-out baselines. Empirically, our correction also reduces the root mean squared error of doubly robust value estimates computed from the same data used to fit and select the reward models.
Related
- A stochastic convex integration scheme for the intermittent Onsager theorem in two-dimensional domains
- On the Onsager--Machlup functional of the infinite-dimensional angular central Gaussian distribution and related measures
- An Anisotropic Onsager Criterion for the Two-Dimensional Navier--Stokes Equations with Horizontal Viscosity
- A Prediction--Correction Analysis of Two-Way Block Splitting in Distributed Learning
- Correction to the article "Global well-posedness and scattering in weighted space for nonlinear Schrödinger equations below the Strauss exponent without gauge-invariance"
BibTeX
@misc{jain2026highdimensional,
title = {{High-Dimensional Off-Policy Evaluation without Sample Splitting using an Onsager Correction}},
author = {Saksham Jain and Guillaume Braun and Masaaki Imaizumi},
howpublished = {arXiv},
year = {2026},
month = nov,
doi = {10.48550/arxiv.2611.08794},
eprint = {2611.08794},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2611.08794},
}
RIS
TY - UNPB
TI - High-Dimensional Off-Policy Evaluation without Sample Splitting using an Onsager Correction
AU - Saksham Jain
AU - Guillaume Braun
AU - Masaaki Imaizumi
PY - 2026
DA - 2026/11/05
PB - arXiv
DO - 10.48550/arxiv.2611.08794
UR - https://arxiv.org/abs/2611.08794
AB - Reward models used for off-policy evaluation (OPE) are commonly selected using held-out data. When sample size and context dimension grow proportionally, the difference in test risk between models fitted on the full sample and on a fixed fraction of it need not vanish. Held-out evaluation can therefore consistently estimate the test risk of the model trained on the smaller sample, but because the relative ranking of candidate reward models can change after full-data refitting, held-out selection can be suboptimal. We develop an Onsager-corrected, target-policy-weighted criterion that estimates the full-data reward model's out-of-sample risk under the target policy using only its training observations. This enables early stopping along the full-data gradient-descent trajectory without sample splitting. For neural networks with fixed depth and hidden width under Gaussian contexts, single-index rewards, context-dependent multi-index target policies, and known context-independent logging, we prove uniform consistency of our criterion over a fixed finite set of gradient-descent iterates and a near-oracle guarantee for the selected reward model. The same correction yields a doubly robust estimator that consistently estimates the true policy value without sample splitting, even when reward-model prediction error does not vanish. Experiments show decreasing risk-estimation error as dimension grows at a fixed sample-to-dimension ratio, and improved reward-model selection relative to held-out baselines. Empirically, our correction also reduces the root mean squared error of doubly robust value estimates computed from the same data used to fit and select the reward models.
ER -