Learning Action Translator for Meta
Reinforcement Learning on Sparse-Reward Tasks

Deep reinforcement learning (DRL) methods achieved remarkable success in solving complex tasks (Mnih et al. 2015; Silver et al. 2016; Schulman et al. 2017). While conventional DRL methods learn an individual policy for each task, meta reinforcement learning (meta-RL) algorithms (Finn, Abbeel, and Levine 2017; Duan et al. 2016; Mishra et al. 2017) learn the shared structure across a distribution of tasks so that the agent can quickly adapt to unseen related tasks in the test phase. Unlike most of the existing meta-RL approaches working on tasks with dense rewards, we instead focus on the sparse-reward training tasks, which are more common in real-world scenarios without access to carefully designed reward functions in the environments. Recent works in meta-RL propose off-policy algorithms (Rakelly et al. 2019; Fakoor et al. 2019) and model-based algorithms (Nagabandi, Finn, and Levine 2018; Nagabandi et al. 2018) to improve the sample efficiency in meta-training procedures. However, it remains challenging to efficiently solve multiple tasks that require reasoning over long horizons with sparse rewards. In these tasks, the scarcity of positive rewards exacerbates the issue of sample efficiency, which plagues meta-RL algorithms and makes exploration difficult due to a lack of guidance signals.

Intuitively, we hope that solving one task facilitates learning of other related tasks since the training tasks share a common structure. However, it is often not the case in practice (Rusu et al. 2015; Parisotto, Ba, and Salakhutdinov 2015). Previous works (Teh et al. 2017; Yu et al. 2020a) point out that detrimental gradient interference might cause an imbalance in policy learning on multiple tasks. Policy distillation (Teh et al. 2017) and gradient projection (Yu et al. 2020a) are developed in meta-RL algorithms to alleviate this issue. However, this issue might become more severe in the sparse-reward setting because it is hard to explore each task to obtain meaningful gradient signals for policy updates. Good performance in one task does not automatically help exploration on the other tasks since the agent lacks positive rewards on the other tasks to learn from.

In this work, we aim to fully exploit the highly-rewarding transitions occasionally discovered by the agent in the exploration. The good experiences in one task should not only improve the policy on this task but also benefit the policy on other tasks to drive deeper exploration. Specifically, once the agent learns from the successful trajectories in one training task, we transfer the good policy in this task to other tasks to get more positive rewards on other training tasks. In Fig. 1, if the learned policy $\pi$ performs better on task $\smash{\mathcal{T}^{(2)}}$ than other tasks, then our goal is to transfer the good policy $\smash{\pi(\cdot,\mathcal{T}^{(2)})}$ to other tasks $\smash{\mathcal{T}^{(1)}}$ and $\smash{\mathcal{T}^{(3)}}$. To enable such transfer, we propose to learn an action translator among multiple training tasks. The objective function forces the translated action to behave on the target task similarly to the source action on the source task. We consider the policy transfer for any pair of source and target tasks in the training task distribution (see the colored arrows in Fig. 1). The agent executes actions following the transferred policy if the transferred policy attains higher rewards than the learned policy on the target task in recent episodes. This approach enables the agent to leverage relevant data from multiple training tasks, encourages the learned policy to perform similarly well on multiple training tasks, and thus leads to better performance when applying the well-trained policy to test tasks.

We summarize the contributions: (1) We introduce a novel objective function to transfer any policy from a source Markov Decision Process (MDP) to a target MDP. We prove a theoretical guarantee that the transferred policy can achieve the expected return on the target MDP close to the source policy on the source MDP. The difference in expected returns is (approximately) upper bounded by our loss function with a constant multiplicative factor. (2) We develop an off-policy RL algorithm called Meta-RL with Context-conditioned Action Translator (MCAT), applying a policy transfer mechanism in meta-RL to help exploration across multiple sparse-rewards tasks. (3) We empirically demonstrate the effectiveness of MCAT on a variety of simulated control tasks with the MuJoCo physics engine (Todorov, Erez, and Tassa 2012), showing that policy transfer improves the performance of context-based meta-RL algorithms.

Context-based Meta-RL Meta reinforcement learning has been extensively studied in the literature (Finn, Abbeel, and Levine 2017; Stadie et al. 2018; Sung et al. 2017; Xu, van Hasselt, and Silver 2018) with many works developing the context-based approaches (Rakelly et al. 2019; Ren et al. 2020; Liu et al. 2020). Duan et al. (2016); Wang et al. (2016); Fakoor et al. (2019) employ recurrent neural networks to encode context transitions and formulate the policy conditioning on the context variables. The objective function of maximizing expected return trains the context encoder and policy jointly. Rakelly et al. (2019) leverage a permutation-invariant encoder to aggregate experiences as probabilistic context variables and optimizes it with variational inference. The posterior sampling is beneficial for exploration on sparse-reward tasks in the adaptation phase, but there is access to dense rewards during training phase. Li, Pinto, and Abbeel (2020) considers a task-family of reward functions. Lee et al. (2020); Seo et al. (2020) trains the context encoder with forward dynamics prediction. These model-based meta-RL algorithms assume the reward function is accessible for planning. In the sparse-reward setting without ground-truth reward functions, they may struggle to discover non-zero rewards and accurately estimating the reward for model-based planning may be problematic as well.

Policy Transfer in RL Policy transfer studies the knowledge transfer in target tasks given a set of source tasks and their expert policies. Policy distillation (Rusu et al. 2015; Yin and Pan 2017; Parisotto, Ba, and Salakhutdinov 2015) minimize the divergence of action distributions between the source policy and the learned policy on the target task. Along this line of works, Teh et al. (2017) create a centroid policy in multi-task reinforcement learning and distills the knowledge from the task-specific policies to this centroid policy. Alternatively, inter-task mapping between the source and target tasks (Zhu, Lin, and Zhou 2020) can assist the policy transfer. Most of these works (Gupta et al. 2017; Konidaris and Barto 2006; Ammar and Taylor 2011) assume existence of correspondence over the state space and learn the state mapping between tasks. Recent work (Zhang et al. 2020c) learns the state correspondence and action correspondence with dynamic cycle-consistency loss. Our method differs from this approach, in that we enable action translation among multiple tasks with a simpler objective function. Importantly, our approach is novel to utilize the policy transfer for any pair of source and target tasks in meta-RL.

Bisimulation for States in MDPs Recent works on state representation learning (Ferns, Panangaden, and Precup 2004; Zhang et al. 2020a; Agarwal et al. 2021) investigate the bismilarity metrics for states on multiple MDPs and consider how to learn a representation for states leading to almost identical behaviors under the same action in diverse MDPs. In multi-task reinforcement learning and meta reinforcement learning problems, Zhang et al. (2020a, b) derives transfer and generalization bounds based on the task and state similarity. We also bound the value of policy transfer across tasks but our approach is to establish action equivalence instead of state equivalence.

In this section, we first describe our approach to learn a context encoder capturing the task features and learn a forward dynamics model predicting next state distribution given the task context (Sec. 3.2). Then we introduce an objective function to train an action translator so that the translated action on the target task behaves equivalently to the source action on the source task. The action translator can be conditioned on the task contexts and thus it can transfer a good policy from any arbitrary source task to any other target task in the training set (Sec. 3.3). Finally, we propose to combine the action translator with a context-based meta-RL algorithm to transfer the good policy from any one task to the others. During meta-training, this policy transfer approach helps exploit the good experiences encountered on any one task and benefits the data collection and further policy optimization on other sparse-reward tasks (Sec. 3.4). Fig. 2 provides an overview of our approach MCAT.

In this section, we theoretically support our objective function (Equation 3) to learn the action translator. Given $s$ on two MDPs with the same state and action space, we define that action $a^{(i)}$ on $\smash{\mathcal{T}^{(i)}}$ is equivalent to action $a^{(j)}$ on $\smash{\mathcal{T}^{(j)}}$ if the actions yielding exactly the same next state distribution and reward, i.e. $\smash{p^{(i)}(\cdot|s,a^{(i)})=p^{(j)}(\cdot|s,a^{(j)})}$ and $\smash{r^{(i)}(s,a^{(i)})=r^{(j)}(s,a^{(j)})}$ . Ideally, the equivalent action always exists on the target MDP $\smash{\mathcal{T}^{(i)}}$ for any state-action pair on the source MDP $\smash{\mathcal{T}^{(j)}}$ and there exists an action translator function $\smash{H:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{A}}$ to identify the exact equivalent action. Starting from state $s$, the translated action $\smash{\tilde{a}=H(s,a)}$ on the task $\smash{\mathcal{T}^{(i)}}$ generates reward and next state distribution the same as action $a$ on the task $\smash{\mathcal{T}^{(j)}}$ (i.e. $\smash{\tilde{a}B_{s}a}$). Then any deterministic policy $\smash{\pi^{(j)}}$ on the source task $\smash{\mathcal{T}^{(j)}}$ can be perfectly transferred to the target task $\smash{\mathcal{T}^{(i)}}$ with $\smash{\pi^{(i)}(s)=H(s,\pi^{(j)}(s))}$. The value of the policy $\smash{\pi^{(j)}}$ on the source task $\smash{\mathcal{T}^{(j)}}$ is equal to the value of transferred policy $\smash{\pi^{(i)}}$ on the target task $\smash{\mathcal{T}^{(i)}}$.

Without the assumption of existence of a perfect correspondence for each action, given any two deterministic policies $\smash{\pi^{(j)}}$ on $\smash{\mathcal{T}^{(j)}}$ and $\smash{\pi^{(i)}}$ on $\smash{\mathcal{T}^{(i)}}$, we prove that the difference in the policy value is upper bounded by a scalar $\frac{d}{1-\gamma}$ depending on L1-distance between reward functions $\smash{|r^{(i)}(s,\pi^{(i)}(s))-r^{(j)}(s,\pi^{(j)}(s))|}$ and total-variation distance between next state distributions $\smash{D_{TV}(p^{(i)}(\cdot|s,\pi^{(i)}(s)),p^{(j)}(\cdot|s,\pi^{(j)}(s)))}$. Detailed theorem (Theorem 1) and proof are in Appendix A.

For a special case where reward function $\smash{r(s,a,s^{\prime})}$ only depends on the current state $s$ and next state $s^{\prime}$, the upper bound of policy value difference is only related to the distance in next state distributions.

According to Proposition 1, if we can optimize the action translator $\smash{H}$ to minimize $d$ for policy $\smash{\pi^{(j)}}$ and $\smash{\pi^{(i)}(s)=H(s,\pi^{(j)}(s))}$, the value of the transferred policy $\smash{\pi^{(i)}}$ on the target task can be close to the value of source policy $\smash{\pi^{(j)}}$. In many real-world scenarios, especially sparse-reward tasks, the reward heavily depends on the state and next state instead of action. For example, robots running forward receive rewards according to their velocity (i.e. the location difference between the current and next state within one step); robot arms manipulating various objects earn positive rewards only when they are in the target positions. Thus, our approach focuses on the cases with reward function approximately as $r(s,s^{\prime})$ under the assumption of Proposition 1. For any state $\smash{s\in\mathcal{S}}$, we minimize the total-variation distance between two next state distributions $\smash{D_{TV}(p^{(j)}(\cdot|s_{t},\pi^{(j)}(s_{t})),p^{(i)}(\cdot|s_{t},\pi^{(i)}(s_{t})))}$ on source and target MDPs. Besides, we discuss the policy transfer for tasks with a general reward function in Appendix C.3.

There is no closed-form solution of $\smash{D_{TV}}$ and $\smash{D_{TV}}$ is related with Kullback–Leibler (KL) divergence $\smash{D_{KL}}$ by the inequality $\smash{D_{TV}(p\|q)^{2}\leq D_{KL}(p\|q)}$ Thus, we instead consider minimizing $\smash{D_{KL}}$ between two next state distributions. $\smash{D_{KL}(p^{(j)}||p^{(i)})}$ is $\smash{-\sum_{s^{\prime}}p^{(j)}(s^{\prime})\log p^{(i)}(s^{\prime})}+\smash{\sum_{s^{\prime}}p^{(j)}(s^{\prime})\log p^{(j)}(s^{\prime})}$. The second term does not involve $H$ and thus can be viewed as a constant term when optimizing $H$. We focus on minimizing the first term $\smash{-\sum_{s^{\prime}}p^{(j)}(s^{\prime})\log p^{(i)}(s^{\prime})}$. $F$ is a forward model approximating $\smash{p^{(i)}(s^{\prime})}$. We sample transitions $s,\smash{\pi^{(j)}(s)},s^{\prime}$ from the source task. $s^{\prime}$ follows the distribution $\smash{p^{(j)}(s^{\prime})}$. Thus, minimizing the negative log-likelihood of observing the next state $\smash{L_{trans}=-\log F(s^{\prime}|s,\pi^{(i)}(s))}$ is to approximately minimize $D_{KL}$. Experiments in Sec. 5.1 suggest that this objective function works well for policy transfer across two MDPs. Sec. 3.3 explains the motivation behind $\smash{\mathcal{L}_{trans}}$ (Equation 3) to learn an action translator among multiple MDPs instead of only two MDPs.

We design and conduct experiments to answer the following questions: (1) Does the transferred policy perform well on the target task (Tab. 1, Fig. 4)? (2) Can we transfer the good policy for any pair of source and target tasks (Fig. 5)? (3) Does policy transfer improve context-based Meta-RL algorithms (Fig. 3, Tab. 2, Tab. 3)? (4) Is the policy transfer more beneficial when the training tasks have sparser rewards (Tab. 4)? Experimental details can be found in Appendix C.

The scope of MCAT is for tasks with varying dynamics, same as many prior works (Yu et al. 2017; Nagabandi et al. 2018; Nagabandi, Finn, and Levine 2018; Zhou, Pinto, and Gupta 2019). our theory and method of policy transfer can be extended to more general cases (1) tasks with varying reward functions (2) tasks with varying state & action spaces.

Following the idea in Sec. 4, on two general MDPs, we are interested in equivalent state-action pairs achieving the same reward and transiting to equivalent next states. Similar to Proposition 1, we can prove that, on two general MDPs, for two correspondent states $\smash{s^{(i)}}$ and $\smash{s^{(j)}}$, the value difference $\smash{|V^{\pi^{(i)}}(s^{(i)},\mathcal{T}^{(i)})-V^{\pi^{(j)}}(s^{(j)},\mathcal{T}^{(j)})|}$ is upper bounded by $\frac{d}{1-\gamma}$, where $d$ depends on $\smash{D_{TV}}$ between the next state distribution on source task and the probability distribution of correspondent next state on target task. As an extension, we learn a state translator jointly with our action translator to capture state and action correspondence. Compared with Zhang et al. (2020c) learning both state and action translator, we simplify the objective function training action translator and afford the theoretical foundation. For (1) tasks with varying reward functions, we conduct experiments on MetaWorld moving the robot arm to a goal location. The reward at each step is inversely proportional to its distance from the goal location. We fix a goal location on source task and set target tasks with distinct goal locations. Furthermore, we evaluate our approach on 2-leg and 3-leg HalfCheetah. We can test our idea on (2) tasks with varying state and action spaces of different dimensions because the agents have different numbers of joints on the source and target task. Experiments demonstrate that ours with a simpler objective function than the baseline (Zhang et al. 2020c) can transfer the source policy to perform well on the target task. Details of theorems, proofs, and experiments are in Appendix E.

Meta-RL with long-horizon, sparse-reward tasks is challenging because an agent can rarely obtain positive rewards, and handling multiple tasks simultaneously requires massive samples from distinctive tasks. We propose a simple yet effective objective function to learn an action translator for multiple tasks and provide the theoretical ground. We develop a novel algorithm MCAT using the action translator for policy transfer to improve the performance of off-policy, context-based meta-RL algorithms. We empirically show its efficacy in various environments and verify that our policy transfer can offer substantial gains in sample complexity.

This work was supported in part by NSF CAREER IIS-1453651, NSF SES 2128623, and LG AI Research. Any opinions, findings, conclusions or recommendations expressed here are those of the authors and do not necessarily reflect views of the sponsor.