This code accompanies the paper "Scalable Multi-Agent Model-Based Reinforcement Learning".

The repository contains MAMBA implementation as well as fine-tuned hyperparameters in configs/dreamer/optimal folder.

python3 train.py --n_workers 2 --env flatland --env_type 5_agents

Two environments are supported for env flag: flatland and starcraft.

The code for the environment can be found at https://github.com/oxwhirl/smac

The original code for the environment can be found at https://github.com/jbr-ai-labs/NeurIPS2020-Flatland-Competition-Solution

• agent contains implementation of MAMBA
  • controllers contains logic for inference
  • learners contains logic for learning the agent
  • memory contains buffer implementation
  • models contains architecture of MAMBA
  • optim contains logic for optimizing loss functions
  • runners contains logic for running multiple workers
  • utils contains helper functions
  • workers contains logic for interacting with environment
• env contains environment logic
• networks contains neural network architectures