This repository implements a multi-agent reinforcement learning (RL) system to identify MNIST images using the custom Gymnasium environment Multi-Agent MNIST Reconstruction. The RL agents work together to reconstruct and identify an MNIST image in as few steps as possible. A computer vision attention model is used to evaluate the reconstructed image and guide the training process.
The project aims to train multiple RL agents to collaboratively identify an MNIST image. The agents are placed on random parts of the image and have limited views. They navigate the image using actions (up, down, left, right), and their observations are combined to form a reconstructed image. A computer vision attention model is then used to evaluate how well the reconstructed image matches the original MNIST image. The RL system is trained to minimize the number of steps needed to accurately identify the image.
- Multi-Agent RL: Each agent learns to navigate the image optimally, contributing to the overall image reconstruction.
- Attention Model: A computer vision attention model evaluates the reconstructed image, providing feedback for the RL agents.
- Efficient Training: The system is designed to minimize the number of steps needed for accurate image identification.
- Python 3.8 or higher
- Gymnasium
- PyTorch (for the attention model)
- Stable-Baselines3 (for RL algorithms)
- NumPy
- OpenCV (optional, for visualization)
-
Clone the repository:
git clone https://github.com/varun646/MARL_MNIST.git
-
Install the required packages:
pip install -r requirements.txt
-
Install the custom environment:
pip install -e .