# Trajectory-Class-Aware Multi-Agent Reinforcement Learning

arXiv: https://arxiv.org/abs/2503.01440 · HTML: https://mlsanity.com/paper/2503.01440/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 50 references that are in the map.
- First posted on arXiv 3 Mar 2025.

## Builds on (most cited first, 30 of 37)

- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments](https://mlsanity.com/paper/1706.02275/) (2017)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [QMIX: Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/1803.11485/) (2018)
- [The StarCraft Multi-Agent Challenge](https://mlsanity.com/paper/1902.04043/) (2019)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Value-Decomposition Networks For Cooperative Multi-Agent Learning](https://mlsanity.com/paper/1706.05296/) (2017)
- [QTRAN: Learning to Factorize with Transformation for Cooperative Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/1905.05408/) (2019)
- [QPLEX: Duplex Dueling Multi-Agent Q-Learning](https://mlsanity.com/paper/2008.01062/) (2020)
- [Weighted QMIX: Expanding Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2006.10800/) (2020)
- [Deep Decentralized Multi-task Multi-Agent Reinforcement Learning under Partial Observability](https://mlsanity.com/paper/1703.06182/) (2017)
- [ROMA: Multi-Agent Reinforcement Learning with Emergent Roles](https://mlsanity.com/paper/2003.08039/) (2020)
- [RODE: Learning Roles to Decompose Multi-Agent Tasks](https://mlsanity.com/paper/2010.01523/) (2020)
- [SMACv2: An Improved Benchmark for Cooperative Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2212.07489/) (2022)
- [Opponent Modeling in Deep Reinforcement Learning](https://mlsanity.com/paper/1609.05559/) (2016)
- [Scaling Multi-Agent Reinforcement Learning with Selective Parameter Sharing](https://mlsanity.com/paper/2102.07475/) (2021)
- [Discretizing Continuous Action Space for On-Policy Optimization](https://mlsanity.com/paper/1901.10500/) (2019)
- [Modeling Others using Oneself in Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/1802.09640/) (2018)
- [Agent Modelling under Partial Observability for Deep Reinforcement Learning](https://mlsanity.com/paper/2006.09447/) (2020)
- [Episodic Multi-agent Reinforcement Learning with Curiosity-Driven Exploration](https://mlsanity.com/paper/2111.11032/) (2021)
- [Learning Policy Representations in Multiagent Systems](https://mlsanity.com/paper/1806.06464/) (2018)
- [Hierarchical Cooperative Multi-Agent Reinforcement Learning with Skill Discovery](https://mlsanity.com/paper/1912.03558/) (2019)
- [UPDeT: Universal Multi-agent Reinforcement Learning via Policy Decoupling with Transformers](https://mlsanity.com/paper/2101.08001/) (2021)
- [MASER: Multi-Agent Reinforcement Learning with Subgoals Generated from Experience Replay Buffer](https://mlsanity.com/paper/2206.10607/) (2022)
- [Experiments with Random Projection](https://mlsanity.com/paper/1301.3849/) (2013)
- [Model-Based Opponent Modeling](https://mlsanity.com/paper/2108.01843/) (2021)

## Built on by (most cited first, 2 of 2)

- [Toward Adaptable Multi-Agent Reinforcement Learning: An Assumption-Aware Review](https://mlsanity.com/paper/2507.10142/) (2025)
- [LOD: Latent Objective Discovery in Heterogeneous Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2505.09756/) (2025)

## Abstract (arXiv)

In the context of multi-agent reinforcement learning, generalization is a challenge to solve various tasks that may require different joint policies or coordination without relying on policies specialized for each task. We refer to this type of problem as a multi-task, and we train agents to be versatile in this multi-task setting through a single training process. To address this challenge, we introduce TRajectory-class-Aware Multi-Agent reinforcement learning (TRAMA). In TRAMA, agents recognize a task type by identifying the class of trajectories they are experiencing through partial observations, and the agents use this trajectory awareness or prediction as additional information for action policy. To this end, we introduce three primary objectives in TRAMA: (a) constructing a quantized latent space to generate trajectory embeddings that reflect key similarities among them; (b) conducting trajectory clustering using these trajectory embeddings; and (c) building a trajectory-class-aware policy. Specifically for (c), we introduce a trajectory-class predictor that performs agent-wise predictions on the trajectory class; and we design a trajectory-class representation model for each trajectory class. Each agent takes actions based on this trajectory-class representation along with its partial observation for task-aware execution. The proposed method is evaluated on various tasks, including multi-task problems built upon StarCraft II. Empirical results show further performance improvements over state-of-the-art baselines.