Feature/nsga optimizer

core/optimizers/base.py

@@ -139,7 +139,7 @@ def run(self, world: Optional[World] = None) -> None:
         >>>     # Optional: update or overwrite own context ---
         >>>     ctx.payload.value += 1.0
         >>>     world.update_context(ctx)
-        
+
         """
         raise NotImplementedError
 

core/optimizers/config.py

@@ -93,7 +93,7 @@ def build_optimizer(self) -> Optimizer:
         # TODO Executer to enforce guardrails
         # if world is None:
         #     raise ValueError("Optimizer requires a World instance")
-        
+
         cfg = self.copy()
 
         # TODO : world is passed to optimizer, but also stored in config. must only have one source of truth

core/optimizers/nsga/optimizer.py

@@ -0,0 +1,104 @@
+import numpy as np
+from pymoo.util.nds.non_dominated_sorting import NonDominatedSorting
+from typing import Optional
+
+class NGSAOptimizer:
+    """
+    Simple NSGA-II style optimizer for multi-objective fitness.
+    Tracks population history and Pareto front.
+    TODO: need to modify regulator_env to return matrix instead of list with fitness values
+    """
+
+    def __init__(self, config):
+        # --- hyperparameters ---
+        self.dimension: int = config.dimension
+        self._batch_capacity: int = getattr(config, "batch_capacity", 32)
+        self.rng = np.random.default_rng(getattr(config, "seed", None))
+
+        # --- runtime state ---
+        self.generation: int = 0
+        self.best_fitness: Optional[np.ndarray] = None        # [sustainability, market_value]
+        self.best_candidate: Optional[np.ndarray] = None      # parameters of best candidate
+        self.population_history: list[tuple[np.ndarray, np.ndarray]] = []
+
+        # Optional convergence tracking
+        self.no_improve_steps: int = 0
+        self.convergence_patience: int = getattr(config, "convergence_patience", 10)
+        self.convergence_eps: float = getattr(config, "convergence_eps", 1e-4)
+        self.converged_once: bool = False
+
+        # Placeholder for environment
+        self.env = None
+
+    def environment(self, env, env_config=None):
+        """
+        Set the environment to use.
+        """
+        self.env = env(**(env_config or {}))
+        return self
+
+    @property
+    def batch_capacity(self) -> int:
+        return self._batch_capacity
+
+    @batch_capacity.setter
+    def batch_capacity(self, value: int) -> None:
+        if value <= 0:
+            raise ValueError("batch_capacity must be positive")
+        self._batch_capacity = value
+
+    def _sample_population(self) -> np.ndarray:
+        """
+        Sample a population uniformly in [0,1]^dimension
+        """
+        return self.rng.random((self._batch_capacity, self.dimension), dtype=np.float32)
+
+    def run(self) -> dict:
+        if self.env is None:
+            raise RuntimeError("Optimizer requires an environment")
+
+        # Sample a population
+        population = self._sample_population()  # (pop_size, dimension)
+
+        # Step environment to get multi-objective fitness
+        # Expected shape: (pop_size, 2) -> [sustainability, market_value]
+        _, fitness_matrix, _, _, _ = self.env.step(population)
+        fitness_matrix = np.asarray(fitness_matrix, dtype=np.float32)
+
+        if fitness_matrix.size == 0:
+            return {
+                "converged": False,
+                "best_fitness": None,
+                "best_trajectory": None,
+                "population_history": self.population_history,
+            }
+
+        # Store population and fitness for history
+        self.population_history.append((population.copy(), fitness_matrix.copy()))
+
+        # Compute Pareto front
+        nds = NonDominatedSorting()
+        front_indices = nds.do(fitness_matrix, only_non_dominated_front=True)
+        pareto_solutions = population[front_indices]
+        pareto_fitness = fitness_matrix[front_indices]
+
+        # Pick representative best candidate (sum of objectives)
+        best_idx = np.argmax(fitness_matrix.sum(axis=1))
+        self.best_candidate = population[best_idx].copy()
+        self.best_fitness = fitness_matrix[best_idx]
+
+        # Get best trajectory if available
+        best_trajectory = None
+        if hasattr(self.env, "trajectories") and self.best_candidate is not None:
+            best_trajectory = self.env.trajectories.get(best_idx)
+
+        # Optional convergence logic
+        # (implement if desired based on improvements over generations)
+        converged = False
+
+        return {
+            "converged": converged,
+            "best_fitness": self.best_fitness,
+            "best_trajectory": best_trajectory,
+            "population_history": self.population_history,
+        }

core/world/base.py

@@ -71,11 +71,14 @@ def register_optimizer(self, opt: Optimizer) -> OptimizerID:
 
         This initializes an empty context set for the optimizer.
         """
-        if opt.opt_id is None:
-            opt.opt_id = generate_uuid(registry=self._opt_ctx_map.keys())
-        if opt.opt_id not in self._opt_ctx_map:
-            self._opt_ctx_map[opt.opt_id] = set()
-        return opt.opt_id
+        return self._set_new_opt_id(opt_id=opt.opt_id)
+
+    def _set_new_opt_id(self, opt_id: OptimizerID) -> OptimizerID:
+        if opt_id is None:
+            opt_id = generate_uuid(registry=self._opt_ctx_map.keys())
+        if opt_id not in self._opt_ctx_map:
+            self._opt_ctx_map[opt_id] = set()
+        return opt_id
 
     def set_new_context(self, ctx: Context, singleton: bool = False) -> ContextID:
         """
@@ -92,17 +95,17 @@ def set_new_context(self, ctx: Context, singleton: bool = False) -> ContextID:
         if singleton:
             self._validate_ctx_schema_exists(type(ctx.payload))
 
-        if ctx.id is not None :
+        if ctx.id is not None:
             if ctx.id in self._contexts:
                 raise ValueError(f"ContextID '{ctx.id}' already exists")
         else:
             ctx.id = generate_uuid(registry=self._contexts.keys())
 
         self._contexts[ctx.id] = ctx
-        
+
         if ctx.opt_id is not None:
             if ctx.opt_id not in self._opt_ctx_map:
-                self.set_new_optimizer(ctx.opt_id)
+                self._set_new_opt_id(ctx.opt_id)
             self._opt_ctx_map[ctx.opt_id].add(ctx.id)
 
         return ctx.id