Run pre-commit

examples/aurora.ipynb

@@ -27,7 +27,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -45,7 +45,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [

examples/sac_pbt.ipynb

@@ -429,7 +429,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "aa4b43a1",
+   "id": "16",
    "metadata": {},
    "source": [
     "### Visualize learnt behaviors"

qdax/baselines/sac.py

@@ -402,9 +402,7 @@ def eval_qd_policy_fn(
         true_returns = jnp.nansum(transitions.rewards, axis=0)
         true_return = jnp.mean(true_returns, axis=-1)
 
-        transitions = jax.tree.map(
-            lambda x: jnp.swapaxes(x, 0, 1), transitions
-        )
+        transitions = jax.tree.map(lambda x: jnp.swapaxes(x, 0, 1), transitions)
         masks = jnp.isnan(transitions.rewards)
         descriptors = descriptor_extraction_fn(transitions, masks)
 

qdax/baselines/td3.py

@@ -341,9 +341,7 @@ def eval_qd_policy_fn(
         true_returns = jnp.nansum(transitions.rewards, axis=0)
         true_return = jnp.mean(true_returns, axis=-1)
 
-        transitions = jax.tree.map(
-            lambda x: jnp.swapaxes(x, 0, 1), transitions
-        )
+        transitions = jax.tree.map(lambda x: jnp.swapaxes(x, 0, 1), transitions)
         masks = jnp.isnan(transitions.rewards)
         descriptors = descriptor_extraction_fn(transitions, masks)
 

qdax/core/containers/ga_repertoire.py

@@ -137,9 +137,7 @@ def add(
         survivor_indices = indices[: self.size]
 
         # keep only the best ones
-        new_candidates = jax.tree.map(
-            lambda x: x[survivor_indices], candidates
-        )
+        new_candidates = jax.tree.map(lambda x: x[survivor_indices], candidates)
 
         new_repertoire = self.replace(
             genotypes=new_candidates, fitnesses=candidates_fitnesses[survivor_indices]

qdax/core/containers/mome_repertoire.py

@@ -114,9 +114,7 @@ def sample(self, key: RNGKey, num_samples: int) -> Genotype:
         cells_idx = jax.random.choice(subkey, indices, shape=(num_samples,), p=p)
 
         # get genotypes (front) from the chosen indices
-        pareto_front_genotypes = jax.tree.map(
-            lambda x: x[cells_idx], self.genotypes
-        )
+        pareto_front_genotypes = jax.tree.map(lambda x: x[cells_idx], self.genotypes)
 
         # prepare second sampling function
         sample_in_fronts = jax.vmap(self._sample_in_masked_pareto_front)
@@ -130,9 +128,7 @@ def sample(self, key: RNGKey, num_samples: int) -> Genotype:
         )
 
         # remove the dim coming from pareto front
-        sampled_genotypes = jax.tree.map(
-            lambda x: x.squeeze(axis=1), sampled_genotypes
-        )
+        sampled_genotypes = jax.tree.map(lambda x: x.squeeze(axis=1), sampled_genotypes)
 
         return sampled_genotypes
 
@@ -287,16 +283,12 @@ def _add_one(
             index = index.astype(jnp.int32)
 
             # get current repertoire cell data
-            cell_genotype = jax.tree.map(
-                lambda x: x[index][0], carry.genotypes
-            )
+            cell_genotype = jax.tree.map(lambda x: x[index][0], carry.genotypes)
             cell_fitness = carry.fitnesses[index][0]
             cell_descriptor = carry.descriptors[index][0]
             cell_mask = jnp.any(cell_fitness == -jnp.inf, axis=-1)
 
-            new_genotypes = jax.tree.map(
-                lambda x: jnp.expand_dims(x, axis=0), genotype
-            )
+            new_genotypes = jax.tree.map(lambda x: jnp.expand_dims(x, axis=0), genotype)
 
             # update pareto front
             (

qdax/core/emitters/cma_emitter.py

@@ -216,9 +216,7 @@ def state_update(
         sorted_indices = jnp.flip(jnp.argsort(ranking_criteria))
 
         # sort the candidates
-        sorted_candidates = jax.tree.map(
-            lambda x: x[sorted_indices], genotypes
-        )
+        sorted_candidates = jax.tree.map(lambda x: x[sorted_indices], genotypes)
         sorted_improvements = improvements[sorted_indices]
 
         # compute reinitialize condition

qdax/core/emitters/cma_pool_emitter.py

@@ -155,9 +155,7 @@ def state_update(
         current_index = emitter_state.current_index
         emitter_states = emitter_state.emitter_states
 
-        used_emitter_state = jax.tree.map(
-            lambda x: x[current_index], emitter_states
-        )
+        used_emitter_state = jax.tree.map(lambda x: x[current_index], emitter_states)
 
         # update the used emitter state
         used_emitter_state = self._emitter.state_update(

qdax/core/emitters/mutation_operators.py

@@ -168,9 +168,7 @@ def polynomial_crossover(
     )
     crossover_fn = jax.vmap(crossover_fn)
     # TODO: check that key usage is correct
-    x = jax.tree.map(
-        lambda x1_, x2_: crossover_fn(x1_, x2_, crossover_keys), x1, x2
-    )
+    x = jax.tree.map(lambda x1_, x2_: crossover_fn(x1_, x2_, crossover_keys), x1, x2)
     return x
 
 
@@ -223,8 +221,6 @@ def _variation_fn(x1: jnp.ndarray, x2: jnp.ndarray, key: RNGKey) -> jnp.ndarray:
     keys_tree = jax.tree.unflatten(jax.tree.structure(x1), keys)
 
     # apply isolinedd to each branch of the tree
-    x = jax.tree.map(
-        lambda y1, y2, key: _variation_fn(y1, y2, key), x1, x2, keys_tree
-    )
+    x = jax.tree.map(lambda y1, y2, key: _variation_fn(y1, y2, key), x1, x2, keys_tree)
 
     return x

qdax/core/emitters/omg_mega_emitter.py

@@ -202,9 +202,7 @@ def emit(
         update_grad = jnp.sum(jax.vmap(lambda x, y: x * y)(coeffs, grads), axis=-1)
 
         # update the genotypes
-        new_genotypes = jax.tree.map(
-            lambda x, y: x + y, genotypes, update_grad
-        )
+        new_genotypes = jax.tree.map(lambda x, y: x + y, genotypes, update_grad)
 
         return new_genotypes, {}
 

qdax/environments/wrappers.py

@@ -22,9 +22,7 @@ def reset(self, rng: jp.ndarray) -> State:
         reset_state = self.env.reset(rng)
         reset_state.metrics["reward"] = reset_state.reward
         eval_metrics = CompletedEvalMetrics(
-            current_episode_metrics=jax.tree.map(
-                jp.zeros_like, reset_state.metrics
-            ),
+            current_episode_metrics=jax.tree.map(jp.zeros_like, reset_state.metrics),
             completed_episodes_metrics=jax.tree.map(
                 lambda x: jp.zeros_like(jp.sum(x)), reset_state.metrics
             ),

qdax/tasks/brax_envs.py

@@ -141,9 +141,7 @@ def scoring_function_brax_envs(
 
     # Reset environments
     key, subkey = jax.random.split(key)
-    keys = jax.random.split(
-        subkey, jax.tree.leaves(policies_params)[0].shape[0]
-    )
+    keys = jax.random.split(subkey, jax.tree.leaves(policies_params)[0].shape[0])
     init_states = jax.vmap(play_reset_fn)(keys)
 
     # Step environments

tests/core_test/neuroevolution_test/buffers_test/buffer_test.py

@@ -42,9 +42,7 @@ def test_insert_batch() -> None:
         buffer_size=buffer_size, transition=dummy_transition
     )
 
-    simple_transition = jax.tree.map(
-        lambda x: x.repeat(3, axis=0), dummy_transition
-    )
+    simple_transition = jax.tree.map(lambda x: x.repeat(3, axis=0), dummy_transition)
     simple_transition = simple_transition.replace(rewards=jnp.arange(3))
     data = QDTransition.from_flatten(replay_buffer.data, dummy_transition)
     pytest.assume(
@@ -85,9 +83,7 @@ def test_sample() -> None:
         buffer_size=buffer_size, transition=dummy_transition
     )
 
-    simple_transition = jax.tree.map(
-        lambda x: x.repeat(3, axis=0), dummy_transition
-    )
+    simple_transition = jax.tree.map(lambda x: x.repeat(3, axis=0), dummy_transition)
     simple_transition = simple_transition.replace(rewards=jnp.arange(3))
 
     replay_buffer = replay_buffer.insert(simple_transition)