Estimating Earlier-Stage Q Functions

Why an Earlier-Stage Q-Function Uses a Later-Stage Value

In a multistage treatment problem, the best action now can depend on what happens to the patient afterward. An earlier-stage Q-function captures that dependence by asking: for each current history and current action, what outcome should be expected if treatment is chosen optimally at later stages? Estimating it therefore requires more than predicting the observed outcome under the treatment patients happened to receive.

Consider a two-stage problem. Let $H_1$ be the history available before the first treatment $A_1$, let $H_2$ be the history available before the second treatment $A_2$, and let $Y$ be the final outcome, with larger values preferred. For simplicity, take the final outcome as the only reward. At stage 2, the Q-function gives the conditional mean outcome under each available action. The optimized stage-2 value selects the best such conditional mean at each history.

Definition: For an allowed second-stage action set $\mathcal{A}_2(h_2)$, define $$ Q_2(h_2,a_2)=E[Y\mid H_2=h_2,A_2=a_2],\qquad V_2(h_2)=\max_{a_2\in\mathcal{A}_2(h_2)}Q_2(h_2,a_2). $$ The earlier-stage Q-function is $$ Q_1(h_1,a_1)=E[V_2(H_2)\mid H_1=h_1,A_1=a_1]. $$ It leaves the current action $a_1$ fixed and averages over the later history that follows it, while using the optimized later-stage value at each possible $H_2$.

If there is also a reward $R_1$ earned at the first stage, the corresponding definition is $Q_1(h_1,a_1)=E[R_1+V_2(H_2)\mid H_1=h_1,A_1=a_1]$. The estimation principle is unchanged: regress the appropriate continuation value, including any first-stage reward, on $H_1$ and $A_1$.

Identification and the Regression Target

A causal interpretation requires assumptions about how treatment is assigned. Consistency means the observed history and outcome agree with the potential history and outcome under the treatments actually received. Sequential exchangeability means that, conditional on the measured history at each stage, treatment assignment is independent of the relevant future potential outcomes. Positivity means every action used by the target regime has positive probability, or adequate support, at histories where it is to be applied. The required support applies at both stages.

Under these assumptions, the conditional outcome regression at stage 2 identifies the mean outcome under each available second-stage action. Maximizing those means gives the value of the best supported action for that history. At stage 1, conditioning on $H_1,A_1$ then identifies the expected optimized value of the next history. In an observational study, this interpretation depends on measuring the confounders needed for sequential exchangeability; a regression does not make unmeasured confounding disappear.

Result: Suppose consistency, sequential exchangeability, and positivity hold, and the conditional means are well defined. For a supported first-stage action $a_1$ and the optimal second-stage rule $d_2(h_2)\in\arg\max_{a_2} Q_2(h_2,a_2)$, the expected final outcome under $A_1=a_1$ followed by $d_2$, conditional on $H_1=h_1$, equals $Q_1(h_1,a_1)$.

To verify the result, condition first on the stage-2 history. By consistency and stage-2 exchangeability, the conditional mean outcome under the action $d_2(h_2)$ is $Q_2(h_2,d_2(h_2))=V_2(h_2)$. Now average over the history reached after the first action. By consistency and stage-1 exchangeability, its conditional distribution is the distribution of $H_2$ given $H_1=h_1,A_1=a_1$. Thus

E[Y^{a_1,d_2}\mid H_1=h_1] =E\!\left[E[Y^{a_1,d_2}\mid H_2^{a_1},H_1=h_1]\mid H_1=h_1\right] =E[V_2(H_2)\mid H_1=h_1,A_1=a_1] =Q_1(h_1,a_1).

This argument also explains the order of operations: optimize at each later history, then average over the later histories associated with the current action. In general, maximizing after averaging would answer a different question.

Fitting the Earlier-Stage Regression

In practice, estimate the stage-2 outcome regression $\hat Q_2(h_2,a_2)$ from the observed second-stage data. For each observed $H_{2i}$, calculate $\hat V_2(H_{2i})$ by maximizing the fitted Q-function over the allowed second-stage actions. Then use these calculated values as responses in a regression on $H_{1i}$ and $A_{1i}$. The fitted regression estimates $Q_1$; it does not yet choose the first-stage action.

\hat V_2(H_{2i})=\max_{a_2\in\mathcal{A}_2(H_{2i})}\hat Q_2(H_{2i},a_2), \qquad \hat Q_1\ \text{is fitted by regressing }\hat V_2(H_{2i})\text{ on }(H_{1i},A_{1i}).

For a finite set of actions, the maximization compares the fitted predictions for each permitted action. For a continuous action, it uses the constrained maximization appropriate to the action space. The earlier-stage regression itself must represent how the optimized value varies with both the earlier history and current action; omitting an interaction or using an overly rigid functional form can obscure that relationship.

Result: If $m(X)=E[Z\mid X]$ and $E[Z^2]$ is finite, then $m$ minimizes $E[(Z-g(X))^2]$ over measurable functions $g$ with finite squared error. Indeed, write $Z-g(X)=(Z-m(X))+(m(X)-g(X))$, square and take expectations. The cross-term is zero because $E[Z-m(X)\mid X]=0$. Therefore $$ E[(Z-g(X))^2]=E[(Z-m(X))^2]+E[(m(X)-g(X))^2]\geq E[(Z-m(X))^2]. $$ Equality holds precisely when $g(X)=m(X)$ almost surely. Applying this result with $Z=V_2(H_2)$ and $X=(H_1,A_1)$ justifies the population regression target. $\square$

A finite-dimensional regression generally estimates the best approximation to this conditional mean within its chosen model class, not necessarily the conditional mean itself. The result just proved identifies the target of squared-error regression; it does not guarantee that a particular model or finite sample recovers that target.

Worked Examples

Worked Example: Binary Histories and Treatments

Suppose the fitted stage-2 Q-values are $Q_2(0,0)=0.2$, $Q_2(0,1)=0.4$, $Q_2(1,0)=0.8$, and $Q_2(1,1)=0.6$, where the first argument is binary $H_2$ and the second is binary $A_2$. At $H_2=0$, the larger value is $0.4$; at $H_2=1$, it is $0.8$. Thus $V_2(0)=0.4$ and $V_2(1)=0.8$.

For patients with $H_1=0,A_1=0$, suppose $P(H_2=0\mid H_1=0,A_1=0)=0.75$ and the remaining probability, $0.25$, is on $H_2=1$. Their earlier-stage Q-value is

$$ Q_1(0,0)=0.75(0.4)+0.25(0.8)=0.30+0.20=0.50. $$

If instead $P(H_2=0\mid H_1=0,A_1=1)=0.25$, then $Q_1(0,1)=0.25(0.4)+0.75(0.8)=0.10+0.60=0.70$. These calculations show the target for the earlier-stage regression: the response at each later history is the optimized value, not the observed outcome under the observed second-stage action.

Worked Example: Preserving an Earlier-Stage Interaction

Suppose a small training data set has four equally sized groups defined by binary $H_1=x$ and $A_1=a$. The mean fitted optimized values $\hat V_2(H_2)$ in the groups are 2 when $(x,a)=(0,0)$, 3 when $(0,1)$, 3 when $(1,0)$, and 7 when $(1,1)$. Fit the saturated regression $\hat V_2=\beta_0+\beta_x x+\beta_a a+\beta_{xa}xa+\varepsilon$.

At $(0,0)$, the fitted mean is $\beta_0=2$. At $(1,0)$, it is $\beta_0+\beta_x=3$, so $\beta_x=1$. At $(0,1)$, it is $\beta_0+\beta_a=3$, so $\beta_a=1$. Finally, at $(1,1)$, it is $\beta_0+\beta_x+\beta_a+\beta_{xa}=7$, giving $\beta_{xa}=7-2-1-1=3$. Therefore

\hat Q_1(x,a)=2+x+a+3xa.

The interaction is needed: an additive model would predict 4 in the $(1,1)$ group, not 7. This example concerns the regression of optimized later-stage values; the coefficients alone do not show that the fitted model is adequate beyond these four training groups.

Worked Example: A Continuous Earlier-Stage Action

Let $A_1=a$ be continuous and suppose $H_2=a+U$, where $U$ equals $-1$ or $1$ with probability $1/2$ each. Suppose the later-stage optimized value is $V_2(h)=8+3h-h^2$. Then

Q_1(a)=E[8+3(a+U)-(a+U)^2\mid A_1=a].

Since $E[U]=0$ and $E[U^2]=1$, expand and take expectations: $E[8+3a+3U-a^2-2aU-U^2]=8+3a-a^2-1=7+3a-a^2$. Thus $Q_1(0)=7$, $Q_1(1)=9$, and $Q_1(2)=9$. The calculation illustrates why predicting the optimized value at the realized later history and then averaging is not generally the same as simply substituting the mean later history into the value function: at $a=0$, for instance, $V_2(E[H_2\mid A_1=0])=V_2(0)=8$, whereas $Q_1(0)=7$.

Diagnostics, Alternatives, and Limitations

The method is useful because it translates a sequential decision problem into a sequence of familiar regressions, while retaining the rule that later actions are chosen to optimize outcome. It is a model-based alternative to inverse-probability weighted approaches to regime evaluation, which use treatment probabilities rather than this stagewise outcome-regression construction. Neither approach removes the need for clear treatment histories and adequate support.

Several checks are especially important:

Check treatment support. At histories where a candidate action has little or no data, its fitted Q-value and the resulting maximum may be extrapolations rather than data-supported estimates.
Check the stage-2 fit before maximizing. Maximization can select an action whose prediction is high because of noise or model error. Compare fitted values with observed outcomes across histories and actions, and inspect whether selected actions occur in supported regions.
Check the stage-1 functional form. The regression must allow relevant dependence on both $H_1$ and $A_1$. For continuous treatments, assess whether the chosen functional form represents the response over the supported action range.
Keep the target straight. Regressing $Y$ directly on earlier variables without encoding later optimization generally describes outcomes under the observed later treatment process, not outcomes under an optimal later rule.

The method also has limitations. Error in $\hat Q_2$ is carried into every calculated $\hat V_2(H_2)$, and the maximizing operation can make small prediction errors consequential when action values are close. Flexible models may reduce approximation bias but can overfit; model selection and inference require additional care. Those issues do not change the definition of $Q_1$, but they do affect how reliably it can be estimated.

Common pitfall: Averaging the observed outcome $Y$ over patients with the same earlier history and action treats their observed $A_2$ choices as the future treatment plan. The correction is to estimate $Q_2$, maximize it over allowed later actions at each $H_2$, and regress those optimized values on $H_1,A_1$. Also, do not maximize over $A_1$ before fitting this regression: $Q_1$ compares current actions, so the current action remains a predictor.

Check Your Understanding

Answer each question, then check the brief solutions.

If $V_2(0)=2$, $V_2(1)=6$, and $P(H_2=1\mid H_1=h,A_1=a)=0.3$, what is $Q_1(h,a)$?
Why does the stage-1 regression include $A_1$ as a predictor rather than maximizing over it before regression?
Suppose $H_2=A_1+U$, $E[U]=0$, $E[U^2]=1$, and $V_2(h)=5-h^2$. Find $Q_1(a)$.
Name two assumptions needed for a causal interpretation of the regression target.
What is the problem with using observed $Y$ as the response in the earlier-stage regression?

Solutions. (1) $Q_1(h,a)=0.7(2)+0.3(6)=1.4+1.8=3.2$. (2) $Q_1$ is defined for each fixed current action; optimizing current actions comes after estimating those conditional values. (3) $Q_1(a)=E[5-(a+U)^2]=5-a^2-2aE[U]-E[U^2]=4-a^2$. (4) Examples are sequential exchangeability and positivity; consistency is also required. (5) It generally averages outcomes under the observed later-treatment choices rather than outcomes under the optimized later-stage rule.