Budgeted Multi-Source Counterfactual Annotation for Off-Policy Evaluation

Biao Xiang; Ali Eshragh; Yuexing Li; Kai Wang

cs.LGOct 7, 2026

Budgeted Multi-Source Counterfactual Annotation for Off-Policy Evaluation

Authors: Biao Xiang, Ali Eshragh, Yuexing Li, Kai Wang

Organizations: Pennsylvania State University, State College, PA · Johns Hopkins Carey Business School, Washington, DC · International Computer Science Institute, Berkeley, CA · Georgia Institute of Technology, Atlanta, GA

Abstract

Off-policy evaluation (OPE) estimates the value of a target policy from logged data, but limited behavior-policy coverage can force high-variance reweighting or reward-model extrapolation. Counterfactual annotations can add evidence about unobserved actions, yet practical sources, including domain experts and large language models (LLMs), may be costly, biased, or noisy. We study budgeted acquisition of such annotations for contextual-bandit OPE. Given source-specific costs and error profiles, we formulate an integer allocation problem over context-action pairs and annotation sources to minimize the component of estimator variance that depends on the annotation plan. We characterize when annotations are valuable through a first-annotation threshold and local annotation-value regimes. For the coupled multi-source problem, we develop a majorization-minimization algorithm with dynamic-programming subroutines that monotonically improves the objective. Experiments in synthetic clinical and LLM-annotated education bandits show that our allocation method reduces fixed-profile mean squared error (MSE) by 20.58% and 10.77%, respectively, relative to no annotation.

Figures & tables

Regime	Excess variance	Bias condition	Best local action
I	$\Delta\geq\sigma^{2}$	$\varepsilon^{2}\geq\frac{\eta\sigma_{g}^{2}}{2n_{f}}$	no annotations $(n=0)$
II	$\Delta<\sigma^{2}$	$\varepsilon^{2}>\frac{\eta\sigma_{g}^{2}}{2n_{f}}$	compare $\lfloor n^{}\rfloor,\lceil n^{}\rceil$ , capped by feasibility
III	$\Delta\leq\sigma^{2}$	$\varepsilon^{2}\leq\frac{\eta\sigma_{g}^{2}}{2n_{f}}$	add as many as feasible
IV	$\Delta>\sigma^{2}$	a: $\varepsilon^{2}<\frac{\eta\sigma^{2}}{n_{f}}$	large batch needed to improve over $n=0$
IV	$\Delta>\sigma^{2}$	b: $\frac{\eta\sigma^{2}}{n_{f}}\leq\varepsilon^{2}<\frac{\eta\sigma_{g}^{2}}{2n_{f}}$	no annotations $(n=0)$

Table 1 : Conditions and local allocation decisions for

\eta>0

, excluding the constant-objective case.

Source ID	Type	Bias ( $\varepsilon_{k}$ )	Excess Variance ( $\Delta_{k}$ )	Cost ( $c_{k}$ )
1	Low Bias Low Variance	2.0	6.0	20.0
2	Low Bias High Variance	2.0	48.0	5.0
3	High Bias Low Variance	5.0	3.0	10.0
4	High Bias High Variance	5.0	48.0	1.0

Table 2: Configuration of annotation sources.

Figure 1 : Scaled mean squared error (MSE)

\mathcal{V}

and the allocation-dependent components

J_{1}

J_{2}

, and

J_{3}

versus the number of annotations for the four annotation sources in Table 2 . Shaded bands show

\pm 1

standard deviation over

100

environments.

Figure 2 : Scaled mean squared error (MSE)

\mathcal{V}(\mathbf{n})

for each active source subset. Bars are grouped by the number of available sources; the dashed line is the no-annotation baseline

\mathcal{V}(\mathbf{0})=84.23

Setting	No ann.	Greedy (cost-aware)	Greedy (max-gain)	MM	MM (greedy init)
ASSISTments	88.608	79.198	79.707	79.062	79.062
Original synthetic	84.234	69.275	67.078	66.901	66.8075
I/IV-focused synthetic	84.234	72.644	71.718	70.663	70.720

Table 3: Scaled MSE of allocations from MM and greedy baselines.

Method	Mean gap (%)	Std. gap (%)	Median gap (%)	Max gap (%)
Greedy (max-gain)	0.30	0.30	0.23	0.98
Greedy (cost-aware)	3.67	3.12	3.37	17.38
MM	0.31	0.27	0.26	1.08
MM (greedy init.)	0.05	0.07	0.02	0.26

Table 4: Relative MSE gap to the certified small-instance optima.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3 : Scaled mean squared error (MSE)

\mathcal{V}(\mathbf{n})

for each active source subset under the proposed majorization-minimization (MM) algorithm . Compared with Figure 2 , the single-source bars (red) are now obtained by optimization rather than uniformly random allocation; the multi-source bars are unchanged.

Figure 4 : Convergence trace of scaled mean squared error (MSE)

\mathcal{V}(\mathbf{n}^{(t)})

across majorization-minimization (MM) iterations on the synthetic clinical bandit. Different colors mark the active source subset; the tuple annotated at selected iterations records the allocation

[n_{1},n_{2},n_{3},n_{4}]

across Sources 1–4 in Table 2 .

ID	Problem Type	Original	Median RT	IQR (s)	Correct Rate	Count
0	algebra	No	$11.7$ s	$[5.1,30.1]$	$0.381$	$128{,}221$
1	algebra	Yes	$27.0$ s	$[11.3,63.3]$	$0.678$	$3{,}367{,}980$
2	choose_1	No	$8.1$ s	$[4.3,17.7]$	$0.812$	$127{,}222$
3	choose_1	Yes	$16.1$ s	$[7.5,39.0]$	$0.705$	$1{,}719{,}882$
4	choose_n	No	$11.2$ s	$[5.1,19.4]$	$0.600$	$295$
5	choose_n	Yes	$13.3$ s	$[7.6,27.9]$	$0.788$	$11{,}300$

Appendix

Table 5: Action space of the ASSISTments environment.

Figure 5 : Ground-Truth reward Left: distribution of Ground-Truth Rewards Middle: heatmap of factual context-action pair data Right: influence of Time-efficiency discount on the total reward

Panel of Figure 6 — Figure 6 : Prompt template used for all three large language models (LLMs) on the ASSISTments environment. Bracketed headers [[ ## …## ]] follow the section-delimiter style of Mandyam et al. [4] ; { curly_braces } denote per-pair fields populated from precomputed pandas lookups. Each batch contains $10$ pairs; $20$ independent queries at temperature $0.8$ are issued per $(s,a)$ .

Source	$\operatorname{mean}\|\hat{\varepsilon}_{k}\|$	$\operatorname{median}\|\hat{\varepsilon}_{k}\|$	$\operatorname{mean}\hat{\sigma}_{k}$	Total API cost
Gemini-2.5-Flash	$1.09$	$0.77$	$0.53$	1.14$
GPT-4o-mini	$\mathbf{0.83}$	$0.57$	$0.44$	$\mathbf{\$ 0.48}$
GPT-o3-mini	$0.96$	$0.62$	$0.55$	37.26$
Claude-Haiku-4.5	$0.86$	$\mathbf{0.56}$	$\mathbf{0.28}$	8.11$

Appendix

Table 6 : LLM annotation quality on ASSISTments. Bias and variability are computed after converting predicted response times to rewards. The cost column reports the total API cost for the

20{,}000

predictions used to estimate each model’s empirical bias and variability; the allocation experiment uses the corresponding per-prediction cost.

Figure 7 : Variance objective

J(\mathbf{n})

for each active LLM source subset under uniformly random allocation. Bars are grouped by the number of available sources; the dashed line is the no-annotation baseline

J(\mathbf{0})=88.61

Figure 8 : Variance objective

J(\mathbf{n})

for each active LLM source subset under the proposed MM algorithm

Figure 9 : Convergence trace of

J(\mathbf{n}^{(t)})

across MM iterations on ASSISTments. The tuple at each highlighted iteration denotes the optimal allocation across [ Gemini-2.5-Flash , Claude-Haiku-4.5 , GPT-4o-mini ]. The algorithm heavily favors GPT-4o-mini due to its low empirical bias and high cost-efficiency, while utilizing Claude-Haiku-4.5 more conservatively because of its higher API cost.

Dataset	Method	Gap to Gurobi lower bound	Runtime (s)
Synthetic	Greedy (max-gain)	0.6863%	0.0314
Synthetic	Greedy (cost-aware)	3.9843%	0.0562
Synthetic	MM	1.0474%	8.2390
Synthetic	MM (greedy init.)	0.2807%	5.2657
Synthetic	Gurobi incumbent	0.1258%	3600.6957
ASSISTments	Greedy (max-gain)	0.8216%	1.0222

Appendix

Table 7 : Solver benchmarks and runtimes for the Synthetic and ASSISTments environments.

Method	Synthetic	ASSISTments
Best optimized single source	17.86%	8.30%
Greedy (max-gain)	18.43%	7.86%
Greedy (cost-aware)	14.39%	8.46%
MM	18.67%	8.46%
MM (greedy init.)	19.13%	8.47%

Appendix

Table 8 : Mean relative reduction in empirical mean squared error (MSE) over

30

pilot replications. For each pilot, MSE is estimated from

50

independent final off-policy value estimates.

Explore similar work

Optimal Sequential Annotations for Off-Policy Evaluation
Sep 22, 2026Woojin Chae, Ezinne Nwankwo, Haitong Qin +1Off-Policy Evaluation Learning with Missing Data
University of Southern California · University of California, Berkeley · University of Washington, Seattle
Cross-Domain Off-Policy Evaluation and Learning for Contextual Bandits
Jul 24, 2026Yuta Natsubori, Masataka Ushiku, Yuta SaitoOff-Policy Evaluation Contextual Bandits
Hakuhodo DY Holdings, Inc. · Cornell University
Variance-Optimal Off-Policy Evaluation with Conjunct Effect Modeling
Oct 6, 2026Nicolò Felicioni, Michael Benigni, Maurizio Ferrari Dacrema +1Causal Effect Estimation Off-Policy Evaluation
Spotify · Politecnico di Milano

Sep 22, 2026, stat.MEJ/K move · Enter open · S save

Optimal Sequential Annotations for Off-Policy Evaluation

Woojin Chae, Ezinne Nwankwo, Haitong Qin, Angela Zhou

University of Southern California · University of California, Berkeley · University of Washington, Seattle

Off-Policy Evaluation Learning with Missing Data Human-in-the-Loop Annotation

Open paper AskPDF arXiv

Offline reinforcement learning and off-policy evaluation evaluates dynamic treatment rules based on retrospectively collected data prior to deployment. In recent AI applications, state and reward information is recorded as complex text or image, which recent AI advancements such as LLM-as-a-judge can label with unknown bias. Expert annotation may be available but at a higher cost. For example, safety classification via cheap but imperfect classifiers vs. expensive expert review. We show how a limited budget for ground-truth data-annotation can be used via doubly-robust OPE with missing rewards, and we optimize variance-optimal annotation probabilities for sequential off-policy evaluation, where the target policy value is estimated from annotated data. We characterize the optimal annotation probabilities for sequential forward-monotone annotation protocols, and provide a feasible batch-adaptive implementation. Our work is motivated by a collaboration with a homelessness services nonprofit that writes casenotes for individuals over time. Our method can be used to unlock trustworthy inference from casenote data and answer new inferential questions such as: how does expanding outreach effort over time affect progress towards a housing application and improvement in housing placement? In simulations and on two real datasets - casenotes from the nonprofit and human-preference votes from LMArena - we see reductions in RMSE of 34-65% for housing placement and 17-68% for progress towards a housing application at budgets of 40% of full annotation and above, and by 55-62% at every budget on LMArena.

Abstract

Explore similar work

Optimal Sequential Annotations for Off-Policy Evaluation

Cross-Domain Off-Policy Evaluation and Learning for Contextual Bandits

Variance-Optimal Off-Policy Evaluation with Conjunct Effect Modeling