Towards Offline Reinforcement Learning with Pessimistic Value Priors

Abstract

Offline reinforcement learning (RL) seeks to train agents in sequential decision-making tasks using only previously collected data and without directly interacting with the environment. As the agent tries to improve on the policy present in the dataset, it can introduce distributional shift between the training data and the suggested agent’s policy which can lead to poor performance. To avoid the agent assigning high values to out-of-distribution actions, successful offline RL requires some form of conservatism to be introduced. Here we present a model-free inference framework that encodes this conservatism in the prior belief of the value function: by carrying out policy evaluation with a pessimistic prior, we ensure that only the actions that are directly supported by the offline dataset will be modelled as having a high value. In contrast to other methods, we do not need to introduce heuristic policy constraints, value regularisation or uncertainty penalties to achieve successful offline RL policies in a toy environment. An additional consequence of our work is a principled quantification of Bayesian uncertainty in off-policy returns in model-free RL. While we are able to present an implementation of this framework to verify its behaviour in the exact inference setting with Gaussian processes on a toy problem, the scalability issues that it suffers as the central avenue for further work. We address in more detail these limitations and consider future directions to improve the scalability of this framework beyond the vanilla Gaussian process implementation, proposing a path towards improving offline RL algorithms in a principled way.

Appendices

Appendix 1

Prior mean and covariance derivation

The prior mean and covariance of the observed rewards can be deduced by writing $R_i = \sum _k A_{ik}Q(x_i^{(k)})$, with the same definition of A and x as given in Eq. 4 and 5 and then considering $\mathbb {E}(R_i)$ and $\text {cov}(R_i,R_j)$.

For the mean, we have

$$\begin{aligned} \mathbb {E}(R_i) &= \mathbb {E}\sum _k A_{ik}Q(x_i^{(k)}) \end{aligned}$$

(11)

$$\begin{aligned} & = \sum _k A_{ik}\mathbb {E}Q(x_i^{(k)}) \end{aligned}$$

(12)

$$\begin{aligned} & = 0, \end{aligned}$$

(13)

and for covariance

$$\begin{aligned} \text {cov}(R_i,R_j) &= \text {cov}\left( \sum _k A_{ik}Q(x_i^{(k)}),\sum _jA_{jl}Q(x_j^{(l)})\right) \end{aligned}$$

(14)

$$\begin{aligned} &= \sum _{k,l}A_{ik}A_{jl}\text {cov}(Q(x_i^{(k)}),Q(x_j^{(l)})) \end{aligned}$$

(15)

$$\begin{aligned} &= \sum _{k,l}A_{ik}A_{jl} k_Q(x_i^{(k)},x_j^{(l)}), \end{aligned}$$

(16)

as required.

Validity of resulting kernel

To show that k is a valid kernel for any policy, we will deduce that the matrix with elements i, j given by

$$\begin{aligned} k(\hat{x}_i,\hat{x}_j) = k_{ij} = \sum _{0\le k,l \le n} A_{ik}A_{jl} k_Q(x_{i}^{(k)},x_{j}^{(l)}) \end{aligned}$$

(17)

is positive semidefinite (and, trivially, symmetric in i and j) for any given valid kernel $k_Q$. This is equivalent to the statement that for any vector with elements $v_i$, the dot product $\sum _{i,j}v_i k_{ij} v_j \ge 0$ given that for any $x'$ and $v'$, $\sum _{m,n}v'_mk_Q(x'_{m},x'_{n})v'_n \ge 0$.

$$\begin{aligned} \sum _{i,j}v_i k_{ij} v_j &= \sum _{i,j, k,l}v_i A_{ik}v_jA_{jl} k_Q(x_{i}^{(k)},x_{j}^{(l)}) \end{aligned}$$

(18)

$$\begin{aligned} &= \sum _{(i,j), (k,l)}(v_i A_{ik})(v_jA_{jl}) k_Q(x_{i}^{(k)},x_{j}^{(l)}) \end{aligned}$$

(19)

$$\begin{aligned} &= \sum _{m,n}v'_m v'_n k_Q(x'_{m},x'_{n}) \end{aligned}$$

(20)

$$\begin{aligned} &\ge 0 \end{aligned}$$

(21)

as required, where we have ‘unrolled’ the matrix $v_i A_{ik}$ into the vector $v'_m$ with $x_i^{(k)}$ the corresponding $x'_m$ values and replaced summation over i, k with summation over m and correspondingly summation over j, l with n.

Appendix 2

The GP employed has prior mean equal to 0 and base covariance for the latent (value) variables that factors into an RBF function term in state space and a Kronecker-delta function in action-space (as actions are discrete we assume independence in the value across the different actions) analogously to Engel et al. [2005]:

$$\begin{aligned} \text {cov}(Q(s_i,a_i),Q(s_j,a_j)) = k_Q((s_i,a_i),(s_j,a_j)) = \sigma _p^2\exp (\frac{(s_i-s_j)^2}{2l^2})\delta _{a_i a_j}, \end{aligned}$$

(22)

where we chose $\sigma _p^2 = 1$ and $l=0.25$ for the experiments presented.

The DQN trained to produce Fig1b has two hidden layers of 256 neurons, batch size 128, Adam optimiser [Kingma and Ba, 2014] with learning rate of 0.001 and was trained for 100k gradient steps.

Here we report the same results from Figs. 1c and 1d but for a wider range of the state-space, where we observe that even outside the main region of interest, in the regions where no data is present the posterior value mean and standard deviation tend to the prior mean and standard deviation, i.e. low value and higher standard deviation, as is desirable in the offline RL setting for unsupported regions.