Off-Policy Meta-Reinforcement Learning Based on Feature Embedding Spaces

Deep reinforcement learning (DRL) has shown superhuman performance in several domains, such as computer games and board games (Silver et al. 2017; Berner et al. 2019). However, conventional DRL considers only learning for a single task and does not reuse experience from past tasks. This is a major reason for the sample inefficiency in conventional DRL.

To overcome this problem, the concept of meta-learning has been proposed (Schmidhuber, Zhao, and Wiering 1996). Meta-learning is a class of methods for learning how to efficiently learn with a small amount of data on a new task by utilizing previous experience. An instance of meta-learning consists of two phases: meta-training and meta-testing. In meta-training, the agent prepares itself for learning in meta-testing by using some training tasks. In meta-testing, the agent is evaluated on the basis of its performance on the task to be solved. Although meta-learning aims to improve sample efficiency in meta-testing, the sample efficiency in meta-training is also important from the perspective of computational cost (Mendonca et al. 2019; Rakelly et al. 2019).

One type of meta-learning methods is called gradient-based meta-learning, and they include MAML (Finn, Abbeel, and Levine 2017) and Reptile (Nichol, Achiam, and Schulman 2018). These methods learn to reduce the loss of the model whose parameters (e.g., weights of neural networks) are updated over several steps. Finn, Abbeel, and Levine (2017) have shown that the sample efficiency of MAML in meta-testing is improved compared with that of naive pretraining methods. However, most reinforcement learning (RL) applications of these methods work with on-policy methods (Mendonca et al. 2019), which are less sample efficient than off-policy methods because on-policy methods cannot reuse the data collected by old policies. In addition, the performance of the learned initial parameters in meta-testing can be very poor in some cases until the parameters are updated. For example, there are tasks where an agent aims to reach a goal as fast as possible, and the tasks differ in terms of goal positions. Let us assume that there are two tasks whose goals are in opposite directions from the initial position of the agent; then, the well-trained policies for the tasks require contradicting actions. Thus, in this case, the performance is very poor before the parameters are updated.

Another type of meta-learning methods is called context-based meta-learning. PEARL (Rakelly et al. 2019) is a state-of-the-art meta-RL method that belongs to this type. PEARL learns how to infer task information in meta-training and uses it in meta-testing. Because of this inference, PEARL generally needs less data to improve its performance in meta-testing than methods that update the parameters of neural networks when the tasks in meta-training and meta-testing are similar. In addition, the policy and Q-function in PEARL are trained in an off-policy manner, which can further improve its sample efficiency. Indeed, PEARL was shown to be more sample efficient than MAML (Rakelly et al. 2019).

In this paper, we extend the idea of PEARL and propose a novel meta-RL method, embedding learning and evaluation of uncertainty (ELUE), which has the following features:

Off-policy embedding training

In PEARL, policy training is based on an off-policy method, but the training for task embedding, which is used for calculating distributions over tasks, depends on what the current policy is, i.e., it is on-policy. By removing the dependency on the current policy from the task embedding learning, we propose a fully off-policy method. Thanks to policy-independent embedding, the training objective is expected to be stable, and the data collected by past policies can be reused.

Policy and Q-function conditioned by beliefs over tasks

PEARL introduces distributions over tasks, i.e. beliefs, but both its policy and Q-function depend on a task variable sampled from the distribution. After the task variable is sampled, the variable contains no information on the uncertainty over the tasks. By contrast, ELUE conditions the policy and Q-function on the belief over tasks, which can be used to evaluate uncertainty. The use of such policy and Q-function has been shown lead to more precise exploration (Humplik et al. 2019; Zintgraf et al. 2020). In addition, for learning these functions, we apply the information bottleneck objective (Tishby, Pereira, and Bialek 1999) for generalization, instead of naively maximizing the cumulative reward.

Combination of belief and parameter update

PEARL does not update the parameters of the policy and Q-function in meta-testing. When there are large differences between the tasks in meta-training and in meta-testing, the pretrained belief model, policy, and Q-function may no longer be useful, and PEARL can fail to improve the performance. To alleviate this drawback, our method performs not only inference but also parameter update.

We compare the sample efficiency of state-of-the-art meta-RL methods and ELUE through experiments in the Meta-World (Yu et al. 2019) environment and show that ELUE performs better than those methods.

Markov decision processes (MDPs) are models for reinforcement learning (RL) tasks. An MDP is defined as a tuple $(\mathcal{S},\mathcal{A},T,R,\rho)$, where $\mathcal{S}$ and $\mathcal{A}$ are the state and the action spaces respectively, $R:\mathcal{S}\times\mathcal{A}\times\mathbb{R}\rightarrow[0,1]$ is a reward function that determines the probability of the reward amount and $T:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow[0,1]$ is a transition function that determines the probability of the next state. $\rho$ is the initial state distribution. Let us denote the policy as $\pi$, which determines the probability of choosing an action at each state. The objective of RL is to maximize the expected cumulative reward by changing the

We assume that each task in meta-training and meta-testing can be represented by an MDP, where $\mathcal{S}$ and $\mathcal{A}$ are the same among tasks. In addition, we consider tasks to be the same when they differ only in $\rho$ because the difference in $\rho$ does not change the optimal policy. Thus, a different task means different $T$ or $R$.

In our problem, it is assumed that the reward and the transition function are not observable directly. We treat this problem as a partially observable MDP (POMDP) and introduce a probability over $R$ and $T$, which is called a belief (Humplik et al. 2019; Zintgraf et al. 2020). For clarity, let us assume that $R$ and $T$ are parameterized by $\varphi$ and denote them as $R_{\varphi}$ and $T_{\varphi}$. It is known that a POMDP can be transformed into a belief MDP whose states are based on beliefs and that the optimal policy of the belief MDP is also optimal in the original POMDP (Kaelbling, Littman, and Cassandra 1998). We denote a history as $h_{t}:=(s_{0},a_{0},r_{0},s_{1},\dots,s_{t})$, where $s_{\tau}\in\mathcal{S}$, $a_{\tau}\in\mathcal{A}$, $r_{\tau}\in\mathbb{R}$ are the state, the action, and the reward at time $\tau$, respectively. In our problem, a belief at time $t$ is $\tilde{b}_{t}(\varphi):=P(\varphi|h_{t})$, and the state of the belief MDP at time $t$ is $s^{+}_{t}=(s_{t},\tilde{b}_{t})$, which is often called a hyper-state. The objective of our problem is maximizing $\mathbb{E}_{h_{\infty}}[\sum_{t=0}^{\infty}\gamma^{t}r_{t}]$ by changing a policy that is conditioned on a hyper-state, where $\gamma$ is a discounted factor. In our problem, the belief is updated by observations:

However, in general, the exact calculation of this belief update is intractable, so the existing methods approximate beliefs and avoid the direct calculation (Humplik et al. 2019; Zintgraf et al. 2020; Igl et al. 2018; Kapturowski et al. 2019). In the next section, we introduce the approximated belief model, its update, and other parts of our method.

In this section, we introduce our method, embedding learning and evaluation of uncertainty (ELUE), which learns how to infer tasks and how to use beliefs based on embeddings of task features. In addition, to alleviate the gap between meta-training and testing, which may prevent improvements from being made if only belief updating is done, it also adapts the learned policy and Q-function to the meta-test task through the updating of their parameters.

A sketch of the architecture of our networks is shown in Figure 1.

In this section, we review existing methods related to our method and discuss the differences between them.

Our method is inspired by PEARL, but there are essential differences. First, PEARL has no decoder, and the encoder is trained to minimize the critic loss. It is a simple approach, but its embedding can change depending on the current policy. The performance of PEARL was shown to degrade when used with off-policy (i.e., not recent) data (Rakelly et al. 2019). Therefore, PEARL uses an additional buffer for recent data to avoid the degradation; in contrast, our method can train the embedding using old data and does not need an additional buffer. Second, PEARL uses an encoder that can be represented as $q_{\phi}(z|h_{t})\propto\prod_{\tau=0}^{t-1}\mathcal{N}(z;\mu_{\tau},\sigma_{\tau}^{2})\propto\mathcal{N}(z;\frac{\sum_{\tau=0}^{t-1}\frac{\mu_{\tau}}{\sigma_{\tau}^{2}}}{\sum_{\tau=0}^{t-1}\frac{1}{\sigma_{\tau}^{2}}},\frac{1}{\sum_{\tau=0}^{t-1}\frac{1}{\sigma_{\tau}^{2}}})$, where $\mu_{\tau}$ and $\sigma_{\tau}^{2}$, the mean and variance of a Gaussian distribution, are the outputs of the neural network, $f_{\phi}(\cdot)$. As discussed around Equation (6), this is not a general form for encoder representation in terms of permutation invariance among $c_{\tau}$, while our encoder is represented in a general form. Third, PEARL’s policy and Q-function, $\pi(s,z)$ and $Q(s,a,z)$, where $z$ is sampled from $q(z|h_{t})$, are $z$-conditional, and $z$ itself has no uncertainty information. In comparison, ours are belief-conditional, which has uncertainty information. Fourth, PEARL only considers inference in meta-testing, while our method considers the updating of the parameters of neural networks.

There are other meta-learning methods based on inference; however, these methods are on-policy (Zintgraf et al. 2020; Humplik et al. 2019; Vuorio et al. 2019; Duan et al. 2016). Among these methods, VariBAD (Zintgraf et al. 2020) is the most related to our method. It considers embedding just like ours and beliefs over the embedding space. However, its sample efficiency is not as high as PEARL as shown in Zintgraf et al. (2020). In addition, its encoder is based on recurrent neural networks whose input is simply a history. Moreover, it does not consider updating the parameters of neural networks in meta-testing.

As for off-policy approaches, guided meta policy search (Mendonca et al. 2019) is introduced as an off-policy meta-learning method. Their method is based on MAML. In meta-training, it updates the parameters to imitate expert trajectories (outer updates) after updates by policy gradients (inner updates), which means that the inner updates and the updates in meta-testing need on-policy data. In addition, it is not based on inference. Lin et al. (2020) proposed an off-policy method which uses a decoder and policy conditioned the parameters of the decoder. The parameters corresponds to a belief in our method and they are adapted to a meta-test task by gradient descent. Their method is also not inference. It is based on MAML-like parameter updates.

In this section, we compare the sample efficiency of the existing “context-based” methods (PEARL, variBAD, and RL2 (Duan et al. 2016)) and our method. In the conventional setting of meta-learning experiments, the tasks in meta-testing and meta-training were in the same category and the tasks differed only in goal positions or weights of torso of the agent (Finn, Abbeel, and Levine 2017; Rakelly et al. 2019; Zintgraf et al. 2020; Vuorio et al. 2019). Our experimental settings include not only the conventional ones but also settings where the types of tasks in meta-testing and meta-training are different. The environment of the experiments is Meta-World with MuJoCo 2.0. Meta-World is a collection of Sawyer robot arm tasks, and there are 50 types of tasks and several benchmarks. In a task of Meta-World, each episode length is 150.

First, to examine the sample efficiency of our method, we conducted experiments with the conventional settings. We followed the ML1 benchmark scheme of Meta-World, where the differences in tasks mean differences in goals and the goals are sampled from the same distribution in meta-training and meta-testing. We chose six types of tasks, basket-ball, dial-turn, pick-place, reach, sweep-into, window-open, that were chosen from the types of tasks of the ML10 benchmark¹¹1The names of the types of tasks in the Meta-World paper are different from those in the Meta-World implementation. We refer to the names in the implementation. (descriptions of the types of tasks we used are shown in Table 1 in appendices). For each run, 20 meta-training tasks were generated. We executed five runs with different seeds and evaluated the average episode reward at the first episode in meta-testing. The results are shown in Figure 2. The results show that ELUE achieved higher performance than the other methods.

Second, to see the merits of our belief updates in meta-testing, we examined the performance of ELUE and PEARL in meta-testing with only inferences (i.e. belief updates in ELUE, and posterior updates and sampling the task variable in PEARL) and that of ELUE without any update of a belief, i.e. keeping a belief being a prior (“NoBelUpdate”) for comparison. In the experiment, meta-trained networks at about three million time steps in the first experiment of pick-place were used. The results are shown in Figure 3. The results show that both methods only need small amounts of data to improve the performance. Especially, the performance of ELUE was relatively high from the first episode while that of PEARL was low at first and improved in the next episode.

Third, to examine the sample efficiency in a set of diverse tasks, we conducted experiments, following the ML10 benchmark scheme, where the types of tasks in meta-training and meta-testing were different. In the benchmark, there are 15 types of tasks, ten for meta-training and five for meta-testing. We evaluated the sample efficiency from two perspectives: the performance in meta-testing in the early stages like Figure 2 and their improvement as the amount of data increases. For an experiment in the latter perspective, we modified PEARL and variBAD to make them update the parameters of neural networks in meta-testing, although the original versions of them do not update the parameters. In the experiment, meta-trainings were executed five times. For each run, 40 tasks were generated. For each meta-training, meta-tests were executed six times. The meta-tests were executed for about two million steps. In the meta-tests, we used meta-trained networks after 600 iterations (25 million time steps) for ELUE and PEARL, 4000 iterations (32 million time steps) for variBAD. To clarify the amount of improvement with our method, we also executed a variant of ELUE which learns policies from scratch in meta-testing without using beliefs. The results are shown in Figure 4. The performances of PEARL and ELUE slightly improved as time steps in meta-training increased. As for the learning curves of meta-testing, although the variance of the performance of ELUE was large, ELUE outperformed the other methods on average.

Fourth, to see the effectiveness of applying information bottleneck to our method, we compared the performances of variants of ELUE which was naively combined with soft actor-critic. We changed the coefficients of entropy in soft-actor critic. The other hyper parameters were the same among the variants.

The results (Figure 5) show that ELUE was slightly better than the other variants. This implies that information bottleneck objectives helped generalization of learned policies as shown in Igl et al. (2019).

Fifth, to examine the effectiveness of our updating method in meta-testing, we compared the performances of the different update methods (“BelGrad”, “NoBelGrad” and “Inference”) in the ML10 benchmark. BelGrad means belief updates by gradient descent after inference, which is used in meta-test of ELUE. NoBelGrad means that gradient descent updates are applied to not the belief but the parmeters of neural networks. Inference means no gradient updates. NoBelGrad and Inference were executed three times for each meta-training. We used the same settings and the same meta-trained networks in the experiment of Figure 4. The results (Figure 6) show that the performance of Inference did not improve as the amount of data increased, and that BelGrad was better than NoBelGrad.

We have proposed a novel off-policy meta-learning method, ELUE, which learns the embeddings of task features, beliefs over the embedding space, belief-conditional policies and Q-functions. Because of the belief-conditional policies, the performance can be improved by updating the beliefs, especially when the meta-test task is similar to the meta-training tasks. ELUE also updates the parameters of neural networks in meta-testing, which can alleviate the gap between tasks in meta-testing and those in meta-training. In the experiments with Meta-World benchmarks, we examined the sample efficiency of ELUE and existing methods in the two cases where the tasks in meta-training and meta-testing were similar and diverse. The results show that ELUE outperforms the other methods in both cases. We have also shown merits of the belief-conditional policies and the parameter updates in meta-testing.

In this section, we provide supplementary information about Meta-World, ELUE implementation, additional experimental results, and theoretical derivations.