EMI-Group · BillHuang2001 · Oct 13, 2023 · Oct 10, 2023 · Oct 10, 2023 · Oct 10, 2023
diff --git a/README.md b/README.md
@@ -18,15 +18,15 @@
 EvoX is a distributed GPU-accelerated framework for scalable evolutionary computation. Our primary goal is to push the boundaries of evolutionary computation by significantly enhancing its speed and versatility, enabling its application to complex and computationally intensive tasks.
 
 ## ⭐️ Key Features
-- ⚡️ Fast
+- 🚀 Fast
   - GPU computing for 10x-100x faster optimization.
   - Distributed workflow for even faster optimization.
 - 🌟 Wide support
   - Single-objective and multi-objective optimization.
   - Comprehensive support for commonly used benchmark problems.
   - Extensive coverage of neuroevolution problems.
 - 🎉 Easy to use
-  - Functional programming for easy function ccomposingomposition.
+  - Functional programming for easy function composition.
   - Hierarchical state management for modular programming.
   - Detailed tutorial available [here](https://evox.readthedocs.io/en/latest/guide/index.html).
 

diff --git a/docs/source/api/algorithms/mo/gde3 → docs/source/api/algorithms/mo/ged3.rst b/docs/source/api/algorithms/mo/gde3 → docs/source/api/algorithms/mo/ged3.rst
diff --git a/docs/source/api/algorithms/mo/lmocso.rst b/docs/source/api/algorithms/mo/lmocso.rst
@@ -0,0 +1,6 @@
+=======
+LMOCSO
+=======
+
+.. autoclass:: evox.algorithms.LMOCSO
+    :members:
diff --git a/docs/source/api/problems/index.rst b/docs/source/api/problems/index.rst
@@ -8,5 +8,4 @@ Problems
     :maxdepth: 2
 
     numerical/index
-    neuroevolution/index
-    rl/index
+    neuroevolution/index
diff --git a/docs/source/api/problems/neuroevolution/index.rst b/docs/source/api/problems/neuroevolution/index.rst
@@ -3,6 +3,7 @@ Neuroevolution
 ==============
 
 .. toctree::
-    :maxdepth: 1
+    :maxdepth: 2
 
-    torchvision
+    reinforcement_learning/index
+    supervised_learning/index
diff --git a/docs/source/api/problems/rl/brax.rst → ...evolution/reinforcement_learning/brax.rst b/docs/source/api/problems/rl/brax.rst → ...evolution/reinforcement_learning/brax.rst
@@ -2,5 +2,5 @@
 Brax-based Problem
 ==================
 
-.. autoclass:: evox.problems.neuroevolution.Brax
+.. autoclass:: evox.problems.neuroevolution.reinforcement_learning.Brax
     :members:
diff --git a/docs/source/api/problems/neuroevolution/reinforcement_learning/env_pool.rst b/docs/source/api/problems/neuroevolution/reinforcement_learning/env_pool.rst
@@ -0,0 +1,6 @@
+========
+Env Pool
+========
+
+.. autoclass:: evox.problems.neuroevolution.reinforcement_learning.EnvPool
+    :members:
diff --git a/docs/source/api/problems/neuroevolution/reinforcement_learning/gym.rst b/docs/source/api/problems/neuroevolution/reinforcement_learning/gym.rst
@@ -0,0 +1,6 @@
+===
+Gym
+===
+
+.. autoclass:: evox.problems.neuroevolution.reinforcement_learning.Gym
+    :members:
diff --git a/docs/source/api/problems/neuroevolution/reinforcement_learning/index.rst b/docs/source/api/problems/neuroevolution/reinforcement_learning/index.rst
@@ -0,0 +1,10 @@
+======================
+Reinforcement Learning
+======================
+
+.. toctree::
+    :maxdepth: 1
+
+    brax
+    gym
+    env_pool
diff --git a/docs/source/api/problems/neuroevolution/supervised_learning/index.rst b/docs/source/api/problems/neuroevolution/supervised_learning/index.rst
@@ -0,0 +1,8 @@
+===================
+Supervised Learning
+===================
+
+.. toctree::
+    :maxdepth: 1
+
+    torchvision
diff --git a/docs/source/api/problems/neuroevolution/supervised_learning/torchvision.rst b/docs/source/api/problems/neuroevolution/supervised_learning/torchvision.rst
@@ -0,0 +1,6 @@
+===================
+Torchvision Dataset
+===================
+
+.. autoclass:: evox.problems.neuroevolution.supervised_learning.TorchvisionDataset
+    :members:
diff --git a/docs/source/api/problems/neuroevolution/torchvision.rst b/docs/source/api/problems/neuroevolution/torchvision.rst
diff --git a/docs/source/api/problems/numerical/dtlz.rst b/docs/source/api/problems/numerical/dtlz.rst
@@ -23,8 +23,3 @@ DTLZ Test Suit
 .. autoclass:: evox.problems.numerical.DTLZ7
     :members:
 
-.. autoclass:: evox.problems.numerical.DTLZ8
-    :members:
-
-.. autoclass:: evox.problems.numerical.DTLZ9
-    :members:
diff --git a/docs/source/api/problems/rl/gym.rst b/docs/source/api/problems/rl/gym.rst
diff --git a/docs/source/api/problems/rl/index.rst b/docs/source/api/problems/rl/index.rst
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -29,7 +29,6 @@
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
-    "myst_parser",
     "numpydoc",
     "sphinx.ext.autodoc",
     "sphinx.ext.autosummary",
@@ -38,6 +37,7 @@
     "sphinx_copybutton",
     "sphinx_design",
     "sphinx_favicon",
+    "myst_nb",
 ]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -52,6 +52,7 @@
 autodoc_mock_imports = [
     "brax",
     "chex",
+    "envpool",
     "gymnasium",
     "ray",
     "torch",
@@ -82,7 +83,7 @@
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
 html_static_path = ["_static"]
-html_css_files = ["evox.css"]
+# html_css_files = ["evox.css"]
 
 autodoc_typehints_format = "short"
 autodoc_typehints = "description"
@@ -93,3 +94,4 @@
 numpydoc_show_class_members = False
 autosummary_generate = True
 autosummary_imported_members = True
+nb_execution_mode = "off"
diff --git a/docs/source/guide/advanced/1-state.md b/docs/source/guide/advanced/1-state.md
@@ -0,0 +1,85 @@
+# Working with state in EvoX
+
+EvoX is designed around the stateful computation.
+
+There are two most fundamental classes, namely {class}`Stateful <evox.Stateful>` and {class}`State <evox.State>`.
+
+All class that involves stateful computation are inherented from `Stateful`. In EvoX, `Algorithm`, `Problem`, `Operator` and workflows are all stateful.
+
+## The idea behind the design
+
+```{image} /_static/hierarchical_state.svg
+:alt: hierarchical state
+:width: 400px
+```
+
+Here we have five different objects, and notice that they have a hierarchical structure.
+To work with such structure, at each level we must "lift the state" by managing the states of child components.
+So, the state at the `workflow` level must contains the state of both `algorithm` and `problem`,
+and since the state at the `algorithm` level must contains the state of both operators,
+the state `workflow` level actual need to handle states from all 5 components.
+
+However, it is frustrating to managing the hierarchy manually, and it is not good for modular design.
+To solve this problem, we introduce `Stateful` and `State`.
+
+## An overview of Stateful
+
+In a `Stateful` class,
+all immutable data are initialized in `__init__`,
+the initial mutable state is generated in `setup`,
+besides these two method and private methods(start with "\_"),
+all other methods are wrapped with `use_state`.
+
+```python
+class Foo(Stateful):
+    def __init__(self,): # required
+        pass
+
+    def setup(self, key) -> State: # optional
+        pass
+
+    def stateful_func(self, state, args) -> State: # wrapped with use_state
+        pass
+
+    def _normal_func(self, args) -> vals: # not wrapped
+        pass
+```
+
+will be wrapped with `use_state` decorator. This decorator requires the method have the following signature:
+
+```python
+def func(self, state: State, ...) -> Tuple[..., State]
+```
+
+which is common pattern in stateful computation.
+
+:::{warning}
+Currently, for all user defined private methods, the name of the method should starts with `_`.
+:::
+
+## An overview of State
+
+In EvoX `State` represents a tree of states, which stores the state of the current object and all child objects.
+
+## Combined together
+
+When combined together,
+they will automatically go 1 level down in the tree of states,
+and merge the subtree back to current level.
+
+So you could write code like this.
+
+```python
+class FooWorkflow(Stateful):
+    ...
+    def step(self, state):
+        population, state = self.algorithm.ask(state)
+        fitness, state = self.problem.evaluate(state, population)
+        ...
+```
+
+Notice that, when calling the method `step`,
+`state` is the state of the workflow,
+but when calling `self.algorithm.ask`,
+`state` behaves like the state of the algorithm,
+and after the call, the state of the algorithm is automatically merged back into the state of the workflow.
diff --git a/docs/source/guide/advanced/1-state.rst b/docs/source/guide/advanced/1-state.rst