doc: enable markdowm (myst) doc

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"

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.

docs/source/guide/advanced/2-jit-able.rst → docs/source/guide/advanced/2-jit-able.md

@@ -1,47 +1,37 @@
-===================
-Jit-able components
-===================
+# Jit-able components
 
-A common pitfall in jit
-=======================
+## A common pitfall in jit
 
 In JAX, it's hard to jump out of a jit-compiled function, meaning if you jit-compile one function,
 then all other functions used within this function must also be jit-compiled.
 
 For example, the follow code will result in compilation error
 
-.. code-block:: python
+```python
+@jax.jit
+def foo(x):
+    return bar(x)
 
-    @jax.jit
-    def foo(x):
-        return bar(x)
+def bar(x):
+    return x[x > 0] # dynamic index, not jit-able
+```
 
-    def bar(x):
-        return x[x > 0] # dynamic index, not jit-able
+Even though `bar` is not marked with `jax.jit`, it is still compiled as `foo` calls `bar`.
+And since bar uses dynamic index, which is not compatible with `jax.jit`, an error will occur.
 
-Even though ``bar`` is not marked with ``jax.jit``, it is still compiled as ``foo`` calls ``bar``.
-And since bar uses dynamic index, which is not compatible with ``jax.jit``, an error will occur.
-
-Solution
-========
+## Solution
 
 To solve is problem, it is common practice to jit-compile low level components, thus give high level components more freedom.
 In EvoX, we have some general rules on whether a function should be jit-able or not.
 
-+-----------+----------+
 | Component | jit-able |
-+===========+==========+
+| --------- | -------- |
 | Workflow  | Optional |
-+-----------+----------+
 | Algorithm | Yes      |
-+-----------+----------+
 | Problem   | Optional |
-+-----------+----------+
 | Operators | Yes      |
-+-----------+----------+
 | Monitor   | No       |
-+-----------+----------+
 
 For standard workflow, one can jit compile when not using monitors and working with jit-able problems.
 But even though the workflow can be compiled, there isn't much performance gain.
-For problems, it depends on the task.
+For problems, it depends on the task.