Keywords: causal inference, policy learning, retrieval-augmented generation, vector search, nearest-neighbor matching, plug-in classifier, conditional average treatment effect, minimax rate, in-context learning.

1 Introduction↩︎

Retrieval-augmented generation (RAG) has become a common architecture for action selection in systems built on language models. In a RAG algorithm, a query is embedded in a vector space, and documents or other information are retrieved according to distances in that space. Conditional on the retrieved evidence, a generative model produces an answer, which may include a recommendation or a tool choice.

Despite its widespread use, the decision-theoretic guarantees of RAG have not been fully studied. If a query is treated as a covariate and the action recommended in the generated answer is treated as the decision, RAG-based decision-making can be viewed as a policy learning problem. Action selection and policy learning have long been studied in causal inference and reinforcement learning.

In causal inference, counterfactual decision-making has been studied as a policy learning problem [1]. The goal is to recommend the action with the highest conditional expected outcome, so estimation of the conditional expected outcome is central. Two main approaches are the plug-in approach and the counterfactual risk minimization (CRM) approach [2]; the latter is also called empirical welfare maximization [3]. The plug-in approach first estimates the conditional expected outcome for each action and then chooses the action with the largest estimate. By contrast, the CRM approach estimates a value functional and learns a policy by maximizing the estimated value.

1.1 Contents and Contributions↩︎

We formulate RAG-based action selection under the potential outcome framework, also known as the Neyman–Rubin causal model. The explicit two-step procedure is a plug-in policy rule because it estimates action-specific conditional expected outcomes before choosing an action. We evaluate the one-step output as a policy without imposing the same interpretation on its unobserved internal computation.

We first consider the case in which \(K\) actions are given. Let \(X\) denote covariates that represent the query or decision context supplied to the RAG system, let \(A\in{\mathcal{K}}\mathrel{\vcenter{:}}=\left\{0,1,\dots,K-1\right\}\) denote the action, and let \(Y(a)\) denote the potential outcome under action \(a\). Our goal is to choose the action with the highest conditional expected outcome, and we write the context-specific optimal action as \[a^*(x)\in\mathop{\mathrm{arg\,max}}_{a\in{\mathcal{K}}}{\mathbb{E}}\left[Y(a)\mid X=x\right].\] Because vector search retrieves documents or examples that are similar or relevant to the query, we interpret the RAG system as producing an estimate of \(a^*(x)\) from the context \(x\). When the retrieval database contains covariates, actions, and outcomes, vector search supplies nearest neighbors used to estimate the conditional expected outcomes that determine \(a^*(x)\). Similarity is therefore used to select matched evidence for each action, rather than to choose the action that appears most often among neighboring cases. This operation corresponds to nearest-neighbor matching in causal inference.

We next extend this formulation to settings in which the available action set is not specified in advance. Although the preceding formulation assumes that \(K\) actions are given, the set of available actions is often not known in advance in applications of language models. For example, in workplace applications, a user may submit a query and ask what to do next without providing a candidate action set. To reflect this setting, we propose two forms of RAG-based policy learning. The one-step method receives a query and directly returns a recommended action, while the actions considered internally remain unobserved. The two-step method proceeds as follows:

• We provide a query and ask the RAG system to generate a finite set of candidate actions.

• For each candidate action, the RAG system retrieves matched evidence and evaluates or ranks the action according to its estimated conditional expected outcome.

• We select the action with the largest estimated conditional expected outcome or the highest rank.

The two-step method separates action generation from action choice. This separation allows us to decompose its regret into the loss from the generated candidate set and the loss from expected-outcome estimation or ranking within that set.

This study makes three contributions:

• We formulate RAG-based action selection as policy learning under the potential outcome framework and give an explicit plug-in interpretation for the two-step procedure.

• We interpret action-specific vector search over observed cases as nearest-neighbor matching and distinguish evidence retrieval from action choice.

• We decompose regret into candidate-set regret and within-candidate regret, and we bound the latter using prediction-error guarantees.

1.2 Related Work↩︎

This study is related to three streams of literature: policy learning, RAG, and nonparametric analysis of neural networks and transformers.

1.2.0.1 Policy learning.

We study policy learning in causal inference [1]–[3]. As in supervised learning [4], policy learning can be approached through plug-in estimation or counterfactual risk minimization. The latter is also referred to as empirical welfare maximization (EWM) [3], and a theoretical comparison of the two approaches is given by [3]. Although that comparison gives favorable results for EWM relative to the plug-in approach under its conditions, their relative performance depends on the underlying data-generating process, as discussed by [4]. Because the structure of RAG is more directly compatible with plug-in estimation, we formulate RAG-based policy learning using the plug-in approach.

In particular, we interpret action-specific retrieval over observed cases as a nearest-neighbor matching procedure. Nearest-neighbor match counts have density-ratio limits that are proportional to inverse propensity scores, and nearest-neighbor matching can also be related to Riesz regression [5], [6]. The analyses of [3] and [1] use inverse probability weighting (IPW) or augmented IPW estimators. The representation results in [5] and [6] therefore connect RAG-based policy learning to those analyses.

1.2.0.2 RAG.

Our study is also related to work on RAG. [7] jointly trains a latent document retriever and a language model. [8] combines a dense passage index with a sequence-to-sequence generator and marginalizes over retrieved passages. Dense passage retrieval maps queries and passages into a shared embedding space and ranks passages by inner-product similarity [9]. This retrieval step is computationally related to nearest-neighbor methods. In particular, the \(k\)-nearest-neighbor language model constructs a nonparametric next-token distribution from nearby contextual representations and interpolates it with the distribution produced by a parametric language model [10]. Later studies develop different ways to incorporate retrieved information into generation. Fusion-in-Decoder encodes retrieved passages separately and combines their representations in the decoder [11]. REPLUG augments a frozen black-box language model and trains the retriever using feedback from the language model [12]. Self-RAG learns when to retrieve and uses self-reflection signals to assess whether the retrieved passages support the generated response [13]. These methods select external evidence according to proximity in a representation space, but they do not interpret the retrieved items as matched observations under the potential outcome framework.

Retrieval has also been incorporated into sequential decision-making. [14] augments reinforcement learning agents with direct access to datasets of past experience. [15] uses approximate nearest-neighbor search over a large collection of expert states to support offline reinforcement learning. REGENT conditions a generalist policy on demonstrations retrieved from related tasks [16], whereas STRAP retrieves relevant sub-trajectories before learning a policy for a target task [17]. These studies place retrieval within reinforcement learning or imitation learning rather than plug-in policy learning under the potential outcome framework. In a different direction, CausalRAG incorporates causal graphs into retrieval and generation for knowledge-intensive tasks [18]. Its use of causal information concerns relations among concepts in the retrieved documents rather than counterfactual outcomes under alternative actions.

Unlike these approaches, we formulate RAG-based action selection under the potential outcome framework. For the two-step method, vector search supplies local evidence for outcome regression, the generator estimates action-specific conditional expected outcomes or their contrasts, and the algorithm selects the action with the largest estimate. The one-step method is evaluated at the level of its returned policy because its internal criterion is not observed. The two-step formulation also separates regret due to candidate generation from regret due to expected-outcome estimation or ranking within the generated set.

1.2.0.3 Nonparametric analysis of neural networks and transformers.

The statistical theory of neural networks provides a basis for viewing transformers as nonparametric estimators. Nonparametric regression with deep ReLU networks has been studied under compositional assumptions and over Besov-type function classes [19]–[21]. For transformers, [22] establishes universal approximation results for sequence-to-sequence functions, while [23] and [24] examine approximation and estimation for high-dimensional sequence inputs and data with low-dimensional structure.

A related literature studies in-context learning as a statistical learning procedure. Existing work examines the function classes that transformers can learn from in-context examples and relates their predictions to regression algorithms, gradient descent, and algorithm selection [25]–[28]. More recent studies analyze in-context learning for nonparametric regression and adaptation to low-dimensional target functions [29]–[31]. We use prediction-error guarantees from this literature as inputs to the regret analysis for transformer-based expected-outcome estimation.

2 Setup↩︎

Let \(X\in{\mathcal{X}}\subseteq\mathbb{R}^d\) denote covariates, which represent a query or decision context. The covariates can combine the query with pre-action individual attributes, such as age, gender, and occupation, when those attributes are available and relevant to the decision problem. Let \(A\in{\mathcal{A}}\) denote an action, where the action space \({\mathcal{A}}\) can be finite or uncountably infinite. Let \(Y\in{\mathcal{Y}}\subseteq{\mathbb{R}}\) denote a scalar outcome. We denote the conditional expected outcome of \(Y\) given \(A=a\in{\mathcal{A}}\) and \(X=x\in{\mathcal{X}}\) by \(f_0(a,x)\mathrel{\vcenter{:}}={\mathbb{E}}\left[Y\mid A=a,X=x\right]\).

2.1 Policy and Policy Value↩︎

We focus on deterministic policies. A policy is a measurable function \[\pi\colon{\mathcal{X}}\to{\mathcal{A}}.\] Let \(\Pi\) be a class of such policies. For \(\pi\in\Pi\), define its value by \[V(\pi) \mathrel{\vcenter{:}}= {\mathbb{E}}\left[f_0(\pi(X),X)\right].\]

We denote a learned or otherwise data-dependent policy by \(\widehat a\colon{\mathcal{X}}\to{\mathcal{A}}\), so that \(\widehat a(x)\) is the action selected at covariate value \(x\). Its value is \[V(\widehat a) \mathrel{\vcenter{:}}= {\mathbb{E}}\left[f_0(\widehat a(X),X)\right].\] When \(\widehat a\) depends on training data, the retrieval database, or auxiliary algorithmic randomness, the expectation also averages over this randomness.

For a chosen action \(\widehat a(\cdot)\), we define its value by \[V(\widehat a)\mathrel{\vcenter{:}}={\mathbb{E}}\left[f_0(\widehat a(X),X)\right].\]

2.2 Optimal Policy↩︎

We assume that the maximum of \(f_0(a,x)\) over \(a\in{\mathcal{A}}\) is attained and that a measurable maximizer can be selected. For each \(x\in{\mathcal{X}}\), let \[a^*(x) \in \mathop{\mathrm{arg\,max}}_{a\in{\mathcal{A}}} f_0(a,x).\] The measurable selector \(a^*\colon{\mathcal{X}}\to{\mathcal{A}}\) is the optimal deterministic policy. Its value is \[V(a^*) = {\mathbb{E}}\left[\max_{a\in{\mathcal{A}}} f_0(a,X)\right].\]

2.3 Regret↩︎

The regret of \(\widehat a\) relative to the optimal action rule \(a^*\) is \[R(\widehat a) \mathrel{\vcenter{:}}= V(a^*)-V(\widehat a) = {\mathbb{E}}\left[ f_0(a^*(X),X)-f_0(\widehat a(X),X) \right].\]

3 RAG-PL↩︎

In this section, we describe our proposed method, called RAG-based policy learning (RAG-PL). We consider two forms of RAG-PL. The one-step method returns an action directly from the covariates. The two-step method first generates candidate actions and then selects one action from those candidates. The two-step method has two stages. The first stage constructs the candidate set. The second stage evaluates the candidates using retrieved evidence and selects one of them.

3.1 RAG↩︎

We consider three input-output forms of a RAG system. The first returns a set of candidate actions given the covariates. he second returns an estimated conditional expected outcome for each action or ranks the candidate actions by that quantity. The third returns a recommended action directly from the covariates.

3.1.0.1 Action-set RAG.

Let \(\mathfrak F({\mathcal{A}})\mathrel{\vcenter{:}}=\left\{C\subseteq{\mathcal{A}}:1\leq |C|<\infty\right\}\) denote the collection of all non-empty finite subsets of \({\mathcal{A}}\). We represent an action-set RAG as a set-valued function \(g\colon{\mathcal{X}}\to\mathfrak F({\mathcal{A}})\). For each \(x\in{\mathcal{X}}\), the set \(g(x)\subseteq{\mathcal{A}}\) contains the candidate actions returned by the RAG system.

3.1.0.2 Vector search as nearest-neighbor matching.

Let \(\varphi\colon{\mathcal{X}}\to\mathbb{R}^{d_\varphi}\) be the embedding map used for vector search, and write \(H=\varphi(X)\). Let \(\mathcal{D}=\left\{D_j\right\}_{j=1}^{N_{\mathrm{db}}}\) be a retrieval database. When \({\mathcal{A}}\) is discrete and the database contains observations of covariates, actions, and outcomes, we write \(D_j=(X_j,A_j,Y_j)\) and, for an action \(a\), define the action-specific index set \(\mathcal{I}_a\mathrel{\vcenter{:}}=\left\{j:A_j=a\right\}\). Given \(x\) and \(a\), let \(\mathcal{N}_{k,a}(x)\) be the indices of the \(k\) observations indexed by \(\mathcal{I}_a\) whose embeddings are closest to \(\varphi(x)\), and let \(\mathcal{R}_k(x,a)\mathrel{\vcenter{:}}=\left\{D_j:j\in\mathcal{N}_{k,a}(x)\right\}\) denote the retrieved evidence. We assume that \(|\mathcal{I}_a|\geq k\) for every action evaluated through action-specific retrieval. This operation matches the current context with similar cases within each action, as in nearest-neighbor matching. For a general document corpus, \(\mathcal{R}_k(x,a)\) denotes evidence retrieved by a query that includes the candidate action. In either case, similarity determines which evidence is used to evaluate an action. It is not itself the criterion for choosing the final action, and we do not select an action by the frequency with which it appears among neighboring cases. For continuous action spaces, exact action-specific retrieval based on \(\mathcal{I}_a=\left\{j:A_j=a\right\}\) is generally unavailable. In that case, retrieval must use an action-aware distance or kernel on \({\mathcal{A}}\times{\mathcal{X}}\), or another procedure that borrows information across nearby actions. The nearest-neighbor analysis in Section 5 is restricted to binary actions.

3.1.0.3 Expected-outcome RAG or ranking RAG.

We represent an expected-outcome RAG by a function \(\widehat f\colon{\mathcal{A}}\times{\mathcal{X}}\to{\mathbb{R}}\). Given an action \(a\in{\mathcal{A}}\) and covariates \(x\in{\mathcal{X}}\), it returns an estimate \(\widehat f(a,x)\) of the conditional expected outcome \(f_0(a,x)\). When the generator and retrieved evidence are written explicitly, we use \[\widehat f(a,x) = G_\theta\bigl(x,a,\mathcal{R}_k(x,a)\bigr).\] A ranking RAG instead receives \(x\) and a finite candidate set \(C=\left\{a_1,\dots,a_m\right\}\in\mathfrak F({\mathcal{A}})\). It retrieves evidence for the candidate actions and returns an ordered tuple \(r(x,C)=(a_{(1)},\dots,a_{(m)})\), which is a permutation of the actions in \(C\) from the highest to the lowest estimated conditional expected outcome.

3.1.0.4 Policy RAG.

A policy RAG receives covariates and directly returns a recommended action. We represent it by a function \(\pi_{\mathrm{RAG}}\colon{\mathcal{X}}\to{\mathcal{A}}\).

3.2 One-Step RAG-PL↩︎

The one-step RAG-PL method applies a policy RAG directly. Given \(x\in{\mathcal{X}}\), the RAG system is instructed to recommend an action \(\widehat a_{\mathrm{one}}(x)\) with the highest conditional expected outcome, and we write \(\widehat a_{\mathrm{one}}(x)=\pi_{\mathrm{RAG}}(x)\). The policy implemented by the RAG system is usually not observed separately from its output. In particular, the one-step method does not expose an intermediate candidate set, expected outcomes for the candidate actions, or a separate ranking step. We therefore evaluate the returned policy without imposing a particular interpretation on its internal computation.

3.3 Two-Step RAG-PL↩︎

3.3.0.1 Action-set generation.

Given \(x\in{\mathcal{X}}\), the action-set RAG returns a finite candidate set \[g(x)=\left\{a_1(x),\dots,a_{M(x)}(x)\right\}\subseteq{\mathcal{A}}, \qquad 1\leq M(x)<\infty.\] Thus, even when \({\mathcal{A}}\) is uncountably infinite, the subsequent evaluation or ranking is performed only over the finite set \(g(x)\). The candidate set can contain actions found in matched cases and actions generated by the language model. Its purpose is to include actions with high conditional expected outcomes, rather than to reproduce the empirical distribution of actions in the retrieved evidence.

3.3.0.2 Ranking or expected-outcome estimation.

For every \(a\in g(x)\), the RAG system retrieves \(\mathcal{R}_k(x,a)\). When an expected-outcome RAG is used, it evaluates \(\widehat f(a,x)=G_\theta(x,a,\mathcal{R}_k(x,a))\). When a ranking RAG is used, it returns the ordered tuple \(r(x,g(x))=(a_{(1)}(x),\dots,a_{(M(x))}(x))\), where \(a_{(1)}(x)\) is the candidate with the highest estimated conditional expected outcome.

3.3.0.3 Action choice.

Under expected-outcome estimation, we choose \[\widehat a_{\mathrm{two}}(x)\in\mathop{\mathrm{arg\,max}}_{a\in g(x)}\widehat f(a,x),\] where a fixed rule is used to break ties. Under ranking, we choose \[\widehat a_{\mathrm{two}}(x)\mathrel{\vcenter{:}}= a_{(1)}(x).\] Both methods define a policy from \({\mathcal{X}}\) to \({\mathcal{A}}\). The two-step method makes action generation, retrieval of similar cases, and action choice explicit, which allows their contributions to regret to be studied separately.

4 Regret Analysis↩︎

We analyze the regret of the one-step and two-step methods under the setup above. We write \(R(\widehat a)\) for the regret of a selected action \(\widehat a(\cdot)\). When the candidate set, the expected-outcome estimate, or the selected action is random, the expectation defining \(R(\widehat a)\) and all expectations below also average over this randomness. We first separate the effect of the generated candidate set from the effect of selecting an action within that set. We then connect the second term to nonparametric prediction error. Section 5 studies the case in which vector search is represented explicitly by nearest-neighbor matching.

4.1 Regret Decomposition for Two-Step RAG-PL↩︎

For each \(x\in{\mathcal{X}}\), let \[a_g^*(x)\in\mathop{\mathrm{arg\,max}}_{a\in g(x)}f_0(a,x),\] where the fixed tie-breaking rule is used when needed, and we assume that the resulting selector is measurable. The maximum is attained because \(g(x)\) is finite. We define \[R_{\mathrm{gen}}(g) \mathrel{\vcenter{:}}= {\mathbb{E}}\left[f_0(a^*(X),X)-f_0(a_g^*(X),X)\right]\] and \[R_{\mathrm{choice}}(\widehat a_{\mathrm{two}};g) \mathrel{\vcenter{:}}= {\mathbb{E}}\left[f_0(a_g^*(X),X)-f_0(\widehat a_{\mathrm{two}}(X),X)\right].\] The first term is the regret caused by restricting the choice to \(g(X)\). The second is the regret caused by expected-outcome estimation or ranking within \(g(X)\). We call \(R_{\mathrm{gen}}(g)\) the candidate-set regret and \(R_{\mathrm{choice}}(\widehat a_{\mathrm{two}};g)\) the within-candidate regret.

Adding and subtracting \(f_0(a_g^*(X),X)\) gives \[\label{eq:two-step-decomposition} R(\widehat a_{\mathrm{two}}) = R_{\mathrm{gen}}(g) + R_{\mathrm{choice}}(\widehat a_{\mathrm{two}};g).\tag{1}\] This identity does not require independence between candidate generation and action choice.

4.2 Bounding the Candidate-Set Regret↩︎

For \(\varepsilon\geq0\), define the set of \(\varepsilon\)-optimal actions by \[{\mathcal{A}}_\varepsilon^*(x) \mathrel{\vcenter{:}}= \left\{a\in{\mathcal{A}}:f_0(a^*(x),x)-f_0(a,x)\leq\varepsilon\right\}.\] Exact inclusion of \(a^*(x)\) is not required. It is enough that the generated set contain an action whose conditional expected outcome is close to the optimum.

The third statement controls the value of the best action in \(g(x)\) through its distance from \(a^*(x)\). It does not concern the estimation error of \(\widehat f\). The smoothness condition on \(f_0\) converts distance in the action space into a difference in conditional expected outcomes.

4.3 Expected-Outcome Estimation and Ranking↩︎

Suppose that the two-step method chooses \[\widehat a_{\mathrm{two}}(x) \in \mathop{\mathrm{arg\,max}}_{a\in g(x)}\widehat f(a,x).\] For each \(x\in{\mathcal{X}}\), define \[\delta_f(x;g) \mathrel{\vcenter{:}}= \max_{a\in g(x)}|\widehat f(a,x)-f_0(a,x)|.\] Let \(v_g^*(x)\mathrel{\vcenter{:}}=\max_{a\in g(x)}f_0(a,x)\) and define the smallest positive gap within the candidate set by \[\Delta_g(x) \mathrel{\vcenter{:}}= \min_{a\in g(x):f_0(a,x)<v_g^*(x)} \left\{v_g^*(x)-f_0(a,x)\right\},\] where \(\Delta_g(x)=+\infty\) if every action in \(g(x)\) attains \(v_g^*(x)\).

The theorem is stated for the realized set \(g(X)\) and therefore does not require candidate generation to be independent of \(\widehat f\). A high-probability error bound can be used in the same way. If \(0\leq f_0(a_g^*(X),X)-f_0(a,X)\leq B_f\) holds almost surely for every \(a\in g(X)\), the margin condition holds, and \[\Pr\left(\delta_f(X;g)>r_f\right)\leq\eta_f,\] then \[\label{eq:high-prob-choice} R_{\mathrm{choice}}(\widehat a_{\mathrm{two}};g) \leq C_M(2r_f)^{1+\kappa}+B_f\eta_f.\tag{2}\] For example, suppose that \(|g(X)|\leq M\) and, conditional on \(X\) and \(g(X)\), \[\Pr\left( |\widehat f(a,X)-f_0(a,X)|>b_n+u \mathrel{\Big|} X,g(X) \right) \leq c_1\exp(-c_2a_nu^2)\] for every \(a\in g(X)\). A union bound gives \[\Pr\left(\delta_f(X;g)>b_n+u\right) \leq c_1M\exp(-c_2a_nu^2).\] Thus, the number of candidates enters the simultaneous error through a logarithmic term when \(u\) is chosen to make the right-hand side small.

For a ranking RAG, let \(a_{(1)}(x)\) be the first action in \(r(x,g(x))\). Then, \[R(a_{(1)}) = R_{\mathrm{gen}}(g) + {\mathbb{E}}\left[f_0(a_g^*(X),X)-f_0(a_{(1)}(X),X)\right].\] When the ranking orders the candidates according to \(\widehat f(a,x)\), Theorem 2 applies. If the RAG system returns only an ordering, its analysis requires a direct bound on the true conditional expected-outcome difference between the best candidate and the first-ranked candidate.

4.4 Relation to Policy-Value Maximization↩︎

Let \[\mathcal{S}_g \mathrel{\vcenter{:}}= \left\{s\colon{\mathcal{X}}\to{\mathcal{A}}:s\text{ is measurable and }s(x)\in g(x)\right\}\] and define the value induced by \(\widehat f\) as \[\widehat V_f(s) \mathrel{\vcenter{:}}= {\mathbb{E}}\left[\widehat f(s(X),X)\right].\] If a measurable selector \(\widehat s_f\) satisfies \[\widehat s_f(x) \in \mathop{\mathrm{arg\,max}}_{a\in g(x)}\widehat f(a,x),\] then, for every \(s\in\mathcal{S}_g\), it holds that \(\widehat f(\widehat s_f(x),x)\geq\widehat f(s(x),x)\). Hence, \[\widehat s_f \in \mathop{\mathrm{arg\,max}}_{s\in\mathcal{S}_g}\widehat V_f(s).\] The same pointwise argument applies to an empirical average over target contexts. Thus, when the policy class contains every measurable selector and its value is constructed from \(\widehat f\), pointwise maximization and maximization over the policy class are the same optimization problem. They can differ when the policy class is restricted or when the value estimator uses observed outcomes through IPW or AIPW rather than being constructed from \(\widehat f\) alone.

4.5 Substitution of Nonparametric Rates↩︎

The preceding results allow a prediction rate to be inserted into the regret bound without changing the RAG-PL algorithm. Suppose that \(R_{\mathrm{gen}}(g)\leq r_{g,n}\) and that the expected-outcome error on the generated set is bounded by \(r_{f,n}\). Under the finite-candidate margin condition in Theorem 2, we have \[R(\widehat a_{\mathrm{two}}) \leq r_{g,n}+C_M(2r_{f,n})^{1+\kappa}.\] For example, if a nonparametric estimator has a uniform error bound \(r_{f,n}\asymp n^{-\alpha/(2\alpha+d)}\), the second term is of order \(n^{-\alpha(1+\kappa)/(2\alpha+d)}\). This substitution is valid only when the available prediction result controls the error required by the regret theorem. An integrated mean-squared error cannot be used as a uniform error without an additional argument.

The same distinction matters for transformers. The analysis of [29] gives a mean-squared prediction bound for a transformer composed of a learned neural representation and a linear-attention layer. In the Besov setting, a representative bound has the form \[q_{\mathrm{ICL}} \lesssim N_{\mathrm{rep}}^{-2\alpha/d} + \frac{N_{\mathrm{rep}}\log N_{\mathrm{rep}}}{n_{\mathrm{ctx}}} + \frac{N_{\mathrm{rep}}^2\log N_{\mathrm{rep}}}{T_{\mathrm{pre}}},\] where \(N_{\mathrm{rep}}\) is the representation dimension, \(n_{\mathrm{ctx}}\) is the number of in-context examples, and \(T_{\mathrm{pre}}\) is the number of pretraining tasks. If this bound holds for both binary expected-outcome estimates, Theorem 3 below gives \[R(\widehat a) \lesssim (q_{\mathrm{ICL},0}+q_{\mathrm{ICL},1})^{(1+\kappa)/(2+\kappa)}.\] A sharper exponent \(1+\kappa\) for the root prediction rate requires a uniform error bound or a pointwise exponential-deviation bound. It does not follow from an MSE result alone. The results of [30] and [31] further describe adaptation to low-dimensional target functions and minimax MSE rates, and they enter the regret analysis according to the same distinction between error types.

The general regret decomposition itself has no sample size. For a specified retrieval database, its size and the number of matched observations have a direct statistical meaning. Transformer theory uses different quantities for the representation dimension, in-context examples, and pretraining tasks. For an already trained RAG system whose training data are not observed, the total number of pretraining tokens is not by itself a sample size for conditional expected-outcome estimation. In that case, the regret bound should be stated in terms of a prediction error that is either assumed or evaluated separately.

4.6 Relation Between One-Step and Two-Step RAG-PL↩︎

The regret of the one-step output is \[R(\widehat a_{\mathrm{one}}) = {\mathbb{E}}\left[f_0(a^*(X),X)-f_0(\widehat a_{\mathrm{one}}(X),X)\right].\] For comparison, its observed output can be represented by the singleton set \(g_{\mathrm{one}}(x)=\left\{\widehat a_{\mathrm{one}}(x)\right\}\). This representation does not assert that the RAG system uses a singleton set internally.

Suppose that \(\widehat a_{\mathrm{one}}(x)\in g(x)\) for \(P_X\)-almost every \(x\). Because \(a_g^*(x)\) is at least as good as every action in \(g(x)\), we have \[R_{\mathrm{gen}}(g) \leq R(\widehat a_{\mathrm{one}}).\] Combining this inequality with 1 gives \[R(\widehat a_{\mathrm{two}}) \leq R(\widehat a_{\mathrm{one}}) + R_{\mathrm{choice}}(\widehat a_{\mathrm{two}};g).\] If the two-step method selects the best action in \(g(x)\) according to \(f_0\), its regret is no larger than the regret of the one-step output. With an estimated expected outcome, Theorem 2 instead gives \[R(\widehat a_{\mathrm{two}}) \leq R(\widehat a_{\mathrm{one}}) + 2{\mathbb{E}}\left[\delta_f(X;g)\right].\] Thus, a larger candidate set can contain a better action than the one-step output, but the gain can be lost when expected-outcome estimation or ranking is inaccurate.

4.7 Fixed Action Sets and Binary Actions↩︎

When the action set is fixed and finite, we can set \(g(x)={\mathcal{A}}={\mathcal{K}}\) for every \(x\). Then, \(R_{\mathrm{gen}}(g)=0\), and Theorem 2 gives \[R(\widehat a_{\mathrm{two}}) \leq 2{\mathbb{E}}\left[\max_{a\in{\mathcal{K}}}|\widehat f(a,X)-f_0(a,X)|\right].\]

For the binary action set \({\mathcal{K}}=\left\{0,1\right\}\), define \[\tau_0(x) \mathrel{\vcenter{:}}= f_0(1,x)-f_0(0,x), \qquad \widehat\tau(x) \mathrel{\vcenter{:}}= \widehat f(1,x)-\widehat f(0,x).\] Using the tie-breaking rule that selects action \(1\) at equality, we have \(a^*(x)=\mathbb{1}\left[\tau_0(x)\geq 0\right]\) and \(\widehat a(x)=\mathbb{1}\left[\widehat\tau(x)\geq 0\right]\). For any selected action \(\widehat a(x)\in\left\{0,1\right\}\), \[\label{eq:binary-regret-identity} R(\widehat a) = {\mathbb{E}}\left[|\tau_0(X)|\mathbb{1}\left[\widehat a(X)\neq a^*(X)\right]\right].\tag{3}\] When \(f_0(a,x)={\mathbb{E}}\left[Y(a)\mid X=x\right]\), \(\tau_0(x)\) is the conditional average treatment effect.

We use the margin condition \[\label{eq:binary-margin} \Pr\left(0<|\tau_0(X)|\leq t\right) \leq C_Mt^\kappa, \qquad 0<t\leq t_0.\tag{4}\] It controls the probability of contexts for which the two actions have nearly equal conditional expected outcomes [4].

5 Nearest-Neighbor Matching and Expected-Outcome Regret↩︎

Vector search connects RAG-PL with causal inference by selecting observations whose covariates are close to the current context. This section gives an explicit nearest-neighbor regression model in which the expected-outcome estimate in Section 3 is formed from matched observations and its error can be substituted into the regret bounds above.

5.1 Nearest-Neighbor Matching↩︎

For each \(a\in\left\{0,1\right\}\), let \[\mathcal{D}_a = \left\{(H_{a,j},Y_{a,j})\right\}_{j=1}^{N_a}\] be a database for action \(a\), where \(H_{a,1},\dots,H_{a,N_a}\) are independently and identically distributed according to \(Q_a\) on \(\mathbb{R}^{d_\varphi}\) and \[Y_{a,j}=m_a(H_{a,j})+\xi_{a,j}.\] We assume that the query \(X\) is independent of the action-specific databases. For a query \(x\), write \(h=\varphi(x)\) and let \(\mathcal{N}_{k,a}(h)\) be the indices of the \(k\) nearest embeddings in \(\mathcal{D}_a\). Define the matched expected-outcome estimate by \[\widetilde{f}_k(a,x) \mathrel{\vcenter{:}}= \frac{1}{k}\sum_{j\in\mathcal{N}_{k,a}(h)}Y_{a,j}.\] The following result gives a standard bias and variance calculation for nearest-neighbor regression in the notation used by RAG-PL [32]. Its use as an action-specific matching estimator follows the setup of [33].

The constant \(C\) depends only on the constants in the assumptions.

The term \(b_{\varphi,a}\) is zero when the embedding retains all information needed for the conditional expected outcome, so \(f_0(a,x)=m_a(\varphi(x))\). Otherwise, it records the difference between conditioning on \(X\) and conditioning on its embedding. The local-mass condition requires enough observations under action \(a\) near each query. It complements the positivity condition used for causal identification.

The displayed estimator is a benchmark for the expected-outcome RAG. Suppose that the RAG output satisfies \[\max_{a\in\left\{0,1\right\}}\sup_{x\in{\mathcal{X}}} |\widehat f(a,x)-\widetilde{f}_k(a,x)| \leq \eta_n.\] Suppose that the margin condition 4 holds. For the deviation-based bound below, also suppose that \(|\tau_0(X)|\leq B_\tau\) holds almost surely. Let \(N_{\min}=\min\left\{N_0,N_1\right\}\) and \(b_\varphi=b_{\varphi,0}+b_{\varphi,1}\). Combining Theorem 4 with Theorem 3 gives \[\begin{align} R(\widehat a) &\leq C\left( b_\varphi + \left(\frac{k}{N_{\min}}\right)^{\beta/d_\varphi} + k^{-1/2} + 2\eta_n \right)^{1+\kappa} + 2B_\tau\exp(-k/4), \label{eq:nn-sharp-regret} \end{align}\tag{5}\] provided that the term in parentheses is sufficiently small relative to \(t_0\). The MSE statement gives the alternative bound \[\label{eq:nn-mse-regret} R(\widehat a) \leq C\left( b_\varphi^2 + \left(\frac{k}{N_{\min}}\right)^{2\beta/d_\varphi} + \frac{1}{k} + \eta_n^2 \right)^{(1+\kappa)/(2+\kappa)}.\tag{6}\] Choosing \(k\asymp N_{\min}^{2\beta/(2\beta+d_\varphi)}\) balances the nearest-neighbor bias and variance. In 5 , this gives the root prediction rate \(N_{\min}^{-\beta/(2\beta+d_\varphi)}\) before the margin exponent is applied. In 6 , the term \(\exp(-k/4)\) is absorbed into \(1/k\).

For a general document corpus without observed outcomes, Theorem 4 does not apply directly. In that case, the theorem describes the matching calculation that the RAG output is meant to approximate, and a regret rate requires an assumption or an evaluation of the difference represented by \(\eta_n\).

5.2 Matching Weights and Propensity Scores↩︎

Similarity is used to select comparable evidence, not to choose the action that appears most often in nearby cases. Let \(\mathcal{N}_k^H(h)\) denote the indices of the \(k\) nearest observations to \(h\) in the pooled embedding database. If \[\widehat e_k^H(a\mid h) = \frac{1}{k}\sum_{j\in\mathcal{N}_k^H(h)} \mathbb{1}\left[A_j=a\right],\] then \(\widehat e_k^H(a\mid h)\) estimates \(\Pr(A=a\mid H=h)\) under standard nearest-neighbor conditions. At \(h=\varphi(x)\), this quantity equals \(\Pr(A=a\mid X=x)\) only under an additional condition such as \(A\perp X\mid H\). Maximizing this quantity selects the most common past action near \(x\), not the action with the largest conditional expected outcome. The joint distribution of \((X,A)\) can remain unchanged while the conditional outcomes are changed so that either action is optimal. Thus, observations of contexts and actions alone do not determine the optimal policy.

Nearest-neighbor matching also has a weighting representation. Let \(X_1^{\mathrm{tar}},\dots,X_m^{\mathrm{tar}}\) be target contexts and define \[\widetilde{f}_k(a,X_i^{\mathrm{tar}}) = \frac{1}{k}\sum^{N_a}_{j=1} \mathbb{1}\left[j\in\mathcal{N}_{k,a}(\varphi(X_i^{\mathrm{tar}}))\right]Y_{a,j}.\] Define \[K_{a,j} \mathrel{\vcenter{:}}= \sum^m_{i=1}\mathbb{1}\left[j\in\mathcal{N}_{k,a}(\varphi(X_i^{\mathrm{tar}}))\right].\] Then, changing the order of summation gives \[\frac{1}{m}\sum^m_{i=1}\widetilde{f}_k(a,X_i^{\mathrm{tar}}) = \sum^{N_a}_{j=1}\frac{K_{a,j}}{mk}Y_{a,j}.\] The match count therefore acts as a weight when local predictions are averaged over target contexts. For the propensity-score interpretation, suppose that the action-specific databases are sampled from the same observational population, so \(Q_a=P_{H\mid A=a}\), and that the target embeddings follow \(P_H\). Bayes’ rule then gives \[\frac{p_H(h)}{p_{H\mid A=a}(h)} = \frac{\Pr(A=a)}{\Pr(A=a\mid H=h)}.\] The density ratio is therefore proportional to an inverse propensity score defined conditional on the embedding. The density ratio is proportional to an inverse propensity score. [5] studies the density-ratio limit of nearest-neighbor match counts, and [6] relates nearest-neighbor matching to least-squares density-ratio estimation and Riesz regression. These results connect matching with IPW and AIPW representations when outcomes are available. They do not make the propensity score an expected outcome.

6 Discussion↩︎

6.1 Identification↩︎

The regret analysis is written in terms of \(f_0(a,x)={\mathbb{E}}\left[Y\mid A=a,X=x\right]\). A causal interpretation requires consistency, conditional exchangeability, and positivity. Under these conditions, it holds that \[{\mathbb{E}}\left[Y\mid A=a,X=x\right] = {\mathbb{E}}\left[Y(a)\mid X=x\right].\] The information contained in \(X\) is therefore part of the identification argument. It is natural to write \(X=(Q,W)\), where \(Q\) is the query and \(W\) contains variables observed before the action, such as age, gender, occupation, and prior decisions. If a variable affects both past action selection and the outcome, omitting it can prevent the conditional comparison from having a causal interpretation. Including more variables does not establish conditional exchangeability by itself, and post-action variables should not be used as controls.

Vector search is usually performed on an embedding \(H=\varphi(X)\). Matching on \(H\) is sufficient only when the representation retains the information required for adjustment or for the conditional expected outcomes. The GenAI-Powered Inference framework combines structured covariates with low-dimensional features extracted from unstructured data [34]. This suggests keeping personal attributes explicit while representing the query text through an embedding. Embedding-powered BISG provides an example in which embeddings are used to estimate a missing personal attribute probabilistically [35]. Such an estimate is a probabilistic proxy rather than an observed covariate, so its estimation error and possible distribution shift remain relevant. The variables used for adjustment can also be richer than the variables allowed in the final policy. In that case, expected outcomes are first adjusted using the larger information set and then averaged over variables excluded from the policy.

Instrumental variables provide another identification strategy when conditional exchangeability fails in the data used to estimate expected outcomes. They enter through the source-data model, not by treating the RAG output as an instrument. Existing studies consider optimal treatment regimes under additional IV assumptions [36], [37] and decision rules under partial identification [38]. These problems have different target values. Applying RAG-PL to them would require the expected outcome supplied to the RAG system to be defined for the corresponding target.

6.2 Expected Outcome and Policy-Value Formulations↩︎

When \(\widehat f(a,x)\) is available and the policy class contains all measurable selectors, maximizing \(\widehat f(a,x)\) at each context is equivalent to maximizing \({\mathbb{E}}\left[\widehat f(s(X),X)\right]\) over the policy class. The same equivalence holds for an empirical average over target contexts. A restricted policy class can introduce an approximation loss, but the criterion is still constructed from the same expected-outcome estimate. The relation between EWM and least-squares CATE estimation for a reparameterized binary policy class is studied by [39].

A different analysis is needed when the policy value is constructed from observed outcomes rather than from \(\widehat f\) alone. IPW uses weighted outcomes, and AIPW combines an expected-outcome estimate with an outcome-residual correction [1], [3]. In the main setting of this study, the analyst starts from a trained RAG system and may not observe unit-level outcomes and assignment probabilities. The fact that outcome-related text may have appeared during pretraining does not provide the data required to construct IPW or AIPW. If a retrieval database contains unit-level observations of \((X,A,Y)\), IPW and AIPW can instead be constructed directly from those data. For one-step RAG-PL, the internal criterion is not observed, so we evaluate the returned policy without assigning it to one of these formulations.

6.3 Margin Conditions↩︎

The margin condition is separate from identification. Identification determines whether the contrast has a causal interpretation, while the margin condition controls how often estimation error changes the preferred action. In the binary case, it bounds the probability that \(|\tau_0(X)|\) is close to zero. When few contexts lie near this boundary, the regret decreases faster than the expected-outcome error.

The margin condition depends on the information used to define the conditional expected outcomes. Omitting a relevant attribute can average positive and negative effects into a contrast near zero, while adding an attribute can also reveal finer effects that lie near zero. Replacing \(X\) by an embedding can change the margin for the same reason. For two-step RAG-PL, the margin within \(g(x)\) is also distinct from the quality of the generated set. A large gap among poor candidates does not remove \(R_{\mathrm{gen}}(g)\).

6.4 Action Generation in Large Action Spaces↩︎

When \({\mathcal{A}}\) is finite and given, \(R_{\mathrm{gen}}(g)=0\). For a large or uncountable action space, the generated set determines which actions can be compared. Exact inclusion of \(a^*(x)\) is stronger than necessary; Proposition 1 uses an \(\varepsilon\)-optimal action, a coverage probability, or distance in the action space. These conditions describe the candidate set through the best conditional expected outcome it contains.

Increasing the number of candidates can improve the chance of including a good action, but more expected outcomes must then be estimated or ranked. Equation 1 keeps these effects separate. A useful candidate-set size depends on both the quality of generation and the accuracy of the second step.

7 Simulation Studies↩︎

We evaluate whether action-specific retrieval of cases with similar pre-action information improves policy choice. Each data-generating process first produces a structured pre-action state together with the conditional expected outcome under every available action. We convert the structured states into complete English queries and case reports before executing the notebooks. The RAG procedures observe only the resulting texts, while the numerical state is retained to calculate the optimal action and regret. Because each numerical level is mapped deterministically to a verbal category, the text retains the full discrete state information. Appendix 15 gives the full specifications.

7.0.0.1 Data-generating processes.

DGP 1 and DGP 2 use two actions. Action 0 is temporary workload reduction and action 1 is weekly individual coaching. Let \(X=(X_1,\ldots,X_6)\) describe recent performance; workload; task complexity; experience; schedule flexibility; and team support. Each component is independently and uniformly distributed on \(\left\{-2,-1,0,1,2\right\}\). DGP 1 has a linear conditional effect and observational action assignment. DGP 2 has randomized action assignment and a smooth nonlinear conditional effect whose population mean is zero. The preferred action therefore varies with the current state even though neither constant action is favored on average. DGP 3 contains 24 interventions formed from four intervention types with three intensity levels and two durations. Its conditional expected outcome rewards agreement between the employee’s stated needs and the intervention profile.

7.0.0.2 Methods and implementation.

The main results use GPT-5.4 mini with temperature zero. Retrieval vectors are obtained from text-embedding-3-small. The embedding input contains only a fixed-order description of the pre-action information. It excludes the historical action and the observed outcome. It also excludes the decision question. Cosine distance is used after normalization.

In DGP 1 and DGP 2, One-step RAG-PL and Two-step RAG-PL receive the same current query and the same six nearest reports under each action. The twelve reports are presented in alternating action order. The prompt also gives the exact mean of the six displayed scores under each action. Providing these means allows the model to use exact local sample averages, so the experiment does not test whether it can calculate those averages from the reports. One-step RAG-PL returns an action directly. Two-step RAG-PL returns the two action-specific conditional expected outcomes, and the experiment code selects the larger estimate. RAG without covariates receives neither the current profile nor historical pre-action profiles. It uses six randomly sampled action-and-outcome reports under each action and returns one population-level decision for the retrieval database. That decision is applied to all ten test queries in the repetition. RAG without covariates is therefore a population-level baseline that differs from RAG-PL in both the current-state information supplied and the retrieval rule. In DGP 3, both RAG-PL procedures receive the same twelve nearest reports and the complete catalog of 24 actions. One-step RAG-PL returns one action directly. Two-step RAG-PL generates five distinct candidate actions. It retrieves five reports observed under each candidate and estimates the five candidate-specific conditional expected outcomes. It then selects the candidate with the largest estimate.

7.0.0.3 Evaluation.

Each DGP is repeated 20 times. Each repetition contains ten test queries. DGP 1 and DGP 2 use 500 historical reports, while DGP 3 uses 600. For each repetition, we first average regret and policy value over its ten queries. Table 1 reports the mean of these repetition-level quantities. The regret standard error is calculated across the 20 repetitions. The optimal-action probability counts any action attaining the largest true conditional expected outcome as optimal. Figures 1–3 show the repetition-level mean regret.

7.0.0.4 Results.

In DGP 1, One-step RAG-PL and Two-step RAG-PL both reduce mean regret from \(4.5450\) under RAG without covariates to \(3.3700\). The paired mean-regret reduction is \(1.1750\) for each procedure, with standard errors of \(0.3562\) for One-step RAG-PL and \(0.3635\) for Two-step RAG-PL. The optimal-action probability rises from \(0.4100\) under RAG without covariates to \(0.5500\) under One-step RAG-PL and \(0.5550\) under Two-step RAG-PL.

DGP 2 isolates personalization because historical actions are randomized. One-step RAG-PL has mean regret \(1.5676\) and Two-step RAG-PL has mean regret \(1.6071\), compared with \(2.3638\) for RAG without covariates. The paired reductions relative to RAG without covariates are \(0.7963\) with standard error \(0.2257\) and \(0.7567\) with standard error \(0.2473\), respectively. Both RAG-PL procedures select an optimal action with probability \(0.6400\), compared with \(0.5100\) under RAG without covariates.

In DGP 3, Two-step RAG-PL lowers mean regret from \(1.4985\) to \(1.2154\). The paired difference, Two-step RAG-PL minus One-step RAG-PL, is \(-0.2831\) with standard error \(0.1431\). The regret decomposition for Two-step RAG-PL gives candidate-set regret \(0.4822\) and within-candidate regret \(0.7332\). Thus, both candidate generation and selection among the generated candidates contribute to the remaining regret. The generated set contains an exactly optimal action with probability \(0.4450\) and an action within \(0.5\) of the optimum with probability \(0.6700\). Additional estimation and robustness results are reported in Appendix 15.9.

8 Conclusion↩︎

We proposed one-step and two-step RAG-based policy learning methods. We interpreted the vector-search component of RAG as nearest-neighbor matching and related this matching procedure to propensity-score weighting, density-ratio estimation, and Riesz regression. This interpretation connects RAG-based decision-making to the existing literature on policy learning. We then derived regret upper bounds for the proposed methods by combining policy-learning analysis with nonparametric prediction-error bounds for nearest-neighbor estimators and transformers. In the simulation studies, both RAG-PL methods had lower mean regret than RAG without covariates in the binary settings, and the two-step method had lower mean regret than the one-step method in the 24-action setting.

9 Proofs for the Regret Analysis↩︎

9.1 Regret decomposition and candidate-set bounds↩︎

For every \(x\in{\mathcal{X}}\), adding and subtracting \(f_0(a_g^*(x),x)\) gives \[\begin{align} &f_0(a^*(x),x)-f_0(\widehat a_{\mathrm{two}}(x),x)\\ &=f_0(a^*(x),x)-f_0(a_g^*(x),x) +f_0(a_g^*(x),x)-f_0(\widehat a_{\mathrm{two}}(x),x). \end{align}\] Taking expectations proves 1 .

9.2 Expected-outcome estimation↩︎

To derive 2 , let \(E_f=\left\{\delta_f(X;g)\leq r_f\right\}\). On \(E_f\), the preceding margin argument applies. On \(E_f^c\), the value difference is at most \(B_f\). Hence, \[\begin{align} R_{\mathrm{choice}}(\widehat a_{\mathrm{two}};g) &\leq C_M(2r_f)^{1+\kappa}+B_f\Pr(E_f^c)\\ &\leq C_M(2r_f)^{1+\kappa}+B_f\eta_f. \end{align}\] The finite-candidate tail bound follows from \[\begin{align} &\Pr\left( \max_{a\in g(X)}|\widehat f(a,X)-f_0(a,X)|>b_n+u \mathrel{\Big|} X,g(X) \right)\\ &\leq \sum_{a\in g(X)} \Pr\left( |\widehat f(a,X)-f_0(a,X)|>b_n+u \mathrel{\Big|} X,g(X) \right)\\ &\leq c_1M\exp(-c_2a_nu^2). \end{align}\]

9.3 Relation between the one-step and two-step methods↩︎

Suppose that \(\widehat a_{\mathrm{one}}(x)\in g(x)\). Then, \[f_0(a_g^*(x),x) \geq f_0(\widehat a_{\mathrm{one}}(x),x),\] so \[f_0(a^*(x),x)-f_0(a_g^*(x),x) \leq f_0(a^*(x),x)-f_0(\widehat a_{\mathrm{one}}(x),x).\] Taking expectations gives \(R_{\mathrm{gen}}(g)\leq R(\widehat a_{\mathrm{one}})\). The remaining inequalities in Section 4 follow from 1 and Theorem 2.

9.4 Proof of Theorem 3↩︎

The identity 3 follows pointwise. If \(\widehat a(x)=a^*(x)\), the loss is zero. Otherwise, the selected action has the lower conditional expected outcome, and the difference is \(|\tau_0(x)|\).

For the first statement, suppose that \(\widehat a(x)\neq a^*(x)\) and \(\tau_0(x)\neq0\). Then, \(\tau_0(x)\widehat\tau(x)\leq0\), which implies \[0<|\tau_0(x)| \leq |\widehat\tau(x)-\tau_0(x)| \leq r_\tau.\] Using 3 and the margin condition, \[\begin{align} R(\widehat a) &\leq {\mathbb{E}}\left[|\tau_0(X)|\mathbb{1}\left[0<|\tau_0(X)|\leq r_\tau\right]\right]\\ &\leq r_\tau\Pr\left(0<|\tau_0(X)|\leq r_\tau\right)\\ &\leq C_Mr_\tau^{1+\kappa}. \end{align}\] If each expected-outcome error is at most \(r_f\), then \[|\widehat\tau(x)-\tau_0(x)| \leq |\widehat f(1,x)-f_0(1,x)| + |\widehat f(0,x)-f_0(0,x)| \leq2r_f.\]

For the second statement, put \(T=|\tau_0(X)|\). On a sign error with \(T>0\), it holds that \(|\widehat\tau(X)-\tau_0(X)|\geq T\). The contribution from \(0<T\leq2b_n\) is at most \[2b_n\Pr\left(0<T\leq 2b_n\right) \leq C_M(2b_n)^{1+\kappa}.\] On \(T>2b_n\), the assumed deviation inequality gives, conditional on \(X\), \[\Pr\left(\widehat a(X)\neq a^*(X)\mid X\right) \leq c_1\exp(-c_2a_nT^2/4)+\delta_n.\] Let \(s=a_n^{-1/2}\) and \(c=c_2/4\). The region \(0<T\leq s\) contributes at most \(C_Ms^{1+\kappa}\). For \(j\geq0\), define \[S_j \mathrel{\vcenter{:}}= \left\{2^js<T\leq 2^{j+1}s\right\}.\] Whenever \(2^{j+1}s\leq t_0\), \[\begin{align} {\mathbb{E}}\left[T\exp(-ca_nT^2)\mathbb{1}\left[S_j\right]\right] &\leq 2^{j+1}s\exp(-c4^j)\Pr\left(0<T\leq 2^{j+1}s\right)\\ &\leq C_Ms^{1+\kappa}2^{(j+1)(1+\kappa)}\exp(-c4^j). \end{align}\] The series over \(j\) is finite. On \(T>t_0\), boundedness gives a term of order \(\exp(-ca_nt_0^2)\), which is bounded by a constant multiple of \(s^{1+\kappa}\) under the stated condition. The contribution of \(\delta_n\) is at most \(B_\tau\delta_n\). Therefore, \[R(\widehat a) \leq C_1b_n^{1+\kappa}+C_2a_n^{-(1+\kappa)/2}+B_\tau\delta_n,\] and the displayed result follows.

For the third statement, let \(E(X)=\widehat\tau(X)-\tau_0(X)\) and choose \(t\in(0,t_0]\). By the sign-error containment, \[\begin{align} R(\widehat a) &\leq {\mathbb{E}}\left[|\tau_0(X)|\mathbb{1}\left[0<|\tau_0(X)|\leq t\right]\right]\\ &\quad+ {\mathbb{E}}\left[|\tau_0(X)|\mathbb{1}\left[\widehat a(X)\neq a^*(X),\;|\tau_0(X)|>t\right]\right]. \end{align}\] The first term is at most \(C_Mt^{1+\kappa}\). On the event in the second term, \(|E(X)|\geq|\tau_0(X)|>t\), so \[|\tau_0(X)| \leq \frac{E(X)^2}{t}.\] Consequently, \[R(\widehat a) \leq C_Mt^{1+\kappa}+\frac{q_n}{t}.\] Taking \(t=q_n^{1/(2+\kappa)}\) proves the result.

10 Proofs for Nearest-Neighbor Matching↩︎

10.1 Nearest-neighbor radius↩︎

Fix an action \(a\) and a query embedding \(h\). Let \[R_{a,k}(h) \mathrel{\vcenter{:}}= \max_{j\in\mathcal{N}_{k,a}(h)}\lVert H_{a,j}-h\rVert.\] The event \(R_{a,k}(h)>r_{a,k}\) occurs only if fewer than \(k\) database points fall in \(B(h,r_{a,k})\). The number of points in that ball is binomial with mean at least \[N_aQ_a(B(h,r_{a,k})) \geq N_ac_ar_{a,k}^{d_\varphi} = 2k.\] The multiplicative Chernoff inequality therefore gives \[\label{eq:nn-radius-tail} \Pr\left(R_{a,k}(h)>r_{a,k}\right) \leq \exp(-k/4).\tag{7}\] Since the support has diameter \(D_H\), \[{\mathbb{E}}\left[R_{a,k}(h)^{2\beta}\right] \leq r_{a,k}^{2\beta}+D_H^{2\beta}\exp(-k/4).\]

10.2 Proof of Theorem 4↩︎

On the event \(R_{a,k}(h)\leq r_{a,k}\), decompose \[\begin{align} \widetilde{f}_k(a,x)-f_0(a,x) ={}& \frac{1}{k}\sum_{j\in\mathcal{N}_{k,a}(h)}\left\{m_a(H_{a,j})-m_a(h)\right\}\\ &+ \frac{1}{k}\sum_{j\in\mathcal{N}_{k,a}(h)}\xi_{a,j}\\ &+ m_a(h)-f_0(a,x). \end{align}\] The absolute values of the first and third terms are bounded by \(L_ar_{a,k}^\beta\) and \(b_{\varphi,a}\), respectively. Conditional on the embeddings, the middle term is mean zero and sub-Gaussian with variance proxy \(\sigma_a^2/k\). Hence, \[\Pr\left( \left|\frac{1}{k}\sum_{j\in\mathcal{N}_{k,a}(h)}\xi_{a,j}\right|\geq u \mathrel{\Big|} \left\{H_{a,j}\right\}_{j=1}^{N_a} \right) \leq 2\exp\left(-\frac{ku^2}{2\sigma_a^2}\right).\] Combining this inequality with 7 proves the pointwise deviation bound.

For the MSE bound, condition on the embeddings. The conditional variance is at most \(\sigma_a^2/k\). The squared conditional bias is bounded by a constant multiple of \[b_{\varphi,a}^2+L_a^2R_{a,k}(h)^{2\beta}.\] Taking expectations and using the radius moment bound gives \[{\mathbb{E}}\left[(\widetilde{f}_k(a,x)-f_0(a,x))^2\right] \leq C\left( b_{\varphi,a}^2 + \left(\frac{k}{N_a}\right)^{2\beta/d_\varphi} + \frac{1}{k} + \exp(-k/4) \right).\] Integrating over \(X\) proves the stated MSE result.

10.3 Derivation of the regret bounds↩︎

Let \[\widetilde{\tau}_k(x) = \widetilde{f}_k(1,x)-\widetilde{f}_k(0,x).\] A union bound and Theorem 4 give, for every \(x\) and \(u>0\), \[\begin{align} &\Pr\left( |\widetilde{\tau}_k(x)-\tau_0(x)| \geq b_\varphi+L_0r_{0,k}^\beta+L_1r_{1,k}^\beta+u \right)\\ &\leq 4\exp\left(-\frac{ku^2}{8\sigma_{\max}^2}\right) +2\exp(-k/4), \end{align}\] where \(\sigma_{\max}=\max(\sigma_0,\sigma_1)\). If \(\widehat f\) differs from \(\widetilde{f}_k\) by at most \(\eta_n\) for each action, then \[|\widehat\tau(x)-\widetilde{\tau}_k(x)|\leq2\eta_n.\] The second statement of Theorem 3 applies with a deterministic term of order \[b_\varphi + \left(\frac{k}{N_{\min}}\right)^{\beta/d_\varphi} + 2\eta_n,\] a concentration scale \(k^{-1/2}\), and a failure probability \(2\exp(-k/4)\). This proves 5 .

For the MSE bound, the inequality \((u+v)^2\leq2u^2+2v^2\) gives \[{\mathbb{E}}\left[(\widehat\tau(X)-\tau_0(X))^2\right] \leq C\left( b_\varphi^2 + \left(\frac{k}{N_{\min}}\right)^{2\beta/d_\varphi} + \frac{1}{k} + \eta_n^2 + \exp(-k/4) \right).\] The third statement of Theorem 3 proves 6 . Balancing \((k/N_{\min})^{\beta/d_\varphi}\) and \(k^{-1/2}\) gives \(k\asymp N_{\min}^{2\beta/(2\beta+d_\varphi)}\).

11 Causal Identification with Covariates and Embeddings↩︎

This section records the conditions under which the observational conditional mean can be interpreted as a potential-outcome conditional mean. We use a discrete action space because that is the setting needed for the matching discussion.

Let \(H=\varphi(X)\). The standard balancing-score argument gives the following result [40].

A balancing representation identifies an average conditional on \(H\). Recovering the more detailed target \({\mathbb{E}}\left[Y(a)\mid X=x\right]\) from matching on \(H\) requires the conditional expected outcome to depend on \(X\) through \(H\).

The two representation conditions have different roles. Balancing concerns adjustment for action assignment. The last proposition concerns the heterogeneity needed to retain the same conditional expected outcomes as the original covariates. The construction of [41] combines structured covariates with features extracted from unstructured data. The estimates studied by [35] can supply a probabilistic measurement of an unavailable attribute, but they do not imply either condition without further assumptions.

The variables used for adjustment need not all appear in the final policy. Suppose \(X=(Q,W)\) and the policy is restricted to depend on \(Q\). After the conditional expected outcomes are adjusted using \((Q,W)\), the value relevant to a \(Q\)-only policy is based on \[f_0^Q(a,q) \mathrel{\vcenter{:}}= {\mathbb{E}}\left[f_0(a,Q,W)\mid Q=q\right].\] The optimal policy within this restricted information set selects an action that maximizes \(f_0^Q(a,q)\).

12 Matching Weights and Outcome Information↩︎

Contexts and historical actions alone do not determine the optimal action. To see this, fix any joint distribution of \((X,A)\). One possible conditional-outcome model sets \(f_0(1,x)=1\) and \(f_0(0,x)=0\) for every \(x\), while another sets \(f_0(1,x)=0\) and \(f_0(0,x)=1\). Both models have the same distribution of \((X,A)\), but their optimal actions are opposite. Outcome information or an external model of expected outcomes is therefore needed.

For a neighborhood \(\mathcal{N}_k(x)\), the local action frequency \[\widehat e_k(a\mid x) = \frac{1}{k}\sum_{j\in\mathcal{N}_k(x)}\mathbb{1}\left[A_j=a\right]\] targets the propensity score \(e_0(a\mid x)=\Pr(A=a\mid X=x)\) under standard nearest-neighbor conditions. If this estimator is consistent and the propensity has a unique maximizer, selecting the largest local frequency converges to the most likely historical action. Its limiting regret is \[{\mathbb{E}}\left[f_0(a^*(X),X)-f_0(a_{\mathrm{hist}}(X),X)\right], \qquad a_{\mathrm{hist}}(x)\in\mathop{\mathrm{arg\,max}}_{a\in{\mathcal{A}}}e_0(a\mid x),\] which need not be zero.

The matching-weight identity in Section 5 follows directly from changing the order of summation. If \[K_{a,j} = \sum_{i=1}^m\mathbb{1}\left[ j\in\mathcal{N}_{k,a}(\varphi(X_i^{\mathrm{tar}})) \right],\] then \[\begin{align} \frac{1}{m}\sum_{i=1}^m\widetilde{f}_k(a,X_i^{\mathrm{tar}}) &= \frac{1}{mk}\sum_{i=1}^m\sum_{j=1}^{N_a} \mathbb{1}\left[j\in\mathcal{N}_{k,a}(X_i^{\mathrm{tar}})\right]Y_{a,j}\\ &= \sum_{j=1}^{N_a}\frac{K_{a,j}}{mk}Y_{a,j}. \end{align}\] Bayes’ rule gives \(p_X(x)/p_a(x)=\Pr(A=a)/\Pr(A=a\mid X=x)\). This is the population relation behind the connection among match counts, density-ratio estimation, and inverse propensity weighting [5], [6].

13 Policy-Value Estimation↩︎

13.1 Value induced by an expected-outcome estimate↩︎

Let \(\widehat s_f(x)\in\mathop{\mathrm{arg\,max}}_{a\in g(x)}\widehat f(a,x)\) be measurable. For any \(s\in\mathcal{S}_g\), \[\widehat f(\widehat s_f(x),x) \geq \widehat f(s(x),x)\] for every \(x\). Taking expectations gives \(\widehat V_f(\widehat s_f)\geq\widehat V_f(s)\). Applying the same inequality at each target context proves the empirical version. This is the equivalence used in Section 4.

13.2 Binary EWM and least-squares CATE estimation↩︎

Let \(\mu_a(x)={\mathbb{E}}\left[Y(a)\mid X=x\right]\) and \(\tau(x)=\mu_1(x)-\mu_0(x)\). For a deterministic binary policy \(\pi\), define \(g_\pi(x)=2\pi(x)-1\). Then, \[\begin{align} V(\pi) &={\mathbb{E}}\left[\mu_0(X)+\pi(X)\tau(X)\right]\\ &={\mathbb{E}}\left[\mu_0(X)\right] +\frac{1}{2}{\mathbb{E}}\left[\tau(X)\right] +\frac{1}{2}{\mathbb{E}}\left[g_\pi(X)\tau(X)\right]. \end{align}\]

13.3 IPW and AIPW policy values↩︎

Suppose that binary actions satisfy conditional exchangeability and positivity, and write \(e(x)=\Pr(A=1\mid X=x)\). For a deterministic policy \(\pi\), the IPW value estimator is \[\widehat V_{\mathrm{IPW}}(\pi) = \frac{1}{n}\sum_{i=1}^n \left( \frac{A_iY_i\pi(X_i)}{\widehat e(X_i)} + \frac{(1-A_i)Y_i(1-\pi(X_i))}{1-\widehat e(X_i)} \right).\] With the true propensity score, iterated expectations give \({\mathbb{E}}\left[\widehat V_{\mathrm{IPW}}(\pi)\right]=V(\pi)\). Let \(\widehat\mu_a(x)\) be expected-outcome estimates. The AIPW value estimator is \[\begin{align} \widehat V_{\mathrm{AIPW}}(\pi) =\frac{1}{n}\sum_{i=1}^n \Bigg(&\widehat\mu_{\pi(X_i)}(X_i)\\ &+ \frac{\mathbb{1}\left[A_i=\pi(X_i)\right]}{\widehat e(A_i\mid X_i)} \left(Y_i-\widehat\mu_{A_i}(X_i)\right) \Bigg), \end{align}\] where \(\widehat e(1\mid x)=\widehat e(x)\) and \(\widehat e(0\mid x)=1-\widehat e(x)\). The second term uses the observed outcome to correct the expected-outcome estimate. This is the statistical difference discussed in Section 6.

14 Nonparametric Rates for Transformer Expected-Outcome Estimates↩︎

The regret theorems take a prediction-error bound as an input. This section records the substitution for the transformer results cited in the main text and keeps the representation dimension, in-context sample size, and number of pretraining tasks separate.

The analysis of [29] considers a transformer composed of a deep neural network with \(N_{\mathrm{rep}}\)-dimensional output and one linear-attention layer. In the Besov setting, one of its main risk bounds has the form \[q_{\mathrm{ICL}} \lesssim N_{\mathrm{rep}}^{-2\alpha/d} + \frac{N_{\mathrm{rep}}\log N_{\mathrm{rep}}}{n_{\mathrm{ctx}}} + \frac{N_{\mathrm{rep}}^2\log N_{\mathrm{rep}}}{T_{\mathrm{pre}}}.\] The three terms are the approximation error, the in-context generalization error, and the pretraining generalization error. When \(T_{\mathrm{pre}}\) is sufficiently large and \(N_{\mathrm{rep}}\asymp n_{\mathrm{ctx}}^{d/(2\alpha+d)}\), the first two terms give the minimax MSE rate \(n_{\mathrm{ctx}}^{-2\alpha/(2\alpha+d)}\) up to logarithmic factors. [30] studies adaptation to low-dimensional target functions, while [31] gives a transformer construction attaining the Hölder minimax MSE rate and states a separate requirement on the number of pretraining sequences. Related approximation and estimation results are given by [23] and [24].

Suppose that the two binary expected-outcome estimates satisfy \[{\mathbb{E}}\left[(\widehat f(a,X)-f_0(a,X))^2\right] \leq q_{a,n}, \qquad a\in\left\{0,1\right\}.\] Then, \[{\mathbb{E}}\left[(\widehat\tau(X)-\tau_0(X))^2\right] \leq 2q_{0,n}+2q_{1,n}.\] The MSE statement of Theorem 3 therefore gives \[R(\widehat a) \leq C(q_{0,n}+q_{1,n})^{(1+\kappa)/(2+\kappa)}.\] Substituting the displayed in-context risk yields the corresponding regret bound. If the prediction result is instead a pointwise exponential-deviation inequality with root scale \(r_n\), the second statement of Theorem 3 gives a bound of order \(r_n^{1+\kappa}\). The MSE result alone does not imply that sharper bound.

The sampling model must also match the RAG application. The in-context examples in these analyses are sampled under the task model used for pretraining, while RAG retrieves examples because they are close to the query. A direct application therefore requires the theoretical sampling condition to cover this selected context. Section 5 instead analyzes nearest-neighbor retrieval directly. For a general trained generator, the prediction error relative to the required conditional expected outcome must be evaluated or assumed separately.

15 Details of the Simulation Studies↩︎

15.1 Prepared Natural-Language Queries and Case Reports↩︎

All natural-language documents are completed before notebook execution. Each DGP uses 100 manually written source templates. Within every repetition, DGP 1 and DGP 2 use each source pattern five times in the historical corpus, while DGP 3 uses each source pattern six times. Ten distinct source patterns are used for the ten test queries in a repetition.

Each historical record contains the covariates together with the observed action and outcome. Each test record contains covariates and a decision question. The outcome is written as a final performance score to one decimal place. The retrieval profile is stored separately from the prose document. It lists the attributes in a fixed order but excludes the action and outcome as well as the decision question. The numerical state is not included in the text supplied to the language model.

For DGP 1 and DGP 2, the levels in \(\left\{-2,-1,0,1,2\right\}\) are expressed as categories from very low to very high. Experience is also expressed in years. For DGP 3, the levels in \(\left\{0.1,0.3,0.5,0.7,0.9\right\}\) are expressed as categories from very low to very high. The same completed corpus and test queries are used by every method within a repetition.

15.2 DGP 1: Linear Conditional Effects with Observational Assignment↩︎

Let \(X=(X_1,\ldots,X_6)\), where the components are independently and uniformly distributed on \(\left\{-2,-1,0,1,2\right\}\). They represent recent performance; workload; task complexity; experience; schedule flexibility; and team support. The two actions are \[0:\text{ temporary workload reduction} \qquad 1:\text{ weekly individual coaching}.\] We define \[\begin{align} \mu_0(x) &=65+4x_1-3x_2-2x_3+2x_4+1.5x_6,\\ \tau_0(x) &=5x_1-4x_2+3x_5,\\ \mu_1(x) &=\mu_0(x)+\tau_0(x). \end{align}\] The historical probability of action 1 is \[e_0(x) =\operatorname{clip}_{[0.10,0.90]} \left( \operatorname{logit}^{-1} \left(0.5x_2-0.4x_4+0.3x_6-0.2x_1\right) \right).\] We draw \(A\mid X=x\sim\operatorname{Bernoulli}(e_0(x))\) and generate \[Y=\mu_A(X)+\epsilon, \qquad \epsilon\sim\mathcal{N}(0,3^2).\] All variables entering the propensity and the two conditional expected outcomes are present in the text.

15.3 DGP 2: Smooth Nonlinear Personalization with Randomized Assignment↩︎

DGP 2 uses the same state support and the same two actions. We define \[\begin{align} \mu_0(x) ={}&65+3\sin\left(\frac{\pi x_1}{4}\right)-2x_2-1.2x_3+1.4x_4+1.2x_6\\ &+0.5x_2x_6-0.4x_4^2 \end{align}\] and \[\begin{align} b(x) ={}&1.25x_1-1.50x_2+x_3-x_4+0.75x_5-0.75x_6\\ &+0.80\sin\left(\frac{\pi(x_1-x_2)}{4}\right) +0.60\sin\left(\frac{\pi(x_3-x_4)}{4}\right),\\ \tau_0(x) ={}&8\tanh\left(\frac{b(x)}{4}\right),\\ \mu_1(x) ={}&\mu_0(x)+\tau_0(x). \end{align}\] We draw \(A\sim\operatorname{Bernoulli}(1/2)\) independently of \(X\) and generate \[Y=\mu_A(X)+\epsilon, \qquad \epsilon\sim\mathcal{N}(0,2.5^2).\] The distribution of \(X\) is symmetric and \(b(-x)=-b(x)\). It follows that \({\mathbb{E}}[\tau_0(X)]=0\). Both actions are optimal on substantial parts of the state space, while any policy that receives no current state is restricted to a population-level choice. The conditional effect is smooth in all six variables used to construct the retrieval profile.

15.4 DGP 3: Large Finite Action Space↩︎

Let \[\mathcal{G}=\left\{0.1,0.3,0.5,0.7,0.9\right\}.\] The context is \(X=(N,I,D,B)\), where \(N=(N_1,N_2,N_3,N_4)\in\mathcal{G}^4\) describes the need for technical skill development; workload relief; managerial guidance; and peer support. The variables \(I\), \(D\), and \(B\) in \(\mathcal{G}\) describe suitable intervention intensity; suitable duration; and baseline performance. All seven variables are drawn independently and uniformly from \(\mathcal{G}\).

An action is a triple \(a=(t,\ell,r)\). The intervention type \(t\) is individual coaching; temporary workload reduction; technical training; or peer support. The intensity level is \(\ell\in\left\{1,2,3\right\}\) and the duration is \(r\in\left\{2,4\right\}\) weeks. The four type profiles are \[\begin{align} v_{\mathrm{coaching}}&=(0.25,0.10,1.00,0.20),\\ v_{\mathrm{workload}}&=(0.05,1.00,0.20,0.05),\\ v_{\mathrm{training}}&=(1.00,0.05,0.35,0.15),\\ v_{\mathrm{peer}}&=(0.20,0.15,0.45,1.00). \end{align}\] The complete action space contains 24 actions. We set \[c(a)=0.35\ell+0.15r\] and define \[f_0(a,x) =62+3B+\frac{9}{2.2}N^\top v_t -4\left(I-\frac{\ell}{3}\right)^2 -3\left(D-\frac{r}{4}\right)^2 -c(a).\] Historical actions are generated from the misspecified score \[s(a,x)=2N^\top v_t-0.7c(a).\] The assignment probability is \[\Pr(A=a\mid X=x) =0.20 \frac{\exp(s(a,x))}{\sum_{a'\in{\mathcal{A}}}\exp(s(a',x))} +\frac{0.80}{24}.\] The assignment rule omits suitable intensity and duration as well as baseline performance, even though all three enter \(f_0\). The realized outcome is \[Y=f_0(A,X)+\epsilon, \qquad \epsilon\sim\mathcal{N}(0,2.5^2).\] Each action has a stable identifier from A01 through A24 and a natural-language description.

15.5 Embedding and Retrieval↩︎

The main experiments use the OpenAI embedding model text-embedding-3-small. Every vector is normalized and cosine distance is used. The embedding input is the fixed-order retrieval profile rather than the full prose report. This profile contains all attributes but excludes the historical action and observed outcome as well as the decision question. The generator receives the full prose query and the full retrieved case reports.

For DGP 1 and DGP 2, the experiment retrieves the six nearest reports observed under action 0 and the six nearest reports observed under action 1. One-step RAG-PL and Two-step RAG-PL receive exactly the same query and reports. Reports are displayed by retrieval rank. The action shown first alternates across ranks, and the initial action is randomized by query. The action descriptions and the two score summaries follow the same randomized order. The prompt states that this order has no priority.

For DGP 3, the initial retrieval returns the twelve nearest reports without restricting the historical action. Both procedures receive these reports and the complete action catalog. Two-step RAG-PL then retrieves five reports within each generated candidate action.

15.6 Implementation of the RAG Procedures↩︎

The main experiments use the model identifier gpt-5.4-mini. The temperature is set to \(0\), and the reasoning effort is set to none. For DGPs 1 and 2, the request configuration allows at most 16,384 input tokens and 256 output tokens. For DGP 3, the corresponding limits are 32,768 input tokens and 1,200 output tokens. Each test query is processed independently, and no conversation history is carried across test queries. All model outputs are validated against predefined structured-output schemas.

In DGPs 1 and 2, every method receives six displayed reports under each action and the exact arithmetic mean of the six displayed outcomes under each action. These summaries are deterministic functions of the displayed reports and therefore add no observations. Providing the summaries to every method ensures that differences among the methods are not driven merely by the model’s ability to calculate sample averages from the displayed outcomes.

One-step RAG-PL returns only a selected action. Two-step RAG-PL returns one estimated conditional expected outcome under each action. The experiment code selects action 1 when its estimated conditional expected outcome is greater than or equal to that under action 0. Thus, action 1 is selected under an exact tie.

RAG without covariates receives neither a current query profile nor historical pre-action profiles. Its historical reports contain only the action and observed outcome. Within each repetition, six reports are sampled uniformly without replacement under each action. The model receives the exact arithmetic mean of the six displayed outcomes under each action and returns two population-level expected-outcome estimates. The experiment code applies the same tie-breaking rule as above. Because the input to this baseline does not vary across the test queries within a repetition, the baseline is queried once for each retrieval database, and its selected action is applied to every test query in that repetition.

In DGP 3, both RAG-PL procedures receive the same current query, the same twelve initially retrieved reports, and the complete catalog of 24 actions. One-step RAG-PL returns one valid action identifier from the catalog. Two-step RAG-PL first returns five distinct valid action identifiers. The procedure then retrieves five reports observed under each candidate action. A subsequent candidate-evaluation request supplies the five candidate actions and their matched reports and asks for one conditional expected-outcome estimate for each candidate. The experiment code selects the candidate with the largest estimate, using a fixed rule to break ties.

During candidate evaluation, valid estimates contained in an incomplete response are retained, and only the missing candidate estimates are requested in subsequent attempts. Candidate generation and candidate evaluation each permit at most eight attempts in total. All experiments reported below obtained complete valid outputs within these limits.

The DGP 3 comparison is not matched in either the number of model requests or the amount of retrieved evidence. In particular, Two-step RAG-PL uses an additional candidate-generation stage, candidate-specific retrieval, and a candidate-evaluation stage. We therefore interpret the DGP 3 results as a comparison of the implemented end-to-end pipelines, rather than as an isolated comparison of one-step and two-step action selection under equal computational budgets.

15.7 Evaluation↩︎

Each DGP contains 20 independent repetitions and ten test queries per repetition. DGP 1 and DGP 2 contain 500 historical reports per repetition. DGP 3 contains 600. All methods in a repetition use the same completed corpus and queries. All 20 repetitions were completed for every method and were included in the main summaries.

For each test query, we record the selected and optimal actions. We also record their true values and the resulting regret. An action is counted as optimal whenever its true conditional expected outcome equals the maximum, so ties are handled by value rather than by agreement with one tie-breaking label. Query-level quantities are averaged within a repetition. Means and standard errors are then calculated across repetitions.

For methods that return expected outcomes, we also calculate expected-outcome mean-squared error. In the binary DGPs, we calculate the squared error of the estimated contrast. For DGP 3, we calculate candidate-specific expected-outcome mean-squared error and centered mean-squared error. We also calculate pairwise ranking accuracy.

For DGP 3, let \(g(x)\) be the five generated candidates and let \(a_g^*(x)\) be the best action in this set. We record \[f_0(a^*(x),x)-f_0(a_g^*(x),x)\] and \[f_0(a_g^*(x),x)-f_0(\widehat a_{\mathrm{two}}(x),x).\] Their sum equals the query-level total regret. We also record exact candidate coverage. Approximate coverage is recorded at value-gap thresholds of \(0.1\), \(0.25\), and \(0.5\).

15.8 Language-Model Robustness↩︎

After all three GPT-5.4 mini experiments are completed, DGP 1 is repeated with Qwen 3.5 4B and Gemma 3 4B. The local models use the same completed documents and OpenAI retrieval vectors as the main DGP 1 experiment. They also use the same matched reports and method definitions. They are used only as RAG generators. Results are reported separately by model because the purpose is to check the direction of the method comparison rather than rank language models.

15.9 Additional Results↩︎

Table 2 reports the expected-outcome errors for procedures that return numerical estimates in the binary DGPs. Two-step RAG-PL has lower expected-outcome and contrast mean-squared error than RAG without covariates in both DGPs.

Table 3 separates the regret of Two-step RAG-PL in DGP 3. Candidate generation accounts for mean regret \(0.4822\), while selection within the generated set accounts for \(0.7332\). Pairwise ranking accuracy within the candidate set is \(0.6200\), and the best candidate is selected with probability \(0.3450\).

Table 4 gives the DGP 1 results for the two local generators. For both models, One-step RAG-PL and Two-step RAG-PL have lower mean regret than RAG without covariates. The relative ordering of the one-step and two-step procedures varies with the generator, so these results are used as a robustness check rather than a model ranking.

References↩︎