Model Selection For Q-Learning

Why Q-Model Selection Needs More Than One Score

Q-learning estimates conditional mean outcomes backward through the treatment stages. At each stage, a candidate Q-model links the history and treatment to the relevant outcome or pseudo-outcome. The selected model therefore affects both fitted values and the treatment rule obtained by maximizing those values. A model that predicts outcomes well on average need not distinguish treatments well enough to produce a useful rule.

For two stages, write \(H_1\) for the first-stage history, \(A_1\) for the first treatment, \(H_2\) for the later history, \(A_2\) for the later treatment, and \(Y\) for the final outcome. The terminal-stage Q-function is \(Q_2(h_2,a_2)=E[Y\mid H_2=h_2,A_2=a_2]\). Its maximized value \( \max_{a_2}Q_2(H_2,a_2) \) is used as the stage-one pseudo-outcome. The stage-one Q-function models the conditional mean of that pseudo-outcome given \(H_1,A_1\). Thus, a change in the terminal-stage specification can change the response that is modeled at stage one.

For a causal interpretation of a learned rule, the usual conditions include consistency, no interference between individuals, sequential exchangeability of each treatment given the measured history, and positivity: every treatment action allowed by the target regime must have positive probability at histories in the target population. Estimation and validation also require independent trajectories, or an appropriate method that accounts for their dependence. Model selection cannot repair a failure of these assumptions.

Compare Specifications Within Each Stage

Begin by defining a plausible candidate set separately for each stage. For example, the set might include a model with history main effects and a model that also allows a history-by-treatment interaction. Include scientifically motivated nonlinear terms when appropriate, but compare models using the same analysis observations and outcome definition. At the terminal stage, the modeled response is \(Y\); at an earlier stage it is a pseudo-outcome, not the original \(Y\).

When a candidate is a Gaussian linear regression with a common variance, an information criterion can summarize the trade-off between fit and the number of mean-model parameters. Let \(SSE\) be the residual sum of squares, \(n\) the number of observations used for that stage, and \(d\) the number of fitted mean parameters, including the intercept. Omitting terms common to all candidate models, the criteria are

The smaller score is preferred. To derive these expressions, the Gaussian log-likelihood at its maximum over the variance has a model-dependent part \(-\frac{n}{2}\log(SSE/n)\). Therefore \(-2\ell\), after dropping constants common to candidates, is \(n\log(SSE/n)\). Adding \(2d\) gives the AIC comparison and adding \(d\log n\) gives the BIC comparison. The variance parameter is also counted in the full likelihood, but its penalty is the same for every candidate here, so it cancels in the comparison.

Assess the Treatment-Effect Structure

For decision making, the contrast between treatment-specific Q-values is often more directly relevant than overall prediction error. In a binary-treatment model, the blip at history \(h\) is the difference between the fitted Q-value for treatment 1 and that for treatment 0. The estimated rule chooses treatment 1 where that contrast is positive, subject to the allowed action set. A candidate can fit the overall outcome well while imposing a constant blip that misses important variation across histories.

At stage two, inspect whether candidate treatment terms allow plausible effect modification by the available history. Then carry the selected terminal-stage model into the construction of the stage-one pseudo-outcome. Compare stage-one candidates on that resulting response, and remember that stage-one fit is conditional on the terminal-stage choices. A useful diagnostic is to check whether modest, scientifically plausible changes in the terminal-stage specification substantially alter the pseudo-outcomes or the earlier-stage treatment contrasts.

Validate the Induced Rule

If the main goal is choosing a treatment rule, a direct comparison should evaluate the rules rather than only their Q-model fits. A validation sample should be separate from the data used to select and fit the models. For a known treatment-assignment probability \(p(a\mid h)\), an inverse-probability estimate of a rule \(d\)'s value is

Here \(H_i\) denotes the history at the decision point for a single-stage illustration. With sequential treatments, the corresponding weight involves the probability of following the rule at every stage. The basic justification in the single-stage case is precise: if treatment is randomized with known \(p(a\mid h)>0\), then

The first equality uses the observed-data conditional mean, and the second uses randomization and consistency to identify it with the potential outcome mean under the assigned rule. Taking expectation over \(H\) gives \(E[\widehat V(d)]=V(d)\) for an independent validation sample. This argument requires a positive known assignment probability and finite expectations. For observational data, identifying the same value additionally requires consistency and conditional exchangeability, with correctly specified treatment probabilities or another valid identification and estimation strategy. The displayed estimator itself does not make those assumptions true. \(\square\)

Practical Comparisons, Diagnostics, and Limitations

A defensible selection analysis uses the criterion that matches the goal and checks whether conclusions are stable. Stage-specific residual diagnostics can reveal systematic lack of fit, influential observations, or a poor variance model. For treatment decisions, also inspect estimated treatment contrasts across clinically meaningful histories and whether the candidate rules differ in regions with adequate treatment support. If a rule depends on extrapolation into histories with little observed support, apparent gains may be artifacts of the model.

Information criteria are convenient for comparing regression specifications, but their target is likelihood fit with a complexity penalty. Cross-validation can estimate out-of-sample prediction error, but ordinary prediction error is still not regime value. Policy-value validation is more directly aligned with decision quality, but can be noisy, especially with small validation samples or low treatment probabilities. When several candidates are searched and the best validation result is reported, that same validation set no longer provides an untouched assessment; a separate test sample or an appropriately nested validation procedure is needed.

A further limitation is that the selected candidate is only the best among those considered. If all candidates omit an important feature, or if the treatment and history support is inadequate, selection cannot produce a reliable rule. Nor does a lower AIC or BIC prove that a model is correctly specified. These are model comparison tools, not substitutes for causal assumptions, scientific judgment, or uncertainty assessment.