Add GradientEmitter to support OMG-MEGA and OG-MAP-Elites

HISTORY.md

@@ -6,6 +6,7 @@
 
 #### API
 
+- Add GradientEmitter to support OMG-MEGA and OG-MAP-Elites (#348)
 - Drop Python 3.7 support and upgrade dependencies ({pr}`350`)
 - Add visualization of QDax repertoires ({pr}`353`)
 - Improve cvt_archive_heatmap flexibility ({pr}`354`)

examples/sphere.py

@@ -35,6 +35,9 @@
 - `cma_me_mixed`: GridArchive with EvolutionStrategyEmitter, where half (7) of
   the emitter are using TwoStageRandomDirectionRanker and half (8) are
   TwoStageImprovementRanker.
+- `og_map_elites`: GridArchive with GradientOperatorEmitter, does not use
+  measure gradients.
+- `omg_mega`: GridArchive with GradientOperatorEmitter, uses measure gradients.
 - `cma_mega`: GridArchive with GradientArborescenceEmitter.
 - `cma_mega_adam`: GridArchive with GradientArborescenceEmitter using Adam
   Optimizer.
@@ -81,12 +84,15 @@
 import numpy as np
 import tqdm
 
+# yapf: disable
 from ribs.archives import CVTArchive, GridArchive
 from ribs.emitters import (EvolutionStrategyEmitter, GaussianEmitter,
-                           GradientArborescenceEmitter, IsoLineEmitter)
+                           GradientArborescenceEmitter,
+                           GradientOperatorEmitter, IsoLineEmitter)
 from ribs.schedulers import BanditScheduler, Scheduler
 from ribs.visualize import cvt_archive_heatmap, grid_archive_heatmap
 
+# yapf: enable
 CONFIG = {
     "map_elites": {
         "dim": 20,
@@ -417,6 +423,62 @@
             "kwargs": {}
         }
     },
+    "og_map_elites": {
+        "dim": 1_000,
+        "iters": 10_000,
+        "archive_dims": (100, 100),
+        "use_result_archive": False,
+        "is_dqd": True,
+        "batch_size": 36 // 2,
+        "archive": {
+            "class": GridArchive,
+            "kwargs": {
+                "threshold_min": -np.inf
+            }
+        },
+        "emitters": [{
+            "class": GradientOperatorEmitter,
+            "kwargs": {
+                "sigma0": 0.5,
+                "sigma_g": 0.5,
+                "measure_gradients": False,
+                "normalize_grad": False,
+            },
+            "num_emitters": 1
+        }],
+        "scheduler": {
+            "class": Scheduler,
+            "kwargs": {}
+        }
+    },
+    "omg_mega": {
+        "dim": 1_000,
+        "iters": 10_000,
+        "archive_dims": (100, 100),
+        "use_result_archive": False,
+        "is_dqd": True,
+        "batch_size": 36 // 2,
+        "archive": {
+            "class": GridArchive,
+            "kwargs": {
+                "threshold_min": -np.inf
+            }
+        },
+        "emitters": [{
+            "class": GradientOperatorEmitter,
+            "kwargs": {
+                "sigma0": 0.0,
+                "sigma_g": 10.0,
+                "measure_gradients": True,
+                "normalize_grad": True,
+            },
+            "num_emitters": 1
+        }],
+        "scheduler": {
+            "class": Scheduler,
+            "kwargs": {}
+        }
+    },
     "cma_mega": {
         "dim": 1_000,
         "iters": 10_000,

ribs/emitters/__init__.py

@@ -26,10 +26,12 @@
 from ribs.emitters._gaussian_emitter import GaussianEmitter
 from ribs.emitters._gradient_arborescence_emitter import \
     GradientArborescenceEmitter
+from ribs.emitters._gradient_operator_emitter import GradientOperatorEmitter
 from ribs.emitters._iso_line_emitter import IsoLineEmitter
 
 __all__ = [
     "EvolutionStrategyEmitter",
+    "GradientOperatorEmitter",
     "GradientArborescenceEmitter",
     "GaussianEmitter",
     "IsoLineEmitter",