Deriving Two-Stage Q-Learning

From a Second-Stage Decision to a First-Stage Decision

In a two-stage treatment problem, the first action can affect both the final outcome and the information available when the second action is chosen. The optimal first action therefore cannot generally be found by comparing immediate treatment effects alone. It must account for what will happen after the intermediate covariates are observed and the best second action is selected.

Two-stage Q-learning makes this calculation explicit. Starting with the conditional mean of the terminal outcome at the second stage, it optimizes the second action for each observed history. It then averages that optimized value over the intermediate covariates that would follow each possible first action. The first action with the largest resulting value is optimal.

Notation and Conditions

Let \(X\) denote covariates measured before the first treatment \(A_1\), and let \(L\) denote covariates measured after \(A_1\) but before the second treatment \(A_2\). The history available before \(A_2\) is \(H_2=(X,A_1,L)\), and \(Y\) is the terminal outcome to be maximized. For clarity, assume each treatment has a finite action set and that the conditional means below are finite. Finite action sets ensure that the maxima are attained; ties can be resolved by a fixed rule.

These definitions describe conditional outcome models in the observed data. To interpret their optimization causally, suppose consistency holds (the observed outcome equals the potential outcome under the treatments received), there is no interference between individuals, and sequential exchangeability holds: conditional on \(X\), \(A_1\) is independent of the potential intermediate covariates and outcomes; conditional on \(H_2\), \(A_2\) is independent of the potential outcomes under the possible second actions. Also assume positivity: each action the target regime may choose has positive probability at the histories where it is chosen. These conditions, together with integrability, allow the observed conditional means and covariate distributions to represent their intervention counterparts. If positivity fails, the recursion can still be written mathematically, but the relevant causal quantities are not identified from the observed data at unsupported histories.

Deriving the Two-Stage Recursion

Write \(L(a_1)\) for the intermediate covariates that would follow first action \(a_1\), and \(Y(a_1,a_2)\) for the terminal outcome under the two actions. Under the assumptions above, the second-stage conditional outcome model at a history \((x,a_1,l)\) represents the conditional mean outcome under second action \(a_2\). Denote that causal conditional mean by \(m_2(x,a_1,l,a_2)=E[Y(a_1,a_2)\mid X=x,L(a_1)=l]\). Then the optimized value after observing \(L(a_1)\) is the maximum of these conditional means.

The second equality uses sequential exchangeability and consistency to identify the conditional outcome model and the distribution of \(L(a_1)\) from the observed data among people with \(X=x,A_1=a_1\). The first-stage decision is consequently \(d_1^*(x)\in\arg\max_{a_1}Q_1(x,a_1)\).

Proof. Consider any regime \((d_1,d_2)\). Conditional on \(X=x\) and the intermediate history produced after \(d_1(x)\), its conditional mean terminal outcome is \(Q_2(H_2,d_2(H_2))\). For every such history, \(Q_2(H_2,d_2(H_2))\leq\max_{a_2}Q_2(H_2,a_2)\). Taking the conditional expectation over the intermediate covariates gives a value no greater than \(Q_1(x,d_1(x))\). By definition of \(d_1^*\), this is no greater than \(Q_1(x,d_1^*(x))\). A fixed tie rule selects a maximizing action at each history; because the action set is finite, this selection is measurable. Thus \(d_2^*\) attains the first bound, and \(d_1^*\) attains the second. Averaging over \(X\) proves that the stated regime has at least as large a mean outcome as any other regime of the specified form. \(\square\)

Worked Examples

How to Use the Derivation and Diagnose Problems

The recursion describes the population target of Q-learning: determine the optimized second-stage conditional mean, then calculate its conditional expectation under each first action. In an analysis, conditional means and covariate distributions have to be estimated. The derivation itself does not guarantee that fitted models are accurate, or that the data support the actions being compared.

A common mistake is to reverse the order of averaging and optimization. In general, \(E[\max_{a_2}Q_2(H_2,a_2)\mid X,A_1]\) is not equal to \(\max_{a_2}E[Q_2(H_2,a_2)\mid X,A_1]\). The left side allows the second action to respond to the observed history; the right side chooses a single action before averaging over that history. The third worked example shows a strict difference.

Other checks follow directly from the derivation:

• Check the timing of each variable. A variable used in \(H_2\) must be measured before \(A_2\). Omitting a relevant intermediate covariate can change the second-stage conditional means or the distribution being averaged.
• Check action support. If a candidate action is rarely or never observed at a history, its conditional mean may be unstable or unidentified. Positivity is a condition for causal interpretation, not a detail that the maximization can repair.
• Check optimization and averaging separately. At each history, verify which second action maximizes the fitted conditional mean; then average the resulting values using the distribution of intermediate covariates under each first action.
• Check close comparisons. If fitted action-specific means are nearly tied, small prediction errors can change the selected action. A plausible-looking action rule is not by itself evidence that its value is precisely estimated.

Q-learning is a backward outcome-regression approach: it builds the decision rule from conditional outcome models. Evaluating a prespecified regime by the g-formula or inverse-probability weighting instead targets that regime’s value without necessarily searching over actions. Those approaches address different tasks and rely on their own modeling or identification conditions. The advantage of Q-learning is a direct route to an estimated optimal rule; a limitation is that errors in conditional outcome models can affect both the estimated values and the maximizing actions. The recursion identifies the target under its assumptions, but does not remove finite-sample uncertainty or model dependence.