# Spatiotemporal Attention-Augmented Inverse Reinforcement Learning for Multi-Agent Task Allocation

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Inverse reinforcement learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 11 of its 38 references that are in the map.
- First posted on arXiv 7 Apr 2025.

## Placed in

- Offline RL and imitation learning from data / [Inverse reinforcement learning](https://mlsanity.com/offlinerl/inverse-reinforcement-learning/)

## Builds on (most cited first, 11 of 11)

- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [The Surprising Effectiveness of PPO in Cooperative, Multi-Agent Games](https://mlsanity.com/paper/2103.01955/) (2021)
- [Learning Robust Rewards with Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/1710.11248/) (2017)
- [A Survey of Inverse Reinforcement Learning: Challenges, Methods and Progress](https://mlsanity.com/paper/1806.06877/) (2018)
- [Query-Dependent Prompt Evaluation and Optimization with Offline Inverse RL](https://mlsanity.com/paper/2309.06553/) (2023)
- [Learning Scalable Policies over Graphs for Multi-Robot Task Allocation using Capsule Attention Networks](https://mlsanity.com/paper/2205.03321/) (2022)
- [A Local Information Aggregation based Multi-Agent Reinforcement Learning for Robot Swarm Dynamic Task Allocation](https://mlsanity.com/paper/2411.19526/) (2024)
- [Modelling resource contention in multi-robot task allocation problems with uncertain timing](https://mlsanity.com/paper/1607.04358/) (2016)
- [Discovering Individual Rewards in Collective Behavior through Inverse Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2305.10548/) (2023)
- [Joint Path planning and Power Allocation of a Cellular-Connected UAV using Apprenticeship Learning via Deep Inverse Reinforcement Learning](https://mlsanity.com/paper/2306.10071/) (2023)
- [Inverse Reinforcement Learning from Non-Stationary Learning Agents](https://mlsanity.com/paper/2410.14135/) (2024)

## Abstract (arXiv)

Adversarial inverse reinforcement learning (IRL) for multi-agent task allocation (MATA) is challenged by non-stationary interactions and high-dimensional coordination. Unconstrained reward inference in these settings often leads to high variance and poor generalization. We propose an attention-structured adversarial IRL framework that constrains reward inference via spatiotemporal representation learning. Our method employs multi-head self-attention (MHSA) for long-range temporal dependencies and graph attention networks (GAT) for agent-task relational structures. We formulate reward inference as a low-capacity, adaptive linear transformation of the environment reward, ensuring stable and interpretable guidance. This framework decouples reward inference from policy learning and optimizes the reward model adversarially. Experiments on benchmark MATA scenarios show that our approach outperforms representative MARL baselines in convergence speed, cumulative rewards, and spatial efficiency. Results demonstrate that attention-guided, capacity-constrained reward inference is a scalable and effective mechanism for stabilizing adversarial IRL in complex multi-agent systems.