diff --git a/docs/source/conf.py b/docs/source/conf.py
index cf7076b76..5c7bf2f51 100644
--- 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"
\ No newline at end of file
diff --git a/docs/source/guide/advanced/1-state.md b/docs/source/guide/advanced/1-state.md
new file mode 100644
index 000000000..b27d0c5a4
--- /dev/null
+++ 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
deleted file mode 100644
index 059c27db2..000000000
--- a/docs/source/guide/advanced/1-state.rst
+++ /dev/null
@@ -1,93 +0,0 @@
-==========================
-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``.
-
-.. code-block:: 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:
-
-.. code-block:: 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.
-
-.. code-block:: 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.
\ No newline at end of file
diff --git a/docs/source/guide/advanced/2-jit-able.rst b/docs/source/guide/advanced/2-jit-able.md
similarity index 58%
rename from docs/source/guide/advanced/2-jit-able.rst
rename to docs/source/guide/advanced/2-jit-able.md
index 1d1f1a77f..bea92b3d7 100644
--- a/docs/source/guide/advanced/2-jit-able.rst
+++ b/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.
\ No newline at end of file
+For problems, it depends on the task.
diff --git a/docs/source/guide/advanced/3-custom-alg-pro.md b/docs/source/guide/advanced/3-custom-alg-pro.md
new file mode 100644
index 000000000..7097eef87
--- /dev/null
+++ b/docs/source/guide/advanced/3-custom-alg-pro.md
@@ -0,0 +1,195 @@
+```{eval-rst}
+.. role:: python(code)
+  :language: python
+  :class: highlight
+```
+
+# Custom algorithms and problems in EvoX
+
+In this chapter, we will introduce how to implement your own algorithm in EvoX.
+
+## The Algorithm Class
+
+The {class}`Algorithm <evox.Algorithm>` class is inherented from {class}`Stateful <evox.Stateful>`.
+Besides the things in `Stateful`, your should also implement a `ask` and a `tell` method.
+In total, there are four methods one need to implement.
+
+| Method       | Signature                               | Usage                                                                                                              |
+| ------------ | --------------------------------------- | ------------------------------------------------------------------------------------------------------------------ |
+| \_\_init\_\_ | {python}`(self, ...)`                   | Initialize hyperparameters that are fixed though out the optimization process, for example, the `population size`. |
+| setup        | {python}`(self, RRNGKey) -> State`      | Initialize mutable state, for example the `momentum`.                                                              |
+| ask          | {python}`(self, State) -> Array, State` | Gives a candidate population for evaluation.                                                                       |
+| tell         | {python}`(self, State, Array) -> State` | Receive the fitness for the candidate population and update the algorithm's state.                                 |
+
+### Migrate from traditional EC library
+
+In traditional EC library, algorithm usually calls the objective function internally, which gives the following layout
+
+```
+Algorithm
+|
++--Problem
+```
+
+But in EvoX, we have a flat layout
+
+```
+Algorithm.ask -- Problem.evaluate -- Algorithm.tell
+```
+
+Here is a pseudocode of a genetic algorithm.
+
+```python
+Set hyperparameters
+Generate the initial population
+Do
+    Generate Offspring
+        Selection
+        Crossover
+        Mutation
+    Compute fitness
+    Replace the population
+Until stopping criterion
+```
+
+And Here is what each part of the algorithm correspond to in EvoX.
+
+```python
+Set hyperparameters # __init__
+Generate the initial population # setup
+Do
+    # ask
+    Generate Offspring
+        Mating Selection
+        Crossover
+        Mutation
+
+    # problem.evaluate (not part of the algorithm)
+    Compute fitness
+
+    # tell
+    Survivor Selection
+Until stopping criterion
+```
+
+## The Problem Class
+
+The Problem class is quite simple, beside `__init__` and `setup`, the only required the method is `evaluate`.
+
+### Migrate from traditional EC library
+
+There is one thing to notice here, `evaluate` is a stateful function, meaning it should accept a state and return a new state.
+So, if you are working with numerical benchmark functions, which don't need to statefule,
+you can simply ignore the state, but remember that you still have to use this stateful interface.
+
+```{eval-rst}
++----------+------------------------------------------------+-------------------------------------------------------+
+| Method   | Signature                                      | Usage                                                 |
++----------+------------------------------------------------+-------------------------------------------------------+
+| __init__ | :python:`(self, ...)`                          | Initialize the settings of the problem.               |
++----------+------------------------------------------------+-------------------------------------------------------+
+| setup    | :python:`(self, RRNGKey) -> State`             | Initialize mutable state of this problem.             |
++----------+------------------------------------------------+-------------------------------------------------------+
+| evaluate | :python:`(self, State, Array) -> Array, State` | Evaluate the fitness of the given candidate solution. |
++----------+------------------------------------------------+-------------------------------------------------------+
+```
+
+### More on the problem's state
+
+If you still wonders what the problem's state actually do, here are the explanations.
+
+Unlike numerical benchmark functions, real-life problems are more complex, and may require stateful computations.
+Here are some examples:
+
+- When dealing with ANN training, we often have training, validation and testing phase.
+  This implies that the same solution could have different fitness values during different phases.
+  So clearly, we can't model the `evaluate` as a stateless pure function any more.
+  To implement this mechanism, simple put an value in the state to indicate the phase.
+- Virtual batch norm is a effective trick especially when dealing with RL tasks.
+  To implement this mechanism, the problem must be stateful,
+  as the problem have to remember the initial batch norm parameters during the first run.
+
+## Example
+
+Here we give an exmaple of implementing the OneMax problem, along with a genetic algorithm that solves this problem.
+The problem itself is straight forward, the fitness is defined as the sum of every digits in a fixed-length bitstring.
+For example, "100111" gives 4 and "000101" gives 2.
+
+Let's starts with implementing the OneMax problem.
+In JAX a bitstring can be easily represented with a tensor of type bool.
+
+```python
+import jax.numpy as jnp
+from evox import Problem, jit_class
+
+
+@jit_class
+class OneMax(Problem):
+    def __init__(self, neg_fitness=True) -> None:
+        super().__init__()
+        self.neg_fitess = neg_fitness
+
+    def evaluate(self, state, bitstrings):
+        # bitstrings has shape (pop_size, num_bits)
+        # so sum along the axis 1.
+        fitness = jnp.sum(bitstrings, axis=1)
+        # Since in EvoX, algorithms try to minimize the fitness
+        # so return the negitive value.
+        if self.neg_fitess:
+            fitness = -fitness
+        return fitness, state
+```
+
+Then we implement a genetic algorithm that uses bitflip mutation and one-point crossover.
+
+```python
+@jit_class
+class ExampleGA(Algorithm):
+    def __init__(self, pop_size, ndim, flip_prob):
+        super().__init__()
+        # those are hyperparameters that stay fixed.
+        self.pop_size = pop_size
+        self.ndim = ndim
+        # the probability of fliping each bit
+        self.flip_prob = flip_prob
+
+    def setup(self, key):
+        # initialize the state
+        # state are mutable data like the population, offsprings
+        # the population is randomly initialized.
+        # we don't have any offspring now, but initialize it as a placeholder
+        # because jax want static shaped arrays.
+        key, subkey = random.split(key)
+        pop = random.uniform(subkey, (self.pop_size, self.ndim)) < 0.5
+        return State(
+            pop=pop,
+            offsprings=jnp.empty((self.pop_size * 2, self.ndim)),
+            fit=jnp.full((self.pop_size,), jnp.inf),
+            key=key,
+        )
+
+    def ask(self, state):
+        key, mut_key, x_key = random.split(state.key, 3)
+        # here we do mutation and crossover (reproduction)
+        # for simplicity, we didn't use any mating selections
+        # so the offspring is twice as large as the population
+        offsprings = jnp.concatenate(
+            (
+                mutation.bitflip(mut_key, state.pop, self.flip_prob),
+                crossover.one_point(x_key, state.pop),
+            ),
+            axis=0,
+        )
+        # return the candidate solution and update the state
+        return offsprings, state.update(offsprings=offsprings, key=key)
+
+    def tell(self, state, fitness):
+        # here we do selection
+        merged_pop = jnp.concatenate([state.pop, state.offsprings])
+        merged_fit = jnp.concatenate([state.fit, fitness])
+        new_pop, new_fit = selection.topk_fit(merged_pop, merged_fit, self.pop_size)
+        # replace the old population
+        return state.update(pop=new_pop, fit=new_fit)
+```
+
+Now, you can assemble a workflow and run it.
diff --git a/docs/source/guide/advanced/3-custom-alg-pro.rst b/docs/source/guide/advanced/3-custom-alg-pro.rst
deleted file mode 100644
index 972cfafb1..000000000
--- a/docs/source/guide/advanced/3-custom-alg-pro.rst
+++ /dev/null
@@ -1,209 +0,0 @@
-.. role:: python(code)
-  :language: python
-  :class: highlight
-
-======================================
-Custom algorithms and problems in EvoX
-======================================
-
-In this chapter, we will introduce how to implement your own algorithm in EvoX.
-
-The Algorithm Class
-===================
-
-The :class:`Algorithm <evox.Algorithm>` class is inherented from :class:`Stateful <evox.Stateful>`.
-Besides the things in ``Stateful``, your should also implement a ``ask`` and a ``tell`` method.
-In total, there are four methods one need to implement.
-
-+----------+-----------------------------------------+------------------------------------------------------------------------------------+
-| Method   | Signature                               | Usage                                                                              |
-+==========+=========================================+====================================================================================+
-| __init__ | :python:`(self, ...)`                   | Initialize hyperparameters that are fixed though out the optimization process,     |
-|          |                                         | for example, the ``population size``.                                              |
-+----------+-----------------------------------------+------------------------------------------------------------------------------------+
-| setup    | :python:`(self, RRNGKey) -> State`      | Initialize mutable state, for example the ``momentum``.                            |
-+----------+-----------------------------------------+------------------------------------------------------------------------------------+
-| ask      | :python:`(self, State) -> Array, State` | Gives a candidate population for evaluation.                                       |
-+----------+-----------------------------------------+------------------------------------------------------------------------------------+
-| tell     | :python:`(self, State, Array) -> State` | Receive the fitness for the candidate population and update the algorithm's state. |
-+----------+-----------------------------------------+------------------------------------------------------------------------------------+
-
-
-Migrate from traditional EC library
------------------------------------
-
-In traditional EC library, algorithm usually calls the objective function internally, which gives the following layout
-
-.. code-block::
-
-    Algorithm
-    |
-    +--Problem
-
-But in EvoX, we have a flat layout
-
-.. code-block::
-
-    Algorithm.ask -- Problem.evaluate -- Algorithm.tell
-
-
-Here is a pseudocode of a genetic algorithm.
-
-.. code-block:: python
-
-    Set hyperparameters
-    Generate the initial population
-    Do
-        Generate Offspring
-            Selection
-            Crossover
-            Mutation
-        Compute fitness
-        Replace the population
-    Until stopping criterion
-
-And Here is what each part of the algorithm correspond to in EvoX.
-
-.. code-block:: python
-
-    Set hyperparameters # __init__
-    Generate the initial population # setup
-    Do
-        # ask
-        Generate Offspring
-            Mating Selection
-            Crossover
-            Mutation
-
-        # problem.evaluate (not part of the algorithm)
-        Compute fitness
-
-        # tell
-        Survivor Selection
-    Until stopping criterion
-
-The Problem Class
-=================
-
-The Problem class is quite simple, beside ``__init__`` and ``setup``, the only required the method is ``evaluate``.
-
-Migrate from traditional EC library
------------------------------------
-
-There is one thing to notice here, ``evaluate`` is a stateful function, meaning it should accept a state and return a new state.
-So, if you are working with numerical benchmark functions, which don't need to statefule,
-you can simply ignore the state, but remember that you still have to use this stateful interface.
-
-+----------+------------------------------------------------+-------------------------------------------------------+
-| Method   | Signature                                      | Usage                                                 |
-+----------+------------------------------------------------+-------------------------------------------------------+
-| __init__ | :python:`(self, ...)`                          | Initialize the settings of the problem.               |
-+----------+------------------------------------------------+-------------------------------------------------------+
-| setup    | :python:`(self, RRNGKey) -> State`             | Initialize mutable state of this problem.             |
-+----------+------------------------------------------------+-------------------------------------------------------+
-| evaluate | :python:`(self, State, Array) -> Array, State` | Evaluate the fitness of the given candidate solution. |
-+----------+------------------------------------------------+-------------------------------------------------------+
-
-More on the problem's state
----------------------------
-
-If you still wonders what the problem's state actually do, here are the explanations.
-
-Unlike numerical benchmark functions, real-life problems are more complex, and may require stateful computations.
-Here are some examples:
-
-* When dealing with ANN training, we often have training, validation and testing phase.
-  This implies that the same solution could have different fitness values during different phases.
-  So clearly, we can't model the `evaluate` as a stateless pure function any more.
-  To implement this mechanism, simple put an value in the state to indicate the phase.
-* Virtual batch norm is a effective trick especially when dealing with RL tasks.
-  To implement this mechanism, the problem must be stateful,
-  as the problem have to remember the initial batch norm parameters during the first run.
-
-
-Example
-=======
-
-Here we give an exmaple of implementing the OneMax problem, along with a genetic algorithm that solves this problem.
-The problem itself is straight forward, the fitness is defined as the sum of every digits in a fixed-length bitstring.
-For example, "100111" gives 4 and "000101" gives 2.
-
-Let's starts with implementing the OneMax problem.
-In JAX a bitstring can be easily represented with a tensor of type bool.
-
-.. code-block:: python
-
-    import jax.numpy as jnp
-    from evox import Problem, jit_class
-
-
-    @jit_class
-    class OneMax(Problem):
-        def __init__(self, neg_fitness=True) -> None:
-            super().__init__()
-            self.neg_fitess = neg_fitness
-
-        def evaluate(self, state, bitstrings):
-            # bitstrings has shape (pop_size, num_bits)
-            # so sum along the axis 1.
-            fitness = jnp.sum(bitstrings, axis=1)
-            # Since in EvoX, algorithms try to minimize the fitness
-            # so return the negitive value.
-            if self.neg_fitess:
-                fitness = -fitness
-            return fitness, state
-
-
-Then we implement a genetic algorithm that uses bitflip mutation and one-point crossover.
-
-.. code-block:: python
-
-    @jit_class
-    class ExampleGA(Algorithm):
-        def __init__(self, pop_size, ndim, flip_prob):
-            super().__init__()
-            # those are hyperparameters that stay fixed.
-            self.pop_size = pop_size
-            self.ndim = ndim
-            # the probability of fliping each bit
-            self.flip_prob = flip_prob
-
-        def setup(self, key):
-            # initialize the state
-            # state are mutable data like the population, offsprings
-            # the population is randomly initialized.
-            # we don't have any offspring now, but initialize it as a placeholder
-            # because jax want static shaped arrays.
-            key, subkey = random.split(key)
-            pop = random.uniform(subkey, (self.pop_size, self.ndim)) < 0.5
-            return State(
-                pop=pop,
-                offsprings=jnp.empty((self.pop_size * 2, self.ndim)),
-                fit=jnp.full((self.pop_size,), jnp.inf),
-                key=key,
-            )
-
-        def ask(self, state):
-            key, mut_key, x_key = random.split(state.key, 3)
-            # here we do mutation and crossover (reproduction)
-            # for simplicity, we didn't use any mating selections
-            # so the offspring is twice as large as the population
-            offsprings = jnp.concatenate(
-                (
-                    mutation.bitflip(mut_key, state.pop, self.flip_prob),
-                    crossover.one_point(x_key, state.pop),
-                ),
-                axis=0,
-            )
-            # return the candidate solution and update the state
-            return offsprings, state.update(offsprings=offsprings, key=key)
-
-        def tell(self, state, fitness):
-            # here we do selection
-            merged_pop = jnp.concatenate([state.pop, state.offsprings])
-            merged_fit = jnp.concatenate([state.fit, fitness])
-            new_pop, new_fit = selection.topk_fit(merged_pop, merged_fit, self.pop_size)
-            # replace the old population
-            return state.update(pop=new_pop, fit=new_fit)
-
-Now, you can assemble a workflow and run it.
\ No newline at end of file
diff --git a/docs/source/guide/advanced/4-container.rst b/docs/source/guide/advanced/4-container.md
similarity index 52%
rename from docs/source/guide/advanced/4-container.rst
rename to docs/source/guide/advanced/4-container.md
index 6a991b131..b69787fdd 100644
--- a/docs/source/guide/advanced/4-container.rst
+++ b/docs/source/guide/advanced/4-container.md
@@ -1,22 +1,18 @@
-====================
-Container Algorithms
-====================
+# Container Algorithms
 
 Container algorithms are a special type of algorithms that works by containing other algorithms and cannot work on its own.
 Container algorithms can be used to compose a series of normal algorithms together.
 
-Working with expensive algorithms
-=================================
+## Working with expensive algorithms
 
-Many algorithms are expensive in term of space or time. For example, CMA-ES requires :math:`O(N^2)` space.
+Many algorithms are expensive in term of space or time. For example, CMA-ES requires $O(N^2)$ space.
 Thus, it is costly to run CMA-ES on high-dimension problems.
 Sep-CMA-ES scales better, but sacrifice the performance.
 That's where container algorithm comes in.
-With it, we can easily construct a variant of CMA-ES that uses :math:`O((\frac{N}{M})^2)` space, where :math:`M` is the number of block.
+With it, we can easily construct a variant of CMA-ES that uses $O((\frac{N}{M})^2)$ space, where $M$ is the number of block.
 This variant is a balance between the normal CMA-ES and Sep-CMA-ES.
 
-Working with PyTree
-===================
+## Working with PyTree
 
 Usually, algorithms expect the decision variables to be in the form of a 1D-vector.
 PyTrees are tree-like structures that are not directly compatible with normal algorithms.
@@ -25,15 +21,14 @@ So, there are two solutions out there:
 1. Flatten the PyTree to 1D-vector.
 2. Use a specialized algorithm that work with PyTree directly.
 
-Solution 1 is called ``adapter`` in EvoX, which is quite simple, but we are not talking about this here.
+Solution 1 is called `adapter` in EvoX, which is quite simple, but we are not talking about this here.
 Solution 2 seems more complicated, but the advantage is that the structural information is preserved,
 meaning the algorithm could see the tree structure and apply some type of heuristic here.
 
-Cooperative Coevolution
-=======================
+## Cooperative Coevolution
 
 We offer Cooperative Coevolution (CC) framework for all algorithms.
-Currently, there are two types of CC container in EvoX, :class:`evox.algorithms.Coevolution` and :class:`evox.algorithms.VectorizedCoevolution`.
-The difference is that ``VectorizedCoevolution`` update all sub-populations at the same time in each iteration,
-but ``Coevolution`` follows traditional approach that update a single sub-populations at each iteration.
-Thus ``VectorizedCoevolution`` is faster, but ``Coevolution`` could be better in terms of best fitness with a limited number of evaluations.
+Currently, there are two types of CC container in EvoX, {class}`evox.algorithms.Coevolution` and {class}`evox.algorithms.VectorizedCoevolution`.
+The difference is that `VectorizedCoevolution` update all sub-populations at the same time in each iteration,
+but `Coevolution` follows traditional approach that update a single sub-populations at each iteration.
+Thus `VectorizedCoevolution` is faster, but `Coevolution` could be better in terms of best fitness with a limited number of evaluations.
diff --git a/docs/source/guide/advanced/index.md b/docs/source/guide/advanced/index.md
new file mode 100644
index 000000000..adbaf20f5
--- /dev/null
+++ b/docs/source/guide/advanced/index.md
@@ -0,0 +1,10 @@
+# Advanced Tutorial
+
+```{toctree}
+:maxdepth: 1
+
+1-state
+2-jit-able
+3-custom-alg-pro
+4-container
+```
diff --git a/docs/source/guide/advanced/index.rst b/docs/source/guide/advanced/index.rst
deleted file mode 100644
index e90d7e216..000000000
--- a/docs/source/guide/advanced/index.rst
+++ /dev/null
@@ -1,11 +0,0 @@
-=================
-Advanced Tutorial
-=================
-
-.. toctree::
-    :maxdepth: 1
-
-    1-state
-    2-jit-able
-    3-custom-alg-pro
-    4-container
\ No newline at end of file
diff --git a/docs/source/guide/basics/2-problems.md b/docs/source/guide/basics/2-problems.md
new file mode 100644
index 000000000..ce5946477
--- /dev/null
+++ b/docs/source/guide/basics/2-problems.md
@@ -0,0 +1,61 @@
+# Working with extended applications
+
+Working with extended applications in EvoX is easy.
+
+## Neuroevolution
+
+EvoX currently supports training ANN with torchvision's datasets via {class}`TorchvisionDataset <evox.problems.neuroevolution.TorchvisionDataset>`.
+
+```python
+from evox.problems.neuroevolution import TorchvisionDataset
+
+problem = TorchvisionDataset(
+    root="./", # where to download the dataset
+    forward_func=<your neural network's forward function>,
+    batch_size=512, # the batchsize
+    num_passes=2, # splite the batch and calculate using 2 passes
+    dataset_name="cifar10" # the name of the dataset
+)
+```
+
+In the above example, we created a problem using `cifar10` dataset,
+each individual in the population will be evaluated with a batch size of 512, and since `num_passes` is 2,
+the batch will be calculated in 2 passes and each pass has a mini-batch size of `512 / 2 = 256`.
+
+## RL
+
+EvoX currently supports rl environments from Gym.
+
+```python
+from evox.problems.neuroevolution import Gym
+problem = Gym(
+    env_name="ALE/Pong-v5", # the environment's name
+    env_options={"full_action_space": False}, # the options passes to the environment
+    policy=<your policy>,
+    num_workers=16, # number of processes
+    env_per_worker=4, # the number of environments each process holds
+    controller_options={ # the options that passes to ray
+        "num_cpus": 1,
+        "num_gpus": 0,
+    },
+    worker_options={ # the options that passes to ray
+        "num_cpus": 1, "num_gpus": 1 / 16
+    },
+    batch_policy=False,
+)
+```
+
+:::{tip}
+`controller_options` and `worker_options` are directly passed down to Ray,
+usually, you can leave `controller_options` blank because there isn't much computation happen there.
+`worker_options` can be used to control how much cpu and gpu each worker gets.
+
+For example:
+
+```python
+worker_options={
+    "num_cpus": 2,
+    "num_gpus": 0,
+}
+```
+:::
diff --git a/docs/source/guide/basics/2-problems.rst b/docs/source/guide/basics/2-problems.rst
deleted file mode 100644
index 3f3fc08ac..000000000
--- a/docs/source/guide/basics/2-problems.rst
+++ /dev/null
@@ -1,64 +0,0 @@
-==================================
-Working with extended applications
-==================================
-
-Working with extended applications in EvoX is easy.
-
-Neuroevolution
-==============
-
-EvoX currently supports training ANN with torchvision's datasets via :class:`TorchvisionDataset <evox.problems.neuroevolution.TorchvisionDataset>`.
-
-.. code-block:: python
-
-    from evox.problems.neuroevolution import TorchvisionDataset
-
-    problem = TorchvisionDataset(
-        root="./", # where to download the dataset
-        forward_func=<your neural network's forward function>,
-        batch_size=512, # the batchsize
-        num_passes=2, # splite the batch and calculate using 2 passes
-        dataset_name="cifar10" # the name of the dataset
-    )
-
-In the above example, we created a problem using ``cifar10`` dataset,
-each individual in the population will be evaluated with a batch size of 512, and since ``num_passes`` is 2,
-the batch will be calculated in 2 passes and each pass has a mini-batch size of ``512 / 2 = 256``.
-
-RL
-===
-
-EvoX currently supports rl environments from Gym.
-
-.. code-block:: python
-
-    from evox.problems.neuroevolution import Gym
-    problem = Gym(
-        env_name="ALE/Pong-v5", # the environment's name
-        env_options={"full_action_space": False}, # the options passes to the environment
-        policy=<your policy>,
-        num_workers=16, # number of processes
-        env_per_worker=4, # the number of environments each process holds
-        controller_options={ # the options that passes to ray
-            "num_cpus": 1,
-            "num_gpus": 0,
-        },
-        worker_options={ # the options that passes to ray
-            "num_cpus": 1, "num_gpus": 1 / 16
-        },
-        batch_policy=False,
-    )
-
-.. tip::
-    ``controller_options`` and ``worker_options`` are directly passed down to Ray,
-    usually, you can leave ``controller_options`` blank because there isn't much computation happen there.
-    ``worker_options`` can be used to control how much cpu and gpu each worker gets.
-
-    For example:
-
-    .. code-block:: python
-
-        worker_options={
-            "num_cpus": 2,
-            "num_gpus": 0,
-        }
diff --git a/docs/source/guide/basics/3-distributed.md b/docs/source/guide/basics/3-distributed.md
new file mode 100644
index 000000000..664c90a8a
--- /dev/null
+++ b/docs/source/guide/basics/3-distributed.md
@@ -0,0 +1,26 @@
+# Distribute the workflow
+
+To scale the workflow using multiple machines, use the {class}`DistributedWorkflow <evox.workflows.DistributedWorkflow>` instead of StdWorkflow.
+
+```python
+algorithm = <your algorithm>
+problem = <your problem>
+
+from evox.workflows import DistributedWorkflow
+
+workflow = DistributedWorkflow(
+    algorithm=algorithm,
+    problem=problem,
+    pop_size=100, # the actual population size used by the algorithm
+    num_workers=4, # the number of machines
+    options={ # the options that passes to ray
+        "num_gpus": 1
+    }
+)
+```
+
+:::{tip}
+It is recommanded that one set the environment variable `XLA_PYTHON_CLIENT_PREALLOCATE=false`.
+This variable control disables the GPU memory preallocation, otherwise running multiple JAX processes may cause OOM.
+For more information, please refer to [JAX's documentation](https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html) on this matter.
+:::
diff --git a/docs/source/guide/basics/3-distributed.rst b/docs/source/guide/basics/3-distributed.rst
deleted file mode 100644
index f7c692645..000000000
--- a/docs/source/guide/basics/3-distributed.rst
+++ /dev/null
@@ -1,27 +0,0 @@
-=======================
-Distribute the workflow
-=======================
-
-To scale the workflow using multiple machines, use the :class:`DistributedWorkflow <evox.workflows.DistributedWorkflow>` instead of StdWorkflow.
-
-.. code-block:: python
-
-    algorithm = <your algorithm>
-    problem = <your problem>
-
-    from evox.workflows import DistributedWorkflow
-
-    workflow = DistributedWorkflow(
-        algorithm=algorithm,
-        problem=problem,
-        pop_size=100, # the actual population size used by the algorithm
-        num_workers=4, # the number of machines
-        options={ # the options that passes to ray
-            "num_gpus": 1
-        }
-    )
-
-.. tip::
-    It is recommanded that one set the environment variable ``XLA_PYTHON_CLIENT_PREALLOCATE=false``.
-    This variable control disables the GPU memory preallocation, otherwise running multiple JAX processes may cause OOM.
-    For more information, please refer to `JAX's documentation <https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html>`_ on this matter.
diff --git a/docs/source/guide/basics/index.md b/docs/source/guide/basics/index.md
new file mode 100644
index 000000000..daeb57e03
--- /dev/null
+++ b/docs/source/guide/basics/index.md
@@ -0,0 +1,9 @@
+# Getting Started
+
+```{toctree}
+:maxdepth: 1
+
+1-start
+2-problems
+3-distributed
+```
diff --git a/docs/source/guide/basics/index.rst b/docs/source/guide/basics/index.rst
deleted file mode 100644
index 94ca8e273..000000000
--- a/docs/source/guide/basics/index.rst
+++ /dev/null
@@ -1,10 +0,0 @@
-===============
-Getting Started
-===============
-
-.. toctree::
-    :maxdepth: 1
-
-    1-start
-    2-problems
-    3-distributed
\ No newline at end of file
diff --git a/docs/source/guide/index.md b/docs/source/guide/index.md
new file mode 100644
index 000000000..3677ff7ad
--- /dev/null
+++ b/docs/source/guide/index.md
@@ -0,0 +1,9 @@
+# EvoX's guide!
+
+```{toctree}
+:maxdepth: 2
+
+install
+basics/index
+advanced/index
+```
diff --git a/docs/source/guide/index.rst b/docs/source/guide/index.rst
deleted file mode 100644
index 253fbe984..000000000
--- a/docs/source/guide/index.rst
+++ /dev/null
@@ -1,10 +0,0 @@
-================
-EvoX's guide!
-================
-
-.. toctree::
-    :maxdepth: 2
-
-    install
-    basics/index
-    advanced/index
\ No newline at end of file
diff --git a/docs/source/guide/install.md b/docs/source/guide/install.md
new file mode 100644
index 000000000..550d49ba4
--- /dev/null
+++ b/docs/source/guide/install.md
@@ -0,0 +1,20 @@
+# Install EvoX
+
+EvoX is available at Pypi and can be installed via:
+
+```bash
+pip install evox
+```
+
+To install EvoX with optional dependencies:
+
+```bash
+pip install evox[<feature1>,<feature2>]
+```
+
+available features are `gym`, `neuroevolution`, `distributed`, and `full` which concludes all features.
+For example, to install EvoX with all features, do:
+
+```bash
+pip install evox[full]
+```
diff --git a/docs/source/guide/install.rst b/docs/source/guide/install.rst
deleted file mode 100644
index d9797be0f..000000000
--- a/docs/source/guide/install.rst
+++ /dev/null
@@ -1,22 +0,0 @@
-===============
-Install EvoX
-===============
-
-EvoX is available at Pypi and can be installed via:
-
-.. code:: bash
-
-    pip install evox
-
-To install EvoX with optional dependencies:
-
-.. code:: bash
-
-    pip install evox[<feature1>,<feature2>]
-
-available features are ``gym``, ``neuroevolution``, ``distributed``, and ``full`` which concludes all features.
-For example, to install EvoX with all features, do:
-
-.. code:: bash
-
-    pip install evox[full]
\ No newline at end of file
diff --git a/docs/source/index.rst b/docs/source/index.md
similarity index 56%
rename from docs/source/index.rst
rename to docs/source/index.md
index 0dabc174e..9d83cd394 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.md
@@ -1,19 +1,14 @@
-.. evox documentation master file, created by
-   sphinx-quickstart on Thu Jul 28 19:12:56 2022.
-   You can adapt this file completely to your liking, but it should at least
-   contain the root `toctree` directive.
+# Welcome to EvoX's documentation!
 
-================================
-Welcome to EvoX's documentation!
-================================
+```{toctree}
+:caption: 'Contents:'
+:maxdepth: 1
 
-.. toctree::
-   :maxdepth: 1
-   :caption: Contents:
-
-   User Guide <guide/index>
-   API reference <api/index>
+User Guide <guide/index>
+API reference <api/index>
+```
 
+```{eval-rst}
 .. grid:: 4
    :gutter: 1 1 2 4
    :padding: 1
@@ -42,10 +37,10 @@ Welcome to EvoX's documentation!
       :link: api/problems/index
       :link-type: doc
 
+```
 
-Indices and tables
-==================
+## Indices and tables
 
-* :ref:`genindex`
-* :ref:`modindex`
-* :ref:`search`
+- {ref}`genindex`
+- {ref}`modindex`
+- {ref}`search`
diff --git a/requirements/docs-requirements.txt b/requirements/docs-requirements.txt
index e28687ba4..197e63882 100644
--- a/requirements/docs-requirements.txt
+++ b/requirements/docs-requirements.txt
@@ -7,7 +7,8 @@ pyarrow
 # for docs generation
 pydata-sphinx-theme > 0.13
 sphinx
-numpydoc
+myst-nb
+nbsphinx
 myst-parser
 sphinx_copybutton
 sphinx_design
diff --git a/src/__init__.py b/src/__init__.py
deleted file mode 100644
index e69de29bb..000000000