Paper: Learning Embeddings for Sequential Tasks Using Population of Agents (ICLR-2024)

Reviewer comment (Experiments):
The results included in the paper focus on learning low dimensional embeddings for the tasks - for example, in CARTPOLEVAR the learnt embedding is of dimension 2 or 3 whereas in BASICKAREL it is of dimension 1. The experiments do not consider more difficult tasks, such as MuJoCo tasks considered in [1].

Author response:
**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.

Reviewer comment (Novelty):
There is no discussion of the relatedness / differences with the bisimulation representation learning method in [1] which also learns an embedding of states in RL tasks, and ensures that states which would lead to similar outcomes have similar embeddings. It would help to include a discussion of why it has not been considered as a baseline in the experiments either.

Author response:
**… 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.

Reviewer comment (Writing):
It is a bit confusing to understand the differences between $S_{init}$ and $S$. The authors should consider clarifying in the main paper the differences between a task definition and the MDP states.

Author response:
**… 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.

Reviewer comment (Novelty):
I do not fully understand the PredModel baseline. The authors say it is "inspired by prior work" but there are no citations provided and I may be missing the link to prior work. Could the authors please clarify that?

Author response:
**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.

Reviewer comment (Experiments):
The entire approach seems to depend heavily on the population of agents used to define the learning objectives. For instance, the probability of success (PoS) of a task is taken to be a measure of task difficulty. However, 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. Some of the experimental results seem to be a direct result of the way the agent population is obtained—e.g., 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 which seem to work well for the environments in the paper, but their applicability to new domains is unclear.

Author response:
**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.

Reviewer comment (Theory):
The overall task is assumed to be representable by the initial state. This enables task embedding to be a function of 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.

Author response:
**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.

Reviewer comment (Evaluation):
Some comparisons to baselines seem unfair since the evaluation criterion is based on the population of agents used to learn the embeddings. For instance, in the task selection experiment, task similarity and difficulty (for evaluation purposes) are measured using the same quantities as those used while learning the embeddings. Therefore, the significance of these experiments is unclear.

Author response:
**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.

Reviewer comment (Novelty):
In Section 5.4, PredModel is mentioned to be inspired by prior work. Could you provide a citation for the work this baseline is inspired by?

Author response:
**… 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].

Reviewer comment (Theory):
Why is the start state assumed to represent the task and why is it a reasonable assumption? Are there other ways to represent tasks (so that they can be input to the embedding network) such as natural language descriptions that are better suited here?

Author response:
**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.

Reviewer comment (Reproducibility):
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. Plus, it is almost impossible to get task embedding without exploring a few trajectories of it to get anything meaningful, unlike some rule-description tasks.

Author response:
**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.

Reviewer comment (Experiments):
It only generalizes to variations of a particular environment.

Author response:
**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.

Reviewer comment (Experiments):
How do you guarantee the diversity of the population?

Author response:
**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.

Reviewer comment (Experiments):
Did you test learning task embedding using a single agent?

Author response:
**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.

Reviewer comment (Novelty):
My major issue with the paper surrounds motivation. Creating this task embedding requires a diverse population of agents which together are competent on a broad range of the tasks. This is a vast amount of compute relative to the amount required to solve an individual task or even a reasonably broad range of tasks in the space of tasks. It's therefore not entirely clear to me when such a task embedding would be appropriate. The authors go some way to answering this by demonstrating the usefulness of the embeddings in task prediction and task similarity identification. However, 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. However, I think judging future usefulness and method relevance is very difficult and so do not weight this point too strongly.

Author response:
**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].

Reviewer comment (Experiments):
Because of the large amount of compute required to build these embeddings, the tasks considered are relatively simple. It would be interesting to consider more complex and higher dimensional tasks, such as by embedding levels in ProcGen.

Author response:
**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.

Reviewer comment (Reproducibility):
How much compute is required to generate the population of agents and embeddings for the tasks? I could not find this information, although I may have missed it.

Author response:
**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).


| Environment | Single Agent | Agent Population | Embedding Network |
| ----- | --- | --- | --- |
| MultiKeyNav | $0.5$ | $3$ | $18$ |
| CartPoleVar | $0.5$ | $3$ | $1$ |
| PointMass | $2.5$ | $7.5$ | $40$ |
| Karel | $20$ | $65$ | $240$ |
| BasicKarel | $1$ | $4$ | $6.5$ |
|  |  |  |  |

Reviewer comment (Experiments):
How much variation is there in the embeddings with the population? If I train the population in a different way, can the performance prediction generalise to a different population? For example, can the embedding of a task be used to predict the task performance of an agent trained with a different algorithm?

Author response:
**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.