Robust Offline Reinforcement Learning for Non-Markovian Decision Processes

In this work, we consider finite-horizon episodic (non-Markovian) decision processes with an observation space $\mathcal{O}$, an action space $\mathcal{A}$ and a horizon $H$. The dynamics of the process is determined by a set of distributions $\{\mathcal{T}_{h}(\cdot|o_{1\mathrel{\mathop{\mathchar 58\relax}}h},a_{1\mathrel{\mathop{\mathchar 58\relax}}h})\in\Delta(\mathcal{O})\}_{h=1}^{H-1}$, where $o_{1\mathrel{\mathop{\mathchar 58\relax}}h}\in\mathcal{O}^{h}$ denotes a sequence of observations up to time $h$, $a_{1\mathrel{\mathop{\mathchar 58\relax}}h}\in\mathcal{A}^{h}$ denotes a sequence of actions up to time $h$, and $\mathcal{T}_{h}(o_{h+1}|o_{1\mathrel{\mathop{\mathchar 58\relax}}h},a_{1\mathrel{\mathop{\mathchar 58\relax}}h})$ denotes the probability of observing $o_{h+1}$ given history $o_{1\mathrel{\mathop{\mathchar 58\relax}}h},a_{1\mathrel{\mathop{\mathchar 58\relax}}h}$. A policy $\pi$ is defined by a collection of distributions $\pi=\{\pi_{h}\mathrel{\mathop{\mathchar 58\relax}}\mathcal{O}^{h}\times\mathcal{A}^{h-1}\rightarrow\Delta(\mathcal{A})\}_{h=1}^{H}$ chosen by an agent, with each $\pi_{h}(\cdot|\tau_{h-1},o_{h})$ being a distribution over $\mathcal{A}$. Specifically, at the beginning of each episode, the process starts at a fixed observation $o_{1}\in\mathcal{O}$ at time step 1. After observing $o_{1}$, the agent samples (takes) an action $a_{1}\in\mathcal{A}$ according to the distribution (policy) $\pi_{1}(\cdot|o_{1})$, and the process transits to $o_{2}\in\mathcal{O}$, which is sampled according to the distribution $\mathcal{T}_{1}(\cdot|o_{1},a_{1})$. Then, at any time step $h\geq 2$, due to the non-Markovian nature, the agent samples an action $a_{h}$ based on all past information $(o_{1\mathrel{\mathop{\mathchar 58\relax}}h},a_{1\mathrel{\mathop{\mathchar 58\relax}}h-1})$, and the process transits to $o_{h+1}$, sampled from $\mathcal{T}_{h}(o_{h+1}|o_{1\mathrel{\mathop{\mathchar 58\relax}}h},a_{1\mathrel{\mathop{\mathchar 58\relax}}h})$. The interaction terminates after time step $H$.

Notably, the general non-Markovian decision-making problem subsumes not only fully observable MDPs but also POMDPs as special cases. In MDPs, $\mathcal{T}_{h-1}(o_{h}|\tau_{h-1})=\mathcal{T}_{h-1}(o_{h}|o_{h-1},a_{h-1})$, and in POMDPs, $\mathcal{T}_{h-1}(o_{h}|\tau_{h-1})$ can be factorized as $\mathcal{T}_{h-1}(o_{h}|\tau_{h-1})=\sum_{s}\mathbb{O}_{h}(o_{h}|s_{h})\mathbb{T}_{h-1}(s_{h}|s_{h-1},a_{h-1})\cdots\mathbb{T}_{1}(s_{2}|s_{1},a_{1})$, where $\{s_{t}\}_{t=1}^{H}$ represents unobserved states, $\{\mathbb{O}_{t}\}_{t=1}^{H}$ and $\{\mathbb{T}_{t}\}_{t=1}^{H}$ are called emission distributions and transition distributions, respectively.

We note that an alternative equivalent description of the process dynamics is the set of probabilities $\{\mathcal{P}(o_{1\mathrel{\mathop{\mathchar 58\relax}}H}|a_{1\mathrel{\mathop{\mathchar 58\relax}}H-1})=\mathcal{T}_{H-1}(o_{H}|o_{1\mathrel{\mathop{\mathchar 58\relax}}H-1},a_{1\mathrel{\mathop{\mathchar 58\relax}}H-1})\cdots\mathcal{T}_{1}(o_{2}|o_{1},a_{1}),\forall o_{1\mathrel{\mathop{\mathchar 58\relax}}H},a_{1\mathrel{\mathop{\mathchar 58\relax}}H-1}\}$, which is also frequently used in this paper. In fact, both $\mathcal{P}$ and $\mathcal{T}$ fully determine the dynamics with an additional constraint for $\mathcal{P}$, that is, $\sum_{o_{h+1\mathrel{\mathop{\mathchar 58\relax}}H}}\mathcal{P}(o_{1\mathrel{\mathop{\mathchar 58\relax}}H}|a_{1\mathrel{\mathop{\mathchar 58\relax}}H-1})$ is constant with respect to $a_{h\mathrel{\mathop{\mathchar 58\relax}}H-1}$ for all action sequence $a_{h\mathrel{\mathop{\mathchar 58\relax}}H-1}$. Hence, we use $\mathbb{D}=\{\mathcal{T}_{h},\mathcal{P}\}$ to represent the model dynamics. To make the presentation more explicit, we introduce a unified parameterization $\theta\in\Theta\subset\mathbb{R}^{D}$ for the process dynamics, written as $\mathbb{D}_{\theta}=\{\mathcal{T}_{h}^{\theta},\mathcal{P}^{\theta}\}$. Here, $\Theta$ is the complete set of model parameters and $D$ is usually a large integer. We also use $\theta$ to refer to a model.

For a given reward function $R\mathrel{\mathop{\mathchar 58\relax}}(\mathcal{O}\times\mathcal{A})^{H}\rightarrow[0,1]$, and any policy $\pi$, the value of the policy $\pi$ under the dynamics $\mathbb{D}_{\theta}$ is denoted by $V_{\theta,R}^{\pi}=\mathbb{E}_{\tau_{H}\sim\mathbb{D}_{\theta}^{\pi}}[R(o_{1\mathrel{\mathop{\mathchar 58\relax}}H},a_{1\mathrel{\mathop{\mathchar 58\relax}}H})]$. We remark here that the goal of an agent in standard RL is to find an $\epsilon$-optimal policy $\hat{\pi}$ such that $\max_{\pi}V_{\theta,R}^{\pi}-V_{\theta,R}^{\hat{\pi}}\leq\epsilon$ for some small error $\epsilon>0$ under a single model $\theta$.

While standard RL assumes that the system dynamics remains the same, robust RL diverges from this setting by focusing on the ‘worst-case’ or robust value in situations where the parameter $\theta$ can change within a defined uncertainty set $B\subset\mathbb{R}^{D}$ (where $B$ will be defined below). Mathematically, given a policy $\pi$, the robust value under the uncertainty set $B$ is defined by $V_{B,R}^{\pi}=\min_{\theta\in B}V_{\theta,R}^{\pi}$.

In this paper, we consider two types of uncertainty sets defined respectively on $\mathcal{T}$ and $\mathcal{P}$ via $f$-divergence $\mathtt{D}_{f}(P,P^{\prime})=\int f(dP/dP^{\prime})dP^{\prime}$, where $f$ is a convex function for measuring the uncertainty. Given a nominal model $\theta^{*}$, a $\mathcal{T}$-type uncertainty set under $f$-divergence with center $\theta^{*}$ and radius $\xi$ is defined as

Similarly, a $\mathcal{P}$-type uncertainty set under $f$-divergence with center $\theta^{*}$ and radius $\xi$ is defined as

For simplicity, we set $\xi$ to be a constant throughout the paper. Our proposed algorithm and analysis can be easily extended to the case when $\xi$ is a function of $\tau_{h}$ or $a_{1\mathrel{\mathop{\mathchar 58\relax}}H}$. Notably, the $\mathcal{T}$-type uncertainty set reduces to the $\mathcal{S}\times\mathcal{A}$-rectangular uncertainty set (Iyengar, , 2005) if $\mathbb{D}$ reduces to a Markov decision process, i.e., $\mathcal{T}_{h}(\cdot|\tau_{h})=\mathcal{T}_{h}(\cdot|o_{h},a_{h})$. However, due to the non-Markovian nature of our problem, we also consider the $\mathcal{P}$-type uncertainty set as $\mathcal{P}$ also fully determines the process.

We study two example types of $f$-divergence, which corresponds to two classical divergences. (i) If $f_{1}(t)=|t-1|/2$, then $\mathtt{D}_{f_{1}}$ is called total variation distance; (ii) If $f_{2}(t)=t\log t$, then $\mathtt{D}_{f_{2}}$ is called Kullback–Leibler (KL) divergence. We also use $B^{i}$ to represent $B^{f_{i}}$ for simplicity.

We focus on robust offline RL, where the agent is provided with an offline dataset $\mathcal{D}=\{\tau_{H}^{n}\}_{n=1}^{N}$ that contains $N$ sample trajectories collected via a behavior policy $\rho$ under the nominal model $\theta^{*}$. Given the type of the uncertain set $B\in\{B_{\mathcal{T}}^{i},B_{\mathcal{P}}^{i}\}_{i=1}^{2}$ (without knowing the center $\theta^{*}$), the learning goal of the agent is to find a near-optimal robust policy $\hat{\pi}$ using this offline dataset $\mathcal{D}$ such that

We first introduce some additional notations that are helpful to describe low-rank non-Markovian decision processes. Given a historical trajectory at time step $h$ as $\tau_{h}\mathrel{\mathop{\mathchar 58\relax}}=(o_{1\mathrel{\mathop{\mathchar 58\relax}}h},a_{1\mathrel{\mathop{\mathchar 58\relax}}h})$, we denote a future trajectory as $\omega_{h}\mathrel{\mathop{\mathchar 58\relax}}=(o_{h+1\mathrel{\mathop{\mathchar 58\relax}}H},a_{h+1\mathrel{\mathop{\mathchar 58\relax}}H})$. The set of all $\tau_{h}$ is denoted by $\mathcal{H}_{h}=(\mathcal{O}\times\mathcal{A})^{h}$ and the set of all future trajectories is denoted by $\Omega_{h}=(\mathcal{O}\times\mathcal{A})^{H-h}$. In addition, we add a superscript $o$ or $\mathrm{a}$ to either historical trajectory $\tau$ or a future trajectory to represent the observation sequence and the action sequence contained in that trajectory, respectively. For example, $\omega_{h}^{o}=o_{h+1\mathrel{\mathop{\mathchar 58\relax}}H}$ and $\omega_{h}^{a}=a_{h+1\mathrel{\mathop{\mathchar 58\relax}}H}$. To better describe a policy, we also use the notation $x_{h}=(\tau_{h-1},o_{h})=(o_{1\mathrel{\mathop{\mathchar 58\relax}}h},a_{1\mathrel{\mathop{\mathchar 58\relax}}h-1})$ as the available information for the learning agent at time step $h$. We remark that, having $\tau_{h},x_{h}$ in the same equation implies that $\tau_{h}=(x_{h},a_{h}).$

Recall that a non-Markov decision process $\theta$ can be fully determined by $\mathcal{P}^{\theta}$. We introduce a collection of dynamic matrices $\{\mathbb{M}_{h}^{\theta}\in\mathbb{R}^{|\mathcal{H}_{h}|\times|\Omega_{h}|}\}_{h=1}^{H}$. The entry at the $\tau_{h}$-th row and $\omega_{h}$-th column of $\mathbb{M}_{h}^{\theta}$ is $\mathcal{P}^{\theta}(\omega_{h}^{o},\tau_{h}^{o}|\tau_{h}^{a},\omega_{h}^{a})$.

It has been proved that any low-rank sequential decision-making problem $\theta$ admits Predictive State Representations (PSR) (Liu et al., , 2022) if for each $h$, there exists a set of future trajectories, namely, core tests known to the agent, $\mathcal{Q}_{h}=\{\mathbf{q}_{h,1},\ldots,\mathbf{q}_{h,d}\}\subset\Omega_{h}$, such that the submatrix restricted to these tests (columns) $\mathbb{M}_{h}^{\theta}[\mathcal{Q}_{h}]$ has rank $r$, where $d\geq r$ is a positive integer¹¹1Setting the same $d$ for different $h$ is for notation simplicity.. We consider all low-rank problems with the same set of core tests $\{\mathcal{Q}_{h}\}_{h}^{H}$, and have the following representations.

where $\mathbf{m}(\omega_{h})\in\mathbb{R}^{d}$, and $\psi(\tau_{h})\in\mathbb{R}^{d}$. We denote $\theta=(\phi_{h},\psi(\tau_{h}))$ to be the representation. In particular, $\Theta$ is the set of all low-rank problems with the same core tests $\{\mathcal{Q}_{h}\}_{h=1}^{H}$.

We assume $\psi_{0}$ is known to the agent (Huang et al., , 2023), and employ a bar over a feature $\psi(\tau_{h})$ to denote its normalization $\bar{\psi}(\tau_{h})=\psi(\tau_{h})/\phi_{h}^{\top}\psi(\tau_{h})$. $\bar{\psi}(\tau_{h})$ is also known as the prediction vector (Littman and Sutton, , 2001) or prediction feature of $\tau_{h}$, since $[\bar{\psi}(\tau_{h})]_{\ell}$ is the probability of observing $\mathbf{q}_{h,\ell}^{\mathrm{o}}$ given the history $\tau_{h}$ and taking action sequence $\mathbf{q}_{h,\ell}^{\mathrm{a}}$.

Moreover, let $\mathcal{Q}_{h}^{A}=\{\mathbf{q}_{h,\ell}^{\mathrm{a}}\}_{\ell=1}^{d}$ be the set of action sequences that are part of core tests, constructed by eliminating any repeated action sequence. $\mathcal{Q}_{h}^{A}$ known as the set of core action sequences.

We further assume that PSRs studied in this paper are well-conditioned, as specified in Assumption 3.2. Such an assumption and its variants are commonly adopted in the study of PSRs (Liu et al., , 2022; Chen et al., , 2022; Zhong et al., , 2022; Huang et al., , 2023).

Assumption 3.2 requires that the error of estimating $\theta$ does not significantly blow up when the estimation error $x$ of estimating the probability of core tests is small.

We present our algorithm for robust offline RL when the nominal model enjoys low-rank structure. Our algorithm features two main novel elements: (a) dataset distillation, which allows the estimated prediction features of distilled data being close to the true features, and thus enables us to construct (b) a lower confidence bound design for the robust value $V_{B_{\mathrm{u}}^{i}(\theta^{*}),R}^{\pi}$. In particular, we develop a new dual form for the robust value $V_{B_{\mathcal{P}}^{i}(\theta),R}^{\pi}$ under $\mathcal{P}$-type uncertainty set, that is more amenable to practical implementation. The main steps of the algorithm are explained in detail as follows, and the pseudo-code is presented in Algorithm 1.

Model estimation and dataset reform. First, we perform maximum likelihood estimation of the nominal model $\theta^{*}$ by minimizing the following negative log-likelihood loss function

which can be done using a supervised learning oracle, e.g. stochastic gradient descent. Let $\hat{\theta}=\arg\min_{\theta}\mathcal{L}(\theta|\mathcal{D})$.

Novel dataset distillation. Motivated by Huang et al., (2023), to construct a desirable lower confidence bound for any value function under the nominal model, the estimated model $\hat{\theta}$ needs to assign non-negligible probabilities on every sample trajectory. Differently from Huang et al., (2023) that introduces additional constraints to the optimization oracle, we propose a novel idea that only keeps the “good” sample trajectories whose estimated probability is greater than a pre-defined small value $p_{\min}$ under the estimated model $\hat{\theta}$ and collect those sample trajectories in set $\mathcal{D}^{\mathrm{g}}$, defined as follows.

We prove in Lemma B.1 that, with high probability $|\mathcal{D}^{\mathrm{g}}|=\Omega(N)$, which still guarantees sufficient number of samples for learning. Then, $\mathcal{D}^{\mathrm{g}}$ is randomly and evenly divided into $H$ datasets $\mathcal{D}_{0}^{\mathrm{g}},\ldots,\mathcal{D}_{H-1}^{\mathrm{g}}$, in order to separately learn $\bar{\psi}^{*}(\tau_{h})$ for each $h\in\{0,1,\ldots,H-1\}$.

LCB of robust value. Given the estimated model $\hat{\theta}$ and divergence type $i$, Algorithm 1 constructs a lower confidence bound (LCB) of $V_{B_{\mathrm{u}}^{i}(\theta^{*}),R}^{\pi}$ for all policy $\pi$ through a bonus function $\hat{b}$ defined as:

where $\lambda$ and $\alpha_{i}$ are pre-defined parameters. Recall that $\bar{\hat{\psi}}(\tau_{h})$ is the prediction feature described in Section 3.3. Then, we can show that $V_{B_{\mathrm{u}}^{i}(\hat{\theta}),R}^{\pi}-C_{\mathrm{u}}^{i}V_{\hat{\theta},\hat{b}}^{\pi}$ is an LCB of $V_{B_{\mathrm{u}}^{i}(\theta^{*}),R}^{\pi}$, where $C_{\mathrm{u}}^{i}$ (defined in Section 4.2) is a scaling parameter that depends on the type of uncertainty set.

Novel dual form of robust value. For the $\mathcal{T}$-type uncertainty sets, define robust value function as:

where $\tau_{h}=(x_{h},a_{h})$. Then, the following Bellman-type equations hold for robust value functions.

Therefore, robust value under $\mathcal{T}$-type uncertainty set can be calculated using the recursive formula above. Notably, the calculation of $Q_{B(\theta),R,h}^{\pi}(\tau_{h})$ involves solving an optimization problem in the form of $\inf_{P\mathrel{\mathop{\mathchar 58\relax}}\mathtt{D}_{f}(P,P^{*})}\mathbb{E}_{P}[g(X)]$, where $P$ and $P^{*}$ are two probability distributions over a set $\mathcal{X}$. This optimization problem can be solved efficiently as shown in the literature (Panaganti et al., , 2022; Blanchet et al., , 2023). We also provide its dual form in Appendix A.

On the other hand, when the uncertainty set is $\mathcal{P}$-type, we develop a new dual form of the robust value $V_{B_{\mathcal{P}}^{1}(\theta^{*}),R}^{\pi}$ and $V_{B_{\mathcal{P}}^{2}(\theta^{*}),R}^{\pi}$ in Equation 6 and Equation 7, respectively. Note that, in those equations, $f(x_{H})=\sum_{a_{H}}R(\tau_{H})\pi(\tau_{H-1})$.

Output policy design. Finally, Algorithm 1 outputs a policy $\hat{\pi}=\arg\max_{\pi\in\Pi}\left(V_{B_{\mathrm{u}}^{i}(\hat{\theta}),R}^{\pi}-C_{\mathrm{u}}^{i}V_{\hat{\theta},\hat{b}}^{\pi}\right)$, which maximizes a lower confidence bound of $V_{B_{\mathrm{u}}^{i}(\theta^{*}),R}^{\pi}$.

In this section, we present the theoretical results for Algorithm 1. Before we state the theorem, we need to introduce several definitions and condition coefficients.

Bracketing number. First, we introduce a complexity measure of the model class $\Theta$, which is standard in the literature (Liu et al., , 2022; Huang et al., , 2023).

Note that if $\Theta$ is the parameter space for tabular PSRs with rank $r$, we have $\log\mathcal{N}_{\varepsilon}(\Theta)\leq r^{2}|\mathcal{O}||\mathcal{A}|H^{2}\log\frac{H|\mathcal{O}||\mathcal{A}|}{\varepsilon}$ (see Theorem 4.7 in Liu et al., (2022)).

Novel concentrability coefficient. Under the low-rank structure and PSRs, we introduce the following novel concentrability coefficient $C_{\mathrm{p}}(\pi|\rho)$, namely Type-I concentrability coefficient, of a behavior policy $\rho$ against a robust policy $\pi$.

Wellness condition of the nominal model. When the uncertainty set is $\mathcal{T}$-type, we further introduce a wellness condition number $C_{B}$ of the nominal model $\theta^{*}$ defined as follows.

Intuitively, this condition number is called the maximum likelihood ratio between the distributions $\mathcal{P}^{\theta^{*}}(\cdot|a_{1\mathrel{\mathop{\mathchar 58\relax}}h})$ and $\mathcal{P}^{\theta}(\cdot|a_{1\mathrel{\mathop{\mathchar 58\relax}}h})$ and only depends on the property of the nominal model and the shape of the uncertainty set. In particular, when $B=B_{\mathcal{T}}^{2}$, we must have

which should hold for all policy $\pi$, and thus implies that $C_{B}$ is highly likely to be small.

Scaling parameters. Finally, we briefly introduce the scaling parameter $C_{\mathrm{u}}^{i}$ that is used to construct the LCB of robust value. Equipped with the dual forms, we define $C_{\mathcal{P}}^{1}=1$, $C_{\mathcal{T}}^{1}=C_{B}$. When uncertainty set is determined by KL divergence, we make the following assumption.

Define the scaling parameters $C_{\mathcal{P}}^{2}=3\exp(1/\eta^{*})$, and $C_{\mathcal{T}}^{2}=C_{B}\max\left\{\exp(\xi)/\xi,\lambda^{*}\exp(1/\lambda^{*})\right\}.$

We are now ready to present the theorem characterizing the performance of Algorithm 1.

The theorem states that as long as the type-I concentrability coefficient of $\rho$ against the optimal robust policy $\pi^{*}=\arg\max_{\pi}V_{B(\theta),R}^{\pi}$ is bounded, and the nominal model $\theta^{*}$ has finite wellness condition number $C_{B}$, Algorithm 1 finds an $\epsilon$-optimal robust policy provided that the offline dataset has $\Omega(1/\epsilon^{2})$ sample trajectories.