1. Fix README

2. Separate out the core requirements to minimize package installation

PiperOrigin-RevId: 708825801

README.md

@@ -1,10 +1,14 @@
 # Iris: Synchronous and Distributed Blackbox Optimization at Scale
 
+## Overview
+Iris is a library for performing synchronous and distributed zeroth-order
+optimization at scale. It is meant primarily to train large neural networks with
+evolutionary methods, but can be applied to optimize any high dimensional
+blackbox function.
+
 ## Associated Publications
-* [SARA-RT: Scaling up Robotics Transformers with Self-Adaptive Robust Attention
-](https://arxiv.org/abs/2312.01990) (ICRA 2024 - Best Robotic Manipulation Award)
-* [Embodied AI with Two Arms: Zero-shot Learning, Safety and Modularity
-](https://arxiv.org/abs/2404.03570) (IROS 2024 - Robocup Best Paper Award)
+* [SARA-RT: Scaling up Robotics Transformers with Self-Adaptive Robust Attention](https://arxiv.org/abs/2312.01990) (ICRA 2024 - Best Robotic Manipulation Award)
+* [Embodied AI with Two Arms: Zero-shot Learning, Safety and Modularity](https://arxiv.org/abs/2404.03570) (IROS 2024 - Robocup Best Paper Award)
 * [Agile Catching with Whole-Body MPC and Blackbox Policy Learning](https://arxiv.org/abs/2306.08205) (L4DC 2023)
 * [Discovering Adaptable Symbolic Algorithms from Scratch](https://arxiv.org/abs/2307.16890) (IROS 2023, Best Paper Finalist)
 * [Visual-Locomotion: Learning to Walk on Complex Terrains with Vision](https://proceedings.mlr.press/v164/yu22a.html) (CoRL 2022)
@@ -16,6 +20,6 @@
 * [Provably Robust Blackbox Optimization for Reinforcement Learning](https://arxiv.org/abs/1903.02993) (CoRL 2019)
 * [Structured Evolution with Compact Architectures for Scalable Policy Optimization](https://arxiv.org/abs/1804.02395) (ICML 2018)
 * [Optimizing Simulations with Noise-Tolerant Structured Exploration](https://arxiv.org/abs/1805.07831) (ICRA 2018)
-* [On Blackbox Backpropagation and Jacobian Sensing](https://proceedings.neurips.cc/paper_files/paper/2017/file/9c8661befae6dbcd08304dbf4dcaf0db-Paper.pdf) (Neurips 2017)
+* [On Blackbox Backpropagation and Jacobian Sensing](https://proceedings.neurips.cc/paper_files/paper/2017/file/9c8661befae6dbcd08304dbf4dcaf0db-Paper.pdf) (NeurIPS 2017)
 
 **Disclaimer:** This is not an officially supported Google product.

requirements-extras.txt

@@ -0,0 +1,3 @@
+# Extras for supporting all features of the project.
+pyglove  # For Pyglove evolutionary algorithms.
+ribs  # For PyRibs: https://pyribs.org/

requirements-rl.txt

@@ -0,0 +1,14 @@
+# Packages required for Reinforcement Learning.
+
+# Environment packages.
+gym
+tf-agents  # NOTE: Requires tensorflow>=2.15.0 for TFP compatibility.
+
+# Numerical packages.
+jax  # Use latest version.
+jaxlib  # Use latest version.
+flax  # Use latest version.
+tensorflow  # TODO(team): Resolve version conflicts.
+
+# Configuration + Experimentation
+gin-config>=0.5.0

requirements-test.txt

@@ -0,0 +1 @@
+pytest  # Use the latest version to match github workflow.

requirements.txt

@@ -1,27 +1,9 @@
-# Standard packages.
 typing  # Version dependent on Python version.
-pytest  # Use the latest version to match github workflow.
 absl-py>=1.0.0
-
-# Numerical packages.
 numpy>=1.21.5
-jax  # Use latest version.
-jaxlib  # Use latest version.
-flax  # Use latest version.
-tensorflow  # TODO(team): Resolve version conflicts.
 
 # Distributed systems libraries.
 # NOTE: Requires tensorflow~=2.8.0 to avoid proto issues.
 dm-launchpad[tensorflow]
 dm-reverb[tensorflow]
-
-# Configuration + Experimentation
 ml-collections>=0.1.1
-gin-config>=0.5.0
-
-# Reinforcement Learning
-gym
-tf-agents  # NOTE: Requires tensorflow>=2.15.0 for TFP compatibility.
-
-# Optimization packages
-pyglove

setup.py

@@ -14,6 +14,7 @@
 
 """Setup for pip package."""
 
+import itertools
 import setuptools
 
 
@@ -33,6 +34,16 @@ def _parse_requirements(requirements_txt_path: str) -> list[str]:
     return [l for l in lines if (l and 'github.com' not in l)]
 
 
+extras_require = {
+    'rl': _parse_requirements('requirements-rl.txt'),
+    'extras': _parse_requirements('requirements-extras.txt'),
+    'test': _parse_requirements('requirements-test.txt'),
+}
+
+extras_require['all'] = list(
+    itertools.chain.from_iterable(extras_require.values())
+)
+
 setuptools.setup(
     name='google-iris',
     version='0.0.1.alpha',
@@ -41,5 +52,6 @@ def _parse_requirements(requirements_txt_path: str) -> list[str]:
     author_email='jaindeepali@google.com',
     install_requires=_parse_requirements('requirements.txt'),
     packages=setuptools.find_packages(),
+    extras_require=extras_require,
     python_requires='>=3.10',
 )