Projected State-action Balancing Weights for Offline Reinforcement Learning

Consider a trajectory $\left\{S_{t},A_{t},R_{t}\right\}_{t\geq 0}$, where $(S_{t},A_{t},R_{t})$ denotes the triplet of the state, action and immediate reward, observed at the decision point $t$. Let ${\cal S}$ and ${\cal A}$ be the state and action spaces, respectively. We assume ${\cal A}$ is finite and consider the following two assumptions, which are commonly imposed in infinite-horizon OPE problems.

For the notational convenience, we assume the transition kernel has a density $p$. Assumption 1 requires that given the current state-action pair, future states are independent of past observations. Note that this assumption can be tested using the observed data [e.g., 61]. If one believes the trajectory satisfies some higher-order Markovian properties, one can instead construct a state variable by aggregating all the original state variables over the corresponding number of decision points. We refer the interested readers to [61] for an example.

Assumption 2 states that the current reward is conditionally mean independent of the history given the current state-action pair. One can also regard $R_{t}$ as a part of $S_{t+1}$ if needed. The uniform boundedness assumption on rewards is introduced for simplifying the technical derivation and can be relaxed. In practice, the time-stationarity of the transition density $p$ and the reward function $r$ can be warranted by incorporating time-associated covariates. Under Assumptions 1 and 2, the tuple ${\cal M}=\langle{\cal S},{\cal A},r,p\rangle$ forms a discrete-time homogeneous MDP.

A policy is defined as a way of choosing actions at all decision points. In this work, we use the value criterion (defined in (2) below) to evaluate policies and thus focus on the time-invariant Markovian policy (also called a stationary policy) $\pi$, which is a function mapping from the state space ${\cal S}$ to a probability mass function over the action space ${\cal A}$. More specifically, $\pi(a\,|\,s)$ is the probability of choosing an action $a$ given a state $s$. The sufficiency of considering only the stationary policy is explained in Section 6.2 of [55].

In the offline RL setting, one of the fundamental tasks is to estimate a target stationary policy’s (expected) value function based on the pre-collected (batch) data. Given a stationary policy $\pi$, the value function is defined as

where $\mathbb{E}^{\pi}$ denotes the expectation with respect to the distribution whose actions are generated by $\pi$, and $0\leq\gamma<1$ refers to the discounted factor controlling the trade-off between the future rewards and the immediate rewards. The value function $V^{\pi}(s)$, which is always finite due to Assumption 2, represents the discounted sum of rewards under the target policy given the initial state $s$. Our goal is to estimate the policy value, i.e., the expectation of the value function, defined as

using the pre-collected training data, where $\mathbb{G}$ denotes a reference distribution over ${\cal S}$. In RL literature, $\mathbb{G}$ is typically assumed known.

Now suppose we are given pre-collected training data ${\cal D}_{n}$ consisting of $n$ independent and identically distributed finite-horizon trajectories, denoted by

where $T_{i}$ denotes the termination time for $i$-th trajectory. For simplicity, we assume the same number of time points are observed for all trajectories, i.e. $T_{i}=T$ for $i=1,\dots,n$. This assumption can be relaxed as long as $T_{i}$, $i=1,\dots,n$ are of the same order. We also make the following assumption on the data generating mechanism.

Under Assumptions 1 and 3, $\left\{S_{t},A_{t}\right\}_{t\geq 0}$ forms a discrete time time-homogeneous Markov chain. In the literature, $b$ is usually called the behavior policy, which may not be known. For convenience, we denote $\mathbb{E}$ by $\mathbb{E}^{b}$. Next we define notations for several important probability distributions. For any $t\geq 0$, we define $p_{t}^{\pi}(s,a)$ as the marginal density of $(S_{t},A_{t})$ at $(s,a)\in{\cal S}\times{\cal A}$ under the target policy $\pi$ and reference distribution $\mathbb{G}$. In particular, when $t=0$, $p_{t}^{\pi}(s,a)=\mathbb{G}(s)\pi(a\mid s)$. Similarly, we can define $p_{t}^{b}$ over ${\cal S}\times{\cal A}$ under the behavior policy, where $p^{b}_{0}(s,a)=\mathbb{G}_{0}(s)b(a\mid s)$. With $p_{t}^{\pi}$, the discounted visitation probability density is defined as

which is assumed to be well-defined.

Most existing model-based OPE methods for the above setting can be grouped into three categories. The first category is direct methods, which directly estimate the state-action value function [e.g., 42, 62] defined as

also known as the Q-function. As we can see from (1) and (2),

It is well-known that $Q^{\pi}$ satisfy the following Bellman equation

for $t\geq 0$, $s\in{\cal S}$ and $a\in{\cal A}$. Clearly (6) forms a conditional moment restriction, based on which $Q^{\pi}$ can be estimated by many methods such as generalized method of moments [21] and the nonparametric instrumental variable regression [52, 8]. The second category of OPE methods is motivated by the idea of marginal importance sampling [e.g., 41, 49]. Notice that, with (3), one can rewrite ${\cal V}(\pi)$ as

as long as $d^{\pi}$ is absolutely continuous with respect to $\bar{p}_{T}^{b}=\frac{1}{T}\sum_{v=0}^{T-1}p_{v}^{b}$. We call

the probability ratio function, which is used to adjust for the mismatch between the behavior policy $b$ and the target policy $\pi$. Based on this relationship, one can obtain an estimator of ${\cal V}(\pi)$ via estimating $\omega^{\pi}$. By the so-called backward Bellman equation of $d^{\pi}$, for any $(s,a)\in{\cal S}\times{\cal A}$, one can show that

This implies that

See the detailed derivation for obtaining (10) in Section S2.1 of the Supplementary Material Recall that $\mathbb{G}$ is known. Based on (10), several methods [e.g., 49, 68] can be leveraged to estimate $\omega^{\pi}$. In Section 3.4, we provide more discussion on other weighted estimators as compared with the proposed estimator.

The last category of methods combines direct and marginal importance sampling methods to construct a doubly robust estimator, which is also motivated by the following efficient influence function [e.g., 30]

where $Q^{\pi}$ and $\omega^{\pi}$ are nuisance functions.

Consider a general form of weighted estimators:

where $\omega_{i,t}$’s are some weights constructed from the training data $\mathcal{D}_{n}$. Due to (7), a reasonable strategy to derive such weights is to first estimate the probability ratio function $\omega^{\pi}$, and then evaluate it at the observed state-action pairs $\{(S_{i,t},A_{i,t})\}$ in ${\cal D}_{n}$. This is analogous to the inverse probability weights commonly adopted in the ATE estimation. However, this strategy often produces an unstable weighted estimator due to small weights [31]. Instead, there is a recent surge of interest to directly obtain weights that achieve empirical covariate balance [e.g., 25, 7, 74, 73]. These weights usually produce more stable weighted estimators with superior finite-sample performances for the ATE estimation.

Inspired by the covariate balancing, a natural idea is to choose the weights $\{\omega_{i,t}\}$ that ensure the (approximate) validity of the empirical counterpart of (10), i.e.,

over $f\in\mathrm{span}\{B_{1},\dots,B_{K}\}$ where $B_{1},\dots,B_{K}$ are $K$ pre-specified functions defined on $\mathcal{S}\times\mathcal{A}$. The equality (12) can be viewed as a form of the state-action balance, in contrast to the covariate balance in the ATE estimation. The space $\mbox{span}\{B_{1},\dots,B_{K}\}$ can be viewed as a finite-dimensional approximation of the function space in which the balance should be enforced. In theory, $K$ is expected to increase with $n$ and $T$. Changing the balancing criterion in [73], one can obtain a form of state-action balancing weights via the following mathematical programming:

where the tuning parameters $\delta_{k}\geq 0$ controls the imbalance with respect to $B_{k}$. When $\delta_{k}=0$, the weights achieve the exact balance (12) over $B_{k}$. In practice, the exact balance can be hard to achieve especially when $K$ is large. Allowing $\delta_{k}>0$ leads to approximate balance [1, 74], which introduces flexibility. In addition, one can also constrain the weights to be non-negative. Since non-negativity constraints are not necessary for consistent estimation (as shown in our theoretical analysis), we will not enforce them throughout this paper. Common choices of $\{B_{k}\}$ are constructed based on tensor products of one-dimensional basis functions. Examples of one-dimensional basis functions include spline basis [12] (for a continuous dimension) and indicator functions of levels (for a categorical dimension). The objective function (13a) is introduced to control the magnitude of the weights. Here $h$ is chosen as a non-negative, strictly convex and continuously differentiable function. Examples include $h(x)=(x-1)^{2}$ and $h(x)=(x-1)\log(x)$. In the following, we discuss the issue with the weights defined by (13), and explain the challenge of directly applying this covariate balancing idea in the OPE problem.

Define $\epsilon_{i,t}=R_{i,t}-r(S_{i,t},A_{i,t})$ and write the solution of (13) as $\tilde{\omega}^{\pi}_{i,t}$. Naturally, we can obtain a weighted estimator as $\tilde{\cal V}(\pi)=(nT)^{-1}\sum_{i=1}^{n}\sum_{t=0}^{T-1}\tilde{\omega}_{i,t}^{\pi}R_{i,t}$. For any function $B$ defined on $\mathcal{S}\times\mathcal{A}$, let

The difference between $\tilde{{\cal V}}(\pi)$ and ${\cal V}(\pi)$ yields the following decomposition:

where the first equality is given by (6) and the representation of ${\cal V}(\pi)$ in (5). Clearly, (17)-(18) can be controlled via the balancing constraint (13b) by carefully controlling $\{\delta_{k}\}$ and selecting $\{B_{k}\}$ so that $Q^{\pi}$ can be well approximated. For (19), by assuming that $\{\epsilon_{i,t}\}$ are independent noises of the trajectories $\{S_{i,t},A_{i,t}\}_{i=1,\dots n;t=0,\dots,T-1}$, it may not be hard to obtain an upper bound of order $\sqrt{nT}$ as long as the magnitude of $\tilde{\omega}_{i,t}^{\pi}$ is properly controlled. However, it remains unclear how to control (20) due to the complex dependence between $\tilde{\omega}_{i,t}^{\pi}$ and ${\cal D}_{n}$. Indeed, we observe an “expanded-dimension” issue due to the balancing constraints (13b), which will be explained in details in Section 3.2. This also motivates the development of the novel projected balancing constraints in Section 3.3.

First, we obtain the dual form of (13), which provides an important characterization for the solution of (13). Define $\bm{\Psi}_{K}(s,a,s^{\prime})=[B_{k}(s,a)-\gamma\sum_{a^{\prime}\in\mathcal{A}}\pi(a^{\prime}\mid s^{\prime})B_{k}(s^{\prime},a^{\prime})]_{k=1}^{K}\in\mathbb{R}^{K}$, $\bm{l}_{K}=[\mathbb{E}_{s\sim\mathbb{G}}\{(1-\gamma)\sum_{a\in\mathcal{A}}\pi(a\mid s)B_{k}(s,a)\}]_{k=1,\dots,K}\in\mathbb{R}^{K}$, and $\bm{\delta}_{K}=[\delta_{1},\dots,\delta_{K}]^{\intercal}\in\mathbb{R}^{K}.$

The proof of Theorem 1 is similar to that of Theorem 2 below, and can be found in Section S2.2 of the Supplementary Material. Now the expanded-dimension issue can be easily seen via the following example. Suppose that there are two triplets $(S_{i_{1},t_{1}},A_{i_{1},t_{1}},S_{i_{1},t_{1}+1})$ and $(S_{i_{2},t_{2}},A_{i_{2},t_{2}},S_{i_{2},t_{2}+1})$ such that $S_{i_{1},t_{1}}=S_{i_{2},t_{2}}$, $A_{i_{1},t_{1}}=A_{i_{2},t_{2}}$ and $S_{i_{1},t_{1}+1}\neq S_{i_{2},t_{2}+1}$. As the true probability ratio function $\omega^{\pi}$ is a function of the current state and action variables, we must have $\omega^{\pi}_{i_{1},t_{1}}=\omega^{\pi}_{i_{2},t_{2}}$. However, the solution form (22) in Theorem 1 does not lead to $\tilde{\omega}^{\pi}_{i_{1},t_{1}}=\tilde{\omega}^{\pi}_{i_{2},t_{2}}$ in general, which violates our knowledge of $\omega^{\pi}$.

One may hypothesize that the expanded-dimension issue is a finite-sample property, and the variation of the estimated weights due to the next state may diminish asymptotically under some reasonable conditions. To gain more insight, we show that the solution form (22) indeed induces an implicit restriction on the modeling of the true weight function under finite-state and finite-action settings. Therefore, unless one is willing to make further non-trivial assumptions on the weight function, the hypothesis cannot be true in general.

Notice that

where $f(s,a)=\sum^{K}_{k=1}\lambda_{k}B_{k}(s,a)$. To avoid dealing with the approximation error, we focus on an even more general class $\rho^{\prime}(\mathcal{G})$ of functions on $\mathcal{S}\times\mathcal{A}\times\mathcal{S}$, where

Recall that $\rho^{\prime}$ is the first derivative of $\rho$ defined in Theorem 1. Assume that $\tilde{\omega}^{\pi}(s,a,s^{\prime})\equiv\omega^{\pi}(s,a)$. We would like to know if $\rho^{\prime}(\mathcal{G})$ can model $\tilde{\omega}^{\pi}$ well. Suppose $\tilde{\omega}^{\pi}\in\rho^{\prime}(\mathcal{G}):=\{\rho^{\prime}(g(\cdot)):g\in\mathcal{G}\}$. As $\tilde{\omega}^{\pi}(s,a,s^{\prime})$ is constant with respect to $s^{\prime}$, we have $\tilde{\omega}^{\pi}\in\rho^{\prime}(\mathcal{G}^{\prime})$ where

characterizes the subclass with reduced input dimensions. A key question is whether $\mathcal{G}^{\prime}$ restricts the class of possible weight functions modeled by $\rho^{\prime}(\mathcal{G})$ and, as a result, induces some implicit form of restriction. To see this, we focus on the settings where $|\mathcal{S}|=p_{S}$ and $|\mathcal{A}|=p_{A}$ are finite. In Lemma 1 below, we show that the dimension of $\mathcal{G}^{\prime}$ is $p_{A}p_{S}-p_{S}+1$, which is strictly less than $p_{S}p_{A}-1$ as long as $p_{S}>2$. Note that, due to the natural constraint that $\mathbb{E}[T^{-1}\sum_{t=0}^{T-1}\omega^{\pi}(S_{t},A_{t})]=1$, a general weight function should have $p_{A}p_{S}-1$ free parameters. As $\rho^{\prime}$ is invertible, Lemma 1 suggests a possible implicit restriction and that the solution obtained from (13) may not be a consistent estimator for $\omega^{\pi}$.

To overcome the expanded-dimension issue, we propose an approximate projection step, which is applied to $\sum_{a^{\prime}\in\mathcal{A}}\pi(a^{\prime}\mid S_{i,t+1})B_{k}(S_{i,t+1},a^{\prime})$, $k=1,\dots,K$, to rule out the involvement of $S_{i,t+1}$. To explain the idea, we again focus on the decomposition of $\tilde{\cal V}(\pi)-{\cal V}(\pi)$. From (15) and (16), we would like to choose weights that ideally control

for every $k=1,\dots,K$. However, in practice, $g_{*}^{\pi}(S_{i,t},A_{i,t};B_{k})$ (i.e., $\mathbb{E}\{\sum_{a^{\prime}\in{\cal A}}\pi(a^{\prime}\mid S_{i,t+1})B_{k}(S_{i,t+1},a^{\prime})\mid S_{i,t},A_{i,t}\}$) is unknown to us. As explained in Section 3.2, the idea of replacing it with the empirical counterpart $\sum_{a^{\prime}\in{\cal A}}\pi(a^{\prime}\mid S_{i,t+1})B_{k}(S_{i,t+1},a^{\prime})$ results in a non-trivial expanded-dimension issue. Instead, we propose to estimate the projection term $g_{*}^{\pi}(S_{t},A_{t};B_{k})$ via a more involved optimization problem:

where ${\mathcal{G}}$ is a pre-specified function space that contains $g_{*}^{\pi}$, $\mu\geq 0$ is a tuning parameter and $J_{\mathcal{G}}(\cdot)$ is a regularization functional. In this work we focus on the kernel ridge regression, where ${\mathcal{G}}$ is a subset of a reproducing kernel Hilbert space (RKHS) and $J_{\mathcal{G}}$ is taken as the squared RKHS norm; see Assumption 4’(c) in Section 4.1 of the Supplementary Material In Theorem 3 of Section 4, we establish the finite-sample error bound of $\hat{g}^{\pi}(\cdot,\cdot;B)$ (in scaling of both $n$ and $T$) that holds uniformly for different $B$ (in a reasonably large class $\mathcal{Q}$ specified later). This provides a solid theoretical guarantee for replacing $g^{\pi}_{*}$ by $\hat{g}^{\pi}$ in the construction of the weights.

With the approximate projection (24), we propose to estimate the weights by solving the following optimization problem:

The resulting solution $\{\hat{w}_{i,t}^{\pi}\}_{1\leq i\leq n,0\leq t<T}$ are the proposed weights. As such, the proposed estimator for ${\cal V}(\pi)$ is

Similar to Theorem 1, we derive the dual form of (25), which is shown in Theorem 2 below. For the notational simplicity, we introduce the following notations: $L_{k}(s,a)=B_{k}(s,a)-\gamma g^{\pi}_{*}(s,a;B_{k})$, $\hat{L}_{k}(s,a)=B_{k}(s,a)-\gamma\hat{g}^{\pi}(s,a;B_{k})$, ${\bm{B}}_{K}(s,a)=[{B}_{k}(s,a)]_{k=1}^{K}\in\mathbb{R}^{K}$, ${\bm{L}}_{K}(s,a)=[{L}_{k}(s,a)]_{k=1}^{K}\in\mathbb{R}^{K}$ and $\hat{\bm{L}}_{K}(s,a)=[\hat{L}_{k}(s,a)]_{k=1}^{K}\in\mathbb{R}^{K}$.

We sometimes write $\hat{\omega}^{\pi}_{i,t}=\hat{\omega}^{\pi}_{i,t}(S_{i,t},A_{i,t})$. The proof can be found in Section S2.2 of the Supplementary Material. As seen from the representation (28), $\hat{w}^{\pi}_{i,t}$ do not suffer from the expanded-dimension issue that we see in (21). Besides, (27) can be regarded as an $M$-estimation of $\bm{\lambda}$ with a weighted $\ell_{1}$-norm regularization. Since the estimated weights are parametrized in $\bm{\lambda}$ via (28), Theorem 2 reveals a connection between $\hat{\omega}^{\pi}_{i,t}$ and the shrinkage estimation of the probability ratio function $\omega^{\pi}$. To see this, we consider the objective function $T^{-1}\sum_{t=0}^{T-1}\rho({\bm{L}}_{K}(S_{t},A_{t})^{\intercal}\bm{\lambda})-\bm{\lambda}^{\intercal}\bm{l}_{K}$. By Lemma 1 in [68], the expectation of this loss function is minimized when $\bm{\lambda}$ satisfies $\omega^{\pi}(s,a)=\rho^{\prime}(\bm{L}_{K}(s,a)^{\intercal}\bm{\lambda})$, for every $s\in\mathcal{S}$, $a\in\mathcal{A}$. In Theorem 4 of Section 4, we show the convergence rate of the proposed weights to the true weights in the scaling of both $n$ and $T$.

Next, we discuss the computation of the weights $\hat{\omega}_{i,t}$ in practice. The projection step (24) is a pre-computation step for the optimization (25). In other words, we only need to estimate $\hat{g}^{\pi}(S_{i,t},A_{i,t};B_{k})$, $k=1,\dots,K$, once and there is no need to recompute them within the weights estimation (25). The optimization (25) is a standard convex optimization problem with linear constraints. We outline the proposed algorithm of the weights in Algorithm 1 of the Supplementary Material.

Apart from the proposed projection method, a sensible alternative to avoid the expanded dimension is to directly specify a class $\mathcal{F}_{w}$ of functions over $\mathcal{S}\times\mathcal{A}$ as a model of the weight function (i.e., weights are evaluations of this function), which naturally results in the following form of estimators: