Policy Search, Retrieval, and Composition via Task Similarity in Collaborative Agentic Systems

TL;DR

MOSAIC tackles the challenge of open-ended, agentic learning by enabling decentralized, selective sharing of policies among autonomous learners. It combines Wasserstein-based task embeddings for similarity, modular binary masks for transfer, and a reward-aware, similarity-driven policy integration to accelerate learning. Across ISL, MiniHack MultiRoom, and MiniGrid Crossing, MOSAIC consistently outperforms isolated and global-sharing baselines, achieving faster convergence and solving tasks that alone were intractable. The work demonstrates emergent curricula and improved interpretability, while acknowledging limitations related to simulation, potential interference, and communication bandwidth, and suggests pathways toward scalable lifelong and cross-domain collaboration.

Abstract

Agentic AI aims to create systems that set their own goals, adapt proactively to change, and refine behavior through continuous experience. Recent advances suggest that, when facing multiple and unforeseen tasks, agents could benefit from sharing machine-learned knowledge and reusing policies that have already been fully or partially learned by other agents. However, how to query, select, and retrieve policies from a pool of agents, and how to integrate such policies remains a largely unexplored area. This study explores how an agent decides what knowledge to select, from whom, and when and how to integrate it in its own policy in order to accelerate its own learning. The proposed algorithm, \emph{Modular Sharing and Composition in Collective Learning} (MOSAIC), improves learning in agentic collectives by combining (1) knowledge selection using performance signals and cosine similarity on Wasserstein task embeddings, (2) modular and transferable neural representations via masks, and (3) policy integration, composition and fine-tuning. MOSAIC outperforms isolated learners and global sharing approaches in both learning speed and overall performance, and in some cases solves tasks that isolated agents cannot. The results also demonstrate that selective, goal-driven reuse leads to less susceptibility to task interference. We also observe the emergence of self-organization, where agents solving simpler tasks accelerate the learning of harder ones through shared knowledge.

Figures (20)

• Figure 1: High-level illustration of the main MOSAIC algorithmic steps. (A) A Wasserstein task embedding $v_\tau$ is computed from a SAR batch by Agent A0 (representing any agent in the collective). (B) Periodically, the agent A0 broadcasts a task embedding query (TEQ) to all known peers. (C) Peers send back a query response (QR) that contains their task's $v_\tau, r$, and the corresponding mask ID; (D) Agent A0 selects relevant embeddings using cosine similarity on the Wasserstein embeddings. (E) A0 further selects relevant embeddings using \ref{['eq:criterion_align']} and \ref{['eq:criterion_perf']}. (F) A0 sends mask requests (MR). (G) The contacted agents respond by sending the requested mask through a mask transfer (MTR). (H) Incoming masks from A2 and A3 are incorporated into A0. The training of the agent's policy occurs in parallel (not represented in the figure).
• Figure 2: Performance of communicating MOSAIC agents versus isolated agents on the same tasks. (A) Image sequence learning, 28 tasks, five seeds/task: average of 140 runs with 95% confidence intervals colas2018many. (B) MiniHack Multiroom, 14 tasks, five seeds/task: average of 70 runs with 95% confidence intervals. MOSAIC agents achieve relative gains of 170.8% and 128.2% over the isolated baseline in the image sequence and MiniHack benchmarks, respectively.
• Figure 3: Comparison of MOSAIC to baseline approaches on the 14 task MiniGrid curricula made up of SimpleCrossing and LavaCrossing task variations: average performance for 70 runs with 95% confidence intervals.
• Figure 4: Pairwise cosine similarity and $\hat{\beta}$ statistics in the image sequence learning benchmark. (A) Cosine similarity matrix. (B) $\hat{\beta}$ values indicating policy use per task. (C) Cosine similarities clustered using WPGMA, shown as a dendrogram. (D) Cosine similarity matrix reordered by clustering. (E) Clustered $\hat{\beta}$ values show that similar tasks exhibit the highest policy reuse. Annotated heatmaps of individual clusters are in Appendix \ref{['sec:further_analysis_sim_beta']}, Figure \ref{['fig:zoom-grid-heatmaps']}.
• Figure 5: Performance grouped by task complexity, i.e., seven different levels, corresponding to different graph depths, in the image sequence learning problem. Average performance for 20 runs (four tasks and five seeds/task) with 95% confidence intervals. Communicating MOSAIC agents (top graph) are compared with isolated MOSAIC agents (bottom graph). Communicating agents solve tasks progressively, from the simplest to the hardest tasks. Isolated agents only manage to solve the two simplest tasks, and partially the third, but fail on the four most complex tasks.