Direct Learning of DTRs, Part ii: Smooth Surrogates and DTRESLO

Direct Learning of Dynamic Treatment Regimes: i. Doubly robust value search · ii. DTRESLO · iii. SDSS.

Part i improves value estimation by combining outcome regressions with inverse-probability weighting. However, its criterion still contains indicators of whether recorded treatments match a candidate rule. These indicators make direct optimization discontinuous.

Here we replace them by smooth functions and ask whether the replacement preserves optimal treatment decisions. We begin with the one-stage construction of Zhao et al. (2012), then derive a normalized version of DTR estimation via surrogate loss optimization (DTRESLO) from Laha et al. (2024), Theorem 3. This construction uses inverse weighting without Part i's doubly robust correction.

Recap: The value identity being smoothed

The earlier inverse-weighting calculation expresses a regime's value through logged trajectories:

Here $H_t$ is pretreatment history, $A_t$ the logged treatment, $R$ terminal utility, and $e_t$ its true treatment probability. The identity requires causal identification and positivity. We retain these conditions and smooth the agreement factors, starting with one treatment.

1. One-stage value as a weighted classification problem

Let $H\in\mathcal H$ denote pretreatment covariates, $A\in\{-1,1\}$ a treatment, and $R$ a reward, with larger values preferred. Write $R(a)$ for the potential reward under treatment $a$ and $V(d)=\mathbb E[R\{d(H)\}]$ for a rule's value. We use $\{-1,1\}$ coding for margin losses; it corresponds to $2A-1$ in the Q-learning note's $\{0,1\}$ coding.

Let $(H,A,R)\sim P$ be the logged-data distribution. Unmarked expectations use $P$; $\mathbb E_H$ uses its marginal $H\sim P_H$. At fixed $H=h$, the conditional calculation averages $A\sim e(\cdot\mid h)$ and then $R\sim P_{R\mid H=h,A}$, while $f$ stays fixed. The potential-outcome mean defining $V(d)$ averages over the baseline population, not over its observed treatment assignments. The assumptions below identify it from $P$.

Assume consistency, no interference, exchangeability $A\perp\!\!\!\perp(R(-1),R(1))\mid H$, and logging propensities

For the weighted losses below, assume $0<r_0\le R\le M<\infty$. A measurable score $f:\mathcal H\to\mathbb R$ defines

Write $\mu_a(h)=\mathbb E[R\mid H=h,A=a]$. Conditioning on $H$ and summing over $A$ cancels the logging probabilities:

The first equality uses causal identification; T is the tower property, and P is probability cancellation under positivity. Define weighted misclassification risk

Since agreement and disagreement partition the two treatments,

which is constant in $f$. Thus maximizing value is equivalent to minimizing weighted classification risk. The Bayes treatment is

2. Why the one-stage OWL hinge loss works

Outcome-weighted learning (OWL) replaces the discontinuous classification loss with the hinge loss

The hinge loss is convex but not smooth at $u=1$. Fix $H=h$, put $u=f(h)$, and abbreviate $p=\mu_1(h)>0$, $q=\mu_{-1}(h)>0$. Probability cancellation gives the conditional risk

Proposition 1: Fisher consistency and the hinge regret bound

The conditional minimizers satisfy

Moreover, for every measurable $f$ with finite hinge risk,

First, locate the minimum from the slopes of $C_h$ on the three intervals:

If $p>q$, the risk decreases until $u=1$ and increases thereafter. If $p<q$, it decreases until $u=-1$ and increases thereafter. If $p=q$, it is constant on $[-1,1]$. In every case,

Next, consider a wrong-sign decision. If $d_f(h)\ne d^*(h)$ and $p\ne q$, the least achievable hinge risk on that wrong side has infimum $p+q$, at $u=0$. Hence

If the treatment is optimal, the value loss is zero and the excess hinge risk is nonnegative. Thus, for all $h$,

Finally, integrate over $H$. The left side becomes value regret, and the measurable pointwise minimizer $\widetilde f(h)=d^*(h)$ attains the conditional infimum, so

This proves both the optimal decision sign and the excess-risk implication. It is the one-stage OWL result developed in Zhao et al., Proposition 3.1 and Theorem 3.2.

3. A smooth surrogate with a useful symmetry

For simultaneous treatment-rule learning, use a smooth increasing function $g:\mathbb R\to(0,1)$ satisfying

Our running choice is

The sign change in $g''$ makes this objective smooth but not globally concave. Our choice normalizes the $\phi(u)=1+\tanh(u)$ example in Laha et al., Condition 2 and Example 1. We use $g$ as a treatment-agreement weight, rather than the logistic log-loss $\log(1+e^{-u})$.

For one stage, the smooth value objective has the conditional form

By monotonicity and the limits of $g$,

At a fixed history with unequal treatment means, scores approaching this conditional supremum eventually select the better treatment. A finite score need not attain the supremum. Maximizing $J_1$ is also equivalent to minimizing the smooth loss $\mathbb E[R\{1-g(Af(H))\}/e(A\mid H)]$.

4. Two-stage smooth value and its policy interpretation

Now observe $O=(X,A_1,L,A_2,R)$ with $H_1=X$ and $H_2=(X,A_1,L)$. Both treatments take values in $\{-1,1\}$. Assume the two-stage consistency, no-interference, and sequential exchangeability conditions from the Q-learning note, with $e_t(a\mid h)\ge\kappa>0$ and $0<r_0\le R\le M$.

Here $P$ denotes the joint logged distribution of this two-stage record; $\mathbb E_X$ uses $X\sim P_X$. In particular, the weighted objective below is $\mathbb E_{O\sim P}$, not an unweighted intervention expectation. Given $(X,A_1)=(x,a)$, continuation means integrate $L\sim P_{L\mid X=x,A_1=a}$; given $(H_2,A_2)=(h,a)$, outcome means integrate $R\sim P_{R\mid H_2=h,A_2=a}$.

For $f=(f_1,f_2)$, $f_t:\mathcal H_t\to\mathbb R$, define

The paper uses $R=Y_1+Y_2$ and a product $\phi(u)\phi(v)$ with $\phi(u)+\phi(-u)=C_\phi$. Here $g=\phi/C_\phi$ and $J=V_\psi/C_\phi^2$. We prove the bound for a positive terminal utility $R$; the source's sum-of-rewards formulation is included when that sum is used as the utility and the history retains the first reward.

Equivalently, minimize the weighted joint surrogate loss

In this display, $f_t=f_t(H_t)$ and $e_t=e_t(A_t\mid H_t)$.

Lemma: The surrogate objective is a stochastic-policy value

Set

Positivity and symmetry give $\pi_{f,t}(a\mid h)>0$ and $\sum_a\pi_{f,t}(a\mid h)=1$. For any stochastic regime $\pi$, define its intervention distribution $P^\pi$ by

Only the treatment probabilities change. Under the causal identification assumptions, the conditional covariate and outcome distributions remain the same as in the observed data. Write $V(\pi)=\mathbb E_{O\sim P^\pi}[R]$. The shorthand $\mathbb E_{\pi_1}$ averages over the first three draws $(X,A_1,L)$; $\mathbb E_{d_1}$ uses the same draws with $A_1=d_1(X)$ fixed by the rule. Conditioning on the second-stage history gives

where $Q_2^*(h,a)=\mathbb E[R\mid H_2=h,A_2=a]$. Another application of the tower property and first-stage probability cancellation yields

Thus the smooth objective evaluates the stochastic regime $\pi_f$. The reported deterministic regime $d_f$ chooses its more probable treatment at each history, so its value still needs to be compared with that smooth objective.

5. Main theorem: Surrogate regret controls DTR regret

Let $V^*=\sup_d V(d)$ over all measurable deterministic regimes, and $J^*=\sup_f J(f)$ over all measurable finite-valued scores. Under the stated assumptions,

This is the normalized form of Laha et al. (2024), Theorem 3. We prove it using nonnegative Bellman gaps. First, define

Each Bellman gap measures the conditional outcome lost by choosing a treatment instead of maximizing its Q-function. The gaps are nonnegative by pointwise maximization, and $V^*=\mathbb E[M_1(X)]$ follows from backward induction.

Next, decompose the regret of any stochastic regime $\pi$ by adding and subtracting $\mathbb E_{\pi_1}[M_2(H_2)]$:

The second equality uses the tower property, the definition of $Q_1^*$, and $\sum_a\pi_t(a\mid h)=1$. Both terms are nonnegative, so $V(\pi)\le V^*$ and, in particular, $J(f)\le V^*$.

To show that the surrogate supremum reaches this upper bound, let $d^*$ be a backward-induction optimal regime and set $f_t^{(k)}(h)=k\,d_t^*(h)$. Then

The integrand defining $J(f^{(k)})$ is bounded in absolute value by $M/\kappa^2$. Dominated convergence (DC) followed by causal identification (I) gives

Together with $J(f)\le V^*$, this proves $J^*=V^*$. It does not assert existence of a finite maximizing score.

We can now compare the stochastic and deterministic regrets. Fix $f$, abbreviate $d=d_f$, $\pi=\pi_f$, and write

Applying the stochastic-regime regret decomposition to deterministic treatment choices gives $V^*-V(d)=D_1+D_2$. Since $g(0)=1/2$ and $g$ is increasing,

For the first stochastic regret term, drop the other treatment's nonnegative contribution:

For the second term, first define the conditional average gap

Retain only histories whose first treatment follows $d_1$, then use the displayed bound on $C_2$:

Each inequality uses nonnegativity of $B_t$ and then the lower bound $\pi_t(d_t\mid H_t)\ge1/2$. The two factors of $1/2$ reflect the first-stage history distribution and the second-stage treatment choice.

Finally, combine these bounds with the stochastic-policy identity and the equality $J^*=V^*$:

Rearranging proves the theorem. $\square$

In the paper's scale, $V_\psi=C_\phi^2J$, so the same bound is

Thus $J(f_n)\to J^*$ implies $V(d_{f_n})\to V^*$. This is surrogate consistency: approaching the best smooth objective value also approaches the best regime value. It is a population result; estimation and numerical optimization errors still need separate bounds.

6. From a sample objective to a learned regime

For i.i.d. records and known logging propensities, write

For feature vectors $z_t(h)\in\mathbb R^{p_t}$ and coefficients $\theta_t\in\mathbb R^{p_t}$, let $f_t(h)=\theta_t^\top z_t(h)$ and $u_{ti}=A_{ti}\theta_t^\top z_t(H_{ti})$. Differentiating the finite sum by the product and chain rules gives

The objective is smooth but generally nonconcave. A stationary point need not maximize it.

Algorithm: Simultaneous binary treatment-rule learning

Input: Trajectories $\{O_i\}_{i=1}^n$; true propensities $e_1,e_2$; features $z_1,z_2$; sigmoid $g$; $B\ge1$ initializations; $K\ge1$ iterations; positive learning rates $\{\eta_r\}_{r=0}^{K-1}$.

Output: Binary treatment rules $\widehat d_1,\widehat d_2$.

1. Compute $w_i\gets R_i/\{e_1(A_{1i}\mid H_{1i})e_2(A_{2i}\mid H_{2i})\}$.
2. for $b=1,\ldots,B$ do
3. Initialize the joint coefficient vector $\theta^{(b,0)}=(\theta_1^{(b,0)},\theta_2^{(b,0)})$.
4. for $r=0,\ldots,K-1$ do
5. Evaluate both stage gradients of $\widehat J_n$ at $\theta^{(b,r)}$.
6. Update jointly: $\theta^{(b,r+1)}\gets\theta^{(b,r)}+\eta_r\nabla_\theta\widehat J_n(\theta^{(b,r)})$.
7. end for
8. end for
9. Select $\widehat\theta$ with the largest $\widehat J_n$ among all visited iterates, including initializations; break ties by visit order.
10. Set $\widehat f_t(h)\gets\widehat\theta_t^\top z_t(h)$ for $t=1,2$.
11. return $\widehat d_t(h)=1$ if $\widehat f_t(h)\ge0$, and $-1$ otherwise.

The algorithm updates both treatment-score functions jointly, without backward pseudo-outcome regression. Multiple initializations provide candidate solutions, but do not certify a global maximum.

Corollary: Approximation, sampling, and optimization errors

Let $\mathcal F_n$ be a deterministic class of score pairs. Assume all following suprema and random quantities are measurable, or impose the usual separability conditions. Define

Add and subtract the empirical supremum and fitted empirical value:

Here $|\sup_{\mathcal F_n}J-\sup_{\mathcal F_n}\widehat J_n|\le\mathsf E_n$ and the last difference is at most $\mathsf E_n$. The main theorem yields

The training sample has distribution $\mathcal T=(O_1,\ldots,O_n)\sim P^{\otimes n}$. Conditional on the fitted scores, $J(\widehat f_n)$ and $V(d_{\widehat f_n})$ average over fresh trajectories, not over the training sample again. Convergence in probability refers to repeated training samples (and any randomness in the fitting procedure). If these three errors vanish in that sense, the learned regime is value-consistent. Smoothness alone proves none of those limits. Estimated propensities introduce additional error not covered by the known-propensity expression above.

7. A four-outcome worked example

Take no baseline variation and no intermediate covariate other than $A_1$. Let the deterministic positive outcome be

Then

For $k>0$, choose $f_1=k$ and $f_2(a_1)=-ka_1$, and put $s=g(k)$. The stochastic value is

The deterministic regime is already optimal for every $k>0$; its stochastic counterpart still sometimes selects a suboptimal treatment. This distinguishes the two values without an estimation argument.

If the first score is instead $-k$, leaving $f_2(a_1)=-ka_1$, then $J(f)=3-g(k)\to2$ and $V(d_f)=2$. Increasing score magnitude cannot repair a wrong treatment sign.

8. What must not be conflated

• Hinge OWL versus smooth DTRESLO: the one-stage hinge proof does not prove consistency of an arbitrary joint two-stage surrogate. The symmetry and bounded range above provide the stochastic-policy identity used in the main proof.
• Unrestricted population scores versus fitted classes: the sequence $k d^*$ may not belong to a chosen linear or neural-network class. That restriction appears in $\mathsf A_n$.
• Outcome shifting versus surrogate estimation: if rewards are bounded below, a fixed shift can make them positive without changing the optimal deterministic regime. No finite shift makes an unbounded-below Gaussian reward always positive. Even when population optimal rules are unchanged, shifts can change finite-class fitted objectives.
• Binary versus multiple treatments: this proof is for two stages and binary actions. No multicategory theorem is proved here.

Part iii replaces the binary sign rule with treatment-specific scores. Its product-based surrogate need not sum to one across treatments, so we must recheck the probability interpretation before extending this proof.

References

• Zhao, Y., Zeng, D., Rush, A. J., and Kosorok, M. R. (2012). Estimating Individualized Treatment Rules Using Outcome Weighted Learning. Journal of the American Statistical Association, 107(499), 1106–1118. Article and proofs, DOI.
• Laha, N., Sonabend-W, A., Mukherjee, R., and Cai, T. (2024). Finding the Optimal Dynamic Treatment Regimes Using Smooth Fisher Consistent Surrogate Loss. The Annals of Statistics, 52(2), 679–707. DOI, open preprint. See Condition 2 and Theorem 3; the preprint and journal versions use different supplement section labels.
• Offline Policy Learning from Logged Data: causal identification, Bellman recursion, and Q-learning.
• Part i: Doubly robust value search · Part iii: SDSS.
• Surrogate Loss Functions and Fisher Consistency in Binary Classification: additional background on surrogate risks.