Improved Robustness and Safety for Pre-Adaptation of Meta Reinforcement Learning with Prior Regularization

It is important for trained Artificial Intelligence (AI) to adapt to a new task relatively quickly with a small amount of data. Although Deep Reinforcement Learning (DRL) methods were shown to be powerful for complicated sequential decision making problems, they usually learn a separate policy for each specific task. In recent years, researchers tried to design AI that can solve multiple tasks, by either switching to the corresponding policy it already learned[1], or learning and adapting to the new task. The first approach requires that we can enumerate the tasks, which may not be possible or practical for some real-world applications. Moreover, learning to master each and every task often requires a lot of training data, making it data inefficient.

The second approach, which is to enable quick learning/adaptation as new tasks arise, is named Meta Reinforcement Learning (Meta-RL)[2]. Meta-RL techniques aim to gain “general learning” for better dealing with new tasks by leveraging learned experience from previous related tasks. This requires that these tasks share some common features. For instance, in autonomous driving, merging into a highway from a ramp and entering a round-about both involve the Autonomous Vehicle (AV) agent choosing a gap, adjusting its velocity, and making the move, while the exact road geometry and behaviors of the adjacent vehicles may vary.

How to leverage learning for a family of tasks sharing similar structures, instead of simply doing well in a single task is still an open topic. In [3, 4, 5], the agent is trained to find the parameters $\theta$ of its policy, so that when the agent takes a few gradient steps on that $\theta$, it will converge toward $\theta^{*}$ which is optimal for a new task. In this case, the common structure of different tasks is implicitly represented by the initial parameter $\theta$. Another approach is to learn latent contexts among tasks and the context acts as part of the input to the policy network. For example, in [6] Rakelly et al. developed a data efficient off-policy Meta-RL method called Probabilistic Embeddings for Actor-critic meta-RL (PEARL), performing online probabilistic filtering of the latent task variables to infer how to solve a new task from a small amount of new data. While the off-policy nature of PEARL enable data efficiency when adapting to a new task, it does require on-line collecting the context and corresponding latent variables of a new task with prior policy (the policy conditioned on the prior distribution). As a result, the performance is not guaranteed when encountering a new task for the first time. This poses a serious challenge for applications like autonomous driving where the agent cannot afford to perform very poorly. In other words, it is important for the prior (pre-adaptation) policy to maintain robustness with respect to task distribution shift.

Safety in RL is not a new problem and a lot of research has been done [7, 8], but the safety of Meta-RL before adaptation has not received enough attention. In this paper, we consider a policy as safe if it would not lead to a priorly-defined unsafe state. Robustness in meta-learning has also been studied in multiple recent works. The work in [9] explores robustness to perturbations in the task samples, [10] deals with imbalances in the number of samples per task instance, and [11] proposes a solution to improve the robustness to shifts in the task distribution between meta-training and meta-testing. In this paper our discussion is on ’pre-adaptive task-robustness’ in the sense that it performs well before adaptation despite the task-distributional shift between prior assumption and the real. All of our experimental studies are safety-critical cases and we will consider a prior policy pre-adaptive task-robust if its ‘safety violation’ rate does not exceeds an absolute rate of 0.1%. We develop PEARL+, a Meta-RL method based on PEARL[6] to improve safety and robustness for pre-adaptation of Meta-RL. We introduce a prior regularization term and a new Q-network, which doesn’t back-propagate into the context encoder, to explicitly consider the robustness and safety of the prior policy without sacrificing after-adaptation performance.

The rest of the paper is organized as follows. We first describe the related work in Section II. Then the preliminaries of Meta-RL and the PEARL method will be introduced in Section III. In Section IV, we present the details of our new method PEARL⁺. We study the performance of PEARL⁺ on three applications and the results are shown in Section V. Finally, the paper is concluded in Section VI.

The current state-of-the-art Meta-RL algorithms can be categorized into four types.

RNN-based Meta-RL. In these methods (e.g., [2][12]), a Recurrent Neural Network (RNN) structure is used, the experience from previous tasks are preserved in the RNN parameters. The agent starts with no memory/knowledge when it starts to adapt to a new task. However, RNN based policy has no mathematical convergence proof and is believed to be vulnerable to meta-overfitting [13].

Gradient-based Meta-RL. This type of Meta-RLs learn a policy initialization that can quickly adapt to a new task after few policy gradient steps. They usually separate learning into two loops: inner loop for policy adaptation, and outer loop for policy updates[14]. In most existing work, optimizations are through gradient descents for both loops[3][5]. One notable exception is [15], which implements behavior cloning in the inner loop to reduce the need for task exploration. Due to their on-policy nature, these methods usually need a lot of data and parameter-tuning[13].

Model-based Meta-RL These methods learn a predictive model, represented by a deep neural network, and find model parameters that are sensitive to variations among the tasks[16][17]. The action to take is then determined through optimization methods such as model-predictive control (MPC). Due to their model-based nature, these methods require less data and can adapt more quickly. However, they require relatively high computation cost for the online optimization process to get actions[13].

Context-based Meta-RL In this category, the meta-agents learn the latent variables of a distribution of tasks, which the policy uses to adapt to new tasks [6, 18, 19, 20, 21]. Our method falls into this category. It is worth noting that the approach in this category allows off-policy learning. That is, the RL learning does not require the latest updated policy to interact with the environment. Rather, it can leverage experiences from other policies, such as learning from replaying buffer samples from old policy as used in [19]. Examples of this category include [20][21] where their off-policy Meta-RL algorithms were developed by decoupling the task inference from the policy training.

Regardless of the categories, most of the meta RL literature has focused on data efficiency and fast adaptability. However, in some applications, the performance of the prior policy (i.e. prior to adaptation) may be critical to avoid failures as these cases cannot tolerate mischief. For example, for automated vehicles, the factory installed default setting can be considered as the prior policy, while vehicles may be sold and driven in different cities/countries with a different driving culture/roadmanship, which can be considered as different tasks (MDPs). In this case, the safety performance of the prior policy (factory setting) cannot be overlooked. We were not aware of any Meta-RL methods that explicitly considered the safety of the prior policy at the beginning of the meta testing, or when it is first exposed to a new task in the real-world. This lack of assurance of robust initial performance is a concern of Meta-RL application to safety-critical tasks.

In this section, we will briefly review the formulation of our meta reinforcement learning problem, and introduce the PEARL algorithm which we build upon and enhance.

Most of the Meta-RL algorithms proposed in the literature focused on optimizing the performance adapting to specific tasks, but did not explicitly consider the performance before adaptation. For some applications like robots and automated vehicles, any fatal mistakes are unacceptable, including before adaptation. In this section, we’ll explain details of our algorithm PEARL⁺, which is a direct enhancement to PEARL.

In PEARL’s configuration, the prior over context variable $z$ is denoted as $z_{0}$, which is assumed to follow an identical independent distribution. However, this assumption does not hold for real tasks and environments. Since the previous formulation (Eq. 4) only optimizes on the posterior return, the policy trained in this way doesn’t perform well under on the prior assumption before adaptation with regard to the distribution shift. We propose to solve the problem by introducing a prior regularization term, to calculate which we create a new Q-Network for state-action value estimation under the prior context assumption.

The prior regularization is formulated as the expected return depending on the random context variable following the prior unit Gaussian distribution $p(\mathbf{z_{0}})$. The objective function of PEARL⁺ can then be formulated as:

where $\alpha>0$ is the trade-off coefficient.

The policy is trained to maximize the return which considers both prior (before adaptation) and posterior (after adaptation) of the context variable. Similar to the exploration-exploitation trade-off: emphasizing prior regularization performance will decrease performance of the posterior context. Note that the prior regularization term does not depend on the inference network. Therefore, it will not affect the training of the context encoder $q_{\phi}$. However, we can train the policy to be more robust by considering the mismatch with the real context.

Based on the structure of PEARL, we introduce a new Q-network $\tilde{Q}_{\theta}$ to recover the action-state value with a prior context assumption $\mathbf{z_{0}}$. As shown in Fig. 1, we update the $\tilde{Q}_{\theta}$ network by gradients computed from loss $\tilde{\mathcal{L}}_{critic}$ (red line) with data sampled from the whole replay buffer. Similar to PEARL, the policy of PEARL⁺ $\pi_{\theta}$ is also built on top of the Soft Actor Critic (SAC) algorithm, but now has two optimization sources (both shown in green lines): the original $\mathcal{L}_{actor}(\mathbf{z})$ computed with the posterior context, and an additional $\mathcal{L}_{actor}(\mathbf{z_{0}})$ computed with the prior context. They share the same set of data sampled from the whole replay buffer. The encoder is optimized by $\mathcal{L}_{critic}$ (yellow line) and $D_{KL}$ (purple line), which is the same as PEARL.

We summarize the meta-training procedure of PEARL⁺ in Algorithm 1. The meta-testing procedure is the same as PEARL, and therefore will not be described in this paper.

In this section, we introduce the three experiments used to examine the performance of the PEARL⁺ algorithm. All three examples are in the field of robotics and autonomous vehicles, and their prior safety performance (after training, but before adaptation to a new task) is important.

In this paper, we propose a Meta-RL algorithm that considers pre-adaptive safety and robust performance to unseen tasks. This is an essential feature for safety-critical real-world applications. By introducing an additional regularization term over the expected return of the policy with prior context assumption, and a corresponding additional critic network that returns the Q value conditioned on the prior context, our algorithm PEARL⁺ achieves superior performance (mostly through better safety reward) when it is exposed to a new task for the first time. In addition, it preserves the fast adaptation of PEARL, the off-policy Meta-RL method our algorithm builds upon. Our algorithm has been validated using three example applications: two MuJoCo robotic tasks and a highway merge autonomous driving task. These experiments show that PEARL⁺ achieves better robustness and safety performance, and the ablation study demonstrates the necessity of introducing an extra Q-Network. Furthermore, the prior regularization term is not restricted to context-based Meta-RL algorithms, and actually can be applied to any other Meta-RL frameworks, which would be a good direction for future work