Merge pull request #141 from EMI-Group/dev3-lzy

Fix bugs of MOEA/D

Author: BillHuang2001

src/evox/algorithms/mo/eagmoead.py

@@ -23,8 +23,8 @@
     non_dominated_sort,
     crowding_distance,
 )
-from evox.operators.sampling import LatinHypercubeSampling
-from evox.utils import pairwise_euclidean_dist
+from evox.operators.sampling import LatinHypercubeSampling, UniformSampling
+from evox.utils import pairwise_euclidean_dist, AggregationFunction
 
 
 @partial(jax.jit, static_argnums=[1])
@@ -76,6 +76,7 @@ def __init__(
         if self.crossover is None:
             self.crossover = crossover.SimulatedBinary(type=2)
         self.sample = LatinHypercubeSampling(self.pop_size, self.n_objs)
+        self.aggregate_func = AggregationFunction("weighted_sum")
 
     def setup(self, key):
         key, subkey1, subkey2 = jax.random.split(key, 3)
@@ -160,10 +161,10 @@ def tell(self, state, fitness):
 
         def body_fun(i, vals):
             population, pop_obj = vals
-            g_old = jnp.sum(
-                pop_obj[B[offspring_loc[i], :]] * w[B[offspring_loc[i], :]], axis=1
+            g_old = self.aggregate_func(
+                pop_obj[B[offspring_loc[i], :]], w[B[offspring_loc[i], :]]
             )
-            g_new = w[B[offspring_loc[i], :]] @ jnp.transpose(offspring_obj[i])
+            g_new = self.aggregate_func(offspring_obj[i], w[B[offspring_loc[i], :]])
             idx = B[offspring_loc[i]]
             g_new = g_new[:, jnp.newaxis]
             g_old = g_old[:, jnp.newaxis]

src/evox/algorithms/mo/moead.py

@@ -3,10 +3,6 @@
 #
 # Title: MOEA/D: A Multiobjective Evolutionary Algorithm Based on Decomposition
 # Link: https://ieeexplore.ieee.org/document/4358754
-#
-# 2. This code has been inspired by PlatEMO.
-# More information about PlatEMO can be found at the following URL:
-# GitHub Link: https://github.com/BIMK/PlatEMO
 # --------------------------------------------------------------------------------------
 
 import math
@@ -17,8 +13,7 @@
 from evox import Algorithm, State, jit_class
 from evox.operators import crossover, mutation
 from evox.operators.sampling import UniformSampling
-from evox.utils import pairwise_euclidean_dist
-
+from evox.utils import pairwise_euclidean_dist, AggregationFunction
 
 @jit_class
 class MOEAD(Algorithm):
@@ -33,7 +28,7 @@ def __init__(
         ub,
         n_objs,
         pop_size,
-        type=1,
+        func_name='pbi',
         mutation_op=None,
         crossover_op=None,
     ):
@@ -42,8 +37,8 @@ def __init__(
         self.n_objs = n_objs
         self.dim = lb.shape[0]
         self.pop_size = pop_size
-        self.type = type
-        self.T = 0
+        self.func_name = func_name
+        self.n_neighbor = 0
 
         self.mutation = mutation_op
         self.crossover = crossover_op
@@ -53,54 +48,57 @@ def __init__(
         if self.crossover is None:
             self.crossover = crossover.SimulatedBinary(type=2)
         self.sample = UniformSampling(self.pop_size, self.n_objs)
+        self.aggregate_func = AggregationFunction(self.func_name)
 
     def setup(self, key):
         key, subkey1, subkey2 = jax.random.split(key, 3)
         w, _ = self.sample(subkey2)
         self.pop_size = w.shape[0]
-        self.T = int(math.ceil(self.pop_size / 10))
+        self.n_neighbor = int(math.ceil(self.pop_size / 10))
 
         population = (
             jax.random.uniform(subkey1, shape=(self.pop_size, self.dim))
             * (self.ub - self.lb)
             + self.lb
         )
 
-        B = pairwise_euclidean_dist(w, w)
-        B = jnp.argsort(B, axis=1)
-        B = B[:, : self.T]
+        neighbors = pairwise_euclidean_dist(w, w)
+        neighbors = jnp.argsort(neighbors, axis=1)
+        neighbors = neighbors[:, : self.n_neighbor]
+
         return State(
             population=population,
             fitness=jnp.zeros((self.pop_size, self.n_objs)),
             next_generation=population,
             weight_vector=w,
-            B=B,
-            Z=jnp.zeros(shape=self.n_objs),
-            parent=jnp.zeros((self.pop_size, self.T)).astype(int),
+            neighbors=neighbors,
+            z=jnp.zeros(shape=self.n_objs),
             key=key,
         )
 
     def init_ask(self, state):
         return state.population, state
 
     def init_tell(self, state, fitness):
-        Z = jnp.min(fitness, axis=0)
-        state = state.update(fitness=fitness, Z=Z)
+        z = jnp.min(fitness, axis=0)
+        state = state.update(fitness=fitness, z=z)
         return state
 
     def ask(self, state):
         key, subkey, sel_key, mut_key = jax.random.split(state.key, 4)
         parent = jax.random.permutation(
-            subkey, state.B, axis=1, independent=True
+            subkey, state.neighbors, axis=1, independent=True
         ).astype(int)
+
         population = state.population
         selected_p = jnp.r_[population[parent[:, 0]], population[parent[:, 1]]]
 
         crossovered = self.crossover(sel_key, selected_p)
         next_generation = self.mutation(mut_key, crossovered)
+        next_generation = jnp.clip(next_generation, self.lb, self.ub)
 
         return next_generation, state.update(
-            next_generation=next_generation, parent=parent, key=key
+            next_generation=next_generation, key=key
         )
 
     def tell(self, state, fitness):
@@ -109,84 +107,28 @@ def tell(self, state, fitness):
         offspring = state.next_generation
         obj = fitness
         w = state.weight_vector
-        Z = state.Z
-        parent = state.parent
-
-        out_vals = (population, pop_obj, Z)
-
-        def out_body(i, out_vals):
-            population, pop_obj, Z = out_vals
-            ind_p = parent[i]
-            ind_obj = obj[i]
-            Z = jnp.minimum(Z, obj[i])
-
-            if self.type == 1:
-                # PBI approach
-                norm_w = jnp.linalg.norm(w[ind_p], axis=1)
-                norm_p = jnp.linalg.norm(
-                    pop_obj[ind_p] - jnp.tile(Z, (self.T, 1)), axis=1
-                )
-                norm_o = jnp.linalg.norm(ind_obj - Z)
-                cos_p = (
-                    jnp.sum(
-                        (pop_obj[ind_p] - jnp.tile(Z, (self.T, 1))) * w[ind_p], axis=1
-                    )
-                    / norm_w
-                    / norm_p
-                )
-                cos_o = (
-                    jnp.sum(jnp.tile(ind_obj - Z, (self.T, 1)) * w[ind_p], axis=1)
-                    / norm_w
-                    / norm_o
-                )
-                g_old = norm_p * cos_p + 5 * norm_p * jnp.sqrt(1 - cos_p**2)
-                g_new = norm_o * cos_o + 5 * norm_o * jnp.sqrt(1 - cos_o**2)
-            if self.type == 2:
-                # Tchebycheff approach
-                g_old = jnp.max(
-                    jnp.abs(pop_obj[ind_p] - jnp.tile(Z, (self.T, 1))) * w[ind_p],
-                    axis=1,
-                )
-                g_new = jnp.max(
-                    jnp.tile(jnp.abs(ind_obj - Z), (self.T, 1)) * w[ind_p], axis=1
-                )
-            if self.type == 3:
-                # Tchebycheff approach with normalization
-                z_max = jnp.max(pop_obj, axis=0)
-                g_old = jnp.max(
-                    jnp.abs(pop_obj[ind_p] - jnp.tile(Z, (self.T, 1)))
-                    / jnp.tile(z_max - Z, (self.T, 1))
-                    * w[ind_p],
-                    axis=1,
-                )
-                g_new = jnp.max(
-                    jnp.tile(jnp.abs(ind_obj - Z), (self.T, 1))
-                    / jnp.tile(z_max - Z, (self.T, 1))
-                    * w[ind_p],
-                    axis=1,
-                )
-            if self.type == 4:
-                # Modified Tchebycheff approach
-                g_old = jnp.max(
-                    jnp.abs(pop_obj[ind_p] - jnp.tile(Z, (self.T, 1))) / w[ind_p],
-                    axis=1,
-                )
-                g_new = jnp.max(
-                    jnp.tile(jnp.abs(ind_obj - Z), (self.T, 1)) / w[ind_p], axis=1
-                )
-
-            g_new = g_new[:, jnp.newaxis]
-            g_old = g_old[:, jnp.newaxis]
-            population = population.at[ind_p].set(
-                jnp.where(g_old >= g_new, offspring[ind_p], population[ind_p])
-            )
-            pop_obj = pop_obj.at[ind_p].set(
-                jnp.where(g_old >= g_new, obj[ind_p], pop_obj[ind_p])
-            )
-
-            return (population, pop_obj, Z)
-
-        population, pop_obj, Z = jax.lax.fori_loop(0, self.pop_size, out_body, out_vals)
-
-        state = state.update(population=population, fitness=pop_obj, Z=Z)
-        return state
+
+        z = jnp.minimum(state.z, jnp.min(obj, axis=0))
+        z_max = jnp.max(pop_obj, axis=0)
+        neighbors = state.neighbors
+
+        def scan_body(carry, x):
+            population, pop_obj = carry
+            off_pop, off_obj, indices = x
+
+            f_old = self.aggregate_func(pop_obj[indices], w[indices], z, z_max)
+            f_new = self.aggregate_func(off_obj[jnp.newaxis, :], w[indices], z, z_max)
+
+            update_condition = (f_old > f_new)[:, jnp.newaxis]
+            updated_population = population.at[indices].set(
+                jnp.where(update_condition, jnp.tile(off_pop, (jnp.shape(indices)[0], 1)), population[indices]))
+            updated_pop_obj = pop_obj.at[indices].set(
+                jnp.where(update_condition, jnp.tile(off_obj, (jnp.shape(indices)[0], 1)), pop_obj[indices]))
+
+            return (updated_population, updated_pop_obj), None
+
+        (population, pop_obj), _ = jax.lax.scan(scan_body, (population, pop_obj), (offspring, obj, neighbors))
+
+
+        state = state.update(population=population, fitness=pop_obj, z=z)
+        return state

src/evox/algorithms/mo/moeaddra.py

@@ -17,7 +17,7 @@
 from evox import Algorithm, State, jit_class
 from evox.operators import crossover, mutation, selection
 from evox.operators.sampling import LatinHypercubeSampling
-from evox.utils import pairwise_euclidean_dist
+from evox.utils import pairwise_euclidean_dist, AggregationFunction
 
 
 @jit_class
@@ -55,6 +55,7 @@ def __init__(
         if self.crossover is None:
             self.crossover = crossover.DifferentialEvolve()
         self.sample = LatinHypercubeSampling(self.pop_size, self.n_objs)
+        self.aggergate_func = AggregationFunction("tchebycheff")
 
     def setup(self, key):
         key, subkey1, subkey2 = jax.random.split(key, 3)
@@ -155,10 +156,8 @@ def out_body(i, out_vals):
             ind_obj = off_obj[i]
             Z = jnp.minimum(Z, ind_obj)
 
-            g_old = jnp.max(
-                jnp.abs(pop_obj[p] - jnp.tile(Z, (len(p), 1))) * w[p], axis=1
-            )
-            g_new = jnp.max(jnp.abs(jnp.tile(ind_obj - Z, (len(p), 1))) * w[p], axis=1)
+            g_old = self.aggergate_func(pop_obj[p], w[p], Z)
+            g_new = self.aggergate_func(ind_obj[jnp.newaxis, :], w[p], Z)
 
             g_new = g_new[:, jnp.newaxis]
             g_old = g_old[:, jnp.newaxis]

src/evox/utils/common.py

@@ -266,3 +266,52 @@ def frames2gif(frames, save_path, duration=0.1):
             writer.append_data(formatted_image)
 
     print("Gif saved to: ", save_path)
+
+
+@jit_class
+class AggregationFunction:
+    """
+    Aggregation function: PBI approach, Tchebycheff approach, Tchebycheff approach with normalization,
+    modified Tchebycheff approach, weighted sum approach.
+
+    Args:
+        function_name (str): name of the aggregation function.
+    """
+
+    def __init__(self, function_name):
+        if function_name == "pbi":
+            self.function = self.pbi
+        elif function_name == "tchebycheff":
+            self.function = self.tchebycheff
+        elif function_name == "tchebycheff_norm":
+            self.function = self.tchebycheff_norm
+        elif function_name == "modified_tchebycheff":
+            self.function = self.modified_tchebycheff
+        elif function_name == "weighted_sum":
+            self.function = self.weighted_sum
+        else:
+            raise ValueError("Unsupported function")
+
+    def pbi(self, f, w, z, *args):
+        norm_w = jnp.linalg.norm(w, axis=1)
+        f = f - z
+        d1 = jnp.sum(f * w, axis=1) / norm_w
+        d2 = jnp.linalg.norm(
+            f - (d1[:, jnp.newaxis] * w / norm_w[:, jnp.newaxis]), axis=1
+        )
+        return d1 + 5 * d2
+
+    def tchebycheff(self, f, w, z, *args):
+        return jnp.max(jnp.abs(f - z) * w, axis=1)
+
+    def tchebycheff_norm(self, f, w, z, z_max, *args):
+        return jnp.max(jnp.abs(f - z) / (z_max - z) * w, axis=1)
+
+    def modified_tchebycheff(self, f, w, z, *args):
+        return jnp.max(jnp.abs(f - z) / w, axis=1)
+
+    def weighted_sum(self, f, w, *args):
+        return jnp.sum(f * w, axis=1)
+
+    def __call__(self, *args, **kwargs):
+        return self.function(*args, **kwargs)

tests/test_multi_objective_algorithms.py

@@ -46,7 +46,6 @@ def test_moead():
         ub=jnp.full(shape=(N,), fill_value=UB),
         n_objs=M,
         pop_size=POP_SIZE,
-        type=1,
     )
     run_moea(algorithm)