From e0d40beaffd320aca3fce8fbdff448a25724eb0a Mon Sep 17 00:00:00 2001
From: Bill Huang <bill.huang2001@gmail.com>
Date: Thu, 9 Jan 2025 13:51:14 +0800
Subject: [PATCH 1/3] Add optional test and docs dependencies

---
 pyproject.toml | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index dc08d47fb..71d2112df 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,12 +34,26 @@ vis = [
   "pandas >= 2.0.0",
 ]
 
-gymnasium = ["gymnasium >= 0.29.0"]
+neuroevolution = [
+  "torchvision >= 0.16.0",
+  "brax >= 0.12.0",
+]
 
-envpool = ["envpool >= 0.8.0"]
+test = [
+  "torchvision >= 0.16.0",
+  "brax >= 0.12.0",
+]
 
-neuroevolution = [
-  "brax >= 0.11.0",
+docs = [
+  "sphinx-book-theme > 1.0",
+  "sphinx",
+  "myst-nb",
+  "nbsphinx",
+  "myst-parser",
+  "sphinx_copybutton",
+  "sphinx_design",
+  "sphinx-favicon",
+  "sphinxcontrib-mermaid",
 ]
 
 [project.urls]

From 73849725725b8b65d442e702b98d77d0aac08ce9 Mon Sep 17 00:00:00 2001
From: Bill Huang <bill.huang2001@gmail.com>
Date: Thu, 9 Jan 2025 13:51:31 +0800
Subject: [PATCH 2/3] Update readthedocs build config

---
 .readthedocs.yaml | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/.readthedocs.yaml b/.readthedocs.yaml
index 460c1e610..c5107cb3a 100644
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -7,9 +7,9 @@ version: 2
 
 # Set the version of Python and other tools you might need
 build:
-  os: ubuntu-22.04
+  os: ubuntu-lts-latest
   tools:
-    python: "3.10"
+    python: "3.11"
   jobs:
     pre_build:
       - python docs/fix_notebook_translation.py
@@ -21,4 +21,7 @@ sphinx:
 # Optionally declare the Python requirements required to build your docs
 python:
   install:
-  - requirements: requirements/docs-requirements.txt
\ No newline at end of file
+    - method: pip
+      path: .
+      extra_requirements:
+        - docs

From 05f5436aac01ad24efa961400f183ba7ecb8991a Mon Sep 17 00:00:00 2001
From: Bill Huang <bill.huang2001@gmail.com>
Date: Thu, 9 Jan 2025 13:52:02 +0800
Subject: [PATCH 3/3] Mock brax and torchvision during doc build

---
 docs/source/conf.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 6bb46f7b7..ace83cf20 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -52,13 +52,7 @@
 # mock these modules, so we can build the document without these dependencies.
 autodoc_mock_imports = [
     "brax",
-    "chex",
-    "envpool",
-    "gymnasium",
-    "ray",
-    "tensorflow_datasets",
-    "gpjax",
-    "orbax-checkpoint",
+    "torchvision",
 ]
 
 # -- Options for HTML output -------------------------------------------------