EconBase
← Back to paper

Inference on Optimal Dynamic Policies via Softmax Approximation

Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.

52,465 characters · 9 sections · 36 citation commands

Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.

Inference on Optimal Dynamic Policies via Softmax Approximation

abstractEstimating optimal dynamic policies from offline data is a fundamental problem in dynamic decision making. In the context of causal inference, the problem is known as estimating the optimal dynamic treatment regime. Even though there exists a plethora of methods for estimation, constructing confidence intervals for the value of the optimal regime and structural parameters associated with it is inherently harder, as it involves non-linear and non-differentiable functionals of unknown quantities that need to be estimated. Prior work resorted to sub-sample approaches that can deteriorate the quality of the estimate. We show that a simple soft-max approximation to the optimal treatment regime, for an appropriately fast growing temperature parameter, can achieve valid inference on the truly optimal regime. We illustrate our result for a two-period optimal dynamic regime, though our approach should directly extend to the finite horizon case. Our work combines techniques from semi-parametric inference and $g$-estimation, together with an appropriate triangular array central limit theorem, as well as a novel analysis of the asymptotic influence and asymptotic bias of softmax approximations.

Introduction

In most real world decision making settings, subjects undergo multiple exposures to some treatment or multiple treatments over time; patients receive multiple therapies, plots of lands are treated with multiple seeds, digital economy users are marketed by multiple campaigns. A typical problem that arises frequently: having access to large amounts of data where many subjects underwent multiple treatments sequentially, based on some naturally occurring and unknown treatment policy, can we identify what would have been the optimal dynamic treatment policy. Moreover, can we estimate with confidence the value of that optimal policy and understand whether we can reap large benefits (with statistical significance) if we intervene and change the status quo. The problem is typically termed as off-policy optimization or off-policy reinforcement learning and has also has garnered attention in the field of precision medicine. However, typical approaches in off-policy reinforcement learning primarily focus on getting fast statistical learning rates (see e.g. shi2022statistically, gunn2022adaptive, song2017semiparametric, kallus2020 for recent work), but not uncertainty quantification and construction of confidence intervals, which have been mostly studied in off-policy evaluation but not optimization (see e.g. pmlr-v139-karampatziakis21a for recent work). Our aim is to make a step in filling this gap and address uncertainty quantification and confidence interval construction for off-policy optimization in reinforcement learning settings.

We consider a setting where we have collected data from multiple realizations of a two-period Markovian dynamic decision making process. We focus on two periods for simplicity of exposition, even though our techniques can easily be extended to multiple periods. The data stem from application of an unknown observational dynamic treatment policy on a sampled set of treated units. We assume we have access to n i.i.d. samples, one for each treated unit, of the form $Z=(S,T_1, X, T_2, Y)$, where $S$ is an initial state, $T_1$ a first period treatment, $X$ a second period state, $T_2$ a second period treatment, and $Y$ a final observed outcome of interest. We add subscripts $i$ to denote the realizations of $i$th data point $Z_i= (S_i,T_{1,i}, X_i, T_{2,i}, Y_i), i\in [n].$ The data is assumed to adhere to the causal graph depicted in Figure (ref). Moreover, the treatments $T_1, T_2\in \mathcal{T}$ are discrete and take $K$ possible values, for a constant $K$. We note that our assumption that the decision process is Markovian is without loss of generality, since the state $X$ is of arbitrary dimension and is allowed to capture all past history (for instance, we can set $X=(T_1, S, S')$ for some raw second period state $S'$). This formalism helps simplify notation.

figure[figure omitted — 1,578 chars of source]

Our goal is to estimate and construct a confidence interval for the best dynamic policy in this setting. In particular, if we intervene and instead of the observed dynamic policy we deploy an alternative dynamic policy $\pi$ (depicted in the single world intervention graph (SWIG) in Figure (ref)), then we will observe counterfactual or potential outcomes $Y^{(\pi)}$ and our goal is to find the policy $\pi^*$ that optimizes the mean counterfactual reward and construct a confidence interval for its value $V^*$ defined as:

align*[align* omitted — 110 chars of source]

The majority of prior work in off-policy evaluation and optimization in reinforcement learning has focused on the estimation of good policies with small regret. However, the focus of our work is in the ability to construct confidence intervals for the value of the optimal policy. Constructing confidence intervals is important in many high-stakes domains where we want to understand from offline data whether some candidate optimal policy will produce a positive improvement with statistical significance. If we are not confident on the magnitude of the improvement that the optimal regime will bring, then most decision makers would go with the status quo policy. Hence, construction of valid confidence intervals is important in high-stakes decision making.

Most prior works that provide confidence interval construction is focused on policy evaluation for a particular fixed policy as opposed to the optimal dynamic policy. Constructing confidence intervals for the optimal dynamic policy is an inherently harder problem as the target quantity tends not to be a smooth function of the distribution. The closest prior work to ours is that of chakraborty2010inference which also addresses inference on structural parameters associated with the dynamic optimal policy, but resorts to sub-sampling techniques that can potentially deteriorate the estimation quality of the point estimate. Instead, we take the route of using smooth approximations to the optimal policy and show that the level of smoothness can be tuned appropriately to obtain correct confidence intervals for the truly optimal (non-smooth) policy. Our smooth approximation can be thought as an analogue to soft-Q-learning schulman2017equivalence,nachum2017bridging,haarnoja2017reinforcement,Haarnoja2018, but in the context of the G-estimation framework proposed by robins2004optimal, which we elaborate in Section (ref).

Related Literature

The problem of learning the optimal treatment rule has been well studied, notably in the field of precision medicine (zhao2012estimating, zhang2012robust, fan2017concordance, song2017semiparametric, gunn2022adaptive). This work has focused on finding the optimal treatments when there is only one time period by optimizing estimators of variants of the value functions or contrasts of value functions of treatment decision rules. For the dynamic treatment regime with multiple treatment stages, there have also been efforts to develop Q-learning (watkins1989learning, watkins1992q, zhao2009reinforcement, zhao2011reinforcement, qian2011performance) and A-learning ((murphy2003optimal, robins2004optimal, moodie2007demystifying, lu2013variable)) type algorithms for the estimation of the optimal dynamic treatment regime. A-learning methods are semi-parametrically efficient and tend to have better performance than Q-learning methods when Q-functions are misspecified (chakraborty2010inference, qian2011performance, shi2018high). For example, shi2022statistically proposes an advantage learning process for offline infinite horizon settings to boost statistical efficiency in the estimation of the optimal treatment regime. They provide a finite sample bias guarantee for their process and show that the estimated contrast of Q-functions from their process converges faster than existing Q-learning estimation methods as the number of stages approaches infinity. However, these works focus on improving statistical efficiency rather than constructing confidence intervals for quantities related to the optimal policy.

In the literature of optimal treatment estimation under dynamic treatment regimes, most estimators are in general not asymptotically normal (robins2004optimal,moodie2010estimating,chakraborty2010inference,song2015penalized). This is particularly due to the fact that the first-stage pseudo-outcomes are not differentiable with respect to the second-stage structural parameters. This phenomenon also appears in other subfields such as in supervised classification when classes are not well separated (e.g laber2011adaptive), where the lack of differentiability of the $0$-$1$ loss can lead to non-asymptotic normality. To address this problem, different solutions have been proposed. A first solution consists of adapting the estimation procedure by using a hard threshold estimator or a soft threshold estimator (e.g see chakraborty2010inference,moodie2012q). These estimators shrink down the value of the second-stage structural parameters when close to the regions of non-regularity. Although it has been empirically shown to perform well (e.g see moodie2012q), this approach suffers from a few drawbacks. Firstly, these threshold estimators often depend on hyperparameters that need to be tuned. Moreover, as noted in chakraborty2010inference, theoretical understanding of these estimators is still lacking. Another solution proposes the construction of confidence intervals not through asymptotic normality but instead by relying on adaptive confidence intervals or adaptive bootstrap procedures (berger1994p, robins2004optimal, laber2011adaptive, laber2014dynamic, huang2015characterizing). For example, laber2014dynamic proposes to build different confidence intervals in, respectively, the regions of regularity and non-regularity. In the latter case, they exploit a union-bound argument over all possible values of the second-stage parameters. However, this approach offers conservative intervals that overcover and is difficult to generalize to multiple stages without having the confidence region size become excessively large. Another proposed solution uses the “m out of n bootstrap" procedure (bretagnolle1983lois, swanepoel1986note, shao1989general, bickel2012resampling). chakraborty2013inference proposes to choose the size $m$ adaptively and derives theoretical guarantees for the obtained estimators. However, note that as is chosen to be much smaller than the sample size, i.e. $m=o(n)$, the proposed estimator will not converge at the optimal rate but at a slower rate. In contrast, our estimation algorithm uses a simple softmax approximation technique and we obtain asymptotic coverage guarantees as long as the log-ratio of the temperature parameter in the softmax operator to the sample size lies in a simple range. Even though an alternative softmax operator has also been proposed (asadi2017alternative) with theoretically better convergence properties that have been employed in other statistical problems, we adhere to the widely adopted Boltzmann softmax operator, demonstrating that it still allows for valid inference for our problem of study.

Semi-parametric approaches are frequently employed for estimating optimal dynamic treatment regimes (e.g robins2004optimal,murphy2003optimal). These estimators are built using the generalized method of moments (levit1976efficiency, hasminskii1979nonparametric, ibragimov1981statistical, pfanzagl1982lecture, klaassen1987consistent, robinson1988root, bickel1988estimating, vaar:1991, bickel1993efficient, newey1994asymptotic, robins1995semiparametric, van2000asymptotic, newey1998undersmoothing, ai2003efficient, newey2004twicing, ai2007estimation, ai2012semiparametric, kosorok2007introduction,tsiatis2007semiparametric). The nuisance parameter is often high-dimensional or infinite-dimensional and in general the difficulty in estimating nuisance parameters can negatively affect the estimation of the target parameter. To solve this problem, the efficient semi-parametric inference literature imposes conditions on the moment function under which errors made in the estimation of the nuisance parameters have a reduced impact on the estimation of the target parameter (hasminskii1979nonparametric,bickel1988estimating,zhang2014confidence,belloni2011inference,belloni2014inference,belloni2014uniform,belloni2014pivotal,javanmard2014confidence,javanmard2014hypothesis,javanmard2018debiasing,van2014asymptotically,ning2017general,chernozhukov2015valid,neykov2018unified,ren2015asymptotic,jankova2015confidence,jankova2016confidence,jankova2018semiparametric,bradic2017uniform,zhu2017breaking,zhu2018linear). This approach dates back to the classical work on doubly robust estimation and targeted maximum likelihood (robins1995analysis, robins1995semiparametric, van2006targeted, van2011targeted, luedtke2016statistical, toth2016tmle) as well as to more recent work on locally robust or Neyman orthogonal conditions on the moment function (chernozhukov2022locally,belloni2017program,chernozhukov2018double). Debiased machine learning uses Neyman orthogonal moments, which satisfy that the first partial derivatives with respect to the nuisance parameters are zero. When the nuisance parameters live in some low complexity class (e.g. Donsker class) or when cross-fitting is used (newey2018cross), the target parameter can then be estimated at parametric rates. Neyman orthogonal moments can be constructed using the concept of a nonparametric influence function (chernozhukov2022locally, newey1998undersmoothing, newey2004twicing, bravo2020two) or the related notion of the Riesz representer (chernozhukov2020adversarial, chernozhukov2022automatic, chernozhukov2022debiased, robins2007comment, newey2018cross, athey2018approximate, hirshberg2018debiased, hirshberg2017augmented, rothenhausler2019incremental, singh2019biased). Building on the work of chernozhukov2018double, lewis2020double showcased the use of the double/debiased machine learning framework to dynamic treatment effect models, specifically Structural Nested Mean Models, by leveraging a sequential residualization approach. In our paper, we extend the debiased machine learning approach to estimate the structural parameters of the optimal dynamic treatment regime. We establish asymptotic normality of the value of the optimal treatment policy and of the structural parameters for both the first and second periods. As is standard in the debiased machine learning literature, our results require that the nuisance function spaces exhibit low complexity (with a small critical radius) or that cross-fitting is employed.

Identification of Optimal Dynamic Policies via G-Estimation

First observe that by invoking the Markovianity of the policy and the conditional independencies implied by the intervention graphs, we can characterize the optimal dynamic policy in a backwards induction manner (see Appendix (ref) for a formal argument):

align[align omitted — 207 chars of source]

In essence, we go to the last period and we optimize our second period action given the second period state $X$, without any change to our first period policy. This provides the optimal second period policy $\pi_2^*$. Then we go to the first period and we fantasize that we are continuing with the second period policy that we have just constructed and under this fantasy, we optimize our first period action given the first period state.

A well-known method for estimating an optimal dynamic regime is $g$-estimation robins2004optimal, which is a form of backwards dynamic programming or $Q$-learning and closely related to advantage or $A$-learning murphy2003optimal,blatt2004learning,Moodie2007,schulte2014q. This process generalizes to arbitrary number of periods, but we describe it here for simplicity in the two period case. The key argument in $g$-estimation is that the improvement that any policy $\pi$ brings, as compared to the observed policy, can be decomposed as the sum of a sequence of improvements, one for each decision-making period. Each of these improvement corresponds to removing the effect of the treatment that was given at that period and adding the effect of the treatment that would have been assigned by policy $\pi$ at that period. For such a decomposition, we need an appropriate notion of effect of each treatment. As a side note, appropriate notions of such an effect are also potentially useful in credit assignment, i.e. attributing parts of the outcome to an assigned treatment; hence estimating these effect functions is interesting in its own right.

More formally, we can write (see Appendix (ref) for details on the derivation): $$ \mathbb{E}\left\{Y^{(\pi)} - Y\right\} = \mathbb{E}\left\{\gamma_1^{(\pi_2)}(\pi_1(S), S) - \gamma_1^{(\pi_2)}(T_1, S)\right\} + \mathbb{E}\left\{\gamma_2(\pi_2(X), X) - \gamma_2(T_2, X)\right\}. $$

The functions $\gamma_t$ are known in the literature as the blip effects and correspond to the following quantity: what is the increase in reward if at the current moment we switch from the baseline treatment to some other treatment level $\tau$, and then continue with the target policy $\pi$. More formally we define:

align*[align* omitted — 237 chars of source]

We will denote with $\gamma_t^*$ the blip effects that correspond to the optimal target policy (note that this only matters for the definition of $\gamma_1$).

If we had some way of estimating the blip effect functions for any target policy, then this decomposition allows us to construct the optimal dynamic regime in a recursive manner. We go to the second period and we note that under the conditional independencies implied by the intervention graphs, one can easily argue that we can equivalently define the optimal policy $\pi_2^*$ from Equation (ref) as the policy that maximizes the blip effect $\gamma_2=\gamma_2^*$, i.e.

align*[align* omitted — 73 chars of source]

We can then use this policy as the continuation policy when we define the blip effect for the first period, which we denote with $\gamma_1^*$ (as it is the blip effect that is associated with the optimal target policy). Again, under the conditional independencies implied by the intervention graphs, we can easily argue that the optimal first period policy $\pi_1^*$, as defined in Equation (ref) is the one that optimizes the first period blip effect, i.e.

align*[align* omitted — 72 chars of source]

Note that the conditional expectations that define the blip effects involve counterfactual quantities and therefore they don't correspond to simple regression problems on the observed data. The ingenuity of the $g$ estimation approach is that it provides a way to estimate the blip functions by noticing that they must satisfy a set of conditional moment restrictions. In particular, one can roughly make the following arguments robins2004optimal,lewis2020double: if we subtract from $Y$ the blip effect of the observed last treatment, then the resulting random variable can roughly be thought as the counterfactual outcome $Y^{(T_1, 0)}$, which should be independent of $T_2$ conditional on $X$. This then implies that for all functions $f$:

align*[align* omitted — 139 chars of source]

This continuum of moment conditions can be used to identify the function $\gamma_2^*$. Similarly, if we subtract from the observed outcome the blip effect of the last period, add the blip effect of the optimal second period treatment and subtract the blip effect of the first period treatment, then the resulting variable can roughly be thought as the counterfactual outcome $Y^{(0, \pi_2^*)}$, which is independent of $T_1$ given $S$. This then implies that for all functions $f$:

align*[align* omitted — 188 chars of source]

This continuum of moment conditions can be used to identify the function $\gamma_1^*$.

This approach becomes very practical if one imposes semi-parametric restrictions on the data generating process that imply that the blip functions admit a parametric form. In particular, we will make a typical assumption that the blip effects that correspond to the optimal target policy are linear in some sufficiently expressive feature map (such a modeling assumption is referred to in the literature as a Linear Structural Nested Mean Model; SNMM hernan2010causal), i.e. for some known $d$-dimensional feature map vectors $\phi, \mu$:

align*[align* omitted — 142 chars of source]

with the convention that $\phi(0, x)=\mu(0,s)=0$. Then we can identify the structural parameters $\psi_0$ by simply solving for $\psi$ the following vector of moment restrictions:

align*[align* omitted — 166 chars of source]

then calculating the optimal second period policy:

align*[align* omitted — 86 chars of source]

Similarly, we can define the residual outcome that emulates $Y^{(0,\pi_2^*)}$ as:

equation[equation omitted — 189 chars of source]

and calculate the first period structural parameters $\theta_0$ by solving for $\theta$ the following vector of moment restrictions:

equation[equation omitted — 159 chars of source]

and calculate the optimal first period policy:

align*[align* omitted — 84 chars of source]

If we denote with:

align[align omitted — 237 chars of source]

then value of the optimal policy is then given as:

align*[align* omitted — 88 chars of source]

Challenges in Inference on Optimal Value and Structural Parameters

We can easily translate the identification recipe in the previous section to an estimation strategy. In particular, we can replace the moment equations with their empirical analogues, estimate regression models for the conditional expectations of the feature maps and construct estimates $\hat{\psi}$ and $\hat{\theta}$ for the structural parameters using $n$ samples. There are two caveats with this approach.

The first is that the estimation error of the regression functions $\mathbb{E}\{\phi(T_2,X)\mid X\}$ and $\mathbb{E}\{\mu(T_1,S)\mid S\}$, will have a first order impact on the accuracy of our structural parameters and hence the quality of the resulting optimal policies. Moreover, if we want to use regularization based methods to flexibly estimate these regressions with modern regression approaches, then the resulting estimates $\hat{\theta}, \hat{\psi}$ will be heavily biased, due to regularization bias. However, we can easily address this issue by resorting to existing techniques in the literature. In particular, instead of considering the raw moments in the previous section, we will consider Neyman orthogonal variants of the moments (see e.g. robins2004optimal,lewis2020double). For instance, for the second period parameter $\psi_0$, we solve for $\psi$ an empirical analogue of the moment equation:

align*[align* omitted — 144 chars of source]

Note that we can write:

align*[align* omitted — 130 chars of source]

Thus, we can estimate this term by estimating the following two regression problems:

equation[equation omitted — 123 chars of source]

This moment is the standard residual-on-residual regression moment appearing in the double/debiased machine learning literature. The extra centering terms we added in the first parenthesis, do not alter the solution, but decrease the variance of the moment. This leads to a more efficient estimation method and more importantly, as we will see formally, robust to small influence from estimation errors in the various regression models that appear in the moment. We can perform similar correction to the first stage moment. Before we do that, let us elaborate on the second caveat of the estimation approach, which is more crucial and the main topic of this work.

One route to constructing asymptotically valid confidence intervals is showing that the structural parameter estimates are asymptotically linear (i.e. asymptotically equivalent to a simple average of some fixed function, known as the influence function). Asymptotic linearity implies asymptotic normality which can be used to provide confidence intervals for the parameter and, coupled with the delta method, for any smooth functional of the structural parameters. However, neither the moment that estimates the first period structural parameter $\theta$ is smooth in $\psi$ nor the policy value is smooth in $\theta, \psi$. Both of these functionals contain the maximum operator, which introduces non-smoothness and leads to non-Gaussian asymptotic stochastic behavior and inability to construct confidence intervals.

The main idea of our work is that by replacing the maximum operator with the soft-maximum operator with an appropriate temperature parameter, growing at the right rate, we can obtain valid inference on the truly optimal policy and its corresponding structural parameters. The key bottleneck of the problem is analyzing the asymptotic linearity of the first period parameter estimator $\hat\theta$, which will be our main theorem. Inference on the structural parameters is also interesting in its own right, for instance when one is interested in understanding dimensions of heterogeneity of the effect of the treatment in the first period, under an optimal continuation policy. Building on this step, in Section (ref) we present asymptotic normality and inference on the policy value, which, in essence, is a repetition of the soft-max analysis we present for inference on the first stage structural parameter.

Softmax Approximation

We will define a softmax approximation to the moment conditions in Equation (ref) that define $\theta_0$. To do so, it will be convenient to introduce a soft-max operator:

align*[align* omitted — 210 chars of source]

We can then define the softmax analogue of the residual $\epsilon_1(\theta, \psi)$ defined in Equation (ref) as:

align*[align* omitted — 225 chars of source]

Then the softmax approximation to the Neyman orthogonal moment condition that defines the parameter $\theta_0$ can be written as:

align*[align* omitted — 172 chars of source]

We will denote with $\theta_0^\beta$ the solution to this softmax moment condition. Note that the term $\mathbb{E}\{\epsilon_1^\beta(\theta, \psi_0)\mid S\}$ can be decomposed as:

align*[align* omitted — 256 chars of source]

Therefore we can estimate this term by estimating the following three regression problems:

equation[equation omitted — 338 chars of source]

It will be convenient to denote with $m_\beta(Z;\theta, \psi, q, p_1, p_2)$ the softmax approximation to the Neyman orthogonal moment, parameterized by the aforementioned regression functions, i.e.

equation[equation omitted — 190 chars of source]

and its expected value as:

align*[align* omitted — 108 chars of source]

For any $\beta\in (0, \infty]$, let $\theta_0^\beta$ be the solution to the following moment condition, with respect to $\theta$: $$M(\theta, \psi_0, q^*, p_1^*, p_{2,\beta}^*;\beta) = 0.$$ Note that we have $\theta_0 = \theta_0^\infty.$

Estimation Procedure and Main Theorem

First, we define our estimation process, which consists of two steps of backwards induction.

\paragraph{Step 1: Procedure for estimate $\hat{\psi}$ of second period structural parameter $\psi_0$.}

tabbing\qquad \enspace Construct estimates $\hat{h}, \hat{r}$ of the nuisance functions $h^*$, $r^*$ defined in Equation (ref).\\ \qquad \enspace For $i=1$ to $i=n$ \\ \qquad \qquad Define residuals $\check{Y}_i=Y_i-\hat{h}(X_i)$ and $\check{\Phi}_i=\phi(T_{2,i},X_i) - \hat{r}(X_i)$ .\\ \qquad \enspace Obtain $\hat{\psi}$ by solving the empirical moment equation with respect to $\psi$:
align*[align* omitted — 105 chars of source]

\paragraph{Step 2: Procedure for estimate $\hat{\theta}^\beta$ of first period structural parameter $\theta_0$}

tabbing\qquad \enspace Construct estimates $\hat{q}, \hat{p}_1$ of the nuisance functions $q^*, p_1^*$ defined in Equation (ref).\\ \qquad \enspace Construct estimate $\hat{p}_{2,\beta} \equiv \hat{p}_{2,\beta, \hat{\psi}}$ by regressing $\hat{\psi} ^{\intercal}\phi(T_2, X) - \operatorname*{\ensuremath{\mathrm{smax}^\beta}}_{\tau_2} \hat{\psi} ^{\intercal}\phi(\tau_2, X)$ on $S$. \\ \qquad \enspace For $i=1$ to $i=n$ \\ \qquad \qquad Define residuals: $\hat{Y}_i=Y_i - \hat{q}(S_i)$ and $\hat{M}_i=\mu(T_{1,i}, S_i) - \hat{p}_1(S_i)$.\\ \qquad \qquad Define residual: $\hat{\Phi}_i = \hat{\psi} ^{\intercal}\phi(T_{2,i}, X_i) - \operatorname*{\ensuremath{\mathrm{smax}^\beta}}_{\tau_2} \hat{\psi} ^{\intercal}\phi(\tau_{2}, X_i) - \hat{p}_{2,\beta}(S_i)$\\ \qquad \enspace Obtain $\hat{\theta}^\beta$ by solving the empirical moment equation with respect to $\theta$:
align*[align* omitted — 109 chars of source]

Having defined our estimation process, we are now ready to prove our main result, which shows that we can construct asymptotically valid confidence intervals for the structural parameters $\theta_0, \psi_0$, assuming the temperature parameter $\beta$ of the soft-max grows fast enough ($\omega(n^{1/\{2(1+\delta)\}})$ for some appropriately chosen $\delta\in(0,1]$), but not too fast ($o(n^{1/2})$). We further need additional assumptions on the estimation quality of all the nuisance functions involved in the estimation process and require that the nuisance estimates lie in function classes of small statistical complexity, e.g. with critical radius of $o_p(n^{-1/4})$, where the classical definition of critical radius of a function class can be found in e.g. Section 3 of chernozhukov2020adversarial.

To present the theorem, we will also need to introduce the following norm notation. For any fixed or random vector $v$, we denote with $\|v\|_{p}=\{\sum |v_i|^p\}^{1/p}$ for any $p\ge 1$, and with $\|v\|_{\infty}=\max |v_i|$, where $v_i$ are the individual components of vector $v.$ Moreover, for any vector-valued function $g$ that takes as input a random variable $X$, we define $\|g\|_2 = [\mathbb{E}_X\{\|g(X)\|_2^2\}]^{1/2}$.

theorem[Main Theorem I] Assume that the random variables $\max_{t} \psi_0 ^{\intercal}\phi(t,X) - \psi_0 ^{\intercal}\phi(\tau, X)\ (\tau\in {\mathcal T})$ are almost surely bounded and each admits a density $f_\tau$ on $(0,c)$ for some constant $c>0$, that satisfies $f_\tau(x)\leq H/x^{1-\delta}$, for some $0\leq H<\infty$ and $0<\delta\leq 1$. Suppose that $\beta=\omega(n^{1/\{2(1+\delta)\}})$ and $\beta=o(n^{1/2})$. Moreover, suppose that the nuisance estimates satisfy the rate conditions: $$\|h^*-\hat{h}\|_{2}, \|r^* - \hat{r}\|_2, \|q^* - \hat q\|_2, \|p_1^* - \hat p_1\|_2=o_p(n^{-1/4})$$ Define $p_{2,\beta,\psi}^*(S)=\mathbb{E}[\psi^{\intercal}\phi(T_2, X)-\operatorname*{\ensuremath{\mathrm{smax}^\beta}}_{\tau_2} \psi^{\intercal}\phi(\tau_2, X)\mid S]$ and $\hat{p}_{2,\beta,\psi}$ the outcome of the estimation algorithm when regressing $\psi^{\intercal}\phi(T_2, X)-\operatorname*{\ensuremath{\mathrm{smax}^\beta}}_{\tau_2} \psi^{\intercal}\phi(\tau_2, X)$ on $S$. Assume that there exists a sufficiently small neighborhood $\mathcal{N}$ of $\psi_0$ such that for all $\epsilon>0$, as $n\to\infty$\footnote{We note that such a guarantee would follow easily, for instance, from results on oracle inequalities via localized complexities and a uniform cover argument over the low dimensional parametric space that the parameter $\psi$ lies in.} \begin{align*} \sup_{\beta>0}pr\left(\sup_{\psi\in\mathcal{N}}\|\hat p_{2,\beta,\psi} - p_{2,\beta,\psi}^*\|_2 \ge \epsilon n^{-1/4}\right)\to 0. \end{align*} Further, suppose that the nuisance estimates take values from a function space of critical radius $\delta_n$ with $\delta_n=o_p(n^{-1/4})$. Assume the nuisance estimates $\hat q(S), \hat p_1(S), \hat p_{2,\beta}(S)$ are all almost surely bounded. Moreover, assume the following boundedness conditions: $$\sup_{\beta>0}\|\theta_0^\beta\|_2<\infty$$ and assume the random variables $$\sum_\tau \|\phi(\tau,X)\|_2, \|\tilde{M}\|_2,\tilde{Y},\sup_{\beta>0}\|p_{2,\beta}^*(S)\|_2, \|\mu(T_1,S)\|_2, \sup_{\tau\in \mathcal{T}, \psi \in \mathcal{N}} |\psi^{\intercal} \phi(\tau, X)|$$ are all almost surely bounded, where $\tilde{M}= \mu(T_1,S) - p_1^*(S)$, $\tilde{Y}= Y - q^*(S)$. Assume that the matrix $\mathbb{E}(\tilde{M} \tilde{M} ^{\intercal})$ is bounded and strictly positive definite so that its inverse matrix $\mathbb{E}(\tilde{M} \tilde{M} ^{\intercal})^{-1}$ exists. Then the estimates $\hat{\psi}, \hat{\theta}^\beta$ are asymptotically linear, i.e. \begin{align*} n^{1/2}(\hat{\psi} - \psi_0) = n^{-1/2} \sum_{i=1}^n \rho_{\psi}(Z_i) + o_p(1),\quad n^{1/2}(\hat{\theta}^\beta - \theta_0) = n^{-1/2} \sum_{i=1}^n \rho_{\theta}(Z_i) + o_p(1) \end{align*} for some functions $\rho_\psi, \rho_\theta$ such that $E\{\rho_\psi(Z_i)\} = 0$ and $E\{\rho_\theta(Z_i)\} = 0$. Moreover, asymptotically valid confidence intervals, for $\psi_0$ and $\theta_0$ respectively, with target coverage level $\alpha$, can be constructed via any consistent estimates $\hat{\sigma}_\psi^2, \hat{\sigma}_\theta^2$ of the variances $\sigma_\psi^2 = \mathbb{E}\{\rho_\psi(Z)^2\}$ and $\sigma_\theta^2 = \mathbb{E}\{\rho_\theta(Z)^2\}$ as: \begin{align*} CI_{\psi}(\alpha) = [\hat{\psi} \pm n^{-1/2}z_{1-\alpha/2} \hat{\sigma}_{\psi}], \quad CI_{\theta}(\alpha) = [\hat{\theta}^\beta \pm n^{-1/2}z_{1-\alpha/2} \hat{\sigma}_{\theta}], \end{align*} where $z_{q}$ is the $q$-th quantile of the standard normal distribution.

The more expansive theorems that we define in subsequent sections also provide an exact form of the asymptotic influence functions and the asymptotic variances, which we omit in the main theorem for succinctness of exposition. Furthermore, we note that a simpler sufficient condition that implies the condition that the random variables $\{\max_{t} \psi_0 ^{\intercal}\phi(t,X) - \psi_0 ^{\intercal}\phi(\tau, X)\}_{\tau\in {\mathcal T}}$ each admits a density $f_\tau$, that satisfies $f_\tau(x)\leq H/x^{1-\delta}$, for some $0\leq H<\infty$ and $0<\delta\leq 1$, is that the variables $\{\psi_0 ^{\intercal}\phi(\tau, X)\}_{\tau\in {\mathcal T}}$ admit a joint density, which is a very benign regularity assumption. This follows from the following lemma:

lemmaLet $(U_t)_{t\in \mathcal{T}}$ be a collection of real-valued random variables and write $U_{\rm{max}}=\max_{t\in \mathcal{T}}U_t$. Suppose that the random variables $(U_t)_{t\in \mathcal{T}}$ admit a joint density that is continuous at zero; then for all $t\in \mathcal{T}$, the probability measure of random variable $U_{\rm{max}}-U_t$ on $(0,\infty)$ is absolutely continuous with respect to the Lebesgue measure on $(0,\infty)$. Hence, in particular, this implies that each $U_{\rm{max}}-U_t$ admits a density that $f_\tau$ that satisfies $f_\tau(x)\leq H/x^{1-\delta}$ for some $0\leq H<\infty$ and $\delta=1.$

Furthermore (as we show in Appendix (ref)), the constraint on the critical radius of the nuisance function spaces can be lifted if one performs nested cross-fitting at each step of the recursion, at the cost of reducing sample size when training each regression.

Asymptotic linearity Theorem for the Optimal Policy Value

In previous sections, we have established the asymptotic linearity of the first-period and second-period structural parameter estimators $\hat\theta,\hat\psi$ and thus constructed confidence intervals for the corresponding parameters. In this section, we will further establish asymptotic linearity for an estimator of the optimal policy value and thus construct a confidence interval for the value.

Recall that the value of the optimal policy is identified as $V^* = \mathbb{E}( Y_{\text{adjusted}}^* )$, where $Y_{\text{adjusted}}^*$ is defined in Equation (ref). It will be helpful to define the softmax counterpart:

align*[align* omitted — 305 chars of source]

Then consider the estimator $\hat{V}$, defined as: $$\frac{1}{n}\sum_{i=1}^n \{Y_i - \hat\psi^{\intercal}\phi(T_{2,i}, X_i) - (\hat\theta^\beta)^{\intercal}\mu(T_{1,i}, S_i) + \operatorname*{\ensuremath{\mathrm{smax}^\beta}}_{\tau_2}\hat\psi^{\intercal}\phi(\tau_2, X_i) + \operatorname*{\ensuremath{\mathrm{smax}^\beta}}_{\tau_1}(\hat\theta^\beta)^{\intercal}\mu(\tau_1,S_i)\}.$$

theoremAssume that the following almost sure finiteness conditions hold: $$\sup_S\sum_{\tau\in{\mathcal T}}\|\mu(\tau,S)\|_2,\ \sup_{\beta>0}|V_{\beta}^*|,\ |Y|<\infty \mbox{ a.s.}$$ Then under the conditions of Theorem (ref), we have that $$n^{1/2}(\hat V - V_\beta^*)=n^{-1/2}\sum_{i=1}^n\rho_{V,*}(Z_i) + o_p(1)$$ with influence function \begin{equation} \begin{aligned} \rho_{V,*}(Z)= Y_{adjusted}^* -V^* & + \mathbb{E}\{\phi_{\infty}(X) - \phi(T_2,X)\}^{\intercal}\rho_\psi(Z)\\ & + \mathbb{E}\{\mu_{\infty}(S)-\mu(T_1,S)\}^{\intercal}\rho_\theta(Z), \end{aligned} \end{equation} where \begin{align*} \rho_\psi(Z)= &\mathbb{E}(\tilde{P}\tilde{P}^{\intercal})^{-1}\{Y-\mathbb{E}(Y\mid X) - \psi_0^{\intercal}\tilde{P}\}\tilde{P},\\ \rho_{\theta}(Z) = & \mathbb{E}(\tilde{M}\tilde{M}^{\intercal})^{-1} \left\{m_*(Z; \theta_0, \psi_0,g_0) + J_*^{\intercal}\rho_{\psi}(Z)\right\},\\ m_*(Z;\theta, \psi, g) = & \{\epsilon_1(\theta, \psi) - q(S) + p_2(S) + \theta^\intercal p_1(S) \}\, \{\mu(T_1, S) - p_1(S)\}, \\ J_* = & \mathbb{E}\left\{\left(\phi_\infty(X) - \phi(T_2, X)\right)\tilde{M}^{\intercal}\right\}. \end{align*} and $\tilde{M} = \mu(T_1, S) - \mathbb{E}\{\mu(T_1, S)\mid S\}$ and $\tilde{P}=\phi(T_2, X) - \mathbb{E}\{\phi(T_2, X)\mid X\}$ and, if we denote $\mathcal{M}(X)=\{\tau: \psi_0^{\intercal}\phi(\tau, X) = \max_{t}\psi_0^{\intercal}\phi(t, X)\}$ and ${\cal L}(S)=\{\tau: \theta_0^{\intercal}\mu(\tau, S) = \max_{t}\theta_0^{\intercal}\mu(t, S)\}$, the set of optimal actions at each state, then \begin{align*} \phi_\infty(X)= &|\mathcal{M}(X)|^{-1}\sum_{\tau\in\mathcal{M}(X)}\phi(\tau, X), & \mu_\infty(S)=|\mathcal{L}(S)|^{-1}\sum_{\tau\in\mathcal{L}(S)}\mu(\tau, S). \end{align*}

Theorem (ref) implies that we can build confidence intervals for the value $V_\beta^*$. Moreover, the bias lemma that follows implies that we can build confidence intervals for the true optimal policy value $V^*$. For each treatment $\tau\in\mathcal{T}$, define the random variables

align*[align* omitted — 122 chars of source]

Taking the maxima of $U_\tau^{\psi_0}$ and $U_\tau^{\theta_0}$ over all treatments $\tau\in\mathcal{T}$, we define

align*[align* omitted — 201 chars of source]

We will also need to make the following assumption:

assumptionAssume that there is a constant $c>0$ such that for each treatment $\tau\in\mathcal{T}$, we have that $$\|U_\tau^{\psi_0}\|_{\infty},\|U_\tau^{\theta_0}\|_{\infty}<\infty$$ and that on $(0,c)$ the random variables $U_{\max}^{\psi_0} - U_{\tau}^{\psi_0}$, ${U}_{\max}^{\theta_0} - {U}_{\tau}^{\theta_0}$ admit densities $f_\tau$ and $\tilde{f}_\tau$ that satisfy $f_\tau(x)\leq H/x^{1-\delta}$ and $\tilde f_\tau(x)\leq \tilde H/x^{1-\tilde\delta}$, for some $0\leq H, \tilde H<\infty$ and $0<\delta,\tilde\delta\leq 1$.
lemma[Softmax Bias Control for Optimal Policy Value] Suppose that Assumption (ref) holds. Assume that $\|\tilde{M}\|_{2}$ is uniformly bounded and that $\mathbb{E}(\tilde{M}\tilde{M}^{\intercal})$ is strictly positive definite. If $\beta=\omega(n^{1/\{2(1+\min\{\delta,\tilde \delta\})\}})$ then we get that as $n\to \infty$: $$n^{1/2}(V_\beta^* - V^*) = o(1).$$

Combining Theorem (ref) and Lemma (ref) gives us the following final result:

corollary[Main Corollary] Under the conditions of Theorem (ref) and Lemma (ref), the estimator $\hat V$ is asymptotically linear around $V^*$: $$n^{1/2}(\hat V - V^*)=n^{-1/2}\sum_{i=1}^n\rho_{V,*}(Z_i) + o_p(1)$$ with influence function $\rho_{V,*}$ as defined in (ref). Letting $z_{q}$ be the $q$-th quantile of the standard normal distribution, asymptotically valid confidence intervals for $V^*$, with target coverage level $\alpha$, can be constructed via any consistent estimate $\hat{\sigma}_V^2$ of the variance $\sigma_V^2 = \mathbb{E}\{\rho_{V,*}(Z)^2\}$ as: \begin{align*} CI_{V}(\alpha) = & [\hat{V} \pm z_{1-\alpha/2} \hat{\sigma}_{V}/\sqrt{n}], \end{align*} where $z_q$ is the $q$-th quantile of the standard normal distribution.

Monte Carlo Experiments

We investigate performance of our inference method on a simple data generating process and demonstrate the efficacy of our method and the importance of choosing the softmax parameter $\beta$, based on our theorem. We consider observational data from the structural equation model:

align*[align* omitted — 241 chars of source]

where $S_1, X_1$ denote the first components of the random vectors $S$ and $X$ respectively. As we show in Appendix (ref), the second period blip effect $\gamma_2$ is of the form:

align*[align* omitted — 139 chars of source]

while the blip effect in the first period, under the optimal continuation policy can be very well approximated as a linear function of the feature map $\mu(T_1, S) = (T_1, S_1 T_1)$. Moreover, $\alpha_1$ controls the magnitude of the blip effect of the first period treatment.

For each $n, \alpha_1, \alpha_2$ and softmax parameter $\beta$ specification, we run $100$ monte-carlo experiments and calculate coverage of the estimated 95% and 90% confidence intervals (and the standard error of the coverage) for the optimal policy value. We used Random Forest regression for estimating all the nuisance functions in the estimation process and no sample splitting.\footnote{Code for replicating this experiment can be found in this \href{https://colab.research.google.com/github/syrgkanislab/optimal_dynamic_regime/blob/main/FinalForPaperSoftmaxInference.ipynb}{Jupyter Notebook} from this \href{https://github.com/syrgkanislab/optimal_dynamic_regime}{Github Repository}.} We report results in Figures (ref) and (ref). We chose pairs of $\alpha_1, \alpha_2$ that identify starkly different estimation regimes, thereby capturing different failure modes of poorly constructed confidence intervals due to irregularity of the optimal policy estimand. We see that $\beta\approx n^{1/4+\epsilon}$ (as indicated by Theorem (ref) with $\delta=1$) produces consistently good coverage above 90% across all specifications, while $\beta$ that is either below $n^{1/4}$ or above $n^{1/2}$, undercovers in at least one specification (highlighted). For $n=10k$ the recommended specification always achieves almost nominal coverage (up to standard error), while the other two specifications fail blatantly in at least one specification.

figure[figure omitted — 1,340 chars of source]
figure[figure omitted — 1,377 chars of source]

Acknowledgement

This work was partially supported by a 2023 Amazon Research Award.