From 6518d2735326c84b11e597bd7669c0783f82cbd3 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 24 Aug 2022 00:52:34 -0700
Subject: [PATCH 01/53] revive old implementation of add_single (w/o numba)

---
 .editorconfig                  |  2 +-
 ribs/archives/_archive_base.py | 78 +++++++++++++++++++++++++++++-----
 2 files changed, 68 insertions(+), 12 deletions(-)

diff --git a/.editorconfig b/.editorconfig
index 2b5338a8c..4794c6bab 100644
--- a/.editorconfig
+++ b/.editorconfig
@@ -4,7 +4,7 @@ root = true
 
 [*]
 indent_style = space
-indent_size = 2
+indent_size = 4
 trim_trailing_whitespace = true
 insert_final_newline = true
 charset = utf-8
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 446bf1960..4cc8b6af8 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -134,6 +134,8 @@ def __init__(self,
                  solution_dim,
                  cells,
                  measure_dim,
+                 learning_rate=1.0,
+                 threshold_min=0.0,
                  seed=None,
                  dtype=np.float64):
 
@@ -155,6 +157,11 @@ def __init__(self,
                                       dtype=self.dtype)
         self._metadata_arr = np.empty(self._cells, dtype=object)
 
+        # For CMA-MAE
+        self._learning_rate = learning_rate
+        self._threshold_min = threshold_min
+        self._threshold_arr = np.empty(self._cells, dtype=self.dtype)
+
         self._stats = None
         self._stats_reset()
 
@@ -434,22 +441,22 @@ def add(self,
         index_batch = self.index_of(measures_batch)
 
         # Copy old objectives since we will be modifying the objectives storage.
-        old_objective_batch = np.copy(self._objective_arr[index_batch])
+        old_objective_arr = np.copy(self._objective_arr[index_batch])
 
         # Compute the statuses -- these are all boolean arrays of length
         # batch_size.
         already_occupied = self._occupied_arr[index_batch]
         is_new = ~already_occupied
         improve_existing = (objective_batch >
-                            old_objective_batch) & already_occupied
+                            old_objective_arr) & already_occupied
         status_batch = np.zeros(batch_size, dtype=np.int32)
         status_batch[is_new] = 2
         status_batch[improve_existing] = 1
 
         # Since we set the new solutions in old_objective_batch to have
         # value 0.0, the values for new solutions are correct here.
-        old_objective_batch[is_new] = 0.0
-        value_batch = objective_batch - old_objective_batch
+        old_objective_arr[is_new] = 0.0
+        value_batch = objective_batch - old_objective_arr
 
         ## Step 3: Insert solutions into archive. ##
 
@@ -466,7 +473,7 @@ def add(self,
         measures_batch_can = measures_batch[can_insert]
         index_batch_can = index_batch[can_insert]
         metadata_batch_can = metadata_batch[can_insert]
-        old_objective_batch_can = old_objective_batch[can_insert]
+        old_objective_batch_can = old_objective_arr[can_insert]
 
         # Retrieve indices of solutions that should be inserted into the
         # archive. Currently, multiple solutions may be inserted at each
@@ -569,18 +576,67 @@ def add_single(self, solution, objective, measures, metadata=None):
             the add operation. Refer to :meth:`add` for the meaning of the
             status and value.
         """
+        self._state["add"] += 1
+
         solution = np.asarray(solution)
         objective = np.asarray(objective, dtype=self.dtype)  # 0-dim array.
         measures = np.asarray(measures)
+        index = self.index_of_single(measures)
 
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
 
-        status_batch, value_batch = self.add(solution[None],
-                                             np.array([objective]),
-                                             measures[None],
-                                             np.array([metadata], dtype=object))
-        return (status_batch[0], value_batch[0])
+        old_objective = self._objective_arr[index]
+
+        was_occupied = self._occupied_arr[index]
+        if not was_occupied or self._objective_arr[index] < objective:
+            # Set this index to "occupied" -- important that we do this before
+            # inserting the solution.
+            self._occupied_arr[index] = True
+
+            # Insert into the archive.
+            self._objective_arr[index] = objective
+            self._measures_arr[index] = measures
+            self._solution_arr[index] = solution
+            self._metadata[index] = metadata
+
+            if was_occupied:
+                status = 1  # IMPROVE_EXISTING
+                value = objective - old_objective
+            else:
+                # Tracks a new occupied index.
+                self._occupied_indices[self._num_occupied] = index
+                self._num_occupied += 1
+                status = 2  # NEW
+                value = objective_value
+
+            # Update archive stats.
+            new_qd_score = self._stats.qd_score + value
+
+            if self._stats.obj_max is None or objective > self._stats.obj_max:
+                new_obj_max = objective
+                self._best_elite = Elite(
+                    readonly(solution),
+                    objective,
+                    readonly(measures),
+                    index,
+                    metadata,
+                )
+            else:
+                new_obj_max = self._stats.obj_max
+
+            self._stats = ArchiveStats(
+                num_elites=len(self),
+                coverage=self.dtype(len(self) / self.cells),
+                qd_score=new_qd_score,
+                obj_max=new_obj_max,
+                obj_mean=new_qd_score / self.dtype(len(self)),
+            )
+        else:
+            status = 0  # NOT_ADDED
+            value = objective_value - old_objective
+
+        return status, value
 
     def elites_with_measures(self, measures_batch):
         """Retrieves the elites with measures in the same cells as the measures
@@ -791,7 +847,7 @@ def as_pandas(self, include_solutions=True, include_metadata=False):
                 be representable in a CSV.
         Returns:
             ArchiveDataFrame: See above.
-        """ # pylint: disable = line-too-long
+        """  # pylint: disable = line-too-long
         data = OrderedDict()
         indices = self._occupied_indices[:self._num_occupied]
 

From bf492ff4fcbd840fae1e67bb3029ec786e363944 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 24 Aug 2022 02:14:08 -0700
Subject: [PATCH 02/53] add threshold

---
 ribs/archives/_archive_base.py | 34 ++++++++++++++++++++++++----------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 4cc8b6af8..b7ede2171 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -101,6 +101,9 @@ class ArchiveBase(ABC):  # pylint: disable = too-many-instance-attributes
         cells (int): Number of cells in the archive. This is used to create the
             numpy arrays described above for storing archive info.
         measure_dim (int): The dimension of the measure space.
+        learning_rate (float): The learning rate of the archive. Described in
+            `Fontaine 2022 <https://arxiv.org/abs/2205.10752>`_.
+        threshold_min (float): The default threshold value for all the cells.
         seed (int): Value to seed the random number generator. Set to None to
             avoid a fixed seed.
         dtype (str or data-type): Data type of the solutions, objectives,
@@ -159,8 +162,10 @@ def __init__(self,
 
         # For CMA-MAE
         self._learning_rate = learning_rate
-        self._threshold_min = threshold_min
-        self._threshold_arr = np.empty(self._cells, dtype=self.dtype)
+        # TODO use np.empty
+        self._threshold_arr = np.full(self._cells,
+                                      threshold_min,
+                                      dtype=self.dtype)
 
         self._stats = None
         self._stats_reset()
@@ -580,38 +585,50 @@ def add_single(self, solution, objective, measures, metadata=None):
 
         solution = np.asarray(solution)
         objective = np.asarray(objective, dtype=self.dtype)  # 0-dim array.
+        print(objective.dtype)
         measures = np.asarray(measures)
         index = self.index_of_single(measures)
 
+        print(objective.dtype)
+
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
 
         old_objective = self._objective_arr[index]
+        # Note that when learning_rate = 1.0, this is equivalent to old_objective
+        old_threshold = self._threshold_arr[index]
 
         was_occupied = self._occupied_arr[index]
-        if not was_occupied or self._objective_arr[index] < objective:
+        status = 0  # NOT_ADDED
+        if not was_occupied or old_threshold < objective:
             # Set this index to "occupied" -- important that we do this before
             # inserting the solution.
             self._occupied_arr[index] = True
 
+            # Update the threshold for this cell by the learning rate if the threshold will increase.
+            self._threshold_arr[index] = old_threshold * \
+                (1.0 - self._learning_rate) + objective * self._learning_rate
+
             # Insert into the archive.
             self._objective_arr[index] = objective
             self._measures_arr[index] = measures
             self._solution_arr[index] = solution
-            self._metadata[index] = metadata
+            self._metadata_arr[index] = metadata
 
             if was_occupied:
                 status = 1  # IMPROVE_EXISTING
-                value = objective - old_objective
             else:
                 # Tracks a new occupied index.
                 self._occupied_indices[self._num_occupied] = index
                 self._num_occupied += 1
                 status = 2  # NEW
-                value = objective_value
 
+        value = objective - old_threshold
+        if status:
             # Update archive stats.
-            new_qd_score = self._stats.qd_score + value
+            if status == 2:
+                old_objective = 0
+            new_qd_score = self._stats.qd_score + (objective - old_objective)
 
             if self._stats.obj_max is None or objective > self._stats.obj_max:
                 new_obj_max = objective
@@ -632,9 +649,6 @@ def add_single(self, solution, objective, measures, metadata=None):
                 obj_max=new_obj_max,
                 obj_mean=new_qd_score / self.dtype(len(self)),
             )
-        else:
-            status = 0  # NOT_ADDED
-            value = objective_value - old_objective
 
         return status, value
 

From 010edbe9c3f93393a95ee388c6aefaab498c8880 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 24 Aug 2022 02:32:38 -0700
Subject: [PATCH 03/53] fix dtype issue

---
 ribs/archives/_archive_base.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index b7ede2171..a37a89b6b 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -584,13 +584,10 @@ def add_single(self, solution, objective, measures, metadata=None):
         self._state["add"] += 1
 
         solution = np.asarray(solution)
-        objective = np.asarray(objective, dtype=self.dtype)  # 0-dim array.
-        print(objective.dtype)
+        objective = self.dtype(objective)
         measures = np.asarray(measures)
         index = self.index_of_single(measures)
 
-        print(objective.dtype)
-
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
 
@@ -627,8 +624,11 @@ def add_single(self, solution, objective, measures, metadata=None):
         if status:
             # Update archive stats.
             if status == 2:
-                old_objective = 0
+                old_objective = self.dtype(0.0)
             new_qd_score = self._stats.qd_score + (objective - old_objective)
+            print(type(objective))
+            print(type(objective - old_objective))
+            print(type(self._stats.qd_score + objective - old_objective))
 
             if self._stats.obj_max is None or objective > self._stats.obj_max:
                 new_obj_max = objective

From 0ea8db85bd831a5e18a4680254a4d57404eeade9 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Thu, 25 Aug 2022 16:30:15 -0700
Subject: [PATCH 04/53] remove print

---
 ribs/archives/_archive_base.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index a37a89b6b..0c1368f78 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -626,9 +626,6 @@ def add_single(self, solution, objective, measures, metadata=None):
             if status == 2:
                 old_objective = self.dtype(0.0)
             new_qd_score = self._stats.qd_score + (objective - old_objective)
-            print(type(objective))
-            print(type(objective - old_objective))
-            print(type(self._stats.qd_score + objective - old_objective))
 
             if self._stats.obj_max is None or objective > self._stats.obj_max:
                 new_obj_max = objective

From ee663ac6d47b0ed9a15a50347a18ca6b0eea5f90 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Thu, 25 Aug 2022 18:33:45 -0700
Subject: [PATCH 05/53] implement threshold for batch_addition

---
 ribs/archives/_archive_base.py | 45 +++++++++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 6 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 0c1368f78..306fbe230 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -263,6 +263,30 @@ def _stats_reset(self):
         self._stats = ArchiveStats(0, self.dtype(0.0), self.dtype(0.0), None,
                                    None)
 
+    def _sum_geometric_series(self, a, r, n):
+        """Compute the sum of a geometric series."""
+        if r == 1.0:
+            return a
+        return (a * (1 - pow(r, n))) / (1 - r)
+
+    def _compute_new_thresholds(self, old_threshold_batch,
+                                objective_batch_insert):
+        """Update thresholds.
+
+            Args:
+                old_threshold_batch (np.ndarray): The threshold of the cells before updating.
+                objective_batch_insert (np.ndarray): The objective values of
+                    the solution that is inserted into the archive.
+            Returns:
+                A batch of new thresholds.
+        """
+        k = len(objective_batch_insert)
+        geometric_sum = self._sum_geometric_series(1, 1.0 - self._learning_rate,
+                                                   k)
+        return (self._learning_rate * np.sum(objective_batch_insert) *
+                geometric_sum / k) + (old_threshold_batch *
+                                      (1.0 - self._learning_rate)**k)
+
     def clear(self):
         """Removes all elites from the archive.
 
@@ -447,21 +471,22 @@ def add(self,
 
         # Copy old objectives since we will be modifying the objectives storage.
         old_objective_arr = np.copy(self._objective_arr[index_batch])
+        old_threshold_arr = np.copy(self._threshold_arr[index_batch])
 
         # Compute the statuses -- these are all boolean arrays of length
         # batch_size.
         already_occupied = self._occupied_arr[index_batch]
         is_new = ~already_occupied
         improve_existing = (objective_batch >
-                            old_objective_arr) & already_occupied
+                            old_threshold_arr) & already_occupied
         status_batch = np.zeros(batch_size, dtype=np.int32)
         status_batch[is_new] = 2
         status_batch[improve_existing] = 1
 
         # Since we set the new solutions in old_objective_batch to have
         # value 0.0, the values for new solutions are correct here.
-        old_objective_arr[is_new] = 0.0
-        value_batch = objective_batch - old_objective_arr
+        old_threshold_arr[is_new] = 0.0
+        value_batch = objective_batch - old_threshold_arr
 
         ## Step 3: Insert solutions into archive. ##
 
@@ -479,6 +504,7 @@ def add(self,
         index_batch_can = index_batch[can_insert]
         metadata_batch_can = metadata_batch[can_insert]
         old_objective_batch_can = old_objective_arr[can_insert]
+        old_threshold_batch_can = old_threshold_arr[can_insert]
 
         # Retrieve indices of solutions that should be inserted into the
         # archive. Currently, multiple solutions may be inserted at each
@@ -508,6 +534,11 @@ def add(self,
         index_batch_insert = index_batch_can[should_insert]
         metadata_batch_insert = metadata_batch_can[should_insert]
         old_objective_batch_insert = old_objective_batch_can[should_insert]
+        old_threshold_batch_insert = old_threshold_batch_can[should_insert]
+
+        # Update the thresholds.
+        self._threshold_arr[index_batch_insert] = self._compute_new_thresholds(
+            old_threshold_batch_insert, objective_batch_insert)
 
         # Set archive storage.
         self._objective_arr[index_batch_insert] = objective_batch_insert
@@ -602,9 +633,11 @@ def add_single(self, solution, objective, measures, metadata=None):
             # inserting the solution.
             self._occupied_arr[index] = True
 
-            # Update the threshold for this cell by the learning rate if the threshold will increase.
-            self._threshold_arr[index] = old_threshold * \
-                (1.0 - self._learning_rate) + objective * self._learning_rate
+            # Update the threshold.
+            # self._threshold_arr[index] = old_threshold * \
+            #     (1.0 - self._learning_rate) + objective * self._learning_rate
+            self._threshold_arr[index] = self._compute_new_thresholds(
+                np.array([old_threshold]), np.array([objective]))
 
             # Insert into the archive.
             self._objective_arr[index] = objective

From 6ccb6b2bc0fcb73b54159f995aad3a11fb004b90 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Thu, 25 Aug 2022 22:03:18 -0700
Subject: [PATCH 06/53] add cma_mae to sphere.py

---
 examples/sphere.py             | 53 ++++++++++++++++++++++++++++++----
 ribs/archives/_grid_archive.py |  7 +++++
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index 2fc499565..1b3345185 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -34,14 +34,19 @@
 - `cma_mega`: GridArchive with GradientAborescenceEmitter.
 - `cma_mega_adam`: GridArchive with GradientAborescenceEmitter using Adam
   Optimizer.
-
-Note: the settings for `cma_mega` and `cma_mega_adam` are consistent with the
-paper (`Fontaine 2021 <https://arxiv.org/abs/2106.03894>`_) in which these
-algorithms are proposed.
+- `cma_mae`: GridArchive (learning_rate = 0.01) with EvolutionStrategyEmitter.
+- `cma_maega`: GridArchive (learning_rate = 0.01) with GradientAborescenceEmitter.
 
 All algorithms use 15 emitters, each with a batch size of 37. Each one runs for
 4500 iterations for a total of 15 * 37 * 4500 ~= 2.5M evaluations.
 
+Exceptions:
+    - `cma_mega` and `cma_mega_adam` uses only one emitter and runs for 10,000
+      iterations. This is to be consistent with the paper (`Fontaine 2021
+      <https://arxiv.org/abs/2106.03894>`_) in which these algorithms were
+      proposed.
+    - `cma_mae` and `cma_maega`
+
 Note that the CVTArchive in this example uses 10,000 cells, as opposed to the
 250,000 (500x500) in the GridArchive, so it is not fair to directly compare
 `cvt_map_elites` and `line_cvt_map_elites` to the other algorithms. However, the
@@ -180,6 +185,15 @@ def create_scheduler(algorithm, dim, seed):
                               dims=(100, 100),
                               ranges=bounds,
                               seed=seed)
+    elif algorithm in ["cma_mae", "cma_maega"]:
+        # Note that the archive is smaller for these algorithms. This is to be
+        # consistent with Fontaine 2022 <https://arxiv.org/abs/2205.10752>.
+        archive = GridArchive(solution_dim=dim,
+                              dims=(100, 100),
+                              ranges=bounds,
+                              learning_rate=0.01,
+                              threshold_min=0,
+                              seed=seed)
     else:
         raise ValueError(f"Algorithm `{algorithm}` is not recognized")
 
@@ -278,6 +292,31 @@ def create_scheduler(algorithm, dim, seed):
                 batch_size=batch_size - 1,  # 1 solution is returned by ask_dqd
                 seed=emitter_seeds[0])
         ]
+    elif algorithm == "cma_mae":
+        emitters = [
+            EvolutionStrategyEmitter(
+                archive=archive,
+                x0=initial_sol,
+                sigma0=0.5,
+                ranker="2imp",
+                selection_rule="mu",
+                restart_rule="basic",
+                batch_size=batch_size,
+                seed=s,
+            ) for s in emitter_seeds
+        ]
+    elif algorithm in ["cma_maega"]:
+        emitters = [
+            GradientAborescenceEmitter(archive,
+                                       initial_sol,
+                                       sigma0=10.0,
+                                       step_size=1.0,
+                                       grad_opt="gradient_ascent",
+                                       restart_rule="basic",
+                                       bounds=None,
+                                       batch_size=batch_size,
+                                       seed=s) for s in emitter_seeds
+        ]
     return Scheduler(archive, emitters)
 
 
@@ -318,17 +357,21 @@ def sphere_main(algorithm,
             and saving heatmap.
         seed (int): Seed for the algorithm. By default, there is no seed.
     """
+    # Use default dim for each algorithm.
     if dim is None:
         if algorithm in ["cma_mega", "cma_mega_adam"]:
             dim = 1_000
+        if algorithm in ["cma_mae", "cma_maega"]:
+            dim = 100
         elif algorithm in [
                 "map_elites", "line_map_elites", "cma_me_imp", "cma_me_imp_mu",
                 "cma_me_rd", "cma_me_rd_mu", "cma_me_opt", "cma_me_mixed"
         ]:
             dim = 20
 
+    # Use default itrs for each algorithm.
     if itrs is None:
-        if algorithm in ["cma_mega", "cma_mega_adam"]:
+        if algorithm in ["cma_mega", "cma_mega_adam", "cma_mae", "cma_maega"]:
             itrs = 10_000
         elif algorithm in [
                 "map_elites", "line_map_elites", "cma_me_imp", "cma_me_imp_mu",
diff --git a/ribs/archives/_grid_archive.py b/ribs/archives/_grid_archive.py
index 766203aea..876a06af1 100644
--- a/ribs/archives/_grid_archive.py
+++ b/ribs/archives/_grid_archive.py
@@ -32,6 +32,9 @@ class GridArchive(ArchiveBase):
             method -- refer to the implementation `here
             <../_modules/ribs/archives/_grid_archive.html#GridArchive.index_of>`_.
             Pass this parameter to configure that epsilon.
+        learning_rate (float): The learning rate of the archive. Described in
+            `Fontaine 2022 <https://arxiv.org/abs/2205.10752>`_.
+        threshold_min (float): The default threshold value for all the cells.
         seed (int): Value to seed the random number generator. Set to None to
             avoid a fixed seed.
         dtype (str or data-type): Data type of the solutions, objectives,
@@ -45,6 +48,8 @@ def __init__(self,
                  solution_dim,
                  dims,
                  ranges,
+                 learning_rate=1.0,
+                 threshold_min=0.0,
                  epsilon=1e-6,
                  seed=None,
                  dtype=np.float64):
@@ -58,6 +63,8 @@ def __init__(self,
             solution_dim=solution_dim,
             cells=np.product(self._dims),
             measure_dim=len(self._dims),
+            learning_rate=learning_rate,
+            threshold_min=threshold_min,
             seed=seed,
             dtype=dtype,
         )

From c2644e7d653e350636373cb854c2f62f5ec98a69 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Mon, 29 Aug 2022 16:51:45 -0700
Subject: [PATCH 07/53] add passive archvie

---
 examples/sphere.py             | 104 ++++++++++++++++++++++-----------
 ribs/archives/_archive_base.py |   2 -
 ribs/schedulers/_scheduler.py  |  27 ++++++++-
 3 files changed, 95 insertions(+), 38 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index 1b3345185..f5d0a8052 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -147,7 +147,12 @@ def sphere(solution_batch):
     )
 
 
-def create_scheduler(algorithm, dim, seed):
+def create_scheduler(algorithm,
+                     solution_dims,
+                     archive_dims,
+                     learning_rate,
+                     use_result_archive=True,
+                     seed=None):
     """Creates a scheduler based on the algorithm name.
 
     Args:
@@ -157,45 +162,33 @@ def create_scheduler(algorithm, dim, seed):
     Returns:
         scheduler: A ribs scheduler for running the algorithm.
     """
-    max_bound = dim / 2 * 5.12
+    max_bound = solution_dims / 2 * 5.12
     bounds = [(-max_bound, max_bound), (-max_bound, max_bound)]
-    initial_sol = np.zeros(dim)
+    initial_sol = np.zeros(solution_dims)
     batch_size = 37
     num_emitters = 15
 
     # Create archive.
-    if algorithm in [
-            "map_elites", "line_map_elites", "cma_me_imp", "cma_me_imp_mu",
-            "cma_me_rd", "cma_me_rd_mu", "cma_me_opt", "cma_me_mixed"
-    ]:
-        archive = GridArchive(solution_dim=dim,
-                              dims=(500, 500),
-                              ranges=bounds,
-                              seed=seed)
-    elif algorithm in ["cvt_map_elites", "line_cvt_map_elites"]:
-        archive = CVTArchive(solution_dim=dim,
+    if algorithm in ["cvt_map_elites", "line_cvt_map_elites"]:
+        archive = CVTArchive(solution_dim=solution_dims,
                              cells=10_000,
                              ranges=bounds,
                              samples=100_000,
                              use_kd_tree=True)
-    elif algorithm in ["cma_mega", "cma_mega_adam"]:
-        # Note that the archive is smaller for these algorithms. This is to be
-        # consistent with Fontaine 2021 <https://arxiv.org/abs/2106.03894>.
-        archive = GridArchive(solution_dim=dim,
-                              dims=(100, 100),
-                              ranges=bounds,
-                              seed=seed)
-    elif algorithm in ["cma_mae", "cma_maega"]:
-        # Note that the archive is smaller for these algorithms. This is to be
-        # consistent with Fontaine 2022 <https://arxiv.org/abs/2205.10752>.
-        archive = GridArchive(solution_dim=dim,
-                              dims=(100, 100),
+    else:
+        archive = GridArchive(solution_dim=solution_dims,
+                              dims=archive_dims,
                               ranges=bounds,
-                              learning_rate=0.01,
+                              learning_rate=learning_rate,
                               threshold_min=0,
                               seed=seed)
-    else:
-        raise ValueError(f"Algorithm `{algorithm}` is not recognized")
+
+    # Create result archive.
+    if use_result_archive:
+        result_archive = GridArchive(solution_dim=solution_dims,
+                                     dims=archive_dims,
+                                     ranges=bounds,
+                                     seed=seed)
 
     # Create emitters. Each emitter needs a different seed, so that they do not
     # all do the same thing.
@@ -317,7 +310,12 @@ def create_scheduler(algorithm, dim, seed):
                                        batch_size=batch_size,
                                        seed=s) for s in emitter_seeds
         ]
-    return Scheduler(archive, emitters)
+    return Scheduler(
+        archive,
+        emitters,
+        result_archive,
+        add_mode="single"
+    )
 
 
 def save_heatmap(archive, heatmap_path):
@@ -343,6 +341,8 @@ def save_heatmap(archive, heatmap_path):
 def sphere_main(algorithm,
                 dim=None,
                 itrs=None,
+                archive_dims=None,
+                learning_rate=None,
                 outdir="sphere_output",
                 log_freq=250,
                 seed=None):
@@ -350,7 +350,7 @@ def sphere_main(algorithm,
 
     Args:
         algorithm (str): Name of the algorithm.
-        dim (int): Dimensionality of solutions.
+        dim (int): Dimensionality of the sphere function.
         itrs (int): Iterations to run.
         outdir (str): Directory to save output.
         log_freq (int): Number of iterations to wait before recording metrics
@@ -361,7 +361,7 @@ def sphere_main(algorithm,
     if dim is None:
         if algorithm in ["cma_mega", "cma_mega_adam"]:
             dim = 1_000
-        if algorithm in ["cma_mae", "cma_maega"]:
+        elif algorithm in ["cma_mae", "cma_maega"]:
             dim = 100
         elif algorithm in [
                 "map_elites", "line_map_elites", "cma_me_imp", "cma_me_imp_mu",
@@ -379,14 +379,41 @@ def sphere_main(algorithm,
         ]:
             itrs = 4500
 
+    # Use default archive_dim for each algorithm.
+    if archive_dims is None:
+        if algorithm in ["cma_mega", "cma_mega_adam", "cma_mae", "cma_maega"]:
+            archive_dims = (100, 100)
+        elif algorithm in [
+                "map_elites", "line_map_elites", "cma_me_imp", "cma_me_imp_mu",
+                "cma_me_rd", "cma_me_rd_mu", "cma_me_opt", "cma_me_mixed"
+        ]:
+            archive_dims = (500, 500)
+
+    # Use default learning_rate for each algorithm.
+    if learning_rate is None:
+        if algorithm in ["cma_mae", "cma_maega"]:
+            learning_rate = 0.01
+        elif algorithm in [
+                "map_elites", "line_map_elites", "cma_me_imp", "cma_me_imp_mu",
+                "cma_me_rd", "cma_me_rd_mu", "cma_me_opt", "cma_me_mixed",
+                "cma_mega", "cma_mega_adam"
+        ]:
+            learning_rate = 1.0
+
     name = f"{algorithm}_{dim}"
     outdir = Path(outdir)
     if not outdir.is_dir():
         outdir.mkdir()
 
-    is_dqd = algorithm in ['cma_mega', 'cma_mega_adam']
+    is_dqd = algorithm in ["cma_mega", "cma_mega_adam"]
+    use_result_archive = algorithm in ["cma_mae", "cma_maega"]
 
-    scheduler = create_scheduler(algorithm, dim, seed)
+    scheduler = create_scheduler(algorithm,
+                                 dim,
+                                 archive_dims,
+                                 learning_rate,
+                                 use_result_archive=use_result_archive,
+                                 seed=seed)
     archive = scheduler.archive
     metrics = {
         "QD Score": {
@@ -426,15 +453,22 @@ def sphere_main(algorithm,
             # Logging and output.
             final_itr = itr == itrs
             if itr % log_freq == 0 or final_itr:
+                result_archive = scheduler.result_archive
+
+                # Save a full archive for analysis
+                df = result_archive.as_pandas(include_solutions=final_itr)
+                df.to_pickle(str(outdir / f"{name}_archive_{itr:08d}.pkl"))
+
                 if final_itr:
                     archive.as_pandas(include_solutions=final_itr).to_csv(
                         outdir / f"{name}_archive.csv")
 
                 # Record and display metrics.
                 metrics["QD Score"]["x"].append(itr)
-                metrics["QD Score"]["y"].append(archive.stats.qd_score)
+                metrics["QD Score"]["y"].append(result_archive.stats.qd_score)
                 metrics["Archive Coverage"]["x"].append(itr)
-                metrics["Archive Coverage"]["y"].append(archive.stats.coverage)
+                metrics["Archive Coverage"]["y"].append(
+                    result_archive.stats.coverage)
                 print(f"Iteration {itr} | Archive Coverage: "
                       f"{metrics['Archive Coverage']['y'][-1] * 100:.3f}% "
                       f"QD Score: {metrics['QD Score']['y'][-1]:.3f}")
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 306fbe230..f2dd80570 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -634,8 +634,6 @@ def add_single(self, solution, objective, measures, metadata=None):
             self._occupied_arr[index] = True
 
             # Update the threshold.
-            # self._threshold_arr[index] = old_threshold * \
-            #     (1.0 - self._learning_rate) + objective * self._learning_rate
             self._threshold_arr[index] = self._compute_new_thresholds(
                 np.array([old_threshold]), np.array([objective]))
 
diff --git a/ribs/schedulers/_scheduler.py b/ribs/schedulers/_scheduler.py
index e14016bc4..b107c480e 100644
--- a/ribs/schedulers/_scheduler.py
+++ b/ribs/schedulers/_scheduler.py
@@ -36,6 +36,10 @@ class Scheduler:
             included for legacy reasons, as it was the only mode of operation in
             pyribs 0.4.0 and before. We highly recommend using "batch" mode
             since it is significantly faster.
+        result_archive (ribs.archives.ArchiveBase): In some algorithms, such as
+            CMA-MAE, the archive does not store all the best-performing
+            solutions. The `result_archive` is a secondary archive where we can
+            store all the best-performing solutions.
     Raises:
         ValueError: The emitters passed in do not have the same solution
             dimensions.
@@ -45,7 +49,11 @@ class Scheduler:
         ValueError: Invalid value for `add_mode`.
     """
 
-    def __init__(self, archive, emitters, add_mode="batch"):
+    def __init__(self,
+                 archive,
+                 emitters,
+                 result_archive=None,
+                 add_mode="batch"):
         if len(emitters) == 0:
             raise ValueError("Pass in at least one emitter to the scheduler.")
 
@@ -75,6 +83,8 @@ def __init__(self, archive, emitters, add_mode="batch"):
         self._emitters = emitters
         self._add_mode = add_mode
 
+        self._result_archive = result_archive
+
         # Keeps track of whether the scheduler should be receiving a call to
         # ask() or tell().
         self._last_called = None
@@ -95,6 +105,16 @@ def emitters(self):
         in this scheduler."""
         return self._emitters
 
+    @property
+    def result_archive(self):
+        """ribs.archives.ArchiveBase: Another archive for storing solutions
+        found in this optimizer.
+        If `result_archive` was not passed to the constructor, this property is
+        the same as :attr:`archive`.
+        """
+        return (self._archive
+                if self._result_archive is None else self._result_archive)
+
     def ask_dqd(self):
         """Generates a batch of solutions by calling ask_dqd() on all DQD
         emitters.
@@ -209,6 +229,11 @@ def _tell_internal(self,
             status_batch = np.asarray(status_batch)
             value_batch = np.asarray(value_batch)
 
+        # Add solution to result_archive
+        if self._result_archive is not None:
+            self._result_archive.add(self._solution_batch, objective_batch,
+                                     measures_batch, metadata_batch)
+
         return (
             objective_batch,
             measures_batch,

From 3d49034cd19f34a045f0f2f0f977d3780d633ae9 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Mon, 29 Aug 2022 18:43:13 -0700
Subject: [PATCH 08/53] Use improvement ranker

---
 examples/sphere.py             |  9 ++-------
 ribs/archives/_archive_base.py | 12 +++++++-----
 ribs/emitters/rankers.py       | 23 +++++++++++++++++++++--
 3 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index f5d0a8052..46d1e4c1c 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -291,7 +291,7 @@ def create_scheduler(algorithm,
                 archive=archive,
                 x0=initial_sol,
                 sigma0=0.5,
-                ranker="2imp",
+                ranker="imp",
                 selection_rule="mu",
                 restart_rule="basic",
                 batch_size=batch_size,
@@ -310,12 +310,7 @@ def create_scheduler(algorithm,
                                        batch_size=batch_size,
                                        seed=s) for s in emitter_seeds
         ]
-    return Scheduler(
-        archive,
-        emitters,
-        result_archive,
-        add_mode="single"
-    )
+    return Scheduler(archive, emitters, result_archive, add_mode="single")
 
 
 def save_heatmap(archive, heatmap_path):
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index f2dd80570..860dbf16d 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -622,8 +622,8 @@ def add_single(self, solution, objective, measures, metadata=None):
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
 
+        # Note that when learning_rate = 1.0, old_threshold === old_objective. 
         old_objective = self._objective_arr[index]
-        # Note that when learning_rate = 1.0, this is equivalent to old_objective
         old_threshold = self._threshold_arr[index]
 
         was_occupied = self._occupied_arr[index]
@@ -634,8 +634,11 @@ def add_single(self, solution, objective, measures, metadata=None):
             self._occupied_arr[index] = True
 
             # Update the threshold.
-            self._threshold_arr[index] = self._compute_new_thresholds(
-                np.array([old_threshold]), np.array([objective]))
+            if old_threshold < objective:
+                # self._threshold_arr[index] = old_threshold * \
+                #     (1.0 - self._learning_rate) + objective * self._learning_rate
+                self._threshold_arr[index] = self._compute_new_thresholds(
+                    np.array([old_threshold]), np.array([objective]))
 
             # Insert into the archive.
             self._objective_arr[index] = objective
@@ -651,7 +654,6 @@ def add_single(self, solution, objective, measures, metadata=None):
                 self._num_occupied += 1
                 status = 2  # NEW
 
-        value = objective - old_threshold
         if status:
             # Update archive stats.
             if status == 2:
@@ -678,7 +680,7 @@ def add_single(self, solution, objective, measures, metadata=None):
                 obj_mean=new_qd_score / self.dtype(len(self)),
             )
 
-        return status, value
+        return status, self.dtype(objective - old_threshold)
 
     def elites_with_measures(self, measures_batch):
         """Retrieves the elites with measures in the same cells as the measures
diff --git a/ribs/emitters/rankers.py b/ribs/emitters/rankers.py
index 1c059a397..3e43ea012 100644
--- a/ribs/emitters/rankers.py
+++ b/ribs/emitters/rankers.py
@@ -117,8 +117,27 @@ def reset(self, emitter, archive, rng):
 
 
 class ImprovementRanker(RankerBase):
-    # TODO Implement this (in another PR)
-    pass
+    """Ranks the solutions based on the improvement in the objective.
+
+    This ranker ranks solutions in a single stage. The solutions are ranked by
+    the improvement "value" described in :meth:`ArchiveBase.add`.
+    """
+
+    def rank(self, emitter, archive, rng, solution_batch, objective_batch,
+             measures_batch, status_batch, value_batch, metadata_batch):
+        # To avoid using an array of tuples, ranking_values is a 2D array
+        # [[status_0, value_0], ..., [status_n, value_n]]
+        ranking_values = np.stack((status_batch, value_batch), axis=-1)
+
+        # Note that lexsort sorts the values in ascending order,
+        # so we use np.flip to reverse the sorted array.
+        return np.flip(np.argsort(value_batch)), value_batch
+
+    rank.__doc__ = f"""
+Generates a list of indices that represents an ordering of solutions.
+
+{_rank_args}
+    """
 
 
 class TwoStageImprovementRanker(RankerBase):

From b5f6587e06c03f958207d47b2b35957ba6ecd96c Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Tue, 30 Aug 2022 18:56:58 -0700
Subject: [PATCH 09/53] history

---
 HISTORY.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/HISTORY.md b/HISTORY.md
index 7ef1c96cc..48d879097 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -6,6 +6,8 @@
 
 #### API
 
+- Implement CMA-MAE archive thresholds (#256)
+  - Revive the old implementation of `add_single` removed in (#221)
 - Add restart timer to `EvolutionStrategyEmitter` and `GradientAborescenceEmitter`(#255)
 - Rename fields and update documentation (#249, #250)
   - **Backwards-incompatible:** rename `Optimizer` to `Scheduler`

From 4c605ca5b18b6595157c1a29d07d82b4f5d374c6 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 31 Aug 2022 00:34:38 -0700
Subject: [PATCH 10/53] tried to fix batch add threshold

---
 examples/sphere.py             | 36 ++++++++++++++------
 ribs/archives/_archive_base.py | 60 ++++++++++++++++++++++------------
 ribs/schedulers/_scheduler.py  |  6 ++--
 3 files changed, 70 insertions(+), 32 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index 46d1e4c1c..b98ddb363 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -157,10 +157,14 @@ def create_scheduler(algorithm,
 
     Args:
         algorithm (str): Name of the algorithm passed into sphere_main.
-        dim (int): Dimensionality of the sphere function.
+        solution_dims (int): Dimensionality of the sphere function.
+        archive_dims (int): Dimensionality of the archive.
+        learning_rate (float): Learning rate of archive.
+        use_result_archive (bool): Whether to use a separate archive to store
+            the results.
         seed (int): Main seed or the various components.
     Returns:
-        scheduler: A ribs scheduler for running the algorithm.
+        ribs.schedulers.Scheduler: A ribs scheduler for running the algorithm.
     """
     max_bound = solution_dims / 2 * 5.12
     bounds = [(-max_bound, max_bound), (-max_bound, max_bound)]
@@ -184,6 +188,7 @@ def create_scheduler(algorithm,
                               seed=seed)
 
     # Create result archive.
+    result_archive = None
     if use_result_archive:
         result_archive = GridArchive(solution_dim=solution_dims,
                                      dims=archive_dims,
@@ -310,6 +315,10 @@ def create_scheduler(algorithm,
                                        batch_size=batch_size,
                                        seed=s) for s in emitter_seeds
         ]
+
+    print(
+        f"Created Scheduler for {algorithm} with learning rate {learning_rate}, "
+        f"using solution dims {solution_dims} and archive dims {archive_dims}.")
     return Scheduler(archive, emitters, result_archive, add_mode="single")
 
 
@@ -410,6 +419,7 @@ def sphere_main(algorithm,
                                  use_result_archive=use_result_archive,
                                  seed=seed)
     archive = scheduler.archive
+    result_archive = scheduler.result_archive
     metrics = {
         "QD Score": {
             "x": [0],
@@ -424,6 +434,7 @@ def sphere_main(algorithm,
     non_logging_time = 0.0
     with alive_bar(itrs) as progress:
         save_heatmap(archive, str(outdir / f"{name}_heatmap_{0:05d}.png"))
+        save_heatmap(result_archive, str(outdir / f"{name}_heatmap_{0:05d}_result.png"))
 
         for itr in range(1, itrs + 1):
             itr_start = time.time()
@@ -440,23 +451,23 @@ def sphere_main(algorithm,
                                    jacobian_batch)
 
             solution_batch = scheduler.ask()
-            objective_batch, _, measure_batch, _ = sphere(solution_batch)
-            scheduler.tell(objective_batch, measure_batch)
+            objective_batch, _, measures_batch, _ = sphere(solution_batch)
+
+            scheduler.tell(objective_batch, measures_batch)
             non_logging_time += time.time() - itr_start
             progress()
 
             # Logging and output.
             final_itr = itr == itrs
             if itr % log_freq == 0 or final_itr:
-                result_archive = scheduler.result_archive
-
                 # Save a full archive for analysis
-                df = result_archive.as_pandas(include_solutions=final_itr)
-                df.to_pickle(str(outdir / f"{name}_archive_{itr:08d}.pkl"))
+                # df = result_archive.as_pandas(include_solutions=final_itr)
+                # df.to_pickle(str(outdir / f"{name}_archive_{itr:08d}.pkl"))
 
                 if final_itr:
-                    archive.as_pandas(include_solutions=final_itr).to_csv(
-                        outdir / f"{name}_archive.csv")
+                    result_archive.as_pandas(
+                        include_solutions=final_itr).to_csv(
+                            outdir / f"{name}_archive.csv")
 
                 # Record and display metrics.
                 metrics["QD Score"]["x"].append(itr)
@@ -471,6 +482,11 @@ def sphere_main(algorithm,
                 save_heatmap(archive,
                              str(outdir / f"{name}_heatmap_{itr:05d}.png"))
 
+                # Save result_archive
+                save_heatmap(
+                    result_archive,
+                    str(outdir / f"{name}_heatmap_{itr:05d}_result.png"))
+
     # Plot metrics.
     print(f"Algorithm Time (Excludes Logging and Setup): {non_logging_time}s")
     for metric in metrics:
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 860dbf16d..bc1860049 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -269,22 +269,23 @@ def _sum_geometric_series(self, a, r, n):
             return a
         return (a * (1 - pow(r, n))) / (1 - r)
 
-    def _compute_new_thresholds(self, old_threshold_batch,
-                                objective_batch_insert):
+    def _compute_new_thresholds(self, old_threshold, objective_batch_insert):
         """Update thresholds.
 
-            Args:
-                old_threshold_batch (np.ndarray): The threshold of the cells before updating.
-                objective_batch_insert (np.ndarray): The objective values of
-                    the solution that is inserted into the archive.
-            Returns:
-                A batch of new thresholds.
+        Args:
+            old_threshold (float): The threshold of the cells before
+                updating.
+            objective_batch_insert (np.ndarray): 1D array of the objective
+                values of the solution that is inserted into the archive for
+                each cell.
+        Returns:
+            A batch of new thresholds.
         """
-        k = len(objective_batch_insert)
+        k = np.size(objective_batch_insert)
         geometric_sum = self._sum_geometric_series(1, 1.0 - self._learning_rate,
                                                    k)
         return (self._learning_rate * np.sum(objective_batch_insert) *
-                geometric_sum / k) + (old_threshold_batch *
+                geometric_sum / k) + (old_threshold *
                                       (1.0 - self._learning_rate)**k)
 
     def clear(self):
@@ -504,7 +505,33 @@ def add(self,
         index_batch_can = index_batch[can_insert]
         metadata_batch_can = metadata_batch[can_insert]
         old_objective_batch_can = old_objective_arr[can_insert]
-        old_threshold_batch_can = old_threshold_arr[can_insert]
+
+        def groupby(a, b):
+            # Get argsort indices, to be used to sort a and b in the next steps
+            sidx = b.argsort(kind='mergesort')
+            a_sorted = a[sidx]
+            b_sorted = b[sidx]
+
+            # Get the group limit indices (start, stop of groups)
+            cut_idx = np.flatnonzero(np.r_[True, b_sorted[1:] != b_sorted[:-1],
+                                           True])
+
+            # Split input array with those start, stop ones
+            out = np.array(
+                [a_sorted[i:j] for i, j in zip(cut_idx[:-1], cut_idx[1:])],
+                dtype=object)
+            return out, np.unique(b_sorted)
+
+        # Group the objectives that can be inserted into the archive by cell index.
+        objective_batch_by_cell, group_by_index = groupby(
+            objective_batch_can, index_batch_can)
+        old_threshold_batch = self._threshold_arr[group_by_index]
+
+        # Update the thresholds.
+        for objective_batch, old_threshold in zip(objective_batch_by_cell,
+                                                  old_threshold_batch):
+            self._threshold_arr[group_by_index] = self._compute_new_thresholds(
+                old_threshold, objective_batch)
 
         # Retrieve indices of solutions that should be inserted into the
         # archive. Currently, multiple solutions may be inserted at each
@@ -534,11 +561,6 @@ def add(self,
         index_batch_insert = index_batch_can[should_insert]
         metadata_batch_insert = metadata_batch_can[should_insert]
         old_objective_batch_insert = old_objective_batch_can[should_insert]
-        old_threshold_batch_insert = old_threshold_batch_can[should_insert]
-
-        # Update the thresholds.
-        self._threshold_arr[index_batch_insert] = self._compute_new_thresholds(
-            old_threshold_batch_insert, objective_batch_insert)
 
         # Set archive storage.
         self._objective_arr[index_batch_insert] = objective_batch_insert
@@ -622,7 +644,7 @@ def add_single(self, solution, objective, measures, metadata=None):
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
 
-        # Note that when learning_rate = 1.0, old_threshold === old_objective. 
+        # Note that when learning_rate = 1.0, old_threshold === old_objective.
         old_objective = self._objective_arr[index]
         old_threshold = self._threshold_arr[index]
 
@@ -635,10 +657,8 @@ def add_single(self, solution, objective, measures, metadata=None):
 
             # Update the threshold.
             if old_threshold < objective:
-                # self._threshold_arr[index] = old_threshold * \
-                #     (1.0 - self._learning_rate) + objective * self._learning_rate
                 self._threshold_arr[index] = self._compute_new_thresholds(
-                    np.array([old_threshold]), np.array([objective]))
+                    old_threshold, np.array([objective]))
 
             # Insert into the archive.
             self._objective_arr[index] = objective
diff --git a/ribs/schedulers/_scheduler.py b/ribs/schedulers/_scheduler.py
index b107c480e..0c18eab5f 100644
--- a/ribs/schedulers/_scheduler.py
+++ b/ribs/schedulers/_scheduler.py
@@ -231,8 +231,10 @@ def _tell_internal(self,
 
         # Add solution to result_archive
         if self._result_archive is not None:
-            self._result_archive.add(self._solution_batch, objective_batch,
-                                     measures_batch, metadata_batch)
+            for sol, obj, mea, meta in zip(self._solution_batch,
+                                           objective_batch, measures_batch,
+                                           metadata_batch):
+                self._result_archive.add_single(sol, obj, mea, meta)
 
         return (
             objective_batch,

From 354386a9098913c576866cb7cc7fe7802f3f5ad3 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 31 Aug 2022 01:48:52 -0700
Subject: [PATCH 11/53] fix batch add

---
 examples/sphere.py             | 16 +++++++---------
 ribs/archives/_archive_base.py |  9 ++++-----
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index b98ddb363..60207a9de 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -171,6 +171,7 @@ def create_scheduler(algorithm,
     initial_sol = np.zeros(solution_dims)
     batch_size = 37
     num_emitters = 15
+    mode = "batch"
 
     # Create archive.
     if algorithm in ["cvt_map_elites", "line_cvt_map_elites"]:
@@ -317,9 +318,9 @@ def create_scheduler(algorithm,
         ]
 
     print(
-        f"Created Scheduler for {algorithm} with learning rate {learning_rate}, "
+        f"Created Scheduler for {algorithm} with learning rate {learning_rate} and add mode {mode}, "
         f"using solution dims {solution_dims} and archive dims {archive_dims}.")
-    return Scheduler(archive, emitters, result_archive, add_mode="single")
+    return Scheduler(archive, emitters, result_archive, add_mode=mode)
 
 
 def save_heatmap(archive, heatmap_path):
@@ -433,8 +434,9 @@ def sphere_main(algorithm,
 
     non_logging_time = 0.0
     with alive_bar(itrs) as progress:
-        save_heatmap(archive, str(outdir / f"{name}_heatmap_{0:05d}.png"))
-        save_heatmap(result_archive, str(outdir / f"{name}_heatmap_{0:05d}_result.png"))
+        save_heatmap(archive, str(outdir / f"{name}_heatmap_{0:05d}_main.png"))
+        save_heatmap(result_archive,
+                     str(outdir / f"{name}_heatmap_{0:05d}_result.png"))
 
         for itr in range(1, itrs + 1):
             itr_start = time.time()
@@ -460,10 +462,6 @@ def sphere_main(algorithm,
             # Logging and output.
             final_itr = itr == itrs
             if itr % log_freq == 0 or final_itr:
-                # Save a full archive for analysis
-                # df = result_archive.as_pandas(include_solutions=final_itr)
-                # df.to_pickle(str(outdir / f"{name}_archive_{itr:08d}.pkl"))
-
                 if final_itr:
                     result_archive.as_pandas(
                         include_solutions=final_itr).to_csv(
@@ -480,7 +478,7 @@ def sphere_main(algorithm,
                       f"QD Score: {metrics['QD Score']['y'][-1]:.3f}")
 
                 save_heatmap(archive,
-                             str(outdir / f"{name}_heatmap_{itr:05d}.png"))
+                             str(outdir / f"{name}_heatmap_{itr:05d}_main.png"))
 
                 # Save result_archive
                 save_heatmap(
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index bc1860049..aa8476def 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -525,13 +525,12 @@ def groupby(a, b):
         # Group the objectives that can be inserted into the archive by cell index.
         objective_batch_by_cell, group_by_index = groupby(
             objective_batch_can, index_batch_can)
-        old_threshold_batch = self._threshold_arr[group_by_index]
 
         # Update the thresholds.
-        for objective_batch, old_threshold in zip(objective_batch_by_cell,
-                                                  old_threshold_batch):
-            self._threshold_arr[group_by_index] = self._compute_new_thresholds(
-                old_threshold, objective_batch)
+        for objective_batch, index in zip(objective_batch_by_cell,
+                                          group_by_index):
+            self._threshold_arr[index] = self._compute_new_thresholds(
+                self._threshold_arr[index], objective_batch)
 
         # Retrieve indices of solutions that should be inserted into the
         # archive. Currently, multiple solutions may be inserted at each

From 76bfaf6a38a3d59883c957e03833718b8f2f865e Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 31 Aug 2022 02:12:07 -0700
Subject: [PATCH 12/53] fix add batch

---
 ribs/schedulers/_scheduler.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/ribs/schedulers/_scheduler.py b/ribs/schedulers/_scheduler.py
index 0c18eab5f..ee18cff2d 100644
--- a/ribs/schedulers/_scheduler.py
+++ b/ribs/schedulers/_scheduler.py
@@ -142,7 +142,7 @@ def ask_dqd(self):
                 self._solution_batch.append(emitter_sols)
                 self._num_emitted[i] = len(emitter_sols)
 
-        # In case the emitters didn't return any solutions
+        # In case the emitters didn't return any solutions.
         self._solution_batch = np.concatenate(
             self._solution_batch, axis=0) if self._solution_batch else np.empty(
                 (0, self._solution_dim))
@@ -173,7 +173,7 @@ def ask(self):
             self._solution_batch.append(emitter_sols)
             self._num_emitted[i] = len(emitter_sols)
 
-        # In case the emitters didn't return any solutions
+        # In case the emitters didn't return any solutions.
         self._solution_batch = np.concatenate(
             self._solution_batch, axis=0) if self._solution_batch else np.empty(
                 (0, self._solution_dim))
@@ -229,12 +229,10 @@ def _tell_internal(self,
             status_batch = np.asarray(status_batch)
             value_batch = np.asarray(value_batch)
 
-        # Add solution to result_archive
+        # Add solutions to result_archive.
         if self._result_archive is not None:
-            for sol, obj, mea, meta in zip(self._solution_batch,
-                                           objective_batch, measures_batch,
-                                           metadata_batch):
-                self._result_archive.add_single(sol, obj, mea, meta)
+            self._result_archive.add(self._solution_batch, objective_batch,
+                                     measures_batch, metadata_batch)
 
         return (
             objective_batch,

From ccbf83292bf0149b1a7cd37f7bb882d28d975b70 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 31 Aug 2022 03:36:05 -0700
Subject: [PATCH 13/53] fix batch add bug

---
 examples/sphere.py             | 16 +++++-----------
 ribs/archives/_archive_base.py |  1 +
 2 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index 60207a9de..846136898 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -406,11 +406,11 @@ def sphere_main(algorithm,
             learning_rate = 1.0
 
     name = f"{algorithm}_{dim}"
-    outdir = Path(outdir)
+    outdir = Path(outdir + "_" + algorithm)
     if not outdir.is_dir():
         outdir.mkdir()
 
-    is_dqd = algorithm in ["cma_mega", "cma_mega_adam"]
+    is_dqd = algorithm in ["cma_mega", "cma_mega_adam", "cma_maega"]
     use_result_archive = algorithm in ["cma_mae", "cma_maega"]
 
     scheduler = create_scheduler(algorithm,
@@ -419,7 +419,6 @@ def sphere_main(algorithm,
                                  learning_rate,
                                  use_result_archive=use_result_archive,
                                  seed=seed)
-    archive = scheduler.archive
     result_archive = scheduler.result_archive
     metrics = {
         "QD Score": {
@@ -434,9 +433,8 @@ def sphere_main(algorithm,
 
     non_logging_time = 0.0
     with alive_bar(itrs) as progress:
-        save_heatmap(archive, str(outdir / f"{name}_heatmap_{0:05d}_main.png"))
         save_heatmap(result_archive,
-                     str(outdir / f"{name}_heatmap_{0:05d}_result.png"))
+                     str(outdir / f"{name}_heatmap_{0:05d}.png"))
 
         for itr in range(1, itrs + 1):
             itr_start = time.time()
@@ -477,13 +475,9 @@ def sphere_main(algorithm,
                       f"{metrics['Archive Coverage']['y'][-1] * 100:.3f}% "
                       f"QD Score: {metrics['QD Score']['y'][-1]:.3f}")
 
-                save_heatmap(archive,
-                             str(outdir / f"{name}_heatmap_{itr:05d}_main.png"))
-
                 # Save result_archive
-                save_heatmap(
-                    result_archive,
-                    str(outdir / f"{name}_heatmap_{itr:05d}_result.png"))
+                save_heatmap(result_archive,
+                             str(outdir / f"{name}_heatmap_{itr:05d}.png"))
 
     # Plot metrics.
     print(f"Algorithm Time (Excludes Logging and Setup): {non_logging_time}s")
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index aa8476def..b673dbbd1 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -486,6 +486,7 @@ def add(self,
 
         # Since we set the new solutions in old_objective_batch to have
         # value 0.0, the values for new solutions are correct here.
+        old_objective_arr[is_new] = 0.0
         old_threshold_arr[is_new] = 0.0
         value_batch = objective_batch - old_threshold_arr
 

From b68f09c63fd4dc9fce57f41eaf330d295e89d96d Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 31 Aug 2022 04:22:58 -0700
Subject: [PATCH 14/53] fix requested changes

---
 examples/sphere.py             | 41 +++++++++++++++++++---------------
 ribs/archives/_archive_base.py | 18 +++++++--------
 ribs/archives/_grid_archive.py |  2 +-
 ribs/emitters/rankers.py       |  4 ----
 ribs/schedulers/_scheduler.py  | 41 +++++++++++++++++++---------------
 5 files changed, 56 insertions(+), 50 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index 846136898..7ddd768e6 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -34,23 +34,25 @@
 - `cma_mega`: GridArchive with GradientAborescenceEmitter.
 - `cma_mega_adam`: GridArchive with GradientAborescenceEmitter using Adam
   Optimizer.
-- `cma_mae`: GridArchive (learning_rate = 0.01) with EvolutionStrategyEmitter.
-- `cma_maega`: GridArchive (learning_rate = 0.01) with GradientAborescenceEmitter.
+- `cma_mae`: GridArchive (learning_rate = 0.01) with EvolutionStrategyEmitter
+  using ImprovementRanker.
+- `cma_maega`: GridArchive (learning_rate = 0.01) with
+  GradientAborescenceEmitter using ImprovementRanker.
 
 All algorithms use 15 emitters, each with a batch size of 37. Each one runs for
 4500 iterations for a total of 15 * 37 * 4500 ~= 2.5M evaluations.
 
-Exceptions:
-    - `cma_mega` and `cma_mega_adam` uses only one emitter and runs for 10,000
+Notes:
+    - `cma_mega` and `cma_mega_adam` use only one emitter and run for 10,000
       iterations. This is to be consistent with the paper (`Fontaine 2021
       <https://arxiv.org/abs/2106.03894>`_) in which these algorithms were
       proposed.
-    - `cma_mae` and `cma_maega`
-
-Note that the CVTArchive in this example uses 10,000 cells, as opposed to the
-250,000 (500x500) in the GridArchive, so it is not fair to directly compare
-`cvt_map_elites` and `line_cvt_map_elites` to the other algorithms. However, the
-other algorithms may be fairly compared because they use the same archive.
+    - `cma_mae` and `cma_maega` run for 10,000 iterations as well.
+    - CVTArchive in this example uses 10,000 cells, as opposed to the 250,000
+      (500x500) in the GridArchive, so it is not fair to directly compare
+      `cvt_map_elites` and `line_cvt_map_elites` to the other algorithms.
+      However, the other algorithms may be fairly compared because they use the
+      same archive.
 
 Outputs are saved in the `sphere_output/` directory by default. The archive is
 saved as a CSV named `{algorithm}_{dim}_archive.csv`, while snapshots of the
@@ -148,7 +150,7 @@ def sphere(solution_batch):
 
 
 def create_scheduler(algorithm,
-                     solution_dims,
+                     solution_dim,
                      archive_dims,
                      learning_rate,
                      use_result_archive=True,
@@ -157,7 +159,7 @@ def create_scheduler(algorithm,
 
     Args:
         algorithm (str): Name of the algorithm passed into sphere_main.
-        solution_dims (int): Dimensionality of the sphere function.
+        solution_dim(int): Dimensionality of the sphere function.
         archive_dims (int): Dimensionality of the archive.
         learning_rate (float): Learning rate of archive.
         use_result_archive (bool): Whether to use a separate archive to store
@@ -166,22 +168,22 @@ def create_scheduler(algorithm,
     Returns:
         ribs.schedulers.Scheduler: A ribs scheduler for running the algorithm.
     """
-    max_bound = solution_dims / 2 * 5.12
+    max_bound = solution_dim / 2 * 5.12
     bounds = [(-max_bound, max_bound), (-max_bound, max_bound)]
-    initial_sol = np.zeros(solution_dims)
+    initial_sol = np.zeros(solution_dim)
     batch_size = 37
     num_emitters = 15
     mode = "batch"
 
     # Create archive.
     if algorithm in ["cvt_map_elites", "line_cvt_map_elites"]:
-        archive = CVTArchive(solution_dim=solution_dims,
+        archive = CVTArchive(solution_dim=solution_dim,
                              cells=10_000,
                              ranges=bounds,
                              samples=100_000,
                              use_kd_tree=True)
     else:
-        archive = GridArchive(solution_dim=solution_dims,
+        archive = GridArchive(solution_dim=solution_dim,
                               dims=archive_dims,
                               ranges=bounds,
                               learning_rate=learning_rate,
@@ -191,7 +193,7 @@ def create_scheduler(algorithm,
     # Create result archive.
     result_archive = None
     if use_result_archive:
-        result_archive = GridArchive(solution_dim=solution_dims,
+        result_archive = GridArchive(solution_dim=solution_dim,
                                      dims=archive_dims,
                                      ranges=bounds,
                                      seed=seed)
@@ -310,6 +312,7 @@ def create_scheduler(algorithm,
                                        initial_sol,
                                        sigma0=10.0,
                                        step_size=1.0,
+                                       ranker="imp",
                                        grad_opt="gradient_ascent",
                                        restart_rule="basic",
                                        bounds=None,
@@ -319,7 +322,7 @@ def create_scheduler(algorithm,
 
     print(
         f"Created Scheduler for {algorithm} with learning rate {learning_rate} and add mode {mode}, "
-        f"using solution dims {solution_dims} and archive dims {archive_dims}.")
+        f"using solution dims {solution_dim} and archive dims {archive_dims}.")
     return Scheduler(archive, emitters, result_archive, add_mode=mode)
 
 
@@ -357,6 +360,8 @@ def sphere_main(algorithm,
         algorithm (str): Name of the algorithm.
         dim (int): Dimensionality of the sphere function.
         itrs (int): Iterations to run.
+        archive_dims (tuple): Dimensionality of the archive.
+        learning_rate (float): The archive learning rate.
         outdir (str): Directory to save output.
         log_freq (int): Number of iterations to wait before recording metrics
             and saving heatmap.
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index b673dbbd1..23c44b04c 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -138,7 +138,7 @@ def __init__(self,
                  cells,
                  measure_dim,
                  learning_rate=1.0,
-                 threshold_min=0.0,
+                 threshold_min=-np.inf,
                  seed=None,
                  dtype=np.float64):
 
@@ -471,24 +471,24 @@ def add(self,
         index_batch = self.index_of(measures_batch)
 
         # Copy old objectives since we will be modifying the objectives storage.
-        old_objective_arr = np.copy(self._objective_arr[index_batch])
-        old_threshold_arr = np.copy(self._threshold_arr[index_batch])
+        old_objective_batch = np.copy(self._objective_arr[index_batch])
+        old_threshold_batch = np.copy(self._threshold_arr[index_batch])
 
         # Compute the statuses -- these are all boolean arrays of length
         # batch_size.
         already_occupied = self._occupied_arr[index_batch]
         is_new = ~already_occupied
         improve_existing = (objective_batch >
-                            old_threshold_arr) & already_occupied
+                            old_threshold_batch) & already_occupied
         status_batch = np.zeros(batch_size, dtype=np.int32)
         status_batch[is_new] = 2
         status_batch[improve_existing] = 1
 
         # Since we set the new solutions in old_objective_batch to have
         # value 0.0, the values for new solutions are correct here.
-        old_objective_arr[is_new] = 0.0
-        old_threshold_arr[is_new] = 0.0
-        value_batch = objective_batch - old_threshold_arr
+        old_objective_batch[is_new] = 0.0
+        old_threshold_batch[is_new] = 0.0
+        value_batch = objective_batch - old_threshold_batch
 
         ## Step 3: Insert solutions into archive. ##
 
@@ -505,7 +505,7 @@ def add(self,
         measures_batch_can = measures_batch[can_insert]
         index_batch_can = index_batch[can_insert]
         metadata_batch_can = metadata_batch[can_insert]
-        old_objective_batch_can = old_objective_arr[can_insert]
+        old_objective_batch_can = old_objective_batch[can_insert]
 
         def groupby(a, b):
             # Get argsort indices, to be used to sort a and b in the next steps
@@ -644,7 +644,7 @@ def add_single(self, solution, objective, measures, metadata=None):
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
 
-        # Note that when learning_rate = 1.0, old_threshold === old_objective.
+        # Note that when learning_rate = 1.0, old_threshold == old_objective.
         old_objective = self._objective_arr[index]
         old_threshold = self._threshold_arr[index]
 
diff --git a/ribs/archives/_grid_archive.py b/ribs/archives/_grid_archive.py
index 876a06af1..bcf995310 100644
--- a/ribs/archives/_grid_archive.py
+++ b/ribs/archives/_grid_archive.py
@@ -49,7 +49,7 @@ def __init__(self,
                  dims,
                  ranges,
                  learning_rate=1.0,
-                 threshold_min=0.0,
+                 threshold_min=-np.inf,
                  epsilon=1e-6,
                  seed=None,
                  dtype=np.float64):
diff --git a/ribs/emitters/rankers.py b/ribs/emitters/rankers.py
index 3e43ea012..e4e9ad821 100644
--- a/ribs/emitters/rankers.py
+++ b/ribs/emitters/rankers.py
@@ -125,10 +125,6 @@ class ImprovementRanker(RankerBase):
 
     def rank(self, emitter, archive, rng, solution_batch, objective_batch,
              measures_batch, status_batch, value_batch, metadata_batch):
-        # To avoid using an array of tuples, ranking_values is a 2D array
-        # [[status_0, value_0], ..., [status_n, value_n]]
-        ranking_values = np.stack((status_batch, value_batch), axis=-1)
-
         # Note that lexsort sorts the values in ascending order,
         # so we use np.flip to reverse the sorted array.
         return np.flip(np.argsort(value_batch)), value_batch
diff --git a/ribs/schedulers/_scheduler.py b/ribs/schedulers/_scheduler.py
index ee18cff2d..b1b545d41 100644
--- a/ribs/schedulers/_scheduler.py
+++ b/ribs/schedulers/_scheduler.py
@@ -79,6 +79,9 @@ def __init__(self,
             raise ValueError("add_mode must either be 'batch' or 'single', but "
                              f"it was '{add_mode}'")
 
+        if archive is result_archive:
+            raise ValueError("archive has same id as result_archive")
+
         self._archive = archive
         self._emitters = emitters
         self._add_mode = add_mode
@@ -212,28 +215,30 @@ def _tell_internal(self,
                 measures_batch,
                 metadata_batch,
             )
+
+            # Add solutions to result_archive.
+            if self._result_archive is not None:
+                self._result_archive.add(self._solution_batch, objective_batch,
+                                         measures_batch, metadata_batch)
         elif self._add_mode == "single":
-            status_batch, value_batch = zip(*[
-                self.archive.add_single(
-                    solution,
-                    objective,
-                    measure,
-                    metadata,
-                ) for solution, objective, measure, metadata in zip(
-                    self._solution_batch,
-                    objective_batch,
-                    measures_batch,
-                    metadata_batch,
-                )
-            ])
+            status_batch = []
+            value_batch = []
+            for solution, objective, measure, metadata in zip(
+                    self._solution_batch, objective_batch, measures_batch,
+                    metadata_batch):
+                status, value = self.archive.add_single(solution, objective,
+                                                        measure, metadata)
+                status_batch.append(status)
+                value_batch.append(value)
+
+                # Add solutions to result_archive.
+                if self._result_archive is not None:
+                    self._result_archive.add(self._solution_batch,
+                                             objective_batch, measures_batch,
+                                             metadata_batch)
             status_batch = np.asarray(status_batch)
             value_batch = np.asarray(value_batch)
 
-        # Add solutions to result_archive.
-        if self._result_archive is not None:
-            self._result_archive.add(self._solution_batch, objective_batch,
-                                     measures_batch, metadata_batch)
-
         return (
             objective_batch,
             measures_batch,

From c781cda9ff055884dba23278f7b5c2a10c9c9de7 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Sat, 3 Sep 2022 23:08:11 -0700
Subject: [PATCH 15/53] add threshold update test

---
 ribs/archives/_archive_base.py                | 107 ++++++++++--------
 ribs/schedulers/_scheduler.py                 |   6 +-
 .../archives/archive_threshold_update_test.py |  85 ++++++++++++++
 3 files changed, 149 insertions(+), 49 deletions(-)
 create mode 100644 tests/archives/archive_threshold_update_test.py

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 23c44b04c..f6d1e1bed 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -162,10 +162,8 @@ def __init__(self,
 
         # For CMA-MAE
         self._learning_rate = learning_rate
-        # TODO use np.empty
-        self._threshold_arr = np.full(self._cells,
-                                      threshold_min,
-                                      dtype=self.dtype)
+        self._threshold_min = threshold_min
+        self._threshold_arr = np.empty(self._cells, dtype=self.dtype)
 
         self._stats = None
         self._stats_reset()
@@ -269,24 +267,58 @@ def _sum_geometric_series(self, a, r, n):
             return a
         return (a * (1 - pow(r, n))) / (1 - r)
 
-    def _compute_new_thresholds(self, old_threshold, objective_batch_insert):
+    def _compute_new_thresholds(self, threshold_arr, objective_batch,
+                                index_batch, learning_rate):
         """Update thresholds.
 
         Args:
-            old_threshold (float): The threshold of the cells before
-                updating.
-            objective_batch_insert (np.ndarray): 1D array of the objective
-                values of the solution that is inserted into the archive for
-                each cell.
+            threshold_arr (np.ndarray): The threshold of the cells before
+                updating. 1D array.
+            objective_batch (np.ndarray): The objective values of the solution
+                that is inserted into the archive for each cell. 1D array.
+            index_batch (np.ndarray): The archive index of the elements in
+                objective batch.
         Returns:
-            A batch of new thresholds.
+            ``nd.array`` of new thresholds where ``new_threshold_batch[i]`` is
+            the new threshold value at cell ``threshold_update_indices[i]``.
+        Raises:
+            ValueError: ``threshold_arr`` or ``objective_batch`` or
+            ``index_batch`` is empty.
         """
-        k = np.size(objective_batch_insert)
-        geometric_sum = self._sum_geometric_series(1, 1.0 - self._learning_rate,
-                                                   k)
-        return (self._learning_rate * np.sum(objective_batch_insert) *
-                geometric_sum / k) + (old_threshold *
-                                      (1.0 - self._learning_rate)**k)
+        if threshold_arr.size == 0 or objective_batch.size == 0 or index_batch.size == 0:
+            raise ValueError("Cannot compute new thresholds when input array is empty")
+        
+        # Compute the number of objectives inserted into each cell.
+        objective_sizes = aggregate(index_batch,
+                                    objective_batch,
+                                    func="len",
+                                    fill_value=0)
+
+        threshold_update_indices = np.unique(np.where(objective_sizes != 0))
+
+        # Compute the sum of the objectives inserted into each cell.
+        objective_sums = aggregate(index_batch,
+                                   objective_batch,
+                                   func="sum",
+                                   fill_value=np.nan)
+
+        objective_sizes = objective_sizes[threshold_update_indices]
+        objective_sums = objective_sums[threshold_update_indices]
+
+        geometric_sums = np.array([
+            self._sum_geometric_series(1, 1 - learning_rate, k)
+            for k in objective_sizes
+        ])
+
+        a = learning_rate * objective_sums * \
+            (geometric_sums / objective_sizes)
+        b = threshold_arr[threshold_update_indices] * np.power(
+            np.full_like(objective_sizes, 1.0 - learning_rate,
+                         dtype=self.dtype), objective_sizes)
+
+        new_threshold_batch = a + b
+
+        return new_threshold_batch, threshold_update_indices
 
     def clear(self):
         """Removes all elites from the archive.
@@ -487,7 +519,7 @@ def add(self,
         # Since we set the new solutions in old_objective_batch to have
         # value 0.0, the values for new solutions are correct here.
         old_objective_batch[is_new] = 0.0
-        old_threshold_batch[is_new] = 0.0
+        old_threshold_batch[is_new] = self._threshold_min
         value_batch = objective_batch - old_threshold_batch
 
         ## Step 3: Insert solutions into archive. ##
@@ -507,31 +539,11 @@ def add(self,
         metadata_batch_can = metadata_batch[can_insert]
         old_objective_batch_can = old_objective_batch[can_insert]
 
-        def groupby(a, b):
-            # Get argsort indices, to be used to sort a and b in the next steps
-            sidx = b.argsort(kind='mergesort')
-            a_sorted = a[sidx]
-            b_sorted = b[sidx]
-
-            # Get the group limit indices (start, stop of groups)
-            cut_idx = np.flatnonzero(np.r_[True, b_sorted[1:] != b_sorted[:-1],
-                                           True])
-
-            # Split input array with those start, stop ones
-            out = np.array(
-                [a_sorted[i:j] for i, j in zip(cut_idx[:-1], cut_idx[1:])],
-                dtype=object)
-            return out, np.unique(b_sorted)
-
-        # Group the objectives that can be inserted into the archive by cell index.
-        objective_batch_by_cell, group_by_index = groupby(
-            objective_batch_can, index_batch_can)
-
         # Update the thresholds.
-        for objective_batch, index in zip(objective_batch_by_cell,
-                                          group_by_index):
-            self._threshold_arr[index] = self._compute_new_thresholds(
-                self._threshold_arr[index], objective_batch)
+        new_thresholds, update_thresholds_indices = self._compute_new_thresholds(
+            self._threshold_arr, objective_batch_can, index_batch_can,
+            self._learning_rate)
+        self._threshold_arr[update_thresholds_indices] = new_thresholds
 
         # Retrieve indices of solutions that should be inserted into the
         # archive. Currently, multiple solutions may be inserted at each
@@ -644,6 +656,10 @@ def add_single(self, solution, objective, measures, metadata=None):
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
 
+        if not self._occupied_arr[index]:
+            self._objective_arr[index] = 0
+            self._threshold_arr[index] = self._threshold_min
+
         # Note that when learning_rate = 1.0, old_threshold == old_objective.
         old_objective = self._objective_arr[index]
         old_threshold = self._threshold_arr[index]
@@ -657,8 +673,9 @@ def add_single(self, solution, objective, measures, metadata=None):
 
             # Update the threshold.
             if old_threshold < objective:
-                self._threshold_arr[index] = self._compute_new_thresholds(
-                    old_threshold, np.array([objective]))
+                self._threshold_arr[index] = old_threshold * \
+                    (1.0 - self._learning_rate) + \
+                    objective * self._learning_rate
 
             # Insert into the archive.
             self._objective_arr[index] = objective
@@ -676,8 +693,6 @@ def add_single(self, solution, objective, measures, metadata=None):
 
         if status:
             # Update archive stats.
-            if status == 2:
-                old_objective = self.dtype(0.0)
             new_qd_score = self._stats.qd_score + (objective - old_objective)
 
             if self._stats.obj_max is None or objective > self._stats.obj_max:
diff --git a/ribs/schedulers/_scheduler.py b/ribs/schedulers/_scheduler.py
index b1b545d41..adf6663b0 100644
--- a/ribs/schedulers/_scheduler.py
+++ b/ribs/schedulers/_scheduler.py
@@ -233,9 +233,9 @@ def _tell_internal(self,
 
                 # Add solutions to result_archive.
                 if self._result_archive is not None:
-                    self._result_archive.add(self._solution_batch,
-                                             objective_batch, measures_batch,
-                                             metadata_batch)
+                    self._result_archive.add_single(solution,
+                                             objective, measure,
+                                             metadata)
             status_batch = np.asarray(status_batch)
             value_batch = np.asarray(value_batch)
 
diff --git a/tests/archives/archive_threshold_update_test.py b/tests/archives/archive_threshold_update_test.py
new file mode 100644
index 000000000..9f13b5978
--- /dev/null
+++ b/tests/archives/archive_threshold_update_test.py
@@ -0,0 +1,85 @@
+"""Tests for theshold update in archive."""
+import numpy as np
+import pytest
+
+from .conftest import get_archive_data
+
+
+@pytest.fixture
+def data():
+    """Data for grid archive tests."""
+    return get_archive_data("GridArchive")
+
+
+def update_threshold(threshold, f_val, learning_rate):
+    return (1.0 - learning_rate) * threshold + learning_rate * f_val
+
+
+def calc_geom(base_value, exponent):
+    if base_value == 1.0:
+        return exponent
+    top = 1 - base_value**exponent
+    bottom = 1 - base_value
+    return top / bottom
+
+
+def calc_expected_threshold(additions, cell_value, learning_rate):
+    k = len(additions)
+    geom = calc_geom(1.0 - learning_rate, k)
+    f_star = sum(additions)
+    term1 = learning_rate * f_star * geom / k
+    term2 = cell_value * (1.0 - learning_rate)**k
+    print("truth", term1, term2)
+    return term1 + term2
+
+
+@pytest.mark.parametrize("learning_rate", [0, 0.001, 0.01, 0.1, 1])
+def test_consistent_single_update(data, learning_rate):
+    archive = data.archive
+
+    threshold_arr = np.array([-3.1])
+    objective_batch = np.array([0.1, 0.3, 0.9, 400.0, 42.0])
+    index_batch = np.array([0, 0, 0, 0, 0])
+
+    result_test, _ = archive._compute_new_thresholds(threshold_arr, objective_batch,
+                                              index_batch, learning_rate)
+    result_true = calc_expected_threshold(objective_batch, threshold_arr[0],
+                                      learning_rate)
+
+    assert pytest.approx(result_test[0]) == result_true
+
+
+@pytest.mark.parametrize("learning_rate", [0, 0.001, 0.01, 0.1, 1])
+def test_empty_array_raise(data, learning_rate):
+    archive = data.archive
+
+    threshold_arr = np.array([])
+    objective_batch = np.array([])
+    index_batch = np.array([])
+
+    with pytest.raises(ValueError):
+        archive._compute_new_thresholds(threshold_arr, objective_batch,
+                                        index_batch, learning_rate)
+
+
+@pytest.mark.parametrize("learning_rate", [0, 0.001, 0.01, 0.1, 1])
+def test_consistent_multi_update(data, learning_rate):
+    archive = data.archive
+
+    update_size = 3
+    old_threshold = [-3.1]
+    objective = [0.1, 0.3, 0.9, 400.0, 42.0]
+
+    threshold_arr = np.array(old_threshold * update_size)
+    objective_batch = np.array(objective * update_size)
+    index_batch = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2])
+
+    result_test, _ = archive._compute_new_thresholds(threshold_arr,
+                                                     objective_batch,
+                                                     index_batch, learning_rate)
+
+    result_true = calc_expected_threshold(objective, old_threshold[0],
+                                          learning_rate)
+
+    for result in result_test:
+        assert pytest.approx(result) == result_true

From 7530023267358935b1f443687509a8930add47f0 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Sun, 4 Sep 2022 00:38:32 -0700
Subject: [PATCH 16/53] fix pytest bugs

---
 ribs/archives/_archive_base.py      | 45 +++++++++++++++--------------
 tests/archives/grid_archive_test.py | 16 ++++++++++
 2 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index f6d1e1bed..a675d002e 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -163,7 +163,7 @@ def __init__(self,
         # For CMA-MAE
         self._learning_rate = learning_rate
         self._threshold_min = threshold_min
-        self._threshold_arr = np.empty(self._cells, dtype=self.dtype)
+        self._threshold_arr = np.full(self._cells, threshold_min, dtype=self.dtype)
 
         self._stats = None
         self._stats_reset()
@@ -286,8 +286,9 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
             ``index_batch`` is empty.
         """
         if threshold_arr.size == 0 or objective_batch.size == 0 or index_batch.size == 0:
-            raise ValueError("Cannot compute new thresholds when input array is empty")
-        
+            raise ValueError(
+                "Cannot compute new thresholds when input array is empty")
+
         # Compute the number of objectives inserted into each cell.
         objective_sizes = aggregate(index_batch,
                                     objective_batch,
@@ -519,7 +520,7 @@ def add(self,
         # Since we set the new solutions in old_objective_batch to have
         # value 0.0, the values for new solutions are correct here.
         old_objective_batch[is_new] = 0.0
-        old_threshold_batch[is_new] = self._threshold_min
+        old_threshold_batch[is_new] = 0.0
         value_batch = objective_batch - old_threshold_batch
 
         ## Step 3: Insert solutions into archive. ##
@@ -656,23 +657,33 @@ def add_single(self, solution, objective, measures, metadata=None):
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
 
-        if not self._occupied_arr[index]:
-            self._objective_arr[index] = 0
-            self._threshold_arr[index] = self._threshold_min
-
         # Note that when learning_rate = 1.0, old_threshold == old_objective.
         old_objective = self._objective_arr[index]
         old_threshold = self._threshold_arr[index]
 
+        if not self._occupied_arr[index]:
+            self._threshold_arr[index] = self._threshold_min
+            old_threshold = self.dtype(0.0)
+            old_objective = self.dtype(0.0)
+
         was_occupied = self._occupied_arr[index]
         status = 0  # NOT_ADDED
-        if not was_occupied or old_threshold < objective:
-            # Set this index to "occupied" -- important that we do this before
-            # inserting the solution.
-            self._occupied_arr[index] = True
+        if not was_occupied or self._threshold_arr[index] < objective:
+
+            if was_occupied:
+                status = 1  # IMPROVE_EXISTING
+            else:
+                # Set this index to "occupied" -- important that we do this before
+                # inserting the solution.
+                self._occupied_arr[index] = True
+
+                # Tracks a new occupied index.
+                self._occupied_indices[self._num_occupied] = index
+                self._num_occupied += 1
+                status = 2  # NEW
 
             # Update the threshold.
-            if old_threshold < objective:
+            if self._threshold_arr[index] < objective:
                 self._threshold_arr[index] = old_threshold * \
                     (1.0 - self._learning_rate) + \
                     objective * self._learning_rate
@@ -683,14 +694,6 @@ def add_single(self, solution, objective, measures, metadata=None):
             self._solution_arr[index] = solution
             self._metadata_arr[index] = metadata
 
-            if was_occupied:
-                status = 1  # IMPROVE_EXISTING
-            else:
-                # Tracks a new occupied index.
-                self._occupied_indices[self._num_occupied] = index
-                self._num_occupied += 1
-                status = 2  # NEW
-
         if status:
             # Update archive stats.
             new_qd_score = self._stats.qd_score + (objective - old_objective)
diff --git a/tests/archives/grid_archive_test.py b/tests/archives/grid_archive_test.py
index 7d57ed545..bc9d00059 100644
--- a/tests/archives/grid_archive_test.py
+++ b/tests/archives/grid_archive_test.py
@@ -119,6 +119,22 @@ def test_add_single_to_archive(data, use_list):
     assert_archive_elite(data.archive_with_elite, data.solution, data.objective,
                          data.measures, data.grid_indices, data.metadata)
 
+@pytest.mark.parametrize("use_list", [True, False], ids=["list", "ndarray"])
+def test_add_single_to_archive_negative(data, use_list):
+    if use_list:
+        status, value = data.archive.add_single(list(data.solution),
+                                                -data.objective,
+                                                list(data.measures),
+                                                data.metadata)
+    else:
+        status, value = data.archive.add_single(data.solution, -data.objective,
+                                                data.measures, data.metadata)
+
+    assert status == AddStatus.NEW
+    assert np.isclose(value, -data.objective)
+    assert_archive_elite(data.archive_with_elite, data.solution, data.objective,
+                         data.measures, data.grid_indices, data.metadata)
+
 
 def test_add_single_with_low_measures(data):
     measures = np.array([-2, -3])

From 08c6a1d06d92c7dbbb03ddb7ada75b082d7dc28a Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Fri, 16 Sep 2022 19:34:05 -0700
Subject: [PATCH 17/53] Add tests for add with single solutions; fix bugs

---
 ribs/archives/_archive_base.py      |  15 +++--
 tests/archives/archive_base_test.py |  75 +++++++++++++++------
 tests/archives/cvt_archive_test.py  |  54 ++++++++++-----
 tests/archives/grid_archive_test.py | 100 +++++++++++++++++++---------
 4 files changed, 171 insertions(+), 73 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index a675d002e..362867781 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -163,7 +163,9 @@ def __init__(self,
         # For CMA-MAE
         self._learning_rate = learning_rate
         self._threshold_min = threshold_min
-        self._threshold_arr = np.full(self._cells, threshold_min, dtype=self.dtype)
+        self._threshold_arr = np.full(self._cells,
+                                      threshold_min,
+                                      dtype=self.dtype)
 
         self._stats = None
         self._stats_reset()
@@ -313,7 +315,10 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
 
         a = learning_rate * objective_sums * \
             (geometric_sums / objective_sizes)
-        b = threshold_arr[threshold_update_indices] * np.power(
+
+        old_threshold = np.copy(threshold_arr[threshold_update_indices])
+        old_threshold[old_threshold == -np.inf] = 0
+        b = old_threshold * np.power(
             np.full_like(objective_sizes, 1.0 - learning_rate,
                          dtype=self.dtype), objective_sizes)
 
@@ -517,7 +522,7 @@ def add(self,
         status_batch[is_new] = 2
         status_batch[improve_existing] = 1
 
-        # Since we set the new solutions in old_objective_batch to have
+        # Since we set the new solutions in old_threshold_batch to have
         # value 0.0, the values for new solutions are correct here.
         old_objective_batch[is_new] = 0.0
         old_threshold_batch[is_new] = 0.0
@@ -661,15 +666,15 @@ def add_single(self, solution, objective, measures, metadata=None):
         old_objective = self._objective_arr[index]
         old_threshold = self._threshold_arr[index]
 
+        # For new solutions, we set the old_threshold and old_objective to
+        # 0 for computing value.
         if not self._occupied_arr[index]:
-            self._threshold_arr[index] = self._threshold_min
             old_threshold = self.dtype(0.0)
             old_objective = self.dtype(0.0)
 
         was_occupied = self._occupied_arr[index]
         status = 0  # NOT_ADDED
         if not was_occupied or self._threshold_arr[index] < objective:
-
             if was_occupied:
                 status = 1  # IMPROVE_EXISTING
             else:
diff --git a/tests/archives/archive_base_test.py b/tests/archives/archive_base_test.py
index 13847e034..281d60e7a 100644
--- a/tests/archives/archive_base_test.py
+++ b/tests/archives/archive_base_test.py
@@ -46,15 +46,21 @@ def test_iteration():
         assert elite.metadata == data.metadata
 
 
-def test_add_during_iteration():
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_during_iteration(add_mode):
     # Even with just one entry, adding during iteration should still raise an
     # error, just like it does in set.
     data = get_archive_data("GridArchive")
     with pytest.raises(RuntimeError):
         for _ in data.archive_with_elite:
-            data.archive_with_elite.add_single(data.solution,
-                                               data.objective + 1,
-                                               data.measures)
+            if add_mode == "single":
+                data.archive_with_elite.add_single(data.solution,
+                                                   data.objective + 1,
+                                                   data.measures)
+            else:
+                data.archive_with_elite.add([data.solution],
+                                            [data.objective + 1],
+                                            [data.measures])
 
 
 def test_clear_during_iteration():
@@ -64,14 +70,20 @@ def test_clear_during_iteration():
             data.archive_with_elite.clear()
 
 
-def test_clear_and_add_during_iteration():
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_clear_and_add_during_iteration(add_mode):
     data = get_archive_data("GridArchive")
     with pytest.raises(RuntimeError):
         for _ in data.archive_with_elite:
             data.archive_with_elite.clear()
-            data.archive_with_elite.add_single(data.solution,
-                                               data.objective + 1,
-                                               data.measures)
+            if add_mode == "single":
+                data.archive_with_elite.add_single(data.solution,
+                                                   data.objective + 1,
+                                                   data.measures)
+            else:
+                data.archive_with_elite.add([data.solution],
+                                            [data.objective + 1],
+                                            [data.measures])
 
 
 #
@@ -90,13 +102,20 @@ def test_stats_dtype(dtype):
     assert isinstance(data.archive_with_elite.stats.obj_mean, dtype)
 
 
-def test_stats_multiple_add():
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_stats_multiple_add(add_mode):
     archive = GridArchive(solution_dim=3,
                           dims=[10, 20],
                           ranges=[(-1, 1), (-2, 2)])
-    archive.add_single([1, 2, 3], 1.0, [0, 0])
-    archive.add_single([1, 2, 3], 2.0, [0.25, 0.25])
-    archive.add_single([1, 2, 3], 3.0, [-0.25, -0.25])
+    if add_mode == "single":
+        archive.add_single([1, 2, 3], 1.0, [0, 0])
+        archive.add_single([1, 2, 3], 2.0, [0.25, 0.25])
+        archive.add_single([1, 2, 3], 3.0, [-0.25, -0.25])
+    else:
+        solution_batch = [[1, 2, 3]] * 3
+        objective_batch = [1.0, 2.0, 3.0]
+        measures_batch = [[0, 0], [0.25, 0.25], [-0.25, -0.25]]
+        archive.add(solution_batch, objective_batch, measures_batch)
 
     assert archive.stats.num_elites == 3
     assert np.isclose(archive.stats.coverage, 3 / 200)
@@ -105,15 +124,22 @@ def test_stats_multiple_add():
     assert np.isclose(archive.stats.obj_mean, 2.0)
 
 
-def test_stats_add_and_overwrite():
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_stats_add_and_overwrite(add_mode):
     archive = GridArchive(solution_dim=3,
                           dims=[10, 20],
                           ranges=[(-1, 1), (-2, 2)])
-    archive.add_single([1, 2, 3], 1.0, [0, 0])
-    archive.add_single([1, 2, 3], 2.0, [0.25, 0.25])
-    archive.add_single([1, 2, 3], 3.0, [-0.25, -0.25])
-    archive.add_single([1, 2, 3], 5.0,
-                       [0.25, 0.25])  # Overwrites the second add.
+    if add_mode == "single":
+        archive.add_single([1, 2, 3], 1.0, [0, 0])
+        archive.add_single([1, 2, 3], 2.0, [0.25, 0.25])
+        archive.add_single([1, 2, 3], 3.0, [-0.25, -0.25])
+        archive.add_single([1, 2, 3], 5.0,
+                           [0.25, 0.25])  # Overwrites the second add.
+    else:
+        solution_batch = [[1, 2, 3]] * 4
+        objective_batch = [1.0, 2.0, 3.0, 5.0]
+        measures_batch = [[0, 0], [0.25, 0.25], [-0.25, -0.25], [0.25, 0.25]]
+        archive.add(solution_batch, objective_batch, measures_batch)
 
     assert archive.stats.num_elites == 3
     assert np.isclose(archive.stats.coverage, 3 / 200)
@@ -122,7 +148,8 @@ def test_stats_add_and_overwrite():
     assert np.isclose(archive.stats.obj_mean, 3.0)
 
 
-def test_best_elite():
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_best_elite(add_mode):
     archive = GridArchive(solution_dim=3,
                           dims=[10, 20],
                           ranges=[(-1, 1), (-2, 2)])
@@ -131,7 +158,10 @@ def test_best_elite():
     assert archive.best_elite is None
 
     # Add an elite.
-    archive.add_single([1, 2, 3], 1.0, [0, 0])
+    if add_mode == "single":
+        archive.add_single([1, 2, 3], 1.0, [0, 0])
+    else:
+        archive.add([[1, 2, 3]], [1.0], [[0, 0]])
 
     assert np.isclose(archive.best_elite.solution, [1, 2, 3]).all()
     assert np.isclose(archive.best_elite.objective, 1.0).all()
@@ -139,7 +169,10 @@ def test_best_elite():
 
     # Add an elite into the same cell as the previous elite -- best_elite should
     # now be overwritten.
-    archive.add_single([4, 5, 6], 2.0, [0, 0])
+    if add_mode == "single":
+        archive.add_single([4, 5, 6], 2.0, [0, 0])
+    else:
+        archive.add([[4, 5, 6]], [2.0], [[0, 0]])
 
     assert np.isclose(archive.best_elite.solution, [4, 5, 6]).all()
     assert np.isclose(archive.best_elite.objective, 2.0).all()
diff --git a/tests/archives/cvt_archive_test.py b/tests/archives/cvt_archive_test.py
index 46a6f5e26..3ab1fd9a6 100644
--- a/tests/archives/cvt_archive_test.py
+++ b/tests/archives/cvt_archive_test.py
@@ -73,15 +73,23 @@ def test_custom_centroids_bad_shape(use_kd_tree):
 
 
 @pytest.mark.parametrize("use_list", [True, False], ids=["list", "ndarray"])
-def test_add_single_to_archive(data, use_list):
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_single_to_archive(data, use_list, add_mode):
+    solution = data.solution
+    objective = data.objective
+    measures = data.measures
+    metadata = data.metadata
+
     if use_list:
-        status, value = data.archive.add_single(list(data.solution),
-                                                data.objective,
-                                                list(data.measures),
-                                                data.metadata)
+        solution = list(data.solution)
+        measures = list(data.measures)
+
+    if add_mode == "single":
+        status, value = data.archive.add_single(solution, objective, measures,
+                                                metadata)
     else:
-        status, value = data.archive.add_single(data.solution, data.objective,
-                                                data.measures, data.metadata)
+        status, value = data.archive.add([solution], [objective], [measures],
+                                         [metadata])
 
     assert status == AddStatus.NEW
     assert np.isclose(value, data.objective)
@@ -89,32 +97,44 @@ def test_add_single_to_archive(data, use_list):
                          data.measures, data.centroid, data.metadata)
 
 
-def test_add_single_and_overwrite(data):
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_single_and_overwrite(data, add_mode):
     """Test adding a new solution with a higher objective value."""
     arbitrary_sol = data.solution + 1
     arbitrary_metadata = {"foobar": 12}
     high_objective = data.objective + 1.0
 
-    status, value = data.archive_with_elite.add_single(arbitrary_sol,
-                                                       high_objective,
-                                                       data.measures,
-                                                       arbitrary_metadata)
+    if add_mode == "single":
+        status, value = data.archive_with_elite.add_single(
+            arbitrary_sol, high_objective, data.measures, arbitrary_metadata)
+    else:
+        status, value = data.archive_with_elite.add([arbitrary_sol],
+                                                    [high_objective],
+                                                    [data.measures],
+                                                    [arbitrary_metadata])
+
     assert status == AddStatus.IMPROVE_EXISTING
     assert np.isclose(value, high_objective - data.objective)
     assert_archive_elite(data.archive_with_elite, arbitrary_sol, high_objective,
                          data.measures, data.centroid, arbitrary_metadata)
 
 
-def test_add_single_without_overwrite(data):
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_single_without_overwrite(data, add_mode):
     """Test adding a new solution with a lower objective value."""
     arbitrary_sol = data.solution + 1
     arbitrary_metadata = {"foobar": 12}
     low_objective = data.objective - 1.0
 
-    status, value = data.archive_with_elite.add_single(arbitrary_sol,
-                                                       low_objective,
-                                                       data.measures,
-                                                       arbitrary_metadata)
+    if add_mode == "single":
+        status, value = data.archive_with_elite.add_single(
+            arbitrary_sol, low_objective, data.measures, arbitrary_metadata)
+    else:
+        status, value = data.archive_with_elite.add([arbitrary_sol],
+                                                    [low_objective],
+                                                    [data.measures],
+                                                    [arbitrary_metadata])
+
     assert status == AddStatus.NOT_ADDED
     assert np.isclose(value, low_objective - data.objective)
     assert_archive_elite(data.archive_with_elite, data.solution, data.objective,
diff --git a/tests/archives/grid_archive_test.py b/tests/archives/grid_archive_test.py
index bc9d00059..2b047d4c9 100644
--- a/tests/archives/grid_archive_test.py
+++ b/tests/archives/grid_archive_test.py
@@ -104,31 +104,48 @@ def test_properties_are_correct(data):
 
 
 @pytest.mark.parametrize("use_list", [True, False], ids=["list", "ndarray"])
-def test_add_single_to_archive(data, use_list):
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_single_to_archive(data, use_list, add_mode):
+    solution = data.solution
+    objective = data.objective
+    measures = data.measures
+    metadata = data.metadata
+
     if use_list:
-        status, value = data.archive.add_single(list(data.solution),
-                                                data.objective,
-                                                list(data.measures),
-                                                data.metadata)
+        solution = list(data.solution)
+        measures = list(data.measures)
+
+    if add_mode == "single":
+        status, value = data.archive.add_single(solution, objective, measures,
+                                                metadata)
     else:
-        status, value = data.archive.add_single(data.solution, data.objective,
-                                                data.measures, data.metadata)
+        status, value = data.archive.add([solution], [objective], [measures],
+                                         [metadata])
 
     assert status == AddStatus.NEW
     assert np.isclose(value, data.objective)
     assert_archive_elite(data.archive_with_elite, data.solution, data.objective,
                          data.measures, data.grid_indices, data.metadata)
 
+
 @pytest.mark.parametrize("use_list", [True, False], ids=["list", "ndarray"])
-def test_add_single_to_archive_negative(data, use_list):
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_single_to_archive_negative(data, use_list, add_mode):
+    solution = data.solution
+    objective = -data.objective
+    measures = data.measures
+    metadata = data.metadata
+
     if use_list:
-        status, value = data.archive.add_single(list(data.solution),
-                                                -data.objective,
-                                                list(data.measures),
-                                                data.metadata)
+        solution = list(data.solution)
+        measures = list(data.measures)
+
+    if add_mode == "single":
+        status, value = data.archive.add_single(solution, objective, measures,
+                                                metadata)
     else:
-        status, value = data.archive.add_single(data.solution, -data.objective,
-                                                data.measures, data.metadata)
+        status, value = data.archive.add([solution], [objective], [measures],
+                                         [metadata])
 
     assert status == AddStatus.NEW
     assert np.isclose(value, -data.objective)
@@ -136,52 +153,75 @@ def test_add_single_to_archive_negative(data, use_list):
                          data.measures, data.grid_indices, data.metadata)
 
 
-def test_add_single_with_low_measures(data):
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_single_with_low_measures(data, add_mode):
     measures = np.array([-2, -3])
     indices = (0, 0)
-    status, _ = data.archive.add_single(data.solution, data.objective, measures,
-                                        data.metadata)
+    if add_mode == "single":
+        status, _ = data.archive.add_single(data.solution, data.objective,
+                                            measures, data.metadata)
+    else:
+        status, _ = data.archive.add([data.solution], [data.objective],
+                                     [measures], [data.metadata])
+
     assert status
     assert_archive_elite(data.archive, data.solution, data.objective, measures,
                          indices, data.metadata)
 
 
-def test_add_single_with_high_measures(data):
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_single_with_high_measures(data, add_mode):
     measures = np.array([2, 3])
     indices = (9, 19)
-    status, _ = data.archive.add_single(data.solution, data.objective, measures,
-                                        data.metadata)
+    if add_mode == "single":
+        status, _ = data.archive.add_single(data.solution, data.objective,
+                                            measures, data.metadata)
+    else:
+        status, _ = data.archive.add([data.solution], [data.objective],
+                                     [measures], [data.metadata])
     assert status
     assert_archive_elite(data.archive, data.solution, data.objective, measures,
                          indices, data.metadata)
 
 
-def test_add_single_and_overwrite(data):
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_single_and_overwrite(data, add_mode):
     """Test adding a new solution with a higher objective value."""
     arbitrary_sol = data.solution + 1
     arbitrary_metadata = {"foobar": 12}
     high_objective = data.objective + 1.0
 
-    status, value = data.archive_with_elite.add_single(arbitrary_sol,
-                                                       high_objective,
-                                                       data.measures,
-                                                       arbitrary_metadata)
+    if add_mode == "single":
+        status, value = data.archive_with_elite.add_single(
+            arbitrary_sol, high_objective, data.measures, arbitrary_metadata)
+    else:
+        status, value = data.archive_with_elite.add([arbitrary_sol],
+                                                    [high_objective],
+                                                    [data.measures],
+                                                    [arbitrary_metadata])
+
     assert status == AddStatus.IMPROVE_EXISTING
     assert np.isclose(value, high_objective - data.objective)
     assert_archive_elite(data.archive_with_elite, arbitrary_sol, high_objective,
                          data.measures, data.grid_indices, arbitrary_metadata)
 
 
-def test_add_single_without_overwrite(data):
+@pytest.mark.parametrize("add_mode", ["single", "batch"])
+def test_add_single_without_overwrite(data, add_mode):
     """Test adding a new solution with a lower objective value."""
     arbitrary_sol = data.solution + 1
     arbitrary_metadata = {"foobar": 12}
     low_objective = data.objective - 1.0
 
-    status, value = data.archive_with_elite.add_single(arbitrary_sol,
-                                                       low_objective,
-                                                       data.measures,
-                                                       arbitrary_metadata)
+    if add_mode == "single":
+        status, value = data.archive_with_elite.add_single(
+            arbitrary_sol, low_objective, data.measures, arbitrary_metadata)
+    else:
+        status, value = data.archive_with_elite.add([arbitrary_sol],
+                                                    [low_objective],
+                                                    [data.measures],
+                                                    [arbitrary_metadata])
+
     assert status == AddStatus.NOT_ADDED
     assert np.isclose(value, low_objective - data.objective)
     assert_archive_elite(data.archive_with_elite, data.solution, data.objective,

From 393b0ea2b3a3ff8e0105d3b47552db5813c8ebe3 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Fri, 16 Sep 2022 19:41:47 -0700
Subject: [PATCH 18/53] raise error when new threhold is nan.

---
 ribs/archives/_archive_base.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 362867781..58e880dc8 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -323,6 +323,9 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
                          dtype=self.dtype), objective_sizes)
 
         new_threshold_batch = a + b
+        if np.any(new_threshold_batch == np.nan):
+            raise ValueError(
+                "Computed new threshold to be nan. Please raise an issue at https://github.com/icaros-usc/pyribs/issues.")
 
         return new_threshold_batch, threshold_update_indices
 

From 4554be0b97d70d55be2645be50fc84c2209a40a1 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Fri, 16 Sep 2022 20:05:02 -0700
Subject: [PATCH 19/53] history

---
 HISTORY.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/HISTORY.md b/HISTORY.md
index 48d879097..5859cbad4 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -8,6 +8,7 @@
 
 - Implement CMA-MAE archive thresholds (#256)
   - Revive the old implementation of `add_single` removed in (#221)
+  - Add separate tests for `add_single` and `add` with single solution
 - Add restart timer to `EvolutionStrategyEmitter` and `GradientAborescenceEmitter`(#255)
 - Rename fields and update documentation (#249, #250)
   - **Backwards-incompatible:** rename `Optimizer` to `Scheduler`

From 86cda6ccbb362740efcf1db60aa5ca9797dffbc6 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 27 Sep 2022 17:15:05 -0700
Subject: [PATCH 20/53] More detailed scheduler error msg

---
 ribs/schedulers/_scheduler.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/ribs/schedulers/_scheduler.py b/ribs/schedulers/_scheduler.py
index adf6663b0..bb3ddaf0b 100644
--- a/ribs/schedulers/_scheduler.py
+++ b/ribs/schedulers/_scheduler.py
@@ -1,6 +1,5 @@
 """Provides the Scheduler."""
 import numpy as np
-
 from ribs.emitters import DQDEmitterBase
 
 
@@ -80,7 +79,10 @@ def __init__(self,
                              f"it was '{add_mode}'")
 
         if archive is result_archive:
-            raise ValueError("archive has same id as result_archive")
+            raise ValueError("`archive` has same id as `result_archive` -- "
+                             "Note that `Scheduler.result_archive` already "
+                             "defaults to be the same as `archive` if you pass "
+                             "`result_archive=None`")
 
         self._archive = archive
         self._emitters = emitters
@@ -233,9 +235,8 @@ def _tell_internal(self,
 
                 # Add solutions to result_archive.
                 if self._result_archive is not None:
-                    self._result_archive.add_single(solution,
-                                             objective, measure,
-                                             metadata)
+                    self._result_archive.add_single(solution, objective,
+                                                    measure, metadata)
             status_batch = np.asarray(status_batch)
             value_batch = np.asarray(value_batch)
 

From 67f6e722b2594c0b9fb26cda11de57077f74c9b2 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 4 Oct 2022 16:34:35 -0700
Subject: [PATCH 21/53] Batch threshold computation

---
 ribs/archives/_archive_base.py | 58 +++++++++++++++++-----------------
 1 file changed, 29 insertions(+), 29 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 58e880dc8..337a4222c 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -4,7 +4,6 @@
 
 import numpy as np
 from numpy_groupies import aggregate_nb as aggregate
-
 from ribs._utils import (check_1d_shape, check_batch_shape, check_is_1d,
                          check_solution_batch_dim)
 from ribs.archives._archive_data_frame import ArchiveDataFrame
@@ -263,12 +262,6 @@ def _stats_reset(self):
         self._stats = ArchiveStats(0, self.dtype(0.0), self.dtype(0.0), None,
                                    None)
 
-    def _sum_geometric_series(self, a, r, n):
-        """Compute the sum of a geometric series."""
-        if r == 1.0:
-            return a
-        return (a * (1 - pow(r, n))) / (1 - r)
-
     def _compute_new_thresholds(self, threshold_arr, objective_batch,
                                 index_batch, learning_rate):
         """Update thresholds.
@@ -281,51 +274,58 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
             index_batch (np.ndarray): The archive index of the elements in
                 objective batch.
         Returns:
-            ``nd.array`` of new thresholds where ``new_threshold_batch[i]`` is
-            the new threshold value at cell ``threshold_update_indices[i]``.
-        Raises:
-            ValueError: ``threshold_arr`` or ``objective_batch`` or
-            ``index_batch`` is empty.
+            `new_threshold_batch` (A self.dtype array of new
+            thresholds) and `threshold_update_indices` (A boolean
+            array indicating which entries in `threshold_arr` should
+            be updated.
         """
-        if threshold_arr.size == 0 or objective_batch.size == 0 or index_batch.size == 0:
-            raise ValueError(
-                "Cannot compute new thresholds when input array is empty")
+        if objective_batch.size == 0 or index_batch.size == 0:
+            return np.array([], dtype=self.dtype), np.array([], dtype=bool)
 
         # Compute the number of objectives inserted into each cell.
         objective_sizes = aggregate(index_batch,
                                     objective_batch,
                                     func="len",
-                                    fill_value=0)
+                                    fill_value=0,
+                                    size=threshold_arr.size)
 
-        threshold_update_indices = np.unique(np.where(objective_sizes != 0))
+        # These indices are with respect to the archive, so we can directly pass
+        # them to threshold_arr.
+        threshold_update_indices = objective_sizes > 0
 
         # Compute the sum of the objectives inserted into each cell.
         objective_sums = aggregate(index_batch,
                                    objective_batch,
                                    func="sum",
-                                   fill_value=np.nan)
+                                   fill_value=np.nan,
+                                   size=threshold_arr.size)
 
+        # Throw away indices that we do not care about.
         objective_sizes = objective_sizes[threshold_update_indices]
         objective_sums = objective_sums[threshold_update_indices]
 
-        geometric_sums = np.array([
-            self._sum_geometric_series(1, 1 - learning_rate, k)
-            for k in objective_sizes
-        ])
+        # Sum of geometric series (1 - learning_rate)^i from i = 0 to i = n - 1
+        # See https://en.wikipedia.org/wiki/Geometric_series#Sum
+        ratio = self.dtype(1.0 - learning_rate)
+        if ratio == 1.0:
+            geometric_sums = objective_sizes
+        else:
+            geometric_sums = (1 - ratio**objective_sizes) / (1 - ratio)
 
-        a = learning_rate * objective_sums * \
-            (geometric_sums / objective_sizes)
+        update = (learning_rate * (objective_sums / objective_sizes) *
+                  geometric_sums)
 
         old_threshold = np.copy(threshold_arr[threshold_update_indices])
+        # TODO: Fix this based on new CMA-ME behavior if needed?
         old_threshold[old_threshold == -np.inf] = 0
-        b = old_threshold * np.power(
-            np.full_like(objective_sizes, 1.0 - learning_rate,
-                         dtype=self.dtype), objective_sizes)
+        prev = old_threshold * ratio**objective_sizes
 
-        new_threshold_batch = a + b
+        new_threshold_batch = prev + update
+        # TODO: When does this happen? Remove if possible.
         if np.any(new_threshold_batch == np.nan):
             raise ValueError(
-                "Computed new threshold to be nan. Please raise an issue at https://github.com/icaros-usc/pyribs/issues.")
+                "Computed new threshold to be nan. Please raise an issue at https://github.com/icaros-usc/pyribs/issues."
+            )
 
         return new_threshold_batch, threshold_update_indices
 

From 41750f2bc4a2aa52cc821eeb2791351b3aae774b Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Tue, 11 Oct 2022 01:11:42 -0700
Subject: [PATCH 22/53] fix requested changes

---
 examples/sphere.py                            | 10 ++++--
 ribs/archives/_archive_base.py                | 35 ++++++++++---------
 ribs/archives/_cvt_archive.py                 |  7 ++++
 .../archives/archive_threshold_update_test.py | 13 -------
 4 files changed, 32 insertions(+), 33 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index 7ddd768e6..aa1adfb76 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -174,6 +174,10 @@ def create_scheduler(algorithm,
     batch_size = 37
     num_emitters = 15
     mode = "batch"
+    threshold_min = -np.inf  # default
+
+    if algorithm in ["cma_mae", "cma_maega"]:
+        threshold_min = 0
 
     # Create archive.
     if algorithm in ["cvt_map_elites", "line_cvt_map_elites"]:
@@ -187,7 +191,7 @@ def create_scheduler(algorithm,
                               dims=archive_dims,
                               ranges=bounds,
                               learning_rate=learning_rate,
-                              threshold_min=0,
+                              threshold_min=threshold_min,
                               seed=seed)
 
     # Create result archive.
@@ -369,9 +373,9 @@ def sphere_main(algorithm,
     """
     # Use default dim for each algorithm.
     if dim is None:
-        if algorithm in ["cma_mega", "cma_mega_adam"]:
+        if algorithm in ["cma_mega", "cma_mega_adam", "cma_maega"]:
             dim = 1_000
-        elif algorithm in ["cma_mae", "cma_maega"]:
+        elif algorithm in ["cma_mae"]:
             dim = 100
         elif algorithm in [
                 "map_elites", "line_map_elites", "cma_me_imp", "cma_me_imp_mu",
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 337a4222c..7160b27cf 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 from numpy_groupies import aggregate_nb as aggregate
+
 from ribs._utils import (check_1d_shape, check_batch_shape, check_is_1d,
                          check_solution_batch_dim)
 from ribs.archives._archive_data_frame import ArchiveDataFrame
@@ -159,9 +160,13 @@ def __init__(self,
                                       dtype=self.dtype)
         self._metadata_arr = np.empty(self._cells, dtype=object)
 
-        # For CMA-MAE
-        self._learning_rate = learning_rate
-        self._threshold_min = threshold_min
+        # threshold min can only be -np.inf if the learning rate is 1.0
+        if learning_rate != 1.0 and threshold_min == -np.inf:
+            raise ValueError(
+                "If learning_rate != 1.0, threshold min cannot be -np.inf (default)."
+            )
+        self._learning_rate = self._dtype(learning_rate)
+        self._threshold_min = self._dtype(threshold_min)
         self._threshold_arr = np.full(self._cells,
                                       threshold_min,
                                       dtype=self.dtype)
@@ -321,11 +326,6 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
         prev = old_threshold * ratio**objective_sizes
 
         new_threshold_batch = prev + update
-        # TODO: When does this happen? Remove if possible.
-        if np.any(new_threshold_batch == np.nan):
-            raise ValueError(
-                "Computed new threshold to be nan. Please raise an issue at https://github.com/icaros-usc/pyribs/issues."
-            )
 
         return new_threshold_batch, threshold_update_indices
 
@@ -528,7 +528,7 @@ def add(self,
         # Since we set the new solutions in old_threshold_batch to have
         # value 0.0, the values for new solutions are correct here.
         old_objective_batch[is_new] = 0.0
-        old_threshold_batch[is_new] = 0.0
+        old_threshold_batch[is_new] = 0.0 if self._threshold_min == -np.inf else self._threshold_min
         value_batch = objective_batch - old_threshold_batch
 
         ## Step 3: Insert solutions into archive. ##
@@ -658,22 +658,23 @@ def add_single(self, solution, objective, measures, metadata=None):
         self._state["add"] += 1
 
         solution = np.asarray(solution)
+        check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
+
         objective = self.dtype(objective)
-        measures = np.asarray(measures)
-        index = self.index_of_single(measures)
 
-        check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
+        measures = np.asarray(measures)
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
+        index = self.index_of_single(measures)
 
         # Note that when learning_rate = 1.0, old_threshold == old_objective.
         old_objective = self._objective_arr[index]
         old_threshold = self._threshold_arr[index]
 
         # For new solutions, we set the old_threshold and old_objective to
-        # 0 for computing value.
+        # 0 for computing value only if threshold min is not set.
         if not self._occupied_arr[index]:
-            old_threshold = self.dtype(0.0)
-            old_objective = self.dtype(0.0)
+            old_objective = self.dtype(
+                0.0) if old_threshold == -np.inf else self._threshold_min
 
         was_occupied = self._occupied_arr[index]
         status = 0  # NOT_ADDED
@@ -692,7 +693,7 @@ def add_single(self, solution, objective, measures, metadata=None):
 
             # Update the threshold.
             if self._threshold_arr[index] < objective:
-                self._threshold_arr[index] = old_threshold * \
+                self._threshold_arr[index] = old_objective * \
                     (1.0 - self._learning_rate) + \
                     objective * self._learning_rate
 
@@ -726,7 +727,7 @@ def add_single(self, solution, objective, measures, metadata=None):
                 obj_mean=new_qd_score / self.dtype(len(self)),
             )
 
-        return status, self.dtype(objective - old_threshold)
+        return status, self.dtype(objective - old_objective)
 
     def elites_with_measures(self, measures_batch):
         """Retrieves the elites with measures in the same cells as the measures
diff --git a/ribs/archives/_cvt_archive.py b/ribs/archives/_cvt_archive.py
index b97691d71..fde2f0391 100644
--- a/ribs/archives/_cvt_archive.py
+++ b/ribs/archives/_cvt_archive.py
@@ -60,6 +60,9 @@ class CVTArchive(ArchiveBase):
             (inclusive), and the second dimension should have bounds
             :math:`[-2,2]` (inclusive). ``ranges`` should be the same length as
             ``dims``.
+        learning_rate (float): The learning rate of the archive. Described in
+            `Fontaine 2022 <https://arxiv.org/abs/2205.10752>`_.
+        threshold_min (float): The default threshold value for all the cells.
         seed (int): Value to seed the random number generator as well as
             :func:`~sklearn.cluster.k_means`. Set to None to avoid a fixed seed.
         dtype (str or data-type): Data type of the solutions, objectives,
@@ -93,6 +96,8 @@ def __init__(self,
                  solution_dim,
                  cells,
                  ranges,
+                 learning_rate=1.0,
+                 threshold_min=-np.inf,
                  seed=None,
                  dtype=np.float64,
                  samples=100_000,
@@ -106,6 +111,8 @@ def __init__(self,
             solution_dim=solution_dim,
             cells=cells,
             measure_dim=len(ranges),
+            learning_rate=learning_rate,
+            threshold_min=threshold_min,
             seed=seed,
             dtype=dtype,
         )
diff --git a/tests/archives/archive_threshold_update_test.py b/tests/archives/archive_threshold_update_test.py
index 9f13b5978..d1be8b5e5 100644
--- a/tests/archives/archive_threshold_update_test.py
+++ b/tests/archives/archive_threshold_update_test.py
@@ -49,19 +49,6 @@ def test_consistent_single_update(data, learning_rate):
     assert pytest.approx(result_test[0]) == result_true
 
 
-@pytest.mark.parametrize("learning_rate", [0, 0.001, 0.01, 0.1, 1])
-def test_empty_array_raise(data, learning_rate):
-    archive = data.archive
-
-    threshold_arr = np.array([])
-    objective_batch = np.array([])
-    index_batch = np.array([])
-
-    with pytest.raises(ValueError):
-        archive._compute_new_thresholds(threshold_arr, objective_batch,
-                                        index_batch, learning_rate)
-
-
 @pytest.mark.parametrize("learning_rate", [0, 0.001, 0.01, 0.1, 1])
 def test_consistent_multi_update(data, learning_rate):
     archive = data.archive

From cfb997f7025e6518cb763fbec191d3bddca69067 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Tue, 11 Oct 2022 01:35:24 -0700
Subject: [PATCH 23/53] fix error merge conflict

---
 examples/sphere.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index a03da4078..b52bf122d 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -441,7 +441,7 @@ def sphere_main(algorithm,
     }
 
     non_logging_time = 0.0
-    save_heatmap(archive, str(outdir / f"{name}_heatmap_{0:05d}.png"))
+    save_heatmap(result_archive, str(outdir / f"{name}_heatmap_{0:05d}.png"))
 
     for itr in tqdm.trange(1, itrs + 1):
         itr_start = time.time()
@@ -464,19 +464,21 @@ def sphere_main(algorithm,
         final_itr = itr == itrs
         if itr % log_freq == 0 or final_itr:
             if final_itr:
-                archive.as_pandas(include_solutions=final_itr).to_csv(
+                result_archive.as_pandas(include_solutions=final_itr).to_csv(
                     outdir / f"{name}_archive.csv")
 
             # Record and display metrics.
             metrics["QD Score"]["x"].append(itr)
-            metrics["QD Score"]["y"].append(archive.stats.qd_score)
+            metrics["QD Score"]["y"].append(result_archive.stats.qd_score)
             metrics["Archive Coverage"]["x"].append(itr)
-            metrics["Archive Coverage"]["y"].append(archive.stats.coverage)
+            metrics["Archive Coverage"]["y"].append(
+                result_archive.stats.coverage)
             print(f"Iteration {itr} | Archive Coverage: "
                   f"{metrics['Archive Coverage']['y'][-1] * 100:.3f}% "
                   f"QD Score: {metrics['QD Score']['y'][-1]:.3f}")
 
-            save_heatmap(archive, str(outdir / f"{name}_heatmap_{itr:05d}.png"))
+            save_heatmap(result_archive,
+                         str(outdir / f"{name}_heatmap_{itr:05d}.png"))
 
     # Plot metrics.
     print(f"Algorithm Time (Excludes Logging and Setup): {non_logging_time}s")

From 105f2ef6d1347552143ac8a89d0e662557bd0029 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Tue, 11 Oct 2022 02:00:57 -0700
Subject: [PATCH 24/53] add github action tests

---
 tests/examples.sh | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/tests/examples.sh b/tests/examples.sh
index 65c741d3d..a1b632922 100644
--- a/tests/examples.sh
+++ b/tests/examples.sh
@@ -21,16 +21,28 @@ export OMP_NUM_THREADS=1
 
 # sphere.py - CVT excluded since it takes a while to build the archive.
 SPHERE_OUTPUT="${TMPDIR}/sphere_output"
-python examples/sphere.py map_elites 20 10 "${SPHERE_OUTPUT}"
-python examples/sphere.py line_map_elites 20 10 "${SPHERE_OUTPUT}"
+python examples/sphere.py map_elites --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py line_map_elites --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
 # python examples/sphere.py cvt_map_elites 20 10 "${SPHERE_OUTPUT}"
 # python examples/sphere.py line_cvt_map_elites 20 10 "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_imp 20 10 "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_imp_mu 20 10 "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_rd 20 10 "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_rd_mu 20 10 "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_opt 20 10 "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_mixed 20 10 "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_imp --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_imp_mu --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_rd --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_rd_mu --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_opt --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_mixed --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_mega --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_mega_adam --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 1 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0.1 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0.01 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0.001 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 1 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0.1 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0.01 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0.001 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0 --outdir "${SPHERE_OUTPUT}"
 
 # lunar_lander.py
 LUNAR_LANDER_OUTPUT="${TMPDIR}/lunar_lander_output"

From 088723802dcd7188bc16d1aa63ff84e555b23f9a Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Tue, 11 Oct 2022 02:05:31 -0700
Subject: [PATCH 25/53] Simplify examples.sh by removing redundence default
 argument.

---
 tests/examples.sh | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/tests/examples.sh b/tests/examples.sh
index a1b632922..d2d176482 100644
--- a/tests/examples.sh
+++ b/tests/examples.sh
@@ -21,23 +21,26 @@ export OMP_NUM_THREADS=1
 
 # sphere.py - CVT excluded since it takes a while to build the archive.
 SPHERE_OUTPUT="${TMPDIR}/sphere_output"
-python examples/sphere.py map_elites --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py line_map_elites --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
-# python examples/sphere.py cvt_map_elites 20 10 "${SPHERE_OUTPUT}"
-# python examples/sphere.py line_cvt_map_elites 20 10 "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_imp --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_imp_mu --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_rd --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_rd_mu --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_opt --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_me_mixed --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py map_elites --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py line_map_elites --itrs 10 --outdir "${SPHERE_OUTPUT}"
+# python examples/sphere.py cvt_map_elites 10 "${SPHERE_OUTPUT}"
+# python examples/sphere.py line_cvt_map_elites 10 "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_imp --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_imp_mu --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_rd --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_rd_mu --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_opt --itrs 10 --outdir "${SPHERE_OUTPUT}"
+python examples/sphere.py cma_me_mixed --itrs 10 --outdir "${SPHERE_OUTPUT}"
+
 python examples/sphere.py cma_mega --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_mega_adam --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
+
 python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 1 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0.1 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0.01 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0.001 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0 --outdir "${SPHERE_OUTPUT}"
+
 python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 1 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0.1 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0.01 --outdir "${SPHERE_OUTPUT}"

From b0861205a13f7968f7e8f1a068c56bc07f4690aa Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 11 Oct 2022 17:04:17 -0700
Subject: [PATCH 26/53] Tweak comments

---
 examples/sphere.py | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index 7ddd768e6..08c35ae51 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -43,16 +43,13 @@
 4500 iterations for a total of 15 * 37 * 4500 ~= 2.5M evaluations.
 
 Notes:
-    - `cma_mega` and `cma_mega_adam` use only one emitter and run for 10,000
-      iterations. This is to be consistent with the paper (`Fontaine 2021
-      <https://arxiv.org/abs/2106.03894>`_) in which these algorithms were
-      proposed.
-    - `cma_mae` and `cma_maega` run for 10,000 iterations as well.
-    - CVTArchive in this example uses 10,000 cells, as opposed to the 250,000
-      (500x500) in the GridArchive, so it is not fair to directly compare
-      `cvt_map_elites` and `line_cvt_map_elites` to the other algorithms.
-      However, the other algorithms may be fairly compared because they use the
-      same archive.
+- `cma_mega` and `cma_mega_adam` use only one emitter and run for 10,000
+  iterations. This is to be consistent with the paper (`Fontaine 2021
+  <https://arxiv.org/abs/2106.03894>`_) in which these algorithms were proposed.
+- `cma_mae` and `cma_maega` run for 10,000 iterations as well.
+- CVTArchive in this example uses 10,000 cells, as opposed to the 250,000
+  (500x500) in the GridArchive, so it is not fair to directly compare
+  `cvt_map_elites` and `line_cvt_map_elites` to the other algorithms.
 
 Outputs are saved in the `sphere_output/` directory by default. The archive is
 saved as a CSV named `{algorithm}_{dim}_archive.csv`, while snapshots of the
@@ -86,7 +83,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from alive_progress import alive_bar
-
 from ribs.archives import CVTArchive, GridArchive
 from ribs.emitters import (EvolutionStrategyEmitter, GaussianEmitter,
                            GradientAborescenceEmitter, IsoLineEmitter)

From 26424c2cb83f4cb60f961528ec056c215654f378 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 11 Oct 2022 17:05:20 -0700
Subject: [PATCH 27/53] Format strings

---
 examples/sphere.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index 08c35ae51..40ebcd83c 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -317,8 +317,9 @@ def create_scheduler(algorithm,
         ]
 
     print(
-        f"Created Scheduler for {algorithm} with learning rate {learning_rate} and add mode {mode}, "
-        f"using solution dims {solution_dim} and archive dims {archive_dims}.")
+        f"Created Scheduler for {algorithm} with learning rate {learning_rate} "
+        f"and add mode {mode}, using solution dims {solution_dim} and archive "
+        f"dims {archive_dims}.")
     return Scheduler(archive, emitters, result_archive, add_mode=mode)
 
 

From 8d6eb2806708a79b2108582e2b2b48887dc96f52 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 11 Oct 2022 17:12:03 -0700
Subject: [PATCH 28/53] Choose new sphere tests

---
 tests/examples.sh | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/tests/examples.sh b/tests/examples.sh
index d2d176482..3ff0fb88c 100644
--- a/tests/examples.sh
+++ b/tests/examples.sh
@@ -19,12 +19,15 @@ fi
 export OPENBLAS_NUM_THREADS=1
 export OMP_NUM_THREADS=1
 
-# sphere.py - CVT excluded since it takes a while to build the archive.
+# sphere.py
 SPHERE_OUTPUT="${TMPDIR}/sphere_output"
 python examples/sphere.py map_elites --itrs 10 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py line_map_elites --itrs 10 --outdir "${SPHERE_OUTPUT}"
+
+# CVT excluded since it takes a while to build the archive.
 # python examples/sphere.py cvt_map_elites 10 "${SPHERE_OUTPUT}"
 # python examples/sphere.py line_cvt_map_elites 10 "${SPHERE_OUTPUT}"
+
 python examples/sphere.py cma_me_imp --itrs 10 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_me_imp_mu --itrs 10 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_me_rd --itrs 10 --outdir "${SPHERE_OUTPUT}"
@@ -35,17 +38,8 @@ python examples/sphere.py cma_me_mixed --itrs 10 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_mega --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_mega_adam --dim 20 --itrs 10 --outdir "${SPHERE_OUTPUT}"
 
-python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 1 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0.1 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0.01 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0.001 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_mae --dim 20 --itrs 10 --learning_rate 0 --outdir "${SPHERE_OUTPUT}"
-
-python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 1 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0.1 --outdir "${SPHERE_OUTPUT}"
 python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0.01 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0.001 --outdir "${SPHERE_OUTPUT}"
-python examples/sphere.py cma_maega --dim 20 --itrs 10 --learning_rate 0 --outdir "${SPHERE_OUTPUT}"
 
 # lunar_lander.py
 LUNAR_LANDER_OUTPUT="${TMPDIR}/lunar_lander_output"

From d52a1b370d8542cd7e83eafa9a1e07b831470b22 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 11 Oct 2022 17:15:43 -0700
Subject: [PATCH 29/53] Remove print

---
 tests/archives/archive_threshold_update_test.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/archives/archive_threshold_update_test.py b/tests/archives/archive_threshold_update_test.py
index d1be8b5e5..cc5ac5446 100644
--- a/tests/archives/archive_threshold_update_test.py
+++ b/tests/archives/archive_threshold_update_test.py
@@ -29,7 +29,6 @@ def calc_expected_threshold(additions, cell_value, learning_rate):
     f_star = sum(additions)
     term1 = learning_rate * f_star * geom / k
     term2 = cell_value * (1.0 - learning_rate)**k
-    print("truth", term1, term2)
     return term1 + term2
 
 
@@ -41,10 +40,11 @@ def test_consistent_single_update(data, learning_rate):
     objective_batch = np.array([0.1, 0.3, 0.9, 400.0, 42.0])
     index_batch = np.array([0, 0, 0, 0, 0])
 
-    result_test, _ = archive._compute_new_thresholds(threshold_arr, objective_batch,
-                                              index_batch, learning_rate)
+    result_test, _ = archive._compute_new_thresholds(threshold_arr,
+                                                     objective_batch,
+                                                     index_batch, learning_rate)
     result_true = calc_expected_threshold(objective_batch, threshold_arr[0],
-                                      learning_rate)
+                                          learning_rate)
 
     assert pytest.approx(result_test[0]) == result_true
 

From 93940a3a0bce9c5207cb5bec0a088f6bd5362a58 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 11 Oct 2022 17:17:30 -0700
Subject: [PATCH 30/53] Change test name

---
 tests/archives/archive_threshold_update_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/archives/archive_threshold_update_test.py b/tests/archives/archive_threshold_update_test.py
index cc5ac5446..8ce491f7f 100644
--- a/tests/archives/archive_threshold_update_test.py
+++ b/tests/archives/archive_threshold_update_test.py
@@ -33,7 +33,7 @@ def calc_expected_threshold(additions, cell_value, learning_rate):
 
 
 @pytest.mark.parametrize("learning_rate", [0, 0.001, 0.01, 0.1, 1])
-def test_consistent_single_update(data, learning_rate):
+def test_threshold_update_for_one_cell(data, learning_rate):
     archive = data.archive
 
     threshold_arr = np.array([-3.1])

From d3edccd0ffe8aba3ea0f2d54a9e1d115458ec562 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 11 Oct 2022 17:19:29 -0700
Subject: [PATCH 31/53] more naming, ignore pylint

---
 tests/archives/archive_threshold_update_test.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/archives/archive_threshold_update_test.py b/tests/archives/archive_threshold_update_test.py
index 8ce491f7f..966335d40 100644
--- a/tests/archives/archive_threshold_update_test.py
+++ b/tests/archives/archive_threshold_update_test.py
@@ -4,6 +4,8 @@
 
 from .conftest import get_archive_data
 
+# pylint: disable = redefined-outer-name
+
 
 @pytest.fixture
 def data():
@@ -40,6 +42,7 @@ def test_threshold_update_for_one_cell(data, learning_rate):
     objective_batch = np.array([0.1, 0.3, 0.9, 400.0, 42.0])
     index_batch = np.array([0, 0, 0, 0, 0])
 
+    # pylint: disable = protected-access
     result_test, _ = archive._compute_new_thresholds(threshold_arr,
                                                      objective_batch,
                                                      index_batch, learning_rate)
@@ -50,7 +53,7 @@ def test_threshold_update_for_one_cell(data, learning_rate):
 
 
 @pytest.mark.parametrize("learning_rate", [0, 0.001, 0.01, 0.1, 1])
-def test_consistent_multi_update(data, learning_rate):
+def test_threshold_update_for_multiple_cells(data, learning_rate):
     archive = data.archive
 
     update_size = 3
@@ -61,6 +64,7 @@ def test_consistent_multi_update(data, learning_rate):
     objective_batch = np.array(objective * update_size)
     index_batch = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2])
 
+    # pylint: disable = protected-access
     result_test, _ = archive._compute_new_thresholds(threshold_arr,
                                                      objective_batch,
                                                      index_batch, learning_rate)

From 05f28de54a2ba1c9714edb34382f1b4b6d4b86e9 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 11 Oct 2022 17:29:04 -0700
Subject: [PATCH 32/53] More complex multiple cell threshold test

---
 ribs/archives/_archive_base.py                |  7 +++++--
 .../archives/archive_threshold_update_test.py | 21 ++++++++++---------
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 7160b27cf..bb9992149 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -275,7 +275,9 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
             threshold_arr (np.ndarray): The threshold of the cells before
                 updating. 1D array.
             objective_batch (np.ndarray): The objective values of the solution
-                that is inserted into the archive for each cell. 1D array.
+                that is inserted into the archive for each cell. 1D array. We
+                assume that the objective values are all higher than the
+                thresholds of their respective cells.
             index_batch (np.ndarray): The archive index of the elements in
                 objective batch.
         Returns:
@@ -528,7 +530,8 @@ def add(self,
         # Since we set the new solutions in old_threshold_batch to have
         # value 0.0, the values for new solutions are correct here.
         old_objective_batch[is_new] = 0.0
-        old_threshold_batch[is_new] = 0.0 if self._threshold_min == -np.inf else self._threshold_min
+        old_threshold_batch[
+            is_new] = 0.0 if self._threshold_min == -np.inf else self._threshold_min
         value_batch = objective_batch - old_threshold_batch
 
         ## Step 3: Insert solutions into archive. ##
diff --git a/tests/archives/archive_threshold_update_test.py b/tests/archives/archive_threshold_update_test.py
index 966335d40..2a3e45037 100644
--- a/tests/archives/archive_threshold_update_test.py
+++ b/tests/archives/archive_threshold_update_test.py
@@ -56,12 +56,11 @@ def test_threshold_update_for_one_cell(data, learning_rate):
 def test_threshold_update_for_multiple_cells(data, learning_rate):
     archive = data.archive
 
-    update_size = 3
-    old_threshold = [-3.1]
-    objective = [0.1, 0.3, 0.9, 400.0, 42.0]
-
-    threshold_arr = np.array(old_threshold * update_size)
-    objective_batch = np.array(objective * update_size)
+    threshold_arr = np.array([-3.1, 0.4, 2.9])
+    objective_batch = np.array([
+        0.1, 0.3, 0.9, 400.0, 42.0, 0.44, 0.53, 0.51, 0.80, 0.71, 33.6, 61.78,
+        81.71, 83.48, 41.18
+    ])
     index_batch = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2])
 
     # pylint: disable = protected-access
@@ -69,8 +68,10 @@ def test_threshold_update_for_multiple_cells(data, learning_rate):
                                                      objective_batch,
                                                      index_batch, learning_rate)
 
-    result_true = calc_expected_threshold(objective, old_threshold[0],
-                                          learning_rate)
+    result_true = [
+        calc_expected_threshold(objective_batch[5 * i:5 * (i + 1)],
+                                threshold_arr[i], learning_rate)
+        for i in range(3)
+    ]
 
-    for result in result_test:
-        assert pytest.approx(result) == result_true
+    assert np.all(np.isclose(result_test, result_true))

From 0a27c7762b20c10207397c54116a95aff0f5bb15 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Tue, 11 Oct 2022 17:39:06 -0700
Subject: [PATCH 33/53] Formatting

---
 ribs/archives/_archive_base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index bb9992149..3fbbead51 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -676,8 +676,8 @@ def add_single(self, solution, objective, measures, metadata=None):
         # For new solutions, we set the old_threshold and old_objective to
         # 0 for computing value only if threshold min is not set.
         if not self._occupied_arr[index]:
-            old_objective = self.dtype(
-                0.0) if old_threshold == -np.inf else self._threshold_min
+            old_objective = (self.dtype(0.0) if old_threshold == -np.inf else
+                             self._threshold_min)
 
         was_occupied = self._occupied_arr[index]
         status = 0  # NOT_ADDED

From d2e075599f13ebc8af783108710ea44a107813e2 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 02:12:06 -0700
Subject: [PATCH 34/53] Simplify test fixture

---
 tests/archives/archive_base_test.py |  5 -----
 tests/archives/conftest.py          |  6 ++++++
 tests/archives/cvt_archive_test.py  |  3 ---
 tests/archives/grid_archive_test.py |  6 ------
 tests/schedulers/scheduler_test.py  | 10 ++++++----
 5 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/tests/archives/archive_base_test.py b/tests/archives/archive_base_test.py
index 281d60e7a..5f583f788 100644
--- a/tests/archives/archive_base_test.py
+++ b/tests/archives/archive_base_test.py
@@ -46,7 +46,6 @@ def test_iteration():
         assert elite.metadata == data.metadata
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_during_iteration(add_mode):
     # Even with just one entry, adding during iteration should still raise an
     # error, just like it does in set.
@@ -70,7 +69,6 @@ def test_clear_during_iteration():
             data.archive_with_elite.clear()
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_clear_and_add_during_iteration(add_mode):
     data = get_archive_data("GridArchive")
     with pytest.raises(RuntimeError):
@@ -102,7 +100,6 @@ def test_stats_dtype(dtype):
     assert isinstance(data.archive_with_elite.stats.obj_mean, dtype)
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_stats_multiple_add(add_mode):
     archive = GridArchive(solution_dim=3,
                           dims=[10, 20],
@@ -124,7 +121,6 @@ def test_stats_multiple_add(add_mode):
     assert np.isclose(archive.stats.obj_mean, 2.0)
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_stats_add_and_overwrite(add_mode):
     archive = GridArchive(solution_dim=3,
                           dims=[10, 20],
@@ -148,7 +144,6 @@ def test_stats_add_and_overwrite(add_mode):
     assert np.isclose(archive.stats.obj_mean, 3.0)
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_best_elite(add_mode):
     archive = GridArchive(solution_dim=3,
                           dims=[10, 20],
diff --git a/tests/archives/conftest.py b/tests/archives/conftest.py
index ce471f206..e25ef093c 100644
--- a/tests/archives/conftest.py
+++ b/tests/archives/conftest.py
@@ -31,6 +31,12 @@ def use_kd_tree(request):
     return request.param
 
 
+@pytest.fixture(params=["single", "batch"])
+def add_mode(request):
+    """Whether to use the KD Tree in CVTArchive."""
+    return request.param
+
+
 #
 # Helpers for generating archive data.
 #
diff --git a/tests/archives/cvt_archive_test.py b/tests/archives/cvt_archive_test.py
index 3ab1fd9a6..b45afac81 100644
--- a/tests/archives/cvt_archive_test.py
+++ b/tests/archives/cvt_archive_test.py
@@ -73,7 +73,6 @@ def test_custom_centroids_bad_shape(use_kd_tree):
 
 
 @pytest.mark.parametrize("use_list", [True, False], ids=["list", "ndarray"])
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_single_to_archive(data, use_list, add_mode):
     solution = data.solution
     objective = data.objective
@@ -97,7 +96,6 @@ def test_add_single_to_archive(data, use_list, add_mode):
                          data.measures, data.centroid, data.metadata)
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_single_and_overwrite(data, add_mode):
     """Test adding a new solution with a higher objective value."""
     arbitrary_sol = data.solution + 1
@@ -119,7 +117,6 @@ def test_add_single_and_overwrite(data, add_mode):
                          data.measures, data.centroid, arbitrary_metadata)
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_single_without_overwrite(data, add_mode):
     """Test adding a new solution with a lower objective value."""
     arbitrary_sol = data.solution + 1
diff --git a/tests/archives/grid_archive_test.py b/tests/archives/grid_archive_test.py
index 2b047d4c9..bc81dc8ab 100644
--- a/tests/archives/grid_archive_test.py
+++ b/tests/archives/grid_archive_test.py
@@ -104,7 +104,6 @@ def test_properties_are_correct(data):
 
 
 @pytest.mark.parametrize("use_list", [True, False], ids=["list", "ndarray"])
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_single_to_archive(data, use_list, add_mode):
     solution = data.solution
     objective = data.objective
@@ -129,7 +128,6 @@ def test_add_single_to_archive(data, use_list, add_mode):
 
 
 @pytest.mark.parametrize("use_list", [True, False], ids=["list", "ndarray"])
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_single_to_archive_negative(data, use_list, add_mode):
     solution = data.solution
     objective = -data.objective
@@ -153,7 +151,6 @@ def test_add_single_to_archive_negative(data, use_list, add_mode):
                          data.measures, data.grid_indices, data.metadata)
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_single_with_low_measures(data, add_mode):
     measures = np.array([-2, -3])
     indices = (0, 0)
@@ -169,7 +166,6 @@ def test_add_single_with_low_measures(data, add_mode):
                          indices, data.metadata)
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_single_with_high_measures(data, add_mode):
     measures = np.array([2, 3])
     indices = (9, 19)
@@ -184,7 +180,6 @@ def test_add_single_with_high_measures(data, add_mode):
                          indices, data.metadata)
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_single_and_overwrite(data, add_mode):
     """Test adding a new solution with a higher objective value."""
     arbitrary_sol = data.solution + 1
@@ -206,7 +201,6 @@ def test_add_single_and_overwrite(data, add_mode):
                          data.measures, data.grid_indices, arbitrary_metadata)
 
 
-@pytest.mark.parametrize("add_mode", ["single", "batch"])
 def test_add_single_without_overwrite(data, add_mode):
     """Test adding a new solution with a lower objective value."""
     arbitrary_sol = data.solution + 1
diff --git a/tests/schedulers/scheduler_test.py b/tests/schedulers/scheduler_test.py
index f7b5ce905..af359280e 100644
--- a/tests/schedulers/scheduler_test.py
+++ b/tests/schedulers/scheduler_test.py
@@ -23,6 +23,12 @@ def scheduler_fixture():
     return Scheduler(archive, emitters), solution_dim, num_solutions
 
 
+@pytest.fixture(params=["single", "batch"])
+def add_mode(request):
+    """Whether to use the KD Tree in CVTArchive."""
+    return request.param
+
+
 def test_init_fails_with_no_emitters():
     # arbitrary sol_dim
     archive = GridArchive(10, [100, 100], [(-1, 1), (-1, 1)])
@@ -69,8 +75,6 @@ def test_ask_fails_when_called_twice(scheduler_fixture):
         scheduler.ask()
 
 
-@pytest.mark.parametrize("add_mode", ["batch", "single"],
-                         ids=["batch_add", "single_add"])
 @pytest.mark.parametrize("tell_metadata", [True, False],
                          ids=["metadata", "no_metadata"])
 def test_tell_inserts_solutions_into_archive(add_mode, tell_metadata):
@@ -98,8 +102,6 @@ def test_tell_inserts_solutions_into_archive(add_mode, tell_metadata):
     )
 
 
-@pytest.mark.parametrize("add_mode", ["batch", "single"],
-                         ids=["batch_add", "single_add"])
 @pytest.mark.parametrize("tell_metadata", [True, False],
                          ids=["metadata", "no_metadata"])
 def test_tell_inserts_solutions_with_multiple_emitters(add_mode, tell_metadata):

From 160c10fe8a11a44a32761868fb8e997590914114 Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 12 Oct 2022 02:21:49 -0700
Subject: [PATCH 35/53] Fix: change old_objective to old_threshold

---
 ribs/archives/_archive_base.py                | 23 +++++++++-------
 .../archives/archive_threshold_update_test.py | 26 +++++++++++++++++++
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 3fbbead51..55e8c9212 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -527,11 +527,12 @@ def add(self,
         status_batch[is_new] = 2
         status_batch[improve_existing] = 1
 
-        # Since we set the new solutions in old_threshold_batch to have
-        # value 0.0, the values for new solutions are correct here.
-        old_objective_batch[is_new] = 0.0
-        old_threshold_batch[
-            is_new] = 0.0 if self._threshold_min == -np.inf else self._threshold_min
+        # If threshold_min is -inf, then we want CMA-ME behavior, which
+        # will compute the improvement value w.r.t zero. Otherwise, we will
+        # use compute w.r.t. threshold_min.
+        old_objective_batch[is_new] = self.dtype(0)
+        old_threshold_batch[is_new] = (self.dtype(0) if self._threshold_min == -np.inf
+                                       else self._threshold_min)
         value_batch = objective_batch - old_threshold_batch
 
         ## Step 3: Insert solutions into archive. ##
@@ -676,8 +677,12 @@ def add_single(self, solution, objective, measures, metadata=None):
         # For new solutions, we set the old_threshold and old_objective to
         # 0 for computing value only if threshold min is not set.
         if not self._occupied_arr[index]:
-            old_objective = (self.dtype(0.0) if old_threshold == -np.inf else
-                             self._threshold_min)
+            old_objective = self.dtype(0)
+            # If threshold_min is -inf, then we want CMA-ME behavior, which
+            # will compute the improvement value w.r.t zero. Otherwise, we will
+            # use compute w.r.t. threshold_min.
+            old_threshold = (self.dtype(0) if self._threshold_min == -np.inf
+                             else self._threshold_min)
 
         was_occupied = self._occupied_arr[index]
         status = 0  # NOT_ADDED
@@ -694,9 +699,9 @@ def add_single(self, solution, objective, measures, metadata=None):
                 self._num_occupied += 1
                 status = 2  # NEW
 
-            # Update the threshold.
+            # Update the threshold if new objective is greater than the old threshold.
             if self._threshold_arr[index] < objective:
-                self._threshold_arr[index] = old_objective * \
+                self._threshold_arr[index] = old_threshold * \
                     (1.0 - self._learning_rate) + \
                     objective * self._learning_rate
 
diff --git a/tests/archives/archive_threshold_update_test.py b/tests/archives/archive_threshold_update_test.py
index 2a3e45037..b3053d9d9 100644
--- a/tests/archives/archive_threshold_update_test.py
+++ b/tests/archives/archive_threshold_update_test.py
@@ -2,6 +2,8 @@
 import numpy as np
 import pytest
 
+from ribs.archives import GridArchive
+
 from .conftest import get_archive_data
 
 # pylint: disable = redefined-outer-name
@@ -75,3 +77,27 @@ def test_threshold_update_for_multiple_cells(data, learning_rate):
     ]
 
     assert np.all(np.isclose(result_test, result_true))
+
+
+def test_init_learning_rate_and_threshold_min():
+    # Setting threshold_min while not setting the learning_rate should not
+    # raise an error.
+    _ = GridArchive(solution_dim=2,
+                    dims=[10, 20],
+                    ranges=[(-1, 1), (-2, 2)],
+                    threshold_min=0)
+
+    # Setting both learning_rate and threshold_min should not raise an error.
+    _ = GridArchive(solution_dim=2,
+                    dims=[10, 20],
+                    ranges=[(-1, 1), (-2, 2)],
+                    learning_rate=0.1,
+                    threshold_min=0)
+
+    # Setting learning_rate while not setting the threshold_min should raise an
+    # error.
+    with pytest.raises(ValueError):
+        _ = GridArchive(solution_dim=2,
+                        dims=[10, 20],
+                        ranges=[(-1, 1), (-2, 2)],
+                        learning_rate=0.1)

From 8dcdb542e0b7217d96ca968ff20a68fb2359e10c Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 12 Oct 2022 02:23:43 -0700
Subject: [PATCH 36/53] docstring

---
 tests/archives/conftest.py         | 2 +-
 tests/schedulers/scheduler_test.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/archives/conftest.py b/tests/archives/conftest.py
index e25ef093c..78ce0d818 100644
--- a/tests/archives/conftest.py
+++ b/tests/archives/conftest.py
@@ -33,7 +33,7 @@ def use_kd_tree(request):
 
 @pytest.fixture(params=["single", "batch"])
 def add_mode(request):
-    """Whether to use the KD Tree in CVTArchive."""
+    """Single or batch add."""
     return request.param
 
 
diff --git a/tests/schedulers/scheduler_test.py b/tests/schedulers/scheduler_test.py
index af359280e..278c1d77e 100644
--- a/tests/schedulers/scheduler_test.py
+++ b/tests/schedulers/scheduler_test.py
@@ -25,7 +25,7 @@ def scheduler_fixture():
 
 @pytest.fixture(params=["single", "batch"])
 def add_mode(request):
-    """Whether to use the KD Tree in CVTArchive."""
+    """Single or batch add."""
     return request.param
 
 

From 6f288315fc2d41b9e622e5c1f9ff790eb1346f24 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 02:31:47 -0700
Subject: [PATCH 37/53] Rename test

---
 tests/archives/grid_archive_test.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/archives/grid_archive_test.py b/tests/archives/grid_archive_test.py
index bc81dc8ab..bce9603c3 100644
--- a/tests/archives/grid_archive_test.py
+++ b/tests/archives/grid_archive_test.py
@@ -128,7 +128,9 @@ def test_add_single_to_archive(data, use_list, add_mode):
 
 
 @pytest.mark.parametrize("use_list", [True, False], ids=["list", "ndarray"])
-def test_add_single_to_archive_negative(data, use_list, add_mode):
+def test_add_single_to_archive_negative_objective(data, use_list, add_mode):
+    """Same as test_add_single_to_archive, but negative objective since there
+    are some weird cases when handling value calculations."""
     solution = data.solution
     objective = -data.objective
     measures = data.measures

From 3b0506d675c31d8baf8659ec63dd257978fe08be Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 16:57:41 -0700
Subject: [PATCH 38/53] Fix add_single

---
 ribs/archives/_archive_base.py | 54 ++++++++++++++++++----------------
 1 file changed, 28 insertions(+), 26 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 55e8c9212..e155f3320 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -162,9 +162,8 @@ def __init__(self,
 
         # threshold min can only be -np.inf if the learning rate is 1.0
         if learning_rate != 1.0 and threshold_min == -np.inf:
-            raise ValueError(
-                "If learning_rate != 1.0, threshold min cannot be -np.inf (default)."
-            )
+            raise ValueError("If learning_rate != 1.0, threshold min cannot "
+                             "be -np.inf (default).")
         self._learning_rate = self._dtype(learning_rate)
         self._threshold_min = self._dtype(threshold_min)
         self._threshold_arr = np.full(self._cells,
@@ -531,8 +530,8 @@ def add(self,
         # will compute the improvement value w.r.t zero. Otherwise, we will
         # use compute w.r.t. threshold_min.
         old_objective_batch[is_new] = self.dtype(0)
-        old_threshold_batch[is_new] = (self.dtype(0) if self._threshold_min == -np.inf
-                                       else self._threshold_min)
+        old_threshold_batch[is_new] = (self.dtype(0) if self._threshold_min
+                                       == -np.inf else self._threshold_min)
         value_batch = objective_batch - old_threshold_batch
 
         ## Step 3: Insert solutions into archive. ##
@@ -664,46 +663,49 @@ def add_single(self, solution, objective, measures, metadata=None):
         solution = np.asarray(solution)
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
 
+        # TODO: Check for inf and nan in objective and measures.
         objective = self.dtype(objective)
 
         measures = np.asarray(measures)
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
         index = self.index_of_single(measures)
 
-        # Note that when learning_rate = 1.0, old_threshold == old_objective.
+        # Only used for computing QD score.
         old_objective = self._objective_arr[index]
+
+        # Used for computing improvement value.
         old_threshold = self._threshold_arr[index]
 
-        # For new solutions, we set the old_threshold and old_objective to
-        # 0 for computing value only if threshold min is not set.
-        if not self._occupied_arr[index]:
+        # New solutions require special settings for old_objective and
+        # old_threshold.
+        was_occupied = self._occupied_arr[index]
+        if not was_occupied:
             old_objective = self.dtype(0)
-            # If threshold_min is -inf, then we want CMA-ME behavior, which
-            # will compute the improvement value w.r.t zero. Otherwise, we will
-            # use compute w.r.t. threshold_min.
+            # If threshold_min is -inf, then we want CMA-ME behavior, which will
+            # compute the improvement value w.r.t. zero for new solutions.
+            # Otherwise, we will compute w.r.t. threshold_min.
             old_threshold = (self.dtype(0) if self._threshold_min == -np.inf
                              else self._threshold_min)
 
-        was_occupied = self._occupied_arr[index]
         status = 0  # NOT_ADDED
-        if not was_occupied or self._threshold_arr[index] < objective:
+        # In the case where we want CMA-ME behavior, the old threshold is -inf
+        # for new cells, which satisfies this if condition.
+        if self._threshold_arr[index] < objective:
             if was_occupied:
                 status = 1  # IMPROVE_EXISTING
             else:
-                # Set this index to "occupied" -- important that we do this before
-                # inserting the solution.
+                # Set this index to be occupied.
                 self._occupied_arr[index] = True
-
-                # Tracks a new occupied index.
                 self._occupied_indices[self._num_occupied] = index
                 self._num_occupied += 1
+
                 status = 2  # NEW
 
-            # Update the threshold if new objective is greater than the old threshold.
-            if self._threshold_arr[index] < objective:
-                self._threshold_arr[index] = old_threshold * \
-                    (1.0 - self._learning_rate) + \
-                    objective * self._learning_rate
+            # This calculation works in the case where threshold_min is -inf
+            # because old_threshold will be set to 0.0 instead.
+            self._threshold_arr[index] = (old_threshold *
+                                          (1.0 - self._learning_rate) +
+                                          objective * self._learning_rate)
 
             # Insert into the archive.
             self._objective_arr[index] = objective
@@ -718,9 +720,9 @@ def add_single(self, solution, objective, measures, metadata=None):
             if self._stats.obj_max is None or objective > self._stats.obj_max:
                 new_obj_max = objective
                 self._best_elite = Elite(
-                    readonly(solution),
+                    readonly(np.copy(self._solution_arr[index])),
                     objective,
-                    readonly(measures),
+                    readonly(np.copy(self._measures_arr[index])),
                     index,
                     metadata,
                 )
@@ -735,7 +737,7 @@ def add_single(self, solution, objective, measures, metadata=None):
                 obj_mean=new_qd_score / self.dtype(len(self)),
             )
 
-        return status, self.dtype(objective - old_objective)
+        return status, objective - old_threshold
 
     def elites_with_measures(self, measures_batch):
         """Retrieves the elites with measures in the same cells as the measures

From 2fe3287230f65aca287c4a25e4b13bb73cf0c80d Mon Sep 17 00:00:00 2001
From: David Lee <dhlee@usc.edu>
Date: Wed, 12 Oct 2022 16:59:19 -0700
Subject: [PATCH 39/53] add test; typos

---
 ribs/archives/_archive_base.py                  |  9 ++++-----
 tests/archives/archive_threshold_update_test.py | 16 ++++++++++++++++
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 55e8c9212..1cdca8fd3 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -160,10 +160,10 @@ def __init__(self,
                                       dtype=self.dtype)
         self._metadata_arr = np.empty(self._cells, dtype=object)
 
-        # threshold min can only be -np.inf if the learning rate is 1.0
+        # threshold min can only be -np.inf if the learning rate is 1.
         if learning_rate != 1.0 and threshold_min == -np.inf:
             raise ValueError(
-                "If learning_rate != 1.0, threshold min cannot be -np.inf (default)."
+                "If learning_rate != 1.0, threshold min cannot be -inf (default)."
             )
         self._learning_rate = self._dtype(learning_rate)
         self._threshold_min = self._dtype(threshold_min)
@@ -311,8 +311,8 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
         objective_sizes = objective_sizes[threshold_update_indices]
         objective_sums = objective_sums[threshold_update_indices]
 
-        # Sum of geometric series (1 - learning_rate)^i from i = 0 to i = n - 1
-        # See https://en.wikipedia.org/wiki/Geometric_series#Sum
+        # Sum of geometric series (1 - learning_rate)^i from i = 0 to i = n - 1.
+        # See https://en.wikipedia.org/wiki/Geometric_series#Sum.
         ratio = self.dtype(1.0 - learning_rate)
         if ratio == 1.0:
             geometric_sums = objective_sizes
@@ -670,7 +670,6 @@ def add_single(self, solution, objective, measures, metadata=None):
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
         index = self.index_of_single(measures)
 
-        # Note that when learning_rate = 1.0, old_threshold == old_objective.
         old_objective = self._objective_arr[index]
         old_threshold = self._threshold_arr[index]
 
diff --git a/tests/archives/archive_threshold_update_test.py b/tests/archives/archive_threshold_update_test.py
index b3053d9d9..1a243eb92 100644
--- a/tests/archives/archive_threshold_update_test.py
+++ b/tests/archives/archive_threshold_update_test.py
@@ -79,6 +79,22 @@ def test_threshold_update_for_multiple_cells(data, learning_rate):
     assert np.all(np.isclose(result_test, result_true))
 
 
+def test_threshold_update_for_empty_objective_and_index(data):
+    archive = data.archive
+
+    threshold_arr = np.array([-3.1, 0.4, 2.9])
+    objective_batch = np.array([])  # Empty objective.
+    index_batch = np.array([])  # Empty index.
+
+    # pylint: disable = protected-access
+    new_threshold_batch, threshold_update_indices = (
+        archive._compute_new_thresholds(threshold_arr, objective_batch,
+                                        index_batch, 0.1))
+
+    assert new_threshold_batch.size == 0
+    assert threshold_update_indices.size == 0
+
+
 def test_init_learning_rate_and_threshold_min():
     # Setting threshold_min while not setting the learning_rate should not
     # raise an error.

From 82dbb491d27c06603107027011ced7a69c856fe8 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 17:44:51 -0700
Subject: [PATCH 40/53] Add inf and nan checks

---
 ribs/_utils.py                 | 14 ++++++++++++++
 ribs/archives/_archive_base.py | 12 +++++++++---
 ribs/archives/_cvt_archive.py  |  3 ++-
 ribs/archives/_grid_archive.py |  3 ++-
 4 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/ribs/_utils.py b/ribs/_utils.py
index eb9e4f97d..c6ce453ad 100644
--- a/ribs/_utils.py
+++ b/ribs/_utils.py
@@ -1,4 +1,18 @@
 """Miscellaneous internal utilities."""
+import numpy as np
+
+
+def check_finite(x, name):
+    """Checks that x is finite (i.e. not infinity or NaN).
+
+    `x` must be either a scalar or NumPy array.
+    """
+    if not np.all(np.isfinite(x)):
+        if np.isscalar(x):
+            raise ValueError(f"{name} must be finite (infinity "
+                             "and NaN values are not supported).")
+        raise ValueError(f"All elements of {name} must be finite (infinity "
+                         "and NaN values are not supported).")
 
 
 def check_batch_shape(array, array_name, dim, dim_name, extra_msg=""):
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index e155f3320..1844140ef 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -5,8 +5,8 @@
 import numpy as np
 from numpy_groupies import aggregate_nb as aggregate
 
-from ribs._utils import (check_1d_shape, check_batch_shape, check_is_1d,
-                         check_solution_batch_dim)
+from ribs._utils import (check_1d_shape, check_batch_shape, check_finite,
+                         check_is_1d, check_solution_batch_dim)
 from ribs.archives._archive_data_frame import ArchiveDataFrame
 from ribs.archives._archive_stats import ArchiveStats
 from ribs.archives._elite import Elite, EliteBatch
@@ -381,6 +381,7 @@ def index_of_single(self, measures):
         """
         measures = np.asarray(measures)
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
+        check_finite(measures, "measures")
         return self.index_of(measures[None])[0]
 
     _ADD_WARNING = (" Note that starting in pyribs 0.5.0, add() takes in a "
@@ -473,6 +474,8 @@ def add(self,
         """
         self._state["add"] += 1
 
+        # TODO: Check inf and nan.
+
         ## Step 1: Validate input. ##
 
         solution_batch = np.asarray(solution_batch)
@@ -663,11 +666,12 @@ def add_single(self, solution, objective, measures, metadata=None):
         solution = np.asarray(solution)
         check_1d_shape(solution, "solution", self.solution_dim, "solution_dim")
 
-        # TODO: Check for inf and nan in objective and measures.
         objective = self.dtype(objective)
+        check_finite(objective, "objective")
 
         measures = np.asarray(measures)
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
+        check_finite(measures, "measures")
         index = self.index_of_single(measures)
 
         # Only used for computing QD score.
@@ -793,6 +797,7 @@ def elites_with_measures(self, measures_batch):
         measures_batch = np.asarray(measures_batch)
         check_batch_shape(measures_batch, "measures_batch", self.measure_dim,
                           "measure_dim")
+        check_finite(measures_batch, "measures_batch")
 
         index_batch = self.index_of(measures_batch)
         occupied_batch = self._occupied_arr[index_batch]
@@ -860,6 +865,7 @@ def elites_with_measures_single(self, measures):
         """
         measures = np.asarray(measures)
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
+        check_finite(measures, "measures")
 
         elite_batch = self.elites_with_measures(measures[None])
         return Elite(
diff --git a/ribs/archives/_cvt_archive.py b/ribs/archives/_cvt_archive.py
index fde2f0391..353ba7925 100644
--- a/ribs/archives/_cvt_archive.py
+++ b/ribs/archives/_cvt_archive.py
@@ -5,7 +5,7 @@
 from scipy.spatial import cKDTree  # pylint: disable=no-name-in-module
 from sklearn.cluster import k_means
 
-from ribs._utils import check_batch_shape
+from ribs._utils import check_batch_shape, check_finite
 from ribs.archives._archive_base import ArchiveBase
 
 
@@ -244,6 +244,7 @@ def index_of(self, measures_batch):
         measures_batch = np.asarray(measures_batch)
         check_batch_shape(measures_batch, "measures_batch", self.measure_dim,
                           "measure_dim")
+        check_finite(measures_batch, "measures_batch")
 
         if self._use_kd_tree:
             return np.asarray(
diff --git a/ribs/archives/_grid_archive.py b/ribs/archives/_grid_archive.py
index bcf995310..6c77e342e 100644
--- a/ribs/archives/_grid_archive.py
+++ b/ribs/archives/_grid_archive.py
@@ -1,7 +1,7 @@
 """Contains the GridArchive."""
 import numpy as np
 
-from ribs._utils import check_batch_shape, check_is_1d
+from ribs._utils import check_batch_shape, check_finite, check_is_1d
 from ribs.archives._archive_base import ArchiveBase
 
 
@@ -168,6 +168,7 @@ def index_of(self, measures_batch):
         measures_batch = np.asarray(measures_batch)
         check_batch_shape(measures_batch, "measures_batch", self.measure_dim,
                           "measure_dim")
+        check_finite(measures_batch, "measures_batch")
 
         # Adding epsilon accounts for floating point precision errors from
         # transforming measures. We then cast to int32 to obtain integer

From d3d9b9e0aa3a74de1d8270e1ff436bc2c265d288 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 18:52:03 -0700
Subject: [PATCH 41/53] Add test for nonfinite inputs

---
 tests/archives/conftest.py          |  2 +-
 tests/archives/grid_archive_test.py | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/tests/archives/conftest.py b/tests/archives/conftest.py
index 78ce0d818..71ff553b8 100644
--- a/tests/archives/conftest.py
+++ b/tests/archives/conftest.py
@@ -77,7 +77,7 @@ def get_archive_data(name, dtype=np.float64):
     ARCHIVE_NAMES.
     """
     # Characteristics of a single solution to insert into archive_with_elite.
-    solution = np.array([1, 2, 3])
+    solution = np.array([1., 2., 3.])
     objective = 1.0
     measures = np.array([0.25, 0.25])
     metadata = {"metadata_key": 42}
diff --git a/tests/archives/grid_archive_test.py b/tests/archives/grid_archive_test.py
index bce9603c3..a26fdc3de 100644
--- a/tests/archives/grid_archive_test.py
+++ b/tests/archives/grid_archive_test.py
@@ -501,3 +501,21 @@ def test_values_go_to_correct_bin(dtype):
     # Upper bound and above belong in last bin.
     assert archive.index_of_single([0.1]) == 9
     assert archive.index_of_single([0.11]) == 9
+
+
+def test_nonfinite_inputs(data):
+    data.solution[0] = np.inf
+    data.measures[0] = np.nan
+
+    with pytest.raises(ValueError):
+        data.archive.add([data.solution], -np.inf, [data.measures])
+    with pytest.raises(ValueError):
+        data.archive.add_single(data.solution, -np.inf, data.measures)
+    with pytest.raises(ValueError):
+        data.archive.elites_with_measures([data.measures])
+    with pytest.raises(ValueError):
+        data.archive.elites_with_measures_single(data.measures)
+    with pytest.raises(ValueError):
+        data.archive.index_of([data.measures])
+    with pytest.raises(ValueError):
+        data.archive.index_of_single(data.measures)

From f8a1f897c8f0c0407853994d039a54fbfc5b3a50 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 19:00:28 -0700
Subject: [PATCH 42/53] Finite checks in add

---
 ribs/archives/_archive_base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 233a8160e..bb4313c65 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -474,8 +474,6 @@ def add(self,
         """
         self._state["add"] += 1
 
-        # TODO: Check inf and nan.
-
         ## Step 1: Validate input. ##
 
         solution_batch = np.asarray(solution_batch)
@@ -490,6 +488,7 @@ def add(self,
                                  batch_size,
                                  is_1d=True,
                                  extra_msg=self._ADD_WARNING)
+        check_finite(objective_batch, "objective_batch")
 
         measures_batch = np.asarray(measures_batch)
         check_batch_shape(measures_batch, "measures_batch", self.measure_dim,
@@ -499,6 +498,7 @@ def add(self,
                                  batch_size,
                                  is_1d=False,
                                  extra_msg=self._ADD_WARNING)
+        check_finite(measures_batch, "measures_batch")
 
         metadata_batch = (np.empty(batch_size, dtype=object) if
                           metadata_batch is None else np.asarray(metadata_batch,

From 4feac42c3fcf8d383cf854948b1af1fe1f889db2 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 19:07:59 -0700
Subject: [PATCH 43/53] Error messages

---
 ribs/archives/_archive_base.py | 10 ++++++++++
 ribs/archives/_cvt_archive.py  |  1 +
 ribs/archives/_grid_archive.py |  1 +
 3 files changed, 12 insertions(+)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index bb4313c65..3bbccb5c5 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -378,6 +378,7 @@ def index_of_single(self, measures):
             storage arrays.
         Raises:
             ValueError: ``measures`` is not of shape (:attr:`measure_dim`,).
+            ValueError: ``measures`` has non-finite values (inf or NaN).
         """
         measures = np.asarray(measures)
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
@@ -388,6 +389,7 @@ def index_of_single(self, measures):
                     "batch of solutions unlike in pyribs 0.4.0, where add() "
                     "only took in a single solution.")
 
+    # TODO: Update value_batch documentation here and in add_single.
     def add(self,
             solution_batch,
             objective_batch,
@@ -471,6 +473,8 @@ def add(self,
                 solution.
         Raises:
             ValueError: The array arguments do not match their specified shapes.
+            ValueError: ``objective_batch`` or ``measures_batch`` has non-finite
+                values (inf or NaN).
         """
         self._state["add"] += 1
 
@@ -656,6 +660,10 @@ def add_single(self, solution, objective, measures, metadata=None):
                     array-like objects as metadata may lead to unexpected
                     behavior. However, the metadata may be a dict or other
                     object which *contains* arrays.
+        Raises:
+            ValueError: The array arguments do not match their specified shapes.
+            ValueError: ``objective`` is non-finite (inf or NaN) or ``measures``
+                has non-finite values.
         Returns:
             tuple: 2-element tuple of (status, value) describing the result of
             the add operation. Refer to :meth:`add` for the meaning of the
@@ -793,6 +801,7 @@ def elites_with_measures(self, measures_batch):
         Raises:
             ValueError: ``measures_batch`` is not of shape (batch_size,
                 :attr:`measure_dim`).
+            ValueError: ``measures_batch`` has non-finite values (inf or NaN).
         """
         measures_batch = np.asarray(measures_batch)
         check_batch_shape(measures_batch, "measures_batch", self.measure_dim,
@@ -862,6 +871,7 @@ def elites_with_measures_single(self, measures):
             described in :meth:`elites_with_measures`.
         Raises:
             ValueError: ``measures`` is not of shape (:attr:`measure_dim`,).
+            ValueError: ``measures`` has non-finite values (inf or NaN).
         """
         measures = np.asarray(measures)
         check_1d_shape(measures, "measures", self.measure_dim, "measure_dim")
diff --git a/ribs/archives/_cvt_archive.py b/ribs/archives/_cvt_archive.py
index 353ba7925..4339ad72d 100644
--- a/ribs/archives/_cvt_archive.py
+++ b/ribs/archives/_cvt_archive.py
@@ -240,6 +240,7 @@ def index_of(self, measures_batch):
         Raises:
             ValueError: ``measures_batch`` is not of shape (batch_size,
                 :attr:`measure_dim`).
+            ValueError: ``measures_batch`` has non-finite values (inf or NaN).
         """
         measures_batch = np.asarray(measures_batch)
         check_batch_shape(measures_batch, "measures_batch", self.measure_dim,
diff --git a/ribs/archives/_grid_archive.py b/ribs/archives/_grid_archive.py
index 6c77e342e..77554a1a7 100644
--- a/ribs/archives/_grid_archive.py
+++ b/ribs/archives/_grid_archive.py
@@ -164,6 +164,7 @@ def index_of(self, measures_batch):
         Raises:
             ValueError: ``measures_batch`` is not of shape (batch_size,
                 :attr:`measure_dim`).
+            ValueError: ``measures_batch`` has non-finite values (inf or NaN).
         """
         measures_batch = np.asarray(measures_batch)
         check_batch_shape(measures_batch, "measures_batch", self.measure_dim,

From d53dceb2ae80efac753103ebe9d9b3c8fbe93c80 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 19:25:41 -0700
Subject: [PATCH 44/53] Test add single with threshold updates

---
 tests/archives/grid_archive_test.py | 44 +++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/tests/archives/grid_archive_test.py b/tests/archives/grid_archive_test.py
index a26fdc3de..d9194cf78 100644
--- a/tests/archives/grid_archive_test.py
+++ b/tests/archives/grid_archive_test.py
@@ -224,6 +224,50 @@ def test_add_single_without_overwrite(data, add_mode):
                          data.measures, data.grid_indices, data.metadata)
 
 
+def test_add_single_with_threshold_update(add_mode):
+    archive = GridArchive(
+        solution_dim=3,
+        dims=[10, 10],
+        ranges=[(-1, 1), (-1, 1)],
+        threshold_min=-1.0,
+        learning_rate=0.1,
+    )
+    solution = [1, 2, 3]
+    measures = [0.1, 0.1]
+
+    # Add a new solution to the archive.
+    if add_mode == "single":
+        status, value = archive.add_single(solution, 0.0, measures)
+    else:
+        status_batch, value_batch = archive.add([solution], [0.0], [measures])
+        status, value = status_batch[0], value_batch[0]
+
+    assert status == 2  # NEW
+    assert np.isclose(value, 1.0)  # 0.0 - (-1.0)
+
+    # Threshold should now be (1 - 0.1) * -1.0 + 0.1 * 0.0 = -0.9
+
+    # These solutions are below the threshold and should not be inserted.
+    if add_mode == "single":
+        status, value = archive.add_single(solution, -0.95, measures)
+    else:
+        status_batch, value_batch = archive.add([solution], [-0.95], [measures])
+        status, value = status_batch[0], value_batch[0]
+
+    assert status == 0  # NOT_ADDED
+    assert np.isclose(value, -0.05)  # -0.95 - (-0.9)
+
+    # These solutions are above the threshold and should be inserted.
+    if add_mode == "single":
+        status, value = archive.add_single(solution, -0.8, measures)
+    else:
+        status_batch, value_batch = archive.add([solution], [-0.8], [measures])
+        status, value = status_batch[0], value_batch[0]
+
+    assert status == 1  # IMPROVE_EXISTING
+    assert np.isclose(value, 0.1)  # -0.8 - (-0.9)
+
+
 def test_add_batch_all_new(data):
     status_batch, value_batch = data.archive.add(
         # 4 solutions of arbitrary value.

From 28c163eba9b6df377f92f59177e2fbdf9c836d17 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 19:59:09 -0700
Subject: [PATCH 45/53] Note on add_single

---
 ribs/archives/_archive_base.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 3bbccb5c5..caae53335 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -647,6 +647,12 @@ def add_single(self, solution, objective, measures, metadata=None):
         The solution is only inserted if it has a higher ``objective``
         than the elite previously in the corresponding cell.
 
+        .. note::
+            To make it more amenable to modifications, this method's
+            implementation is designed to be readable at the cost of
+            performance, e.g., none of its operations are modified. If you need
+            performance, we recommend using :meth:`add`.
+
         Args:
             solution (array-like): Parameters of the solution.
             objective (float): Objective function evaluation of the solution.

From 918602864fc34e1180b089c172f83b0e7532ef26 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 20:05:47 -0700
Subject: [PATCH 46/53] Revert sphere outdir name

---
 examples/sphere.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index 27c03c0fe..a8a4fd333 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -413,7 +413,7 @@ def sphere_main(algorithm,
             learning_rate = 1.0
 
     name = f"{algorithm}_{dim}"
-    outdir = Path(outdir + "_" + algorithm)
+    outdir = Path(outdir)
     if not outdir.is_dir():
         outdir.mkdir()
 

From d2eeac637290b4ea08e659c48b71ca5642ed44d2 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 20:25:23 -0700
Subject: [PATCH 47/53] TODO

---
 ribs/archives/_archive_base.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index caae53335..6e20eb872 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -54,6 +54,7 @@ def __next__(self):
         )
 
 
+# TODO: Document threshold behavior.
 class ArchiveBase(ABC):  # pylint: disable = too-many-instance-attributes
     """Base class for archives.
 

From 5d66944e808501005eb9014ea948046f1813062f Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 20:34:15 -0700
Subject: [PATCH 48/53] Modify err msg

---
 ribs/archives/_archive_base.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 6e20eb872..6a84e1030 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -161,10 +161,9 @@ def __init__(self,
                                       dtype=self.dtype)
         self._metadata_arr = np.empty(self._cells, dtype=object)
 
-        # threshold min can only be -np.inf if the learning rate is 1.
-        if learning_rate != 1.0 and threshold_min == -np.inf:
-            raise ValueError("If learning_rate != 1.0, threshold min cannot "
-                             "be -np.inf (default).")
+        if threshold_min == -np.inf and learning_rate != 1.0:
+            raise ValueError("threshold_min can only be -np.inf if "
+                             "learning_rate is 1.0")
         self._learning_rate = self._dtype(learning_rate)
         self._threshold_min = self._dtype(threshold_min)
         self._threshold_arr = np.full(self._cells,

From 025807b1859aa1d0a6b05a4b3177409bec0887e2 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Wed, 12 Oct 2022 20:36:15 -0700
Subject: [PATCH 49/53] Grammar

---
 examples/sphere.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/sphere.py b/examples/sphere.py
index a8a4fd333..3772d3e35 100644
--- a/examples/sphere.py
+++ b/examples/sphere.py
@@ -323,7 +323,7 @@ def create_scheduler(algorithm,
 
     print(
         f"Created Scheduler for {algorithm} with learning rate {learning_rate} "
-        f"and add mode {mode}, using solution dims {solution_dim} and archive "
+        f"and add mode {mode}, using solution dim {solution_dim} and archive "
         f"dims {archive_dims}.")
     return Scheduler(archive, emitters, result_archive, add_mode=mode)
 

From e308c397f8be63cad0e9091f013059c5ad473d0e Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Sun, 16 Oct 2022 03:57:50 -0700
Subject: [PATCH 50/53] Update add_batch implementation

---
 ribs/archives/_archive_base.py | 50 ++++++++++++++++++++++++----------
 1 file changed, 36 insertions(+), 14 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 6a84e1030..dbd707f6c 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -285,6 +285,10 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
             array indicating which entries in `threshold_arr` should
             be updated.
         """
+        # Even though we do this check, it should not be possible to have
+        # empty objective_batch or index_batch in the add() method since
+        # we check that at least one cell is being updated by seeing if
+        # can_insert has any True values.
         if objective_batch.size == 0 or index_batch.size == 0:
             return np.array([], dtype=self.dtype), np.array([], dtype=bool)
 
@@ -321,9 +325,12 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
         update = (learning_rate * (objective_sums / objective_sizes) *
                   geometric_sums)
 
+        # Unlike in add_single, we do not need to worry about
+        # old_threshold having -np.inf here as a result of threshold_min
+        # being -np.inf. This is because the case with threshold_min =
+        # -np.inf is handled separately since we compute the new
+        # threshold based on the max objective in each cell in that case.
         old_threshold = np.copy(threshold_arr[threshold_update_indices])
-        # TODO: Fix this based on new CMA-ME behavior if needed?
-        old_threshold[old_threshold == -np.inf] = 0
         prev = old_threshold * ratio**objective_sizes
 
         new_threshold_batch = prev + update
@@ -533,10 +540,13 @@ def add(self,
         status_batch[is_new] = 2
         status_batch[improve_existing] = 1
 
+        # New solutions require special settings for old_objective and
+        # old_threshold.
+        old_objective_batch[is_new] = self.dtype(0)
+
         # If threshold_min is -inf, then we want CMA-ME behavior, which
         # will compute the improvement value w.r.t zero. Otherwise, we will
-        # use compute w.r.t. threshold_min.
-        old_objective_batch[is_new] = self.dtype(0)
+        # compute w.r.t. threshold_min.
         old_threshold_batch[is_new] = (self.dtype(0) if self._threshold_min
                                        == -np.inf else self._threshold_min)
         value_batch = objective_batch - old_threshold_batch
@@ -558,12 +568,6 @@ def add(self,
         metadata_batch_can = metadata_batch[can_insert]
         old_objective_batch_can = old_objective_batch[can_insert]
 
-        # Update the thresholds.
-        new_thresholds, update_thresholds_indices = self._compute_new_thresholds(
-            self._threshold_arr, objective_batch_can, index_batch_can,
-            self._learning_rate)
-        self._threshold_arr[update_thresholds_indices] = new_thresholds
-
         # Retrieve indices of solutions that should be inserted into the
         # archive. Currently, multiple solutions may be inserted at each
         # archive index, but we only want to insert the maximum among these
@@ -608,6 +612,24 @@ def add(self,
                                    index_batch_insert[is_new_and_inserted])
         self._num_occupied += n_new
 
+        # Update the thresholds.
+        if self._threshold_min == -np.inf:
+            # Here we want regular archive behavior, so the thresholds
+            # should just be the maximum objective.
+            self._threshold_arr[index_batch_insert] = objective_batch_insert
+        else:
+            # Here we compute the batch threshold update described in the
+            # appendix of Fontaine 2022 https://arxiv.org/abs/2205.10752
+            # This computation is based on the mean objective of all
+            # solutions in the batch that could have been inserted into
+            # each cell. This method is separated out to facilitate
+            # testing.
+            (new_thresholds,
+             update_thresholds_indices) = self._compute_new_thresholds(
+                 self._threshold_arr, objective_batch_can, index_batch_can,
+                 self._learning_rate)
+            self._threshold_arr[update_thresholds_indices] = new_thresholds
+
         ## Step 4: Update archive stats. ##
 
         # Since we set the new solutions in the old objective batch to have
@@ -622,9 +644,9 @@ def add(self,
         if self._stats.obj_max is None or max_obj_insert > self._stats.obj_max:
             new_obj_max = max_obj_insert
             self._best_elite = Elite(
-                readonly(solution_batch_insert[max_idx]),
+                readonly(np.copy(solution_batch_insert[max_idx])),
                 objective_batch_insert[max_idx],
-                readonly(measures_batch_insert[max_idx]),
+                readonly(np.copy(measures_batch_insert[max_idx])),
                 index_batch_insert[max_idx],
                 metadata_batch_insert[max_idx],
             )
@@ -706,8 +728,8 @@ def add_single(self, solution, objective, measures, metadata=None):
                              else self._threshold_min)
 
         status = 0  # NOT_ADDED
-        # In the case where we want CMA-ME behavior, the old threshold is -inf
-        # for new cells, which satisfies this if condition.
+        # In the case where we want CMA-ME behavior, threshold_arr[index]
+        # is -inf for new cells, which satisfies this if condition.
         if self._threshold_arr[index] < objective:
             if was_occupied:
                 status = 1  # IMPROVE_EXISTING

From 50b7e7b3287fba6966bb955a0d3c8aeb49f2d2dc Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Sun, 16 Oct 2022 21:18:19 -0700
Subject: [PATCH 51/53] Documentation and tidying up

---
 docs/tutorials.md                   |  2 +
 examples/tutorials/cma_mae.ipynb    | 35 ++++++++++
 ribs/archives/_archive_base.py      | 99 ++++++++++++++++++++---------
 ribs/archives/_cvt_archive.py       | 10 ++-
 ribs/archives/_grid_archive.py      | 10 ++-
 tests/archives/archive_base_test.py |  8 +++
 6 files changed, 127 insertions(+), 37 deletions(-)
 create mode 100644 examples/tutorials/cma_mae.ipynb

diff --git a/docs/tutorials.md b/docs/tutorials.md
index 21f6a4c5f..ae2e4e39a 100644
--- a/docs/tutorials.md
+++ b/docs/tutorials.md
@@ -7,6 +7,7 @@ tutorials/lunar_lander
 tutorials/lsi_mnist
 tutorials/arm_repertoire
 tutorials/fooling_mnist
+tutorials/cma_mae
 ```
 
 Tutorials are Python notebooks with detailed explanations of pyribs usage. They
@@ -23,3 +24,4 @@ needs for execution.
 | {doc}`tutorials/lsi_mnist`      | {class}`~ribs.archives.GridArchive` | {class}`~ribs.emitters.EvolutionStrategyEmitter` | {class}`~ribs.schedulers.Scheduler` |
 | {doc}`tutorials/arm_repertoire` | {class}`~ribs.archives.CVTArchive`  | {class}`~ribs.emitters.EvolutionStrategyEmitter` | {class}`~ribs.schedulers.Scheduler` |
 | {doc}`tutorials/fooling_mnist`  | {class}`~ribs.archives.GridArchive` | {class}`~ribs.emitters.GaussianEmitter`          | {class}`~ribs.schedulers.Scheduler` |
+| {doc}`tutorials/cma_mae`        | {class}`~ribs.archives.GridArchive` | {class}`~ribs.emitters.EvolutionStrategyEmitter` | {class}`~ribs.schedulers.Scheduler` |
diff --git a/examples/tutorials/cma_mae.ipynb b/examples/tutorials/cma_mae.ipynb
new file mode 100644
index 000000000..13a23a6f1
--- /dev/null
+++ b/examples/tutorials/cma_mae.ipynb
@@ -0,0 +1,35 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "69e0c236-9407-496f-be0d-9562f8191ce4",
+   "metadata": {},
+   "source": [
+    "# CMA-MAE and Archive Thresholds\n",
+    "\n",
+    "Coming soon!"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index dbd707f6c..fa3e917ab 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -11,6 +11,10 @@
 from ribs.archives._archive_stats import ArchiveStats
 from ribs.archives._elite import Elite, EliteBatch
 
+_ADD_WARNING = (" Note that starting in pyribs 0.5.0, add() takes in a "
+                "batch of solutions unlike in pyribs 0.4.0, where add() "
+                "only took in a single solution.")
+
 
 def readonly(arr):
     """Sets an array to be readonly."""
@@ -54,18 +58,19 @@ def __next__(self):
         )
 
 
-# TODO: Document threshold behavior.
 class ArchiveBase(ABC):  # pylint: disable = too-many-instance-attributes
     """Base class for archives.
 
     This class assumes all archives use a fixed-size container with cells that
     hold (1) information about whether the cell is occupied (bool), (2) a
     solution (1D array), (3) objective function evaluation of the solution
-    (float), (4) measure space coordinates of the solution (1D array), and (5)
-    any additional metadata associated with the solution (object). In this
-    class, the container is implemented with separate numpy arrays that share
-    common dimensions. Using the ``solution_dim``, ``cells`, and ``measure_dim``
-    arguments in ``__init__``, these arrays are as follows:
+    (float), (4) measure space coordinates of the solution (1D array), (5)
+    any additional metadata associated with the solution (object), and (6) a
+    threshold which determines how high an objective value must be for a
+    solution to be inserted into a cell (float). In this class, the container is
+    implemented with separate numpy arrays that share common dimensions. Using
+    the ``solution_dim``, ``cells`, and ``measure_dim`` arguments in
+    ``__init__``, these arrays are as follows:
 
     +------------------------+----------------------------+
     | Name                   |  Shape                     |
@@ -80,6 +85,8 @@ class ArchiveBase(ABC):  # pylint: disable = too-many-instance-attributes
     +------------------------+----------------------------+
     | ``_metadata_arr``      |  ``(cells,)``              |
     +------------------------+----------------------------+
+    | ``_threshold_arr``     |  ``(cells,)``              |
+    +------------------------+----------------------------+
 
     All of these arrays are accessed via a common integer index. If we have
     index ``i``, we access its solution at ``_solution_arr[i]``, its measure
@@ -97,14 +104,18 @@ class ArchiveBase(ABC):  # pylint: disable = too-many-instance-attributes
     .. note:: Attributes beginning with an underscore are only intended to be
         accessed by child classes (i.e. they are "protected" attributes).
 
+    .. note:: The idea of archive thresholds was introduced in `Fontaine 2022
+        <https://arxiv.org/abs/2205.10752>`_. Refer to our `CMA-MAE tutorial
+        <../../tutorials/cma_mae.html>`_ for more info on thresholds, including
+        the ``learning_rate`` and ``threshold_min`` parameters.
+
     Args:
         solution_dim (int): Dimension of the solution space.
         cells (int): Number of cells in the archive. This is used to create the
             numpy arrays described above for storing archive info.
         measure_dim (int): The dimension of the measure space.
-        learning_rate (float): The learning rate of the archive. Described in
-            `Fontaine 2022 <https://arxiv.org/abs/2205.10752>`_.
-        threshold_min (float): The default threshold value for all the cells.
+        learning_rate (float): The learning rate for threshold updates.
+        threshold_min (float): The initial threshold value for all the cells.
         seed (int): Value to seed the random number generator. Set to None to
             avoid a fixed seed.
         dtype (str or data-type): Data type of the solutions, objectives,
@@ -126,6 +137,8 @@ class ArchiveBase(ABC):  # pylint: disable = too-many-instance-attributes
             coordinates of each solution.
         _metadata_arr (numpy.ndarray): Object array storing the metadata
             associated with each solution.
+        _threshold_arr (numpy.ndarray): Float array storing the threshold for
+            insertion into each cell.
         _occupied_indices (numpy.ndarray): A ``(cells,)`` array of integer
             (``np.int32``) indices that are occupied in the archive. This could
             be a list, but for efficiency, we make it a fixed-size array, where
@@ -222,6 +235,16 @@ def solution_dim(self):
         """int: Dimensionality of the solutions in the archive."""
         return self._solution_dim
 
+    @property
+    def learning_rate(self):
+        """float: The learning rate for threshold updates."""
+        return self._learning_rate
+
+    @property
+    def threshold_min(self):
+        """float: The initial threshold value for all the cells."""
+        return self._threshold_min
+
     @property
     def stats(self):
         """:class:`ArchiveStats`: Statistics about the archive.
@@ -392,11 +415,6 @@ def index_of_single(self, measures):
         check_finite(measures, "measures")
         return self.index_of(measures[None])[0]
 
-    _ADD_WARNING = (" Note that starting in pyribs 0.5.0, add() takes in a "
-                    "batch of solutions unlike in pyribs 0.4.0, where add() "
-                    "only took in a single solution.")
-
-    # TODO: Update value_batch documentation here and in add_single.
     def add(self,
             solution_batch,
             objective_batch,
@@ -404,12 +422,21 @@ def add(self,
             metadata_batch=None):
         """Inserts a batch of solutions into the archive.
 
-        Each solution is only inserted if it has a higher objective than the
-        elite previously in the corresponding cell. If multiple solutions in the
-        batch end up in the same cell, we only insert the solution with the
-        highest objective. If multiple solutions end up in the same cell and tie
-        for the highest objective, we insert the solution that appears first in
-        the batch.
+        Each solution is only inserted if it has a higher ``objective`` than the
+        threshold of the corresponding cell. For the default values of
+        ``learning_rate`` and ``threshold_min``, this threshold is simply the
+        objective value of the elite previously in the cell.  If multiple
+        solutions in the batch end up in the same cell, we only insert the
+        solution with the highest objective. If multiple solutions end up in the
+        same cell and tie for the highest objective, we insert the solution that
+        appears first in the batch.
+
+        For the default values of ``learning_rate`` and ``threshold_min``, the
+        threshold for each cell is updated by taking the maximum objective value
+        among all the solutions that landed in the cell, resulting in the same
+        behavior as in the vanilla MAP-Elites archive. However, for other
+        settings, the threshold is updated with the batch update rule described
+        in the appendix of `Fontaine 2022 <https://arxiv.org/abs/2205.10752>`_.
 
         .. note:: The indices of all arguments should "correspond" to each
             other, i.e. ``solution_batch[i]``, ``objective_batch[i]``,
@@ -465,8 +492,10 @@ def add(self,
               status_batch]``.
 
             - **value_batch** (:attr:`dtype`): An array with values for each
-              solution in the batch. The meaning of each ``value`` depends on
-              the corresponding ``status``:
+              solution in the batch. With the default values of ``learning_rate
+              = 1.0`` and ``threshold_min = -np.inf``, the meaning of each value
+              depends on the corresponding ``status`` and is identical to that
+              in CMA-ME (`Fontaine 2020 <https://arxiv.org/abs/1912.02400>`_):
 
               - ``0`` (not added): The value is the "negative improvement," i.e.
                 the objective of the solution passed in minus the objective of
@@ -478,6 +507,11 @@ def add(self,
                 of the elite previously in the archive.
               - ``2`` (new cell): The value is just the objective of the
                 solution.
+
+              In contrast, for other values of ``learning_rate`` and
+              ``threshold_min``, each value is equivalent to the objective value
+              of the solution minus the threshold of its corresponding cell in
+              the archive.
         Raises:
             ValueError: The array arguments do not match their specified shapes.
             ValueError: ``objective_batch`` or ``measures_batch`` has non-finite
@@ -489,37 +523,37 @@ def add(self,
 
         solution_batch = np.asarray(solution_batch)
         check_batch_shape(solution_batch, "solution_batch", self.solution_dim,
-                          "solution_dim", self._ADD_WARNING)
+                          "solution_dim", _ADD_WARNING)
         batch_size = solution_batch.shape[0]
 
         objective_batch = np.asarray(objective_batch, self.dtype)
-        check_is_1d(objective_batch, "objective_batch", self._ADD_WARNING)
+        check_is_1d(objective_batch, "objective_batch", _ADD_WARNING)
         check_solution_batch_dim(objective_batch,
                                  "objective_batch",
                                  batch_size,
                                  is_1d=True,
-                                 extra_msg=self._ADD_WARNING)
+                                 extra_msg=_ADD_WARNING)
         check_finite(objective_batch, "objective_batch")
 
         measures_batch = np.asarray(measures_batch)
         check_batch_shape(measures_batch, "measures_batch", self.measure_dim,
-                          "measure_dim", self._ADD_WARNING)
+                          "measure_dim", _ADD_WARNING)
         check_solution_batch_dim(measures_batch,
                                  "measures_batch",
                                  batch_size,
                                  is_1d=False,
-                                 extra_msg=self._ADD_WARNING)
+                                 extra_msg=_ADD_WARNING)
         check_finite(measures_batch, "measures_batch")
 
         metadata_batch = (np.empty(batch_size, dtype=object) if
                           metadata_batch is None else np.asarray(metadata_batch,
                                                                  dtype=object))
-        check_is_1d(metadata_batch, "metadata_batch", self._ADD_WARNING)
+        check_is_1d(metadata_batch, "metadata_batch", _ADD_WARNING)
         check_solution_batch_dim(metadata_batch,
                                  "metadata_batch",
                                  batch_size,
                                  is_1d=True,
-                                 extra_msg=self._ADD_WARNING)
+                                 extra_msg=_ADD_WARNING)
 
         ## Step 2: Compute status_batch and value_batch ##
 
@@ -666,8 +700,11 @@ def add(self,
     def add_single(self, solution, objective, measures, metadata=None):
         """Inserts a single solution into the archive.
 
-        The solution is only inserted if it has a higher ``objective``
-        than the elite previously in the corresponding cell.
+        The solution is only inserted if it has a higher ``objective`` than the
+        threshold of the corresponding cell. For the default values of
+        ``learning_rate`` and ``threshold_min``, this threshold is simply the
+        objective value of the elite previously in the cell.  The threshold is
+        also updated if the solution was inserted.
 
         .. note::
             To make it more amenable to modifications, this method's
diff --git a/ribs/archives/_cvt_archive.py b/ribs/archives/_cvt_archive.py
index 4339ad72d..3148406e8 100644
--- a/ribs/archives/_cvt_archive.py
+++ b/ribs/archives/_cvt_archive.py
@@ -50,6 +50,11 @@ class CVTArchive(ArchiveBase):
     and pass them into ``custom_centroids`` when constructing archives for
     subsequent experiments.
 
+    .. note:: The idea of archive thresholds was introduced in `Fontaine 2022
+        <https://arxiv.org/abs/2205.10752>`_. Refer to our `CMA-MAE tutorial
+        <../../tutorials/cma_mae.html>`_ for more info on thresholds, including
+        the ``learning_rate`` and ``threshold_min`` parameters.
+
     Args:
         solution_dim (int): Dimension of the solution space.
         cells (int): The number of cells to use in the archive, equivalent to
@@ -60,9 +65,8 @@ class CVTArchive(ArchiveBase):
             (inclusive), and the second dimension should have bounds
             :math:`[-2,2]` (inclusive). ``ranges`` should be the same length as
             ``dims``.
-        learning_rate (float): The learning rate of the archive. Described in
-            `Fontaine 2022 <https://arxiv.org/abs/2205.10752>`_.
-        threshold_min (float): The default threshold value for all the cells.
+        learning_rate (float): The learning rate for threshold updates.
+        threshold_min (float): The initial threshold value for all the cells.
         seed (int): Value to seed the random number generator as well as
             :func:`~sklearn.cluster.k_means`. Set to None to avoid a fixed seed.
         dtype (str or data-type): Data type of the solutions, objectives,
diff --git a/ribs/archives/_grid_archive.py b/ribs/archives/_grid_archive.py
index 77554a1a7..cc0e6150f 100644
--- a/ribs/archives/_grid_archive.py
+++ b/ribs/archives/_grid_archive.py
@@ -15,6 +15,11 @@ class GridArchive(ArchiveBase):
     solution that `maximizes` the objective function for the measures in that
     cell.
 
+    .. note:: The idea of archive thresholds was introduced in `Fontaine 2022
+        <https://arxiv.org/abs/2205.10752>`_. Refer to our `CMA-MAE tutorial
+        <../../tutorials/cma_mae.html>`_ for more info on thresholds, including
+        the ``learning_rate`` and ``threshold_min`` parameters.
+
     Args:
         solution_dim (int): Dimension of the solution space.
         dims (array-like of int): Number of cells in each dimension of the
@@ -32,9 +37,8 @@ class GridArchive(ArchiveBase):
             method -- refer to the implementation `here
             <../_modules/ribs/archives/_grid_archive.html#GridArchive.index_of>`_.
             Pass this parameter to configure that epsilon.
-        learning_rate (float): The learning rate of the archive. Described in
-            `Fontaine 2022 <https://arxiv.org/abs/2205.10752>`_.
-        threshold_min (float): The default threshold value for all the cells.
+        learning_rate (float): The learning rate for threshold updates.
+        threshold_min (float): The initial threshold value for all the cells.
         seed (int): Value to seed the random number generator. Set to None to
             avoid a fixed seed.
         dtype (str or data-type): Data type of the solutions, objectives,
diff --git a/tests/archives/archive_base_test.py b/tests/archives/archive_base_test.py
index 5f583f788..ee0c7313b 100644
--- a/tests/archives/archive_base_test.py
+++ b/tests/archives/archive_base_test.py
@@ -236,6 +236,14 @@ def test_solution_dim_correct(data):
     assert data.archive.solution_dim == len(data.solution)
 
 
+def test_learning_rate_correct(data):
+    assert data.archive.learning_rate == 1.0  # Default value.
+
+
+def test_threshold_min_correct(data):
+    assert data.archive.threshold_min == -np.inf  # Default value.
+
+
 def test_basic_stats(data):
     assert data.archive.stats.num_elites == 0
     assert data.archive.stats.coverage == 0.0

From 4349b0f8df07b958862142ae64d3d93b7609b21f Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Sun, 16 Oct 2022 22:55:28 -0700
Subject: [PATCH 52/53] More tests; fix bugs

---
 ribs/archives/_archive_base.py      |  24 ++---
 tests/archives/grid_archive_test.py | 158 +++++++++++++++++++++++++---
 2 files changed, 152 insertions(+), 30 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index fa3e917ab..14a147616 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -337,26 +337,16 @@ def _compute_new_thresholds(self, threshold_arr, objective_batch,
         objective_sizes = objective_sizes[threshold_update_indices]
         objective_sums = objective_sums[threshold_update_indices]
 
-        # Sum of geometric series (1 - learning_rate)^i from i = 0 to i = n - 1.
-        # See https://en.wikipedia.org/wiki/Geometric_series#Sum.
-        ratio = self.dtype(1.0 - learning_rate)
-        if ratio == 1.0:
-            geometric_sums = objective_sizes
-        else:
-            geometric_sums = (1 - ratio**objective_sizes) / (1 - ratio)
-
-        update = (learning_rate * (objective_sums / objective_sizes) *
-                  geometric_sums)
-
         # Unlike in add_single, we do not need to worry about
         # old_threshold having -np.inf here as a result of threshold_min
         # being -np.inf. This is because the case with threshold_min =
         # -np.inf is handled separately since we compute the new
         # threshold based on the max objective in each cell in that case.
         old_threshold = np.copy(threshold_arr[threshold_update_indices])
-        prev = old_threshold * ratio**objective_sizes
 
-        new_threshold_batch = prev + update
+        ratio = self.dtype(1.0 - learning_rate)**objective_sizes
+        new_threshold_batch = (ratio * old_threshold +
+                               (objective_sums / objective_sizes) * (1 - ratio))
 
         return new_threshold_batch, threshold_update_indices
 
@@ -567,9 +557,11 @@ def add(self,
         # Compute the statuses -- these are all boolean arrays of length
         # batch_size.
         already_occupied = self._occupied_arr[index_batch]
-        is_new = ~already_occupied
-        improve_existing = (objective_batch >
-                            old_threshold_batch) & already_occupied
+        # In the case where we want CMA-ME behavior, threshold_arr[index]
+        # is -inf for new cells, which satisfies the condition for can_be_added.
+        can_be_added = objective_batch > old_threshold_batch
+        is_new = can_be_added & ~already_occupied
+        improve_existing = can_be_added & already_occupied
         status_batch = np.zeros(batch_size, dtype=np.int32)
         status_batch[is_new] = 2
         status_batch[improve_existing] = 1
diff --git a/tests/archives/grid_archive_test.py b/tests/archives/grid_archive_test.py
index d9194cf78..bba09d743 100644
--- a/tests/archives/grid_archive_test.py
+++ b/tests/archives/grid_archive_test.py
@@ -224,7 +224,7 @@ def test_add_single_without_overwrite(data, add_mode):
                          data.measures, data.grid_indices, data.metadata)
 
 
-def test_add_single_with_threshold_update(add_mode):
+def test_add_single_threshold_update(add_mode):
     archive = GridArchive(
         solution_dim=3,
         dims=[10, 10],
@@ -268,6 +268,21 @@ def test_add_single_with_threshold_update(add_mode):
     assert np.isclose(value, 0.1)  # -0.8 - (-0.9)
 
 
+def test_add_single_wrong_shapes(data):
+    with pytest.raises(ValueError):
+        data.archive.add_single(
+            solution=[1, 1],  # 2D instead of 3D solution.
+            objective=0,
+            measures=[0, 0],
+        )
+    with pytest.raises(ValueError):
+        data.archive.add_single(
+            solution=[0, 0, 0],
+            objective=0,
+            measures=[1, 1, 1],  # 3D instead of 2D measures.
+        )
+
+
 def test_add_batch_all_new(data):
     status_batch, value_batch = data.archive.add(
         # 4 solutions of arbitrary value.
@@ -415,19 +430,134 @@ def test_add_batch_first_solution_wins_in_ties(data):
     )
 
 
-def test_add_single_wrong_shapes(data):
-    with pytest.raises(ValueError):
-        data.archive.add_single(
-            solution=[1, 1],  # 2D instead of 3D solution.
-            objective=0,
-            measures=[0, 0],
-        )
-    with pytest.raises(ValueError):
-        data.archive.add_single(
-            solution=[0, 0, 0],
-            objective=0,
-            measures=[1, 1, 1],  # 3D instead of 2D measures.
-        )
+def test_add_batch_not_inserted_if_below_threshold_min():
+    archive = GridArchive(
+        solution_dim=3,
+        dims=[10, 10],
+        ranges=[(-1, 1), (-1, 1)],
+        threshold_min=-10.0,
+        learning_rate=0.1,
+    )
+
+    status_batch, value_batch = archive.add(
+        solution_batch=[[1, 2, 3]] * 4,
+        objective_batch=[-20.0, -20.0, 10.0, 10.0],
+        measures_batch=[[0.0, 0.0]] * 4,
+    )
+
+    # The first two solutions should not have been inserted since they did not
+    # cross the threshold_min of -10.0.
+    assert (status_batch == [0, 0, 2, 2]).all()
+    assert np.isclose(value_batch, [-10.0, -10.0, 20.0, 20.0]).all()
+
+    assert_archive_elite_batch(
+        archive=archive,
+        batch_size=1,
+        solution_batch=[[1, 2, 3]],
+        objective_batch=[10.0],
+        measures_batch=[[0.0, 0.0]],
+        metadata_batch=[None],
+        grid_indices_batch=[[5, 5]],
+    )
+
+
+def test_add_batch_threshold_update():
+    archive = GridArchive(
+        solution_dim=3,
+        dims=[10, 10],
+        ranges=[(-1, 1), (-1, 1)],
+        threshold_min=-1.0,
+        learning_rate=0.1,
+    )
+    solution = [1, 2, 3]
+    measures = [0.1, 0.1]
+    measures2 = [-0.1, -0.1]
+
+    # Add new solutions to the archive in two cells determined by measures and
+    # measures2.
+    status_batch, value_batch = archive.add(
+        [solution, solution, solution, solution, solution, solution],
+        # The first three solutions are inserted since they cross
+        # threshold_min, but the last solution is not inserted since it does not
+        # cross threshold_min.
+        [0.0, 1.0, 2.0, 10.0, 100.0, -10.0],
+        [measures, measures, measures, measures2, measures2, measures2],
+    )
+
+    assert (status_batch == [2, 2, 2, 2, 2, 0]).all()
+    assert np.isclose(
+        value_batch, [1.0, 2.0, 3.0, 11.0, 101.0, -9.0]).all()  # [...] - (-1.0)
+
+    # Thresholds based on batch update rule should now be
+    # (1 - 0.1)**3 * -1.0 + (0.0 + 1.0 + 2.0) / 3 * (1 - (1 - 0.1)**3) = -0.458
+    # and
+    # (1 - 0.1)**2 * -1.0 + (10.0 + 100.0) / 2 * (1 - (1 - 0.1)**2) = 9.64
+
+    # Mix between solutions which are inserted and not inserted.
+    status_batch, value_batch = archive.add(
+        [solution, solution, solution, solution],
+        [-0.95, -0.457, 9.63, 9.65],
+        [measures, measures, measures2, measures2],
+    )
+
+    assert (status_batch == [0, 1, 0, 1]).all()
+    # [-0.95 - (-0.458), -0.458 - (-0.457), 9.63 - 9.64, 9.65 - 9.64]
+    assert np.isclose(value_batch, [-0.492, 0.001, -0.01, 0.01]).all()
+
+    # Thresholds should now be
+    # (1 - 0.1)**1 * -0.458 + (-0.457) / 1 * (1 - (1 - 0.1)**1) = -0.4579
+    # and
+    # (1 - 0.1)**1 * 9.64 + (9.65) / 1 * (1 - (1 - 0.1)**1) = 9.641
+
+    # Again mix between solutions which are inserted and not inserted.
+    status_batch, value_batch = archive.add(
+        [solution, solution, solution, solution],
+        [-0.4580, -0.4578, 9.640, 9.642],
+        [measures, measures, measures2, measures2],
+    )
+
+    assert (status_batch == [0, 1, 0, 1]).all()
+    assert np.isclose(value_batch, [-0.0001, 0.0001, -0.001, 0.001]).all()
+
+
+def test_add_batch_threshold_update_inf_threshold_min():
+    # These default values of threshold_min and learning_rate induce special
+    # CMA-ME behavior for threshold updates.
+    archive = GridArchive(
+        solution_dim=3,
+        dims=[10, 10],
+        ranges=[(-1, 1), (-1, 1)],
+        threshold_min=-np.inf,
+        learning_rate=1.0,
+    )
+    solution = [1, 2, 3]
+    measures = [0.1, 0.1]
+    measures2 = [-0.1, -0.1]
+
+    # Add new solutions to the archive.
+    status_batch, value_batch = archive.add(
+        [solution, solution, solution, solution, solution, solution],
+        [0.0, 1.0, 2.0, -10.0, 10.0, 100.0],
+        [measures, measures, measures, measures2, measures2, measures2],
+    )
+
+    # Value is same as objective since these are new cells.
+    assert (status_batch == [2, 2, 2, 2, 2, 2]).all()
+    assert np.isclose(value_batch, [0.0, 1.0, 2.0, -10.0, 10.0, 100.0]).all()
+
+    # Thresholds are updated based on maximum values in each cell, i.e. 2.0 and
+    # 100.0.
+
+    # Mix between solutions which are inserted and not inserted.
+    status_batch, value_batch = archive.add(
+        [solution, solution, solution, solution],
+        [1.0, 10.0, 99.0, 101.0],
+        [measures, measures, measures2, measures2],
+    )
+
+    assert (status_batch == [0, 1, 0, 1]).all()
+    # [1.0 - 2.0, 10.0 - 2.0, 99.0 - 100.0, 101.0 - 100.0]
+    assert np.isclose(value_batch, [-1.0, 8.0, -1.0, 1.0]).all()
 
 
 def test_add_batch_wrong_shapes(data):

From fda24bc7d6cc5aebc04432b2f13b499e403524b4 Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Mon, 17 Oct 2022 02:08:31 -0700
Subject: [PATCH 53/53] Comment

---
 ribs/archives/_archive_base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ribs/archives/_archive_base.py b/ribs/archives/_archive_base.py
index 14a147616..cd9b13893 100644
--- a/ribs/archives/_archive_base.py
+++ b/ribs/archives/_archive_base.py
@@ -571,8 +571,8 @@ def add(self,
         old_objective_batch[is_new] = self.dtype(0)
 
         # If threshold_min is -inf, then we want CMA-ME behavior, which
-        # will compute the improvement value w.r.t zero. Otherwise, we will
-        # compute w.r.t. threshold_min.
+        # will compute the improvement value of new solutions w.r.t zero.
+        # Otherwise, we will compute w.r.t. threshold_min.
         old_threshold_batch[is_new] = (self.dtype(0) if self._threshold_min
                                        == -np.inf else self._threshold_min)
         value_batch = objective_batch - old_threshold_batch