Add minimum version numbers for jax packages, and add extras requirem…

…ents for dopamine and acme.

PiperOrigin-RevId: 427008871

README.md

@@ -20,7 +20,15 @@ Note: The BLE requires python >= 3.7
 The BLE can easily be installed with pip:
 
 ```
-pip install --upgrade pip && pip install balloon_learning_environment
+$ pip install --upgrade pip
+$ pip install balloon_learning_environment
+```
+
+To install with `Dopamine` and/or `ACME` packages:
+
+```
+$ pip install --upgrade pip
+$ pip install balloon_learning_environment[dopamine,acme]
 ```
 
 Once the package has been installed, you can test it runs correctly by
@@ -34,6 +42,14 @@ python -m balloon_learning_environment.eval.eval \
   --output_dir=/tmp/ble/eval
 ```
 
+To install from GitHub directly, run the following commands from the root
+directory where you cloned the repository:
+
+```
+$ pip install --upgrade pip
+$ pip install .[dopamine,acme]
+```
+
 ## Ensure the BLE is Using Your GPU/TPU
 
 The BLE contains a VAE for generating winds, which you will probably want

docs/getting_started.rst

@@ -21,6 +21,19 @@ and then install the balloon_learning_environment package:
 
    pip install balloon_learning_environment
 
+To install with Dopamine and/or ACME packages, run:
+
+.. code-block:: console
+
+   pip install balloon_learning_environment[dopamine,acme]
+
+To install from GitHub directly, run the following commands from the root
+directory where you cloned the repository:
+
+.. code-block:: console
+
+   pip install .[dopamine,acme]
+
 
 Ensure the BLE is Using Your GPU/TPU
 ####################################

setup.py

@@ -20,6 +20,18 @@
 current_directory = pathlib.Path(__file__).parent
 description = (current_directory / 'README.md').read_text()
 
+dopamine_requirements = [
+    'dopamine-rl >= 4.0.0',
+]
+
+acme_requirements = [
+    'dm-acme',
+    'dm-haiku',
+    'dm-reverb',
+    'dm-sonnet',
+    'rlax',
+]
+
 setuptools.setup(
     name='balloon_learning_environment',
     long_description=description,
@@ -28,11 +40,12 @@
     packages=setuptools.find_packages(),
     install_requires=[
         'absl-py',
-        'dopamine-rl >= 4.0.0',
         'flax',
         'gcsfs',
         'gin-config',
         'gym',
+        'jax >= 0.2.28",
+        'jaxlib >= 0.1.76",
         'opensimplex <= 0.3.0',
         's2sphere',
         'scikit-learn',
@@ -42,6 +55,10 @@
         'transitions',
         'zarr',
     ],
+    extras_require={
+        'dopamine': dopamine_requirements,
+        'acme': acme_requirements,
+    },
     package_data={
         '': ['*.msgpack', '*.pb', '*.gin'],
     },