MOReL : Model-Based Offline Reinforcement Learning

Introduction

The fields of computer vision and NLP have seen tremendous advances by utilizing large-scale offline datasets for training and deploying deep learning models . In contrast, reinforcement learning (RL) is typically viewed as an online learning process. The RL agent iteratively collects data through interactions with the environment while learning the policy. Unfortunately, a direct embodiment of this trial and error learning is often inefficient and feasible only with a simulator . Similar to progress in other fields of AI, the ability to learn from offline datasets may hold the key to unlocking the sample efficiency and widespread use of RL agents.

Offline RL, also known as batch RL , involves learning a highly rewarding policy using only a static offline dataset collected by one or more data logging (behavior) policies. Since the data has already been collected, offline RL abstracts away data collection or exploration, and allows prime focus on data-driven learning of policies. This abstraction is suitable for safety sensitive applications like healthcare and industrial automation where careful oversight by a domain expert is necessary for taking exploratory actions or deploying new policies . Additionally, large historical datasets are readily available in domains like autonomous driving and recommendation systems, where offline RL may be used to improve upon currently deployed policies.

Due to use of static dataset, offline RL faces unique challenges. Over the course of learning, the agent has to evaluate and reason about various candidate policy updates. This offline policy evaluation is particularly challenging due to deviation between the state visitation distribution of the candidate policy and the logging policy. Furthermore, this difficulty is exacerbated over the course of learning as the candidate policies increasingly deviate from the logging policy. This change in distribution, as a result of policy updates, is typically called distribution shift and constitutes a major challenge in offline RL. Recent studies show that directly using off-policy RL algorithms with an offline dataset yields poor results due to distribution shift and function approximation errors . To overcome this, prior works have proposed modifications like Q-network ensembles and regularization towards the data logging policy . Most notably, prior work in offline RL has been confined almost exclusively to model-free methods .

Model-based RL (MBRL) presents an alternate set of approaches involving the learning of approximate dynamics models which can subsequently be used for policy search. MBRL enables the use of generic priors like smoothness and physics for model learning, and a wide variety of planning algorithms . As a result, MBRL algorithms have been highly sample efficient for online RL . However, direct use of MBRL algorithms with offline datasets can prove challenging, again due to the distribution shift issue. In particular, since the dataset may not span the entire state-action space, the learned model is unlikely to be globally accurate. As a result, planning using a learned model without any safeguards against model inaccuracy can result in “model exploitation” , yielding poor results . In this context, we study the pertinent question of how to effectively regularize and adapt model-based methods for offline RL.

Our Contributions: The principal contribution of our work is the development of MOReL (Model-based Offline Reinforcement Learning), a novel model-based framework for offline RL (see figure 1 for an overview). MOReL enjoys rigorous theoretical guarantees, enables transparent algorithm design, and offers state of the art (SOTA) results on widely studied offline RL benchmarks.

MOReL consists of two modular steps: (a) learning a pessimistic MDP (P-MDP) using the offline dataset; and (b) learning a near-optimal policy for the P-MDP. For any policy, the performance in the true MDP (environment) is approximately lower bounded by the performance in the P-MDP, making it a suitable surrogate for purposes of policy evaluation and learning. This also guards against model exploitation, which often plagues MBRL.

The P-MDP partitions the state space into “known” and “unknown” regions, and uses a large negative reward for unknown regions. This provides a regularizing effect during policy learning by heavily penalizing policies that visit unknown states. Such a regularization in the space of state visitations, afforded by a model-based approach, is particularly well suited for offline RL. In contrast, model-free algorithms are forced to regularize the policies directly towards the data logging policy, which can be overly conservative.

Theoretically, we establish upper bounds for the sub-optimality of a policy learned with MOReL, and a lower-bound for the sub-optimality of a policy learnable by any offline RL algorithm. We find that these bounds match upto log factors, suggesting that MOReL is nearly minimax optimal.

We evaluate MOReL on standard benchmark tasks used for offline RL. MOReL obtains SOTA results in $12$ out of $20$ environment-dataset configurations, and performs competitively in the rest. In contrast, the best prior algorithm obtains SOTA results in only $5$ (out of $20$) configurations.

In addition, this version of the paper extends the results presented at NeurIPS 2020 through addition of results in the D4RL benchmark suite and by expanding the scope of Lemma 3.

Related Work

Offline RL dates to at least the work of Lange et al. , and has applications in healthcare , recommendation systems , dialogue systems , and autonomous driving . Algorithms for offline RL typically fall under three categories. The first approach utilizes importance sampling and is popular in contextual bandits . For full offline RL, Liu et al. perform planning with learned importance weights while using a notion of pessimism for regularization. However, Liu et al. don’t explicitly consider generalization and their guarantees become degenerate if the logging policy does not span the support of the optimal policy. In contrast, our approach accounts for generalization, leads to stronger theoretical guarantees, and obtains SOTA results on challenging offline RL benchmarks. The second, and perhaps most popular approach is based on approximate dynamic programming (ADP). Recent works have proposed modification to standard ADP algorithms towards stabilizing Bellman targets with ensembles and regularizing the learned policy towards the data logging policy . ADP-based offline RL has also be studied theoretically . However, these works again don’t study the impact of support mismatch between logging policy and optimal policy. Finally, model-based RL has been explored only sparsely for offline RL in literature (see appendix for details). The work of Ross and Bagnell considered a straightforward approach of learning a model from offline data, followed by planning. They showed that this can have arbitrarily large sub-optimality. In contrast, our work develops a new framework utilizing the notion of pessimism, and shows both theoretically and experimentally that MBRL can be highly effective for offline RL. Concurrent to our work, Yu et al. also study a model-based approach to offline RL.

A cornerstone of MOReL is the P-MDP which partitions the state space into known and unknown regions. Such a hard partitioning was considered in early works like $E^{3}$ , R-MAX , and metric-$E^{3}$ , but was not used to encourage pessimism. Similar ideas have been explored in related settings like online RL and imitation learning . Our work differs in its focus on offline RL, where we show the P-MDP construction plays a crucial role. Moreover, direct practical instantiations of $E^{3}$ and metric-$E^{3}$ with function approximation have remained elusive.

Problem Formulation

To avoid notation clutter, we suppress the dependence on $\rho_{0}$ when understood from context, i.e. $J(\pi,\mathcal{M})\equiv J_{\rho_{0}}(\pi,\mathcal{M})$. We denote the optimal policy using $\pi^{*}:=\arg\max_{\pi}J_{\rho_{0}}(\pi,\mathcal{M})$. Typically, a class of parameterized policies $\pi_{\theta}\in\Pi(\Theta)$ are considered, and the parameters $\theta$ are optimized.

In offline RL, we are provided with a static dataset of interactions with the environment consisting of $\mathcal{D}=\{(s_{i},a_{i},r_{i},s_{i}^{\prime})\}_{i=1}^{N}$. The data can be collected using one or more logging (or behavioral) policies denoted by $\pi_{b}$. We do not assume logging policies are known in our formulation. Given $\mathcal{D}$, the goal in offline RL is to output a $\pi_{\text{out}}$ with minimal sub-optimality, i.e. $J(\pi^{*},\mathcal{M})-J(\pi_{\text{out}},\mathcal{M})$. In general, it may not be possible to learn the optimal policy with a static dataset (see section 4.1). Thus, we aim to design algorithms that would result in as low sub-optimality as possible.

Model-Based RL (MBRL) involves learning an MDP $\hat{\mathcal{M}}=\{S,A,r,\hat{P},\hat{\rho_{0}},\gamma\}$ which uses the learned transitions $\hat{P}$ instead of the true transition dynamics $P$. In this paper, we assume the reward function $r$ is known and use it in $\hat{M}$. If $r(\cdot)$ is unknown, it can also be learned from data. The initial state distribution $\hat{\rho_{0}}$ can either be learned from the data or $\rho_{0}$ can be used if known. Analogous to $\mathcal{M}$, we use $J_{\hat{\rho_{0}}}(\pi,\hat{\mathcal{M}})$ or simply $J(\pi,\hat{\mathcal{M}})$ to denote performance of $\pi$ in $\hat{M}$.

Algorithmic Framework

For ease of exposition and clarity, we first begin by presenting an idealized version of MOReL, for which we also establish theoretical guarantees. Subsequently, we describe a practical version of MOReL that we use in our experiments. Algorithm 1 presents the broad framework of MOReL. We now study each component of MOReL in greater detail.

Learning the dynamics model: The first step involves using the offline dataset to learn an approximate dynamics model $\hat{P}(\cdot|s,a)$. This can be achived through maximum likelihood estimation or other techniques from generative and dynamics modeling . Since the offline dataset may not span the entire state space, the learned model may not be globally accurate. So, a naïve MBRL approach that directly plans with the learned model may over-estimate rewards in unfamiliar parts of the state space, resulting in a highly sub-optimal policy . We overcome this with the next step.

Unknown state-action detector (USAD): We partition the state-action space into known and unknown regions based on the accuracy of learned model as follows.

($\alpha$-USAD) Given a state-action pair $(s,a)$, define an unknown state action detector as:

Here $D_{TV}\left(\hat{P}(\cdot|s,a),P(\cdot|s,a)\right)$ denotes the total variation distance between $\hat{P}(\cdot|s,a)$ and $P(\cdot|s,a)$. Intuitively, USAD provides confidence about where the learned model is accurate. It flags state-actions for which the model is guarenteed to be accurate as “known”, while flagging state-actions where such a guarantee cannot be ascertained as “unknown”. Note that USAD is based on the ability to guarantee the accuracy, and is not an inherent property of the model. In other words, there could be states where the model is actually accurate, but flagged as unknown due to the agent’s inability to guarantee accuracy. Two factors contribute to USAD’s effectiveness: (a) data availability: having sufficient data points “close” to the query; (b) quality of representations: certain representations, like those based on physics, can lead to better generalization guarantees. This suggests that larger datasets and research in representation learning can potentially enable stronger offline RL results.

Pessimistic MDP construction: We now construct a pessimistic MDP (P-MDP) using the learned model and USAD, which penalizes policies that venture into unknown parts of state-action space.

The $(\alpha,\kappa)$-pessimistic MDP is described by $\hat{\mathcal{M}}_{p}:=\{S\cup\textrm{HALT},A,r_{p},\hat{P}_{p},\hat{\rho}_{0},\gamma\}$. Here, $S$ and $A$ are states and actions in the MDP $\mathcal{M}$. HALT is an additional absorbing state we introduce into the state space of $\hat{\mathcal{M}}_{p}$. $\hat{\rho}_{0}$ is the initial state distribution learned from the dataset $\mathcal{D}$. $\gamma$ is the discount factor (same as $\mathcal{M}$). The modified reward and transition dynamics are given by:

Planning: The final step in MOReL is to perform planning in the P-MDP defined above. For simplicity, we assume a planning oracle that returns an $\epsilon_{\pi}$-sub-optimal policy in the P-MDP. A number of algorithms based on MPC , search-based planning , dynamic programming , or policy optimization can be used to approximately realize this..

In order to state our results, we begin by defining the notion of hitting time.

(Hitting time) Given an MDP $\mathcal{M}$, starting state distribution $\rho_{0}$, state-action pair $(s,a)$ and a policy $\pi$, the hitting time $T_{(s,a)}^{\pi}$ is defined as the random variable denoting the first time action $a$ is taken at state $s$ by $\pi$ on $\mathcal{M}$, and is equal to $\infty$ if $a$ is never taken by $\pi$ from state $s$. For a set of state-action pairs $\mathcal{S}\subseteq S\times A$, we define $T_{\mathcal{S}}^{\pi}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\min_{(s,a)\in\mathcal{S}}T_{(s,a)}^{\pi}$.

We are now ready to present our main result with the proofs deferred to the appendix.

(Policy value with pessimism) The value of any policy $\pi$ on the original MDP $\mathcal{M}$ and its $(\alpha,R_{\max})$-pessimistic MDP $\hat{\mathcal{M}}_{p}$ satisfies:

where $T_{\mathcal{U}}^{\pi}$ denotes the hitting time of unknown states $\mathcal{U}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\left\{(s,a):U^{\alpha}(s,a)=\textrm{TRUE}\right\}$ by $\pi$ on $\mathcal{M}$.

Theorem 1 can be used to bound the suboptimality of output policy $\pi_{\textrm{out}}$ of Algorithm 1.

Suppose PLANNER in Algorithm 1 returns an $\epsilon_{\pi}$ sub-optimal policy. Then, we have

(Upper bound; MOReL improves over the behavioral policy) Suppose $\rho_{0,\textrm{min}}>0$, $p_{\textrm{min}}>0$ and $d^{\pi_{b}}_{\textrm{min}}>0$ are the smallest non-zero elements of initial distribution $\rho_{0}$, state transition probabilities $P(\cdot|s,a)$, and discounted state probability distribution $d^{\pi_{b},\mathcal{M}}(s,a)$ respectively. If the dataset $\mathcal{D}$ consists of $n\geq\frac{C}{\left(d^{\pi_{b}}_{\textrm{min}}\right)^{2}}\cdot\log\frac{1}{\delta d^{\pi_{b}}_{\textrm{min}}}$ independent trajectories sampled according to a behavior policy $\pi_{b}$ with initial distribution $\rho_{0}$, then the output $\pi_{\textrm{out}}$ of Algorithm 1 satisfies:

with probability at least $1-C\delta$, where $C$ is a large enough constant and

is an error term related to finite samples that goes to zero as $n\rightarrow\infty$.

The bound consists of three terms: (i) a sampling error term $\epsilon_{n}$ which decreases with larger dataset sizes that is typical of offline RL; (ii) an optimization error term $\epsilon_{\pi}$ that can be made small with additional compute to find the optimal policy in the learned model; and (iii) a distribution shift term that depends on the coverage of the offline dataset and overlap with the optimal policy.

Prior results assume that $d^{\pi^{*},\mathcal{M}}(\mathcal{U}_{D})=0$, where $\mathcal{U}_{D}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\left\{(s,a)|(s,a,r,s^{\prime})\notin\mathcal{D}\right\}\supseteq\mathcal{U}$ is the set of state action pairs that don’t occur in the offline dataset, and guarantee finding an optimal policy under this assumption. Our result significantly improves upon these in three ways: i) $\mathcal{U}_{D}$ is replaced by a smaller set $\mathcal{U}$, leveraging the generalization ability of learned dynamics model, ii) the sub-optimality bound is extended to the setting where full support coverage is not satisfied i.e., $d^{\pi^{*},\mathcal{M}}(\mathcal{U})>0$, and iii) the sub-optimality bound on $\pi_{\textrm{out}}$ is stated in terms of unknown state hitting time $T_{\mathcal{U}}^{\pi^{*}}$, which can be significantly better than a bound that depends only on $d^{\pi^{*},\mathcal{M}}(\mathcal{U})$. To further strengthen our results, the following proposition shows that Lemma 3 is tight up to $\log$ factors.

(Lower bound) For any discount factor $\gamma\in[0.95,1)$, support mismatch $\epsilon\in\left(0,\frac{1-\gamma}{\log\frac{1}{1-\gamma}}\right]$ and reward range $[-R_{\textrm{max}},R_{\textrm{max}}]$, there is an MDP $\mathcal{M}$, starting state distribution $\rho_{0}$, optimal policy $\pi^{*}$ and a dataset collection policy $\pi_{b}$ such that i) $d^{\pi^{*},\mathcal{M}}(\mathcal{U}_{D})\leq\epsilon$, and ii) any policy $\hat{\pi}$ that is learned solely using the dataset collected with $\pi_{b}$ satisfies:

where $\mathcal{U}_{D}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\left\{(s,a):(s,a,r,s^{\prime})\notin\mathcal{D}\mbox{ for any }r,s^{\prime}\right\}$ denotes state action pairs not in the dataset $\mathcal{D}$.

We see that for $\epsilon<(1-\gamma)/(\log\tfrac{1}{1-\gamma})$, the lower bound obtained by Proposition 4 on the suboptimality of any offline RL algorithm matches the asymptotic (as $n\rightarrow\infty$) upper bound of Lemma 3 up to an additional log factor. For $\epsilon>(1-\gamma)/(\log\tfrac{1}{1-\gamma})$, Proposition 4 also implies (by choosing $\epsilon^{\prime}=(1-\gamma)/(\log\tfrac{1}{1-\gamma})<\epsilon$) that any offline algorithm must suffer at least constant factor suboptimality in the worst case. Finally, we note that as the size of dataset $\mathcal{D}$ increases to $\infty$, Theorem 1 and the optimality of PLANNER (i.e., $\epsilon_{\pi}=0$) together imply that $J_{\rho_{0}}(\pi_{\textrm{out}},\mathcal{M})\geq J_{\rho_{0}}(\pi_{b},\mathcal{M})$.

2 Practical Implementation Of MOReL

We now present a practical instantiation of MOReL (algorithm 1) utilizing a recent model-based NPG approach . The principal difference is the specialization to offline RL and construction of the P-MDP using an ensemble of learned dynamics models.

Dynamics model learning: We consider Gaussian dynamics models , $\hat{P}(\cdot|s,a)\equiv\mathcal{N}\left(f_{\phi}(s,a),\Sigma\right)$, with mean $f_{\phi}(s,a)=s+\sigma_{\Delta}\ \textrm{MLP}_{\phi}\left((s-\mu_{s})/\sigma_{s},(a-\mu_{a})/\sigma_{a}\right)$, where $\mu_{s},\sigma_{s},\mu_{a},\sigma_{a}$ are the mean and standard deviations of states/actions in $\mathcal{D}$; $\sigma_{\Delta}$ is the standard deviation of state differences, i.e. $\Delta=s^{\prime}-s,(s,s^{\prime})\in\mathcal{D}$; this parameterization ensures local continuity since the MLP learns only the state differences. The MLP parameters are optimized using maximum likelihood estimation with mini-batch stochastic optimization using Adam .

Experiments

Through our experimental evaluation, we aim to answer the following questions:

Comparison to prior work: How does MOReL compare to prior SOTA offline RL algorithms in commonly studied benchmark tasks?

Quality of logging policy: How does the quality (value) of the data logging (behavior) policy, and by extension the dataset, impact the quality of the policy learned by MOReL?

Importance of pessimistic MDP: How does MOReL compare against a naïve model-based RL approach that directly plans in a learned model without any safeguards?

Transfer from pessimistic MDP to environment: Does learning progress in the P-MDP, which we use for policy learning, effectively translate or transfer to learning progress in the environment?

To answer the above questions, we consider commonly studied benchmark tasks from OpenAI gym simulated with MuJoCo . Our experimental setup closely follows prior work . The tasks considered include Hopper-v2, HalfCheetah-v2, Ant-v2, and Walker2d-v2, which are illustrated in Figure 2. We consider five different logged data-sets for each environment, totalling 20 environment-dataset combinations. Datasets are collected based on the work of Wu et al. , with each dataset containing the equivalent of 1 million timesteps of environment interaction. We first partially train a policy $(\pi_{p})$ to obtain values around 1000, 4000, 1000, and 1000 respectively for the four environments. The first exploration strategy, Pure, involves collecting the dataset solely using $\pi_{p}$. The four other datasets are collected using a combination of $\pi_{p}$, a noisy variant of $\pi_{p}$, and an untrained random policy. The noisy variant of $\pi_{p}$ utilizes either epsilon-greedy or Gaussian noise, resulting in configurations $\texttt{eps-1},\texttt{eps-3},\texttt{gauss-1},\texttt{gauss-3}$ that signify various types and magnitudes of noise added to $\pi_{p}$. Please see appendix for additional experimental details.

We parameterize the dynamics model using 2-layer ReLU-MLPs and use an ensemble of 4 dynamics models to implement USAD as described in Section 4.2. We parameterize the policy using a 2-layer tanh-MLP, and train it using model-based NPG . We evaluate the learned policies using rollouts in the (real) environment, but these rollouts are not made available to the algorithm in any way for purposes of learning. This is similar to evaluation protocols followed in prior work . We present all our results averaged over $5$ different random seeds. Note that we use the same hyperparameters for all random seeds. In contrast, the prior works whose results we compare against tune hyper-parameters separately for each random seed .

We compare results of MOReL with prior SOTA algorithms like BCQ, BEAR, and all variants of BRAC. The results are summarized in Table 1. For fairness of comparison, we reproduce results from prior work and do not run the algorithms ourselves. We provide a more expansive table with additional baseline algorithms in the appendix. Our algorithm, MOReL, achives SOTA results in $12$ out of the $20$ environment-dataset combinations, overlaps in error bars for $3$ other combinations, and is competitive in the remaining cases. In contrast, the next best approach (a variant of BRAC) achieves SOTA results in only $5$ out of $20$ configurations.

Comparison of MOReL’s performance in the D4RL benchmark suite

The D4RL benchmark suite for offline RL was introduced in concurrent work. We also study the performance of MOReL in this benchmark suite. We find that MOReL achieves the highest (normalized) score in $5$ out of $12$ domains studied, while the next best algorithm (CQL) achieves the highest score in only $3$ out of $12$ domains. Furthermore, we observe that MOReL is often very competitive with the best performing algorithm in any given domain even if it doesn’t achieve the top score. However, in many domains, MOReL significantly improves over the state of the art (e.g. hopper-medium-replay and hopper-random). To aggregate results across multiple domains, we consider the average of the normalized scores as a proxy, and observe that MOReL significantly outperforms prior algorithms.

Importance of Pessimistic MDP To highlight the importance of P-MDP, we again consider the Pure-partial dataset outlined above. We compare MOReL with a naiv̈e MBRL approach that first learns a dynamics model using the offline data, followed by running model-based NPG without any safeguards against model inaccuracy. The results are summarized in Figure 3. We observe that the naiv̈e MBRL approach already works well, achieving results comparable to prior algorithms like BCQ and BEAR. However, MOReL clearly exhibits more stable and monotonic learning progress. This is particularly evident in Hopper-v2, HalfCheetah-v2, and Walker2d-v2, where an uncoordinated set of actions can result in the agent falling over. Furthermore, in the case of naiv̈e MBRL, we observe that performance can quickly degrade after a few hundred steps of policy improvement, such as in case of Hopper-v2, HalfCheetah-v2 and Walker2d-v2. This suggests that the learned model is being over-exploited. In contrast, with MOReL, we observe that the learning curve is stable and nearly monotonic even after many steps of policy improvement.

Quality of logging policy

Section 4.1 indicates that it is not possible for any offline RL algorithm to learn a near-optimal policy when faced with support mismatch between the dataset and optimal policy. To verify this experimentally for MOReL, we consider two datasets (of the same size) collected using the Pure strategy. The first uses a partially trained policy $\pi_{p}$ (called Pure-partial), which is the same as the Pure dataset studied in Table 1. The second dataset is collected using an untrained random Gaussian policy (called Pure-random). Table 3 compares the results of MOReL using these two datasets. We observe that the value of policy learned with Pure-partial dataset far exceeds the value with the Pure-random dataset. Thus, the value of policy used for data logging plays a crucial role in the performance achievable with offline RL.

Transfer from P-MDP to environment Finally, we study how the learning progress in P-MDP relates to the progress in the environment. Our theoretical results (Theorem 1) suggest that the value of a policy in the P-MDP cannot substantially exceed the value in the environment. This makes the value in the P-MDP an approximate lower bound on the true performance, and a good surrogate for optimization. In Figure 4, we plot the value or return of the policy in the P-MDP and environment over the course of learning. Note that the policy is being learned in the P-MDP, and as a result we observe a clear monotonic learning curve for value in the P-MDP, consistent with the monotonic improvement theory of policy gradient methods . We observe that the value in the true environment closely correlates with the value in P-MDP. In particular, the P-MDP value never substantially exceeds the true performance, suggesting that the pessimism helps to avoid model exploitation.

Conclusions

We introduced MOReL, a new model-based framework for offline RL. MOReL incorporates both generalization and pessimism (or conservatism). This enables MOReL to perform policy improvement in known states that may not directly occur in the static offline dataset, but can nevertheless be predicted using the dataset by leveraging the power of generalization. At the same time, due to the use of pessimism, MOReL ensures that the agent does not drift to unknown states where the agent cannot predict accurately using the static dataset.

Theoretically, we obtain bounds on the suboptimality of MOReL which improve over those in prior work. We further showed that this suboptimality bound cannot be improved upon by any offline RL algorithm in the worst case. Experimentally, we evaluated MOReL in the standard continuous control benchmarks in OpenAI gym and showed that it achieves state of the art results. The modular structure of MOReL comprising of model learning, uncertainty estimation, and model-based planning allows the use of a variety of approaches such as multi-step prediction for model learning, abstention for uncertainty estimation, or model-predictive control for action selection. In future work, we hope to explore these directions.

Acknowledgements

The authors thank Prof. Emo Todorov for generously providing the MuJoCo simulator for use in this paper. Rahul Kidambi thanks Mohammad Ghavamzadeh and Rasool Fakoor for pointers to related works and other valuable discussions/pointers about offline RL. Aravind Rajeswaran thanks Profs. Sham Kakade and Emo Todorov for valuable discussions. The authors also thank Prof. Nan Jiang and Anirudh Vemula for pointers to related work. Rahul Kidambi acknowledges funding from NSF Award $\text{CCF}-1740822$ and computing resources from the Cornell “Graphite” cluster. Part of this work was completed when Aravind held dual affiliations with the University of Washington and Google Brain. Aravind acknowledges financial support through the JP Morgan PhD Fellowship in AI. Thorsten Joachims acknowledges funding from NSF Award $\text{IIS}-1901168$. All content represents the opinion of the authors, which is not necessarily shared or endorsed by their respective employers and/or sponsors.

Broader Impact

This paper studies offline RL, which allows for data driven policy learning using pre-collected datasets. The ability to train policies offline can expand the range of applications where RL can be applied as well as the sample efficiency of any downstream online learning. Since the dataset has already been collected, offline RL enables us to abstract away the exploration or data collection challenge. Safe exploration is crucial for applications like robotics and healthcare, where poorly designed exploratory actions can have harmful physical consequences. Avoiding online exploration by an autonomous agent, and working with a safely collected dataset, can have the broader impact of alleviating safety challenges in RL. That said, the impact of RL agents to the society at large is highly dependent on the design of the reward function. If the reward function is designed by malicious actors, any RL agent, be it offline or not, can present negative consequences. Therefore, the design of reward functions requires checks, vetting, and scrutiny to ensure RL algorithms are aligned with societal norms.

References

Appendix A Theoretical Results: Proofs For Section 4.1

In this section, we present the proofs of our main results Theorem 1 and Proposition 4.

We wish to show the following two inequalities.

The proof of this theorem is inspired by the simulation lemma of , with some additional modifications due to pessimism, and goes through the pessimistic MDP ${\mathcal{M}}_{p}$, which is the same as $\hat{\mathcal{M}}_{p}$ except that the starting state distribution is $\rho_{0}$ instead of $\hat{\rho}_{0}$ and the transition probability from a known state-action pair $(s,a)$ is $P(s^{\prime}|s,a)$ instead of $\widehat{P}(s^{\prime}|s,a)$. More concretely, ${\mathcal{M}}_{p}$ is described by $\{S\cup\textrm{HALT},A,r_{p},{P}_{p},{\rho}_{0},\gamma\}$, where HALT is an additional absorbing state we introduce similar to what we did for $\hat{\mathcal{M}}_{p}$. The modified reward and transition dynamics are given by:

The main idea is to couple the evolutions of any given policy on the pessimistic MDP ${\mathcal{M}}_{p}$ and the model-based pessimistic MDP $\hat{\mathcal{M}}_{p}$ so that $(s_{t-1},a_{t-1})\stackrel{{\scriptstyle\textrm{def}}}{{=}}(s_{t-1}^{{\mathcal{M}}_{p}},a_{t-1}^{{\mathcal{M}}_{p}})=(s_{t-1}^{\hat{\mathcal{M}}_{p}},a_{t-1}^{\hat{\mathcal{M}}_{p}})$.

Assuming that such a coupling can be performed in the first step, since $\left\|P(s,a)-\hat{P}(s,a)\right\|_{1}\leq\alpha$, this coupling can be performed at each subsequent step with probability $1-\alpha$. The probability that the coupling is not valid at time $t$ is at most $1-(1-\alpha)^{t}$. So the total difference in the values of the policy $\pi$ on the two MDPs can be upper bounded as:

For the second part, consider any policy $\pi$ and let it evolve on the MDP $\mathcal{M}$ as $\left(s,a,s^{\prime}_{\mathcal{M}}\right)$. Simulate an evolution of the same policy $\pi$ on ${\mathcal{M}}_{p}$, $\left(s,a,s^{\prime}_{{\mathcal{M}}_{p}}\right)$, as follows: if $(s,a)\in SA_{k}$, then $s^{\prime}_{{\mathcal{M}}_{p}}=s^{\prime}_{\mathcal{M}}$ and if $(s,a)\in\mathcal{U}$, then $s^{\prime}_{{\mathcal{M}}_{p}}=\textrm{HALT}$. We see that the rewards obtained by $\pi$ on each transition in ${\mathcal{M}}_{p}$ is less than or equal to that obtained by $\pi$ on the same transition in $\mathcal{M}$. This proves the second part of the lemma. ∎

The proof is rather straightforward. We have

We consider the MDP in Figure 5, where we set $k=10\log\frac{1}{1-\gamma}$. The MDP has $k+1$ states, with three actions $a_{1},a_{2}$ and $a_{3}$ at each state. The rewards (shown on the transition arrows) are all except for the action $a_{1}$ taken in state $k+1$, in which case it is $1$. Note that the rewards can be scaled by $R_{\textrm{max}}$ but for simplicity, we consider the setting with $R_{\textrm{max}}=1$. It is clear that the optimal policy $\pi^{*}$ is to take the action $a_{1}$ in all the states. The starting state distribution $\rho_{0}$ is state $1$ with probability $p_{0}\stackrel{{\scriptstyle\textrm{def}}}{{=}}\frac{\epsilon}{(1-\gamma)\log\frac{1}{1-\gamma}}$ and state $k+1$ with probability $1-p_{0}$. The actions taken by the data collection policy are shown in blue. Since the dataset consists only of (state, action, reward, next state) pairs $(1,a_{1},0,2),(2,a_{2},0,1)$ and $(k+1,a_{1},1,k+1)$ we see that $\mathcal{U}_{D}=(S\times A)\setminus\left\{(1,a_{1}),(2,a_{2}),(k+1,a_{1})\right\}$ and $d^{\pi^{*},\mathcal{M}}(\mathcal{U}_{D})=(1-\gamma)\cdot\sum_{t=1}^{k-1}\gamma^{t}\cdot p_{0}\leq(1-\gamma)\cdot(k-1)\cdot p_{0}\leq\epsilon$ proving the first claim. Since none of the states and actions in $\mathcal{U}_{D}$ are seen in the dataset, after permuting the actions if necessary, the expected time taken by any policy learned from the dataset, to reach state $k+1$ starting from state $1$ is at least $\exp\left(k/5\right)\geq(1-\gamma)^{-2}$. So, the value of any policy $\hat{\pi}$ learned from the dataset is at most $\frac{1-p_{0}}{1-\gamma}+\frac{p_{0}\cdot\gamma^{(1-\gamma)^{-2}}}{1-\gamma}=\frac{1}{1-\gamma}-p_{0}\cdot\frac{1-\gamma^{(1-\gamma)^{-2}}}{1-\gamma}\leq\frac{1}{1-\gamma}-\frac{3p_{0}}{4(1-\gamma)}$, where we used $\gamma\in[0.95,1)$ in the last step. On the other hand, the value of $\pi^{*}$ is at least $\frac{1-p_{0}}{1-\gamma}+p_{0}\cdot\left(\frac{1}{1-\gamma}-k\right)$. So the suboptimality of any learned policy is at least $p_{0}\cdot\left(\frac{3}{4(1-\gamma)}-k\right)=p_{0}\cdot\left(\frac{3}{4(1-\gamma)}-10\log\frac{1}{1-\gamma}\right)\geq\frac{p_{0}}{4(1-\gamma)}$, where we again used $\gamma\in[0.95,1)$ in the last step. Substituting the value of $p_{0}$ proves the proposition. ∎

We first note that the empirical starting distribution $\hat{\rho_{0}}$ satisfies $D_{TV}(\rho_{0},\hat{\rho_{0}})\leq\frac{C}{\rho_{0,\textrm{min}}}\cdot\sqrt{\frac{\log\frac{1}{\delta\rho_{0,\textrm{min}}}}{n}}$, for a large enough constant $C$. This is because for each state $s$ in the support of $\rho_{0}$, its empirical frequency in $\mathcal{D}$ satisfies:

with probability at least $1-\delta\rho_{0,\textrm{min}}$ using Chernoff’s bound, where $C$ is an absolute numerical constant. Using union bound over at most $\frac{1}{\rho_{0,\textrm{min}}}$ states in the support of $\rho_{0}$, we see that with probability at least $1-\delta$, we have $D_{TV}(\rho_{0},\hat{\rho_{0}})\leq\frac{C}{\rho_{0,\textrm{min}}}\cdot\sqrt{\frac{\log\frac{1}{\delta\rho_{0,\textrm{min}}}}{n}}$.

Similarly, for any state action pair $(s,a)$, denoting $n_{(s,a)}$ as the number of times $(s,a)$ appears in $\mathcal{D}$, we have that:

with probability at least $1-{d^{\pi_{b}}_{\textrm{min}}\delta}{}$. Again using a union bound over all state-action pairs in the support of $d^{\pi_{b},\mathcal{M}}(\cdot)$, we see that:

for every $(s,a)$ in the support of $d^{\pi_{b},\mathcal{M}}(\cdot)$ with probability at least $1-\delta$. The asssumption on the size of $n$ then implies that $n_{(s,a)}\geq\frac{d^{\pi_{b}}_{\textrm{min}}\cdot n}{2}$. Using a similar Chernoff bound argument, we see that $D_{TV}(P(\cdot|s,a),\hat{P}(\cdot|s,a))\leq\frac{C}{p_{\textrm{min}}}\cdot\sqrt{\frac{\log\frac{1}{\delta p_{\textrm{min}}d^{\pi_{b}}_{\textrm{min}}}}{n_{(s,a)}}}$ for every $(s,a)$ in the support of $d^{\pi_{b},\mathcal{M}}(\cdot)$ with probability at least $1-\delta$. By choosing $\alpha=\frac{C}{p_{\textrm{min}}}\cdot\sqrt{\frac{\log\frac{1}{\delta p_{\textrm{min}}d^{\pi_{b}}_{\textrm{min}}}}{n_{(s,a)}}}$, we see that $\mathcal{U}\cap\textrm{Supp}(d^{\pi_{b},\mathcal{M}})=\emptyset$ and hence $T_{\mathcal{U}}^{\pi_{b}}=\infty$. By Theorem 1, we have that for any policy $\pi$, we have:

Plugging $\pi=\pi_{b}$ gives us the first assertion and plugging $\pi=\pi^{*}$ and using Lemma 5 gives us the second assertion. ∎

Appendix B Detailed Related Work

Our work takes a model-based approach to offline RL. We review related work pertaining to both of these domains in this section.

Offline RL dates at least to the work of Lange et al. . In this setting, an RL agent is provided access to a typically large offline dataset, using which it has to produce a highly rewarding policy. This has direct applications in fields like healthcare , recommendation systems , dialogue systems , and autonomous driving . We refer the readers to the review paper of Levine et al. for an overview of potential applications. On the algorithmic front, prior work in offline RL can be broadly categorized into three groups as described below.

The first approach to offline RL is through importance sampling. In this approach, trajectories from the offline dataset are directly used to estimate the policy gradient, which is subsequently corrected using importance weights. This approach is particularly common in contextual bandits literature where the importance weights are relatively easier to estimate due to the non-sequential nature of the problem. For MDPs, Liu et al. present an importance sampling based off-policy policy gradient method by estimating state distribution weights . The work of Liu et al. also utilizes the notion of pessimism by optimizing only over a subset of states visited by the behavioral policy. They utilize importance weighted policy gradient (with estimated importance weights) to optimize this MDP. However, their work does not naturally capture a notion of generalization over the state space. Moreover, their results require strong assumptions on the data collecting policy in the sense of ensuring support on states visited by the optimal policy. Our framework, MOReL, provides the same guarantees under identical assumptions, but we also show that the performance of MOReL degrades gracefully when these assumptions aren’t satisfied.

Dynamic programming

The overwhelming majority of recent algorithmic work in offline RL is through the paradigm of approximate dynamic programming. In principle, any off-policy algorithm based on Q-learning or actor-critic architectures can be used with a static offline dataset. However, recent empirical studies confirm that such a direct extension leads to poor results due to the challenges of overestimation bias in generalization and distribution shift. To address overestimation bias, prior work has proposed approaches like ensembles of Q-networks . As for distribution shift, the principle approach used is to regularize the learned policy towards the data logging policy . Different regularization schemes, such as those based on KL-divergence and maximum mean discrepancy (MMD), have been considered in the past. Wu et al. perform a comparative study of such regularization schemes and find that they all perform comparably. ADP-based offline RL has also be studied theoretically , with Chen and Jiang providing an information-theoretic lower bound on sample complexity. However, these works again don’t study the impact of support mismatch between logging policy and optimal policy. Finally, a recent line of work focuses on obtaining provably convergent methods for minimizing the (one-step) Bellman error using Duality theory. While they show promising results in continuous control tasks in the online RL setting, their performance in the offline RL setting is yet to be studied.

Model-based RL

The interplay between model-based methods and offline RL has only been sparsely explored. The work of Ross & Bagnell theoretically studied the performance of MBRL in the batch setting. In particular, the algorithm they analyzed involves learning a dynamics model using the offline dataset, and subsequently planning in the learned model without any additional safeguards. Their theoretical results are largely negative for this algorithm, suggesting that in the worst case, this algorithm could have arbitrarily large sub-optimality. In addition, their sub-optimality bounds become pathologically loose when the data logging distribution does not share support with the distribution of the optimal policy. Model-based offline RL methods from a safe policy improvement perspective have also been considered . In contrast to both these works, we present a novel algorithmic framework that constructs and pessimistic MDP, and show that this is crucial for better empirical results and sharper theoretical analysis.

B.2 Advances in Model-Based RL

Since our work utilizes model-based RL, we review the most directly related work in the online RL setting. Classical works in MBRL have focused extensively on tabular MDPs and linear quadratic regulartor (LQR). For tabular MDPs (in the online RL setting), the first known polynomial time algorithms were the model-based algorithms of $E^{3}$ and R-MAX . More recent work suggests that model-based methods are minimax optimal for tabular MDPs when equipped with a wide restart state distribution . However, these works critically rely on the tabular nature of the problem. Since each table entry is typically considered to be independent, and updates to any entry to do not affect other entries, tabular MDPs do not afford any notion of generalization. The metric-$E^{3}$ algorithm aims to overcome this challenge by considering an underlying metric space for state-actions that enables generalization. While this work provides a strong theoretical basis, it does not directly provide a practical algorithm that can be used with function approximation. Our work is perhaps conceptually closest to $E^{3}$ and metric-$E^{3}$ which partitions the state space into known and unknown regions. A cornerstone of MOReL is the P-MDP which partitions the state space into known and unknown regions, as in, $E^{3}$ and R-MAX , but these constructions were not developed to encourage pessimism. However, all of these works primarily deal with the standard (online) RL setting. Our work differs in its focus on offline RL, where we show the P-MDP construction plays a crucial role. Moreover, direct practical instantiations of $E^{3}$ and metric-$E^{3}$ with function approximation have remained elusive.

In recent years, along with an explosion of interest in deep RL, MBRL has emerged as a powerful class of approaches for sample efficient learning. Modern MBRL methods (typically in the online RL setting) can support the use of flexible function approximators like neural networks, as well as generic priors like smoothness and approximate knowledge of physics , enabling the learning of accurate models. Furthermore, MBRL can draw upon the rich literature on model-based planning including model predictive control (MPC) , search based planning , dynamic programming , and policy optimization . These advances in MBRL have enabled highly sample efficient learning in widely studied benchmark tasks , as well as in a number of challenging robotic control tasks like aggressive driving , dexterous hand manipulation , and quadrupedal locomotion . Among these works, the recent work of Rajeswaran et al. demonstrated state of the art results with MBRL in a range of benchmark tasks, and forms the basis for our practical implementation. In particular, our model learning and policy optimization subroutines are extended from the MAL framework in Rajeswaran et al. . However, our work crucially differs from it due to the pessimistic MDP construction, which we show is important for success in the offline RL setting.

Appendix C Additional Experimental Details And Setup

As mentioned before, following recent efforts in offline RL , we consider four continuous control tasks: Hopper-v2, HalfCheetah-v2, Ant-v2, Walker2d-v2 from OpenAI gym simulated with MuJoCo . As normally done in MBRL literature with OpenAI gym tasks , we reduce the planning horizon for the environments to 400 or 500. Similar to , we append our state parameterization with center of mass velocity to compute the reward from observations. Mirroring realistic settings, we assume access to data collected using a partially trained (sub-optimal) policy interacting with the environment. To obtain a partially trained policy $\pi_{p}$ , we run (online) TRPO until the policy reaches a value of $1000$, $4000$, $1000$, $1000$ respectively for these environments. This policy in conjunction with exploration strategies are used to collect the datasets (see below for more details). All our results are obtained by averaging runs of five random seeds (for the planning algorithm), with the seed values being $123,246,369,492,615$. Each of our experiments are run with 1 NVidia GPU and 2 CPUs using a total of 16GB of memory.

C.2 Dynamics Model, Policy Network And Evaluation

We use 2 hidden layer MLPs with 512 (for Hopper-v2, Walker2d-v2, Ant-v2) or 1024 (for HalfCheetah-v2) ReLU activated nodes each for representing the dynamics model, use an ensemble of four such models for building the USAD, and our policy is represented with a 2 hidden layer MLP with 32 tanh activated nodes in each layer. The dynamics model is learnt using Adam and the policy parameters are learnt using model-based NPG steps . We set hyper-parameters and track policy learning curve by performing rollouts in the real environment; these rollouts aren’t used for other purposes in the learning procedure. Similar protocols are used in prior work.

C.3 Description Of Types Of Policies

We build off the experimental setup of . Towards this, we first go over some notation. Firstly, let $\pi_{b}$ represent the behavior policy, $\pi_{r}$ is a random policy that picks actions according to a certain probability distribution (for e.g., Gaussian $\pi_{r}^{g}$/Uniform $\pi_{r}^{u}$ etc.), $\pi_{p}$ a partially-trained policy, which one can assume is better than a random policy in value. Let $\pi_{b}^{u}(q)$ be a policy that plays random actions with probability $q$, and sampled actions from $\pi_{b}$ with probability $1-q$. Let $\pi_{b}^{g}(\beta)$ be a policy that adds zero-mean Gaussian noise with standard deviation $\beta$ to actions sampled from $\pi_{b}$. Consider a behavior policy which, for instance, can be a partially trained data logging policy $\pi_{b}$. We consider five different exploration strategies, each corresponding to adding different kinds of exploratory noise to $\pi_{b}$, as described below.

C.4 Datasets And Exploration Strategies

Pure: The entire dataset is collected with the data logging (behavioral) policy $\pi_{b}$.

Eps-1: $40\%$ of the dataset is collected with $\pi_{b}$, another $40\%$ collected with $\pi_{b}^{u}(0.1)$, and the final $20\%$ is collected with a random policy $\pi_{r}$.

Eps-3: $40\%$ of the dataset is collected with $\pi_{b}$, another $40\%$ collected with $\pi_{b}^{u}(0.3)$, and the final $20\%$ is collected with a random policy $\pi_{r}$.

Gauss-1: $40\%$ of the dataset is collected with $\pi_{b}$, another $40\%$ collected with $\pi_{b}^{g}(0.1)$, and the final $20\%$ is collected with a random policy $\pi_{r}$.

Gauss-3: $40\%$ of the dataset is collected with $\pi_{b}$, another $40\%$ collected with $\pi_{b}^{g}(0.3)$, and the final $20\%$ is collected with a random policy $\pi_{r}$.

C.5 Hyperparameter Selection

Refer to table 4 for details with regards to parameters of MOReL. For all environments and data collection strategies, we learn two-layer MLP based dynamics models with ReLU activations by minimizing the one-step prediction errors using Adam and utilize four of these models for defining the USAD. The negative reward for defining the absorbing unknown state is set as the minimum reward in the dataset $\mathcal{D}$ offsetted by a value that is searched over $\{30,50,100,200\}$.

Ascertaining unknown state-action pairs: In order to ascertain unknown state-action pairs, we compute the model disagreement as: $\text{disc}(s,a)=\max_{i\neq j}||f_{\phi_{i}}(s,a)-f_{\phi_{j}}(s,a)||_{2}$, where, $f_{\phi_{i}}$ and $f_{\phi_{j}}$ are members of the ensemble of learnt dynamics model. Specifically, we compute $\text{disc}(s,a)$ over all state-action pairs that occur in the static dataset $\mathcal{D}$. Next, we can compute the mean $\mu_{d}$, standard deviation $\sigma_{d}$ and the max $m_{d}$ of the disagreements evaluated for every state-action pair occuring in the dataset. Then, we utilize an upper-confidence inspired strategy by defining a threshold $\text{thresh}=\mu_{d}+\beta\cdot\sigma_{d}$. The value of beta is tuned between to $\beta_{\max}=(m_{d}-\mu_{d})/\sigma_{d}$ in steps of $5$. For any model-based rollout encountered during planning, if the discrepancy of the state-action pair at a given timestep exceeds thresh, the rollout is truncated at this timestep and is assigned a large negative reward. We emphasize that for every environment, all hyper-parameters (except for $\beta$) is maintained at the same value across all exploration settings.

With regards to the policy and the planning algorithm, we consider a $(32,32)$ tanh MLP optimized using normalized model-based NPG steps (see, for instance, the work of Rajeswaran et al. for the model-based NPG algorithm). Parameters of model-based NPG is described in table 5.

C.6 Ablation Study with the Pure-partial dataset

C.7 Hyperparameter Guidelines and Ablations

We did not have resources to perform a thorough hyperparamter search, and largely used our intuitions to guide the choice of hyperparameters. We believe that better results are possible with hyperparameter optimization. First, we present the influence of the discrepancy threshold for differentiating known and unknown states. We first define the maximum discipancy in the dataset:

where $\mathcal{D}$ denotes offline dataset, and $f_{i}$ denotes $i^{th}$ dynamics model in the ensemble.

Our general observations and guidelines for hyperparameters are:

A high degree of pessimism makes policy optimization in the P-MDP difficult. The optimization process may be slow or highly noisy. This is due to non-smoothness introduced in the dynamics and reward due to abrupt changes involving early episode terminations. If difficulty in policy optimization is observed in the P-MDP, we recommend considering reducing the degree of pessimism.

With a lack or low degree of pessimism, policy optimization is typically easier, but the performance in the true MDP might degrade. If it is observed that the value in the P-MDP overestimates the value in the true MDP substantially, then we recommend increasing the degree of pessimism.

For the tasks considered in this work, positive rewards indicate progress towards the goal. Most of the locomotion tasks involve forward velocity as the primary component of the the reward term. In these cases, we observed that the choice of reward penalty for going into unknown regions did not play a crucial role, as long as it was $\leq 0$. The degree of influence of this parameter in other environments is yet to be determined, and beyond the scope of our empirical study.