Update MOs and Add LMOCSO (#76)

* fix bug of DTLZ7

* add grid sampling

* update rvea,ibea and hype

* modify dtlz.rst

* add lmocso

* modify docs

* reformat the code

* modify docs

* add docstring

* add docs

* fix bug of cos that close to 0

* readme: use 🚀 to represent fast

* dev: enable multi-objective optimization for gym

* fix: avoid import unnessary packages

* fix: typo in readme

* flake.lock: Update

Flake lock file updates:

• Updated input 'flake-compat':
    'github:edolstra/flake-compat/35bb57c0c8d8b62bbfd284272c928ceb64ddbde9' (2023-01-17)
  → 'github:edolstra/flake-compat/0f9255e01c2351cc7d116c072cb317785dd33b33' (2023-10-04)
• Updated input 'nixpkgs':
    'github:NixOS/nixpkgs/d680ded26da5cf104dd2735a51e88d2d8f487b4d' (2023-08-19)
  → 'github:NixOS/nixpkgs/f99e5f03cc0aa231ab5950a15ed02afec45ed51a' (2023-10-09)
• Updated input 'utils':
    'github:numtide/flake-utils/919d646de7be200f3bf08cb76ae1f09402b6f9b4' (2023-07-11)
  → 'github:numtide/flake-utils/ff7b65b44d01cf9ba6a71320833626af21126384' (2023-09-12)

* doc: enable markdowm (myst) doc

* dev: StdSOMonitor: get_min_fitness -> get_best_fitness

* dev: standard the way to specific optimization direction

* fix: envpool with 64-bit return values, and add test cases

* dev: add support for brax >= 0.9.0

* dev: add envpool to test requirements

---------

Co-authored-by: Bill Huang <[email protected]>

Author: Zhenyu2Liang

docs/source/api/algorithms/mo/gde3 renamed to docs/source/api/algorithms/mo/ged3.rst

@@ -0,0 +1,6 @@
+=======
+LMOCSO
+=======
+
+.. autoclass:: evox.algorithms.LMOCSO
+    :members:

docs/source/api/algorithms/mo/lmocso.rst

@@ -8,5 +8,4 @@ Problems
     :maxdepth: 2
 
     numerical/index
-    neuroevolution/index
-    rl/index
+    neuroevolution/index

docs/source/api/problems/index.rst

@@ -3,6 +3,7 @@ Neuroevolution
 ==============
 
 .. toctree::
-    :maxdepth: 1
+    :maxdepth: 2
 
-    torchvision
+    reinforcement_learning/index
+    supervised_learning/index

docs/source/api/problems/neuroevolution/index.rst

@@ -2,5 +2,5 @@
 Brax-based Problem
 ==================
 
-.. autoclass:: evox.problems.neuroevolution.Brax
+.. autoclass:: evox.problems.neuroevolution.reinforcement_learning.Brax
     :members:

docs/source/api/problems/rl/brax.rst renamed to docs/source/api/problems/neuroevolution/reinforcement_learning/brax.rst

@@ -0,0 +1,6 @@
+========
+Env Pool
+========
+
+.. autoclass:: evox.problems.neuroevolution.reinforcement_learning.EnvPool
+    :members:

docs/source/api/problems/neuroevolution/reinforcement_learning/env_pool.rst

@@ -0,0 +1,6 @@
+===
+Gym
+===
+
+.. autoclass:: evox.problems.neuroevolution.reinforcement_learning.Gym
+    :members:

docs/source/api/problems/neuroevolution/reinforcement_learning/gym.rst

@@ -0,0 +1,10 @@
+======================
+Reinforcement Learning
+======================
+
+.. toctree::
+    :maxdepth: 1
+
+    brax
+    gym
+    env_pool

docs/source/api/problems/neuroevolution/reinforcement_learning/index.rst

@@ -0,0 +1,8 @@
+===================
+Supervised Learning
+===================
+
+.. toctree::
+    :maxdepth: 1
+
+    torchvision

docs/source/api/problems/neuroevolution/supervised_learning/index.rst

@@ -0,0 +1,6 @@
+===================
+Torchvision Dataset
+===================
+
+.. autoclass:: evox.problems.neuroevolution.supervised_learning.TorchvisionDataset
+    :members:

docs/source/api/problems/neuroevolution/supervised_learning/torchvision.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:

docs/source/api/problems/neuroevolution/torchvision.rst

@@ -14,3 +14,5 @@
 from .sra import SRA
 from .tdea import TDEA
 from .bce_ibea import BCEIBEA
+from .lmocso import LMOCSO
+

docs/source/api/problems/numerical/dtlz.rst

@@ -17,47 +17,36 @@
 from evox.operators import selection, mutation, crossover, non_dominated_sort
 
 
-@partial(jax.jit, static_argnums=[0, 1])
-def calculate_alpha(N, k):
-    alpha = jnp.zeros(N)
-
-    for i in range(1, k + 1):
-        num = jnp.prod((k - jnp.arange(1, i)) / (N - jnp.arange(1, i)))
-        alpha = alpha.at[i - 1].set(num / i)
-    return alpha
-
-
 @partial(jax.jit, static_argnums=[2, 3])
 def cal_hv(points, ref, k, n_sample, key):
     n, m = jnp.shape(points)
-    alpha = calculate_alpha(n, k)
+
+    # hit in alpha relevant partition
+    alpha = jnp.cumprod(
+        jnp.r_[1, (k - jnp.arange(1, n)) / (n - jnp.arange(1, n))]
+    ) / jnp.arange(1, n + 1)
 
     f_min = jnp.min(points, axis=0)
 
-    s = jax.random.uniform(key, shape=(n_sample, m), minval=f_min, maxval=ref)
+    samples = jax.random.uniform(key, shape=(n_sample, m), minval=f_min, maxval=ref)
 
-    pds = jnp.zeros((n, n_sample), dtype=bool)
+    # update hypervolume estimates
     ds = jnp.zeros((n_sample,))
-
-    def body_fun1(i, vals):
-        pds, ds = vals
-        x = jnp.sum((jnp.tile(points[i, :], (n_sample, 1)) - s) <= 0, axis=1) == m
-        pds = pds.at[i].set(jnp.where(x, True, pds[i]))
-        ds = jnp.where(x, ds + 1, ds)
-        return pds, ds
-
-    pds, ds = jax.lax.fori_loop(0, n, body_fun1, (pds, ds))
-    ds = ds - 1
-
-    f = jnp.zeros((n,))
-
-    def body_fun2(pd):
-        temp = jnp.where(pd, ds, -1).astype(int)
+    pds = jax.vmap(
+        lambda x: jnp.sum((jnp.tile(x, (n_sample, 1)) - samples) <= 0, axis=1) == m,
+        in_axes=0,
+        out_axes=0,
+    )(points)
+    ds = jnp.sum(jnp.where(pds, ds + 1, ds), axis=0)
+    ds = jnp.where(ds == 0, ds, ds - 1)
+
+    def cal_f(val):
+        temp = jnp.where(val, ds, -1).astype(int)
         value = jnp.where(temp != -1, alpha[temp], 0)
         value = jnp.sum(value)
         return value
-    
-    f = jax.vmap(body_fun2)(pds)
+
+    f = jax.vmap(cal_f, in_axes=0, out_axes=0)(pds)
     f = f * jnp.prod(ref - f_min) / n_sample
 
     return f

docs/source/api/problems/rl/gym.rst

@@ -37,6 +37,9 @@ class IBEA(Algorithm):
     """IBEA algorithm
 
     link: https://link.springer.com/chapter/10.1007/978-3-540-30217-9_84
+
+    Args:
+        kappa: fitness scaling factor. Default: 0.05
     """
 
     def __init__(
@@ -111,27 +114,31 @@ def _tell_normal(self, state, fitness):
         merged_pop = jnp.concatenate([state.population, state.next_generation], axis=0)
         merged_obj = jnp.concatenate([state.fitness, fitness], axis=0)
 
-        n = jnp.shape(merged_pop)[0]
         merged_fitness, I, C = cal_fitness(merged_obj, self.kappa)
 
-        next_ind = jnp.arange(n)
-        vals = (next_ind, merged_fitness)
-
-        def body_fun(i, vals):
-            next_ind, merged_fitness = vals
-            x = jnp.argmin(merged_fitness)
-            merged_fitness += jnp.exp(-I[x, :] / C[x] / self.kappa)
-            merged_fitness = merged_fitness.at[x].set(jnp.max(merged_fitness))
-            next_ind = next_ind.at[x].set(-1)
-            return (next_ind, merged_fitness)
-
-        next_ind, merged_fitness = jax.lax.fori_loop(0, self.pop_size, body_fun, vals)
-
-        ind = jnp.where(next_ind != -1, size=n, fill_value=-1)[0]
-        ind_n = ind[0 : self.pop_size]
-
-        survivor = merged_pop[ind_n]
-        survivor_fitness = merged_obj[ind_n]
+        # Different from the original paper, the selection here is directly through fitness.
+        next_ind = jnp.argsort(-merged_fitness)[0: self.pop_size]
+
+        # The following code is from the original paper's implementation
+        # and is kept for reference purposes but is not being used in this version.
+        # n = jnp.shape(merged_pop)[0]
+        # next_ind = jnp.arange(n)
+        # vals = (next_ind, merged_fitness)
+        # def body_fun(i, vals):
+        #     next_ind, merged_fitness = vals
+        #     x = jnp.argmin(merged_fitness)
+        #     merged_fitness += jnp.exp(-I[x, :] / C[x] / self.kappa)
+        #     merged_fitness = merged_fitness.at[x].set(jnp.max(merged_fitness))
+        #     next_ind = next_ind.at[x].set(-1)
+        #     return (next_ind, merged_fitness)
+        #
+        # next_ind, merged_fitness = jax.lax.fori_loop(0, self.pop_size, body_fun, vals)
+        #
+        # next_ind = jnp.where(next_ind != -1, size=n, fill_value=-1)[0]
+        # next_ind = next_ind[0: self.pop_size]
+
+        survivor = merged_pop[next_ind]
+        survivor_fitness = merged_obj[next_ind]
 
         state = state.update(population=survivor, fitness=survivor_fitness)