Since this repository is a clone of the original Mava, please refer to the official Mava repository at https://github.com/instadeepai/Mava for the original source and credits. Mava allows researchers to experiment with multi-agent reinforcement learning (MARL) at lightning speed. The single-file JAX implementations are built for rapid research iteration, modify, and test new ideas fast. Created for researchers, by The Research Team at InstaDeep.
Based on the paper: A-MARL: Agile Multi-Agent Reinforcement Learning for Soft Real-Time Task Scheduling in Edge Computing (IEEE CASCON 2025).
Presentation Slides: For a visual overview of EPD please refer to the Presentation Slides.
Modern Soft Real-Time Applications (SRTAs) impose heavy computational demands on embedded devices. Offloading workloads to Edge Computing (EC) resources is an attractive solution, but task scheduling remains challenging due to:
- Strict timing constraints
- Vast combinatorial search spaces
- Multiple conflicting optimization objectives
- Highly dynamic and unpredictable environments
Conventional heuristic and meta-heuristic approaches struggle to adapt. Single-agent Reinforcement Learning (RL) converges too slowly on medium- and large-scale problems due to enormous action spaces.
A-MARL solves this by enhancing Multi-Agent PPO with entropy-guided rule-based exploration, using the Shortest Processing Time (SPT) algorithm to steer agents toward promising regions of the action space during early training, and seamlessly transitioning to a learned policy as confidence grows.
- Entropy-Guided Exploration: Adaptively switches between SPT-based informed exploration and policy-based action selection depending on current policy entropy.
- Pseudo-Action Masking: Prunes the effective action space during high-entropy phases, focusing exploration on domain-informed (SPT) choices.
- Tailored Actor-Critic Architecture: Fuses a residual MLP for feature embedding with a lightweight self-attention block to capture cross-dependencies between operations and edge resources.
- A-MARL converges 61% faster compared to standard MARL baseline.
- A-MARL achieves 57% better sample efficiency, which means fewer environment interactions required to reach comparable solution quality.
- A-MARL reduces convergence-time by 846.3 seconds.
A-MARL
├── Environment: JobShop (discrete action space)
│ └── Each action = (Task ID, Edge-Server ID)
│
├── Agents: One RL-agent per Edge-Server
│ └── Independent policy learning in a shared environment
│
├── Actor-Critic Network
│ ├── Residual MLP ──── embeds SRTA & edge-server features
│ ├── Self-Attention ─── captures cross-operation/resource dependencies
│ ├── Actor Head ──────── outputs action distribution
│ └── Critic Head ─────── estimates value function
│
└── Entropy-Guided Mechanism
├── High Entropy → SPT-based informed exploration + pseudo-action masking
└── Low Entropy → Learned policy exploitation
The switching condition is defined as in Eq.(14) and below:
where
During high-entropy phases (early training), the policy is immature and action probabilities are nearly uniform. A-MARL applies entropy-conditioned pruning: the action space is constrained to only SPT-guided actions, dramatically reducing exploration overhead and focusing agents on domain-informed choices.
When entropy surpasses
As training progresses and entropy decreases, A-MARL seamlessly transitions to full policy-based exploitation; combining the benefits of SPT-guided initialization and learned optimization.
At the moment Mava is not meant to be installed as a library, but rather to be used as a research tool. Mava developers recommend cloning the Mava repo and installing dependencies using uv as follows:
# Clone the repository
git clone https://github.com/instadeepai/Mava.git
cd Mava
# Create a virtual environment and install all dependencies
uv sync
# Activate the virtual environment
source .venv/bin/activateAlternatively with pip, create a virtual environment and then:
pip install -e .Mava developers have tested Mava on Python 3.11 and 3.12, but earlier versions may also work. Specifically, they use Python 3.10 for the Quickstart notebook on Google Colab since Colab uses Python 3.10 by default. Note that because the installation of JAX differs depending on your hardware accelerator,
we advise users to explicitly install the correct JAX version (see the official installation guide). For more in-depth installation guides including Docker builds and virtual environments, please see their detailed installation guide.
To run the code, start training, and view the results of A-MARL and MARL in the JobShop environment, please execute the following commands:
# for running MARL (IPPO)
python mava/systems/ppo/anakin/ff_ippo.py env=job_shop
# for running A-MARL (IPPO+SPT)
python mava/systems/ppo/anakin/informed_ff_ippo.py env=job_shopIn order to see the default system configs please see the mava/configs/ directory.
A benefit of Hydra is that configs can either be set in config yaml files or overwritten from the terminal on the fly.
Additionally, Mava developers also have a [Quickstart notebook][quickstart] that can be used to quickly create and train your first multi-agent system.
If you find this work useful in your research, please cite:
@INPROCEEDINGS{11344342,
author={Avan, Amin and Azim, Akramul and Mahmoud, Qusay H.},
booktitle={2025 IEEE International Conference on Collaborative Advances in Software and COmputiNg (CASCON)},
title={A-MARL: Agile Multi-Agent Reinforcement Learning for Soft Real-Time Task Scheduling in Edge Computing},
year={2025},
volume={},
number={},
pages={275-284},
keywords={Schedules;Processor scheduling;Reinforcement learning;Dynamic scheduling;Search problems;Real-time systems;Software;Timing;Space exploration;Edge computing;Reinforcement Learning;Multi-Agent Reinforcement Learning;Edge Computing;Task Scheduling;Real-Time Systems;Real-Time Applications},
doi={10.1109/CASCON66301.2025.00053}}
Since this repository is a clone of the original Mava, please refer to the official Mava repository at https://github.com/instadeepai/Mava for the original source and credits.