Merge branch 'main' into develop

Author: Lookatator

README.md

@@ -114,10 +114,14 @@ centroids = compute_euclidean_centroids(
 key, subkey = jax.random.split(key)
 repertoire, emitter_state, metrics = map_elites.init(init_variables, centroids, subkey)
 
+# Jit the update function for faster iterations
+update_fn = jax.jit(map_elites.update)
+
 # Run MAP-Elites loop
 for i in range(num_iterations):
     key, subkey = jax.random.split(key)
-    (repertoire, emitter_state, metrics,) = map_elites.update(
+
+    (repertoire, emitter_state, metrics,) = update_fn(
         repertoire,
         emitter_state,
         subkey,
@@ -146,7 +150,7 @@ QDax currently supports the following algorithms:
 | [Multi-Objective MAP-Elites (MOME)](https://arxiv.org/abs/2202.03057)                                                         | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/mome.ipynb)       |
 | [MAP-Elites Evolution Strategies (MEES)](https://dl.acm.org/doi/pdf/10.1145/3377930.3390217)                                  | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/mees.ipynb)       |
 | [MAP-Elites PBT (ME-PBT)](https://openreview.net/forum?id=CBfYffLqWqb)                                                        | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/me_sac_pbt.ipynb) |
-| [MAP-Elites Low-Spread (ME-LS)](https://dl.acm.org/doi/abs/10.1145/3583131.3590433)                                           | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/me_ls.ipynb)      |
+| [MAP-Elites Low-Spread (ME-LS)](https://dl.acm.org/doi/abs/10.1145/3583131.3590433)                                           | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/mels.ipynb)      |
 
 
 ## QDax baseline algorithms
@@ -200,6 +204,9 @@ QDax was developed and is maintained by the [Adaptive & Intelligent Robotics Lab
 <a href="https://github.com/maxiallard" title="Maxime Allard"><img src="https://github.com/maxiallard.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/Lookatator" title="Luca Grilloti"><img src="https://github.com/Lookatator.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/manon-but-yes" title="Manon Flageat"><img src="https://github.com/manon-but-yes.png" height="auto" width="50" style="border-radius:50%"></a>
+<a href="https://github.com/maxencefaldor" title="Maxence Faldor"><img src="https://github.com/maxencefaldor.png" height="auto" width="50" style="border-radius:50%"></a>
+<a href="https://github.com/hannah-jan" title="Hannah Janmohamed"><img src="https://github.com/hannah-jan.png" height="auto" width="50" style="border-radius:50%"></a>
+<a href="https://github.com/LisaCoiffard" title="Lisa Coiffard"><img src="https://github.com/LisaCoiffard.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/Aneoshun" title="Antoine Cully"><img src="https://github.com/Aneoshun.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/felixchalumeau" title="Felix Chalumeau"><img src="https://github.com/felixchalumeau.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/ranzenTom" title="Thomas Pierrot"><img src="https://github.com/ranzenTom.png" height="auto" width="50" style="border-radius:50%"></a>

examples/aurora.ipynb

@@ -370,7 +370,7 @@
     "    model, subkey, (1, *observations_dims)\n",
     ")\n",
     "\n",
-    "print(jax.tree_map(lambda x: x.shape, model_params))\n",
+    "print(jax.tree.map(lambda x: x.shape, model_params))\n",
     "\n",
     "# Define the encoder function\n",
     "encoder_fn = jax.jit(\n",

examples/me_sac_pbt.ipynb

@@ -236,10 +236,10 @@
     "\n",
     "def scoring_function(genotypes, key):\n",
     "    population_size = jax.tree.leaves(genotypes)[0].shape[0]\n",
-    "    first_states = jax.tree_map(\n",
+    "    first_states = jax.tree.map(\n",
     "        lambda x: jnp.expand_dims(x, axis=0), eval_env_first_states\n",
     "    )\n",
-    "    first_states = jax.tree_map(\n",
+    "    first_states = jax.tree.map(\n",
     "        lambda x: jnp.repeat(x, population_size, axis=0), first_states\n",
     "    )\n",
     "    population_returns, population_descriptors, _, _ = eval_policy(genotypes, first_states)\n",
@@ -378,7 +378,7 @@
     "    repertoire, emitter_state, metrics = update_fn(\n",
     "        repertoire, emitter_state, keys\n",
     "    )\n",
-    "    metrics_cpu = jax.tree_map(\n",
+    "    metrics_cpu = jax.tree.map(\n",
     "        lambda x: jax.device_put(x, jax.devices(\"cpu\")[0])[0], metrics\n",
     "    )\n",
     "    timelapse = time.time() - start_time\n",
@@ -401,7 +401,7 @@
    "outputs": [],
    "source": [
     "# Create the performance evolution plots and visualize final grid\n",
-    "repertoire_cpu = jax.tree_map(\n",
+    "repertoire_cpu = jax.tree.map(\n",
     "    lambda x: jax.device_put(x, jax.devices(\"cpu\")[0])[0], repertoire\n",
     ")\n",
     "num_loops_with_init = num_loops + 1\n",
@@ -510,7 +510,7 @@
     "key, subkey = jax.random.split(key)\n",
     "env_state = jax.jit(env.reset)(rng=subkey)\n",
     "\n",
-    "training_state, env_state = jax.tree_map(\n",
+    "training_state, env_state = jax.tree.map(\n",
     "    lambda x: jnp.expand_dims(x, axis=0), (training_state, env_state)\n",
     ")\n",
     "\n",
@@ -529,7 +529,7 @@
    "outputs": [],
    "source": [
     "rollout = [\n",
-    "    jax.tree_map(lambda x: jax.device_put(x[0], jax.devices(\"cpu\")[0]), env_state)\n",
+    "    jax.tree.map(lambda x: jax.device_put(x[0], jax.devices(\"cpu\")[0]), env_state)\n",
     "    for env_state in rollout\n",
     "]"
    ]

examples/me_td3_pbt.ipynb

@@ -238,10 +238,10 @@
     "\n",
     "def scoring_function(genotypes, key):\n",
     "    population_size = jax.tree_leaves(genotypes)[0].shape[0]\n",
-    "    first_states = jax.tree_map(\n",
+    "    first_states = jax.tree.map(\n",
     "        lambda x: jnp.expand_dims(x, axis=0), eval_env_first_states\n",
     "    )\n",
-    "    first_states = jax.tree_map(\n",
+    "    first_states = jax.tree.map(\n",
     "        lambda x: jnp.repeat(x, population_size, axis=0), first_states\n",
     "    )\n",
     "    population_returns, population_descriptors, _, _ = eval_policy(genotypes, first_states)\n",
@@ -379,7 +379,7 @@
     "    repertoire, emitter_state, metrics = update_fn(\n",
     "        repertoire, emitter_state, keys\n",
     "    )\n",
-    "    metrics_cpu = jax.tree_map(lambda x: jax.device_get(x)[0], metrics)\n",
+    "    metrics_cpu = jax.tree.map(lambda x: jax.device_get(x)[0], metrics)\n",
     "    timelapse = time.time() - start_time\n",
     "\n",
     "    # log metrics\n",
@@ -391,7 +391,7 @@
     "            all_metrics[k] = v\n",
     "\n",
     "    if i % save_repertoire_freq == 0:\n",
-    "        repertoires.append(jax.tree_map(lambda x: jax.device_get(x)[0], repertoire))"
+    "        repertoires.append(jax.tree.map(lambda x: jax.device_get(x)[0], repertoire))"
    ]
   },
   {

qdax/__init__.py

@@ -1 +1 @@
-__version__ = "0.4.1"
+__version__ = "0.5.0"

qdax/tasks/brax/v2/wrappers/eval_metrics_wrapper.py

@@ -22,10 +22,8 @@ def reset(self, rng: jnp.ndarray) -> State:
         reset_state = self.env.reset(rng)
         reset_state.metrics["reward"] = reset_state.reward
         eval_metrics = CompletedEvalMetrics(
-            current_episode_metrics=jax.tree_util.tree_map(
-                jnp.zeros_like, reset_state.metrics
-            ),
-            completed_episodes_metrics=jax.tree_util.tree_map(
+            current_episode_metrics=jax.tree.map(jnp.zeros_like, reset_state.metrics),
+            completed_episodes_metrics=jax.tree.map(
                 lambda x: jnp.zeros_like(jnp.sum(x)), reset_state.metrics
             ),
             completed_episodes=jnp.zeros(()),
@@ -46,16 +44,17 @@ def step(self, state: State, action: jnp.ndarray) -> State:
         completed_episodes_steps = state_metrics.completed_episodes_steps + jnp.sum(
             nstate.info["steps"] * nstate.done
         )
-        current_episode_metrics = jax.tree_util.tree_map(
+
+        current_episode_metrics = jax.tree.map(
             lambda a, b: a + b, state_metrics.current_episode_metrics, nstate.metrics
         )
         completed_episodes = state_metrics.completed_episodes + jnp.sum(nstate.done)
-        completed_episodes_metrics = jax.tree_util.tree_map(
+        completed_episodes_metrics = jax.tree.map(
             lambda a, b: a + jnp.sum(b * nstate.done),
             state_metrics.completed_episodes_metrics,
             current_episode_metrics,
         )
-        current_episode_metrics = jax.tree_util.tree_map(
+        current_episode_metrics = jax.tree.map(
             lambda a, b: a * (1 - nstate.done) + b * nstate.done,
             current_episode_metrics,
             nstate.metrics,