On the Convergence Theory of Meta Reinforcement Learning with Personalized Policies

First, we propose a personalization framework to improve the performance of meta-RL on distinct tasks, which learns a meta-policy and personalized policies for all tasks and specific tasks, respectively. In particular, by learning multiple task-specific personalized policies, the performance on each task can be improved. Meanwhile, to aggregate these differentiated personalized policies to obtain a meta-policy applicable to all tasks, our framework addresses the gradient conflict problem by adopting personalization constraint in the objective function. Under the tabular setting, we constrain the meta Q-table and personalized Q-tables by formulating a differentiable function to encourage each task to pursue its personalized Q-tables around the meta Q-table. Under the deep RL setting, we obtain the personalization performance by constraining the personalized and meta networks (e.g., actor, critic, and inference networks), so that the networks can refer to each other during training.

Second, we propose an alternating minimization algorithm, named pMeta-RL, to solve the personalized meta-RL problem. The personalized policies are updated based on the value iteration, while the auxiliary policies used to aggregate the meta-policy are updated based on the gradient descent. More importantly, theoretical analysis shows that the proposed algorithm can converge well, and the convergence speed is basically linear with the iteration number. And we give an upper bound on the difference between the personalized policies and meta-policy, which is influenced by the regularization parameter and task diversity. Moreover, we extend the proposed pMeta-RL algorithm to a deep network version (named deep pMeta-RL) based on soft actor-critic (SAC), making it suitable for continuous control tasks.

Finally, we evaluate the performance of pMeta-RL and deep pMeta-RL on the Gridworld environment and the MuJoCo suites, respectively. The experimental results show that the proposed pMeta-RL can achieve better performance than the model averaging method, and verify the conclusion of the convergence analysis. Compared to other previous Meta-RL algorithms, our deep pMeta-RL algorithm can achieve the highest average return on the MuJoCo locomotion tasks.

Existing meta-RL methods learn the dynamics models and policies that can quickly adapt to unseen tasks, which is built on the meta-learning framework in the context of RL [7, 15, 8]. $\text{RL}^{2}$ [15] formulate the meta-learning problem as a second RL procedure, it represent a “fast” RL algorithm as a recurrent neural network and learn it from data but update the parameters at a “slow” pace. MAML [8] meta-learns model parameters by differentiating the learning process for fast adaptation on unseen tasks. E-MAML and E-$\text{RL}^{2}$ [9] explicitly optimize the per-task sampling distributions during adaptation with respect to the expected future returns, which is closely related to the MAML algorithm. ProMP [16] theoretically analyses the MAML formulation and addresses the biased meta-gradient issue. To achieve sample efficiency, PEARL [10] performs structured exploration by reasoning about uncertainty over tasks and enables fast adaptation by accumulating experience online. VariBAD [11] learns a policy that conditions on this posterior belief, which can trade off exploration and exploitation under task uncertainty. Other meta-RL works, such as MAME [17] and MetaCURE [12], learn a separate exploration policy for general dense and sparse-reward tasks.

Multi-task RL methods have demonstrated that learning with different tasks can benefit each other in robotics and other fields [18, 19, 20]. However, the gradient conflict problem still exists in multi-task RL, various approaches based on compositional models [21, 22, 23, 24], gradients similarity [25, 26], or policy distillation [27, 14, 28] have been proposed to solve this challenge. In particular, multi-head multi-task SAC [29] [13] and experience sharing networks [22, 23] training with diverse robot manipulation tasks jointly with a sharing network backbone and multiple task-specific heads for actions. Hard routing [21] and Soft modularization [24] learn a routing network to reconfigure different modules for different tasks automatically without explicitly specifying the policy structure, which can avoid the gradients conflicts effectively. Other works such as gradient surgery [25] and normalization [26], leverage the similarity between gradients from different tasks to enhance the learning process and mitigate interference. Existing algorithms have shown superior performance in multiple tasks, however, they need to share trajectories between tasks and are difficult to generalize to unseen tasks.

Another related work is the federated reinforcement learning (FRL), these methods encourage multiple agents to federatively build a better decision-making policy in a privacy-preserving manner, which have been applied in robot navigation and Internet of things [30, 31, 32]. [31] analyzes the communication overhead in FRL and gives the convergence bound. [32] studies adversarial attacks in FRL systems in detail and gives theoretical guarantees in the event of participant failure. However, exist FRL algorithms usually only consider the same tasks and lack the ability to handle multiple tasks.

Reinforcement learning (RL) algorithms are designed to solve sequential decision problems [2, 1, 3]. Generally speaking, the agent makes action decisions at each discrete time step based on its observation of the environment $\mathcal{E}$. After performing an action, the environment generates a reward and transfers the agent to a new state. This can be formulated as a standard Markov decision process (MDP), consisting of a tuple $\mathcal{G}=\left\langle\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\gamma\right\rangle$, where $\mathcal{S}$ is a finite set of states, $\mathcal{A}$ is a finite set of actions, $r\in\mathcal{R}(s,a):\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ is the bounded reward function with $\mathcal{R}(s,a)\leq\mathcal{R}_{\max}$, $\mathcal{P}(s^{\prime}\left|s\right.,a):\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow\left[0,1\right]$ denotes the state transition function, and $\gamma\in\left[0,1\right)$ is the discount factor. The goal of RL is to find the optimal policy $\pi^{*}$ by maximizing the expected long-term return.

Most gradient-based meta-RL algorithms are model-agnostic meta-learners built upon MAML, performing vanilla policy gradient steps towards maximizing the performance of the policy for each task. These methods generally consider tasks drawn from a distribution $\mathcal{T}_{i}\sim p\left(\mathcal{T}\right)$, where each task $\mathcal{T}_{i}$ is a different MDP. During the meta-training time, the goal of Meta-RL is to learn a policy $\pi(a|s)$, which can adapt to the multiple tasks. In particular, this policy $\pi(a|s)$ can be optimized by maximizing the average expected return across the task distribution,

where $\rho_{0}$ is the initial state-distribution, and $Q^{\pi}_{i}(s,a)$ is the state-action value function.

Existing meta-RL algorithms leverage the transitions from similar tasks to improve the sample efficiency and achieve good performance on future new tasks. However, most gradient-based methods suppose that the task distribution is “homogeneous”, which means that all tasks are from the same domain, (i.e. a single dataset, multiple datasets with same input feature space [33] or robot control with different transition functions but similar dynamics [8]). This limits the application of these algorithms to distinct or “hetergeneous” tasks, i.e., different tasks may vary in environments or transition functions with a large interval [34]. Figure 1(a) shows several distinct tasks, such as navigating to two completely opposite directions on Ant robot or controlling the speed of the Half-Cheetah to reach two endpoints of a speed interval. In this setting, optimizing some tasks in meta-RL algorithms may negatively affect others.

Figure 1(b) illustrates these problem. For distinct tasks, gradient aggregation makes the values of different state-action pairs similar, i.e., a recommended action may not be given accurately. Other Meta-RL methods, such as PEARL [10], distinguish different tasks by learning a task-relevant context $\mathbf{z}$ based on the history. These algorithms maintain an inference network $q_{\zeta}(\mathbf{z}|\mathbf{c})$, parameterized by $\zeta$, to encode salient information about tasks from context $\mathbf{c}$, which consists of the transitions of different tasks. Nevertheless, the gradient conflict problem still exists due to the shared policy network parameters. To solve this challenge, we aim to learn a meta-policy that adapts to multiple distinct tasks, while optimizing a personalized policy for each task. In this subsection, we first give the definition of personalized Meta-RL under the tabular and deep network settings, while the corresponding algorithm for solving personalized Meta-RL is proposed in the next section.

Definition 1 describes a personalization approach based on Q-learning, where the personalized constraint enables each task optimizes its own policy around $Q^{\pi}(s,a)$. Since Q learning is difficult to adapt to large-scale problems, we use a function approximator to represent the Q-values, and extend Definition 1 to deep personalized meta-RL based on SAC [29] and inference network $q_{\zeta}$.

Definition 2 considers the personalization approach to three models $(\phi_{i},\theta_{i},\zeta_{i})$, and establishes the relationship between them to optimize the personalized policy based on the meta-policy. In the next section, we will propose the corresponding algorithms to solve these problems.

To solve personalized Meta-RL, our pMeta-RL algorithm updates $Q_{i}^{\pi_{i}}$ and $Q^{\pi}$ by alternately minimizing the two subproblems of pMeta-RL. More specifically, we first update $Q_{i}^{\pi_{i}}$ by solving

then we introduce an auxiliary meta Q-values $Q_{i}^{\pi}$ for each task, which can be updated by solving

After that, all auxiliary meta Q-values $Q_{i}^{\pi},i=1,\cdots,N$ can be aggregated to update the $Q^{\pi}$.

We first focus on solving (3). Note that in the model-free RL, the standard Q-learning can be used to solve $\arg\min_{\pi_{i}}-L_{i}(\pi_{i})$. For each state-action pair $(s,a)$, $Q_{i}^{\pi_{i}}(s,a)$ is estimated by iteratively applying the Bellman optimal operator $\mathcal{B}^{*}\left[Q_{i}^{\pi_{i}}(s,a)\right]=\mathcal{R}_{i}(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim\mathcal{P}(s^{\prime}|s,a)}\left[\max_{a^{\prime}}Q_{i}^{\pi_{i}}(s^{\prime},a^{\prime})\right]$, i.e.,

where $Q^{\pi_{i}}_{i,k}(s,a)$ is the Q-table at the $k$-th iteration. Then the exact or an approximate maximization scheme is used to recover the greedy policy. This inspires us to update $Q^{\pi_{i}}_{i}(s_{i},a_{i})$ by performing the following iterations

where

We solve (5) via one step gradient descent based on the gradient $\nabla\Xi_{i,k}\left(Q_{i}^{\pi_{i}}\right)$ as the following

where $\nabla\varphi_{i}=2\left(Q_{i}^{\pi_{i}}(s_{i},a_{i})-Q_{i}^{\pi}(s_{i},a_{i})\right)$ and $\eta^{k}$ is a learning rate. However, it is not straightforward to update $Q_{i}^{\pi_{i}}$ based on (6) since the system transition function $\mathcal{P}_{i}$ is usually unknown. In most model-free RL algorithms, we need to sample transition data to update Q-value. We hence leverage several samples $(s_{i},a_{i},r_{i},s_{i}^{\prime})$ to obtain an approximated Q-table $\widetilde{Q}_{i}^{\pi_{i}}$ as the following,

The iterations in (6) and (IV-A) go until the maximum iteration number $K$ is reached. Then we have the following Theorem 1, which is proved in Appendix A-C. Theorem 1 indicates that the error generated by using the approximated Q-table $\widetilde{Q}_{i}^{\pi_{i}}$ is minimal after sufficient iterations. We hence set $Q_{i}^{\pi_{i}}(s_{i},a_{i})\leftarrow\widetilde{Q}_{i,K}^{\pi_{i}}(s_{i},a_{i})$.

Once $Q_{i}^{\pi_{i}}(s_{i},a_{i})$ is available, we update the auxiliary Q values $Q_{i}^{\pi}$ in (4) as the following

where $\eta$ is a learning rate.

We solve these two subproblems alternately until the maximum number $R$ of iterations is reached and then update the meta Q-values. Note that since $s\in{\mathcal{S}}=\cup_{i\in N}\mathcal{S}_{i},a\in{\mathcal{A}}=\cup_{i\in N}\mathcal{A}_{i}$, the states and actions in the meta Q-table may not necessarily be in a specific personalized Q-table. Hence all the auxiliary meta Q-tables $Q_{i,R}^{\pi}(s_{i},a_{i}),i=1,...,N$ are collected to perform the following aggregation

where $\beta$ is a weight parameter, $\mathbb{I}\{\cdot\}$ is an indicator operator, and $N_{s}$ denotes the number of tasks that satisfy $Q_{i}^{\pi}(s,a)\neq\emptyset$, i.e., $|\mathbb{I}\{Q_{i}^{\pi}(s,a)\neq\emptyset\}|$. We repeat the above process until the meta-policy converges or reaches the maximum number of iterations $C$.

In this subsection, we present the convergence analysis of our pMeta-RL algorithm. First, we present the following assumption, which quantifies the diversity of the reward and transition function among tasks.

Then we have the following Theorem 2 based on Assumption 1, which is proved in Appendix A-D.

In this subsection, we extend the meta-policy iteration algorithm to a deep learning version with neural network function approximators. Taking the actor model $\phi_{i}$ as an example, the deep pMeta-RL updates the personalized policy by solving the following two sub-problems:

where $\omega_{i}$ is the auxiliary model, which is used to update the meta-model $\omega$.

To solve (10), similar with Theorem 1, we sample a mini-batch data $D_{i}$ to obtain an approximated $\nabla J_{\pi_{i}}(\phi_{i})$ by calculating $\mathbb{E}\left[\nabla J_{\pi_{i}}(\phi_{i};D_{i})\right]$. Then, the personalized policy is updated based on the stochastic gradient descent by

where we use the $\ell_{2}$-norm constraint such that $\nabla\varphi_{i}(\phi_{i},\omega_{i})=2(\phi_{i}-\omega_{i})$ and $\alpha$ is the personalized policy learning rate.

For the subproblem (11), $\omega_{i}$ is updated by

where $\eta$ denotes the auxiliary model learning rate. Moreover, the critic and inference models are also updated similarly, which we omit here for brevity.

After $R$ iterations, the meta-policy can be updated as the following

where $\beta$ is a weight parameter and $c$ represents the current number of update rounds of the meta-policy. Finally, we obtain a well-trained meta-policy $(\omega^{C},\vartheta^{C},\varpi^{C})$. We summarize pMeta-RL and deep pMeta-RL algorithms in Algorithm 1 for clarification.

The deep pMeta-RL in Algorithm 1 proposes a general meta-training procedure, and the meta-testing phase is consistent with PEARL. By training an inference network, deep pMeta-RL can estimate the state-value function $Q_{\vartheta}(s,a,\mathbf{z})$ under different tasks and generalize to unseen tasks.

For the Gym suite [35], we consider two classical control environments, “CartPole”, and “MountainCar” ¹¹1The IDs of these environments in the OpenAI Gym library are: CartPole-v0 and MonutainCar-v0. For each environment, we modify its physical parameters to induce different transitions as different tasks.

CartPole: The goal of CartPole is to balance a pole on the top of the cart. The agent can get the position and velocity of the cart, the angle and angular velocity of the pole as observations, which is a four-dimensional continuous space. We can make the cart move left or right as the action space. This task ends when the pole falls over or the cart out of bounds. The system parameters contain the force, the mass of pole (masspole) and the length of pole (lengthpole).

MountainCar: The goal of MounatianCar is to reach a specific position on a one dimensional track between two mountains. The agent learns to drive back and forth to obtain enough power to allow the car to reach its goal. The states are the position $\dot{p}$ and velocity $\dot{v}$ of the car and the action consist of pushing left, right and no push. The system transition function can be written as

where $a=0,1,2$ represent the push left,right and no push, respectively, $\hat{F}$ is the magnitude of the force, $\hat{G}$ is the gravity, and $\zeta$ controls the inclination of the track. We modify $\zeta$ to generate different tasks, which is shown in table II.

For the mojuco suite, we consider the same meta-learning domains as [36, 12].

Ant-Fwd-Back: Move forward or backward (2 train tasks, 2 test tasks).

Half-Cheetah-Dir: Move forward or backward (2 train tasks, 2 test tasks).

Half-Cheetah-Vel: Reach a target velocity (10 train tasks, 3 test tasks).

Walker-2D-Params: The agent is randomly initialized with different system dynamics parameters and keep walking (10 train tasks, 3 test tasks).

We implement our personalized method based on DQN with a prioritized experience replay buffer, all the details can be found in Table II. Moreover, our code is based on Tianshou [37] framework²²2https://github.com/thu-ml/tianshou, and all experiments were conducted on a NVIDIA Quadro RTX 6000 environment.

Then we present all the hyperparameter settings in the Gym and MuJoCo tasks. Table II shows the hyperparameters used on CartPole and MountainCar. All hyperparameters are basically set to be the same as in [38]. Table III shows the hyperparameters used on MuJoCo suite, where all hyperparameters are basically set to be the same as in [10, 12]. For all the compared algorithms, we set the number of training steps to be the same as that used in our pMeta-RL. The rest of the hyperparameters in their algorithm are consistent with the settings recommended in their papers [10, 12].

To evaluate the performance of the proposed pMeta-RL, we construct a simple Gridworld environment with a finite state-action space based on [39]. In this Gridworld environment, the goal of each task is to reach the landmark in the grid and the reward is set as the distance between the agent and landmark. The number of tasks is $N=5$, and the tasks vary in different grid sizes. We use the model averaging method of the Q-values as the baseline, whose performance of this method has been validated in [38]. In particular, this method only performs weighted average aggregation on the Q tables of each task to obtain the multi-task Q-table. We set $C=10,R=3,K=1,\lambda=10$ in the experiments. Figure 2 shows the convergence rate of pMeta-RL, the average returns obtained by the personalized policy and meta-policy are better than that obtained by the model-averaging method. Note that the performance of the model-averaging method can gradually approach the personalized model as the training epoch increases. This is because the multi-task Q-table also retains all the state-action values of each task, which can adapt to different tasks. Furthermore, we directly extend the Q-learning-based approach to DQN and test the performance in the Gym environments, whose tasks vary the transition function. We can find that the performance of the personalized policy is better than the averaged-DQN, which is not well adapted to distinct tasks from Figure 2.

Figure 2 shows the performance of pMeta-RL in Gridworld environment with different values of $\lambda$, where we set $C=10,R=3,K=1$. When $\lambda=10$, pMeta-RL achieves the best performance. However, a larger $\lambda$ may degrade performance which mean this parameter need to be chosen carefully for various environments.

To evaluate the proposed deep pMeta-RL algorithm, we conducted extensive experiments on the continuous control tasks in the MuJoCo suits [40], which are benchmarks commonly used by meta-RL algorithms. These tasks vary in either the reward function (target velocity for Half-Cheetah-Vel, and walking direction for Half-Cheetah-Fwd-Back, Ant-Fwd-Back) or transition function (Walker-Params). We compare deep pMeta-RL against several representative meta-RL algorithms, including PEARL [10], ProMP [16], MetaCURE [12] and MAML-TRPO [8]. We also compare the recurrence-based policy gradient $\text{RL}^{2}$ method as [10]. For a fair comparison, we set the maximum episode length to 200 for all the above tasks, which is similar to PEARL. In addition, we set the same number of training steps for each epoch for the above algorithms.

Figure 4 shows that the personalized policy of deep pMeta-RL based on SAC can outperform prior meta-RL methods across all domains in terms of the average returns. The meta-policy achieves comparable performance to other algorithms, albeit with a drop in performance in the Half-Cheetah-Vel environment. This may be mainly due to the increase in the number of tasks in this environment and the increase in the similarity between tasks, resulting in a weakened personalization ability and a slower meta-policy convergence rate. Furthermore, we also observe that MetaCURE is unable to obtain reasonable results in Walker-2D-Params under the same setting. MetaCURE is an efficient algorithm for solving the sparse rewards problem. However, in the context of dense rewards, the intrinsic rewards proposed in the original paper may conflict with dense rewards, resulting in performance degradation. The average returns obtained using the personalized policies are much better than other meta-RL algorithms (e.g., $12.59\%$ and $11.86\%$ than PEARL in Ant-Fwd-Back and Cheetah-Fwd-Back environments, respectively), but only slightly gain in the Cheetah-Vel and Walker-2D-Params environments. And we found PEARL can obtains almost the same performance as the personalized policy in Walker-2D-Params. Summarizing the above, we can infer that our personalization method is better suited for more distinct tasks, e.g., navigating to the exact opposite direction.

We visualized the trajectories to illustrate the superiority of the personalized policy on the Sparse-Point-Robot, which is a 2D navigation task in which a point robot must navigate to different goal locations on edge of a semicircle. A reward is given only when the robot is within a certain radius of the goal, which is set as 0.2 in our experiments. By randomly sampling 10 goals, we compare PEARL with different context $\mathbf{c}_{1:5}$ and $\mathbf{c}_{1:15}$, whose subscripts denote the number of trajectories contained in the context. Specifically, the first trajectory is collected with probabilistic context variable $\mathbf{z}$ sampled from the prior $p(\mathbf{z})$ and the subsequent trajectories are collected with $\mathbf{z}\sim q_{\zeta}(\mathbf{z}|\mathbf{c})$ where context is aggregated over all collected trajectories. As shown in Figure 5, we can observe that that the agent with personalized policy can navigate to more targets than PEARL and achieve higher returns. Moreover, the robot can navigate more accurately with more context information, for example, a robot making decisions with $\mathbf{c}_{1:15}$ can navigate to the center of the goals, while using $\mathbf{c}_{1:5}$ may go beyond the goals radius area.

We also validate the performance of the meta-policy on the unseen tasks in Figure 6, which is almost consistent with the performance in Figure 4. We find that the meta-policy can adapt well to multiple tasks in most environments, except for the slower convergence on Half-Cheetah-vel. Combining Figure 4 and Figure 6, we conclude that the performance of the personalized policy is promising, such that our algorithm can be more suitable for decentralized multi-task learning, each task maintains its personalized policy while learning a meta-policy for multiple tasks without sharing trajectories.