PLAS: Latent Action Space for Offline Reinforcement Learning

As is common in reinforcement learning, we define the environment as a Markov Decision Process (MDP) represented as the tuple $\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},r,\gamma)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $\mathcal{P}:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow[0,1]$ is the transition probability function, $r:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow\mathbb{R}$ is the reward function, and $\gamma$ is the discount factor. The general objective of RL is to find a policy that maximizes the expectation of the return $\smash{G_{t}=\sum_{k=0}^{\infty}\gamma^{k}r(s_{t+k},a_{t+k},s_{t+k+1})}$.

Given a policy $\pi$, the action-value function, or Q-function, is defined as $Q^{\pi}(s,a)=\mathbb{E}_{\pi}[G_{t}|S_{t}=s,A_{t}=a]$. Our method builds on top of the commonly used off-policy actor-critic procedure with a deterministic policy [21, 22]. The Q-function of the deterministic policy $\pi$ is estimated based on the Bellman Operator:

The policy $\pi_{\theta}$ is updated following the Deterministic Policy Gradient [23]:

In this section, we will discuss the objectives for offline reinforcement learning and the limitations of existing methods that build on top of off-policy RL.

In offline RL, we are given a fixed dataset $\mathcal{D}=\{(s_{t},a_{t},r_{t},s_{t+1})_{i}\}$ with a finite number of transitions. The difficulty comes from the fact that the static dataset does not cover the entire state space and action space of the MDP. This is especially true when the state and action spaces are continuous. The objective of offline RL is typically to find the policy that maximizes the cumulative reward during its deployment in the environment. However, the performance of the policy will be limited by our knowledge over the MDP, which is inferred from a limited set of transitions.

Reconsidering this problem, another reasonable objective for offline RL is to maximize the cumulative reward of the MDP under the transitions that have been visited in the dataset. Following [13], we may assume a pessimistic MDP such that $r(s,a)$ is significantly small for any unvisited $(s,a)$. Optimizing under such a pessimistic MDP is an intuitive surrogate objective. In addition, [13] proves that the performance of any policy for such a pessimistic MDP is a lower bound in the true MDP.

An additional reason to optimize for this pessimistic MDP is the extrapolation error of approximated Q-functions [7]. In off-policy algorithms, we bootstrap $Q(s_{t},a_{t})$ by using $Q(s_{t+1},\pi(s_{t+1}))$ according to the Bellman operator as in Equation (1). If $(s_{t+1},\pi(s_{t+1}))$ is not in the dataset, $Q(s_{t+1},\pi(s_{t+1}))$ can be arbitrarily wrong. This error caused by out-of-distribution actions will be accumulated and exacerbated by the policy update. (Note that out-of-distribution states do not occur during training.) Thus, optimizing the policy under the pessimistic MDP is equivalent to forcing the policy to select known actions that avoid any error accumulation. This is the motivation for constraining the policy to be within the support of the dataset. On the other hand, the constraint should not be overly restrictive and should not be affected by the distribution of the dataset as proposed in [15]. The policy should have the full flexibility to choose actions within the support.

Existing offline RL methods enforce such a constraint in different ways. BCQ [7] constrains the policy by sampling from the behavior policy. However, the policy is then restricted by the distribution of the behavior policy. BEAR [15] and BRAC [16] incorporate the constraint on the policy as a regularization term into the optimization process for the policy or the Q-function. This regularization term is calculated by a divergence metric, such as KL-divergence or MMD. There are two practical challenges to these approaches. First, this additional loss term creates a trade-off between optimizing the original objective and satisfying the constraint. Although BEAR uses a Lagrangian multiplier to solve a constrained optimization problem, in practice, the constraint is almost never satisfied [16]. Second, the choice of the divergence metric, such as KL and MMD, might be overly restrictive given datasets that have diverse actions. We provide further discussion on MMD constraint in Appendix E. These past approaches have shown that properly enforcing an explicit policy constraint is difficult; this observation motivates our approach.

Since our method uses a conditional variational autoencoder, we include a brief background of VAE in this section in its most general form based on [24] and [25]. Given a dataset $X=\{x^{(i)}\}_{i=1}^{N}$, the goal of a VAE is to generate samples that are from the same distribution as the data points, in other words, to maximize $p(x)$ for all $x^{(i)}$. This is achieved by introducing a latent variable $z$ sampled from a prior distribution $p(z)$ and modeling a decoder $p_{\theta}(x|z)$ with parameter $\theta$. Directly maximizing the marginal likelihood $p_{\theta}(x)=\int p(z)p_{\theta}(x|z)dz$ is intractable. Instead, Kingma and Welling [24] propose to approximate the true posterior $p_{\theta}(z|x)$ by training an encoder $q_{\phi}(z|x)$. In this way, they derive the following evidence lower bound (ELBO) on the log-likelihood of the data:

The term $\log p_{\theta}(x|z)$ (where $z$ is sampled from $q_{\phi}(z|x)$) represents the reconstruction loss. The second term is the KL-divergence between the encoder output and the prior of $z$, which is usually set to be $\mathcal{N}(0,1)$. Thus, optimizing for this objective enables us to train a model that generates samples similar to the data distribution by sampling $z$ and then passing it into the decoder.

Given a static dataset, we use a conditional variational autoencoder (CVAE) to model the behavior policy $p(a|s)$, as in other recent methods [7, 15, 16]. The CVAE is trained to reconstruct actions conditioned on the states. Converting Equation 3 into our problem formulation, the objective of the CVAE is to maximize $\log p(a|s)$ by maximizing its lower bound:

where $z$ is the latent variable, $\alpha$ and $\beta$ are the parameters of the encoder and the decoder, respectively. This is similar to Equation 3, except that all terms are conditioned on the state $s$. A trained decoder $p_{\beta}(a|s,z)$ provides a mapping from the latent space to the action space, conditioned on the state.

In order to constrain the policy to be within the support of the dataset, we propose to train a deterministic policy $z=\pi(s)$ to map from a state $s$ to a “latent action” $z$; we then use the pretrained decoder $p_{\beta}(a|s,z)$ to project the latent action into the actual action space as shown in Figure 2. This is significantly different from BCQ which samples from a fixed range of the latent space. The latent policy has the flexibility of choosing the latent actions and thus avoids being affected by the density of the dataset distribution. The CVAE is trained to maximize $\log p(a|s)$, equivalent to maximizing the expected output of the decoder $\mathbb{E}_{p(z|s)}[p_{\beta}(a|s,z)]$ as shown:

Thus, for values of $z$ that have a high probability under the prior $p(z|s)$, the decoder $p_{\beta}(a|s,z)$ will output a high probability action under the behavior policy distribution $p(a|s)$ in expectation.

If we constrain the latent policy $z=\pi(s)$ to output a latent action $z$ that has a high probability under the prior $p(z|s)$, then the full policy, formed by $p_{\beta}(a|s,z=\pi(s))$, is likely to have a high probability under the behavior policy $p(a|s)$. Fortunately, this constraint is simple to enforce; since the prior $p(z|s)$ is set to a normal distribution $\mathcal{N}(0,1)$, we simply define $z=\pi(s)$ such that $z_{i}\in[-\sigma,\sigma]$ for each latent dimension $i$ for some hyperparameter $\sigma$ (see Section 4.3 for details). Furthermore, for each state $s$, the latent policy has the flexibility to choose any latent action $z$ in this constrained latent action space, leading to an easier optimization compared to previous work with explicit constraints.

The latent policy provides a natural constraint to stay within the support of the dataset. However, in some cases when the Q-function can generalize well, we may relax the pessimistic objective and allow the policy to select out-of-distribution actions to improve its performance. The benefit of the out-of-distribution actions is more likely to happen when the environment (both transition probabilities and the reward function) is smooth and when the dataset has limited quality and diversity. To do so, we add a perturbation layer to the output of the decoder that outputs a residual over the action. The residual is limited to a specific range $[-\epsilon,\epsilon]$, where $\epsilon$ is a hyperparameter. Mathematically, this enforces the final output action to be close to the actions within the dataset in terms of the $L_{\infty}$ norm. This perturbation layer is inspired by BCQ. However, BCQ forms a policy by sampling from the generative model; the perturbation layer is used to prevent “sampling from the generative model for a prohibitive number of times” [7]. In our case, we don’t perform any sampling since our policy is deterministic; the perturbation layer is instead specifically designed for out-of-distribution generalization. In the experiments, we will demonstrate that when the dataset has enough coverage in the state-action space, this additional layer is not necessary.

The full algorithm is summarized in Algorithm 1. Our algorithm can be built on top of off-policy algorithm such as DDPG [21] or TD3 [22]. We use a deterministic policy $z=\pi(s)$ to output a latent action. The policy uses a tanh activation at the output layer to limit the max latent action. This limit is set to 2 by default, which corresponds to $2\sigma$ for the latent variable $p(z)=\mathcal{N}(0,1)$. We use a soft Clipped Double Q-learning with parameter $\lambda$ to weight the two Q-functions. Following common practice, we use target networks to stabilize training with hyperparameter $\tau$. The code for our algorithm is based on the BCQ repository, and we mostly follow the hyperparameters of BCQ. Further implementation details and hyperparameters can be found in Appendix A.

Real-Robot Experiment: The task for the real-robot experiment is to slide along the edge of the cloth as far as possible with a tactile sensor. The experiment setup is shown in Figure 3(a) and an example of the tactile sensor reading is shown in Figure 3(b). More details of this experiment can be found in Appendix F. The dataset consists of the replay buffer from a previous online RL experiment with around 7000 timesteps of transitions and 5 episodes (around 300 timesteps) of expert demonstrations from a trained policy. We train the policy for 2400 steps in total for each experiment with evaluations every 150 steps over 5 episodes.

Locomotion Datasets: We mainly focus on the locomotion environments from the d4rl datasets in this section including Walker2d-v2, Hopper-v2, and Halfcheetah-v2. For each environment, there are four types of datasets: random, medium, medium expert, and medium replay datasets. Random datasets are generated by randomly initialized policies. Medium datasets are generated from rollouts of a “medium” performance policy trained with Soft Actor-Critic up to a certain performance. Medium-expert datasets are generated by combining the medium datasets and expert datasets. Medium-replay datasets are the replay buffers created during the training of the medium policies. Note that medium-replay datasets are much smaller than the other types of datasets, making it more challenging to obtain stable training performance. We train the policy for 500 epochs and each epoch has 1000 training steps. The policy is evaluated every 1 epoch over 10 episodes.

Figure 3(c) shows the evaluation performance of different methods across the training process for the cloth sliding task. The brown dashed line indicates the behavior cloning (BC) policy. It fails as expected because the average quality of the dataset is poor. The “Offline-TD3” baseline is to directly run TD3 over the offline dataset without any extra online data collection. The performance is even worse than BC, which shows the necessity of designing offline RL algorithms that can utilize fixed datasets for real-world applications. BCQ also doesn’t perform well, possibly because it is overly constrained by the dataset distribution, resulting in similar performance as behavior cloning. BEAR achieves reasonable performance at the beginning of training but it drops soon after. Our method outperforms all the baselines, and the final performance is similar to the expert policy.

To more systematically benchmark the performance of our method with the other offline RL algorithms, we ran experiments on the d4rl benchmarks. We focus the discussions and analysis on the locomotion environments here; full results on the d4rl datasets can be found in Appendix B including Locomotion, Maze2d, AntMaze, Adroit Hand, Franka Kitchen environments. To highlight the performance of our method over the medium-expert datasets and the medium-replay datasets, we include the training curves in Figure 4. These two types of datasets are especially important because they have diverse coverage over states and actions generated by a mixture of policies. These kinds of diverse datasets are also more likely to appear in real-world applications with data collected from different sources. The diversity allows the potential to learn a good policy from the data; on the other hand, diversity also introduces difficulties for the policy constraints. Figure 4 shows that we consistently achieve performance that is similar to or better than the best baselines on these datasets, demonstrating the effectiveness of our method in fully utilizing the datasets.

We analyze the quality of the learned Q-function in detail with the Walker2d medium-expert dataset for different algorithms. By definition, the Q-value $Q_{\pi}(s_{t},a_{t})$ is equal to the expected return starting from state $s_{t}$ following action $a_{t}$; our learned Q-function attempts to estimate this value. Thus, we evaluate the Q-values by comparing them to the true returns for the transitions during rollouts. The true return is calculated by the cumulative discounted reward until termination or up to $r_{t+1000}$, since the reward after 1000 steps is negligible due to the discount factor. We define the estimation error to be $Q(s_{t},a_{t})-G(s_{t},a_{t})$, where $G(s_{t},a_{t})$ is the empirical return. Positive error corresponds to overestimation bias and negative error corresponds to underestimation bias. In Figure 5, we show four metrics on the quality of the Q-function over $N$ transitions from 50 episodes:

• •

  Mean Squared Error (MSE) - measures the overall quality: $\sum_{i}(Q(s_{i},a_{i})-G(s_{i},a_{i}))^{2}$)

• •

  Positive Error Percentage - percentage of overestimation: $\frac{1}{N}\sum_{i}(\mathbbm{1}(Q(s_{i},a_{i})-G(s_{i},a_{i}))>0))$

• •

  Positive Error Mean - the mean value of the positive errors, which indicates the average magnitude of over-estimation: Average of $Q(s_{i},a_{i})-G(s_{i},a_{i})$ for $Q(s_{i},a_{i})-G(s_{i},a_{i})>0$

• •

  Negative Error Mean - the mean value of the negative errors, which indicates the average magnitude of under-estimation: Average of $Q(s_{i},a_{i})-G(s_{i},a_{i})$ for $Q(s_{i},a_{i})-G(s_{i},a_{i})<0$

MSE measures the overall estimation bias, and the other three metrics capture the direction of the bias. Our method achieves consistently low MSE during training compared with the baselines. Note that BRAC has low MSE during the beginning of training because the return is close to 0, as shown in the training curves in Figure 4. Although both our method and BRAC achieve similar performance at the end of training on the Walker2d medium-expert dataset (though our method converges faster), MSE indicates that our method results in a better Q-function. In terms of the direction of the bias, BEAR has a large overestimation bias and BRAC has a large underestimation bias. Overestimation bias is usually considered more harmful than underestimation for Q-learning based algorithms [22]. As a result, although BRAC has a higher MSE than BEAR, the evaluation performance is still better. Our method does not have significant underestimation or overestimation in this case.

In Section 3.2, we mentioned that when the Q-function generalizes well, allowing the policy to select some out-of-distribution actions might be helpful. This motivates us to introduce an additional perturbation layer as mentioned in Section 4.2. We evaluate the benefit of this optional perturbation layer with different max perturbation limits, with $\epsilon=0$ being the latent policy alone. The results for a selective set of environments are summarized in Table 1. Note that the action space in these tasks is defined to be $(-1,1)$; thus $\epsilon=0.5$ allows a very high range of perturbation. We found that the importance of the perturbation layer depends on both the dataset and the environment. Allowing out-of-distribution actions often leads to improved performance for random datasets. The “medium” datasets tend to have peak performance with smaller values of $\epsilon$; a larger value likely leads to errors in the Q-function evaluation due to out-of-distribution state-action pairs. The full results including medium-expert and medium-replay datasets are in Appendix D. We found that medium-expert and medium-replay datasets usually do not benefit from the perturbation layer.