[Paper Review] Towards Efficient Multi-Agent Learning Systems

Kailash Gogineni, Wei Peng+2 more|arXiv (Cornell University)|May 22, 2023

Data Stream Mining Techniques4 citations

TL;DR

This paper proposes a locality-aware neighbor sampling strategy to optimize the computationally intensive mini-batch sampling phase in Multi-Agent Reinforcement Learning (MARL), significantly improving cache efficiency. The method achieves 26.66–27.39% performance gains in sampling time and 10.2% end-to-end training time reduction on 12-agent MARL workloads without degrading mean reward, demonstrating a key systems-level optimization for scalable MARL training.

ABSTRACT

Multi-Agent Reinforcement Learning (MARL) is an increasingly important research field that can model and control multiple large-scale autonomous systems. Despite its achievements, existing multi-agent learning methods typically involve expensive computations in terms of training time and power arising from large observation-action space and a huge number of training steps. Therefore, a key challenge is understanding and characterizing the computationally intensive functions in several popular classes of MARL algorithms during their training phases. Our preliminary experiments reveal new insights into the key modules of MARL algorithms that limit the adoption of MARL in real-world systems. We explore neighbor sampling strategy to improve cache locality and observe performance improvement ranging from 26.66% (3 agents) to 27.39% (12 agents) during the computationally intensive mini-batch sampling phase. Additionally, we demonstrate that improving the locality leads to an end-to-end training time reduction of 10.2% (for 12 agents) compared to existing multi-agent algorithms without significant degradation in the mean reward.

Motivation & Objective

• To identify and analyze performance bottlenecks in Multi-Agent Reinforcement Learning (MARL) training from a systems perspective.
• To understand the computational intensity of key MARL training phases, especially the mini-batch sampling and update phases.
• To improve data access efficiency in MARL by optimizing memory locality during the sampling phase.
• To reduce end-to-end training time in MARL without sacrificing learning performance.
• To enable broader real-world deployment of MARL by addressing scalability issues tied to growing agent counts.

Proposed method

• Proposes a neighbor sampling strategy that reorders transitions in the experience replay buffer to improve spatial and temporal locality.
• Applies locality-aware indexing to group transitions from neighboring agents together during mini-batch sampling.
• Integrates the optimized sampling loop into existing MARL frameworks (MADDPG and MASAC) with minimal architectural changes.
• Uses hardware performance counters to profile and characterize computational bottlenecks in MARL training workloads on GPU architectures (e.g., A100, RTX 3090).
• Employs a centralized critic with joint observation-action space, while maintaining decentralized actors for action selection.
• Validates the optimization across 3 to 48 agents in a Predator-Prey environment, measuring both performance and learning stability.

Experimental results

Research questions

• RQ1 Which MARL training phases are the primary performance bottlenecks in terms of computation and memory access?
• RQ2 How does data access pattern in mini-batch sampling affect cache efficiency and training performance in MARL?
• RQ3 Can locality-aware neighbor sampling reduce training time without degrading learning quality in MARL?
• RQ4 How does the performance improvement scale with increasing numbers of agents in MARL systems?
• RQ5 What is the impact of neighbor sampling on end-to-end training time and model convergence in MARL?

Key findings

• The 'Update all trainers' phase dominates training time, accounting for 35% to 85% of total time as agent count increases from 3 to 48.
• The mini-batch sampling phase exhibits significant performance degradation due to poor cache locality, especially with large observation-action spaces.
• Neighbor sampling improves cache locality, reducing sampling phase runtime by 26.66% (3 agents) to 27.39% (12 agents).
• End-to-end training time is reduced by 10.2% for 12-agent MARL workloads using the proposed optimization.
• The optimization maintains high learning performance, with no significant degradation in mean episode reward across all tested agent counts.
• The performance gains are consistent across both MADDPG and MASAC frameworks, indicating generalizability across MARL algorithms.