Multi-Task Reinforcement Learning with Context-based Representations

Introduction

Reinforcement learning (RL) has made large strides over the last several years (Mnih et al., 2013; Silver et al., 2017; Radford et al., 2019). While these improvements are significant, much of this success has been restricted to the single task setting (Teh et al., 2017; Yu et al., 2020b). In contrast, humans are adept at multi-tasking by acquiring new skills and composing known skills to solve complex tasks (Parascandolo et al., 2018; Allen et al., 2020). For autonomous agents to adapt effectively in the real world, they need to master multiple tasks in a sample efficient manner. Multi-task reinforcement learning (MTRL) is a promising approach to train effective real-world agents (Tanaka & Yamamura, 2003; Rusu et al., 2016; Borsa et al., 2016; Rajeswaran et al., 2016; El Bsat et al., 2017; Andreas et al., 2017; Igl et al., 2020; D’Eramo et al., 2020; Yu et al., 2020a).

One limitation of existing MTRL methods is the inability to leverage side information (or metadata), like the description of a task, to learn generalizable skills and transfer common knowledge across tasks. Such metadata is often available in real-world tasks but is not leveraged in the typical MTRL setting. This metadata can take the form of natural language task descriptions or instructions, which are often incorporated only for human usage to communicate information about tasks. These natural language descriptions have been utilized in single-task RL setups like goal-oriented RL (Chevalier-Boisvert et al., 2019; Luketina et al., 2019; Jiang et al., 2019). We show that this information can also be used to learn context-dependent, composable representations shared across a family of tasks, with the metadata acting as the context. Specifically, we show that the metadata can be used to learn a prior over a collection of encoders and can be leveraged to select the encoder(s) for any given task. The learned encoders can specialize to different aspects of the tasks, allowing for more efficient sharing of knowledge across the tasks. An important additional benefit is that the learned representations are interpretable.

The default formulation of MTRL environments is as a family of Markov Decision Processes (MDPs) (Bellman, 1957; Puterman, 1995). This framework does not provide a natural way to incorporate metadata. We therefore turn to the Contextual Markov Decision Process formulation (Hallak et al., 2015) to define our setting (Figure 1). We assume access to a context (the metadata) that contains additional task-specific information which is easily available and can be used to improve per-task performance. This metadata is more useful in a multi-task setting where it can be used to infer the relationship between different tasks coming from a specified family. As an example, given two contexts, “pour tea into my cup” and “pour water into the mug kept on the table”, the agent can infer the relations between the two tasks and identify the common contexts (and subtasks) like “pour” (Figure 1, right). Such context is less likely to be useful to accelerate learning in the single task setup.

Sharing representations across tasks can be an effective way for RL agents to transfer knowledge across tasks. However, not all knowledge transfers are positive. Some aspects of a given task could be meaningful only for that task and irrelevant (or even detrimental) for the other tasks. This effect is commonly known as negative interference (Parisotto et al., 2015; Teh et al., 2017). In general, choosing which information/knowledge to transfer across the tasks, or even deciding which tasks should be learned together is an open problem (Standley et al., 2019). We posit that in the CMDP setup, the context can be used to inform which representations and information should be shared across the tasks.

In this work, we propose a novel approach for the contextual multi-task RL setting where we encode an input observation into multiple representations (corresponding to different skills or objects) using a mixture of encoders. The learning agent can use the context to decide which representation(s) it uses for any given task, giving the agent a fine-grained control over what information is shared across tasks, thus alleviating negative interference. We call our method Contextual Attention-based REpresentation learning, or CARE for short.

Key contributions of this work are 1) a simple, yet effective, way to incorporate task metadata, or contextual information, to improve sample efficiency and asymptotic performance, 2) a new representation learning algorithm for MTRL that leverages a mixture of interpretable encoders which encodes task and object-specific information about each state space, and 3) state-of-the-art results on a challenging multi-task RL benchmark, Meta-World (Yu et al., 2020b). For example videos see https://sites.google.com/view/mtrl-care. The implementation of the algorithms is available at https://github.com/facebookresearch/mtrl.

Preliminaries

Contextual Markov Decision Processes were first proposed by Hallak et al. (2015) as an augmented form of Markov Decision Processes that utilize side information as a form of context, similar to in contextual bandits.

A contextual Markov decision process (CMDP) is defined by tuple $\langle\mathcal{C},\mathcal{S},\mathcal{A},\mathcal{M}\rangle$ where $\mathcal{C}$ is the context space, $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space. $\mathcal{M}$ is a function which maps a context $c\in\mathcal{C}$ to MDP parameters $\mathcal{M}(c)=\{R^{c},T^{c}\}$.

Contexts can be applied in the multi-task setting, where we define a family of MDPs where each MDP has a shared state space $\mathcal{S}$. However, the agent only has access to a partial state space $\mathcal{S}^{c}$ (either low-dimensional or rich, like pixels) that is a subspace of the original state space $\mathcal{S}$, focusing only on objects relevant to the task at hand. Different MDPs can involve different combinations of objects and skills, hence the state space $\mathcal{S}^{c}$ and reward function $R^{c}$ can differ across MDPs. However, the objects are shared across tasks, i.e., the object-specific dynamics remain consistent across tasksNote that the dynamics $T^{c}$ for each MDP are still different because the state spaces are different.. In this work, we focus on the low-dimensional setting where $\mathcal{S}^{c}$ is a strict subset of the dimensions in $\mathcal{S}$.

We attach this additional relaxation of the original CMDP definition to define a new setting, a Block Contextual MDP (BC-MDP)This is not the partial observability setting because we have access to a task id or description that uniquely identifies the task, and therefore what objects are referred to by the task-specific state space. (Du et al., 2019; Zhang et al., 2020):

A block contextual Markov decision process (BC-MDP) is defined by tuple $\langle\mathcal{C},\mathcal{S},\mathcal{A},\mathcal{M}^{\prime}\rangle$ where $\mathcal{C}$ is the context space, $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space. $\mathcal{M}^{\prime}$ is a function which maps a context $c\in\mathcal{C}$ to MDP parameters and observation space $\mathcal{M}(c)=\{R^{c},T^{c},\mathcal{S}^{c}\}$.

This brings us to our setting for evaluation, Meta-World (Yu et al., 2020b)We use the following commit from MetaWorld for our experiments: af8417bfc82a3e249b4b02156518d775f29eb289, as a natural instantiation of a BC-MDP. Meta-World proposes a benchmark for meta-RL and multi-task RL, consisting of 50 distinct robotics manipulation tasks, with some example tasks shown in Figure 2. The state space across all tasks is of the same dimensionality, but those dimensions have different semantics across tasks. For example, the same subset of dimensions can refer to a goal position in one task and some object’s position in the other task. We assume access to a family of $N$ MDPs consisting of potentially different reward functions and state spaces of consistent dimensionality, but not necessarily with the same semantic meaning. Unlike previous works, which generally focus on very narrow task distributions, Meta-World provides a diverse task distribution with 50 different tasks involving objects like doors, cups, windows, drawers, etc. and skills like push, pull, open, close, etc. while still providing a shared state and action space. Evaluating on a broad task distribution provides a better estimate of the generalization capabilities of MTRL algorithms.

A Method for Learning Contextual Attention-based Representations

In the multi-task reinforcement learning setting described in Section 2, we propose to factorize the state representation into sub-components that are common across the MDPs within a defined BC-MDP family. While each task has its own state space, there are commonalities across tasks. For example, this can take the form of objects like “drawer” or “door”, and skills like “open” or “close”. In this example, the goal would be to disentangle the state into object-specific and skill-specific representations. We train a universal policy (i.e. one shared policy for all the tasks) that uses the task-specific metadata (or context) to choose a functional representation (e.g. of objects and skills) for any given task. We introduce this compositionality by incorporating a mixture of encoders where different encoders specialize to different aspects for the given family of tasks. In our example, given three tasks “open a door”, “open a drawer”, and “close a drawer”, the encoders could specialize to “open” and “close” skills and “door” and “drawer” objects.

In this section, we describe how we use the metadata to train the different components in CARE. It is important to note that the role of CARE is to learn a representation that enables the incorporation of metadata and functional abstraction. For end-to-end reinforcement learning, it must be paired with a policy optimization algorithm. In the scope of this work, we use Soft Actor-Critic (SAC, (Haarnoja et al., 2018)), but CARE can be paired with any policy optimization method.

In the BC-MDP setting, the different tasks share objects and skills across the family of MDPs, but the state spaces across tasks are context-dependent, and therefore not the same, $\mathcal{S}^{c}$. Our goal is to reconstruct the universal state space $\mathcal{S}$. Of course, we do not have access to the true state space, but a useful inductive bias for the learning agent would be to learn composable representations for objects and attend over these representations for different tasks. One major challenge to this approach is that knowing which objects are relevant for each task requires object-level supervision, which is not commonly available. We propose to sidestep this problem by conditioning the attention on the task context, which is modelled using the easily available task metadata or description. Note that this metadata can be high-level, under-specified, and unstructured. It does not have to explain “how to perform the task”; it can simply describe the task. Even in cases where this metadata is not readily available, it can be easily constructed. An example of a task description from Meta-World could be “Reach a goal position”. Another example is MuJoCo tasks from Deepmind Control (Tassa et al., 2018) and OpenAI Gym (Brockman et al., 2016), where humans identify the tasks by descriptive names like “HalfCheetah Run” and “Maze Solver Ant”, as opposed to task ids $\{1,2,...,n\}$.

Given such a high-level task description, we focus on the case where the task context is captured using pre-trained language models. Specifically, we use the Roberta model (Liu et al., 2019b) to obtain a $768$-dimensional representation of the task description. This representation is projected to a lower-dimensional space using feedforward layers, and the resulting representation is denoted as the context $z_{context}$. The context is used to condition the different components of the policy, as described below.

2 Contextual Attention based Representations

We posit that a useful inductive bias for the training agent is to learn contextual representations for different tasks by learning multiple representations and attending over those representations using the task context. Given $N$ tasks, we use a mixture of $k$ encoders to learn $k$ state-representations. Here, $k$ is a hyperparameter and, in practice, it is much smaller than the number of tasks in the family of MDPs. Note that unlike some work on object-oriented learning, we do not assume access to privileged information in terms of which objects are present in the input observation or useful to encode in a task. Moreover, while our design encourages the specialization of different encoders to combinations of different objects and skills, we note that this setup is incorporating a softer inductive bias than object-oriented learning (Greff et al., 2019; Locatello et al., 2020). Specifically, while we are conditioning the policy on object representations, we do not explicitly model the interactions between the encoders/objects, as is done in recent works like Goyal et al. (2019). We factorize the representation in terms of reusable components, unlike methods that first perform object detection and then model higher-order interactions between the objects using attention-based mechanisms or graph neural networks (Kipf et al., 2018; Pathak et al., 2019; Li et al., 2020). We see our design choice as a trade-off between imposing more useful structure on the algorithm versus requiring less access to privileged information.

Given the $k$ encoders, we compute the $k$ representations $z_{enc}^{i}\forall i\in\{1,\cdots,k\}$. Given the context, $z_{context}$, we compute the normalized soft-attention weights for the encoder representations (denoted as $\alpha_{i}\forall i\in\{1,\cdots,k\}$). We pool the $k$ encoder representations into a fixed-size representation by performing a weighted sum using the soft-attention weights. The resulting encoder representation ($z_{enc}$) is computed as $z_{enc}=\sum_{i}^{k}\alpha_{i}\times z_{enc}^{i}$. We concatenate the encoder representation ($z_{enc}$) with the context representation ($z_{context}$) to obtain the state encoding ($z_{s}$). This state encoding is used as an input to the policy network, and the entire setup is trained end-to-end. Note that the language model is not updated during training and the context representation $z_{context}$ is detached from the computation graph before computing the attention weights. $z_{context}$ is updated using the policy loss directly, as it is a part of the state encoding.

3 Downstream Evaluation

We use Soft Actor-Critic (SAC,( Haarnoja et al. (2018) for downstream evaluation of the learned representations. SAC is an off-policy actor-critic method that uses the maximum entropy framework for soft policy iteration. At each iteration, SAC performs soft policy evaluation and improvement steps. For more details on SAC, refer Haarnoja et al. (2018).

The overall algorithm is described in Algorithm 1, with the sub-function details available in the Appendix (Algorithms 2, 3, 4). We note that steps 3 to 11 can be run concurrently for multiple tasks (as is done in our implementation). The architecture diagram is shown in Figure 3 and the Appendix contains additional implementation details (Appendix A) and hyper-parameters (Appendix B).

Experiments

We now empirically evaluate the effectiveness of the proposed CARE model on Meta-World (Yu et al., 2020b) – a multi-task RL benchmark with 50 tasks. We design our experiments to answer the following questions: i) Is learning contextual attention-based representations an effective mechanism for knowledge transfer in multi-task RL? Does it perform better than methods that do not utilize this context? ii) Is the metadata useful only when learning compositional representations? iii) Does the metadata help to learn factored, specialized representations? iv) Does the metadata help in zero-shot generalization to unseen environments?

Existing works in multi-task RL come in two flavours: i) Extend a single task RL baseline for multi-task by using task-specific parameters (like one policy-head-per-task or per-task entropy regularization). These approaches are generally algorithm-agnostic. ii) Specialized multi-task algorithms. In this second category, we compare against the PCGrad (Yu et al., 2020a) algorithm which is specifically proposed for the Meta-World benchmark and achieves state-of-the-art results, outperforming multi-task algorithms like GradNorm (Chen et al., 2018) and Orthogonal Gradients (CosReg) (Suteu & Guo, 2019). We also compare with the Soft Modularization approach (Yang et al., 2020) that performs routing in a shared policy network to learn different policies for different tasks and provides state-of-the-art results for Meta-World benchmark. Additionally, we compare with a popular and general-purpose conditioning method called FiLM (Perez et al., 2018). We use FiLM layers to condition the encoder on the context (generated using a context encoder, just like in CARE). While FiLM is not a multi-task RL algorithm, it is used effectively in the language-conditioned RL setups (Chevalier-Boisvert et al., 2019; Zhong et al., 2020).

The evaluation performance of the agent is computed as follows: At regular intervals, the agent is evaluated 5 times on each test environment, and the mean of the 5 runs is taken as the success rate for the corresponding environment. These success rates are averaged across the environments to obtain the mean success rate at every interval. A time-series of the mean success rates is obtained by evaluating the agent at regular intervals. The agent is trained for multiple seeds (10 in our case) resulting in 10 different time-series (one per seed) of mean success rates. These 10 time-series are averaged to compute the mean of the success rates (mean over the seeds). The best mean (across the time-series) is reported as the evaluation performance.

We highlight some challenges in comparing the performance of different models on the Meta-World suite. Generally, RL agents are evaluated for continuous-valued episodic rewards (Brockman et al., 2016; Tassa et al., 2018) while Meta-World uses a binary-valued success signal. While the use of success signals is not entirely unheard of (Chevalier-Boisvert et al., 2019), we find that with Meta-World, we can improve the agent’s performance just by increasing the frequency of evaluation (as shown in Table 19 in Appendix). Evaluating the agent more often makes it more likely for the agent to solve a given task, making it harder to compare results across different works. For example, consider the extreme case where the agent is evaluated after every single update. Since we report the best (max) of the mean(success), evaluating more frequently could improve the performance since we are computing the max over a larger set. Thus, evaluation frequency acts like an implicit hyperparameter. We control for this issue in our setup by evaluating every agent at a fixed frequency (once every 10K environment steps, per task). Second, we find that the number of seeds (for evaluation) plays a big role in a model’s reported performance. For example, we report that for MT10, the standard Multi-headed SAC achieves a mean success of 61% (10 seeds), whereas Yu et al. (2020b) reports a success rate of 88% (1 seed), leading to a 44% change in performance. Some other models (Yang et al., 2020) use just 3 seeds. We account for the stochasticity of the evaluation process and ensure a fair comparison of all the models by running all the experiments with 10 seeds. We also report if our model’s improvements are statistically significant or not. Additional details about testing for statistical significance can be found in Appendix D.

We use the same setup as other works. At regular interval, agent is evaluated on all envs

We note that these challenges are not inherent limitations of Meta-World and issues related to seeds affect the evaluation of RL algorithms in general (Henderson et al., 2017). However, we believe that some of these challenges can be alleviated by standardizing the evaluation protocol. Given the usefulness of Meta-World as a multi-task RL benchmark, these challenges should be highlighted to ensure they are considered in the subsequent works.

Meta-World benchmark provides two setups - MT10: a suite of 10 tasks and MT50: a suite of 50 tasks (a superset of MT10). We use both setups for evaluation. In Table 1, we compare the performance of the CARE model for MT10 with the different baselines and report the performance after 2M steps. Following the setup in Kaiser et al. (2019); Srinivas et al. (2020), we also compare the performance in the low-sample regime with 100K steps per task in Table 2 and 500K steps per task in Table 16 (in the Appendix). We note that the proposed CARE model consistently outperforms the other models for the MT10 task. We also note that among specialized multi-task algorithms, PCGrad (Yu et al., 2020a) performs quite poorly in the low-sample regime.

Similarly, in Table 3 and Table 4, we compare the performance of the models for MT50 after 2M and 100K steps respectively. The corresponding table for 500K steps is in the Appendix (Table 17). We note that not only does CARE outperform the other baselines, it is much more sample efficient than other baselines in the low-sample regime.

One additional benefit of the CARE model is that it can be easily combined with more powerful policies and learning algorithms. For example, we find that using multi-headed SAC with the CARE model significantly improves the performance on the MT50 setup (mean success of $0.61$, with a standard error of $0.0287$). Since the CARE model focuses on learning representations, it can benefit from the improvements in policy optimisation for multi-task RL.

In Section C.1 we consider some ablations with the CARE model. First, we vary the number of encoders, showing that increasing the number of encoders hurts performance when shared information is no longer leveraged. We also try using only top-$k$ encoders (i.e. encoders with top-$k$ highest attention scores) with hard attention, confirming the robustness of our method to different aggregation techniques. Finally, we design an experiment where we hardcode the mapping between the tasks and the encoders, showing that mapping encoders to specific objects and skills helps performance.

2 Is the metadata useful only when learning compositional representations?

In Section 4.1, we showed that the metadata is useful for learning compositional representations in the CARE model. A follow-up question is whether the metadata is also useful when using a single encoder. We address this question in Tables 5 and 6 where we compare the performance of the CARE model in the absence of metadata or in the absence of mixture of encoders. We note that removing the metadata (first row) hurts the models more than removing the mixture of encoders (second row). While using metadata with just a single encoder (second row) provides comparable performance to the other baselines, using it with a mixture of encoders leads to even better performance.

3 Interpreting the specialized representations

We perform a visual investigation of what information is being shared across tasks and the role of metadata in the information sharing. Specifically, in Figure 4 we show cosine similarity between the context representations under the 10 tasks of MT10. For each task, we feed the context through the pretrained language model (RoBERTa) and refer to the resulting representation as the pretrained task embedding. In Figure 4(b), we show the cosine similarity between just these pretrained task embeddings, which are the input to the MLP of the context encoder (Figure 3). In Figure 4(c), we show the cosine similarity between the context representations from the proposed CARE model (with 6 encoders) after training. We observe a similar structure in both similarity matrices, that tasks with semantically similar descriptions also have similar context encodings. For example, task 2 and 3 require the agent to interact with the same object. Similarly tasks 4, 5, and 6 are related by the skill “open” and object “drawer”. We now want to observe if these similarities can be found from just the tasks, without access to metadata. In Figure 4(d), we plot the cosine similarity between the context representations from the CARE model (with 6 encoders) without the use of task metadata. Some similarities are present, but nothing as strongly correlated as seen in Figure 4(c). This result further supports our hypothesis that task metadata plays an important role in inferring similarity between tasks, and shows the interpretability of the CARE representation.

4 Zero-shot generalization to unseen environments

Given the compositionality present in language, we want to evaluate if CARE can be used for zero-shot generalization to unseen environments. We train the agents on 8 environments from MT10 and evaluate on two held-out environments, “drawer-open-v1” and “window-open-v1”. There are three training environments that directly relate to these test environments, “drawer-close-v1”, “window-close-v1”, “door-open-v1”, thus allowing for the possibility of zero-shot generalization. In Table 7, we observe that CARE exhibits some promising performance. We note that the comparison is unfair to PCGrad and Soft Modularization as they do not have any means for generalizing to the unseen task, but our motivation is to highlight the potential benefits of using metadata for zero-shot generalization to unseen environments. We further note that CARE generalizes better than FiLM, which also leverages metadata.

Related Work

Multi-task learning holds the promise of accelerating learning across multiple tasks by sharing useful information (Caruana, 1997; Zhang et al., 2014; Kokkinos, 2017; Radford et al., 2019; Rajeswaran et al., 2016; Ruder, 2017; Liu et al., 2019a; Mott et al., 2019; Vithayathil Varghese & Mahmoud, 2020). Multi-task reinforcement learning (MTRL) has been extensively studied with the focus on assumptions around shared properties and structures of different tasks. (Calandriello et al., 2014; Borsa et al., 2016; Maurer et al., 2016) assume that in the multi-task settings, tasks share a common, low dimensional representation, and therefore advocate for learning a shared representation space across all the tasks. (Bräm et al., 2019) considered the setup when the action space between different tasks is not aligned. Zhang et al. (2021) describes multi-task learning algorithms where the tasks have different dynamics but a shared reward structure and makes the assumption of a universal dynamics model. All of these methods for the MTRL setting only utilize a simplistic context in the form of an ordinal task id. Richer contexts in the form of task embeddings are learned online from the differences in reward and dynamics across tasks. We instead focus on the setting where side information is available and can be utilized as a richer context than task ids.

Several works have focused on the problem of negative interference (Du et al., 2018; Suteu & Guo, 2019; Yu et al., 2020a) where the gradients corresponding to the different tasks interfere negatively with each other. Along with slowing down training, conflicting gradients could cause the agent to unlearn knowledge of one task to learn another task. Despite some successes, the proposed approaches are unsatisfactory, either because they increase the computational/memory overhead (for example, Yu et al. (2020a) introduces an $\mathcal{O}(n^{2})$ complexity, where $n$ is the number of tasks) or because they require the training model to ignore some components of the gradients, thus slowing down learning and deteriorating sample efficiency. We propose to use the context for deciding which information should be shared across tasks, thus alleviating negative interference.

Contextual MDPs have been previously defined and analyzed as a setting where side information is exploited for transfer across tasks (Hallak et al., 2015; Modi et al., 2017). Hallak et al. (2015) first defined the contextual MDP setting, drawing connections to contextual multi-arm bandits (Lai & Robbins, 1985; Langford & Zhang, 2008). However, they assume that the state spaces across contexts are the same. Modi et al. (2017) requires an assumption of smoothness in the MDP parameters with respect to the context and examine the online learning scenario, providing an extension of the Rmax algorithm with PAC bounds. They also assume that the context is given. Klink et al. (2019) adopt the contextual MDP setting but with an additional assumption that the agent can control the context, and therefore the task distribution. They propose an algorithm to generate a curriculum that allows the agent to gradually progress to a target context distribution. In our work, we relax several of the assumptions made by these prior works and extend the type of context explored.

Multi-task learning with metadata has been used in You et al. (2016); Zheng et al. (2019) where the focus is on Task Relation Discovery in the context of supervised learning. In Reinforcement Learning, our work has close ties to language-conditioned RL, where natural language phrases have been used as part of task descriptions in the context of several single task RL setups like goal-oriented RL (Chao et al., 2011; Chevalier-Boisvert et al., 2019), grounded language acquisition (Hermann et al., 2017; Chaplot et al., 2017), and instruction following (Tellex et al., 2011; Chen & Mooney, 2011; Williams et al., 2018). Similar to our BC-MDP setting, recent work (Zhong et al., 2020) also defines meta-environments specific to the language-conditioned setting. In all these works, the language description is a part of the problem specification and not just extra information, while our work proposes a way to incorporate auxiliary side information to improve sample efficiency and generalization in multi-task RL. Many of these works make additional assumptions about the language description. For example, (Shu et al., 2017) uses “a two-word tuple template consisting of a skill and an item” to describe the task while the metadata for CARE can be high-level, under-specified, and unstructured. Further, we note that the contextual setting we propose includes information beyond just text, and can also include images or features.

Several works have also focused on learning compositional models for multi-task learning. Liu et al. (2019a) trains a network with task-specific soft-attention modules. Devin et al. (2017) decompose the policy into two components – “task-specific” (shared across all robots) and “robot-specific” (shared across all tasks). Chang et al. (2018) consider a family of algorithmic tasks, with different levels of complexity. Yang et al. (2020) performs routing in a base policy network to generate different policies for different tasks. Unlike our work, these works do not leverage metadata to learn a context on which the module decomposition can be conditioned on.

Discussion

In this work, we highlight an under-explored setting of using contextual information to improve performance in multi-task reinforcement learning. We show that metadata, which is often present in multi-task settings, can be leveraged within the MDP framework to improve performance by enabling more efficient sharing of information across tasks. Further, we define a new setting, the block contextual Markov decision process (BC-MDP), to handle settings where the state space can differ across tasks. Finally, we show that a mixture of encoders can effectively learn a context-dependent representation that can be used with a single policy to solve a family of tasks. We showcase our method, CARE, on Meta-World and achieve new state-of-the-art results.

There are two dimensions along which this work can be extended. Here, we only explore a specific type of context in the form of text descriptions for each task. Other forms of context can also be considered, such as people or places. As examples, in personalized medicine or recommendation systems, the user can be used as a form of context to adapt and share knowledge from other users to transfer the policy. Similarly, in household robotics, the location and layouts of homes will differ although the task is the same.

The relevance of this work also carries over to rich observation settings, or the high-dimensional version of the block contextual MDP setting we examine here (Mozifian et al., 2020). This setting is relevant in robotics, where the agent typically has a first person view of the world. The entire state space (the world) cannot be captured in each task. Instead, the agent only focuses on a relevant subspace of the entire state space at a time. This extension significantly enriches the set of multi-task problems we can tackle and brings us closer to the intractable partial observability (POMDP) setting (Kaelbling et al., 1998; Zhang et al., 2019), but offers more flexibility that allows for additional tractability by limiting the partial observability settings we consider.

Acknowledgements

We thank Edward Grefenstette, Tim Rocktäschel, Danielle Rothermel and Olivier Delalleau for feedback that improved this paper.

References

Appendix A Additional Implementation Details

We use the following open-source libraries: PyTorch (Paszke et al., 2019)https://pytorch.org/, Hydra (Yadan, 2019)https://github.com/facebookresearch/hydra, MetaWorld (Yu et al., 2020b)https://github.com/rlworkgroup/metaworld, MTEnv (Sodhani et al., 2021)https://github.com/facebookresearch/mtenv, MTRL (Sodhani & Zhang, 2021)https://github.com/facebookresearch/mtrl, Numpy (Harris et al., 2020)https://numpy.org/ and Pandas (pandas development team, 2020)https://pandas.pydata.org/. For MetaWorld, we use the following commit-id to run our experiments: https://github.com/rlworkgroup/metaworld/commit/af8417bfc82a3e249b4b02156518d775f29eb289

A.2 SAC Algorithm

We use the SAC policy algorithm (Haarnoja et al., 2018) to learn representation for the CARE model. We provide the pseduo-code for SAC, along with the key equations. For more details, refer to Haarnoja et al. (2018).

The soft value function is trained to minimize the squared residual error

The gradient of Equation 1 is estimated with an unbiased estimator

where the actions are sampled according to the current policy, instead of the replay buffer. The update uses a target value network $V_{\bar{\psi}}$, where $\bar{\psi}$ is an exponentially moving average of the value network weights, which has been shown to stabilize training (Mnih et al., 2015).

The soft Q-function parameters are trained to minimize the soft Bellman residual

which is optimized with stochastic gradients

The policy parameters are learned by minimizing:

A.3 CARE Components

The CARE algorithm uses two components: (i) A context encoder network and (ii) mixture of $k$ encoders (shown in Figure 3). Both components are trained using the policy loss as described in Algorithm 3 and Algorithm 4 respectively. We note that for computing the encoder representation $z_{enc}$, the context encoding $z_{context}$ is detached from the computational graph.

Appendix B Hyperparameter Details

In this section, we provide hyper-parameter values for each of the methods in our experimental evaluation. In Table 8, we provide the hyperparameter values that are common acrss all the methods.

Appendix C Additional Results

In Table 18, we consider some ablations with the CARE model for MT10. First, we vary the number of encoders ($k$). We note that having too many encoders can hurt the performance. Second, we consider the case where we use initialise $m$ encoders but use only top-$k$ encoders at each timestep. Top-$k$ encoders are the ones which have the top-$k$ highest attention scores. The remaining $m-k$ are not used and are considered to be inactive. The attention scores $\alpha$, corresponding to the selected $k$ encoders, are re-normalized to sum to 1. In this case,

C.2 Effect of frequency of evaluation

We find that we can improve the agent’s performance just by increasing the frequency of evaluation as shown in Table 19. In Meta-World, the agent is evaluated for a binary-valued success signal and evaluating the agent more often makes it more likely for the agent to solve a given task. This effect makes it harder to compare the performance of the algorithms from different works. We control for this issue in our setup by evaluating every agent at a fixed frequency (once every 10K environment steps, per task).

Appendix D Testing for statistical significance

We perform a two-tailed, Student’s $t$-distribution test (Student, 1908) under equal sample sizes, unequal variance setup (also called Welch’s $t$-test). The null hypothesis is: the mean performance of the two models (CARE and any baseline) are equal. The significance level ($p$) is set to 0.05.