Feat: remove save and load methods (#225)

* Removing save and load methods from MapElitesRepertoire and children classes, as they were tiresome to maintain and implement for new types of repertoires.

* Instead, users can use their own custom way to save and load repertoire, such as pickle and orbax.

* fix brax v2 wrapper inheritance causing: _n_frames=1

Author: Lookatator

examples/aurora.ipynb

@@ -21,8 +21,7 @@
     "- how to create an AURORA instance\n",
     "- which functions must be defined before training\n",
     "- how to launch a certain number of training steps\n",
-    "- how to visualise the optimization process\n",
-    "- how to save/load a repertoire"
+    "- how to visualise the optimization process"
    ]
   },
   {

examples/mapelites.ipynb

@@ -21,8 +21,7 @@
     "- how to create a Map-elites instance\n",
     "- which functions must be defined before training\n",
     "- how to launch a certain number of training steps\n",
-    "- how to visualise the optimization process\n",
-    "- how to save/load a repertoire"
+    "- how to visualise the optimization process"
    ]
   },
   {
@@ -368,77 +367,6 @@
     "fig, axes = plot_map_elites_results(env_steps=env_steps, metrics=metrics, repertoire=repertoire, min_descriptor=min_descriptor, max_descriptor=max_descriptor)"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# How to save/load a repertoire\n",
-    "\n",
-    "The following cells show how to save or load a repertoire of individuals and add a few lines to visualise the best performing individual in a simulation."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Load the final repertoire"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "repertoire_path = \"./last_repertoire/\"\n",
-    "os.makedirs(repertoire_path, exist_ok=True)\n",
-    "repertoire.save(path=repertoire_path)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Build the reconstruction function"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Init population of policies\n",
-    "key, subkey = jax.random.split(key)\n",
-    "fake_batch = jnp.zeros(shape=(env.observation_size,))\n",
-    "fake_params = policy_network.init(subkey, fake_batch)\n",
-    "\n",
-    "_, reconstruction_fn = ravel_pytree(fake_params)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Use the reconstruction function to load and re-create the repertoire"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "repertoire = MapElitesRepertoire.load(reconstruction_fn=reconstruction_fn, path=repertoire_path)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Get the best individual of the repertoire"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -469,6 +397,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# select the parameters of the best individual\n",
     "my_params = jax.tree.map(\n",
     "    lambda x: x[best_idx],\n",
     "    repertoire.genotypes\n",

examples/mapelites_asktell.ipynb

@@ -21,8 +21,7 @@
     "- how to create a Map-elites instance\n",
     "- which functions must be defined before training\n",
     "- how to launch a certain number of training steps\n",
-    "- how to visualise the optimization process\n",
-    "- how to save/load a repertoire"
+    "- how to visualise the optimization process"
    ]
   },
   {

examples/mapelites_brax_v2.ipynb

@@ -368,70 +368,6 @@
     "fig, axes = plot_map_elites_results(env_steps=env_steps, metrics=metrics, repertoire=repertoire, min_descriptor=min_descriptor, max_descriptor=max_descriptor)"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# How to save/load a repertoire\n",
-    "\n",
-    "The following cells show how to save or load a repertoire of individuals and add a few lines to visualise the best performing individual in a simulation."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Load the final repertoire"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "repertoire_path = \"./last_repertoire/\"\n",
-    "os.makedirs(repertoire_path, exist_ok=True)\n",
-    "repertoire.save(path=repertoire_path)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Build the reconstruction function"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Init population of policies\n",
-    "key, subkey = jax.random.split(key)\n",
-    "fake_batch = jnp.zeros(shape=(env.observation_size,))\n",
-    "fake_params = policy_network.init(subkey, fake_batch)\n",
-    "\n",
-    "_, reconstruction_fn = ravel_pytree(fake_params)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Use the reconstruction function to load and re-create the repertoire"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "repertoire = MapElitesRepertoire.load(reconstruction_fn=reconstruction_fn, path=repertoire_path)"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -519,7 +455,8 @@
    "outputs": [],
    "source": [
     "def save_rollout_html(env, pipeline_env_list, file_name: str):\n",
-    "  rollout_html = html.render(env.sys.replace(dt=env.dt), pipeline_env_list)\n",
+    "  sys = env.sys.tree_replace({'opt.timestep': env.dt})\n",
+    "  rollout_html = html.render(sys, pipeline_env_list)\n",
     "  with open(file_name, 'w') as f:\n",
     "      f.write(rollout_html)\n",
     "\n",

examples/mels.ipynb

@@ -20,8 +20,7 @@
     "- how to create an ME-LS instance\n",
     "- which functions must be defined before training\n",
     "- how to launch a certain number of training steps\n",
-    "- how to visualise the optimization process\n",
-    "- how to save/load a repertoire"
+    "- how to visualise the optimization process"
    ]
   },
   {
@@ -383,70 +382,6 @@
     "fig, axes = plot_map_elites_results(env_steps=env_steps, metrics=metrics, repertoire=repertoire, min_descriptor=min_descriptor, max_descriptor=max_descriptor)"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# How to save/load a repertoire\n",
-    "\n",
-    "The following cells show how to save or load a repertoire of individuals and add a few lines to visualise the best performing individual in a simulation."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Load the final repertoire"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "repertoire_path = \"./last_repertoire/\"\n",
-    "os.makedirs(repertoire_path, exist_ok=True)\n",
-    "repertoire.save(path=repertoire_path)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Build the reconstruction function"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Init population of policies\n",
-    "key, subkey = jax.random.split(key)\n",
-    "fake_batch = jnp.zeros(shape=(env.observation_size,))\n",
-    "fake_params = policy_network.init(subkey, fake_batch)\n",
-    "\n",
-    "_, reconstruction_fn = ravel_pytree(fake_params)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Use the reconstruction function to load and re-create the repertoire"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "repertoire = MELSRepertoire.load(reconstruction_fn=reconstruction_fn, path=repertoire_path)"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -488,6 +423,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# select the parameters of the best individual\n",
     "my_params = jax.tree.map(\n",
     "    lambda x: x[best_idx],\n",
     "    repertoire.genotypes\n",

examples/pga_aurora.ipynb

@@ -21,8 +21,7 @@
     "- how to create an AURORA instance and mix it with the right emitter to define PGA-AURORA\n",
     "- which functions must be defined before training\n",
     "- how to launch a certain number of training steps\n",
-    "- how to visualise the optimization process\n",
-    "- how to save/load a repertoire"
+    "- how to visualise the optimization process"
    ]
   },
   {
@@ -426,8 +425,6 @@
     "    model, subkey, (1, *observations_dims)\n",
     ")\n",
     "\n",
-    "print(jax.tree_map(lambda x: x.shape, model_params))\n",
-    "\n",
     "# Define the encoder function\n",
     "encoder_fn = jax.jit(\n",
     "    functools.partial(\n",
@@ -459,8 +456,6 @@
     "    model, subkey, (1, *observations_dims)\n",
     ")\n",
     "\n",
-    "print(jax.tree_map(lambda x: x.shape, model_params))\n",
-    "\n",
     "# define arbitrary observation's mean/std\n",
     "mean_observations = jnp.zeros(observations_dims[-1])\n",
     "std_observations = jnp.ones(observations_dims[-1])\n",

mypy.ini

@@ -18,6 +18,3 @@ strict_equality = True
 explicit_package_bases = True
 follow_imports = skip
 ignore_missing_imports = True
-
-[mypy-tensorflow_probability.*]
-ignore_missing_imports = True

qdax/core/containers/ga_repertoire.py

@@ -2,12 +2,11 @@
 
 from __future__ import annotations
 
-from typing import Callable, Optional, Tuple
+from typing import Optional, Tuple
 
 import flax
 import jax
 import jax.numpy as jnp
-from jax.flatten_util import ravel_pytree
 
 from qdax.core.containers.repertoire import Repertoire
 from qdax.core.emitters.repertoire_selectors.selector import GARepertoireT, Selector
@@ -45,46 +44,6 @@ def size(self) -> int:
         first_leaf = jax.tree.leaves(self.genotypes)[0]
         return int(first_leaf.shape[0])
 
-    def save(self, path: str = "./") -> None:
-        """Saves the repertoire.
-
-        Args:
-            path: place to store the files. Defaults to "./".
-        """
-
-        def flatten_genotype(genotype: Genotype) -> jnp.ndarray:
-            flatten_genotype, _ = ravel_pytree(genotype)
-            return flatten_genotype
-
-        # flatten all the genotypes
-        flat_genotypes = jax.vmap(flatten_genotype)(self.genotypes)
-
-        jnp.save(path + "genotypes.npy", flat_genotypes)
-        jnp.save(path + "scores.npy", self.fitnesses)
-
-    @classmethod
-    def load(cls, reconstruction_fn: Callable, path: str = "./") -> GARepertoire:
-        """Loads a GA Repertoire.
-
-        Args:
-            reconstruction_fn: Function to reconstruct a PyTree
-                from a flat array.
-            path: Path where the data is saved. Defaults to "./".
-
-        Returns:
-            A GA Repertoire.
-        """
-
-        flat_genotypes = jnp.load(path + "genotypes.npy")
-        genotypes = jax.vmap(reconstruction_fn)(flat_genotypes)
-
-        fitnesses = jnp.load(path + "fitnesses.npy")
-
-        return cls(
-            genotypes=genotypes,
-            fitnesses=fitnesses,
-        )
-
     def select(
         self,
         key: RNGKey,