# Reinforcement Learning with Heterogeneous Data: Estimation and Inference

Elynn Y. Chen  Stern School of Business, New York University  Rui Song
Department of Statistics, North Carolina State University  Michael I. Jordan
Department of Computer Science, University of California, Berkeley

(January 18, 2022)

###### Abstract

Reinforcement Learning (RL) has the promise of providing data-driven support
for decision-making in a wide range of problems in healthcare, education,
business, and other domains. Classical RL methods focus on the mean of the
total return and, thus, may provide misleading results in the setting of the
heterogeneous populations that commonly underlie large-scale datasets. We
introduce the $K$-Heterogeneous Markov Decision Process ($K$-Hetero MDP) to
address sequential decision problems with population heterogeneity. We propose
the Auto-Clustered Policy Evaluation (ACPE) for estimating the value of a
given policy, and the Auto-Clustered Policy Iteration (ACPI) for estimating
the optimal policy in a given policy class. Our auto-clustered algorithms can
automatically detect and identify homogeneous sub-populations, while
estimating the $Q$ function and the optimal policy for each sub-population. We
establish convergence rates and construct confidence intervals for the
estimators obtained by the ACPE and ACPI. We present simulations to support
our theoretical findings, and we conduct an empirical study on the standard
MIMIC-III dataset. The latter analysis show evidence of value heterogeneity
and confirms the advantages of our new method.

##  1 Introduction

Many real world problems involve making decisions sequentially based on data
that is influenced by previous decisions. For example, in clinical practice,
physicians make a series of treatment decisions over the course of a patient’s
disease based on his/her baseline and evolving characteristics (Schulte et
al., 2014). In education, human and automated tutors attempt to choose
pedagogical activities that will maximize a student’s learning, informed by
their estimates of the student’s current knowledge (Rafferty et al., 2016;
Reddy et al., 2017). The framework of reinforcement learning (RL), powered by
large-scale datasets, has promise to provide valuable decision-making support
in a wide range of domains, from healthcare, education, and business to
scientific research (Schulte et al., 2014; Mandel, 2017; Zhou et al., 2017).

Classical RL methods estimate the _value function_ , which is defined as the
expectation of accumulated future rewards conditioned on the state, or a
state-action pair, at the time when a decision is made. An _optimal policy_ is
obtained by maximizing the value function at each decision point. Such mean-
value-based RL methods have been the subject of a large literature (See Sutton
and Barto (2018) and references therein for a complete review). However, naive
applications of the mean-value-based RL methods to large-scale datasets may
generate misleading results because real-world datasets are often created via
aggregating many heterogeneous data sources. Each data source may contain a
different sub-population with different associations between state variable,
action and outcomes (such as state transition, immediate rewards, and return).
Consider, for example, the MIMIC-III dataset (Johnson et al., 2016; Komorowski
et al., 2018). In Figure 1 (a), we plot the results of learning a single
linear $Q$ function (i.e., a conditional expectation of future return on
state-action pairs) with the entire sepsis cohort data. The plot shows the
estimation residual. It is obvious that the dataset consists of two sub-
populations with different mean values. The mean-value-based $Q$ learning
procedure fails to capture the heterogeneity in the dataset. In general,
results based on integrating over the entire super-population can be
misleading, sometimes even dangerous in high-risk applications such as health
care, finance and autonomous driving.

(a) Residual distribution of simple linear regression is clearly bi-modal,
with two peaks at around -4 and 4.

(b) The residual distribution for our algorithm is bell-shaped, with one peak
at around zero. It also performs better in terms of prediction error.

Figure 1: Heterogeneity in the MIMIC-III dataset and the improvement achieved
by ACPI.

In this paper, we introduce a framework that we refer to as _$K$
-heterogeneous MDP_ ($K$-hetero MDP) to model dynamic decision making problems
with heterogeneous populations. We propose an algorithm, Auto-Clustered Policy
Iteration (ACPI), that can automatically detect and identify homogeneous sub-
populations from a dataset while learning the $Q$ function and the optimal
policy for each sub-population. We establish statistical guarantees for the
estimators produced by the ACPI. Specifically, we obtain convergence rates and
we construct confidence intervals (CIs) for the estimated model parameters and
value functions. These theoretical results are especially important in high-
risk applications. We apply the proposed method on the standard MIMIC-III
dataset. Figure 1 (b) presents the residuals after applying the ACPI to the
same MIMIC-III dataset. It is clear that the ACPI captures the heterogeneity
in the population. We present additional results in the experiments section
that further confirm the favorable performance of our new method.

The contributions of our work are summarized as follows. Methodologically, we
introduce a formal definition of $K$-hetero MDP and propose the new ACPE and
ACPI algorithms to deal with this heterogeneity automatically. The efficacy of
the proposed model and method is validated through empirical evidences from a
synthetic and a real large-scale dataset. While this paper focuses on batch-
mode estimation, the framework of $K$-hetero MDP is generally applicable and
the ACPE and ACPI can be extended to deal with online estimation.
Theoretically, we establish convergence rates and construct confidence
intervals (CIs) under a bidirectional asymptotic framework that allows either
$N$ or $T$ to go to infinity. Our CIs are valid as long as either the number
of trajectories $N$ or the length of trajectory $T$ diverge to infinity. Our
approach covers a wide variety of real applications and is especially useful
when the number of trajectories in a sub-population is small. We also offer
theoretical guidance to practitioners on the choice of tuning parameters in
ACPI. Technically, in order to study the properties of the estimators by the
ACPE and ACPI, we study the temporal difference error and derive its
convergence in the infinity norm. This result is of independent importance,
especially in studying RL algorithms with penalization (Song et al., 2015).

###  1.1 Related work

This paper is closely related to three different strands of research in
machine learning and statistics, namely distributional RL, dynamic treatment
regime and mixture models.

Distributional RL. In contrast to classical mean-value-based RL, Bellemare et
al. (2017) introduce distributional RL (DRL) that focuses on the full
distribution of the random return $Y_{t}$. Rowland et al. (2019) present an
example of multimodality in value distributions. The main approaches to DRL
include learning discrete categorical distribution (Bellemare et al., 2017;
Rowland et al., 2018; Barth-Maron et al., 2018), learning distribution
quantiles (Dabney et al., 2018; Yang et al., 2019), and learning distribution
expectiles (Rowland et al., 2019). This work differs from ours in that it
mainly focuses on computational issues and empirical performance. No
statistical properties of the estimators are studied, partly because of the
fundamental difficulty of estimating a full distribution.

A useful concept introduced in the DRL literature is the _Bellman closedness_
of a statistic. Lemma 4.2 and Theorem 4.3 in Rowland et al. (2019) show that
collections of moments are effectively the only finite set of statistics that
are Bellman closed. This inspires us to estimate the sub-population means
instead of a full distribution of the super-population: the means are Bellman
closed but the discretized distribution or the quantile is not. In the middle
of the spectrum from the mean-value-based RL to the distributional RL, our
method is shown to capture heterogeneity while avoiding the sample
inefficiency inherent in estimating a full distribution.

Dynamic treatment regime. Researchers in dynamic treatment regime (DTR) have
used the RL framework and MDP to derive a set of sequential treatment decision
rules, or treatment regimes, to optimize patients’ clinical outcomes over a
fixed period of time. See Murphy (2003); Robins (2004); Schulte et al. (2014);
Luedtke and Van Der Laan (2016); Ertefaie and Strawderman (2018); Zhu et al.
(2019); Shi et al. (2020) and references therein. Most methods in DTR are
designed for finite horizon $T$ and are implemented through a backward
recursive fitting procedure that is related to dynamic programming (Schulte et
al., 2014). For infinite horizon $T$, Ertefaie and Strawderman (2018) propose
a variant of greedy gradient $Q$ learning (GGQ) to estimate optimal dynamic
treatment regimes. This requires modeling a non-smooth function of the data,
which creates complications. Luckett et al. (2019) propose a $V$-learning
algorithm that models the state value function and directly searches for the
optimal policy among a pre-specified class of policies. They show that the
estimators of $V$-learning converge at the rate of the order $N^{-1/2}$, where
$N$ is the number of decision trajectories in a dataset. Although their
setting allows for infinite horizon, the information accumulated over horizon
$T$ does not enter explicitly in their results. In contrast, we show that the
convergence rate of our estimators are of the order $\left(NT\right)^{-1/2}$
under similar assumptions. Shi et al. (2020) obtain a convergence rate of the
order $\left(NT\right)^{-1/2}$ and derive confidence intervals for policy
evaluation. But they do not consider finding the optimal policy or data
heterogeneity.

While the personalized treatment regime in DTR aims to offer individualized
treatment to each patient $i$ according to different state variable $X_{i,t}$,
the model coefficients are the same for all the patients. In contrast, the
heterogeneity in this paper stems from different model coefficients. Our
problem is more challenging since we need to estimate the model coefficients
and, at the same time, to group them into clusters.

Mixture models. In a linear regression setting, researchers have studied the
problem of supervised clustering that explores homogeneous effects of
covariates along coordinates and across samples. The mixture-model-based
approach needs to specify an underlying distribution for data, and also
requires specification of the number of mixture components in the population.
An alternative approach to model-based clustering employs grouping
penalization on pairwise differences of the data points. For example, assuming
parameter homogeneity over individuals, Tibshirani et al. (2005), Bondell and
Reich (2008), Shen and Huang (2010) and Ke et al. (2013) employ different
penalty functions for each pair of coordinates of the coefficient vector to
group similar-effect covariates along coordinates. Assuming parameter
heterogeneity across individuals, Hocking et al. (2011), Pan et al. (2013),
and Ma and Huang (2017) adopt a fusion-type penalty with either an
$L_{p}$-shrinkage or a non-convex penalty function to formulate clusters
across samples with the same regression coefficients.

This model-based approach is difficult to deploy in the RL setting because the
distribution of the value function is generally not likely to be characterized
by a Gaussian or mixture of Gaussian distribution (Sobel, 1982). Our ACPE and
ACPI algorithms therefore are developed based on pairwise-fusion-type
penalties. However, our problem in RL is more challenging than that in linear
regression. First, while it is common to assume the noise in regression is
i.i.d. sub-Gaussian, the observations over $T$ in RL are not uncorrelated.
This raises significant challenges for our analysis. Second, we consider the
semi-parametric approximation with a diverging number of basis functions
instead of the linear models. Lastly, our ultimate goal is to learn an optimal
policy, while the mixture regression framework focuses entirely on parameter
estimation.

###  1.2 Notation and organization

We use $x$, $\boldsymbol{x}$, $\boldsymbol{X}$ to represent scalars, vectors
and matrices, respectively. Capital letter $X$ represents a random variable.
For any matrix $\boldsymbol{X}$, we use $\boldsymbol{x}_{i\cdot}$,
$\boldsymbol{x}_{\cdot j}$, and $x_{ij}$ to refer to its $i$-th row, $j$-th
column, and $ij$-th entry, respectively. All vectors are column vectors and
row vectors are written as $\boldsymbol{x}^{\top}$ for any vector
$\boldsymbol{x}$. We denote the matrix $\ell_{2}$-norm as
$\left\lVert\boldsymbol{X}\right\rVert_{2}$ and max norm as
$\left\lVert\boldsymbol{X}\right\rVert_{\max}\triangleq\underset{i,j}{\max}\left\lvert
x_{ij}\right\rvert$. When $\boldsymbol{X}$ is a square matrix, we denote by
$\operatorname{Tr}\left(\boldsymbol{X}\right)$,
$\lambda_{max}\left(\boldsymbol{X}\right)$, and
$\lambda_{min}\left(\boldsymbol{X}\right)$ the trace, maximum and minimum
eigenvalues of $\boldsymbol{X}$, respectively. We denote
${\textsc{mat}}\left(\boldsymbol{x}\right)$ as the matricization of a vector
$\boldsymbol{x}$ and ${\textsc{vec}}\left(\boldsymbol{X}\right)$ as the
vectorization of a matrix $\boldsymbol{X}$. We let $C,c,C_{0},c_{0},\ldots$
denote generic constants whose actual values may vary.

The rest of this paper is organized as follows. Section 2 introduces the
problem setting and a few important concepts. Section 3 proposes the ACPI.
Section 5 describes the computation procedure. Section 4 establishes the
statistical properties of the estimators. Section 6 and Section 7 present
empirical results with synthetic and real datasets. Section 8 concludes. All
proofs and technical lemmas are relegated to the supplementary material.

##  2 Statistical Framework

###  2.1 Markov decision process

The problem of dynamic decision making can be formalized mathematically as a
_Markov Decision Process_ (MDP), ${\cal M}=\left\\{{\cal X},{\cal
A},\Pr,r,\gamma\right\\}$. Let $t\in[T]$ index the times in the decision
process at which an action is taken, and let $i\in[N]$ index the set of
individuals interacting with the environment. The $i$-th individual may
correspond to the $i$-th round of game, the $i$-th student, and the $i$-th
patient in the setting of game playing, education, and health care,
respectively. At each time point $t$, the current status of the $i$-th
individual, $i\in[N]$, is characterized by covariates $X_{i,t}\in{\cal X}$ and
a finite set of all possible actions ${\cal A}$ is available. An action
$A_{i,t+1}\in{\cal A}$ is selected and the state of the $i$-th individual
changes to $X_{i,t+1}$ according to the transition probability
$\Pr(X_{i,t+1}\mid X_{i,t},A_{i,t})$. An immediate random scalar reward
$R_{i,t}$ is observed. Over time, these comprise a sample trajectory for the
$i$-th individual and we have in total $N$ sample trajectories,
$\left\\{\left(X_{i,t},A_{i,t},R_{i,t}\right)\right\\}_{0\leq t\leq T}$ for
$i\in[N]$ .

We assume that the data-generating model is a Time-Homogeneous Markov process
(TH Markov) and we also assume that the observed rewards satisfy a Conditional
Mean Independence (CMI) assumption.

######  Assumption 1 (Data generating MDP.).

The sample trajectories
$\left\\{\left\\{\left(X_{i,t},A_{i,t},R_{i,t}\right)\right\\}_{t\geq
0}\right\\}_{i\in[N]}$ are generated from (possibly $N$ different) Markov
decision processes satisfying

  1. (a)

(Markovianity) For any $i\in[N]$ and $t\in[T]$, it holds that $X_{i,t}$ is
independent of $\left\\{\left(X_{i,s},A_{i,s},Y_{i,s}\right)\right\\}$, for
any $0\leq s\leq t-1$, conditional on $\left(X_{i,t-1},A_{i,t-1}\right)$.

  2. (b)

(Time-Homogeneity) The conditional density $\Pr(X_{i,t+1}\mid
X_{i,t},A_{i,t})$ is the same over $t$ for any $i\in[N]$.

  3. (c)

(Conditional Mean Independence (CMI)) For any $i\in[N]$ and $t\in[T]$, it
holds that

${\mathbb{E}}\left[R_{i,t}|X_{i,t}=\boldsymbol{x},A_{i,t}=a,\left\\{X_{i,s},A_{i,s},R_{i,s}\right\\}_{0\leq
s<t}\right]={\mathbb{E}}\left[R_{i,t}|X_{i,t}=\boldsymbol{x},A_{i,t}=a\right]=r_{i}\left(\boldsymbol{x},a\right),$

for some bounded reward function $r_{i}\left(\boldsymbol{x},a\right)$.

######  Remark 1.

The assumptions of the TH Markov and CMI are common in the literature.
However, Assumption 1 (c) is different because
$r_{i}\left(\boldsymbol{x},a\right)$ can be different across $N$ trajectories.

######  Remark 2.

Assumption 1 implies that
$R_{i,t}=r_{i}\left(X_{i,t},A_{i,t}\right)+\eta_{i,t}$ where
${\mathbb{E}}\left[\eta_{i,t}\mid X_{i,t},A_{i,t}\right]=0$. The $\ell_{2}$
convergence and the asymptotic normality of the oracle estimators do not need
any additional assumption on the noise term $\eta_{i,t}$. However, the uniform
convergence that is needed for the feasible estimators requires $\eta_{i,t}$
to be a sub-Gaussian random variable.

A policy
$\pi\left(\boldsymbol{x}\right)=\left[\pi(1\mid\boldsymbol{x}),\;\cdots\;,\pi(M\mid\boldsymbol{x})\right]$
is a function that maps the state space ${\cal X}$ to probability mass
functions on the action space ${\cal A}:=\\{1,\cdots,M\\}$. Under a policy
$\pi$, an agent chooses action $A_{i,t}$ at state $X_{i,t}$ with probability
$\pi\left(A_{i,t}\mid X_{i,t}\right)$ and receives an immediate reward
$R_{i,t}$. The overall return at step $t$ is the accumulated discounted reward
defined as

$Y_{i,t}=\sum_{s=t}^{T}\gamma^{t-s}R_{i,s},$ (1)

where the discount rate $\gamma\in[0,1)$ reflects a trade-off between
immediate and future rewards. Classic RL measures the goodness of a policy
$\pi$ by value functions defined as follows. The state-value function is the
conditional expectation of the return conditioned on the starting state
$\boldsymbol{x}$:

$V\left(\pi,\boldsymbol{x}\right)={\mathbb{E}}\left[\sum_{t=0}^{T}\gamma^{t}R_{i,t}\;\bigg{|}\;X_{i,0}=\boldsymbol{x},\;A_{i,t}\sim\pi(\cdot|X_{i,t}),\;X_{i,t+1}\sim\Pr(\cdot|X_{i,t},A_{i,t})\right],$
(2)

where the expectation is taken under the trajectory distribution generated by
policy $\pi$. The action-value function or $Q$ function is the conditional
expectation of the return conditioned on the starting state-action pair
$(\boldsymbol{x},a)$:

$Q\left(\pi,\boldsymbol{x},a\right)={\mathbb{E}}\left[\sum_{t=0}^{T}\gamma^{t}R_{i,t}\;\bigg{|}\;X_{i,0}=\boldsymbol{x},\;A_{i,0}=a,\;X_{i,t+1}\sim\Pr(\cdot|X_{i,t},A_{i,t}),\;A_{i,t+1}\sim\pi(\cdot|X_{i,t+1})\right],$
(3)

where the expectation is taken by assuming that the dynamic system follows the
given policy $\pi$ afterwards. Let ${\cal R}$ be a reference distribution of
the covariate space ${\cal X}$. The integrated state-value function of a given
policy $\pi$ is defined as

$V_{\cal R}(\pi)=\int_{\boldsymbol{x}\in{\cal
X}}V\left(\pi,\boldsymbol{x}\right){\cal
R}\left(d\boldsymbol{x}\right)=\int\sum_{a\in{\cal
A}}Q\left(\pi,\boldsymbol{x},a\right)\pi\left(a\mid\boldsymbol{x}\right){\cal
R}\left(d\boldsymbol{x}\right).$ (4)

The reference distribution can be viewed as a distribution of initial states
or stationary states and can be estimated from the dataset.

For the purpose of finding the optimal policy, we focus on a finite action set
${\cal A}=\left\\{1,\cdots,M\right\\}$ and a parametric class of policies
$\Pi$ defined by

$\displaystyle\pi\left(\boldsymbol{x};\boldsymbol{\alpha}\right):=\left[\pi\left(1,\boldsymbol{x};\boldsymbol{\alpha}\right),\cdots,\pi\left(M,\boldsymbol{x};\boldsymbol{\alpha}\right)\right],\quad\text{where}$
(5)
$\displaystyle\pi\left(j,\boldsymbol{x};\boldsymbol{\alpha}\right)=\begin{cases}\frac{\exp(\boldsymbol{x}^{\top}\boldsymbol{\alpha}_{j})}{1+\sum_{j=1}^{M-1}\exp(\boldsymbol{x}^{\top}\boldsymbol{\alpha}_{j})},\quad\text{for}\quad
k=1,\cdots,M-1,\\\
\frac{1}{1+\sum_{j=1}^{M-1}\exp(\boldsymbol{x}^{\top}\boldsymbol{\alpha}_{M})},\quad\text{otherwise,}\end{cases}$

where we use $\boldsymbol{\alpha}_{j}$ to denote a vector of parameters for
the $j$-th action and
$\boldsymbol{\alpha}=\left[\boldsymbol{\alpha}_{1}^{\top},\cdots,\boldsymbol{\alpha}_{M}^{\top}\right]^{\top}$.
The optimal policy $\pi^{*}_{\cal R}\in\Pi$ satisfies $V_{{\cal
R}}\left(\pi_{\cal R}^{*}\right)\geq V_{{\cal R}}\left(\pi\right)$ for all
$\pi\in\Pi$. That is,

$\pi_{{\cal R}}^{*}=\underset{\pi\in\Pi}{\arg\max}\;V_{{\cal
R}}\left(\pi\right).$ (6)

The goal is to estimate $\pi_{{\cal R}}^{*}$ using data collected from $N$
trajectories, each of which is assumed to have length $T_{i}=T$ for simplicity
of presentation, and where $T$ may be finite or diverge to infinity.
Throughout, we assume that this reference distribution is fixed and the same
for each homogeneous sub-population.

###  2.2 Heterogeneous MDP and Bellman consistent equation

Since the learning objective is optimizing the total return, we define the
heterogeneity in terms of the value functions—the expectation of return (1) in
homogeneous sub-populations. Different from the classic RL, we allow different
values, and thus different optimal policies, across a heterogeneous super-
population.

######  Definition 1 ($K$-Hetero MDP).

A $K$-heterogeneous MDP is defined as a dynamic system with a latent variable
$\omega\in[K]$ such that

  1. (i)

For $\forall k\in[K]$, $\Pr\left(\omega=k\right)=P^{(k)}$ and it satisfies
that $P^{(k)}\in(0,1)$ and $\sum_{k=1}^{K}P^{(k)}=1$.

  2. (ii)

The value functions for each sub-population are conditioned additionally on
$w_{i}=k$:

$\displaystyle V^{(k)}\left(\pi,\boldsymbol{x}\right)$
$\displaystyle:={\mathbb{E}}\left[\sum_{t=0}^{T}\gamma^{t}R_{i,t}\;\bigg{|}\;X_{i,0}=\boldsymbol{x},\;A_{i,t}\sim\pi(\cdot|X_{i,t}),\;X_{i,t+1}\sim\Pr(\cdot|X_{i,t},A_{i,t}),w_{i}=k\right],$
$\displaystyle Q^{(k)}\left(\pi,\boldsymbol{x},a\right)$
$\displaystyle:=\mathbb{E}\left[\sum_{t=0}^{T}\gamma^{t}R_{i,t}\;\bigg{|}\;X_{i,0}=\boldsymbol{x},\;A_{i,0}=a,\;X_{i,t+1}\sim\Pr(\cdot|X_{i,t},A_{i,t}),\;\right.$
$\displaystyle\hskip
80.00012pt\bigg{.}A_{i,t+1}\sim\pi(\cdot|X_{i,t+1}),w_{i}=k\bigg{]}.$

  3. (iii)

For any $k\neq k^{\prime}$, the value functions of a given policy $\pi$ are
different, that is, $V^{(k)}\left(\pi,\boldsymbol{x}\right)\neq
V^{(k^{\prime})}\left(\pi,\boldsymbol{x}\right)$ and
$Q^{(k)}\left(\pi,\boldsymbol{x},a\right)\neq
Q^{(k^{\prime})}\left(\pi,\boldsymbol{x},a\right)$ for any $\pi\in\Pi$,
$\boldsymbol{x}\in{\cal X}$, and $a\in{\cal A}$.

######  Remark 3.

Heterogeneity may also be defined in terms of MDP tuple ${\cal M}^{(k)}$, or
in terms of trajectory distribution ${\cal H}^{(k)}(\pi)$. However, Definition
1 provides a clearer framework to work with for the purpose of estimating an
optimal policy that maximizes the expected return. Since MDPs with different
${\cal M}^{(k)}$ or ${\cal H}^{(k)}(\pi)$ may have the same value functions,
heterogeneity defined by ${\cal M}^{(k)}$ or ${\cal H}^{(k)}(\pi)$
unnecessarily complicates the problem.

Now we allow for heterogeneity over $i\in[N]$ on trajectories
$\left\\{X_{i,t},A_{i,t},R_{i,t}\right\\}$, $t\in[T]$, $i\in[N]$. By the
$K$-Hetero MDP assumption in Definition 1, the latent variable $w_{i}$
partitions $N$ trajectories into $K$ groups: ${\cal G}=\left({\cal
G}^{(1)},\cdots,{\cal G}^{(K)}\right)$. Without knowing the true groups, it is
safe for now to use different action-value function
$Q^{i}\left(\pi,\boldsymbol{x},a\right)$ for different trajectory $i\in[N]$.
We have
$Q^{i}\left(\pi,\boldsymbol{x},a\right)=Q^{(k)}\left(\pi,\boldsymbol{x},a\right)$
for all $i\in{\cal G}^{(k)}$. Our goal is to estimate $K$ and identify sub-
populations that respond differently to a policy $\pi$. The following Bellman
consistency equation is a first-moment condition that we will use to construct
an auto-clustered estimation of the action-value function.

######  Lemma 1.

Under the assumptions of TH Markov and CMI, we have the first-moment condition
of the action-value function $Q^{i}$ for any $i\in[N]$,

$\displaystyle{\mathbb{E}}\left[R_{i,t}+\gamma\sum_{a^{\prime}\in{\cal
A}}Q^{i}\left(\pi,X_{i,t+1},a^{\prime}\right)\pi\left(a^{\prime}|X_{i,t+1}\right)-Q^{i}\left(\pi,X_{i,t},A_{i,t}\right)\bigg{|}X_{i,t},A_{i,t}\right]=0,$
(7)
$\displaystyle{\mathbb{E}}\left[\left(R_{i,t}+\gamma\sum_{a^{\prime}\in{\cal
A}}Q^{i}\left(\pi,X_{i,t+1},a^{\prime}\right)\pi(a^{\prime}|X_{i,t+1})-Q^{i}\left(\pi,X_{i,t},A_{i,t}\right)\right)\psi\left(X_{i,t},A_{i,t}\right)\right]=0,$
(8)

where $\psi:{\cal X}\times{\cal A}\mapsto\mathbb{R}$ is any real function
supported on ${\cal X}\times{\cal A}$.

##  3 Estimation

The Auto-Clustered Policy Evaluation (ACPE) and the Auto-Clustered Policy
Iteration (ACPI) are summarized in Algorithm 1 and 2, respectively. Algorithm
1 deals with value evaluation of a given policy on the $K$-Hetero MDP.
Algorithm 2 obtains optimal policies for different groups by using Algorithm 1
and group-wise policy improvement.

###  3.1 Auto-Clustered Policy Evaluation

Let $Q\left(\boldsymbol{x},a;\boldsymbol{\beta}^{\pi,i}\right)$ denote a model
for $Q^{i}(\pi,\boldsymbol{x},a)$ indexed by
$\boldsymbol{\beta}^{\pi,i}=\left[(\boldsymbol{\beta}_{1}^{\pi,i})^{\top},\cdots,(\boldsymbol{\beta}_{M}^{\pi,i})^{\top}\right]^{\top}$,
$i\in[N]$, where $\boldsymbol{\beta}_{m}^{\pi,i}\in\mathbb{R}^{J}$ is the
coefficient for the $m$-th action in the action set ${\cal A}$. We assume that
the map $\boldsymbol{\beta}^{\pi,i}\mapsto
Q\left(\boldsymbol{x},a;\boldsymbol{\beta}^{\pi,i}\right)$ is differentiable
everywhere for each fixed $\boldsymbol{x}$, $a$ and $\pi$. We stack all
individual coefficients in a long vector
$\boldsymbol{\beta}^{\pi}=\left[(\boldsymbol{\beta}^{\pi,1})^{\top}\cdots(\boldsymbol{\beta}^{\pi,N})^{\top}\right]^{\top}$
of length $JMN$ and define a trajectory indicator
$\boldsymbol{\Lambda}^{i}=[\boldsymbol{0}\cdots\boldsymbol{0}\;\boldsymbol{I}\;\boldsymbol{0}\cdots\boldsymbol{0}]$
such that
$\boldsymbol{\beta}^{\pi,i}=\boldsymbol{\Lambda}^{i}\boldsymbol{\beta}^{\pi}$.
Let
$\nabla\,Q\left(\boldsymbol{x},a;\boldsymbol{\Lambda}^{i}\boldsymbol{\beta}^{\pi}\right)$
denote the gradient of
$Q\left(\boldsymbol{x},a;\boldsymbol{\Lambda}^{i}\boldsymbol{\beta}^{\pi}\right)$
with respect to $\boldsymbol{\beta}^{\pi}$. With observed trajectories
$\left\\{X_{i,t},A_{i,t},R_{i,t}\right\\}$, we define a sample-version
quantity:

$\displaystyle\boldsymbol{G}\left(\pi,\boldsymbol{\beta}^{\pi}\right)$
$\displaystyle=\frac{1}{NTJ}\sum_{i=1}^{N}\sum_{t=1}^{T}\bigg{(}R_{i,t}+\gamma\sum_{a\in{\cal
A}}Q\left(X_{i,t},a;\boldsymbol{\Lambda}^{i}\boldsymbol{\beta}^{\pi}\right)\pi\left(a|X_{i,t+1}\right)$
$\displaystyle\hskip
40.00006pt-Q\left(X_{i,t},A_{i,t};\boldsymbol{\Lambda}^{i}\boldsymbol{\beta}^{\pi}\right)\bigg{)}\cdot\nabla\,Q\left(X_{i,t},A_{i,t};\boldsymbol{\Lambda}^{i}\boldsymbol{\beta}^{\pi}\right).$
(9)

By Lemma 1, under certain mild conditions, there exist some
$\mathring{\boldsymbol{\beta}}^{\pi}$ that satisfy
${\mathbb{E}}\left[\boldsymbol{G}^{\pi}\left(\mathring{\boldsymbol{\beta}}^{\pi}\right)\right]=\boldsymbol{0}$.
Thus, an estimator of $\mathring{\boldsymbol{\beta}}^{\pi}$ can be obtained by
minimizing
$\boldsymbol{G}\left(\pi,\boldsymbol{\beta}^{\pi}\right)^{\top}\boldsymbol{G}^{\pi}\left(\pi,\boldsymbol{\beta}^{\pi}\right)$.
At the same time, the true model coefficients should satisfy
$\mathring{\boldsymbol{\beta}}^{\pi,i}=\mathring{\boldsymbol{\beta}}^{\pi,j}$
if $i,j\in{\cal G}^{(k)}$, for any $k\in[K]$, under the $K$-hetero MDP
assumption. When the grouping ${\cal G}^{(k)}$ is unknown, it is crucial and
beneficial to encourage grouping individuals with the same model coefficients
together to achieve efficient estimation. Given a penalty function
$p:\mathbb{R}\mapsto\mathbb{R}$, the ACPE algorithm estimates
$\mathring{\boldsymbol{\beta}}^{\pi}$ for a given policy $\pi$ as:

$\widehat{}\boldsymbol{\beta}^{\pi}=\underset{\boldsymbol{\beta}=[(\boldsymbol{\beta}^{1})^{\top}\cdots(\boldsymbol{\beta}^{N})^{\top}]^{\top}}{\arg\min}\;\left[\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)^{\top}\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)+\frac{1}{N^{2}}\sum_{1\leq
i<j\leq
N}p\left((JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right\rVert_{2},\lambda\right)\right].$
(10)

The choices of the approximation model for $Q$ function and the penalty
function $p$ will be presented in the sequel. The fusion-type penalty shrinks
some of the pairs
$\widehat{}\boldsymbol{\beta}^{\pi,i}-\widehat{}\boldsymbol{\beta}^{\pi,j}$ to
zero, based on which we can partition the sample into subgroups simultaneously
when we estimate the coefficients. The choice of tuning parameter $\lambda$ is
given in Theorem 2.

The ACPE requires a class of models for the $Q$ function indexed by parameter
$\boldsymbol{\beta}^{\pi}$. We use a basis function approximation. Let
$\boldsymbol{\phi}(\cdot)=\left[\phi_{1}(\cdot),\cdots,\phi_{J}(\cdot)\right]^{\top}$
be a vector of pre-specified basis functions such as Gaussian basis function
or splines. We allow $J$ to grow with the sample size to reduce the bias of
the resulting estimates. Given a policy $\pi$, the $Q$ function becomes

$Q\left(X_{i,t},A_{i,t};\boldsymbol{\Lambda}^{i}\boldsymbol{\beta}^{\pi}\right)\approx
Z_{i,t}^{\top}\boldsymbol{\Lambda}^{i}\boldsymbol{\beta}^{\pi},$ (11)

where $Z_{i,t}:=\boldsymbol{z}\left(X_{i,t},A_{i,t}\right)$ and

$\boldsymbol{z}\left(\boldsymbol{x},a\right):=\left[\boldsymbol{\phi}(\boldsymbol{x})^{\top}\mathbbm{1}(a=1),\cdots,\boldsymbol{\phi}(\boldsymbol{x})^{\top}\mathbbm{1}(a=M)\right]^{\top}.$
(12)

Plugging (11) in (3.1), we have

$\boldsymbol{G}\left(\pi,\boldsymbol{\beta}^{\pi}\right)=\frac{1}{NTJ}\sum_{i=1}^{N}\sum_{t=1}^{T}(\boldsymbol{\Lambda}^{i})^{\top}Z_{i,t}\left(R_{i,t}-\left(Z_{i,t}-\gamma
U^{\pi}_{i,t+1}\right)^{\top}\boldsymbol{\Lambda}^{i}\boldsymbol{\beta}^{\pi}\right),$
(13)

where $U^{\pi}_{i,t+1}:=\boldsymbol{u}\left(\pi,X_{i,t+1}\right)$ and

$\boldsymbol{u}\left(\pi,\boldsymbol{x}\right):=\left[\boldsymbol{\phi}(\boldsymbol{x})^{\top}\pi(1|\boldsymbol{x}),\cdots,\boldsymbol{\phi}(\boldsymbol{x})^{\top}\pi(M|\boldsymbol{x})\right]^{\top}.$
(14)

The ACPE also requires a penalty function $p\left(\cdot\right)$. We consider
concave penalties satisfying Assumption 6 that can produce unbiased estimates.
Examples includes SCAD (Fan and Li, 2001) and MCP (Zhang et al., 2010), which
are defined respectively as

$\displaystyle p_{\eta}^{MCP}(t,\lambda)$
$\displaystyle=\lambda\int_{0}^{t}\left(1-x/(\eta\lambda)\right)_{+}dx,\quad\eta>1,$
(15) $\displaystyle p_{\eta}^{SCAD}(t,\lambda)$
$\displaystyle=\lambda\int_{0}^{t}\min\left\\{1,\left(\gamma-x/\lambda\right)_{+}/(\gamma-1)\right\\}dx,\quad\eta>2,$
(16)

where $\eta$ is a parameter that controls the concavity of the penalty
function. In particular, both penalties converge to the $L_{1}$ penalty as
$\eta\rightarrow\infty$.

The procedure of ACPE is summarized in Algorithm 1. Under a semi-parametric
$Q$ function approximation, we obtain $\widehat{}\boldsymbol{\beta}^{\pi}$ by
solving (10) with instantiated
$\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)$ given in (13) and a
concave penalty $p\left(\cdot\right)$ satisfying Assumption 6. Then, the
individual-wise value functions can be estimated by

$\widehat{Q}^{i}\left(\pi,\boldsymbol{x},a\right)=\boldsymbol{z}\left(\boldsymbol{x},a\right)^{\top}\widehat{}\boldsymbol{\beta}^{\pi,i},\quad\text{and}\quad\widehat{V}^{i}\left(\pi,\boldsymbol{x},a\right)=\boldsymbol{u}\left(\pi,\boldsymbol{x}\right)^{\top}\widehat{}\boldsymbol{\beta}^{\pi,i},\quad\text{for}\quad
i\in[N],$ (17)

where $\boldsymbol{z}\left(\boldsymbol{x},a\right)$ and
$\boldsymbol{u}\left(\pi,\boldsymbol{x}\right)$ are defined in (12) and (14),
respectively.

Input: Sample trajectories
$\left\\{X_{i,t},A_{i,t},R_{i,t},X_{i,t+1}\right\\}$, $t\in[T],i\in[N]$, a
policy $\pi\left(a\mid\boldsymbol{x}\right)$ to be evaluated, and a tuning
parameter $\lambda$.

Output: Individual $Q$ coefficients $\widehat{}\boldsymbol{\beta}^{\pi,i}$ for
$i\in[N]$, number of groups $K$, a partition of $N$ trajectories, ${\cal
G}=\left\\{{\cal G}^{(1)},\cdots,{\cal G}^{(K)}\right\\}$ and group
coefficients $\widehat{}\boldsymbol{\theta}^{\pi,(k)}$, $k\in[K]$.

1

2Estimate $\widehat{}\boldsymbol{\beta}^{\pi}$ by solving
$\widehat{}\boldsymbol{\beta}^{\pi}=\underset{\boldsymbol{\beta}=[(\boldsymbol{\beta}^{1})^{\top}\cdots(\boldsymbol{\beta}^{N})^{\top}]^{\top}}{\arg\min}\;\left[\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)^{\top}\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)+\frac{1}{N^{2}}\sum_{1\leq
i<j\leq
N}p\left((JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right\rVert_{2},\lambda\right)\right],$
with instantiated $\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)$ given in
(13) and a concave penalty $p\left(\cdot\right)$ satisfying Assumption 6. Two
candidates for $p\left(\cdot\right)$ are given in (15) and (16).

3Let
${\textsc{mat}}\left(\widehat{}\boldsymbol{\beta}^{\pi}\right)=\left[\widehat{}\boldsymbol{\beta}^{\pi,1},\cdots,\widehat{}\boldsymbol{\beta}^{\pi,N}\right]$
be the matricization of $\widehat{}\boldsymbol{\beta}^{\pi}$.

4Apply a chosen clustering algorithm to cluster $N$ trajectories to $K$ groups
based on $\left\\{\widehat{}\boldsymbol{\beta}^{\pi,i}\right\\}_{i\in[N]}$ and
obtain the $N\times K$ group membership matrix $\widehat{}\boldsymbol{W}$.

5Calculate the group-wise coefficients by
$\widehat{}\boldsymbol{\theta}^{\pi,(k)}={\textsc{mat}}\left(\widehat{}\boldsymbol{\beta}^{\pi}\right)\,\widehat{}\boldsymbol{w}_{\cdot
k}/\widehat{N}^{(k)},$ where
$\widehat{N}^{(k)}=\widehat{}\boldsymbol{w}_{\cdot
k}^{\top}\widehat{}\boldsymbol{w}_{\cdot k}$ estimates the number of
trajectories in ${\cal G}^{(k)}$.

Algorithm 1 Auto-Clustered Policy Evaluation (ACPE)

Given estimated $\widehat{}\boldsymbol{\beta}^{\pi,i}$, we can estimate the
number of groups $K$ and the $N\times K$ group membership matrix
$\widehat{}\boldsymbol{W}=\left[w_{i,k}\right]_{1\leq i\leq N,1\leq k\leq K}$
by using any chosen clustering algorithm. The $k$-th group coefficient is
estimated as

$\widehat{}\boldsymbol{\theta}^{\pi,(k)}={\textsc{mat}}\left(\widehat{}\boldsymbol{\beta}^{\pi}\right)\widehat{}\boldsymbol{w}_{\cdot
k}/\widehat{N}^{(k)},$ (18)

where $\widehat{N}^{(k)}=\widehat{}\boldsymbol{w}_{\cdot
k}^{\top}\widehat{}\boldsymbol{w}_{\cdot k}$ estimates the number of
trajectories in ${\cal G}^{(k)}$ and
${\textsc{mat}}\left(\widehat{}\boldsymbol{\beta}^{\pi}\right)\in\mathbb{R}^{MJ\times
N}$ is the matricization of $\widehat{}\boldsymbol{\beta}^{\pi}$. When $N$ or
$T$ are sufficiently large, by the oracle property in Theorem 2, we have
$\Pr\left(\widehat{}\boldsymbol{\beta}^{\pi,i}=\widehat{}\boldsymbol{\beta}^{\pi,j}\right)\rightarrow
1$ if $i$ and $j$ belong to the same group. The group coefficients
$\widehat{}\boldsymbol{\theta}^{\pi,(k)}$, $1\leq k\leq K$, can be chosen as
the distinct values of $\widehat{}\boldsymbol{\beta}^{\pi,i}$, $i\in[N]$.
Accordingly, we have $\Pr\left(\widehat{K}=K\right)\rightarrow 1$. The value
functions for ${\cal G}^{(k)}$ can be estimated by

$\widehat{Q}^{(k)}\left(\pi,\boldsymbol{x},a\right)=\boldsymbol{z}\left(\boldsymbol{x},a\right)^{\top}\widehat{}\boldsymbol{\theta}^{\pi,(k)},\quad\text{and}\quad\widehat{V}^{(k)}\left(\pi,\boldsymbol{x}\right)=\boldsymbol{u}\left(\pi,\boldsymbol{x}\right)^{\top}\widehat{}\boldsymbol{\theta}^{\pi,(k)},\quad\text{for}\quad
1\leq k\leq K,$ (19)

where $\boldsymbol{z}\left(\boldsymbol{x},a\right)$ and
$\boldsymbol{u}\left(\pi,\boldsymbol{x}\right)$ are defined in (12) and (14),
respectively. Let ${\cal R}$ be a reference distribution on the covariate
space ${\cal X}_{k}$ of the $k$-th group. The integrated value function for
group $k$ is estimated by

$\widehat{V}^{(k)}_{\cal
R}\left(\pi\right)=\int\widehat{V}^{(k)}\left(\pi,\boldsymbol{x}\right){\cal
R}\left(d\boldsymbol{x}\right).$ (20)

###  3.2 Auto-Clustered Policy Improvement

For the purpose of finding the optimal policy, we focus on the policy class
defined in (5). Because of the heterogeneity across different sub-populations,
we expect their respect optimal polices are different. For each group $k$, its
policy is denoted as $\pi\left(\boldsymbol{\alpha}^{(k)}\right)$, where the
group-wise policy is indexed by parameter $\boldsymbol{\alpha}^{(k)}$,
$k\in[K]$.

The idea of Auto-Clustered Policy Iteration (ACPI) is to use ACPE (Algorithm
1) to estimate the value of any policy, obtain a $K$-partition of $N$
trajectories, and maximize the integrated state-value function over the class
of policies (5) for each group $k\in[K]$. The procedure is summarized in
Algorithm 2.

Input: Sample trajectories
$\left\\{\left(X_{i,t},A_{i,t},R_{i,t},X_{i,t+1}\right)\right\\}$,
$t\in[T],i\in[N]$, and tuning parameter $\lambda$.

Output: A partition of $N$ trajectories into $K$ groups, and an optimal policy
for each group $k\in[K]$ in parametric family (5), with parameter
$\boldsymbol{\alpha}^{(k)}$.

1

2At step $s=1$, we initialize the number of groups $K=1$ (the value of $K$
will be updated immediately in the ACPE step) and policy coefficients
$\alpha_{s}^{{(k)}}$, $1\leq k\leq K$.

3while  _Not converged_ do

4 for  _each group $1\leq k\leq K$_ do

5 Let $\pi_{s}^{(k)}=\pi\left(\alpha_{s}^{{(k)}}\right)$, we apply ACPE
(Algorithm 1) to obtain
$\widehat{}\boldsymbol{\beta}^{\pi_{s}^{(k)},1},\cdots,\widehat{}\boldsymbol{\beta}^{\pi_{s}^{(k)},N}$.

6

7 Apply a chosen clustering algorithm to cluster $N$ trajectories to $K$
groups based on
$\left\\{\widehat{}\boldsymbol{\beta}^{\pi_{s}^{(k)},i}\right\\}_{i\in[N]}$
and obtain the $N\times K$ group membership matrix $\widehat{}\boldsymbol{W}$.

8 Calculate the group-wise coefficients
$\widehat{}\boldsymbol{\theta}^{\pi_{s}^{(k)},(k)}={\textsc{mat}}\left(\widehat{}\boldsymbol{\beta}^{\pi_{s}^{(k)}}\right)\,\widehat{}\boldsymbol{w}_{\cdot
k}/\widehat{N}^{(k)},$ where $N^{(k)}=\widehat{}\boldsymbol{w}_{\cdot
k}^{\top}\widehat{}\boldsymbol{w}_{\cdot k}$ estimates the number of
trajectories in ${\cal G}^{(k)}$.

9 for  _each group $1\leq k\leq K$_ do

10 Let $\pi_{s}^{(k)}=\pi\left(\boldsymbol{\alpha}^{(k)}_{s}\right)$ be the
fixed index, we update
$\boldsymbol{\alpha}^{(k)}_{s+1}\leftarrow\arg\underset{\boldsymbol{\alpha}}{\max}\;\widehat{V}_{{\cal
R}}\left(\boldsymbol{\alpha}\right)$ where the integrated value function
$\widehat{V}_{{\cal R}}\left(\boldsymbol{\alpha}\right)$ is defined as
$\widehat{V}_{{\cal
R}}\left(\boldsymbol{\alpha}\right)=\int\boldsymbol{u}\left(\pi\left(\boldsymbol{\alpha}\right),\boldsymbol{x}\right)^{\top}\widehat{}\boldsymbol{\theta}^{\pi_{s}^{(k)},(k)}d{\cal
R}(\boldsymbol{x}),$ and $\pi\left(\boldsymbol{\alpha}\right)$ and
$\boldsymbol{u}\left(\pi\left(\boldsymbol{\alpha}\right),\boldsymbol{x}\right)$
are defined in (14) and (5), respectively.

11 $s\leftarrow s+1$.

Algorithm 2 Auto-Clustered Policy Iteration (ACPI)

##  4 Theory

In this section, we lay out the theoretical framework for ACPE and ACPI in a
double-divergence structure, which allows either sample size $N$ or decision
horizon $T$ to go to infinity. We establish asymptotic properties of the
offline estimation of the coefficients, value functions and optimal policy.
Let
$\boldsymbol{W}=\left[\boldsymbol{w}_{1\cdot}\cdots\boldsymbol{w}_{N\cdot}\right]^{\top}\in\mathbb{R}^{N\times
K}$ be the group membership matrix where $w_{ik}=1$ for $i\in{\cal G}^{(k)}$
and $w_{ik}=0$ otherwise. We first present theoretical results for the oracle
estimator of ACPE when the true $\boldsymbol{W}$ is known a priori. Section
4.1 derives the oracle properties, including $\ell_{2}$ and $\ell_{\infty}$
convergence rates, when the true sub-population information $\left\\{{\cal
G}^{(k)},1\leq k\leq K\right\\}$ is known. Then, in Section 4.2, we establish
the $\ell_{2}$ convergence rate and the asymptotic normality of the parameters
estimated from (10) when $\boldsymbol{W}$ is not known. Lastly, in Section
4.3, we present that the estimated optimal policies for each group ${\cal
G}_{k}$ convergences to their respective true optimal policies. As a
convention, we denote the true, oracle and estimated values as
$\mathring{\beta}$, $\widetilde{\beta}$, and $\widehat{\beta}$, respectively.
The proofs of all theorems are provided in the supplementary materials.

###  4.1 Properties of the oracle estimators

We denote
$\boldsymbol{\theta}^{\pi}=\left[{\boldsymbol{\theta}^{\pi,(1)}}^{\top}\cdots{\boldsymbol{\theta}^{\pi,(K)}}^{\top}\right]^{\top}$
as the group coefficient matrix where $\boldsymbol{\theta}^{\pi,(k)}$ is the
coefficient of the $Q$ function for ${\cal G}^{(k)}$. When the
$\boldsymbol{W}$ is known, the oracle estimator refers to the estimator that
minimizes the objective function in (10) with respect to
$\boldsymbol{\theta}^{\pi}$ without penalty, that is

$\widetilde{}\boldsymbol{\theta}^{\pi}=\underset{\boldsymbol{\theta}}{\arg\min}\;\widetilde{}\boldsymbol{G}\left(\pi,\boldsymbol{\theta}\right)^{\top}\widetilde{}\boldsymbol{G}\left(\pi,\boldsymbol{\theta}\right),$
(21)

where
$\widetilde{}\boldsymbol{G}\left(\pi,\boldsymbol{\theta}\right)=\boldsymbol{G}\left(\pi,(\boldsymbol{W}\otimes\boldsymbol{I}_{JM})\boldsymbol{\theta}\right)$,
or equivalently,

$\widetilde{}\boldsymbol{G}\left(\pi,\boldsymbol{\theta}\right)=\frac{1}{NTJ}\sum_{i=1}^{N}\sum_{t=1}^{T}\left((\widetilde{}\boldsymbol{\Lambda}^{i})^{\top}Z_{i,t}R_{i,t}-(\widetilde{}\boldsymbol{\Lambda}^{i})^{\top}\left(Z_{i,t}Z_{i,t}^{\top}-\gamma
Z_{i,t}(U_{i,t+1}^{\pi})^{\top}\right)\widetilde{}\boldsymbol{\Lambda}^{i}\boldsymbol{\theta}\right),$
(22)

and
$\widetilde{}\boldsymbol{\Lambda}^{i}=\boldsymbol{\Lambda}^{i}(\boldsymbol{W}\otimes\boldsymbol{I}_{JM})$,
$\widetilde{}\boldsymbol{\beta}^{\pi}=(\boldsymbol{W}\otimes\boldsymbol{I}_{JM})\widetilde{}\boldsymbol{\theta}^{\pi}$,
$Z_{i,t}$ and $U^{\pi}_{i,t}$ are defined in (12) and (14), respectively.

We first detail the assumptions that are necessary for the convergence and
asymptotic normality of estimates obtained by the ACPE.

######  Definition 2 ($\kappa$-Smooth functions).

Let $f\left(\cdot\right)$ be an arbitrary function on ${\cal
X}\in\mathbb{R}^{p}$. For a $p$-tuple
$\boldsymbol{\alpha}=(\alpha_{1},\cdots,\alpha_{p})$ fo non-negative integers,
let $D^{\alpha}$ denote the differential operator:

$D^{\boldsymbol{\alpha}}f(\boldsymbol{x})=\frac{{\partial}^{\left\lVert\boldsymbol{\alpha}\right\rVert_{1}}f(\boldsymbol{x})}{\partial
x_{1}^{\alpha_{1}}\cdots\partial x_{p}^{\alpha_{p}}},$

where $\boldsymbol{x}=(x_{1},\cdots,x_{p})^{\top}$. The class of
$\kappa$-smooth functions is defined as

${\cal
H}(\kappa,c)=\left\\{f:\underset{\left\lVert\boldsymbol{\alpha}\right\rVert_{1}\leq\lfloor\kappa\rfloor}{\sup}\;\underset{\boldsymbol{x}\in{\cal
X}}{\sup}\;\left\lvert
D^{\boldsymbol{\alpha}}f(\boldsymbol{x})\right\rvert\leq c;\text{ and
}\underset{\left\lVert\boldsymbol{\alpha}\right\rVert_{1}\leq\lfloor\kappa\rfloor}{\sup}\;\underset{\boldsymbol{x}_{1},\boldsymbol{x}_{2}\in{\cal
X},\boldsymbol{x}_{1}\neq\boldsymbol{x}_{2}}{\sup}\;\frac{\left\lvert
D^{\boldsymbol{\alpha}}f(\boldsymbol{x}_{1})-D^{\boldsymbol{\alpha}}f(\boldsymbol{x}_{2})\right\rvert}{\left\lVert\boldsymbol{x}_{1}-\boldsymbol{x}_{2}\right\rVert_{2}^{\kappa-\lfloor\kappa\rfloor}}\right\\}.$

######  Assumption 2.

There exists some $\kappa$, $c>0$ such that $r(\boldsymbol{x},a)$,
$\Pr(\boldsymbol{x}^{\prime}|\boldsymbol{x},a)$ belong to the class of
$\kappa$-smooth function of $\boldsymbol{x}$ for any $a\in{\cal A}$ and
$\boldsymbol{x}^{\prime}\in{\cal X}$.

Lemma 1 in Shi et al. (2020) shows that under Assumption 2, there exists some
constant $c^{\prime}>0$ such that $Q\left(\pi,\boldsymbol{x},a\right)$ belongs
to the class of $\kappa$-smooth function of $\boldsymbol{x}$ for any policy
$\pi$ and any action $a\in{\cal A}$. This implies that the $Q$ function has
bounded derivatives up to order $\lfloor\kappa\rfloor$. To approximate the $Q$
function, we restrict our attention to two particular type of Sieve basis
functions. Let $BSpline(J,r)$ denote a tensor-product B-spline basis of
dimension $J$ and of degree $r$ on $[0,1]^{p}$ and $Wav(J,r)$ denote a tensor-
product Wavelet basis of regularity $r$ and dimension $J$ on $[0,1]^{p}$. The
sieve $\boldsymbol{\phi}_{J}$ we use is either $BSpline(J,r)$ or $Wav(J,r)$
with $r>\max\left(\kappa,1\right)$. This, together with Assumption 2, implies
that there exists a set of vectors
$\\{\mathring{\boldsymbol{\theta}}^{\pi}_{a}\\}_{a\in{\cal A}}$ that satisfies
$\underset{\boldsymbol{x}\in{\cal X},a\in{\cal A}}{\sup}\left\lvert
Q(\pi,\boldsymbol{x},a)-\Phi_{J}^{\top}(\boldsymbol{x})\mathring{\boldsymbol{\theta}}^{\pi}_{a}\right\rvert=\mathcal{O}\left(J^{-\kappa/p}\right)$.

######  Assumption 3.

Let $\mu$ be the limiting density function of $X_{i,t}$ and $\nu_{0}$ be the
initial state density. The density function $\mu$ and $\nu_{0}$ are uniformly
bounded away from zero and infinity on ${\cal X}$.

######  Assumption 4.

There exists some constant $C_{1}>0$ such that

$\underset{\pi\in\Pi}{\inf}\;\lambda_{\min}\left(T^{-1}\sum_{t=0}^{T-1}{\mathbb{E}}\left[Z_{i,t}Z_{i,t}^{\top}-\gamma^{2}\overline{}\boldsymbol{u}(\pi,X_{i,t},A_{i,t})\overline{}\boldsymbol{u}(\pi,X_{i,t},A_{i,t})^{\top}\right]\right)\geq
C_{1},$

where
$\overline{}\boldsymbol{u}(\pi,\boldsymbol{x},a)={\mathbb{E}}\left[\boldsymbol{u}\left(\pi,X_{i,t+1}\right)\,\mid\,X_{i,t}=\boldsymbol{x},A_{i,t}=a\right]$
and $\lambda_{\min}(\cdot)$ denotes the minimal eigenvalue of a matrix.

######  Assumption 5.

Under the setting that $T\rightarrow\infty$, the Markov chain
$\left\\{X_{i,t}\right\\}_{t\geq 0}$ is geometrically ergodic, that is, there
exists some function $f(\boldsymbol{x})$ on ${\cal X}$ and some constant
$c\leq 1$ such that $\int_{\boldsymbol{x}\in{\cal
X}}f(\boldsymbol{x})\mu(\boldsymbol{x})d\boldsymbol{x}<+\infty$ and

$\left\lVert
P_{t}(\cdot|\boldsymbol{x})-\mu(\boldsymbol{x})\right\rVert_{TV}\leq
f(\boldsymbol{x})\rho^{t},\qquad\forall t\geq 0,$

where $\mu(\cdot)$ is the limiting density function and
$\left\lVert\cdot\right\rVert_{TV}$ denotes the total variation norm.

Assumption 3 is weak, it does not require the limiting density function $\mu$
to be equal to the limit state density $\nu_{0}$. When $T$ is finite,
Assumption 4 guarantees that $\widetilde{}\boldsymbol{\Sigma}$ defined in (37)
is invertible with probability approaching one when $N\rightarrow\infty$. When
$T$ is infinite, Assumption 5 (i) guarantees that the matrix
${\mathbb{E}}\left[\widetilde{}\boldsymbol{\Sigma}\right]$ is invertible and
(ii) enables us to derive matrix concentration inequalities for
$\widetilde{}\boldsymbol{\Sigma}$. Together, they imply that
$\widetilde{}\boldsymbol{\Sigma}$ is invertible with probability approaching
one. Assumption 4 and 5 are needed to show the existence of a unique
$\\{\mathring{\boldsymbol{\theta}}^{\pi}_{a}\\}_{a\in{\cal A}}$ uniformly over
$\Pi$. They are mild and can be verified empirically by checking that certain
data-dependent matrices are invertible. More discussion on similar assumptions
can be found in Luckett et al. (2019) and Shi et al. (2020).

Let
$\mathring{\boldsymbol{\theta}}^{\pi}=\left[\mathring{\boldsymbol{\theta}}^{\pi,(1)\top}\;\cdots\;\mathring{\boldsymbol{\theta}}^{\pi,(K)\top}\right]^{\top}$
where
$\mathring{\boldsymbol{\theta}}^{\pi,(k)}=\left[\mathring{\boldsymbol{\theta}}_{1}^{\pi,(k)\top}\;\cdots\;\mathring{\boldsymbol{\theta}}_{M}^{\pi,(k)\top}\right]^{\top}$
for $k\in[K]$. Proposition 1 and Theorem 5 in the supplemental material
establish the $\ell_{2}$ convergence and the asymptotic normality of the
oracle estimator $\widetilde{}\boldsymbol{\theta}^{\pi}$. Corollary 2 in the
supplemental material shows the the asymptotic normality of the integrated
value of a given policy for each sub-population. They imply that the oracle
group-wise estimator converges at a rate of $\left(TNP_{k}\right)^{-1/2}$ for
$\forall k\in[K]$. The estimators in Luckett et al. (2019); Jiang and Li
(2016); Thomas et al. (2015) typically converge at a rate of
$\left(NP_{k}\right)^{-1/2}$ and are not suitable for settings when one sub-
population has only few trajectories. Our estimation procedure also aggregates
information along horizon $T$ and achieves same rates as in Shi et al. (2020)
when group membership $\boldsymbol{W}$ are known a priori.

To finally establish the large sample theory for the ACPE when group
membership are unknown, we need a stronger uniform consistency result for the
oracle estimators when either $N\rightarrow\infty$ or $T\rightarrow\infty$.

######  Theorem 1 (Oracle estimator uniform convergence).

Suppose Assumption 2 – 5 hold. Let $N_{min}=\underset{1\leq k\leq
K}{\min}N_{k}$ and $N_{max}=\underset{1\leq k\leq K}{\max}N_{k}$. If $K={\rm
o}\left(N_{\min}T\right)$, $J\ll\sqrt{N_{\min}T}/\log(N_{\min}T)$,
$J^{-\kappa/p}\ll 1/\sqrt{N_{\max}T}$, we have with probability at least
$1-2JMK(N_{\min}T)^{-2}-\mathcal{O}\left((N_{\min}T)^{-2}\right)$ that

$\underset{\pi\in\Pi}{\sup}\left\lVert\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right\rVert_{\infty}\leq\phi_{NT},$

where

$\phi_{NT}=6cC^{-1}\frac{N_{\max}}{N_{\min}}\sqrt{2J\frac{\log(N_{\min}T)}{N_{\min}T}}.$
(23)

######  Remark 4.

By Definition 1, $N_{\max}/N_{\min}=P_{\max}/P_{\min}$ are bounded away from
infinity. The bound (23) can be simplified to
${\mathcal{O}_{p}}\left(\sqrt{J\frac{\log(NT)}{NT}}\right)$.

###  4.2 Properties of the feasible estimator

We now study the theoretical properties of the feasible estimator when the
true group membership $\boldsymbol{W}$ is not known. We introduce the
following assumptions on the penalty function and the minimum signal
difference between groups.

######  Assumption 6.

The penalty function $p\left(x,\lambda\right)$ is a symmetric function of $x$,
and it is non-decreasing and concave in $x$ for $x\in[0,+\infty)$. Let
$\rho(x)=\lambda^{-1}p\left(x,\lambda\right)$, there exists a constant
$0<c<\infty$ such that $\rho(0)=0$ and $\rho(x)$ is a constant for all $x\geq
c\lambda$. Its derivative $\rho^{\prime}(x)$ exists and is continuous except
for a finite number of $x$ and $\rho^{\prime}(0+)=1$.

######  Assumption 7.

For $K>2$, define the minimal difference of the common values between any pair
of groups as

$d_{NT}=(JM)^{-1/2}\;\underset{k\neq
l}{\min}\;\left\lVert\mathring{\boldsymbol{\theta}}_{k}-\mathring{\boldsymbol{\theta}}_{l}\right\rVert_{2}.$

We assume that $d_{NT}\gg\phi_{NT}$ where $\phi_{NT}$ is given in (23).

Assumption 6 is commonly assumed in high-dimensional settings. It is satisfied
by concave penalties such as MCP (15) and SCAD (16). Assumption 7 is the
separability condition on the minimum signal difference between groups that is
needed to recover the true groups.

######  Theorem 2 (Feasible estimator).

Suppose the conditions in Theorem 1 and Assumption 6 and 7 hold and $K\geq 2$.
Let $N_{min}=\underset{1\leq k\leq K}{\min}N_{k}$ and $N_{max}=\underset{1\leq
k\leq K}{\max}N_{k}$. If $d_{NT}\geq C\lambda$ and $\lambda\gg\phi_{NT}$ where
$C$, $d_{NT}$, and $\phi_{NT}$ are defined in Assumption 6, 7, and Equation
(23). Then, uniformly over $\Pi$, there exists a local minimizer
$\widehat{}\boldsymbol{\beta}^{\pi}$ of the objective function ${\cal L}_{NT}$
given in (10) satisfying

$\Pr\left(\widehat{}\boldsymbol{\beta}^{\pi}=\widetilde{}\boldsymbol{\beta}^{\pi}\right)\rightarrow
1.$

Recall that we define $\widehat{}\boldsymbol{W}$ as an estimator of
$\boldsymbol{W}$ that is obtained by applying any clustering method on the
column vector of ${\textsc{mat}}\left(\boldsymbol{\beta}^{\pi}\right)$. A
direct conclusion from Theorem 2 is that
$\Pr\left(\widehat{}\boldsymbol{W}=\boldsymbol{W}\right)\rightarrow 1$ since
the oracle estimator
$\widetilde{}\boldsymbol{\beta}^{\pi}_{i}=\widetilde{}\boldsymbol{\beta}^{\pi}_{j}$
for any $i,j\in{\cal G}^{(k)}$.

The oracle property in Theorem 2 together with Corollary 2 in Appendix A
directly leads to the asymptotic distribution of $\widehat{\beta}^{\pi}$,
which is presented in the following theorem.

######  Theorem 3 (Asymptotic normality of integrated value estimator).

Suppose the conditions in Theorem 1 and 2 hold. If
$J^{-\kappa/p}\ll\left(N_{\max}T\left(1+\left\lVert\int\boldsymbol{\phi}_{J}(\boldsymbol{x}){\cal
R}(d\boldsymbol{x})\right\rVert_{2}^{-2}\right)\right)^{-1/2}$, as either
$N_{\min}\rightarrow\infty$ or $T\rightarrow\infty$, we have for any
$\boldsymbol{x}\in{\cal X}$,

$\displaystyle\sqrt{NT}\cdot\widehat{\sigma}_{{\cal
R}}^{(k)}(\pi)^{-1}\left(\widehat{V}_{{\cal R}}^{(k)}\left(\pi\right)-V_{{\cal
R}}^{(k)}\left(\pi\right)\right)$ $\displaystyle\stackrel{{\scriptstyle{\cal
D}}}{{\longrightarrow}}{\cal N}\left(0,1\right)\,,$

for any $k\in[K]$, where

$\widehat{\sigma}_{{\cal
R}}^{(k)}\left(\pi\right)^{2}=\left(\int\boldsymbol{u}\left(\pi,\boldsymbol{x}\right){\cal
R}(d\boldsymbol{x})\right)^{\top}\widehat{}\boldsymbol{\Sigma}^{\pi,(k)}\widehat{}\boldsymbol{\Omega}^{\pi,(k)}(\widehat{}\boldsymbol{\Sigma}^{\pi,(k)\top})^{-1}\left(\int\boldsymbol{u}\left(\pi,\boldsymbol{x}\right){\cal
R}(d\boldsymbol{x})\right),$

where $\widehat{}\boldsymbol{\Sigma}^{\pi,(k)}$ and
$\widehat{}\boldsymbol{\Omega}^{\pi,(k)}$ are defined as

$\displaystyle\widehat{}\boldsymbol{\Sigma}^{\pi,(k)}$
$\displaystyle=\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=0}^{T-1}(\widehat{}\boldsymbol{\Lambda}^{i,(k)})^{\top}Z_{i,t}\left(Z_{i,t}-\gamma
U_{i,t+1}^{\pi}\right)^{\top}\widehat{}\boldsymbol{\Lambda}^{i,(k)},$ (24)
$\displaystyle\widehat{}\boldsymbol{\Omega}^{\pi,(k)}$
$\displaystyle=\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=0}^{T-1}(\widehat{}\boldsymbol{\Lambda}^{i,(k)})^{\top}Z_{i,t}Z_{i,t}^{\top}\widehat{}\boldsymbol{\Lambda}^{i,(k)}\left(R_{i,t}-\left(Z_{i,t}-\gamma
U_{i,t+1}^{\pi}\right)^{\top}\widehat{}\boldsymbol{\Lambda}^{i,(k)}\widehat{}\boldsymbol{\theta}^{\pi}\right)^{2},$
(25)

and
$\widehat{}\boldsymbol{\Lambda}^{i,(k)}=\boldsymbol{\Lambda}^{i}(\widehat{}\boldsymbol{w}_{\cdot
k}\otimes\boldsymbol{I}_{JM})$, $Z_{i,t}$ and $U^{\pi}_{i,t}$ are defined in
(12) and (14), respectively.

Theorem 3 imply that the group-wise estimator converges at a rate of
$\left(NTP_{k}\right)^{-1/2}$, $k\in[K]$. The estimators in Luckett et al.
(2019); Jiang and Li (2016); Thomas et al. (2015) typically converge at a rate
of $\left(NP_{k}\right)^{-1/2}$ and are not suitable for settings when one
sub-population has only a few trajectories. Our estimation procedure also
aggregates information along horizon $T$. The theoretical property along $T$
is obtained by treating finite and infinite $T$ separately, using the matrix
concentration and martingale center limit theorem.

The dimension of the unknown parameters $\boldsymbol{\beta}$ in (10) is $NJM$,
which will diverge as sample size $N$ increases. Without the $K$-hetero MDP
assumption and the penalty term, the $\ell_{2}$ convergence of the parameters
can be shown to be $1/\sqrt{T}$ under the TH Markov and CMI assumption,
employing arguments similar to those in Luckett et al. (2019) and Shi et al.
(2020). The information accumulated along $N$ does not help. In contrast, the
ACPE obtain $1/\sqrt{TNP_{\min}}$ convergence rate where
$P_{\min}=\min\left\\{P^{(1)},\cdots,P^{(k)}\right\\}$. The estimation in
Luckett et al. (2019) is based on the state value $V(\cdot)$ and the
calculation necessitates a correct specification of the behavior policy for
the importance ratio that shows up in the Bellman equation of $V(\cdot)$. In
contrast, our method estimates the action value function $Q(\cdot)$ and derive
the corresponding state value estimators. As a result, we do not need to
specify the behavior policy, nor do we need to estimate it from the observed
dataset. Our method can be viewed as an implicit importance weighting.

###  4.3 Optimal policy

In this section, we establish a convergence result for the estimated optimal
policy for each homogeneous sub-population assuming that the parametric class
$\Pi$ satisfies the following properties.

######  Assumption 8.

The map $\boldsymbol{\alpha}\rightarrow V_{{\cal
R}}\left(\pi\left(\boldsymbol{\alpha}\right)\right)$ has a unique and well-
separated maximum $\boldsymbol{\alpha}^{*}$ in the interior of the support of
$\boldsymbol{\alpha}$, where $V_{{\cal
R}}\left(\pi\left(\boldsymbol{\alpha}\right)\right)$ is defined in (4).

######  Assumption 9.

For any $\boldsymbol{x}\in{\cal X}$ and $a\in{\cal A}$, we have as
$\delta\downarrow 0$,

$\underset{\left\lVert\boldsymbol{\alpha}_{1}-\boldsymbol{\alpha}_{2}\right\rVert_{2}\leq\delta}{\sup}\;\mathbb{E}\left\lVert\pi\left(a,\boldsymbol{x};\boldsymbol{\alpha}_{1}\right)-\pi\left(a,\boldsymbol{x};\boldsymbol{\alpha}_{1}\right)\right\rVert\rightarrow
0.$

Assumption 8 requires that the true optimal decision in each state is unique
(see also the Assumption A.8 of Ertefaie and Strawderman (2018) and Assumption
6 of Luckett et al. (2019)) and is a standard assumption in M-estimation.
Assumption 9 requires smoothness on the class of the polices. The following
lemma shows that the parametric class of policies $\Pi$ defined in (5)
satisfies Assumption 8 and 9.

######  Lemma 2.

The parametric class of policies $\Pi$ defined in (5) satisfies Assumption 8
and 9.

Theorem 4 establishes that the estimated optimal policy for each homogeneous
sub-population converges in probability to the true optimal policy over $\pi$
and that the estimated value of the estimated optimal policy converges to the
true value of the estimated optimal policy.

######  Theorem 4.

Suppose the conditions in Theorem 1 and 2 and Assumption 8, 9 hold. Let
$\widehat{}\boldsymbol{\alpha}^{(k)\,*}=\arg\max_{\boldsymbol{\alpha}}\widehat{V}_{{\cal
R}}^{(k)}\left(\pi(\alpha)\right)$ and
$\boldsymbol{\alpha}^{(k)\,*}=\arg\max_{\boldsymbol{\alpha}}V_{{\cal
R}}^{(k)}\left(\pi(\alpha)\right)$, we have as either
$N_{\min}\rightarrow\infty$ or $T\rightarrow\infty$,

  1. (i)

$\left\lVert\widehat{}\boldsymbol{\alpha}^{(k)\,*}-\boldsymbol{\alpha}^{(k)\,*}\right\rVert_{2}\stackrel{{\scriptstyle{\cal
P}}}{{\longrightarrow}}0$.

  2. (ii)

$\left\lvert V_{\cal
R}\left(\pi(\widehat{}\boldsymbol{\alpha}^{(k)\,*})\right)-V_{\cal
R}\left(\pi(\boldsymbol{\alpha}^{(k)\,*})\right)\right\rvert\stackrel{{\scriptstyle{\cal
P}}}{{\longrightarrow}}0$.

##  5 Computation

The optimization problem of (10) is challenging because of the coupling of
$\boldsymbol{\beta}_{i}$ and $\boldsymbol{\beta}_{j}$ in the penalty term.
This problem is similar to that in Chi and Lange (2015) and can be solved by
the alternating direction method of multipliers (ADMM) (Boyd et al., 2011),
the alternating minimization algorithm (AMA) (Tseng, 1991), and the general
iterative shrinkage and thresholding method (Gong et al., 2013). All three
approaches employ variable splitting to handle the shrinkage penalties in
(10). In this section, we summarize the ADMM procedure to solve (10); readers
are referred to Boyd et al. (2011); Chi and Lange (2015); Gong et al. (2013)
and references therein for more information on alternative optimization
algorithms. For brevity, we suppress the superscript $\pi$ in
$\boldsymbol{\beta}^{\pi}$ in this section.

We first recast (10) as an equivalent constrained problem:

$\displaystyle\min_{\boldsymbol{\beta},\boldsymbol{\delta}}$
$\displaystyle{\cal
L}\left(\boldsymbol{\beta},\boldsymbol{\delta}\right):=\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)^{\top}\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)+\frac{1}{N^{2}}\sum_{1\leq
i<j\leq
N}p\left((JM)^{-1/2}\left\lVert\boldsymbol{\delta}_{ij}\right\rVert_{2},\lambda\right)$
(26) $\displaystyle s.t.$
$\displaystyle\boldsymbol{\delta}_{ij}=\boldsymbol{\beta}_{i}-\boldsymbol{\beta}_{j},\quad\forall\\{i,j\\}:1\leq
i<j\leq N,$

where
$\boldsymbol{\delta}=\left[\boldsymbol{\delta}_{ij}^{\top},i<j\right]^{\top}$.
Applying the augmented Lagrangian method (ALM), the solution of (26) can be
obtained by minimizing

${\cal
L}_{AL}\left(\boldsymbol{\beta},\boldsymbol{\delta},\boldsymbol{\nu}\right):={\cal
L}\left(\boldsymbol{\beta},\boldsymbol{\delta}\right)+\frac{1}{N^{2}}\sum_{i<j}\langle\boldsymbol{\nu}_{ij},\boldsymbol{\beta}_{i}-\boldsymbol{\beta}_{j}-\boldsymbol{\delta}_{ij}\rangle+\frac{\rho}{2JMN^{2}}\sum_{i<j}\left\lVert\boldsymbol{\beta}_{i}-\boldsymbol{\beta}_{j}-\boldsymbol{\delta}_{ij}\right\rVert_{2}^{2},$
(27)

where $\boldsymbol{\nu}=\left[\boldsymbol{\nu}_{ij}^{\top},i<j\right]^{\top}$
is a vector of Lagrangian multipliers and $\rho$ is a non-negative tuning
parameter. ADMM minimizes the augmented Lagrangian (27) by iteratively
optimizing over blocks of variables $\boldsymbol{\beta}$,
$\boldsymbol{\delta}$, and the dual parameter $\boldsymbol{\nu}$.
Specifically, we start at initial values
$\left(\boldsymbol{\beta}^{0},\boldsymbol{\delta}^{0},\boldsymbol{\nu}^{0}\right)$
and update
$\left(\boldsymbol{\beta}^{s},\boldsymbol{\delta}^{s},\boldsymbol{\nu}^{s}\right)$
at the $s$-th iteration as follows.

  1. Step 1.

Given $\left(\boldsymbol{\delta}^{s},\boldsymbol{\nu}^{s}\right)$, update
$\boldsymbol{\beta}^{s+1}$ by solving $\boldsymbol{\beta}$ from
$\nabla_{\boldsymbol{\beta}}\;{\cal
L}_{AL}\left(\boldsymbol{\beta},\boldsymbol{\delta}^{s},\boldsymbol{\nu}^{s}\right)=0$.

  2. Step 2.

Given $\boldsymbol{\beta}^{s+1}$, update $\boldsymbol{\delta}^{s+1}$ using the
analytical forms, for MCP, with

$\boldsymbol{\delta}_{ij}^{s+1}=\begin{cases}\frac{{\cal
S}\left(\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s},\lambda/\rho\right)}{1-1/(\gamma\rho)}&\text{if
}\|\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s}\|\leq\gamma\lambda\\\
\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s}&\text{if
}\|\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s}\|>\gamma\lambda\end{cases}$

and, for SCAD, with

$\boldsymbol{\delta}_{ij}^{s+1}=\begin{cases}{\cal
S}\left(\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s},\lambda/\rho\right)&\text{if
}\|\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s}\|\leq\lambda+\lambda/\rho\\\
\frac{{\cal
S}\left(\boldsymbol{\beta}_{i}^{s+1}-\beta_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s},\gamma\lambda/((\gamma-1)\rho)\right)}{1=1/((\gamma-1)/\rho)}&\text{if
}\lambda+\lambda/\rho<\|\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s}\|\leq\gamma\lambda\\\
\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s}&\text{if
}\|\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}+\rho^{-1}\boldsymbol{\nu}_{ij}^{s})\|>\gamma\lambda\par\end{cases},$

where ${\cal S}\left(x,c\right)={\rm sign}(x)\left(\left\lvert
x\right\rvert-c\right)_{+}$ is the soft thresholding rule and
$\left(x\right)_{+}=x$ if $x>0$ and $0$ otherwise.

  3. Step 3.

Update the dual parameter
$\boldsymbol{\nu}_{ij}^{s+1}=\boldsymbol{\nu}_{ij}^{s}+\rho(\boldsymbol{\beta}_{i}^{s+1}-\boldsymbol{\beta}_{j}^{s+1}-\boldsymbol{\delta}_{ij}^{s+1})$.

The iteration terminates when the norm of the primal residual is smaller than
some pre-specified small tolerance $\epsilon$, that is, when
$\|\boldsymbol{\beta}_{i}-\boldsymbol{\beta}_{j}-\boldsymbol{\delta}_{ij}\|<\epsilon$.
The algorithmic convergence of ADMM has been established in Boyd et al. (2011)
and Chi and Lange (2015).

##  6 Simulations

In this section, we compare the performance of ACPE and ACPI with their mean-
value-based counterparts (i.e., policy evaluation and iteration) on simulated
data. We generate the initial state variable $X_{i,1}$ from a standard normal
distribution ${\cal N}\left(\boldsymbol{0}_{2},\boldsymbol{I}_{2}\right)$. The
available action set is ${\cal A}=\left\\{0,1\right\\}$. The system evolves
according to

$X_{i,t+1}=\begin{bmatrix}0.75(2A_{i,t}-1)&0\\\
0&0.75(1-2A_{i,t})\end{bmatrix}X_{i,t}+\boldsymbol{\varepsilon}_{i,t},$

where $\boldsymbol{\varepsilon}_{i,t}\overset{i.i.d}{\sim}{\cal
N}\left(\boldsymbol{0},\boldsymbol{I}_{2}/4\right)$. The data-generating
actions are i.i.d. Bernoulli random variables with expectation $0.5$ and are
independent of $\boldsymbol{x}_{i,t}$ for any $t\geq 1$. We consider $K=2$
homogeneous groups. The immediate reward $R_{i,t}$ is defined by

$R_{i,t}=X_{i,t}^{\top}\boldsymbol{b}_{k}-0.25(2A_{i,t}-1),\quad\text{for}\quad\forall
i\in{\cal G}_{k},\;k\in 1,2,$

where $\boldsymbol{b}_{1}=[2,-1]$ and $\boldsymbol{b}_{2}=[-2,1]$. Note that
$\boldsymbol{b}_{k}$ is not the coefficient of Q function. Unless otherwise
specified, we use $\gamma=0.6$ for discount factor, the MCP (15) with
$\eta=1.5$ for penalization throughout all experiments. We present results for
the consistency and asymptotic distribution in policy estimation and solving
for optimal policies.

###  6.1 Coefficients of the value function of a given policy $\pi$

The target policy $\pi$ to be evaluated is specified as

$\pi\left(a|\boldsymbol{x}\right)=\begin{cases}0,&x_{1}>0\text{ and
}x_{2}>0;\\\ 1,&\text{otherwise},\end{cases}$

where $x_{i}$ denotes the $i$-th element of a vector $\boldsymbol{x}$. For
each group, we simulate $N_{k}=100$ trajectories of length $T=10$. Thus, we
have in total $N=200$ trajectories. Recall that we use $Q$ function model (11)
with features arranged in (12). The parameter $\boldsymbol{\beta}^{i}$
contains $M$ blocks of coefficients and each block corresponds to one action
from ${\cal A}=\\{a_{1},\cdots,a_{M}\\}$. In this simulation setting, we have
$\boldsymbol{\beta}^{i}=[\boldsymbol{\beta}_{1}^{i\top},\boldsymbol{\beta}_{2}^{i\top}]^{\top}\in\mathbb{R}^{4}$
where $\boldsymbol{\beta}_{1}^{i}$ and
$\boldsymbol{\beta}_{2}^{i}\in\mathbb{R}^{2}$ are coefficients associated with
action $a_{1}=0$ and $a_{2}=1$, respectively. Figure 2 plots the
$\widehat{}\boldsymbol{\beta}_{1}^{i}$ and
$\widehat{}\boldsymbol{\beta}_{2}^{i}$ for all $200$ trajectories. We further
run K-means algorithms to cluster $\\{\boldsymbol{\beta}^{i}\\}_{i\in[N]}$
into two groups and calculate the cluster centroids
$\widehat{}\boldsymbol{\theta}_{1}^{(k)}=[\widehat{}\boldsymbol{\theta}_{1}^{(k)\top},\widehat{}\boldsymbol{\theta}_{2}^{(k)\top}]^{\top}$
for $k=1,2$. The centroids $\widehat{}\boldsymbol{\theta}_{1}^{(k)}$ and
$\widehat{}\boldsymbol{\theta}_{2}^{(k)}$ for $k=1,2$ are also plotted in the
figure as the red dots. The left and right columns of Figure 2 corresponds to
using $\lambda=0.1$ or $0.05$, respectively, in the MCP (15) penalization.

As shown in Figure 2, the proposed algorithm is able to recover mostly correct
coefficients with an appropriate choice of $\lambda$ even when the group
centroids are relatively close to each other. The tuning parameter $\lambda$
controls the “focus of the group.” Smaller $\lambda$ encourages heterogeneity,
that is, different values for different trajectories, while large $\lambda$
enforce homogeneity. When $\lambda$ is too small, we lose efficiency by
grouping trajectories together. When $\lambda$ is too big, the problem reduces
to homogeneous policy evaluation with the same coefficients and hence loses
the heterogeneity that we are seeking.

(a) $\\{\widehat{}\boldsymbol{\beta}_{1}^{i}\\}_{i\in[N]}$ for action
$a_{1}=0$ with $\lambda=0.05$

(b) $\\{\widehat{}\boldsymbol{\beta}_{1}^{i}\\}_{i\in[N]}$ for action
$a_{1}=0$ with $\lambda=0.1$

(c) $\\{\widehat{}\boldsymbol{\beta}_{2}^{i}\\}_{i\in[N]}$ for action
$a_{2}=1$, $\lambda=0.05$

(d) $\\{\widehat{}\boldsymbol{\beta}_{2}^{i}\\}_{i\in[N]}$ for action
$a_{2}=1$, $\lambda=0.1$

Figure 2: Coefficients of $Q$ functions and the centroids estimated by ACPE.
We use blue and gray colors to denote the true membership of each point. We
use MCP (15) penalization with $\eta=1.5$. Plots (a) and (c) on the left
correspond to the results with $\lambda=0.05$; plots (b) and (d) on the right
correspond to the results with $\lambda=0.1$. We can see that with an
appropriate choice of the tuning parameter $\lambda$, the majority of fitted
coefficients fall in a very small neighborhood of the centroids and can be
hence correctly classified by simple clustering algorithms.

###  6.2 Confidence intervals of the integrated value of a given policy $\pi$

We compare the integrated value estimated by ACPE and classical mean-value
policy evaluation (MVPE). The true integrated value $\mathring{V}^{\pi}_{\cal
R}$ is computed by Monte Carlo approximation. Specifically, we simulate
$N=5,000$ independent trajectories (2,500 trajectories for each group) with
initial states distributed according to ${\cal R}$. For each trajectory, we
simulate $T$ steps to obtain $\left\\{Y_{i,t}^{(k)}\right\\}$ for $1\leq i\leq
N$, $1\leq t\leq T$, and $k=1,2$. For each sub-group, the true value is
approximated by $\mathring{V}_{{\cal R}}^{\pi,(k)}\approx
N_{k}^{-1}\sum_{i=1}^{N_{k}}\sum_{t=1}^{T}\gamma^{t}Y_{i,t}^{(k)}$. For the
MVPE, the estimated $\widehat{V}^{MV}_{\cal R}(\pi)$ and the corresponding CI
are obtained using equation (3.9) and (3.10) in Shi et al. (2020). For the
ACPE, we use MCP (15) with $\eta=1.5$ and $\lambda=0.1$. The estimated
$\widehat{V}^{AC}_{\cal R}(\pi)$ and the corresponding CI are obtained using
results in Theorem 3. We compare the true value with our estimated values with
$n=20,50,100$ and $T=10,30,40$. The coverage probability is plotted in Figure
3 for trajectories from group 1 and 2 respectively in subplot (a) and (b). We
can see that heterogeneous estimation performs consistently better in terms of
coverage probability.

(a) For trajectories from group 1

(b) For trajectories from group 2

Figure 3: Empirical covering probabilities (ECP) of the 95% confidence
interval. ACPE represents our auto-clustered policy evaluation and its
corresponding CI in Theorem 3. MVPE represent the classic mean-value-based
policy evaluation and its corresponding CI are constructed using equation
(3.9) and (3.10) in Shi et al. (2020). Subplot (a) and (b) correspond to the
ECP’s for trajectories from two different groups $k=1$ and $2$, respectively.

###  6.3 Parametric optimal policy

We train 100 steps of Algorithm 2 (ACPI) to obtain two policies, which we
denote $\pi(\widehat{}\boldsymbol{\alpha}^{(1)*}_{ACPI})$ and
$\pi(\widehat{}\boldsymbol{\alpha}^{(2)*}_{ACPI})$, respectively for estimated
group 1 and 2. We follow the same protocol to obtain a policy
$\pi(\widehat{}\boldsymbol{\alpha}^{*}_{MVPI})$ via standard mean-value-based
policy iteration (MVPI).

Firstly, we show how different are the actions given by
$\pi(\widehat{}\boldsymbol{\alpha}^{(1)*})$,
$\pi(\widehat{}\boldsymbol{\alpha}^{(2)*})$, and
$\pi(\widehat{}\boldsymbol{\alpha}^{*}_{MVPI})$ on a specific sample. We take
a $j$-th test sample with $X_{j,1}=\begin{bmatrix}1.2506\\\
0.77477\end{bmatrix}$, $X_{j,2}=\begin{bmatrix}1.0277\\\
-0.52615\end{bmatrix}$, $A_{j,1}=1$, $Y_{j,1}=1.6559$. The optimal actions at
$t=2$, i.e. $A_{j,2}$, suggested by ACPI for group 1 and 2 and by MVPI are
given in Table 1. This trajectory is in fact in group 1, which is correctly
clustered by ACPE. Thus, the optimal action should be chosen by following
$\pi\left(\widehat{}\boldsymbol{\alpha}^{(1)*}_{ACPI}\right)$ which prefers
action $A_{j,2}=1$. If we ignore heterogeneity and naively adopt MVPI, we will
wrongly weight both actions evenly, yielding a lower expected outcome.

Estimated optimal policy | $\Pr(A=0)$ | $\Pr(A=1)$
---|---|---
$\pi\left(\widehat{}\boldsymbol{\alpha}^{(1)*}_{ACPI}\right)$ | 0.14 | 0.86
$\pi\left(\widehat{}\boldsymbol{\alpha}^{(2)*}_{ACPI}\right)$ | 0.75 | 0.25
$\pi\left(\widehat{}\boldsymbol{\alpha}_{MVPI}\right)$ | 0.58 | 0.42
Table 1: The optimal action given by ACPI for group 1 and 2 and by MVPI on a
random trajectory. We use $\eta=1.5$ and $\lambda=0.1$ in ACPI.

We further generate $N=500$ samples distributed according to ${\cal
N}\left(\boldsymbol{0},\boldsymbol{I}_{2}\right)$ (250 samples for each
group). We estimate the true value of all three policies using Monte Carlo
simulation $\mathring{V}_{{\cal R}}\left(\pi\right)\approx
N^{-1}\sum_{i=1}^{N}\sum_{t=0}^{T-1}\gamma^{t}\boldsymbol{Y}_{i,t}$ at $T=50$.
We repeat 10 times and obtain the average values of all three policies on
different groups. The results are presented in Table 2. It shows that policies
produced by ACPI are indeed better in terms of expected discounted total sum
of rewards.

Estimated optimal policy | Integrated value $\mathring{V}_{\cal R}(\pi)$
---|---
True group 1 | True group 2
$\pi\left(\widehat{}\boldsymbol{\alpha}^{(1)*}_{ACPI}\right)$ | 0.119 | -
$\pi\left(\widehat{}\boldsymbol{\alpha}^{(2)*}_{ACPI}\right)$ | - | 0.104
$\pi\left(\widehat{}\boldsymbol{\alpha}_{MVPI}\right)$ | 0.002 | 0.002
Table 2: The integrated values of the optimal policies given by ACPI for group
1 and 2 and by MVPI. We use $\eta=1.5$ and $\lambda=0.1$ in ACPI.

##  7 Real data application: MIMIC-III

In this section, we illustrate the advantages of the ACPI on the Medical
Information Mart for Intensive Care version III (MIMIC-III) Database (Johnson
et al., 2016), which is a freely available source of de-identified critical
care data from 2001–2012 in six ICUs at a Boston teaching hospital.

We consider a cohort of sepsis patients, following the same data processing
procedure as in Komorowski et al. (2018). Each patient in the cohort is
characterized by a set of 47 variables, including demographics, Elixhauser
premorbid status, vital signs, and laboratory values. Since we do not consider
high-dimensional issue in the present paper, we conduct a dimension reduction
using principal component analysis and choose top $10$ principal components as
our state features, which explain around 97% of the total variance.. Patients’
data were coded as multidimensional discrete time series
$X_{i,t}\in\mathbb{R}^{10}$ for $1\leq i\leq N$ and $1\leq t\leq T_{i}$ with
four-hour time steps. The actions of interests are the total volume of
intravenous (IV) fluids and maximum dose of vasopressors administrated over
each four-hour period. We discretize two action variables into three levels,
respectively. Our low corresponds to 1 - 2, medium corresponds to 3 and high
corresponds to 4 - 5 in Komorowski et al. (2018). The combination of the two
drugs makes $M=3\times 3=9$ possible actions in total. In the final processed
dataset, we sampled $1000$ unique adult ICU admissions, corresponding to
unique trajectories to be fed into our algorithms. The observation length
$T_{i}$ varies across trajectories.

The reward signal $R_{i,t}$ is important and is crafted carefully in real
applications. For the final reward, we follow Komorowski et al. (2018) and use
hospital mortality or 90-day mortality. Specifically, when a patient survived,
a positive reward was released at the end of each patient’s trajectory; a
negative reward was issued if the patient died. For the intermediate rewards,
we follow Prasad et al. (2017) and associate reward to the health measurement
of a patient. The detailed description of the data pre-processing is presented
in Section E of the supplemental material. More information about the dataset
can also be found in Komorowski et al. (2018) and Prasad et al. (2017).

The final processed data contains 12,987 tuples
$\left\\{\left(X_{i,t},A_{i,t},R_{i,t},X_{i,t+1}\right)\right\\}$ with
$i\in[1000]$ and $t\in[T_{i}]$. We first evaluate the value of a random policy
$\pi(a|\boldsymbol{x})=1/9$ for $a\in[9]$ on the data set using ACPE and MVPE.
Note both ACPE and MVPV try to minimize the Bellman residual, which are the
sample counterpart of the left hand side of (7), while ACPE accounts for
heterogeneity over different trajectories. We fit linear $Q$ functions and
plot the histograms of the sample Bellman residuals of MVPE and ACPE,
respectively, in (a) and (b) of Figure 4. As demonstrated in Figure 4 (a), we
see that the Bellman residual plot of MVPE is bi-modal. This hints that there
are at least two different coefficient groups in this population. Figure 4 (b)
is the residual plot of ACPE. We can see that the residual plot for our
algorithm is well-shaped and performs better in terms of fitting error. We
perform a principal component analysis on our coefficient matrix (averaged
across actions) as shown in Figure 5. We can observe that there are two clear
separate clusters within this coefficient matrix, which is a confirmation of
the bi-modality of the residuals and a sign that simple $k$-means clustering
would work well when we cluster the $\boldsymbol{\beta}$’s into groups.

We further use ACPI to obtain optimal policies
$\pi(\widehat{}\boldsymbol{\alpha}^{(1)*}_{ACPI})$ and
$\pi(\widehat{}\boldsymbol{\alpha}^{(2)*}_{ACPI})$, respectively for estimated
group 1 and 2. We follow the same protocol to obtain a optimal policy
$\pi(\widehat{}\boldsymbol{\alpha}^{*}_{MVPI})$ via standard mean-value-based
policy iteration (MVPI). Table 3 summarizes the estimated values of
$\pi(\widehat{}\boldsymbol{\alpha}^{(1)*}_{ACPI})$,
$\pi(\widehat{}\boldsymbol{\alpha}^{(2)*}_{ACPI})$, and
$\pi(\widehat{}\boldsymbol{\alpha}^{*}_{MVPI})$ with different values of
discount factor $\gamma$. These policies are evaluated on two groups
separately on the two sub-groups returned by ACPI. We can see that the policy
obtained by ACPI outperforms the one returned by MVPI by around 30 % on both
group 1 and group 2.

We then demonstrate two sample patients in Table 4 and 5 to show the
effectiveness of estimated policies. Patient 1 eventually died in hospital,
and this policy is evaluated at two time points before his/her death. At this
time point, many of his/her indices are extremely abnormal. For example,
his/her Glasgow Coma Scale (GCS) decreased sharply from 15 to 6. This is a
critical sign of urgent treatment. Indeed, policy trained under heterogeneous
estimation ACPI recommends high intravenous volume and high/medium
vasopressors, while policy trained under homogeneous estimation MVPI still
puts significant weights on medium intravenous volume and medium vasopressors.
This conservative policy recommended by homogeneous estimation would be
dangerous. Patient 2 survives and all his/her indices are fairly normal at the
evaluated time point, except for his/her relatively high respiratory rate
(RR). He/she is also one of the youngest patients in the sample. Thus, the
physician decides to inject little to no amounts of both intravenous and
vasopressors. The heterogeneously estimated policy by ACPI recommends the same
treatment as the physician does. However, the homogeneously estimated policy
by MVPI also suggests a medium volume of intravenous, which is unnecessary for
such a patient.

(a) Residual distribution of simple linear regression is clearly bi-model,
with two peaks at around -4 and 4.

(b) Residual of our algorithm is bell-shaped, with one peak at around 0. It
also performs better in terms of fitting error.

Figure 4: Bellman residuals of MVPE and ACPE of the random policy on the MIMIC-III data set. Figure 5: Plot of the first two principal components of the coefficient matrix (averaged across actions) produced by ACPE of the random policy. We can see there is a clear separation between two groups. Estimated optimal policy | Estimated Group 1 | Estimated Group 2
---|---|---
$\gamma=0.5$ | $\gamma=0.7$ | $\gamma=0.5$ | $\gamma=0.7$
$\pi\left(\widehat{}\boldsymbol{\alpha}^{(1)*}_{ACPI}\right)$ | -2.2140 | -3.1851 | - | -
$\pi\left(\widehat{}\boldsymbol{\alpha}^{(2)*}_{ACPI}\right)$ | - | - | -1.4433 | -1.7793
$\pi\left(\widehat{}\boldsymbol{\alpha}^{*}_{MVPI}\right)$ | -3.0964 | -4.5735 | -2.0634 | -2.6487
ACPI Improvement | 28.50% | 30.36% | 30.05% | 32.82%
Table 3: Estimated values on two subgroups of the heterogeneously estimated policies by ACPI and the homogeneously estimated policy by MVPI. The groups are estimated by ACPI.  Optimal policy | ACPI | MVPI | Expert
---|---|---|---
iv low, vaso low | $<0.0001$ | $<0.0001$ | 0
iv low, vaso med | $<0.0001$ | 0.071 | 0
iv low, vaso high | $<0.0001$ | <0.0001 | 0
iv med, vaso low | $<0.0001$ | $0.10$ | 0
iv med, vaso med | $<0.0001$ | 0.14 | 0
iv med, vaso high | $0.042$ | 0.21 | 0
iv high, vaso low | $0.23$ | $<0.0001$ | 0
iv high, vaso med | $0.36$ | 0.26 | 0
iv high, vaso high | 0.37 | 0.22 | 1
Table 4: Recommended action by policy estimated by ACPI, MVPI and the expert for a sample patient from group 1. Optimal policy | ACPI | MVPI | Expert
---|---|---|---
iv low, vaso low | $0.72$ | $0.36$ | 1
iv low, vaso med | $0.18$ | 0.018 | 0
iv low, vaso high | $<0.0001$ | <0.0001 | 0
iv med, vaso low | $<0.0001$ | $0.43$ | 0
iv med, vaso med | $<0.0001$ | <0.0001 | 0
iv med, vaso high | $<0.0001$ | 0.11 | 0
iv high, vaso low | $0.09$ | $<0.0001$ | 0
iv high, vaso med | $<0.001$ | 0.075 | 0
iv high, vaso high | <0.001 | <0.0001 | 0
Table 5: Recommended action by policy estimated by ACPI, MVPI and the expert
for a sample patient from group 2.

##  8 Summary

Classical RL methods model and optimize the (action-)value function defined as
the expectation of total return. Direct applications of such mean-value based
RL to large-scale datasets may generate misleading results because of data
heterogeneity. In this paper, we go beyond classical mean-value based RL to
allow for heterogeneity which is characterized by different values across sub-
populations. We proposed ACPE and ACPI for both the policy evaluation and
control. We establish convergence rates and construct confidence intervals
(CIs) for the estimators obtained by the ACPE and ACPI. Our theoretical
findings are validated on synthetic and real datasets. Particularly, the
experiments on the standard MIMIC-III dataset shows evidence of data
heterogeneity and confirms the advantage of our new method.

Statistical analysis of policy evaluation and optimal control under the
framework of RL and MDP have great potential to facilitate dynamic decision
making in a variety of real applications. We have demonstrated the importance
of data heterogeneity when combining RL and large-scale datasets. There are
several interesting directions for future research in this area. First, our
method is based on the Bellman consistency equation to approximate the
$Q^{\pi}$ function. Its performance is guaranteed under the assumption of
bounded state distribution shift, which is caused by the discrepancy between
the behavior policy and the target policy. However, this assumption may not
hold generally in real applications. It is of great interest to investigate
various ways to relax such an assumption. Second, the theoretical results in
this paper are developed for offline batch estimation. Developing online RL
estimation with heterogeneity is an interesting topic for future research.
Lastly, we estimate the optimal policy by a variation of $Q$ learning. It
would also be worthwhile to investigate the ways to incorporate heterogeneity
in other RL methods such as policy gradient.

## References

  * Barth-Maron et al. (2018) Barth-Maron, G., M. W. Hoffman, D. Budden, W. Dabney, D. Horgan, D. TB, A. Muldal, N. Heess, and T. Lillicrap (2018).  Distributed distributional deterministic policy gradients.  In Proceedings of the International Conference on Learning Representations (ICLR).
  * Bellemare et al. (2017) Bellemare, M. G., W. Dabney, and R. Munos (2017).  A distributional perspective on reinforcement learning.  In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pp. 449–458. JMLR. org.
  * Bondell and Reich (2008) Bondell, H. D. and B. J. Reich (2008).  Simultaneous regression shrinkage, variable selection, and supervised clustering of predictors with oscar.  Biometrics 64(1), 115–123.
  * Boyd et al. (2011) Boyd, S., N. Parikh, and E. Chu (2011).  Distributed optimization and statistical learning via the alternating direction method of multipliers.  Now Publishers Inc.
  * Burman et al. (1989) Burman, P., K.-W. Chen, et al. (1989).  Nonparametric estimation of a regression function.  The Annals of Statistics 17(4), 1567–1596.
  * Chen and Christensen (2015) Chen, X. and T. M. Christensen (2015).  Optimal uniform convergence rates and asymptotic normality for series estimators under weak dependence and weak conditions.  Journal of Econometrics 188(2), 447–465.
  * Chi and Lange (2015) Chi, E. C. and K. Lange (2015).  Splitting methods for convex clustering.  Journal of Computational and Graphical Statistics 24(4), 994–1013.
  * Dabney et al. (2018) Dabney, W., M. Rowland, M. G. Bellemare, and R. Munos (2018).  Distributional reinforcement learning with quantile regression.  In Thirty-Second AAAI Conference on Artificial Intelligence.
  * Ertefaie and Strawderman (2018) Ertefaie, A. and R. L. Strawderman (2018).  Constructing dynamic treatment regimes over indefinite time horizons.  Biometrika 105(4), 963–977.
  * Fan and Li (2001) Fan, J. and R. Li (2001).  Variable selection via nonconcave penalized likelihood and its oracle properties.  Journal of the American statistical Association 96(456), 1348–1360.
  * Gong et al. (2013) Gong, P., C. Zhang, Z. Lu, J. Z. Huang, and J. Ye (2013).  A general iterative shrinkage and thresholding algorithm for non-convex regularized optimization problems.  In Proceedings of the 30th International Conference on Machine Learning, ICML’13. JMLR.org.
  * Hocking et al. (2011) Hocking, T. D., A. Joulin, F. Bach, and J.-P. Vert (2011).  Clusterpath: An algorithm for clustering using convex fusion penalties.  In 28th International Conference on Machine Learning, pp. 1.
  * Huang et al. (1998) Huang, J. Z. et al. (1998).  Projection estimation in multiple regression with application to functional anova models.  The Annals of Statistics 26(1), 242–272.
  * Jiang and Li (2016) Jiang, N. and L. Li (2016).  Doubly robust off-policy value evaluation for reinforcement learning.  In International Conference on Machine Learning, pp. 652–661.
  * Johnson et al. (2016) Johnson, A. E., T. J. Pollard, L. Shen, H. L. Li-wei, M. Feng, M. Ghassemi, B. Moody, P. Szolovits, L. A. Celi, and R. G. Mark (2016).  MIMIC-III, a freely accessible critical care database.  Scientific Data 3, 160035.
  * Ke et al. (2013) Ke, T., J. Fan, and Y. Wu (2013).  Homogeneity in regression.  arXiv preprint arXiv:1303.7409.
  * Komorowski et al. (2018) Komorowski, M., L. A. Celi, O. Badawi, A. C. Gordon, and A. A. Faisal (2018).  The artificial intelligence clinician learns optimal treatment strategies for sepsis in intensive care.  Nature Medicine 24(11), 1716–1720.
  * Luckett et al. (2019) Luckett, D. J., E. B. Laber, A. R. Kahkoska, D. M. Maahs, E. Mayer-Davis, and M. R. Kosorok (2019).  Estimating dynamic treatment regimes in mobile health using v-learning.  Journal of the American Statistical Association, 1–34.
  * Luedtke and Van Der Laan (2016) Luedtke, A. R. and M. J. Van Der Laan (2016).  Statistical inference for the mean outcome under a possibly non-unique optimal treatment strategy.  Annals of Statistics 44(2), 713.
  * Ma and Huang (2017) Ma, S. and J. Huang (2017).  A concave pairwise fusion approach to subgroup analysis.  Journal of the American Statistical Association 112(517), 410–423.
  * Mandel (2017) Mandel, T. S. (2017).  Better Education Through Improved Reinforcement Learning.  Ph. D. thesis.
  * Murphy (2003) Murphy, S. A. (2003).  Optimal dynamic treatment regimes.  Journal of the Royal Statistical Society: Series B (Statistical Methodology) 65(2), 331–355.
  * Pan et al. (2013) Pan, W., X. Shen, and B. Liu (2013).  Cluster analysis: unsupervised learning via supervised learning with a non-convex penalty.  The Journal of Machine Learning Research 14(1), 1865–1889.
  * Prasad et al. (2017) Prasad, N., L. F. Cheng, C. Chivers, M. Draugelis, and B. E. Engelhardt (2017).  A reinforcement learning approach to weaning of mechanical ventilation in intensive care units.  In 33rd Conference on Uncertainty in Artificial Intelligence, UAI 2017.
  * Rafferty et al. (2016) Rafferty, A. N., E. Brunskill, T. L. Griffiths, and P. Shafto (2016).  Faster teaching via pomdp planning.  Cognitive Science 40(6), 1290–1332.
  * Reddy et al. (2017) Reddy, S., S. Levine, and A. Dragan (2017).  Accelerating human learning with deep reinforcement learning.  In NIPS Workshop: Teaching Machines, Robots, and Humans.
  * Robins (2004) Robins, J. M. (2004).  Optimal structural nested models for optimal sequential decisions.  In Proceedings of the Second Seattle Symposium in Biostatistics, pp. 189–326. Springer.
  * Rowland et al. (2018) Rowland, M., M. Bellemare, W. Dabney, R. Munos, and Y. W. Teh (2018).  An analysis of categorical distributional reinforcement learning.  In International Conference on Artificial Intelligence and Statistics, pp. 29–37.
  * Rowland et al. (2019) Rowland, M., R. Dadashi, S. Kumar, R. Munos, M. G. Bellemare, and W. Dabney (2019).  Statistics and samples in distributional reinforcement learning.  In International Conference on Machine Learning, pp. 5528–5536.
  * Schulte et al. (2014) Schulte, P. J., A. A. Tsiatis, E. B. Laber, and M. Davidian (2014).  Q-and Q-learning methods for estimating optimal dynamic treatment regimes.  Statistical Science 29(4), 640.
  * Shen and Huang (2010) Shen, X. and H.-C. Huang (2010).  Grouping pursuit through a regularization solution surface.  Journal of the American Statistical Association 105(490), 727–739.
  * Shi et al. (2020) Shi, C., S. Zhang, W. Lu, and R. Song (2020).  Statistical inference of the value function for reinforcement learning in infinite horizon settings.  arXiv preprint arXiv:2001.04515.
  * Sobel (1982) Sobel, M. J. (1982).  The variance of discounted Markov decision processes.  Journal of Applied Probability 19(4), 794–802.
  * Song et al. (2015) Song, R., W. Wang, D. Zeng, and M. R. Kosorok (2015).  Penalized Q-learning for dynamic treatment regimens.  Statistica Sinica 25(3), 901.
  * Sutton and Barto (2018) Sutton, R. S. and A. G. Barto (2018).  Reinforcement Learning: An Introduction.  MIT Press.
  * Thomas et al. (2015) Thomas, P. S., G. Theocharous, and M. Ghavamzadeh (2015).  High-confidence off-policy evaluation.  In Twenty-Ninth AAAI Conference on Artificial Intelligence.
  * Tibshirani et al. (2005) Tibshirani, R., M. Saunders, S. Rosset, J. Zhu, and K. Knight (2005).  Sparsity and smoothness via the fused lasso.  Journal of the Royal Statistical Society: Series B (Statistical Methodology) 67(1), 91–108.
  * Tseng (1991) Tseng, P. (1991).  Applications of a splitting algorithm to decomposition in convex programming and variational inequalities.  SIAM Journal on Control and Optimization 29(1), 119–138.
  * Yang et al. (2019) Yang, D., L. Zhao, Z. Lin, T. Qin, J. Bian, and T.-Y. Liu (2019).  Fully parameterized quantile function for distributional reinforcement learning.  In Advances in Neural Information Processing Systems, pp. 6190–6199.
  * Zhang et al. (2010) Zhang, C.-H. et al. (2010).  Nearly unbiased variable selection under minimax concave penalty.  The Annals of Statistics 38(2), 894–942.
  * Zhou et al. (2017) Zhou, Z., X. Li, and R. N. Zare (2017).  Optimizing chemical reactions with deep reinforcement learning.  ACS Central Science 3(12), 1337–1344.
  * Zhu et al. (2019) Zhu, W., D. Zeng, and R. Song (2019).  Proper inference for value function in high-dimensional Q-learning for dynamic treatment regimes.  Journal of the American Statistical Association 114(527), 1404–1417.

Supplementary materials for “Heterogeneous Reinforcement Learning with Offline
Data: Estimation and Inference”

Elynn Y. Chen, Rui Song, and Michael I. Jordan

##  Appendix A Proofs: Oracle Estimators

The solution of (21) can be obtained by solving the estimating equation:

$\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=1}^{T}\left((\widetilde{\boldsymbol{\Lambda}}^{i})^{\top}Z_{i,t}R_{i,t}-(\widetilde{\boldsymbol{\Lambda}}^{i})^{\top}Z_{i,t}\left(Z_{i,t}-\gamma
U^{\pi}_{i,t+1}\right)^{\top}\widetilde{\boldsymbol{\Lambda}}^{i}\boldsymbol{\theta}^{\pi}\right)=0.$
(28)

The resulting oracle estimator of the group coefficients obtained by (28) has
the following decomposition:

$\left(\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right)=(\widetilde{}\boldsymbol{\Sigma}^{\pi})^{-1}\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=1}^{T}\left((\widetilde{\boldsymbol{\Lambda}}^{i})^{\top}Z_{i,t}(\varepsilon_{1,i,t}^{\pi}+\varepsilon_{2,i,t}^{\pi})\right)=(\widetilde{}\boldsymbol{\Sigma}^{\pi})^{-1}\widetilde{}\boldsymbol{\zeta}_{1}^{\pi}+(\widetilde{}\boldsymbol{\Sigma}^{\pi})^{-1}\widetilde{}\boldsymbol{\zeta}_{2}^{\pi},$
(29)

where

$\displaystyle\widetilde{}\boldsymbol{\Sigma}^{\pi}$
$\displaystyle=\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=1}^{T}(\widetilde{\boldsymbol{\Lambda}}^{i})^{\top}Z_{i,t}\left(Z_{i,t}-\gamma
U^{\pi}_{i,t+1}\right)^{\top}\widetilde{\boldsymbol{\Lambda}}^{i},$ (30)
$\displaystyle\widetilde{}\boldsymbol{\zeta}_{1}^{\pi}$
$\displaystyle=\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=1}^{T}(\widetilde{\boldsymbol{\Lambda}}^{i})^{\top}Z_{i,t}\varepsilon_{1,i,t},$
(31) $\displaystyle\widetilde{}\boldsymbol{\zeta}_{2}^{\pi}$
$\displaystyle=\sum_{i=1}^{N}\sum_{t=1}^{T}(\widetilde{\boldsymbol{\Lambda}}^{i})^{\top}Z_{i,t}\varepsilon_{2,i,t},$
(32) $\displaystyle\varepsilon_{1,i,t}^{\pi}$
$\displaystyle=R_{i,t}+\gamma\cdot\sum_{a\in{\cal
A}}Q(\pi,X_{i,t+1},a)\pi(a\mid X_{i,t+1})-Q(\pi,X_{i,t},A_{i,t})$ (33)
$\displaystyle\varepsilon_{2,i,t}^{\pi}$
$\displaystyle=\gamma\cdot\sum_{a\in{\cal A}}e(\pi,X_{i,t+1},a)\pi(a\mid
X_{i,t+1})-\sum_{a\in{\cal A}}e(\pi,X_{i,t},a)\mathbbm{1}(A_{i,t}=a),$ (34)
$\displaystyle e(\pi,X_{i,t},a)$
$\displaystyle=Q(\pi,X_{i,t},a)-\boldsymbol{\phi}(X_{i,t})^{\top}\widetilde{\boldsymbol{\Lambda}}^{i}\mathring{\boldsymbol{\theta}}^{\pi}.$
(35)

In the proofs, we will omit the superscript $\pi$ in
$\widetilde{}\boldsymbol{\Sigma}^{\pi}$,
$\widetilde{}\boldsymbol{\zeta}_{1}^{\pi}$,
$\widetilde{}\boldsymbol{\zeta}_{2}^{\pi}$, $\varepsilon_{1,i,t}^{\pi}$,
$\varepsilon_{2,i,t}^{\pi}$, $\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)}$,
$\widetilde{}\boldsymbol{\zeta}_{1}^{\pi,(k)}$, and
$\widetilde{}\boldsymbol{\zeta}_{2}^{\pi,(k)}$ for brevity.

######  Proposition 1 (Oracle estimator $\ell_{2}$ convergence).

Suppose Assumption 2 – 5 hold. $N_{min}=\underset{1\leq k\leq K}{\min}N_{k}$
and $N_{max}=\underset{1\leq k\leq K}{\max}N_{k}$. If $K={\rm
o}\left(N_{\min}T\right)$, $J\ll\sqrt{N_{\min}T}/\log(N_{\min}T)$,
$J^{\kappa/p}\gg\sqrt{N_{\max}T}$, as either $N_{\min}\rightarrow\infty$ or
$T\rightarrow\infty$,

$\left\lVert\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right\rVert_{2}={\mathcal{O}_{p}}\left(J^{1/2}(TN_{\min})^{-1/2}\right)+{\mathcal{O}_{p}}\left(J^{-\kappa/p}\right).$

If we have balanced groups, that is, $N_{\min}\asymp N_{\max}\asymp N$, then

$\left\lVert\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right\rVert_{2}={\mathcal{O}_{p}}\left(J^{1/2}(TN)^{-1/2}\right)+{\mathcal{O}_{p}}\left(J^{-\kappa/p}\right).$

###### Proof.

Recall that
$\widetilde{\boldsymbol{\Lambda}}^{i}=\boldsymbol{\Lambda}^{i}(\boldsymbol{W}\otimes\boldsymbol{I}_{JM})$.
Let $P^{(k)}=N^{(k)}/N$, $\widetilde{}\boldsymbol{\Sigma}^{(k)}$ and
$\widetilde{}\boldsymbol{\zeta}_{i}^{(k)}$ are defined in (50) and (52),
respectively. We have

$\widetilde{}\boldsymbol{\Sigma}={\rm
bdiag}\left(P^{(k)}\widetilde{}\boldsymbol{\Sigma}^{(k)}\right)_{1\leq k\leq
K},\quad\text{and}\quad\widetilde{}\boldsymbol{\zeta}_{i}=\begin{bmatrix}P^{(1)}(\widetilde{}\boldsymbol{\zeta}_{i}^{(1)})^{\top}&\cdots&P^{(K)}(\widetilde{}\boldsymbol{\zeta}_{i}^{(K)})^{\top}\end{bmatrix}^{\top},\quad\text{for
}i=1,2.$ (36)

Then, by decomposition (29) and Proposition 3, we have

$\left\lVert\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right\rVert_{2}=\left(\sum_{k=1}^{K}\left\lVert(\widetilde{}\boldsymbol{\Sigma}^{(k)})^{-1}\widetilde{}\boldsymbol{\zeta}_{1}^{(k)}+(\widetilde{}\boldsymbol{\Sigma}^{(k)})^{-1}\widetilde{}\boldsymbol{\zeta}_{2}^{(k)}\right\rVert_{2}^{2}\right)^{1/2}={\mathcal{O}_{p}}\left(\sqrt{\frac{KJ}{TN_{min}}}\right)+{\mathcal{O}_{p}}\left(KJ^{-\kappa/p}\right).$

The desired result follows by noticing that $K$ is fixed and $P^{(k)}$ is
bounded away from zero under our assumption. ∎

######  Theorem 5 (Asymptotic normality of oracle coefficient estimator).

Suppose Assumption 2 – 5 hold. Let $N_{min}=\underset{1\leq k\leq
K}{\min}N_{k}$ and $N_{max}=\underset{1\leq k\leq K}{\max}N_{k}$. If $K={\rm
o}\left(N_{\min}T\right)$, $J={\rm
o}\left(\sqrt{N_{\min}T}/\log(N_{\min}T)\right)$. For any
$\boldsymbol{\nu}\in\mathbb{R}^{JMK}$ satisfying
$\left\lVert\boldsymbol{\nu}\right\rVert_{2}=\mathcal{O}\left(1\right)$,
$\left\lVert\boldsymbol{\nu}\right\rVert_{2}^{-1}=\mathcal{O}\left(1\right)$
and
$J^{-\kappa/p}\ll\left(N_{\min}T\left(1+\left\lVert\boldsymbol{\nu}\right\rVert_{2}^{-2}\right)\right)^{-1/2}$,
we have, as either $N_{\min}\rightarrow\infty$ or $T\rightarrow\infty$,

$\sqrt{NT}\cdot\widetilde{\sigma}_{\theta}\left(\pi,\boldsymbol{\nu}\right)^{-1}\cdot\boldsymbol{\nu}^{\top}\left(\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right)\stackrel{{\scriptstyle{\cal
D}}}{{\longrightarrow}}{\cal N}\left(0,1\right),$

where
$\widetilde{\sigma}_{\theta}\left(\pi,\boldsymbol{\nu}\right)=\sqrt{\boldsymbol{\nu}^{\top}(\widetilde{}\boldsymbol{\Sigma}^{\pi})^{-1}\widetilde{}\boldsymbol{\Omega}^{\pi}(\widetilde{}\boldsymbol{\Sigma}^{\pi\top})^{-1}\boldsymbol{\nu}}$,

$\displaystyle\widetilde{}\boldsymbol{\Sigma}^{\pi}$
$\displaystyle=\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=1}^{T}(\widetilde{}\boldsymbol{\Lambda}^{i})^{\top}Z_{i,t}\left(Z_{i,t}-\gamma
U_{i,t+1}^{\pi}\right)^{\top}\widetilde{}\boldsymbol{\Lambda}^{i},$ (37)
$\displaystyle\widetilde{}\boldsymbol{\Omega}^{\pi}$
$\displaystyle=\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=1}^{T}(\widetilde{}\boldsymbol{\Lambda}^{i})^{\top}Z_{i,t}Z_{i,t}^{\top}\widetilde{}\boldsymbol{\Lambda}^{i}\left(R_{i,t}-\left(Z_{i,t}-\gamma
U_{i,t+1}^{\pi}\right)^{\top}\widetilde{}\boldsymbol{\Lambda}^{i}\widetilde{}\boldsymbol{\theta}^{\pi}\right)^{2}.$
(38)

###### Proof.

Suppose that
$J^{-\kappa/p}\ll\left(N_{\min}T\left(1+\left\lVert\boldsymbol{\nu}\right\rVert_{2}^{-2}\right)\right)^{-1/2}$.
Combining Proposition 1 and the fact that
$J^{-\kappa/p}\ll\left(N_{\min}T\right)^{-1/2}$, we have, for any
$\boldsymbol{\nu}=\left[\boldsymbol{\nu}^{(1)},\cdots,\boldsymbol{\nu}^{(k)}\right]\in\mathbb{R}^{JMK}$,
$\boldsymbol{\nu}^{(k)}\in\mathbb{R}^{JM}$ for $k\in[K]$, that

$\displaystyle\sqrt{NT}\cdot\boldsymbol{\nu}^{\top}\left(\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right)$
$\displaystyle=\sqrt{NT}\cdot\boldsymbol{\nu}^{\top}\widetilde{}\boldsymbol{\Sigma}^{-1}\widetilde{}\boldsymbol{\zeta}_{1}+{\rm
o_{p}}\left(1\right)$
$\displaystyle=\sum_{k\in[K]}\sqrt{N_{k}T/P^{(k)}}\cdot(\boldsymbol{\nu}^{(k)})^{\top}(\widetilde{}\boldsymbol{\Sigma}^{(k)})^{-1}\widetilde{}\boldsymbol{\zeta}_{1}^{(k)}+{\rm
o_{p}}\left(1\right)$ $\displaystyle\stackrel{{\scriptstyle{\cal
D}}}{{\longrightarrow}}{\cal
N}\left(0,\sum_{k\in[K]}(\boldsymbol{\nu}^{(k)})^{\top}(\widetilde{}\boldsymbol{\Sigma}^{(k)})^{-1}\widetilde{}\boldsymbol{\Omega}^{(k)}(\widetilde{\boldsymbol{\Sigma}}^{(k)^{\top}})^{-1}\boldsymbol{\nu}^{(k)}/P^{(k)}\right),$

where the second equality follows from (36) and the distribution is obtained
by applying Theorem 7. Note that $\widetilde{}\boldsymbol{\Omega}^{\pi}={\rm
bdiag}\left(P^{(k)}\widetilde{}\boldsymbol{\Omega}^{\pi,(k)}\right)$ where
$\widetilde{}\boldsymbol{\Omega}^{\pi,(k)}$ is defined in (56). Combining this
with (36), we have for any $\boldsymbol{\nu}\in\mathbb{R}^{JMK}$ satisfying
$\left\lVert\boldsymbol{\nu}\right\rVert_{2}=\mathcal{O}\left(1\right)$ and
$\left\lVert\boldsymbol{\nu}\right\rVert_{2}^{-1}=\mathcal{O}\left(1\right)$,

$\frac{\sqrt{NT}\boldsymbol{\nu}^{\top}\left(\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right)}{\sqrt{\boldsymbol{\nu}^{\top}\widetilde{}\boldsymbol{\Sigma}^{-1}\widetilde{}\boldsymbol{\Omega}(\widetilde{}\boldsymbol{\Sigma}^{\top})^{-1}\boldsymbol{\nu}}}\stackrel{{\scriptstyle{\cal
D}}}{{\longrightarrow}}{\cal N}\left(0,1\right).$

Under the condition in Theorem 5, we can show that
$\widetilde{\sigma}_{\theta}\left(\pi,\boldsymbol{\nu}\right)$ converges
almost surely to
$\sigma_{\theta}\left(\pi,\boldsymbol{\nu}\right)=\sqrt{\boldsymbol{\nu}^{\top}(\boldsymbol{\Sigma}^{\pi})^{-1}\boldsymbol{\Omega}^{\pi}(\boldsymbol{\Sigma}^{\pi\,\top})^{-1}\boldsymbol{\nu}}$,
where
$\boldsymbol{\Omega}^{\pi}=\mathbb{E}[\widetilde{}\boldsymbol{\Omega}^{\pi}]$
and
$\boldsymbol{\Sigma}^{\pi}=\mathbb{E}[\widetilde{}\boldsymbol{\Sigma}^{\pi}]$.
∎

Applying Theorem 5 to the value function estimators defined in (19), we
establish the asymptotic distribution for state and action value functions
$V(\pi,\boldsymbol{x})$ and $Q(\pi,\boldsymbol{x},a)$ and the asymptotic
distribution for integrated value $V_{\cal R}(\pi)$ for a given distribution
${\cal R}$ in the next two corollaries.

######  Corollary 1 (Asymptotic normality of oracle value estimator.).

Suppose Assumption 2 – 5 hold. Let $N_{min}=\underset{1\leq k\leq
K}{\min}N_{k}$ and $N_{max}=\underset{1\leq k\leq K}{\max}N_{k}$. If $K={\rm
o}\left(N_{\min}T\right)$, $J={\rm
o}\left(\sqrt{N_{\min}T}/\log(N_{\min}T)\right)$, and
$J^{-\kappa/p}\ll\left(N_{\max}T\left(1+\left\lVert\boldsymbol{u}(\pi,\boldsymbol{x})\right\rVert_{2}^{-2}\right)\right)^{-1/2}$,
we have, for any $\boldsymbol{x}\in{\cal X}$ and $a\in{\cal A}$, that as
either $N_{\min}\rightarrow\infty$ or $T\rightarrow\infty$,

$\displaystyle\sqrt{NT}\,\widetilde{\sigma}_{v}^{(k)}\left(\pi,\boldsymbol{x}\right)^{-1}\left(\widetilde{V}^{(k)}\left(\pi,\boldsymbol{x}\right)-V^{(k)}\left(\pi,\boldsymbol{x}\right)\right)\stackrel{{\scriptstyle{\cal
D}}}{{\longrightarrow}}{\cal N}\left(0,1\right),$
$\displaystyle\sqrt{NT}\,\widetilde{\sigma}_{q}^{(k)}\left(\pi,\boldsymbol{x},a\right)^{-1}\left(\widetilde{Q}^{(k)}\left(\pi,\boldsymbol{x},a\right)-Q^{(k)}\left(\pi,\boldsymbol{x},a\right)\right)\stackrel{{\scriptstyle{\cal
D}}}{{\longrightarrow}}{\cal N}\left(0,1\right),$

for any $k\in[K]$, where

$\displaystyle\widetilde{\sigma}_{v}^{(k)}\left(\pi,\boldsymbol{x}\right)$
$\displaystyle=\sqrt{\boldsymbol{u}\left(\pi,\boldsymbol{x}\right)^{\top}(\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)})^{-1}\widetilde{}\boldsymbol{\Omega}^{\pi,(k)}(\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)\top})^{-1}\boldsymbol{u}\left(\pi,\boldsymbol{x}\right)},\quad$
$\displaystyle\widetilde{\sigma}_{q}^{(k)}\left(\pi,\boldsymbol{x},a\right)$
$\displaystyle=\sqrt{\boldsymbol{z}\left(\boldsymbol{x},a\right)^{\top}(\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)})^{-1}\widetilde{}\boldsymbol{\Omega}^{\pi,(k)}(\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)\top})^{-1}\boldsymbol{z}\left(\boldsymbol{x},a\right)},$

where the matrices $\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)}$ and
$\widetilde{}\boldsymbol{\Omega}^{\pi,(k)}$ are defined as

$\displaystyle\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)}$
$\displaystyle=\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=1}^{T}(\widetilde{}\boldsymbol{\Lambda}^{i,(k)})^{\top}Z_{i,t}\left(Z_{i,t}-\gamma
U_{i,t+1}^{\pi}\right)^{\top}\widetilde{}\boldsymbol{\Lambda}^{i,(k)},$ (39)
$\displaystyle\widetilde{}\boldsymbol{\Omega}^{\pi,(k)}$
$\displaystyle=\frac{1}{NT}\sum_{i=1}^{N}\sum_{t=1}^{T}(\widetilde{}\boldsymbol{\Lambda}^{i,(k)})^{\top}Z_{i,t}Z_{i,t}^{\top}\widetilde{}\boldsymbol{\Lambda}^{i,(k)}\left(R_{i,t}-\left(Z_{i,t}-\gamma
U_{i,t+1}^{\pi}\right)^{\top}\widetilde{}\boldsymbol{\Lambda}^{i,(k)}\widetilde{}\boldsymbol{\theta}^{\pi}\right)^{2},$
(40)

and
$\widetilde{}\boldsymbol{\Lambda}^{i,(k)}=\boldsymbol{\Lambda}^{i}(\boldsymbol{w}_{\cdot
k}\otimes\boldsymbol{I}_{JM})$.

###### Proof.

The proof uses the result in Corollary 3 and is similar to that of Theorem 5.
∎

######  Corollary 2 (Asymptotic normality of oracle integrated value
estimator).

Suppose Assumption 2 – 5 hold. Let $N_{min}=\underset{1\leq k\leq
K}{\min}N_{k}$ and $N_{max}=\underset{1\leq k\leq K}{\max}N_{k}$. If $K={\rm
o}\left(N_{\min}T\right)$, $J={\rm
o}\left(\sqrt{N^{(k)}T}/\log(N^{(k)}T)\right)$,
$J^{-\kappa/p}\ll\left(N_{\max}T\left(1+\left\lVert\int\boldsymbol{\phi}_{J}(\boldsymbol{x}){\cal
R}(d\boldsymbol{x})\right\rVert_{2}^{-2}\right)\right)^{-1/2}$, as either
$N_{\min}\rightarrow\infty$ or $T\rightarrow\infty$, we have for any
$\boldsymbol{x}\in{\cal X}$,

$\displaystyle\sqrt{NT}\cdot\widetilde{\sigma}_{{\cal
R}}^{(k)}(\pi)^{-1}\left(\widetilde{V}_{{\cal
R}}^{(k)}\left(\pi\right)-V_{{\cal R}}^{(k)}\left(\pi\right)\right)$
$\displaystyle\stackrel{{\scriptstyle{\cal D}}}{{\longrightarrow}}{\cal
N}\left(0,1\right)\,,$

for any $k\in[K]$, where

$\widetilde{\sigma}_{{\cal
R}}^{(k)}\left(\pi\right)^{2}=\left(\int\boldsymbol{u}\left(\pi,\boldsymbol{x}\right){\cal
R}(d\boldsymbol{x})\right)^{\top}(\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)})^{-1}\widetilde{}\boldsymbol{\Omega}^{\pi,(k)}(\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)\top})^{-1}\left(\int\boldsymbol{u}\left(\pi,\boldsymbol{x}\right){\cal
R}(d\boldsymbol{x})\right),$

$\widetilde{}\boldsymbol{\Sigma}^{\pi,(k)}$ and
$\widetilde{}\boldsymbol{\Omega}^{\pi,(k)}$ are given in (39) and (40).

###### Proof.

The proof uses the result in Corollary 4 and is similar to that of Theorem 5.
∎

###  A.1 Proof of Theorem 1

###### Proof.

By decomposition (29), we have

$\left\lVert\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right\rVert_{\infty}\leq\left\lVert\left(\widetilde{}\boldsymbol{\Sigma}\right)^{-1}\right\rVert_{\infty}\left\lVert\widetilde{}\boldsymbol{\zeta}_{1}\right\rVert_{\infty}+\left\lVert\left(\widetilde{}\boldsymbol{\Sigma}\right)^{-1}\right\rVert_{\infty}\left\lVert\widetilde{}\boldsymbol{\zeta}_{2}\right\rVert_{\infty}$

We bound each term on the right hand side as follows.

Using (36) and Lemma 10 (iii), we have

$\left\lVert(\widetilde{}\boldsymbol{\Sigma})^{-1}\right\rVert_{\infty}\leq\underset{k}{\max}\left\lVert(P^{(k)}\widetilde{}\boldsymbol{\Sigma}^{(k)})^{-1}\right\rVert_{\infty}=6C^{-1}\sqrt{J}(N/N_{\min}),$

with probability at least $1-\mathcal{O}\left((N_{\min}T)^{-2}\right)$.

By Lemma 12 (iii), we have

$\displaystyle\Pr\left(\left\lVert\widetilde{}\boldsymbol{\zeta}_{1}^{(k)}\right\rVert_{\infty}>c\sqrt{2\log(N_{k}T)/(N_{k}T)}\right)\leq
2JM(N_{k}T)^{-2},$

for some positive constant $c$. Lemma 12 (iv) shows that
$\left\lVert\widetilde{}\boldsymbol{\zeta}_{2}^{(k)}\right\rVert_{\infty}\leq
c_{1}J^{-\kappa/p}\ll c_{1}\left(N_{k}T\right)^{-1/2}$ almost surely. By a
union bound and Lemma 12 (iii), we have

$\displaystyle\Pr\left(\left\lVert\widetilde{}\boldsymbol{\zeta}_{1}\right\rVert_{\infty}>c(N_{\max}/N)\sqrt{2\log(N_{\min}T)/N_{\min}T}\right)$
$\displaystyle\leq\sum_{k=1}^{K}\Pr\left(\left\lVert
P^{(k)}\widetilde{}\boldsymbol{\zeta}_{1}^{(k)}\right\rVert_{\infty}>(N_{\max}/N)\sqrt{2\log(N_{\min}T)/N_{\min}T}\right)$
$\displaystyle\leq\sum_{k=1}^{K}\Pr\left(\left\lVert
P^{(k)}\widetilde{}\boldsymbol{\zeta}_{1}^{(k)}\right\rVert_{\infty}>P^{(k)}\sqrt{2\log(N_{k}T)/N_{k}T}\right)$
$\displaystyle=\sum_{k=1}^{K}2JM(N_{k}T)^{-2}$ $\displaystyle\leq
2JMK(N_{\min}T)^{-2},$

for some positive constant $c$. Similarly, using Lemma 12 (iv), we have

$\left\lVert\widetilde{}\boldsymbol{\zeta}_{2}\right\rVert_{\infty}=\underset{1\leq
k\leq K}{\max}\;\left\lVert
P^{(k)}\widetilde{}\boldsymbol{\zeta}_{2}^{(k)}\right\rVert_{\infty}\leq
c_{1}J^{-\kappa/p}N_{max}/N\ll
c_{1}N_{max}/N(N_{\max}T)^{-1/2},\quad\text{almost surely}.$

Thus, by a union bound, we have that

$\left\lVert\widetilde{}\boldsymbol{\theta}^{\pi}-\mathring{\boldsymbol{\theta}}^{\pi}\right\rVert_{\infty}\leq
6cC^{-1}(N_{\max}/N_{\min})\sqrt{2J\log(N_{\min}T)/N_{\min}T}$

holds with probability at least
$1-2JMK(N_{\min}T)^{-2}-\mathcal{O}\left((N_{\min}T)^{-2}\right)$. ∎

##  Appendix B Feasible estimator

###  B.1 Proof of Theorem 2

###### Proof.

Note that all coefficients are with respect to a given policy $\pi\in\Pi$,
thus we drop superscript $\pi$ for brevity. We use
$\boldsymbol{B}=\left[\boldsymbol{\beta}_{1},\cdots,\boldsymbol{\beta}_{N}\right]$
and $\boldsymbol{\beta}={\textsc{vec}}\left(\boldsymbol{B}\right)$ (also
$\boldsymbol{\Theta}=\left[\boldsymbol{\theta}_{1},\cdots,\boldsymbol{\theta}^{(k)}\right]$
and $\boldsymbol{\theta}={\textsc{vec}}\left(\boldsymbol{\Theta}\right)$)
interchangeably. It is easy to see that
$\boldsymbol{B}=\boldsymbol{\Theta}\boldsymbol{W}^{\top}$,
$\boldsymbol{\beta}^{i}=\boldsymbol{\Theta}\boldsymbol{w}_{i\cdot}=\boldsymbol{\Lambda}_{i}\boldsymbol{\beta}$,
and
$\boldsymbol{\beta}=(\boldsymbol{W}\otimes\boldsymbol{I}_{p})\boldsymbol{\theta}$.
Define

$\displaystyle L\left(\boldsymbol{B}\right)$
$\displaystyle=\boldsymbol{G}\left(\pi,\boldsymbol{B}\right)^{\top}\boldsymbol{G}\left(\pi,\boldsymbol{B}\right),\quad
P\left(\boldsymbol{B},\lambda\right)=\frac{1}{N^{2}}\sum_{1\leq i<j\leq
N}p\left((JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right\rVert_{2},\lambda\right),$
$\displaystyle\widetilde{L}\left(\boldsymbol{\Theta}\right)$
$\displaystyle=\widetilde{}\boldsymbol{G}(\pi,\boldsymbol{\Theta})^{\top}\widetilde{}\boldsymbol{G}\left(\pi,\boldsymbol{\Theta}\right),\quad\widetilde{P}\left(\boldsymbol{\Theta},\lambda\right)=P\left(\boldsymbol{\Theta}\boldsymbol{W}^{\top},\lambda\right),$

where

$\displaystyle
G\left(\pi,\boldsymbol{B}\right)=\frac{1}{NTJ}\sum_{i=1}^{N}\sum_{t=1}^{T}\boldsymbol{\Lambda}_{i}^{\top}\boldsymbol{Z}_{i,t}R_{i,t}-\boldsymbol{\Lambda}_{i}^{\top}\boldsymbol{Z}_{i,t}\left(\boldsymbol{Z}_{i,t}-\gamma\boldsymbol{U}_{\pi,i,t+1}\right)^{\top}\boldsymbol{\Lambda}_{i}\boldsymbol{\beta},$
$\displaystyle\widetilde{G}\left(\pi,\boldsymbol{\Theta}\right)=\frac{1}{NTJ}\sum_{i=1}^{N}\sum_{t=1}^{T}\widetilde{}\boldsymbol{\Lambda}_{i}^{\top}\boldsymbol{Z}_{i,t}R_{i,t}-\widetilde{}\boldsymbol{\Lambda}_{i}^{\top}\boldsymbol{Z}_{i,t}\left(\boldsymbol{Z}_{i,t}-\gamma\boldsymbol{U}_{\pi,i,t+1}\right)^{\top}\widetilde{}\boldsymbol{\Lambda}_{i}\boldsymbol{\theta}.$

Then the loss function in the optimization problem (10) and (21) can be
rewritten, respectively, as,

${\cal
L}\left(\boldsymbol{B}\right):=L\left(\boldsymbol{B}\right)+P\left(\boldsymbol{B},\lambda\right),\text{and}\qquad\widetilde{\cal
L}\left(\boldsymbol{\Theta}\right):=\widetilde{L}\left(\boldsymbol{\Theta}\right)+\widetilde{P}\left(\boldsymbol{\Theta},\lambda\right).$
(41)

Let ${\cal M}_{{\cal G}}$ be the subspace of $\mathbb{R}^{JM\times N}$,
defined as

${\cal M}_{{\cal G}}=\left\\{\boldsymbol{B}\in\mathbb{R}^{JM\times
N}:\boldsymbol{\beta}^{i}=\boldsymbol{\beta}^{j},\text{ for any }i,j\in{\cal
G}_{k},1\leq k\leq K\right\\}.$

Recall that $\boldsymbol{W}$ is the true $N\times K$ membership matrix, then
for each $\boldsymbol{B}\in{\cal M}_{{\cal G}}$, it can be written as
$\boldsymbol{B}=\boldsymbol{\Theta}\boldsymbol{W}^{\top}$ for some
$\boldsymbol{\Theta}\in\mathbb{R}^{JM\times K}$. Also by matrix calculation,
we have
$\boldsymbol{W}^{\top}\boldsymbol{W}=\operatorname{diag}\left\\{N_{1},\cdots,N_{K}\right\\}$
where $N_{k}$ denotes the number of trajectories in ${\cal G}_{k}$.

Let ${\cal T}:{\cal M}_{G}\rightarrow\mathbb{R}^{JM\times K}$ be the mapping
such that ${\cal T}(\boldsymbol{B})$ is the $JM\times K$ matrix whose $k$-th
column equals to the common value of $\boldsymbol{\beta}^{i}$ for $i\in{\cal
G}_{k}$. Let $\overline{{\cal T}}:\mathbb{R}^{JM\times
N}\rightarrow\mathbb{R}^{JM\times K}$ be the mapping such that
$\overline{{\cal
T}}(\boldsymbol{B})=\boldsymbol{B}\boldsymbol{W}(\boldsymbol{W}^{\top}\boldsymbol{W})^{-1}$.
Clearly, when $\boldsymbol{B}\in{\cal M}_{\cal G}$, ${\cal
T}(\boldsymbol{B})=\overline{{\cal T}}(\boldsymbol{B})$. By calculation, we
have $P\left(\boldsymbol{B},\lambda\right)=\widetilde{P}\left({\cal
T}(\boldsymbol{B}),\lambda\right)=0$ for every $\boldsymbol{B}\in{\cal
M}_{{\cal G}}$ and $P\left({\cal
T}^{-1}(\boldsymbol{\Theta}),\lambda\right)=\widetilde{P}\left(\boldsymbol{\Theta},\lambda\right)$
for every $\boldsymbol{\Theta}\in\mathbb{R}^{JM\times K}$.

Consider the neighborhood of $\mathring{\boldsymbol{B}}$:

$\mathring{{\cal N}}=\left\\{\boldsymbol{B}\in\mathbb{R}^{JM\times
N},\left\lVert{\textsc{vec}}\left(\boldsymbol{B}-\mathring{\boldsymbol{B}}\right)\right\rVert_{\infty}\leq\phi_{NT}\right\\}.$

According to Theorem 1, there is an event $E_{1}$ such that on the event
$E_{1}$,

$\left\lVert{\textsc{vec}}\left(\widetilde{}\boldsymbol{B}-\mathring{\boldsymbol{B}}\right)\right\rVert_{\infty}\leq\phi_{NT},$

and $\Pr(E_{1}^{C})\leq
2JMK(N_{min}T)^{-2}+\mathcal{O}\left((N_{\min}T)^{-2}\right)$. Hence,
$\widetilde{}\boldsymbol{B}\in\mathring{{\cal N}}$ on the event $E_{1}$.

For any $\boldsymbol{B}\in\mathbb{R}^{JM\times N}$, let
$\overline{\boldsymbol{\Theta}}=\overline{{\cal T}}(\boldsymbol{B})$,
$\overline{\boldsymbol{B}}={\cal
T}^{-1}\left(\overline{\boldsymbol{\Theta}}\right)$. Lemma 3 shows that, on
event $E_{1}$, ${\cal L}\left(\widetilde{}\boldsymbol{B}\right)<{\cal
L}\left(\overline{\boldsymbol{B}}\right)$ for any $\boldsymbol{B}$ whose
$\overline{\boldsymbol{B}}\neq\widetilde{}\boldsymbol{B}$. Lemma 4 shows that
there is an event $E_{2}$ such that $\Pr\left(E_{2}^{C}\right)\leq
2((NT)^{-1})$. On $E_{1}\cap E_{2}$, there is a neighborhood of
$\widetilde{}\boldsymbol{B}$, denote by $\widetilde{\cal N}$, such that ${\cal
L}\left(\overline{\boldsymbol{B}}\right)<{\cal L}\left(\boldsymbol{B}\right)$
for any $\boldsymbol{B}\in\mathring{{\cal N}}\cap\widetilde{\cal N}$ for
sufficient large $N$ or $T$. Therefore, we have ${\cal
L}\left({\textsc{vec}}\left(\widetilde{}\boldsymbol{B}\right)\right)<{\cal
L}\left({\textsc{vec}}\left(\boldsymbol{B}\right)\right)$ for any
$\boldsymbol{B}\in\mathring{{\cal N}}\cap\widetilde{\cal N}$ and
$\boldsymbol{B}\neq\widetilde{}\boldsymbol{B}$, so that
$\widetilde{}\boldsymbol{B}$ is a strict local minimizer of ${\cal
L}({\textsc{vec}}\left(\boldsymbol{B}\right))$ given in (10) on the event
$E_{1}\cap E_{2}$ with $\Pr\left(E_{1}\cap E_{2}\right)\geq
1-2(K+JM+1)(NT)^{-1}$ for sufficient large $N$ or $T$. ∎

######  Lemma 3.

Suppose the assumptions in Theorem 2 hold. For any
$\boldsymbol{B}\in\mathbb{R}^{JM\times N}$, let
$\overline{\boldsymbol{B}}={\cal T}^{-1}\left(\overline{\cal
T}(\boldsymbol{B})\right)$. On event $E_{1}$, ${\cal
L}\left(\widetilde{}\boldsymbol{B}\right)<{\cal
L}\left(\overline{\boldsymbol{B}}\right)$ for any
$\boldsymbol{B}\in\mathbb{R}^{JM\times N}$ whose
$\overline{\boldsymbol{B}}\neq\widetilde{}\boldsymbol{B}$.

###### Proof.

By definition (41), we have ${\cal L}(\cdot)=L(\cdot)+P(\cdot)$ and consider
each term on the RHS. First, we have
$L\left(\widetilde{}\boldsymbol{B}\right)<L\left(\overline{\boldsymbol{B}}\right)$
by the following argument. Since $\widetilde{}\boldsymbol{\Theta}$ is the
unique global minimizer of $\widetilde{L}\left(\boldsymbol{\Theta}\right)$,
then
$\widetilde{L}\left(\widetilde{}\boldsymbol{\Theta}\right)<\widetilde{L}\left(\overline{{\cal
T}}(\boldsymbol{B})\right)$ for all $\overline{{\cal
T}}(\boldsymbol{B})\neq\widetilde{}\boldsymbol{\Theta}$. By definition we have
$\widetilde{L}\left(\widetilde{}\boldsymbol{\Theta}\right)=L\left(\widetilde{}\boldsymbol{B}\right)$
and $\widetilde{L}\left(\overline{{\cal
T}}(\boldsymbol{B})\right)=L\left(\overline{\boldsymbol{B}}\right)$ and thus
the result.

Now we study the penalty term $P\left(\widetilde{}\boldsymbol{B}\right)$ and
$P\left(\overline{\boldsymbol{B}}\right)$. The group-wise coefficients satisfy

$\left\lVert\boldsymbol{\theta}-\mathring{\boldsymbol{\theta}}\right\rVert_{\infty}=\underset{k}{\sup}\left\lVert
N_{k}^{-1}\sum_{i\in{\cal
G}}\left(\boldsymbol{\beta}^{i}-\mathring{\boldsymbol{\beta}}^{i}\right)\right\rVert_{\infty}\leq\left\lVert\boldsymbol{\beta}-\mathring{\boldsymbol{\beta}}\right\rVert_{\infty}.$
(42)

Then, by Assumption 7, we have for any pair of groups $k\neq l$

$(JM)^{-1/2}\left\lVert\boldsymbol{\theta}^{(k)}-\boldsymbol{\theta}^{(l)}\right\rVert_{2}\geq(JM)^{-1/2}\left\lVert\mathring{\boldsymbol{\theta}}^{(k)}-\mathring{\boldsymbol{\theta}}^{(l)}\right\rVert_{2}-2\left\lVert\boldsymbol{\theta}-\mathring{\boldsymbol{\theta}}\right\rVert_{\infty}\geq
d_{NT}-2\phi_{NT}>c\lambda.$

By Assumption 6, we have
$\widetilde{P}\left(\widetilde{}\boldsymbol{\Theta}\right)=\widetilde{P}\left(\overline{\boldsymbol{\Theta}}\right)=C_{N}$.
Further we have
$P\left(\widetilde{}\boldsymbol{B}\right)=P\left(\overline{\boldsymbol{B}}\right)=C_{N}$
since by definition we have
$P\left(\widetilde{}\boldsymbol{B}\right)=\widetilde{P}\left(\widetilde{}\boldsymbol{\Theta}\right)$
and
$P\left(\overline{\boldsymbol{B}}\right)=\widetilde{P}\left(\overline{\boldsymbol{\Theta}}\right)$
where $\overline{\boldsymbol{\Theta}}=\overline{{\cal T}}(\boldsymbol{B})$. As
a result,

$\displaystyle{\cal
L}\left(\widetilde{}\boldsymbol{B}\right)=L\left(\widetilde{}\boldsymbol{B}\right)+P\left(\widetilde{}\boldsymbol{B}\right)<L\left(\overline{\boldsymbol{B}}\right)+P\left(\overline{\boldsymbol{B}}\right)={\cal
L}\left(\overline{\boldsymbol{B}}\right).$

∎

######  Lemma 4.

Suppose the assumptions in Theorem 2 hold. For any
$\boldsymbol{B}\in\mathbb{R}^{JM\times N}$, let
$\overline{\boldsymbol{\Theta}}=\overline{{\cal T}}(\boldsymbol{B})$ and
$\overline{\boldsymbol{B}}={\cal
T}^{-1}\left(\overline{\boldsymbol{\Theta}}\right)$. There is an event $E_{2}$
such that $\Pr\left(E_{2}^{C}\right)\leq 2((NT)^{-1})$. On $E_{1}\cap E_{2}$,
there is a neighborhood $\widetilde{\cal N}$ of $\widetilde{}\boldsymbol{B}$
such that ${\cal L}\left(\overline{\boldsymbol{B}}\right)<{\cal
L}\left(\boldsymbol{B}\right)$ for any $\boldsymbol{B}\in\mathring{{\cal
N}}\cap\widetilde{\cal N}$ for sufficiently large $N$ or $T$.

###### Proof.

For a positive sequence $\delta_{NT}$, let $\widetilde{\cal
N}=\left\\{\boldsymbol{\beta}:\left\lVert\boldsymbol{\beta}-\widetilde{}\boldsymbol{\beta}\right\rVert_{2}\leq\delta_{NT}\right\\}$
be a neighborhood of $\widetilde{}\boldsymbol{\beta}$. Similar to (42), we
have, for any $\boldsymbol{\beta}\in\widetilde{\cal N}\cap\mathring{{\cal
N}}$,

$\displaystyle\left\lVert\overline{}\boldsymbol{\theta}-\mathring{\boldsymbol{\theta}}\right\rVert_{\infty}$
$\displaystyle\leq\left\lVert\overline{}\boldsymbol{\beta}-\mathring{\boldsymbol{\beta}}\right\rVert_{\infty}\leq\phi_{NT},$
$\displaystyle\left\lVert\overline{}\boldsymbol{\theta}-\widetilde{}\boldsymbol{\theta}\right\rVert_{\infty}$
$\displaystyle\leq\left\lVert\overline{}\boldsymbol{\beta}-\widetilde{}\boldsymbol{\beta}\right\rVert_{\infty}\leq\delta_{NT},$

thus $\overline{}\boldsymbol{\beta}\in\widetilde{\cal N}\cap\mathring{{\cal
N}}$. Let
$\boldsymbol{\beta}^{(\iota)}=\iota\boldsymbol{\beta}+(1-\iota)\overline{}\boldsymbol{\beta}$
for some $\iota\in(0,1)$, we have

$\left\lVert\boldsymbol{\beta}^{(\iota)}-\mathring{\boldsymbol{\beta}}\right\rVert_{\infty}\leq\phi_{NT},$

By a Taylor expansion, we have

$\displaystyle{\cal L}\left(\boldsymbol{\beta}\right)-{\cal
L}\left(\overline{}\boldsymbol{\beta}\right)$
$\displaystyle=\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)^{\top}\boldsymbol{G}\left(\pi,\boldsymbol{\beta}\right)-\boldsymbol{G}\left(\pi,\overline{}\boldsymbol{\beta}\right)^{\top}\boldsymbol{G}\left(\pi,\overline{}\boldsymbol{\beta}\right)$
$\displaystyle~{}~{}~{}+P\left(\boldsymbol{\beta},\lambda\right)-P\left(\overline{}\boldsymbol{\beta},\lambda\right)$
$\displaystyle=I_{1}+I_{2},$

where, for some
$\boldsymbol{\beta}^{(\iota)}=\iota\boldsymbol{\beta}+(1-\iota)\overline{}\boldsymbol{\beta}$
with $\iota\in(0,1)$,

$I_{1}=2\boldsymbol{G}\left(\pi,\boldsymbol{\beta}^{(\iota)}\right)^{\top}\boldsymbol{\Sigma}\left(\boldsymbol{\beta}-\overline{}\boldsymbol{\beta}\right),\quad\text{and}\quad
I_{2}=\frac{1}{N^{2}}\sum_{1\leq i\leq N}\left(\frac{\partial
P\left(\boldsymbol{\beta},\lambda\right)}{\partial\boldsymbol{\beta}^{(\iota)}}\bigg{|}_{\boldsymbol{\beta}=\boldsymbol{\beta}^{(\iota)}}\right)^{\top}\left(\boldsymbol{\beta}^{i}-\overline{}\boldsymbol{\beta}^{i}\right).$

Firstly, we consider $I_{2}$. Recall that by Assumption 6, Term $I_{2}$ can be
rewritten as

$\displaystyle N^{2}I_{2}$ $\displaystyle=\sum_{1\leq i\leq
N}\left(\frac{\partial
P\left(\boldsymbol{\beta},\lambda\right)}{\partial\boldsymbol{\beta}^{(\iota)}}\bigg{|}_{\boldsymbol{\beta}=\boldsymbol{\beta}^{(\iota)}}\right)^{\top}\left(\boldsymbol{\beta}^{i}-\overline{}\boldsymbol{\beta}^{i}\right)$
$\displaystyle=\lambda\sum_{i<j}\rho^{\prime}\left((JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i,\iota}-\boldsymbol{\beta}^{j,\iota}\right\rVert_{2}\right)(JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i,\iota}-\boldsymbol{\beta}^{j,\iota}\right\rVert_{2}^{-1}\left(\boldsymbol{\beta}^{i,\iota}-\boldsymbol{\beta}^{j,\iota}\right)^{\top}\left(\boldsymbol{\beta}^{i}-\overline{}\boldsymbol{\beta}^{i}\right)$
$\displaystyle~{}~{}~{}+\lambda\sum_{j>i}\rho^{\prime}\left((JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{j,\iota}-\boldsymbol{\beta}^{i,\iota}\right\rVert_{2}\right)(JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{j,\iota}-\boldsymbol{\beta}^{i,\iota}\right\rVert_{2}^{-1}\left(\boldsymbol{\beta}^{j,\iota}-\boldsymbol{\beta}^{i,\iota}\right)^{\top}\left(\boldsymbol{\beta}^{j}-\overline{}\boldsymbol{\beta}^{j}\right)$
$\displaystyle=\lambda\sum_{i<j}\rho^{\prime}\left((JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i,\iota}-\boldsymbol{\beta}^{j,\iota}\right\rVert_{2}\right)(JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i,\iota}-\boldsymbol{\beta}^{j,\iota}\right\rVert_{2}^{-1}\left(\boldsymbol{\beta}^{i,\iota}-\boldsymbol{\beta}^{j,\iota}\right)^{\top}\left(\left(\boldsymbol{\beta}^{i}-\overline{}\boldsymbol{\beta}^{i}\right)-\left(\boldsymbol{\beta}^{j}-\overline{}\boldsymbol{\beta}^{j}\right)\right).$

Note that, for any $i,j\in{\cal G}_{k}$,
$\overline{}\boldsymbol{\beta}^{i}=\overline{}\boldsymbol{\beta}^{j}$ and
$\boldsymbol{\beta}^{i,\iota}-\boldsymbol{\beta}^{j,\iota}=\iota\left(\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right)$,
while, for any $i\in{\cal G}_{k}$, $j\in{\cal G}_{l}$, $k\neq l$,

$(JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i,\iota}-\boldsymbol{\beta}^{j,\iota}\right\rVert_{2}\geq\underset{k\neq
l}{\min}\;(JM)^{-1/2}\left\lVert\mathring{\boldsymbol{\theta}}^{(k)}-\mathring{\boldsymbol{\theta}}^{(l)}\right\rVert_{2}-2\left\lVert\boldsymbol{\beta}-\mathring{\boldsymbol{\beta}}\right\rVert_{\infty}\geq
d_{NT}-2\phi_{NT}>a\lambda,$

and thus
$\rho^{\prime}\left((JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i,\iota}-\boldsymbol{\beta}^{j,\iota}\right\rVert_{2}\right)$
= 0. As a result, we have

$I_{2}=\frac{1}{N^{2}}\cdot\lambda\sum_{k=1}^{K}\sum_{i,j\in{\cal
G}_{k},i<j}\rho^{\prime}\left((JM)^{-1/2}\iota\left\lVert\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right\rVert_{2}\right)(JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right\rVert_{2}.$

Note that, since $\overline{}\boldsymbol{\beta}\in\widetilde{\cal
N}\cap\mathring{{\cal N}}$, we have

$(JM)^{-1}\left\lVert\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right\rVert_{2}^{2}\leq
2\left\lVert\boldsymbol{\beta}-\overline{}\boldsymbol{\beta}\right\rVert_{\infty}^{2}\leq
2\left\lVert\boldsymbol{\beta}-\widetilde{}\boldsymbol{\beta}\right\rVert_{\infty}^{2}+2\left\lVert\overline{}\boldsymbol{\beta}-\widetilde{}\boldsymbol{\beta}\right\rVert_{\infty}^{2}\leq
4\delta_{NT}.$

Hence, by concavity of $\rho(\cdot)$ and $0<\iota<1$, we have

$I_{2}\geq\frac{1}{N^{2}}\cdot\lambda\sum_{k=1}^{K}\sum_{i,j\in{\cal
G}_{k},i<j}\rho^{\prime}\left(4\delta_{NT}\right)(JM)^{-1/2}\left\lVert\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right\rVert_{2}.$
(43)

Now we consider $I_{1}$, which can be rewritten as

$\displaystyle-I_{1}$
$\displaystyle=\boldsymbol{G}\left(\pi,\boldsymbol{\beta}^{(\iota)}\right)^{\top}\boldsymbol{\Sigma}\left(\boldsymbol{\beta}-\overline{}\boldsymbol{\beta}\right)$
$\displaystyle=\frac{1}{N^{2}}\sum_{i=1}^{N}(\boldsymbol{G}_{i}^{\iota})^{\top}\boldsymbol{\Sigma}_{i}\left(\boldsymbol{\beta}^{i}-\overline{}\boldsymbol{\beta}^{i}\right)$
$\displaystyle=\frac{1}{N^{2}}\sum_{1\leq k\leq K}\sum_{i,j\in{\cal
G}_{k}}(N_{k})^{-1}\left(\boldsymbol{\Sigma}_{i}^{\top}\boldsymbol{G}_{i}^{\iota}\right)^{\top}\left(\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right)$
$\displaystyle=\frac{1}{N^{2}}\sum_{1\leq k\leq K}\sum_{i,j\in{\cal
G}_{k},i<j}(N_{k})^{-1}\left(\boldsymbol{\Sigma}_{i}^{\top}\boldsymbol{G}_{i}^{\iota}-\boldsymbol{\Sigma}_{j}^{\top}\boldsymbol{G}_{j}^{\iota}\right)^{\top}\left(\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right),$

where the third equation follows from the fact that $\sum_{j\in{\cal
G}_{k}}\left(\boldsymbol{\beta}^{j}-\overline{}\boldsymbol{\beta}^{j}\right)=0$
and

$\displaystyle\boldsymbol{G}_{i}^{\iota}=(TJ)^{-1}\sum_{t=1}^{T}Z_{i,t}\left(R_{i,t}-(Z_{i,t}-\gamma
U_{i,t+1}^{\pi})\boldsymbol{\beta}^{i,\iota}\right)=\boldsymbol{\zeta}_{i,1}+\boldsymbol{\zeta}_{i,2}+\boldsymbol{\Sigma}_{i}(\boldsymbol{\beta}^{i,\iota}-\mathring{\boldsymbol{\beta}}^{i})$
$\displaystyle\boldsymbol{\Sigma}_{i}=(TJ)^{-1}\sum_{t=1}^{T}Z_{i,t}\left(Z_{i,t}-\gamma
U_{i,t+1}^{\pi}\right)^{\top},\quad\boldsymbol{\zeta}_{i,1}=(TJ)^{-1}\sum_{t=1}^{T}Z_{i,t}\varepsilon_{1,i,t},\quad\boldsymbol{\zeta}_{i,2}=(TJ)^{-1}\sum_{t=1}^{T}Z_{i,t}\varepsilon_{2,i,t},$

and $\varepsilon_{1,i,t}^{\iota}$ and $\varepsilon_{2,i,t}^{\iota}$ are
defined in (33) and (34). Hence, we have

$\displaystyle\left\lvert I_{1}\right\rvert$
$\displaystyle=\left\lvert\sum_{1\leq k\leq K}\sum_{i,j\in{\cal
G}_{k},i<j}(N_{k})^{-1}\left(\boldsymbol{\Sigma}_{i}\boldsymbol{G}_{i}^{\iota}-\boldsymbol{\Sigma}_{j}\boldsymbol{G}_{j}^{\iota}\right)^{\top}\left(\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right)\right\rvert$
(44)
$\displaystyle\leq(N_{\min})^{-1}\cdot\max_{i,j}\left\lVert\boldsymbol{\Sigma}_{i}\boldsymbol{G}_{i}^{\iota}-\boldsymbol{\Sigma}_{j}\boldsymbol{G}_{j}^{\iota}\right\rVert_{2}\cdot\sum_{1\leq
k\leq K}\sum_{i,j\in{\cal
G}_{k},i<j}\left\lVert\boldsymbol{\beta}^{i}-\boldsymbol{\beta}^{j}\right\rVert_{2}.$

By Lemma 6 and 7, we have that, for some positive constant $c_{1}$,

$\max_{1\leq i\leq
N}\left\lVert\boldsymbol{\Sigma}_{i}\right\rVert_{2}^{2}\leq\max_{1\leq i\leq
N}(TJ)^{-2}T\cdot\sum_{t=1}^{T}\left(\left\lVert
Z_{i,t}Z_{i,t}^{\top}\right\rVert_{2}^{2}+\gamma\left\lVert
U_{\pi,i,t+1}Z_{i,t}^{\top}\right\rVert_{2}^{2}\right)\leq 2c_{1}^{4}J^{-1}.$

By Lemma 7 and 8, we have that,

$\displaystyle\max_{1\leq i\leq
N}\left\lVert\boldsymbol{\zeta}_{i,1}+\boldsymbol{\zeta}_{i,2}\right\rVert_{2}^{2}\leq\max_{1\leq
i\leq N}(TJ)^{-2}T\cdot\sum_{t=1}^{T}\left\lVert
Z_{i,t}\right\rVert_{2}^{2}\cdot\left(\left\lVert\varepsilon_{1,i,t}^{\iota}\right\rVert_{2}^{2}+\left\lVert\varepsilon_{2,i,t}^{\iota}\right\rVert_{2}^{2}\right)\leq
J^{-1}\left(c_{2}^{2}+(2c_{3}J^{-\kappa/p})^{2})\right).$ $\max_{1\leq i\leq
N}\left\lVert\boldsymbol{\Sigma}_{i}(\boldsymbol{\beta}^{i,\iota}-\mathring{\boldsymbol{\beta}}^{i})\right\rVert_{2}^{2}\leq\max_{1\leq
i\leq N}\left\lVert\boldsymbol{\Sigma}_{i}\right\rVert_{2}^{2}\max_{1\leq
i\leq
N}\left\lVert\boldsymbol{\beta}^{i,\iota}-\mathring{\boldsymbol{\beta}}\right\rVert_{2}^{2}\leq
2c_{1}^{4}\phi_{NT}^{2}.$

Thus, we have

$\displaystyle\max_{i,j}\left\lVert\boldsymbol{\Sigma}_{i}\boldsymbol{G}_{i}^{\iota}-\boldsymbol{\Sigma}_{j}\boldsymbol{G}_{j}^{\iota}\right\rVert_{2}$
$\displaystyle\leq 2\max_{1\leq i\leq
N}\left\lVert\boldsymbol{\Sigma}_{i}\boldsymbol{G}_{i}^{\iota}\right\rVert_{2}$
$\displaystyle\leq 4c_{1}^{2}\cdot
J^{-1}\cdot(c_{2}+2c_{3}J^{-\kappa/p})+4c_{1}^{2}J^{-1/2}\phi_{NT},$ (45)