From b0ce86fc591559843966ae9265e19b8a6d87f3c2 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <55539777+matteobettini@users.noreply.github.com>
Date: Thu, 25 Apr 2024 12:56:42 +0100
Subject: [PATCH] [CI] Fix CI  (#79)

* amend

* [Quality] Automatically get version

* amend

* amend
---
 .github/unittest/install_dependencies.sh         |  2 +-
 .github/unittest/install_dependencies_nightly.sh |  2 +-
 docs/source/usage/installation.rst               | 15 +++++++++++++++
 setup.py                                         |  1 +
 test/test_models.py                              |  2 +-
 5 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/.github/unittest/install_dependencies.sh b/.github/unittest/install_dependencies.sh
index 85e80372..aeef298d 100644
--- a/.github/unittest/install_dependencies.sh
+++ b/.github/unittest/install_dependencies.sh
@@ -1,7 +1,7 @@
 
 
 python -m pip install --upgrade pip
-python -m pip install flake8 pytest pytest-cov hydra-core tqdm torch torch_geometric
+python -m pip install flake8 pytest pytest-cov hydra-core tqdm torch torch_geometric torchvision av
 
 if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
 
diff --git a/.github/unittest/install_dependencies_nightly.sh b/.github/unittest/install_dependencies_nightly.sh
index 6f90c4f9..bd71bb8d 100644
--- a/.github/unittest/install_dependencies_nightly.sh
+++ b/.github/unittest/install_dependencies_nightly.sh
@@ -5,7 +5,7 @@ python -m pip install flake8 pytest pytest-cov hydra-core tqdm torch_geometric
 
 if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
 
-python -m pip install torch
+python -m pip install torch torchvision av
 # Not using nightly torch
 # python -m pip install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu --force-reinstall
 
diff --git a/docs/source/usage/installation.rst b/docs/source/usage/installation.rst
index 19c55a27..7499ec60 100644
--- a/docs/source/usage/installation.rst
+++ b/docs/source/usage/installation.rst
@@ -30,6 +30,21 @@ Or also clone it locally to access the configs and scripts
     git clone https://github.com/facebookresearch/BenchMARL.git
     pip install -e BenchMARL
 
+Install optional packages
+-------------------------
+
+By default, BenchMARL has only the core requirements.
+Here are some optional packages you may want to install.
+
+Logging
+^^^^^^^
+
+You may want to install the following rendering and logging tools
+
+.. code-block:: console
+
+   pip install wandb moviepy torchvision av
+
 Install environments
 --------------------
 
diff --git a/setup.py b/setup.py
index 1dcda9a1..a218dbe2 100644
--- a/setup.py
+++ b/setup.py
@@ -52,6 +52,7 @@ def get_version():
         "pettingzoo": ["pettingzoo[all]>=1.24.3"],
         "meltingpot": ["dm-meltingpot"],
         "gnn": ["torch_geometric"],
+        "logging": ["moviepy", "wandb", "torchvision"],
     },
     packages=find_packages(),
     include_package_data=True,
diff --git a/test/test_models.py b/test/test_models.py
index 4cff2d59..8867b05a 100644
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -170,6 +170,6 @@ def test_models_forward_shape(
         agent_group="agents",
         action_spec=None,
     )
-    input_td = input_spec.rand()
+    input_td = input_spec.expand(batch_size).rand()
     out_td = model(input_td)
     assert output_spec.expand(batch_size).is_in(out_td)