Merge branch 'main' of github.com:FLAIROx/JaxMARL

.github/workflows/tests.yml

@@ -0,0 +1,35 @@
+name: Tests
+on: [push, pull_request]
+
+jobs:
+  build:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: true
+      max-parallel: 15
+      matrix:
+        # os: [ubuntu-latest, macos-latest, windows-latest, macos-13-xlarge]
+        # For Apple Silicon: https://github.com/actions/runner-images/issues/8439
+        os: [ubuntu-latest, macos-latest, windows-latest]
+        python-version: ['3.9']
+    defaults:
+      run:
+        shell: bash
+    steps:
+    - name: Check out repository
+      uses: actions/checkout@v3
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+        cache: 'pip'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install pytest
+        pip install -e '.[dev]'
+
+    - name: Run pytest
+      run: pytest tests

README.md

@@ -65,6 +65,7 @@ We follow CleanRL's philosophy of providing single file implementations which ca
 | IQL | [Paper](https://arxiv.org/abs/1312.5602v1) | [Source](https://github.com/FLAIROx/JaxMARL/tree/main/baselines/QLearning) | 
 | VDN | [Paper](https://arxiv.org/abs/1706.05296)  | [Source](https://github.com/FLAIROx/JaxMARL/tree/main/baselines/QLearning) |
 | QMIX | [Paper](https://arxiv.org/abs/1803.11485) | [Source](https://github.com/FLAIROx/JaxMARL/tree/main/baselines/QLearning) |
+| SHAQ | [Paper](https://arxiv.org/abs/2105.15013) | [Source](https://github.com/FLAIROx/JaxMARL/tree/main/baselines/QLearning) |
 
 <h2 name="install" id="install">Installation 🧗 </h2>
 
@@ -151,3 +152,4 @@ JAX-native environments:
 - [Jumanji](https://github.com/instadeepai/jumanji): A diverse set of environments ranging from simple games to NP-hard combinatorial problems.
 - [Pgx](https://github.com/sotetsuk/pgx): JAX implementations of classic board games, such as Chess, Go and Shogi.
 - [Brax](https://github.com/google/brax): A fully differentiable physics engine written in JAX, features continuous control tasks.
+- [XLand-MiniGrid](https://github.com/corl-team/xland-minigrid): Meta-RL gridworld environments inspired by XLand and MiniGrid.

baselines/IPPO/config/ippo_rnn_smax.yaml

@@ -1,20 +1,20 @@
-"LR": 4e-3
-"NUM_ENVS": 64
+"LR": 0.004
+"NUM_ENVS": 128
 "NUM_STEPS": 128
-"TOTAL_TIMESTEPS": 5e5
-"UPDATE_EPOCHS": 2
-"NUM_MINIBATCHES": 2
+"TOTAL_TIMESTEPS": 1e7
+"UPDATE_EPOCHS": 4
+"NUM_MINIBATCHES": 4
 "GAMMA": 0.99
 "GAE_LAMBDA": 0.95
-"CLIP_EPS": 0.04
+"CLIP_EPS": 0.05
 "SCALE_CLIP_EPS": False
-"ENT_COEF": 0.00
+"ENT_COEF": 0.01
 "VF_COEF": 0.5
-"MAX_GRAD_NORM": 0.5
+"MAX_GRAD_NORM": 0.25
 "ACTIVATION": "relu"
 "ENV_NAME": "HeuristicEnemySMAX"
 "MAP_NAME": "2s3z"
-"SEED": 30
+"SEED": 0
 "ENV_KWARGS": 
   "see_enemy_actions": True
   "walls_cause_death": True
@@ -24,4 +24,4 @@
 # WandB Params
 "ENTITY": ""
 "PROJECT": "jaxmarl-smax"
-"WANDB_MODE" : "disabled"
+"WANDB_MODE" : "disabled"

baselines/MAPPO/config/mappo_homogenous_rnn_smax.yaml

@@ -14,7 +14,7 @@
 "ACTIVATION": "relu"
 "OBS_WITH_AGENT_ID": True
 "ENV_NAME": "HeuristicEnemySMAX"
-"MAP_NAME": "27m_vs_30m"
+"MAP_NAME": "2s3z"
 "SEED": 0
 "ENV_KWARGS": 
   "see_enemy_actions": True
@@ -25,4 +25,4 @@
 # WandB Params
 "WANDB_MODE": "disabled"
 "ENTITY": ""
-"PROJECT": "jaxmarl-smax"
+"PROJECT": "jaxmarl-smax"

baselines/QLearning/README.md

@@ -1,6 +1,13 @@
 # QLearning Baselines
 
-*Pure Jax implementation of **IQL** (Independent Q-Learners), **VDN** (Value Decomposition Network), and **QMix**. These implementations follow the original [Pymarl](https://github.com/oxwhirl/pymarl/blob/master/src/learners/q_learner.py) codebase.*
+
+Pure JAX implementations of:
+* IQL (Independent Q-Learners)
+* VDN (Value Decomposition Network)
+* QMIX
+* SHAQ (Incorporating Shapley Value Theory into Multi-Agent Q-Learning)
+
+The first three are follow the original [Pymarl](https://github.com/oxwhirl/pymarl/blob/master/src/learners/q_learner.py) codebase while SHAQ follows the [paper code](https://github.com/hsvgbkhgbv/shapley-q-learning)
 
 ```
 ⚠️ The implementations were tested with Python 3.9 and Jax 0.4.11. 
@@ -50,9 +57,9 @@ If you have cloned JaxMARL and you are in the repository root, you can run the a
 python baselines/QLearning/iql.py +alg=iql_mpe +env=mpe_speaker_listener
 # VDN with MPE spread
 python baselines/QLearning/vdn.py +alg=vdn_mpe +env=mpe_spread
-# QMix with SMAX
+# QMIX with SMAX
 python baselines/QLearning/qmix.py +alg=qmix_smax +env=smax
-# QMix against pretrained agents
+# QMIX against pretrained agents
 python baselines/QLearning/qmix_pretrained.py +alg=qmix_mpe +env=mpe_tag_pretrained
 ```
 

baselines/QLearning/config/alg/shaq_mpe.yaml

@@ -0,0 +1,30 @@
+"NUM_ENVS": 8
+"BUFFER_SIZE": 5000
+"BUFFER_BATCH_SIZE": 32
+"TOTAL_TIMESTEPS": 2050000
+"AGENT_HIDDEN_DIM": 64
+"AGENT_INIT_SCALE": 2.
+"PARAMETERS_SHARING": True
+"EPSILON_START": 1.0
+"EPSILON_FINISH": 0.05
+"EPSILON_ANNEAL_TIME": 100000
+"MIXER_EMBEDDING_DIM": 32
+"MIXER_HYPERNET_HIDDEN_DIM": 64
+"MIXER_INIT_SCALE": 0.00001
+"MAX_GRAD_NORM": 25
+"TARGET_UPDATE_INTERVAL": 200
+"LR": 0.005
+"LR_LINEAR_DECAY": True
+"EPS_ADAM": 0.001
+"WEIGHT_DECAY_ADAM": 0.00001
+"TD_LAMBDA_LOSS": True
+"TD_LAMBDA": 0.6
+"GAMMA": 0.9
+"VERBOSE": False
+"WANDB_ONLINE_REPORT": True
+"NUM_TEST_EPISODES": 32
+"TEST_INTERVAL": 50000
+"SAMPLE_SIZE": 5
+"MANUAL_ALPHA_ESTIMATES": null
+"LR_ALPHA": 0.001
+"ALG_NAME": shaq

baselines/QLearning/config/alg/shaq_smax.yaml

@@ -0,0 +1,30 @@
+"NUM_ENVS": 8
+"BUFFER_SIZE": 3000
+"BUFFER_BATCH_SIZE": 32
+"TOTAL_TIMESTEPS": 20000000
+"AGENT_HIDDEN_DIM": 256
+"AGENT_INIT_SCALE": 1.
+"PARAMETERS_SHARING": True
+"EPSILON_START": 1.0
+"EPSILON_FINISH": 0.05
+"EPSILON_ANNEAL_TIME": 100000
+"MIXER_EMBEDDING_DIM": 64
+"MIXER_HYPERNET_HIDDEN_DIM": 256
+"MIXER_INIT_SCALE": 0.001
+"MAX_GRAD_NORM": 10
+"TARGET_UPDATE_INTERVAL": 200
+"LR": 0.001
+"LR_LINEAR_DECAY": False
+"EPS_ADAM": 0.00001
+"WEIGHT_DECAY_ADAM": 0.000001
+"TD_LAMBDA_LOSS": False
+"TD_LAMBDA": 0.6
+"GAMMA": 0.99
+"VERBOSE": False
+"WANDB_ONLINE_REPORT": True
+"NUM_TEST_EPISODES": 32
+"TEST_INTERVAL": 100000
+"SAMPLE_SIZE": 1 # suggest choosing from 1/5/10
+"MANUAL_ALPHA_ESTIMATES": null
+"LR_ALPHA": 0.0005
+"ALG_NAME": shaq

baselines/QLearning/config/config.yaml

@@ -5,7 +5,7 @@
 # wandb params
 "ENTITY": ""
 "PROJECT": "jaxMARL"
-"WANDB_MODE": "online"
+"WANDB_MODE": "disabled"
 
 # where to save the params (if None, will not save)
 "SAVE_PATH": "baselines/QLearning/checkpoints"