PaperlayerView exact revision on arXiv

Source-linked AI summary

Prior-Data Fitted Networks for In-Context Off-Policy Evaluation

Shuguang Yu, Ting Li, Fan Zhou

arXiv:2610.34035v1stat.MLcs.LG

TL;DR

Limited behavior-policy data and behavior–target mismatch can make OPE estimates difficult, especially for states and actions rarely observed. InCOPE pretrains a shared value predictor on synthetic OPE tasks with exact target-value labels, then estimates values from new behavior data and known task context without task-specific updates. Across four domains, both InCOPE configurations achieve lower mean MAE than evaluated baselines.

  • Problem

    With limited offline data, target-relevant states and actions may be poorly represented, making target-policy value estimation difficult when behavior and target policies differ.

    Source: §1 · p. 1

  • Method

    InCOPE pretrains a shared predictor on synthetic MDP tasks paired with behavior-policy datasets and exact target-policy values, then uses new behavior data, the target policy, and known task settings as context.

    Source: §1 · pp. 1–2, §4.1 · p. 3

  • Results

    In all four domains, both InCOPE configurations achieve lower mean MAE than every evaluated baseline; specialists reduce error by 19.7%–67.1% relative to the lowest baseline MAE in each domain.

    Source: §5.2 · p. 7

  • Takeaways & Limitations

    Experiments in four domains support more accurate OPE than the evaluated baselines with limited data, while ablations support combining pretraining with current-task observations.

    Source: §6 · p. 10

    The stated test-risk optimality result requires effective-pretraining and test conditional means of target value given context to agree, alongside the corollary’s other assumptions.

    Source: §4.1 · p. 4

Abstract

from arXiv · show

Off-policy evaluation (OPE) is essential for assessing decision policies before deployment when online experimentation is costly or unsafe. Accurate estimation remains difficult when behavior-policy data are limited and the target policy favors states and actions rarely observed in those data. We propose InCOPE, a prior-data fitted network that learns a shared policy-value prediction rule across synthetic OPE tasks. We construct an OPE task prior over Markov decision processes and behavior-target policy pairs, generating behavior-policy datasets and exact target-policy values for pretraining. For a new tabular decision problem, the pretrained network uses the behavior-policy dataset, target policy, and known task settings as context to predict the policy's expected return without task-specific parameter updates. Experiments across four tabular domains, including a semi-synthetic clinical setting, demonstrate InCOPE's effectiveness for policy-value estimation.

1 INTRODUCTION

OPE estimates a target policy’s expected return from behavior-policy data, but limited data and behavior–target mismatch make target-relevant state–action pairs difficult to estimate. InCOPE learns a shared value-prediction rule from synthetic OPE tasks and improves estimation across four tabular task families.

1 INTRODUCTION

  • Limited offline data can leave target-favored states and actions poorly represented, contributing to high-variance importance weighting and inaccurate model- or value-based estimates.These difficulties motivate transferring knowledge from related decision problems.Source: §1 · p. 1
  • InCOPE pretrains on finite behavior-policy datasets paired with exact target-policy values generated from MDPs and policy pairs sampled from a task prior.At evaluation, the network uses the new behavior dataset, target policy, and known task settings to predict return in a forward pass.Source: §1 · pp. 1–2
  • The paper formulates a joint prior over MDPs, behavior and target policies, and observed trajectories, and characterizes the Bayesian policy-value predictor targeted by prior fitting.Source: §1 · p. 2
  • InCOPE improves policy-value estimation across four structurally distinct tabular task families when data are limited and behavior and target policies differ substantially.The pretrained model predicts from historical behavior-policy data without task-specific parameter updates.Source: §1 · p. 2

2 RELATED WORK

Prior OPE work develops estimators within individual tasks and auxiliary-data methods for limited-data settings. InCOPE instead pretrains a reusable value predictor on synthetic MDPs, using behavior-policy data as context and exact target-policy values as supervision.

2 RELATED WORK

  • Existing OPE estimators use value prediction, learned reward and transition models, importance reweighting, marginalized ratios, or combinations such as doubly robust estimation.OPERA combines candidate estimators using bootstrap estimates of joint errors (Nie et al., 2024).Source: §2 · p. 2
  • Auxiliary-information approaches include structural priors for contextual bandits, trajectory augmentation, policy-guided trajectory generation, and auxiliary rollouts for sequential OPE.A complementary route embeds cross-task knowledge in a pretrained value predictor without generating extra trajectories at inference time.Source: §2 · p. 2
  • PFNs learn reusable predictors from synthetic tasks, with prior applications to tabular prediction, causal inference, policy learning, and language-model success prediction.InCOPE addresses sequential OPE by pretraining a shared value predictor on synthetic MDPs, using behavior-policy data as context and exact target-policy values as supervision.Source: §2 · p. 2

3 PROBLEM FORMULATION

The paper studies finite-horizon, time-homogeneous MDPs and estimates a target policy’s expected discounted return using trajectories collected under a behavior policy. The task settings and target policy are known, while the dynamics, rewards, and behavior policy are unknown.

3 PROBLEM FORMULATION

  • The problem is defined on a finite-horizon, time-homogeneous MDP with finite state and action spaces, finite-mean rewards, an initial-state distribution, discount factor, and horizon.Reward and next state are assumed independent conditional on the state and action.Source: §3 · p. 3
  • The dataset contains n independent trajectories generated under a stationary Markov behavior policy, while the target policy is also stationary Markov.Coverage requires behavior-policy visitation of every state–action pair visited with positive probability by the target policy, yet finite data may still omit target-relevant pairs.Source: §3 · p. 3
  • OPE estimates the target policy’s expected discounted return without executing it, given known task settings and target policy but unknown dynamics, rewards, and behavior policy.Source: §3 · p. 3

4 INCOPE: IN-CONTEXT OFF-POLICY EVALUATION

InCOPE predicts target-policy value from a behavior-policy dataset and known task context, with its parameters learned across source OPE tasks and kept fixed at evaluation. Its architecture encodes state–action evidence, propagates information under the queried policy, and pools representations to estimate value.

4 INCOPE: IN-CONTEXT OFF-POLICY EVALUATION

  • InCOPE evaluates a new task in a forward pass from the behavior-policy dataset and target policy, without task-specific parameter updates.The method section covers task priors, network design, pretraining, and evaluation.Source: §4 · p. 3

4.1 PRIOR-FITTED POLICY-VALUE ESTIMATION

  • The source OPE prior varies task properties including trajectory count, horizon, and behavior–target policy pair, and supplies exact target-policy values computed by finite-horizon dynamic programming.The task generator has access to true dynamics and reward distributions; the predictor does not.Source: §4.1 · p. 3
  • The prior can mix multiple task families, with uniform family weights in the experiments; fixed family loss scales reweight errors so larger-scale families do not dominate at comparable relative errors.The resulting effective prior determines the squared-risk objective’s minimizers.Source: §4.1 · pp. 3–4
  • Under the effective pretraining prior, the squared-error-optimal predictor is the conditional mean target value given the context, and it is unique up to zero-probability contexts.It also minimizes test squared error when the effective-pretraining and test conditional means agree under the corollary’s assumptions.Source: §4.1 · p. 4

4.2 NETWORK ARCHITECTURE AND TRAINING

  • The context organizes behavior-data statistics by state–action pair, including visits, reward statistics, policy probabilities, and known task information, alongside a smoothed transition estimate.The network encodes these features before applying action attention within each state.Source: §4.2 · pp. 4–5
  • Policy-conditioned propagation uses target-policy-weighted state representations, state attention, and successor information from each pair’s empirical transition distribution.Weighted pooling over actions under the target policy and states under the initial-state distribution feeds a scalar prediction head and context-dependent output scale.Source: §4.2 · p. 5
  • Training minimizes an empirical squared-error objective over source-training queries paired with exact target-policy values, with source-validation loss used for checkpoint selection.The main comparison uses 3,000 training and 600 validation queries per family per replicate, with disjoint source MDP groups.Source: §4.2 · p. 5

5 EXPERIMENTS

Across four tabular OPE domains, the experiments compare InCOPE with existing estimators and test how data availability, policy mismatch, pretraining-task design, and current-task inputs affect accuracy. InCOPE outperforms evaluated baselines, while results also show that source-task content and current-task observations matter.

5.1 EXPERIMENTAL SETUP

  • The experiments span GridWorld, Tumor, Dispatch, and semi-synthetic Clinical tasks, comparing domain-specific specialists and a pooled model with six OPE estimators using MAE against exact policy values.Specialists use domain-specific synthetic tasks, while the unified model is pretrained on pooled tasks and uses no domain identifier.Source: §5 · p. 5, §5.1 · p. 6

5.2 COMPARISON WITH EXISTING OPE ESTIMATORS

  • Specialist and unified InCOPE models achieve lower mean MAE than every evaluated baseline in all four domains; specialists reduce error by 19.7%–67.1% relative to each domain’s lowest baseline MAE.The advantage persists against MB-OPE, which shares InCOPE’s pooled reward and transition estimates.Source: §5.2 · p. 7, Table 1 · p. 7
  • The unified model beats all baselines across four domains, but specialists perform better on Dispatch and Clinical, similarly on GridWorld, and worse on Tumor.Thus, a shared predictor works across these domains, though domain-specific pretraining remains preferable in some cases.Source: §5.2 · p. 7

5.3 DATA SIZE AND POLICY MISMATCH

  • At the smallest tested data size in every domain, both InCOPE configurations have lower mean MAE than all baselines; at α = 1, both also beat all baselines in every domain.The data-size advantage persists across tested sizes in GridWorld, Tumor, and Clinical; in Dispatch, MB-OPE approaches the specialist at n = 24. The mismatch advantage is clearer as policies diverge in GridWorld and Tumor.Source: §5.3 · pp. 7–8, Figure 2 · p. 7

5.4 EFFECT OF THE PRETRAINING PRIOR

  • Mean test MAE decreases in every domain as more source tasks retain original rewards, favoring original reward mechanisms over the tested randomized-reward alternative at the same source-query budget.The intervention changes reward-dependent inputs and exact value labels, so it compares the resulting pretraining tasks as a whole.Source: §5.4 · p. 8, Table 2 · p. 8
  • The original, broader policy generator yields lower mean MAE in all four domains than the tested narrowed generator.Changing both behavior and target policies changes the OPE problems represented during pretraining; together with the reward comparison, this shows source-task content matters at fixed synthetic training-data quantity.Source: §5.4 · p. 8, Table 2 · p. 8

5.5 CURRENT-TASK INFORMATION AND ESTIMATOR DESIGN

  • Removing current-task observation inputs increases mean MAE in all four domains, with larger differences on Tumor and Dispatch.The restricted model is retrained with the same architecture and supervision, and the results indicate observations add value beyond the learned prior, target policy, and known task specification.Source: §5.5 · p. 9, Table 3 · p. 9
  • Full InCOPE beats every target-policy-only control at every tested trajectory count in all four domains, and has the lowest mean MAE against regressors given complete task inputs in all four domains.The policy-only predictors receive no current-task data; the complete-input comparison matches available information and source supervision, supporting InCOPE’s structured estimator.Source: §5.5 · pp. 9–10, Figure 3 · p. 9, Table 3 · p. 9

6 CONCLUSION

InCOPE is a prior-data fitted network for finite-horizon tabular MDPs that estimates target-policy values from behavior-policy data without task-specific updates. Across four domains, it is more accurate than evaluated baselines with limited data, and ablations support combining pretraining with current-task observations.

6 CONCLUSION

  • Across four domains, InCOPE estimates policy values more accurately than evaluated baselines with limited data.Ablations support combining pretraining with current-task observations.Source: §6 · p. 10
  • InCOPE is pretrained on synthetic OPE tasks with exact value labels and estimates target-policy values from behavior-policy data without task-specific parameter updates.Source: §6 · p. 10
  • The method targets off-policy evaluation in finite-horizon tabular MDPs.Source: §6 · p. 10

AI USE STATEMENT

The statement describes the authors’ use of language models and their review of the work. The appendix notes how fixed family scales affect the effective task prior and limits the claim to the population target.

AI USE STATEMENT

  • Language models assisted across literature search, drafting, translation, method development, experiments, implementation, proofs, and interpretation; authors reviewed the manuscript and manually checked key results.Synthetic samples came from specified simulation models, not directly from language models.Source: AI USE STATEMENT · p. 10

A.1 SQUARED LOSS AND THE EFFECTIVE TASK PRIOR

  • For a specialist or equal family scales, the effective and pretraining task priors are equal.Source: §A.1 · p. 13
  • The reweighting accounts for the training objective without assuming the network observes the family index, but makes no finite-pretraining generalization-bound claim.Source: §A.1 · p. 13

A.2 PROOF OF PROPOSITION 1

The proof decomposes effective risk into a predictor-independent conditional-variance term and a nonnegative error term, yielding the conditional expectation as the unique minimizer up to null sets.

A.2 PROOF OF PROPOSITION 1

  • The proof uses conditional expectation and iterated expectations to show that the cross term between prediction error and the conditional residual vanishes.Source: §A.2 · p. 13
  • Effective risk separates a predictor-independent conditional-variance term from a nonnegative term measuring deviation from the conditional mean.Source: §A.2 · p. 13
  • The minimizer is the conditional expectation of the target value given context, unique up to null sets.Source: §A.2 · p. 13

A.3 PROOF OF COROLLARY 2

Under conditional-mean compatibility and context support, the ideal predictor is the test conditional mean and minimizes test squared risk. The appendices also specify state–action preprocessing and a horizon-conditioned attention and propagation architecture.

A.3 PROOF OF COROLLARY 2

  • Under context support and conditional-mean compatibility, the ideal predictor equals the test conditional mean and minimizes test squared risk.Compatibility concerns conditional means, not complete value distributions or underlying MDPs.Source: §A.3 · pp. 13–14
  • The compatibility condition is sufficient for the ideal prediction rule; it does not require matching complete conditional value distributions or underlying MDPs.Matching marginal value ranges alone does not establish the condition, nor does the condition establish that an experimental prior satisfies it.Source: §A.3 · p. 14

B.1 INPUT FEATURES AND PREPROCESSING

  • Each state–action pair receives features combining visit and reward statistics, policy probabilities, and known task settings, while the estimated transition tensor is retained separately.The 25 input coordinates use the same preprocessing in pretraining and evaluation; task-level scalars are repeated across pairs and reward coordinates retain their original units.Source: Table 4 · p. 15, §B.1 · p. 15
  • Preprocessing smooths behavior probabilities and transitions, assigns uniform estimates to unobserved states or state–action pairs, and propagates target-policy visitation in the empirical model.The feature construction also uses trajectory statistics and a separate availability vector indicating known initial-state distribution and estimated behavior policy.Source: §B.1 · pp. 15–16

B.2 NETWORK ARCHITECTURE

  • The network encodes state–action features, attends over actions within each state, then propagates successor information through horizon-conditioned hidden-vector updates.The architecture uses two action-attention blocks and two propagation blocks; these are not horizon-length scalar Bellman recursions.Source: §B.2 · p. 16

B.3 TRAINING DETAILS

This section describes the training and model-selection protocol and how source and evaluation tasks are generated, including task settings, policies, trajectory data, and exact policy-value labels.

B.3 TRAINING DETAILS

  • Each of five replicates uses 3,000 training and 600 validation queries per family; AdamW training uses effective batches of 32 and retains the lowest-validation-loss checkpoint.Validation uses training-derived scales; early stopping starts at epoch 8 with patience 8, while training runs for at most 48 epochs.Source: §B.3 · p. 17

C.1 TASK PRIORS AND TEST ENVIRONMENTS

  • The task-prior protocol specifies source MDP generators, behavior and target policies, logs, and exact value labels, alongside separate test generators.Source MDP parameters are available to the simulator for generating logs and labels, but are not given to the predictor instead of its observed context.Source: §C.1 · p. 17

C.1.1 GRIDWORLD

  • GridWorld source tasks contain four equally represented mechanisms, while evaluation uses a separate 4×4 generator with four modes; test reference values are exact finite-horizon dynamic-programming values.The main comparison had 2,500 separately generated MDP/policy/log queries, without shared-log query groups.Source: §C.1.1 · pp. 18, 20

C.1.2 TUMOR

  • Tumor source tasks represent burden and accumulated toxicity, whereas test tasks use burden alone; both use 25-step horizons and exact dynamic-programming value labels.Source state spaces have 24, 36, or 48 states, and test states comprise 24 representative burden values.Source: §C.1.2 · pp. 21–22

C.1.3 DISPATCH

  • Dispatch models a two-sided vehicle-service market with supply-incentive and demand-promotion actions; source and test tasks share action semantics but use different zone dynamics and imbalance kernels.Source tasks vary zone counts and graph structures and balance three imbalance mechanisms: matching, queue relaxation, and conserved-fleet coupling.Source: §C.1.3 · p. 23
  • Dispatch rewards combine service surplus, imbalance and intervention costs, and state-dependent benefits, with a fixed transformation adjusting the no-intervention baseline without changing raw action contrasts.The transformation belongs to the source reward generator rather than post-processing predicted policy values.Source: §C.1.3 · p. 24
  • Each replicate contains 3,000 training and 600 validation queries across a balanced factorial design, with exact finite-horizon target-policy values computed by dynamic programming.Queries vary task shape, horizon, trajectory count, and target concentration; each uses its own MDP, behavior log, and target policy, with γ = 0.97 and independently sampled reward noise.Source: §C.1.3 · p. 25
  • The five test batches contain 300 queries each; each main-evaluation query pairs one MDP/log/target-policy triple with an exact dynamic-programming value in the test MDP.Test rollouts use γ = 0.97 and independent reward noise with standard deviation 0.03, separate from fixed reward-table perturbation ξsa. The initial distribution is uniform over zones at neutral imbalance.Source: §C.1.3 · p. 25

C.1.4 CLINICAL

The clinical evaluation uses a semi-synthetic aggregate of the released ICU-Sepsis MDP, while a separate synthetic treatment-like prior supplies pretraining tasks. The appendix also specifies the OPE baseline implementations used in the comparisons.

C.1.4 CLINICAL

  • Clinical evaluation aggregates the released ICU-Sepsis MDP and generates new simulated logs; it is not new clinical data collection or a test of treatment effectiveness.The released MDP was constructed from MIMIC-III records.Source: §C.1.4 · p. 25
  • The separate source prior uses 48 live states and three absorbing outcomes, with 25 actions representing combinations of five fluid and five vasopressor levels.Its synthetic source process does not perturb or resample the ICU-Sepsis transition table.Source: §C.1.4 · p. 25
  • Source rewards vary in concentration while being scaled to a chosen uniform-policy value at horizon 20; infeasible peak constraints are rejected rather than changing the value label.The value is computed by exact dynamic programming, and the source scaling uses synthetic quantities rather than ICU-Sepsis values or evaluation labels.Source: §C.1.4 · p. 27
  • Each source replicate generates training and validation MDPs, multiple behavior-log conditions, and five target-policy queries per log, with exact finite-horizon dynamic-programming labels.Each replicate has 50 training and 10 validation MDPs; each log is queried by five target policies, and trajectories use deterministic expected rewards.Source: §C.1.4 · p. 28

C.2 BASELINE IMPLEMENTATIONS

  • FQE updates separate time-indexed value estimates and assigns zero to unobserved cells, while Tabular-MIS estimates rewards and transitions separately by time and is algebraically equivalent under the same initial distribution.OPERA aggregates four full-log estimators using bootstrap-based weights; its fourth candidate is a scalar, cross-fitted Shaped-PDIS adaptation.Source: §C.2 · pp. 29–30
  • The pooled model-based baseline fits a time-homogeneous empirical MDP using observations pooled across time, with zero reward and a uniform transition row for unobserved state-action pairs.It uses the known initial distribution without source pretraining or importance weighting.Source: §C.2 · p. 30

C.3 MAIN EVALUATION DETAILS

The main evaluation pools 6,700 queries across four domains and reports batch-level error variability. Additional experiments vary trajectory count and policy mismatch while comparing estimators on shared test conditions.

C.3 MAIN EVALUATION DETAILS

  • The main comparison covers 6,700 queries across four domains, with five separately generated test batches per domain and results spanning the listed trajectory counts and horizons.Clinical contributes 1,200 queries from 240 logs, with five target policies queried per log.Source: §C.3 · p. 31, Table 14 · p. 31
  • Reported MAE summaries average five batch-level errors, and their sample SD measures test-batch variability with InCOPE checkpoints held fixed.For each InCOPE configuration, five pretrained checkpoints are evaluated on every batch and their MAEs averaged within that batch.Source: §C.3 · p. 31

C.4 TRAJECTORY COUNT AND POLICY MISMATCH

  • The trajectory-count and policy-mismatch experiment uses 50 test tasks per domain, holding each task fixed across conditions while varying sample size and behavior policy.Clinical tasks within a block share an aggregate environment, while the other domains use a separately generated environment for each task.Source: §C.4 · pp. 31–32
  • Mismatch levels interpolate between each task’s target policy and a behavior anchor selected by policy distance, not estimation error.The levels are fractions of task-specific anchor distance, so absolute distances can differ across tasks and domains.Source: §C.4 · p. 32
  • At each mismatch level, the experiment generates one largest-size log and uses its trajectory prefixes for smaller sample sizes, with estimators sharing each evaluation log and target policy.Figure 2 uncertainty bands use 2,000 bootstrap draws that resample blocks and tasks, plus training replicates for InCOPE.Source: §C.4 · p. 32

C.5 PRETRAINING PRIOR ABLATIONS

These ablations vary source-task rewards and policy range while holding the comparison test queries fixed within each experiment. They also describe how the modified source tasks and replicate-level errors are constructed.

C.5 PRETRAINING PRIOR ABLATIONS

  • Each configuration is evaluated on shared test queries within its experiment, and Table 2 reports mean and sample SD across five training replicates on a fixed test set.Modified priors are separately trained, whereas original-prior comparisons reuse pretrained specialists.Source: §C.5 · p. 33
  • The source-reward intervention replaces rewards while retaining transitions, initial distributions, policies, and logged states, actions, and successors, then recomputes reward-dependent inputs and exact value labels.The source prior mixes 0%, 50%, or 100% original-reward tasks, assigning mixtures by complete predefined groups.Source: §C.5 · p. 33
  • The narrow policy-range setting retains source MDPs, rewards, horizons, and trajectory counts, but replaces both behavior and target policies with mixtures involving a uniform policy.Logs are resampled and value labels and affected inputs recomputed; the test distribution is unchanged.Source: §C.5 · p. 33

C.6 INFORMATION AND ESTIMATOR CONTROLS

The controls compare InCOPE with restricted-input retraining and conventional regressors. Separate experiments test whether current-task log information and richer task inputs distinguish the pretrained predictor from alternatives.

C.6 INFORMATION AND ESTIMATOR CONTROLS

  • The log-free ablation is retrained from scratch with log-derived inputs removed from source training, validation, and evaluation, rather than masking a full-input model only at test time.Its architecture, source-query identities, value labels, seed schedule, optimizer, and stopping rule match full InCOPE.Source: §C.6 · p. 34
  • Policy-only regressors receive only a padded target-policy probability matrix, without logs, initial distribution, discount, horizon, or other task inputs.Their hyperparameters are selected using source-validation MAE, and their final comparison uses a separately generated test set.Source: §C.6 · p. 34
  • Table 16 presents the policy-only comparison over all test queries rather than stratifying results by trajectory count.The policy-only curves grouped by trajectory count use different queries at different sample sizes.Source: Table 16 · p. 35, §C.6 · p. 35
  • Complete-input regressors receive the preprocessed state-action features, empirical transition tensor, masks, and availability indicators used by InCOPE, without raw trajectories or a domain identifier.Their comparison uses a new shared test set and reports replicate-level MAEs across five training replicates.Source: §C.6 · p. 35

Related papers

Loading 2610.34035v1…