Simultaneous installation of JAX and PyTorch agents without tensorboa…

…rd conflicts (#56)

* Instructions to install torch and jax agents at the same time, avoiding tensorboard conflicts.

* test.yml: including freeglut installation before tests

.github/workflows/test.yml

@@ -26,10 +26,13 @@ jobs:
           python-version: '3.7'
       - name: Install dependencies
         run: |
+          sudo apt-get install freeglut3-dev
           python -m pip install --upgrade pip
           pip install torch==1.7.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
-          pip install -e .[test]
+          pip install pytest
+          pip install pytest-cov
           pip install -e .[jax_agents]
+          pip install -e .[torch_agents]
       - name: Test with pytest
         run: |
           export NUMBA_DISABLE_JIT=1

README.md

@@ -103,11 +103,11 @@ year = {2021}
 
 ## Tests
 
-To run tests, install test dependencies with `pip install -e .[test]` and run `pytest`. 
-
-To check coverage, install test dependencies and run 
-
 ```bash 
+$ pip install pytest
+$ pip install pytest-cov
+$ pip install -e .[jax_agents]
+$ pip install -e .[torch_agents]
 $ cd scripts
 $ bash run_testscov.sh
 ```

docs/installation.rst

@@ -18,33 +18,35 @@ Then you have two options for the installation. For a stable version, you can ju
 
 .. code:: bash
 
-    $ pip install rlberry[full]
+    $ pip install rlberry[default]
 
-or, for a basic installation (without heavy libraries like PyTorch):
-
-.. code:: bash
-
-    $ pip install rlberry
 
 For more advanced users who want to try the development version, all you need to do is clone the rlberry_ repository and install:
 
 .. code:: bash
 
     $ git clone https://github.com/rlberry-py/rlberry.git
     $ cd rlberry
-    $ pip install -e .[full]
+    $ pip install -e .[default]
 
-or, for a basic installation:
 
-.. code:: bash
+Installation for Deep RL agents
+===============================
 
-    $ pip install -e .
+Deep RL agents require extra libraries, like PyTorch and JAX.
 
-Full installation includes, for instance:
+* PyTorch agents:
 
-*   `Numba <https://github.com/numba/numba>`_ for just-in-time compilation of algorithms based on dynamic programming
-*   `PyTorch <https://pytorch.org/>`_ for Deep RL agents
-*   `Optuna <https://optuna.org/#installation>`_ for hyperparameter optimization
-*   `ffmpeg-python <https://github.com/kkroening/ffmpeg-python>`_ for saving videos
-*   `PyOpenGL <https://pypi.org/project/PyOpenGL/>`_ for more rendering options
+.. code:: bash
+    $ pip install -e .[torch_agents]
+    $ pip install tensorboard   # only if you're not installing jax_agents too!
+* JAX agents:
+
+.. code:: bash
+    $ pip install -e .[jax_agents]
 
+.. warning::
+    If you're using PyTorch agents *and* JAX agents, do not install tensorboard separately,
+    since `pip install -e .[jax_agents]` installs tensorflow, which already contains
+    tensorboard. Otherwise, there might be a conflict between the two installations
+    and tensorboard will not work properly.

scripts/full_install.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# Install everything!
+
+pip install -e .[default]
+pip install -e .[jax_agents]
+pip install -e .[torch_agents]
+
+pip install pytest
+pip install pytest-cov
+conda install -c conda-forge jupyterlab

setup.py

@@ -2,6 +2,9 @@
 
 packages = find_packages(exclude=['docs', 'notebooks', 'assets'])
 
+#
+# Base installation (interface only)
+#
 install_requires = [
     'numpy>=1.17',
     'pygame',
@@ -14,45 +17,42 @@
     'pyyaml',
 ]
 
-tests_require = [
-    'pytest',
-    'pytest-cov',
-    'numpy>=1.17',
-    'numba',
-    'matplotlib',
-    'pandas',
-    'seaborn',
-    'optuna',
-    'pyvirtualdisplay',
-    'gym',
-]
+#
+# Extras
+#
 
-full_requires = [
+# default installation
+default_requires = [
     'numba',
-    'torch>=1.6.0',
-    'tensorboard',
     'optuna',
     'ffmpeg-python',
     'PyOpenGL',
     'PyOpenGL_accelerate',
     'pyvirtualdisplay',
 ]
 
+# tensorboard must be installed manually, due to conflicts with
+# dm-reverb-nightly[tensorflow] in jax_agents_requires
+torch_agents_requires = default_requires + [
+    'torch>=1.6.0',
+    # 'tensorboard'
+]
+
+jax_agents_requires = default_requires + [
+    'jax[cpu]',
+    'chex',
+    'dm-haiku',
+    'optax',
+    'dm-reverb-nightly[tensorflow]',
+    'dm-tree',
+    'rlax'
+]
+
 extras_require = {
-    'full': full_requires,
-    'test': tests_require,
-    'jax_agents': ['jax[cpu]',
-                   'chex',
-                   'dm-haiku',
-                   'optax',
-                   'dm-reverb-nightly[tensorflow]',
-                   'dm-tree',
-                   'rlax'],
+    'default': default_requires,
+    'jax_agents': jax_agents_requires,
+    'torch_agents': torch_agents_requires,
     'deploy': ['sphinx', 'sphinx_rtd_theme'],
-    'opengl_rendering': ['PyOpenGL', 'PyOpenGL_accelerate'],
-    'torch_agents': ['torch>=1.6.0', 'tensorboard'],
-    'hyperparam_optimization': ['optuna'],
-    'save_video': ['ffmpeg-python'],
 }
 
 with open("README.md", "r") as fh:
@@ -79,7 +79,6 @@
         "Operating System :: OS Independent",
     ],
     install_requires=install_requires,
-    tests_require=tests_require,
     extras_require=extras_require,
     zip_safe=False,
 )