Fix omg_mega emitter

adaptive-intelligent-robotics · Jan 8, 2024 · 9e75ded · 9e75ded
1 parent 84a8f97
commit 9e75ded
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/qdax/core/emitters/omg_mega_emitter.py b/qdax/core/emitters/omg_mega_emitter.py
@@ -118,6 +118,18 @@ def init(
             genotype=gradient_genotype, centroids=self._centroids
         )
 
+        # get gradients out of the extra scores
+        assert "gradients" in extra_scores.keys(), "Missing gradients or wrong key"
+        gradients = extra_scores["gradients"]
+
+        # update the gradients repertoire
+        gradients_repertoire = gradients_repertoire.add(
+            gradients,
+            descriptors,
+            fitnesses,
+            extra_scores,
+        )
+
         return (
             OMGMEGAEmitterState(gradients_repertoire=gradients_repertoire),
             random_key,