Proper multi-gpu support with PPO (#178)

* Add multi-gpu example

* fix pre-commit

* Add documentation and benchmark

* Update documentation

* Quick fix

* Also record world size in the params

* remove trailing space

* revert changes

* Update test cases

* Update test cases

* Fix CI

* Fix tests

* Fix pre-commit

* Fix tests

* Add a note that multi gpu only supported in linux

Author: vwxyzjn

.github/workflows/tests.yaml

@@ -56,10 +56,8 @@ jobs:
           poetry-version: ${{ matrix.poetry-version }}
 
       # atari tests
-      - name: Install core dependencies
-        run: poetry install -E pytest
       - name: Install atari dependencies
-        run: poetry install -E atari
+        run: poetry install -E "pytest atari"
       - name: Downgrade setuptools
         run: poetry run pip install setuptools==59.5.0
       - name: Run atari tests
@@ -156,3 +154,55 @@ jobs:
             libosmesa6-dev patchelf
       - name: Run mujoco tests
         run: poetry run pytest tests/test_mujoco.py
+
+  test-envpool-envs:
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [3.8]
+        poetry-version: [1.1.11]
+        os: [ubuntu-18.04]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Run image
+        uses: abatilo/actions-poetry@v2.0.0
+        with:
+          poetry-version: ${{ matrix.poetry-version }}
+
+      # envpool tests
+      - name: Install envpool dependencies
+        run: poetry install -E "pytest envpool"
+      - name: Downgrade setuptools
+        run: poetry run pip install setuptools==59.5.0
+      - name: Run envpool tests
+        run: poetry run pytest tests/test_envpool.py
+
+  test-atari-multigpu-envs:
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [3.8]
+        poetry-version: [1.1.11]
+        os: [ubuntu-18.04]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Run image
+        uses: abatilo/actions-poetry@v2.0.0
+        with:
+          poetry-version: ${{ matrix.poetry-version }}
+
+      # atari multigpu tests
+      - name: Install atari dependencies
+        run: poetry install -E "pytest atari"
+      - name: Downgrade setuptools
+        run: poetry run pip install setuptools==59.5.0
+      - name: Run atari tests
+        run: poetry run pytest tests/test_atari_multigpu.py

benchmark/ppo.sh

@@ -41,4 +41,11 @@ xvfb-run -a python -m cleanrl_utils.benchmark \
     --env-ids starpilot bossfight bigfish \
     --command "poetry run python cleanrl/ppo_procgen.py --track --capture-video" \
     --num-seeds 3 \
-    --workers 1
+    --workers 1
+
+poetry install -E atari
+xvfb-run -a python -m cleanrl_utils.benchmark \
+    --env-ids PongNoFrameskip-v4 BeamRiderNoFrameskip-v4 BreakoutNoFrameskip-v4 \
+    --command "poetry run torchrun --standalone --nnodes=1 --nproc_per_node=2 cleanrl/ppo_atari_multigpu.py --track --capture-video" \
+    --num-seeds 3 \
+    --workers 1