---
title: "Learning Embeddings for Sequential Tasks Using Population of Agents"
authors:
  - "Mridul Mahajan"
  - "Georgios Tzannetos"
  - "Goran Radanovic"
  - "Adish Singla"
doi: "10.24963/ijcai.2024/523"
url: "https://ijcai.org/proceedings/2024/523"
pdf: "https://ijcai.org/proceedings/2024/0523.pdf"
venue: "Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence"
year: 2024
pages: "4733-4741"
track: "Main Track"
keywords:
  - "Machine Learning: ML: Reinforcement learning"
  - "Machine Learning: ML: Representation learning"
  - "Machine Learning: ML: Multi-task and transfer learning"
---

# Learning Embeddings for Sequential Tasks Using Population of Agents

Mridul Mahajan, Georgios Tzannetos, Goran Radanovic, Adish Singla

*Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence*, pages 4733-4741. DOI: [10.24963/ijcai.2024/523](https://doi.org/10.24963/ijcai.2024/523)

## Abstract

We present an information-theoretic framework to learn fixed-dimensional embeddings for tasks in reinforcement learning. We leverage the idea that two tasks are similar if observing an agent's performance on one task reduces our uncertainty about its performance on the other. This intuition is captured by our information-theoretic criterion which uses a diverse agent population as an approximation for the space of agents to measure similarity between tasks in sequential decision-making settings. In addition to qualitative assessment, we empirically demonstrate the effectiveness of our techniques based on task embeddings by quantitative comparisons against strong baselines on two application scenarios: predicting an agent's performance on a new task by observing its performance on a small quiz of tasks, and selecting tasks with desired characteristics from a given set of options.

## Full text

## 1 Introduction

Embeddings are widely used to represent data points as vectors in a space that captures meaningful relations between them [Sun et al., 2014; Sung et al., 2018; Athar et al., 2020; Mikolov et al., 2013; Pennington et al., 2014; Cer et al., 2018; Zhang et al., 2021]. They could also be utilized as representations for tasks, as studied in various areas such as multi-task learning [Zhang et al., 2018], meta-learning [Achille et al., 2019], and domain-adaptation [Peng et al., 2020]. In reinforcement learning (RL), task embeddings could be used to understand the shared structure in sequential decision-making problems if similar tasks are embedded in close proximity. This enables efficient, one-shot computation of task similarity, eliminating the need for time-consuming policy rollouts.

Despite the potential benefits, prior work on learning task embeddings in RL [Qin et al., 2022; Schäfer et al., 2022; Arnekvist et al., 2018; Yoo et al., 2022; Rakelly et al., 2019; Bing et al., 2023; Gupta et al., 2018; Fu et al., 2020; Li et al., 2021; Lan et al., 2019; Walke et al., 2022; Sodhani et al., 2021b; Vuorio et al., 2019] does not explicitly optimize for task similarity. This could primarily be attributed to the lack of a general framework to measure (and reason about) similarities among sequential tasks.

To this end, we introduce an information-theoretic framework to learn fixed-dimensional embeddings for tasks in RL; the inner product in the embedding space captures similarity between tasks, and the norm of the embedding induces an ordering on the tasks based on their difficulties (see Fig. 1). A critical component of the framework is an agent population exhibiting diverse behaviors, which serves as an approximation for the space of agents. It leverages the idea that two sequential tasks are similar to each other if observing the performance of an agent from this population on one task significantly decreases our uncertainty about its performance on the other.

Essentially, there is an underlying notion of skills required to solve sequential tasks, and several tasks require some skills in common. As these skills are unknown, we measure task similarity by gathering task correlation statistics from an agent population that emerges during the multi-task policy learning process. More concretely, we introduce an informationtheoretic criterion to measure task similarity (Section 4.1). Through this, we construct a set of ordinal constraints on the embeddings (with each such constraint asserting the relative similarity between a triplet of tasks), and propose a training scheme for an embedding network to learn them (Section 4.2).

Besides assessing the learned embedding spaces through visualizations (Section 5) and comparisons with prior work (Section 6), we ground our framework in two downstream scenarios that are inspired by real-world applications (Section 7). Firstly, we show its utility in predicting an agent's performance on a new task given its performance on a small quiz of tasks, which is similar to assessing a student's proficiency in adaptive learning platforms via a compact quiz [He-Yueya and Singla, 2021]. Several methods for curriculum design in RL rely on the agent's value function [Klink et al., 2020b; Eimer et al., 2021; Yengera et al., 2021]. Thus, one concrete application of task embeddings is in black-box curriculum design, where only the agent's performance is observable, not the value function. We demonstrate that task embeddings can be used to the infer the agent's value function by observing its performance on a small quiz. Secondly, we demonstrate the application of our framework in selecting tasks with desired characteristics from a given set of options, such as choosing tasks that are slightly harder than a reference task. This is analogous to selecting desired questions from a pool for a personalized learning experience in online education systems [Ghosh et al., 2022]. Through comparisons with strong baselines on diverse environments, we show the efficacy of our techniques based on task embeddings.

Figure 1: Schematics of our approach. We learn a task embedding function fϕ(.) that maps a task s to its fixed-dimensional representation E. In this illustration, we show the properties of the learned embeddings using the MULTIKEYNAV environment in which tasks require the agent (shown as a black circle) to pick-up certain keys (from the gray segments) to unlock the door (the right-most segment) that has certain requirements (shown in color in the form of gates). A possible solution trajectory is depicted using dotted lines. Keys on this trajectory correspond to the ones that the agent possesses at that point in time. For instance, in task s2, the agent starts off with the yellow key in possession already. ⟨E1, E2⟩ is greater than ⟨E1, E3⟩, since tasks s1 and s2 have a common requirement of picking the blue key, and thus, are similar. Additionally, ∥E2∥2 is less than both ∥E1∥2 and ∥E3∥2 , since task s2 requires picking a single key, while tasks s1 and s3 require picking two keys, which makes them harder than s2.

![Image](https://ijcai.org/proceedings/2024/fulltext/523/image_000000_b179237dbc89798236675b4cdca600aa7b64696449076b70c9c13eea2fab775b.png)

To summarize, our work makes the following contributions:

- I. We introduce an information-theoretic framework to learn task embeddings in RL. As part of the framework, we propose a task similarity criterion which uses a diverse population of agents to measure similarity among sequential tasks (Section 4.1).
- II. We propose a scheme to learn task embeddings by leveraging the ordinal constraints imposed by our similarity criterion (Section 4.2).
- III. To assess our framework, we perform visual assessments of the learned embedding spaces, and introduce two quantitative benchmarks: (a) agent's performance prediction, and (b) task selection with desired characteristics (Sections 5, 6, and 7). 1

## 2 RelatedWork

Task embeddings in RL. Several works in meta-learning and multi-task learning have explored the use of embeddings to model relationships between sequential tasks, where embeddings are either learned explicitly through objectives such as reconstruction [Arnekvist et al., 2018; Yoo et al., 2022; Bing et al., 2023] and trajectory-based contrastive learning [Fu et al., 2020; Li et al., 2021], or implicitly to aid generalization to new tasks [Lan et al., 2019; Walke et al., 2022; Sodhani et al., 2021b; Vuorio et al., 2019]. While these methods integrate embeddings with policies solely to improve performance, we propose a framework to learn general-purpose embeddings that can be used to quantify and analyze task similarities. Furthermore, in our framework, embedding computation is a one-shot operation, unlike prior work that relies on experience data from the policy for the task. These distinctions position our work as complementary to existing methods.

Population-based techniques. Our framework requires a diverse agent population. This is inline with [Furuta et al., 2021; Tylkin et al., 2021; Vinyals and et al., 2019; Jaderberg and et al., 2019; Parker-Holder et al., 2020], which use agent populations in the RL setting. For instance, Furuta et al. [2021] use a randomly generated agent population to empirically estimate policy information capacity, an information-theoretic measure of task difficulty in RL.

[1 GitHub repository: https://github.com/machine-teaching-group/ ijcai2024-task-embeddings-rl.](https://github.com/machine-teaching-group/ijcai2024-task-embeddings-rl)

## 3 Problem Setup

MDP and Tasks. We use the Markov Decision Process (MDP) framework to define an environment. An MDP M is defined as a 6-tuple (S, A, R, T , Sinit, γ), where S is the state space, A is the action space, R : S ×A→R is the reward function, T : S × S × A → [0, 1] is the transition dynamics, and Sinit ⊆ S is the set of initial states. Each state s ∈ Sinit corresponds to a goal-based task (for example, the goal could be to reach a specific destination in a navigation task). Our formalism of tasks as initial states is connected to the contextual MDP setting [Hallak et al., 2015] (see Section 6). The reward is 0 on all transitions but those on which a task gets completed. On task completion, the reward is 1. As an alternative to discounting, at each time step, there is a failure probability of 1 - γ, which incentivises the agent to solve the task quickly. This ensures that the cumulative reward is binary.

Population of agents and task solvability. p(θ) represents a distribution over the population of agents. Concretely, it is a distribution over the agents' policy parameters. We use Θ to represent the random variable that takes on the value θ. Os,Θ ∈ {0, 1} is a Bernoulli random variable that takes on the value 1 if, on a rollout, an agent sampled from p(θ) could successfully solve the task s ∈ Sinit (i.e., the cumulative reward is 1), and 0 otherwise. We call Os,Θ the optimality variable for task s. POS(s) := E[Os,Θ] denotes the probability of success on s, and is the complement of task difficulty.

Task embedding space. Formally, we wish to learn a task embedding function (parameterized by ϕ) fϕ : Sinit → Rn , for an MDP M and a prior over the population of agents p(θ), that maps tasks to n-dimensional representations. The range of fϕ(.) is the task embedding space.

Objective. Our objective is to learn embeddings for sequential tasks with the following properties: (a) the inner product in the embedding space captures task similarity, where the realizations of optimality variables are similar for tasks that are embedded in close proximity, and (b) the norm of the embedding induces an ordering on the tasks based on their difficulties. We formalize these objectives in Section 4.

## 4 Learning Framework

In Sections 4.1, we formally define our information-theoretic criterion to measure task similarity in RL. In Section 4.2, we view the problem of learning task embeddings through the lens of ordinal constraint satisfaction.

## 4.1 Information-Theoretic Task Similarity Measure

Our goal is to measure similarity between sequential tasks. To this end, we propose the mutual information between task optimality variables as a measure of task similarity. This metric captures the intuition that two tasks are similar to each other if observing an agent's performance on one task reduces our uncertainty about its performance on the other. We begin by formally defining performance uncertainty. Thereafter, we provide a formal definition of our task similarity criterion.

Definition 1 (Performance Uncertainty). The entropy of the population with prior p(θ) solving a task s is defined as:

<!-- formula-not-decoded -->

where Os,Θ is the optimality variable for s.

Thus, the similarity between tasks si, sj ∈ Sinit could be measured as the reduction in H(Osi,Θ) by observing Osj ,Θ.

Definition 2 (Task Similarity). Given a prior over the population of agents p(θ), we measure the similarity between two tasks si, sj ∈ Sinit as the mutual information I(.; .) between their optimality variables Osi,Θ, Osj ,Θ:

<!-- formula-not-decoded -->

It quantifies the information obtained about Osi,Θ by observing Osj ,Θ.

A comprehensive algorithm detailing the empirical estimation of I is provided in the longer version of the paper [Mahajan et al., 2023]. Given an MDP M and a prior distribution of the agent parameters p(θ), our algorithm uses N samples to estimate I(Osi,Θ;Osj,Θ). For each sample, it randomly samples θl ∼ p(θ), and performs rollouts of πθl from si and sj to obtain estimates of the probability mass functions required for computing I. The estimation procedure can be invoked with the signature ESTIMATE(si, sj , M, π, p(θ), N).

## 4.2 Learning Task Embeddings

With the task similarity criterion defined, we are interested in learning a task embedding function fϕ : Sinit → Rn (consequently, an embedding space) that satisfies the desiderata introduced in Section 3. To this end, we pose the problem of learning fϕ(.) as an ordinal constraint satisfaction problem. Essentially, the similarity criterion I imposes a set CMI of triplet ordinal constraints on the task embeddings. POS(.) imposes another set CNORM of pairwise ordinal constraints.

Concretely, CMI is a collection of ordered triplets of tasks s.t. for each (s1, s2, s3) ∈ CMI, I(Os1,Θ; Os2,Θ) &gt; I(Os1,Θ; Os3,Θ). Consequently,

## Algorithm 1 Learn the Task Embedding Function (fϕ)

- 1: procedure TRAIN(Set of tasks Sinit, MDPM, Policy π, Prior distribution of the agent parameters p(θ), Number of samples N, Hyperparameter λ, Number of iterations M)

```
2: Initialize ϕ. 3: for i ∈ {1, . . . , M} do 4: Sample task s1, s2, s3 ∼ Sinit. 5: E1, E2, E3 ← fϕ(s1), fϕ(s2), fϕ(s3) 6: ˆ I12 ← ESTIMATE(s1, s2,M, π, p(θ), N) 7: ˆ I13 ← ESTIMATE(s1, s3,M, π, p(θ), N) 8: if ˆ I12 > ˆ I13 then 9: loss ← log(1 + exp(⟨E1, E3⟩ - ⟨E1, E2⟩)) 10: else 11: loss ← log(1 + exp(⟨E1, E2⟩ - ⟨E1, E3⟩)) 12: Sample task s4, s5 ∼ Sinit. 13: E4, E5 ← fϕ(s4), fϕ(s5) 14: if POS(s4) > POS(s5) then 15: loss ← loss + λ log(1 + exp(∥E4∥2 - ∥E5∥2 )) 16: else 17: loss ← loss + λ log(1 + exp(∥E5∥2 - ∥E4∥2 )) 18: Update ϕ to minimize loss. 19: return ϕ
```

we would like to satisfy the constraint ⟨fϕ(s1), fϕ(s2)⟩ &gt; ⟨fϕ(s1), fϕ(s3)⟩. Likewise, CNORM is a collection of ordered tuples of tasks s.t. for each (s1, s2) ∈ CNORM, POS(s1) &gt; POS(s2). Consequently, we would like to satisfy the constraint ∥fϕ(s2)∥2 &gt; ∥fϕ(s1)∥2 (embeddings for easier tasks have smaller norm).

We learn the task embedding function fϕ(.), for an MDP M and a prior over the agent population p(θ), by optimizing the parameters ϕ to maximize the log-likelihood of the ordinal constraints under the Bradley-Terry-Luce (BTL) model [Luce, 1959]. Concretely, given a task triplet (s1, s2, s3), we define:

<!-- formula-not-decoded -->

Similarly, given a tuple of tasks (s1, s2), we define:

<!-- formula-not-decoded -->

Hence, the task embedding function fϕ(.) is learned by solving the following optimization problem:

<!-- formula-not-decoded -->

where Ei denotes fϕ(si), and λ is a hyperparameter. The pseudocode for the proposed algorithm to learn the task embedding function fϕ(.) is given in Algorithm 1.

## 5 Experiments: Visualizing Embedding Spaces

In this section, we visualize the embedding spaces to gather qualitative insights, addressing the following research questions: (i) Can distinct clusters of tasks be identified by visualizing the embedding space? (ii) How does regularization

| Environment                                        | Task Variability                                                                        | Action      | State                                                 | Number of Tasks                        |
|----------------------------------------------------|-----------------------------------------------------------------------------------------|-------------|-------------------------------------------------------|----------------------------------------|
| MULTIKEYNAV CARTPOLEVAR POINTMASS KAREL BASICKAREL | Reward Function Dynamics Dynamics Reward Function + Dynamics Reward Function + Dynamics | 7 2 R2 52 6 | R × {0, 1}6 R5 ×{0, 1}× [200] R7 {0, 1}51840 {0, 1}88 | Infinite Infinite Infinite 73688 24000 |

(a) Comparison of environments' complexity

Figure 2: We evaluate our framework on a diverse set of environments. (a) compares the characteristics of these environments. (b) illustrates these environments for a better understanding of the tasks.

![Image](https://ijcai.org/proceedings/2024/fulltext/523/image_000001_26f79f70c39706c4106d81acd26a33a3af7878960e3e1169cee3c5ddc8c9516d.png)

through CNORM affect the embedding space? (iii) What influence do agent population and environment specification have on the embedding space? We begin by discussing the rationale for environment selection, describing these environments. Subsequently, we provide an overview of the embedding networks' training process, followed by the qualitative results.

## 5.1 Environments

We evaluate our framework on environments with diverse characteristics to demonstrate its generality and scalability to different sequential decision-making problems (see Fig. 2). As the running example, we use MULTIKEYNAV (based on [Devidze et al., 2021]) because of its compositional nature in which the agent needs to compose different actions for picking keys (with four distinct key types, each requiring a specific action to be picked) in a task-specific manner to unlock the door. This also makes it suitable for ablation experiments. Task variability comes from the agent's initial position, the keys that it possesses initially, and the door type (with each type requiring a unique set of keys).

Given that task variability in MULTIKEYNAV comes from the reward function, we use CARTPOLEVAR to highlight our framework's applicability to environments where it comes from the dynamics instead. This environment is a variation of the classic control task from OpenAI gym [Brockman et al., 2016], and also takes inspiration from [Sodhani et al., 2021a] in which the forces could be negative as well. Tasks require keeping a pole attached by an unactuated joint to a cart upright for 200 timesteps by applying forces to the left (action 0) or to the right (action 1) of the cart. Task variability comes from the force F applied on the cart by each action, and the TaskType ∈ {0, 1}. Tasks of Type 0 involve "Pulling" with action 0 pulling the cart from the left and action 1 pulling the cart from the right, while tasks of Type 1 involve "Pushing".

We select POINTMASS (introduced in [Klink et al., 2020a]) to test if our framework can handle continuous action spaces. Here, the agent applies forces to control a point mass inside a walled square. Tasks require reaching a fixed goal position through a gate, with variability arising from the gate width and position, along with the space's kinetic friction coefficient.

Finally, we investigate our framework's scalability using the real-world environment KAREL from [Bunel et al., 2018], which is a challenging environment (with high-dimensional states) that has applications in programming education. Tasks require the agent to synthesize a program, potentially containing control flow constructs such as loops and conditionals, satisfying a given specification comprising input-output examples. This program serves as a controller for an avatar navigating a grid, where each cell could contain marker(s), correspond to a wall, or be empty. The avatar can traverse the grid and manipulate it by picking or placing markers. Thus, an example in the specification comprises the Pre-Grid and the corresponding Post-Grid. In addition, we evaluate it on BASICKAREL [Tzannetos et al., 2023], which is a simpler variant of KAREL that excludes control flow constructs.

## 5.2 Training Process

To learn the embedding function, we first obtain the agent population by taking snapshots while training a neural network policy using either behavioral cloning [Bain and Sammut, 1995] or policy gradient methods [Sutton et al., 1999]. Concretely, a snapshot is recorded if the average performance on a validation set of tasks (denoted as Ssnap) improves by δsnap compared to the previously recorded snapshot. A snapshot of the untrained policy is recorded by default. Different subpopulations, obtained by either masking actions or by using biased task distributions during training, are combined to form the final population. Here, masking an action corresponds to setting its logit to a large negative number. Using biased task distribution during training is another way to inject diversity into the population. In MULTIKEYNAV, for instance, this could correspond to assigning low probability mass to tasks with certain types of doors in the initial state distribution during training. Finally, we parameterize the embedding function fϕ(.) with a neural network, optimizing its parameters as described in Algorithm 1. We provide additional details in the longer version of the paper [Mahajan et al., 2023].

Figure 3: Visualization of the task embedding spaces learned through our framework. Each point represents a task, and the size of the points is proportional to the norm of the embeddings.

![Image](https://ijcai.org/proceedings/2024/fulltext/523/image_000002_02eb7a8079198f3f484d41803ab6e1af95ea41ca991322e7c4579600b461b3ea.png)

Figure 4: Task embedding spaces for the MULTIKEYNAV environment: (a) without CNORM, (b) pickKey actions masked, (c) all doors require KeyA, KeyB, and (d) all doors require KeyA.

![Image](https://ijcai.org/proceedings/2024/fulltext/523/image_000003_53bd97e74c508d75868be90e41bc81c8a11eaa6f6bf843edd5ef64922865e084.png)

## 5.3 Visualizations and Qualitative Results

We visualize the embedding spaces on a 2-dimensional map using t-SNE [van der Maaten and Hinton, 2008] to identify distinct clusters of tasks. Although t-SNE preserves the local structure, it does not necessarily preserve the embeddings' norm. For this reason, we scale the points in proportion to the norm of the embeddings. Additionally, we provide PCA plots in the longer version of the paper [Mahajan et al., 2023].

Visualizations. For MULTIKEYNAV (Fig. 3a), our framework discovers distinct clusters of tasks, with each cluster corresponding to a unique set of keys that need to be picked. The norm of the embeddings is in accordance with the number of keys that need to be picked (with tasks requiring navigation only having the smallest norm). Additionally, tasks in clusters adjacent to each other share a common key requirement. For CARTPOLEVAR (Fig. 3b), it discovers that each task exhibits one of two types of underlying dynamics. In one (+ve F and Type 0, or -ve F and Type 1), action 0 moves the cart to the left, while in the other (-ve F and Type 0, or +ve F and Type 1), it moves the cart to the right. For POINTMASS (Fig. 3c), it discovers three clusters of tasks based on the behavior that the agent needs to exhibit near the gate. The first cluster includes tasks in which the agent need not steer to cross the gate, while the second and third clusters contain tasks in which it must steer left or right to cross the gate, respectively. For KAREL and BASICKAREL (Fig. 3d and 3e), it discovers different clusters based on whether the solution code requires loops or conditionals, and whether the agent needs to pick or put markers in the grid, respectively. In the longer version of the paper [Mahajan et al., 2023], we show that the learned embedding network can generalize to new tasks.

Ablation w.r.t. CNORM. Fig. 4a shows the task embedding space learned without the norm ordinal constraints CNORM (i.e., λ is set to 0). As expected, the norm of the embeddings is not proportional to the number of keys that need to be picked. Instead, the points are nearly uniform in size.

Ablation w.r.t. population specification. To understand the effect of population on the embedding space, we learn the embedding function fϕ(.) for MULTIKEYNAV using an agent population in which pickKey actions are masked (Fig. 4b).

In this case, we obtain two distinct clusters of tasks - one of the clusters contains tasks that cannot be solved (these require picking key(s)), and the other contains tasks that require navigation only. These results emphasize the importance of the population's quality in learning a good task embedding space.

Ablation w.r.t. environment specification. We check the impact of the environment specification on the embedding space through the following MULTIKEYNAV variants: (a) each door requires KeyA and KeyB (Fig. 4c), i.e., they have identical key requirements, and (b) each door requires KeyA only (Fig. 4d). Modifying the environment specification changes the task semantics, thereby impacting the embedding space. Thus, these results are inline with our intuition.

## 6 Experiments: Comparison with PriorWork

In this section, we compare our framework with prior work using silhouette scores, which assess clustering quality in the embedding spaces based on intuitively identified task clusters. To enable this comparison, we introduce the PredModel baseline, inspired by prior work (e.g., PEARL [Rakelly et al., 2019]). We also compare our method against embedding networks with random weights (RandomModel).

Most existing methods use variational inference to learn latent context from task-specific experience data, where the inference network could be trained to reconstruct the MDP for the task through predictive models of reward and dynamics. To adapt this approach to our setting, we connect 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. By including context variables as part of the state, this set of MDPs can be converted into an equivalent MDP, where each initial state represents a task.

The modifications needed for PredModel are as follows: Firstly, since context is observable in our setup, we condition the approximate posterior over the embeddings on the initial state, eliminating the need for experience data. Secondly, we train the predictive models on states with context variables removed, ensuring the utilization of the embedding that the model is conditioned on. We provide additional details in the longer version of the paper [Mahajan et al., 2023].

Results. Table 1 reports the silhouette scores, averaged across 3 random seeds, with 1000 tasks per seed (5000 for BASICKAREL). The scores for the models learned through our framework are consistently better. While PredModel clusters similar tasks together in the embedding space for CARTPOLEVAR, it fails to do so in rest of the environments. In contrast to CARTPOLEVAR, where task variability comes from dense differences in the dynamics, task variability in other environments comes from sparse differences in the reward function and/or dynamics. Thus, we hypothesize that PredModel fails on environments with sparse variability across tasks.

## 7 Experiments: Application Scenarios

In this section, we evaluate our framework on two application scenarios: performance prediction, and task selection. We conduct this evaluation on MULTIKEYNAV and CARTPOLEVAR, as they cover two distinct sources of task variability, namely reward function and dynamics.

Table 1: Comparison of silhouette scores (higher is better) based on intuitively identified clusters of tasks in the learned embedding spaces. The scores for our models are consistently better.

| Environment                                  | RandomModel                                             | PredModel                                               | Ours                                                |
|----------------------------------------------|---------------------------------------------------------|---------------------------------------------------------|-----------------------------------------------------|
| MULTIKEYNAV CARTPOLEVAR POINTMASS BASICKAREL | +0.036 ±0.048 +0.015 ±0.016 +0.104 ±0.026 -0.058 ±0.007 | -0.037 ±0.003 +0.242 ±0.007 -0.010 ±0.004 -0.002 ±0.003 | +0.753±0.001 +0.325±0.009 +0.380±0.019 +0.811±0.019 |

Figure 5: Results for performance prediction using task embeddings. Our method (listed as Ours) is competitive with the OPT baseline, which is the best one could do on this benchmark.

![Image](https://ijcai.org/proceedings/2024/fulltext/523/image_000004_22a8fcf53433bda3345e9231d416f64cfbab2895d11e74d66b0ee98d92ab35c1.png)

## 7.1 Performance Prediction

First, we assess the learned task embeddings by using them to predict an agent's performance on a task stest ∈ Sinit after observing its performance on a quiz Squiz ⊆ Sinit. This requires estimating its probability of success on a task, which is precisely its value function when the reward is binary and sparse. Specifically, we seek to answer the following research question: Would an agent show similar performance on tasks that are in close proximity in the learned embedding space? We begin by creating a benchmark for this application scenario, and then compare our method against various baselines.

Benchmark. Given the realizations of the optimality variables of a set of tasks for an agent θ, we are interested in predicting the most probable realization of the optimality variable of a new task for this agent without observing θ. We create benchmarks for this scenario by generating datasets for quiz sizes ranging from 1 to 20, with 5000 examples for both training and testing. Each example is generated by randomly sampling a quiz Squiz of desired size, along with a task stest from Sinit, and then recording the performance of a randomly sampled agent θ on these tasks. Performance prediction techniques are evaluated by partitioning each dataset into 10 folds, and measuring the mean prediction accuracy across the folds, along with the standard error.

Our approach. We perform soft-nearest neighbor matching of stest with Squiz in the task embedding space to predict performance on stest. Concretely, given the embedding function fϕ(.), the prediction is 1c&gt;0.5, where c equals

∑s∈Squiz os exp(-β∥fϕ(s)-fϕ(stest)∥22 )

∑s∈Squiz exp(-β∥fϕ(s)-fϕ(stest)∥22 ) , os is the realization of

Figure 6: Task selection results using embeddings (dark bars represent Top-3 accuracy and light bars represent Top-1). Our method (listed as Ours) is competitive witĥ OPT50 . Further, it outperforms OurswoNorm on Type-2 queries, highlighting CNORM's significance in our framework.

![Image](https://ijcai.org/proceedings/2024/fulltext/523/image_000005_09a2366b5921d78ff0ac8baa9a4367f9555f600e9429bc36b953c8edea8b69ac.png)

the optimality variable for task s, and β is a hyperparameter.

Baselines. Besides PredModel, we compare against different levels of oracle knowledge: (i) Random: Randomly predicts the agent's performance. (ii) IgnoreTask: Predicts the agent to succeed on stest iff the probability that it succeeds on a random task exceeds 0.5. (iii) IgnoreAgent: Predicts success on stest iff the probability that a random agent succeeds on it exceeds 0.5. (iv) OPT: Predicts the agent to succeed on stest iff the probability that it succeeds on stest exceeds 0.5.

Results. Fig. 5 shows the prediction accuracies of various methods. Our method is competitive with OPT, which is an upper-bound on the prediction accuracy but relies on the unrealistic assumption of full observability of both the agent and task. Additionally, our proposed task similarity criterion, which is used learn the embeddings, captures performance correlations that remain valid for new agents (see the longer version of the paper [Mahajan et al., 2023]).

## 7.2 Task Selection

Next, we assess the learned embeddings by using them to select tasks with desired characteristics. Specifically, we seek to answer the following research questions: (i) Does the inner product in the learned embedding space capture task similarity according to our information-theoretic criterion? (ii) Does the norm of the embedding learned by our framework induce an ordering on the tasks based on their difficulties? We begin by creating a benchmark for this scenario, and then compare our method based on task embeddings against various baselines.

Benchmark. Amongst several options of tasks Soptions, we are interested in choosing the task that best matches the desired characteristics, which we categorize into two query types: Type-1 : Select the task that is the most similar to a given reference task sref . The ground-truth answer to this query is argmaxs∈Soptions I(Osref ,Θ;Os,Θ). Type-2: Select the task that is the most similar to (but harder than) a given reference task sref . Out of all tasks in Soptions that are harder than sref , the ground-truth answer to this query is the task most similar to it. To create benchmarks for this scenario, we generate a dataset of 50 examples. Each example consists of a randomly sampled sref and 10 tasks that form Soptions. The benchmarks also include 5 easy tasks for reference (determined by ranking a randomly sampled pool of 500 tasks). We evaluate task selection methods by reporting mean selection accuracy across 4 randomly sampled datasets, along with the standard error.

Our approach. We use task embeddings to rank the options according to similarity and/or difficulty, based on which the selection is made. We additionally compare our method based on embeddings learned without CNORM (listed as OurswoNorm).

Baselines. Besides PredModel, we compare against the following baselines: (i) Random: Randomly selects answers from Soptions. (ii) StateSim: Measures task similarity based on state representation distances. For queries of type 2, it considers a task s1 to be harder than s2 iff the similarity between s1 and the task most similar to it in the set of easy reference tasks, is less than that for s2. (iii) TrajectorySim: Measures task similarity using the edit distance between expert trajectories. (iv) OPT: Estimates task similarity and difficulty using the entire agent population. Given the variance in the estimation process, this is the best one could do on this benchmark. (v) ̂ OPT50 : Estimates task similarity and difficulty using a randomly sampled 50% of the population.

Results. Fig. 6 compares different methods on the task selection benchmark. Our method outperforms Random, StateSim, TrajectorySim, and PredModel, and is competitive with ̂ OPT50 . This suggests that the inner product in the learned task embedding space successfully captures similarity between tasks. Notably, our technique significantly outperforms OurswoNorm on Type-2 queries, indicating that the norm of the embedding effectively orders tasks by difficulty.

## 8 Concluding Discussion

In this work, we introduced an information-theoretic framework for learning task embeddings in sequential decisionmaking settings. Through experiments on diverse environments, we empirically demonstrated that the inner product in the embedding space captures task similarity, and the norm of the embedding induces an ordering on the tasks based on their difficulties. A limitation of our current framework is the requirement for tasks to be goal-based, which we plan to address in future work. This could involve using the difference between the cumulative reward obtained during the rollout and the maximum achievable cumulative reward for the given task to parameterize the optimality variable. Additionally, it would be interesting to explore more principled methods for constructing the agent population that explicitly optimize for diversity. Further, empirically estimating the proposed similarity criterion by directly estimating the underlying mass functions could be sample-inefficient for some environments. Thus, a promising direction is to construct sample-efficient estimators for it. Moreover, evaluation in multi-agent settings, where the task embedding could encode the behavior of non-ego agents, is another interesting direction.

## Acknowledgments

Funded/Co-funded by the European Union (ERC, TOPS, 101039090). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council. Neither the European Union nor the granting authority can be held responsible for them.

## References

1. [Achille et al., 2019] Alessandro Achille, Michael Lam, Rahul Tewari, Avinash Ravichandran, Subhransu Maji, Charless C. Fowlkes, Stefano Soatto, and Pietro Perona. Task2Vec: Task Embedding for Meta-Learning. In ICCV, 2019.
2. [Arnekvist et al., 2018] Isac Arnekvist, Danica Kragic, and Johannes Andreas Stork. VPE: Variational Policy Embed- ding for Transfer Reinforcement Learning. In ICRA, 2018.
3. [Athar et al., 2020] Ali Athar, Sabarinath Mahadevan, Aljoša Ošep, Laura Leal-Taixé, and Bastian Leibe. STEm-Seg: Spatio-temporal Embeddings for Instance Segmentation in Videos. In ECCV, 2020.
4. [Bain and Sammut, 1995] Michael Bain and Claude Sammut. A Framework for Behavioural Cloning. In Machine Intelli- gence, Intelligent Agents, 1995.
5. [Bing et al., 2023] Z. Bing, D. Lerch, K. Huang, and A. Knoll. Meta-Reinforcement Learning in Non-Stationary and Dy- namic Environments. In IEEE TPAMI, 2023.
6. [Brockman et al., 2016] Greg Brockman, Vicki Cheung, Lud- wig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. OpenAI Gym. CoRR, abs/1606.01540, 2016.
7. [Bunel et al., 2018] Rudy Bunel, Matthew J. Hausknecht, Ja- cob Devlin, Rishabh Singh, and Pushmeet Kohli. Lever- aging Grammar and Reinforcement Learning for Neural Program Synthesis. In ICLR, 2018.
8. [Cer et al., 2018] Daniel Matthew Cer, Yinfei Yang, Sheng yi Kong, Nan Hua, Nicole Limtiaco, Rhomni St. John, Noah Constant, Mario Guajardo-Cespedes, Steve Yuan, Chris Tar, Brian Strope, and Ray Kurzweil. Universal Sentence Encoder for English. In EMNLP, 2018.
9. [Devidze et al., 2021] Rati Devidze, Goran Radanovic, Parameswaran Kamalaruban, and Adish Singla. Expli- cable Reward Design for Reinforcement Learning Agents. In NeurIPS, 2021.
10. [Eimer et al., 2021] Theresa Eimer, André Biedenkapp, Frank Hutter, and Marius Lindauer. Self-Paced Context Evaluation for Contextual Reinforcement Learning. In ICML, 2021.
11. [Fu et al., 2020] Haotian Fu, Hongyao Tang, Jianye Hao, Chen Chen, Xidong Feng, Dong Li, and Wulong Liu. To- wards Effective Context for Meta-Reinforcement Learning: an Approach based on Contrastive Learning. In AAAI, 2020.
12. [Furuta et al., 2021] Hiroki Furuta, Tatsuya Matsushima, Tadashi Kozuno, Yutaka Matsuo, Sergey Levine, Ofir Nachum, and Shixiang Shane Gu. Policy Information Ca- pacity: Information-Theoretic Measure for Task Complex- ity in Deep Reinforcement Learning. In ICML, 2021.
13. [Ghosh et al., 2022] Ahana Ghosh, Sebastian Tschiatschek, Sam Devlin, and Adish Singla. Adaptive Scaffolding in Block-based programming via Synthesizing New Tasks as Pop quizzes. In AIED, 2022.
14. [Gupta et al., 2018] Abhishek Gupta, Russell Mendonca, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- Reinforcement Learning of Structured Exploration Strate- gies. In NIPS, 2018.
15. [Hallak et al., 2015] Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual Markov Decision Processes. CoRR, abs/1502.02259, 2015.
16. [He-Yueya and Singla, 2021] Joy He-Yueya and Adish Ku- mar Singla. Quizzing Policy Using Reinforcement Learn- ing for Inferring the Student Knowledge State. In EDM, 2021.
17. [Jaderberg and et al., 2019] Max Jaderberg and et al. Human- level Performance in 3D Multiplayer Games with Population-based Reinforcement Learning. In Science, 2019.
18. [Klink et al., 2020a] P. Klink, C. D'Eramo, J. Peters, and J. Pajarinen. Self-Paced Deep Reinforcement Learning. In NeurIPS, 2020.
19. [Klink et al., 2020b] Pascal Klink, Carlo D' Eramo, Jan R Peters, and Joni Pajarinen. Self-Paced Deep Reinforcement Learning. In NeurIPS, 2020.
20. [Lan et al., 2019] Lin Lan, Zhenguo Li, Xiaohong Guan, and Pinghui Wang. Meta Reinforcement Learning with Task Embedding and Shared Policy. In IJCAI, 2019.
21. [Li et al., 2021] Lanqing Li, Rui Yang, and Dijun Luo. FO- CAL: Efficient Fully-Offline Meta-Reinforcement Learning via Distance Metric Learning and Behavior Regularization. In ICLR, 2021.
22. [Luce, 1959] R. Duncan Luce. Individual Choice Behavior: A Theoretical analysis. Wiley, 1959.
23. [Mahajan et al., 2023] Mridul Mahajan, Georgios Tzannetos, Goran Radanovic, and Adish Singla. Learning Embeddings for Sequential Tasks Using Population of Agents. CoRR, abs/2306.03311, 2023.
24. [Mikolov et al., 2013] Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed Repre- sentations of Words and Phrases and their Compositionality. In NIPS, 2013.
25. [Parker-Holder et al., 2020] Jack Parker-Holder, Aldo Pac- chiano, Krzysztof M Choromanski, and Stephen J Roberts. Effective Diversity in Population Based Reinforcement Learning. In NeurIPS, 2020.
26. [Peng et al., 2020] Xingchao Peng, Yichen Li, and Kate Saenko. Domain2Vec: Domain Embedding for Unsuper- vised Domain Adaptation. In ECCV, 2020.
27. [Pennington et al., 2014] Jeffrey Pennington, Richard Socher, and Christopher D. Manning. GloVe: Global Vectors for Word Representation. In EMNLP, 2014.
28. [Qin et al., 2022] Rongjun Qin, F. Chen, Tonghan Wang, Lei Yuan, Xiaoran Wu, Zongzhang Zhang, Chongjie Zhang, and Yang Yu. Multi-Agent Policy Transfer via Task Rela- tionship Modeling. CoRR, abs/2203.04482, 2022.
29. [Rakelly et al., 2019] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient Off- Policy Meta-Reinforcement Learning via Probabilistic Con- text Variables. In ICML, 2019.
30. [Schäfer et al., 2022] Lukas Schäfer, Filippos Christianos, Amos J. Storkey, and Stefano V. Albrecht. Learning Task Embeddings for Teamwork Adaptation in Multi-Agent Re- inforcement Learning. CoRR, abs/2207.02249, 2022.
31. [Sodhani et al., 2021a] Shagun Sodhani, Ludovic Denoyer, Pierre-Alexandre Kamienny, and Olivier Delalleau. MTEnv -Environment interface for mulit-task reinforcement learn- ing. GitHub, 2021.
32. [Sodhani et al., 2021b] Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-Task Reinforcement Learning with Context-based Representations. In ICML, 2021.
33. [Sun et al., 2014] Yi Sun, Xiaogang Wang, and Xiaoou Tang. Deep Learning Face Representation from Predicting 10,000 Classes. In CVPR, 2014.
34. [Sung et al., 2018] Flood Sung, Yongxin Yang, Li Zhang, Tao Xiang, Philip H. S. Torr, and Timothy M. Hospedales. Learning to Compare: Relation Network for Few-Shot Learning. In CVPR, 2018.
35. [Sutton et al., 1999] Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy Gradient Methods for Reinforcement Learning with Function Ap- proximation. In NIPS, 1999.
36. [Tylkin et al., 2021] Paul Tylkin, Goran Radanovic, and David C. Parkes. Learning Robust Helpful Behaviors in Two-Player Cooperative Atari Environments. In AAMAS, 2021.
37. [Tzannetos et al., 2023] Georgios Tzannetos, Bárbara Gomes Ribeiro, Parameswaran Kamalaruban, and Adish Singla. Proximal Curriculum for Reinforcement Learning Agents. In TMLR, 2023.
38. [van der Maaten and Hinton, 2008] Laurens van der Maaten and Geoffrey Hinton. Visualizing Data using t-SNE. In JMLR, 2008.
39. [Vinyals and et al., 2019] Oriol Vinyals and et al. Grandmas- ter Level in StarCraft II Using Multi-agent Reinforcement Learning. In Nature, 2019.
40. [Vuorio et al., 2019] Risto Vuorio, Shao-Hua Sun, Hexiang Hu, and Joseph J. Lim. Multimodal Model-Agnostic Meta- Learning via Task-Aware Modulation. In NeurIPS, 2019.
41. [Walke et al., 2022] Homer Walke, Jonathan Yang, Albert Yu, Aviral Kumar, Jedrzej Orbik, Avi Singh, and Sergey Levine. Don't Start From Scratch: Leveraging Prior Data to Auto- mate Robotic Reinforcement Learning. In CoRL, 2022. [Yengera et al., 2021]
42. Gaurav Yengera, Rati Devidze, Parameswaran Kamalaruban, and Adish Singla. Curricu- lum Design for Teaching via Demonstrations: Theory and Applications. In NeurIPS, 2021.
43. [Yoo et al., 2022] Minjong Yoo, Sangwoo Cho, and Honguk Woo. Skills Regularized Task Decomposition for Multi- task Offline Reinforcement Learning. In NeurIPS, 2022.
44. [Zhang et al., 2018] Yu Zhang, Ying Wei, and Qiang Yang. Learning to Multitask. In NeurIPS, 2018.
45. [Zhang et al., 2021] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learn- ing Invariant Representations for Reinforcement Learning without Reconstruction. In ICLR, 2021.