Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning Embeddings for Sequential Tasks Using Population of Agents

Mridul Mahajan, Georgios Tzannetos, Goran Radanovic, Adish Singla

TL;DR

The paper introduces an information-theoretic framework to learn fixed-dimensional task embeddings for sequential decision tasks in reinforcement learning by leveraging a diverse population of agents. Task similarity is quantified via mutual information between task optimality events, and embeddings are learned through ordinal constraints that combine a triplet-based similarity term with a norm-based difficulty ordering under a Bradley-Terry-Luce model. The approach is validated across multiple environments, showing interpretable embedding spaces, superior clustering compared to baselines, and practical utility in two downstream tasks: predicting agent performance on new tasks and selecting tasks with desired characteristics. The results suggest that the inner product of embeddings accurately captures task similarity, while the embedding norm encodes relative difficulty, enabling scalable, one-shot reasoning about sequential tasks with potential for curriculum design and task selection in RL systems.

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.

Learning Embeddings for Sequential Tasks Using Population of Agents

TL;DR

The paper introduces an information-theoretic framework to learn fixed-dimensional task embeddings for sequential decision tasks in reinforcement learning by leveraging a diverse population of agents. Task similarity is quantified via mutual information between task optimality events, and embeddings are learned through ordinal constraints that combine a triplet-based similarity term with a norm-based difficulty ordering under a Bradley-Terry-Luce model. The approach is validated across multiple environments, showing interpretable embedding spaces, superior clustering compared to baselines, and practical utility in two downstream tasks: predicting agent performance on new tasks and selecting tasks with desired characteristics. The results suggest that the inner product of embeddings accurately captures task similarity, while the embedding norm encodes relative difficulty, enabling scalable, one-shot reasoning about sequential tasks with potential for curriculum design and task selection in RL systems.

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.
Paper Structure (29 sections, 6 equations, 11 figures, 1 table, 2 algorithms)

This paper contains 29 sections, 6 equations, 11 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Problem Setup
  4. Learning Framework
  5. Information-Theoretic Task Similarity Measure
  6. Learning Task Embeddings
  7. Experiments: Visualizing Embedding Spaces
  8. Environments
  9. Training Process
  10. Visualizations and Qualitative Results
  11. Experiments: Comparison with Prior Work
  12. Experiments: Application Scenarios
  13. Performance Prediction
  14. Task Selection
  15. Concluding Discussion
  16. ...and 14 more sections

Figures (11)

  • Figure 1: Schematics of our approach. We learn a task embedding function $f_{\phi}(.)$ that maps a task $s$ to its fixed-dimensional representation $\mathrm{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 $s_{2}$, the agent starts off with the yellow key in possession already. ${\langle \mathrm{E}_{1}, \mathrm{E}_{2} \rangle}$ is greater than ${\langle \mathrm{E}_{1}, \mathrm{E}_{3} \rangle}$, since tasks $s_{1}$ and $s_{2}$ have a common requirement of picking the blue key, and thus, are similar. Additionally, $\left\lVert\mathrm{E}_{2}\right\rVert_{2}$ is less than both $\left\lVert\mathrm{E}_{1}\right\rVert_{2}$ and $\left\lVert\mathrm{E}_{3}\right\rVert_{2}$, since task $s_{2}$ requires picking a single key, while tasks $s_{1}$ and $s_{3}$ require picking two keys, which makes them harder than $s_{2}$.
  • 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.
  • 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.
  • Figure 4: Task embedding spaces for the MultiKeyNav environment: (a) without $\mathcal{C}_{\mathrm{NORM}}$, (b) $\mathtt{pickKey}$ actions masked, (c) all doors require $\mathtt{KeyA}$, $\mathtt{KeyB}$, and (d) all doors require $\mathtt{KeyA}$.
  • 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.
  • ...and 6 more figures

Theorems & Definitions (2)

  • Definition 1: Performance Uncertainty
  • Definition 2: Task Similarity