Learning Embeddings for Sequential Tasks Using Population of Agents

We thank the reviewer for carefully reviewing our paper! We greatly appreciate the feedback. Please see below our responses to the comments.

---

… differences between a task definition and the MDP states …

We would like to provide some intuition by connecting our formalism of tasks as initial states to the contextual MDP setting [Hallak et al., 2015], where each context (e.g., MultiKeyNav's context: agent's initial position, possessed keys initially, door type) corresponds to a distinct task represented by a separate MDP with context-dependent transitions and rewards. This set of MDPs can be converted into an equivalent MDP by including context variables as part of the state. In this converted MDP, each initial state corresponds to a task, as it determines the context for the entire episode. The context is observable.

---

… relatedness / differences with the bisimulation representation learning method in [1]

Our work differs significantly from [1] and other state representation learning methods, as our focus is on learning representations for tasks rather than individual states. Although our definition of tasks as initial states might indicate an interest in state representations, our formalism is connected to the contextual MDP setting and we are interested in learning representations for context. Proximity in the context space doesn't necessarily imply task similarity. This is where our work makes its contribution by mapping the context to task embeddings. We will include a discussion on this in the camera-ready version of the paper.

---

The experiments do not consider more difficult tasks, such as MuJoCo tasks considered in [1].

The environments used in [1] are designed to evaluate robustness to task-irrelevant distractors. These distractors, however, do not introduce any variability in the reward function or transition dynamics. This choice is perhaps deliberate, as the objective is to compare methods for learning state representations. In contrast, we are interested in learning representations for tasks, where variability occurs in the reward function and/or dynamics.

In summary, we use MultiKeyNav as the running example because it is compositional (the agent needs to compose different actions for picking keys in a task-specific manner to unlock the door), which also makes it suitable for ablation experiments. Considering that task variability in MultiKeyNav comes from the reward function, we use CartPoleVar to highlight our framework’s applicability to environments in which task variability comes from the dynamics instead. We select PointMass to test if our framework can handle continuous action spaces. Finally, to investigate our framework’s scalability, we use the real-world environment Karel, which is a challenging environment (with high-dimensional states) that has applications in programming education.

---

I do not fully understand the PredModel baseline. … there are no citations provided …

PredModel is inspired by PEARL [Rakelly et al., 2019]. We discuss the technical details of this baseline in Appendix D.

---

[Hallak et al., 2015] Contextual Markov Decision Processes. arXiv:1502.02259 2015.

[Rakelly et al., 2019] Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables. ICML 2019.

---

We hope that our responses address your concerns and are helpful in improving your rating. If you have any other comments or feedback, please let us know! We are looking forward to hearing back from you! Thank you again for the review.

Dear Reviewer 9bsB,

Since the Author/Reviewer Discussion period concludes on Nov. 22, we would be glad to address any remaining concerns or questions.

Thanks,
Authors

We thank the reviewer for carefully reviewing our paper! We greatly appreciate the feedback. Please see below our responses to the comments.

---

The entire approach seems to depend heavily on the population of agents … it is possible that an “easy” task has a lower PoS when compared to a “difficult” task if a policy solving the easy task is absent in the set of agents.

Indeed, the agent population serves as an important component in our framework. Essentially, we gather task correlation statistics from an agent population that emerges during the multi-task policy learning process. The scenario where “a policy solving an easy task is absent from the agent population” highlights the significance of having a good-quality population. For instance, please refer to our ablation experiment w.r.t. the population specification in Section 5.3, where agents in the population fail to learn to pick keys. This experiment shows how a poor-quality population affects both our proposed similarity criterion and the resulting task embedding space (see Fig. 4b). Thus, the population of agents must be created by appropriately choosing a good learning algorithm and a diversity mechanism. A promising research direction is to develop more principled approaches for generating this population, possibly optimizing for specific diversity measures.

---

the clusters corresponding to unique sets of keys in MultiKeyNav could be a result of using biased task distributions to train agents. Some heuristics are suggested for obtaining the population of agents … but their applicability to new domains is unclear.

Please refer to Appendix E for more details on how the agent populations are created. We utilize action masking for MultiKeyNav, where the effective task is determined by the door type (dictating required keys) and the agent's initial key possession (determining remaining keys to pick up). Thus, the task embedding space for this environment is not merely an outcome of action masking; the embedding network needs to reason about the relationship between door types and initial keys for determining the effective task.

We agree that the heuristics for creating a diverse population of agents might require some domain knowledge. Therefore, a promising research direction is to develop more principled approaches for generating this population, possibly optimizing for specific environment-agnostic diversity measures.

---

The overall task is assumed to be representable by the initial state. … This assumption might not hold in general (several tasks could start from the same state and vice-versa). In such cases, the task is represented by the reward function and the proposed approach is not readily applicable.

Our formalism of tasks as initial states is connected to the contextual MDP setting [Hallak et al., 2015], where each context (e.g., MultiKeyNav's context: agent's initial position, possessed keys initially, door type) corresponds to a distinct task represented by a separate MDP with context-dependent transitions and rewards. This set of MDPs can be converted into an equivalent MDP by including context variables as part of the state. In this converted MDP, each initial state corresponds to a task, as it determines the context for the entire episode. The context is observable. Thus, two MDPs that differ in terms of the reward function can be represented within our framework. This scenario is illustrated in MultiKeyNav, where distinct tasks share similar transition dynamics but may have different reward functions.

---

(the response is continued in Part 2)

(continuation of the response from Part 1)

---

Some comparisons to baselines seem unfair since the evaluation criterion is based on the population of agents used to learn the embeddings.

For Section 6, our baselines correspond to different levels of oracle knowledge. For instance, in the performance prediction benchmark (Section 6.1), OPT represents the oracle that can observe both the agent and the task, and Ignore-Agent represents the one that can observe the task but not the agent. Additionally, in Appendix I, we evaluate the performance of the task embedding networks in the performance prediction application scenario for new agent populations different from those used to train the embedding networks.

For the task selection benchmark (Section 6.2), our baselines consist of various oracles, namely OPT and $\widehat{\text{OPT}}\_\text{50}$, along with alternative one-shot measures of “similarity” between tasks, such as StateSim and TrajectorySim. The purpose of comparing with StateSim and TrajectorySim is to evaluate how our definition of task similarity differs from other one-shot measures. The results indicate that our proposed task similarity criterion captures correlations that extend beyond direct measurements of differences in initial state representations or the edit distance between expert trajectories.

PredModel is included for completeness. However, as we acknowledge in Section 2, our work is positioned as complementary to existing methods.

---

… PredModel is mentioned to be inspired by prior work. Could you provide a citation for the work this baseline is inspired by?

PredModel is inspired by PEARL [Rakelly et al., 2019].

---

Are there other ways to represent tasks … such as natural language descriptions …

As discussed previously, our formalism of tasks as initial states is connected to the contextual MDP setting. The context could correspond to natural language descriptions and is covered by our formalism.

---

… transfer learning and multi-task learning … Are the generated embeddings helpful for these applications as well?

While our framework assumes access to policies that can already solve tasks in a multi-task setup, it would be interesting to explore the possibility of jointly learning task embeddings and multi-task policies using our framework.

---

[Hallak et al., 2015] Contextual Markov Decision Processes. arXiv:1502.02259 2015.

[Rakelly et al., 2019] Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables. ICML 2019.

---

We hope that our responses address your concerns and are helpful in improving your rating. If you have any other comments or feedback, please let us know! We are looking forward to hearing back from you! Thank you again for the review.

Dear Reviewer HkeS,

Since the Author/Reviewer Discussion period concludes on Nov. 22, we would be glad to address any remaining concerns or questions.

Thanks,
Authors

We thank the reviewer for carefully reviewing our paper! We greatly appreciate the feedback. Please see below our responses to the comments.

---

It is unclear how task embedding is useful to me. As it requires checkpoints of learned policies that almost solve the task and "difficulty" is vague to an agent's performance as an agent may take a different path to solve the task when there are multiple solutions.

Our framework utilizes an agent population that emerges during the multi-task policy learning process to collect task correlation statistics and estimate task difficulty. The proposed task similarity criterion is based on the collective behavior of the agent population, while task difficulty, defined as the probability of a random agent from the population failing on the task, is an aggregate statistic computed for the entire agent population, not individually per agent. Thus, these quantities are population-level statistics.

The task embedding network learned through our framework offers the following advantages:

1. It enables one-shot computation of task similarity without requiring experience data or access to the agent population.
2. It can generalize to new tasks, as demonstrated in Appendix H. More concretely, we show that it generalizes from a training set to a validation set of tasks not encountered during agent population training or embedding network learning.
3. Our proposed task similarity criterion, which is used to train the embedding network, captures performance correlations that remain valid for new agents. We demonstrate this in Appendix I by assessing the performance of the task embedding networks in the performance prediction application scenario for different agent populations.

---

It only generalizes to variations of a particular environment.

Our formalism of tasks is connected to the contextual MDP setting [Hallak et al., 2015], where each context (e.g., MultiKeyNav's context: agent's initial position, possessed keys initially, door type) corresponds to a distinct task represented by a separate MDP with context-dependent transitions and rewards. It's important to note that proximity in the context space doesn't necessarily imply task similarity. This is where our work makes its contribution by mapping the context to task embeddings. However, we must note that our framework is general enough to cover the scenario where tasks correspond to different environments. The challenge here lies in learning multi-task policies that can jointly solve these environments.

---

How do you guarantee the diversity of the population?

We introduce action masking or biased task distributions, as ways to diversify the agent population. A promising research direction is to develop more principled approaches for generating this population, possibly optimizing for specific diversity measures.

---

Did you test learning task embedding using a single agent?

As previously mentioned, our proposed task similarity criterion is formulated as a population-level statistic. When the agent population is reduced to a singleton, it evaluates to zero by definition, which effectively makes the criterion uninformative.

---

How the agent performance data is collected? What agents did you use? Were they involved in the training of the embedding?

The benchmarks for the agent performance prediction scenario in Section 6.1 were created with the same agent population used to learn the task embedding network. However, our proposed task similarity criterion captures performance correlations that remain valid for new agents. We demonstrate this in Appendix I by assessing the performance of the task embedding network in the performance prediction application scenario for different agent populations.

---

How to test the generalization of the task embedding? (aka generalize across different tasks.)

The learned task embedding network can generalize to new tasks, as demonstrated in Appendix H. More concretely, we show that it generalizes from a training set to a validation set of tasks not encountered during agent population training or embedding network learning.

---

[Hallak et al., 2015] Contextual Markov Decision Processes. arXiv:1502.02259 2015.

---

We hope that our responses address your concerns and are helpful in improving your rating. If you have any other comments or feedback, please let us know! We are looking forward to hearing back from you! Thank you again for the review.

Dear Reviewer Q8ok,

Since the Author/Reviewer Discussion period concludes on Nov. 22, we would be glad to address any remaining concerns or questions.

Thanks,
Authors

Dear authors,

Thank you for providing the response to address my questions! I have to appreciate that I can see the great effort of the authors in the paper and in the rebuttal phase.

I read through other reviews and responses as well to try to look for answers to some of my questions, I think all the reviewers share the following three concerns:

• (1) Motivation of the task embedding.
• (2) Diversity guarantee of the population. (Is "hard" really hard, and "easy" really easy)
• (3) Generalization across environments. (I think I am clear about it at this point that the "task" is defined within an environment so I am not going to be harsh with this.)

For the motivation, I am wondering if the authors could further elaborate on how "task embedding" can actually help with "what". For example, are you using it to improve the curriculum learning task selection or improve other aspects of any algorithms? Does it help with any efficiency of algorithms or is it a solid evaluation or estimation methodology, or other applications? I would like to see that statement as the current explanation of framework advantage is not quite convincing yet. For example, "one-shot computation of task similarity", yes, and does it imply "quicker adaptation of the policy to other tasks"?

I did see another paper https://arxiv.org/abs/2310.07218 that also adopts populations and information-theoretic metrics working on different topics, but their motivation is clearer to me -- estimate the need for multi-agent interactions through a less computationally expensive method (selfplay) to decide whether we need a more complex algorithm (fictitiou co-play https://arxiv.org/abs/2110.08176).

We thank the reviewer for carefully reviewing our paper! We greatly appreciate the feedback. Please see below our responses to the comments.

---

My major issue with the paper surrounds motivation. … it's not clear to me when either of these tasks would be useful compared to training a single agent on a broader task distribution for the same total compute time required to train the population.

We agree that the process of learning the embedding network could be compute-intensive for certain environments. However, post-learning, this network offers the following advantages:

1. It enables one-shot computation of task similarity without requiring experience data or access to the agent population.
2. It can generalize to new tasks, as demonstrated in Appendix H. More concretely, we show that it generalizes from a training set to a validation set of tasks not encountered during agent population training or embedding network learning.
3. Our proposed task similarity criterion, which is used to train the embedding network, captures performance correlations that remain valid for new agents. We demonstrate this in Appendix I by assessing the performance of the task embedding networks in the performance prediction application scenario for different agent populations.

The two downstream scenarios presented in our work are inspired by real-world applications; the performance prediction scenario is similar to assessing a student’s proficiency in adaptive learning platforms via a compact quiz [He-Yueya et al., 2021], and the task selection scenario is analogous to selecting desired questions from a pool for a personalized learning experience in online education systems [Ghosh et al., 2022].

---

Because of the large amount of compute required … the tasks considered are relatively simple.

We selected environments with apparent variability along reward and/or dynamics axes. In summary, we use MultiKeyNav as the running example because it is compositional (the agent needs to compose different actions for picking keys in a task-specific manner to unlock the door), which also makes it suitable for ablation experiments. Considering that task variability in MultiKeyNav comes from the reward function, we use CartPoleVar to highlight our framework’s applicability to environments in which task variability comes from the dynamics instead. We select PointMass to test if our framework can handle continuous action spaces. Finally, to investigate our framework’s scalability, we use the real-world environment Karel, which is a challenging environment (with high-dimensional states) that has applications in programming education.

---

How much compute is required to generate the population of agents and embeddings for the tasks?

Below we report the approximate run times (in hours) to train a single agent, an agent population, and the embedding network for the environments considered in this work. All the experiments were conducted with an Intel Xeon Gold 6134M CPU (clocked at 3.20 Ghz) and an Nvidia Tesla V100 GPU (32 GB VRAM configuration).

---

How much variation is there in the embeddings with the population? … can the performance prediction generalise to a different population?

Indeed, the agent population serves as an important component in our framework. Essentially, we gather task correlation statistics from an agent population that emerges during the multi-task policy learning process. For instance, please refer to our ablation experiment w.r.t. the population specification in Section 5.3, where agents in the population fail to learn to pick keys. This experiment shows how a poor-quality population affects both our proposed similarity criterion and the resulting task embedding space (see Fig. 4b).

As previously mentioned, given a diverse population of agents, our proposed task similarity criterion captures performance correlations that remain valid for new agents. We demonstrate this in Appendix I by assessing the performance of the task embedding network in the performance prediction application scenario for agent populations trained differently.

---

[He-Yueya et al., 2021] Quizzing Policy Using Reinforcement Learning for Inferring the Student Knowledge State. EDM 2021.

[Ghosh et al., 2022] Adaptive Scaffolding in Block-based programming via Synthesizing New Tasks as Pop quizzes. AIED 2022.

---

We hope that our responses address your concerns and are helpful in improving your rating. If you have any other comments or feedback, please let us know! We are looking forward to hearing back from you! Thank you again for the review.

Dear Reviewer viG4,

Since the Author/Reviewer Discussion period concludes on Nov. 22, we would be glad to address any remaining concerns or questions.

Thanks,
Authors