feat: add Brax v2 support and restructure brax tasks organization

* New Directory Structure: Created a new /tasks/brax directory with clear versioning:

* /tasks/brax/v1/: Contains code compatible with Brax v1 API

* /tasks/brax/v2/: Contains code compatible with Brax v2 API

* Relocated environment-related code from /environments to /tasks/brax/v1/envs and /tasks/brax/v2/envs

* Relocated wrapper code to /tasks/brax/v1/wrappers and /tasks/brax/v2/wrappers

* Moved descriptor extractors to /tasks/brax/descriptor_extractors

* Added base environment classes in both v1 and v2 paths

* Fixed all import statements throughout the codebase to reference the new structure

* Updated all relevant test files to use the new imports

* Moved AuroraExtraInfo: Relocated AuroraExtraInfo and AuroraExtraInfoNormalization classes from descriptor_extractors to custom_types.py

* Stop using task-specific types (e.g. brax.env.State JumanjiState) in the main custom types file

Author: Lookatator

.gitignore

@@ -28,6 +28,9 @@ hydra_outputs/
 
 # Usual python .gitignore
 
+# Mac OS X files
+.DS_Store
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -166,3 +169,6 @@ dmypy.json
 
 # Cython debug symbols
 cython_debug/
+
+# HTML files
+*.html

README.md

@@ -111,14 +111,16 @@ centroids = compute_euclidean_centroids(
 )
 
 # Initializes repertoire and emitter state
-repertoire, emitter_state, key = map_elites.init(init_variables, centroids, key)
+key, subkey = jax.random.split(key)
+repertoire, emitter_state, metrics = map_elites.init(init_variables, centroids, subkey)
 
 # Run MAP-Elites loop
 for i in range(num_iterations):
-    (repertoire, emitter_state, metrics, key,) = map_elites.update(
+    key, subkey = jax.random.split(key)
+    (repertoire, emitter_state, metrics,) = map_elites.update(
         repertoire,
         emitter_state,
-        key,
+        subkey,
     )
 
 # Get contents of repertoire

docs/api_documentation/environments.md

@@ -65,21 +65,20 @@
     "\n",
     "from qdax.core.aurora import AURORA\n",
     "from qdax.core.containers.unstructured_repertoire import UnstructuredRepertoire\n",
-    "from qdax import environments\n",
-    "from qdax.tasks.brax_envs import (\n",
+    "import qdax.tasks.brax.v1 as environments\n",
+    "from qdax.tasks.brax.v1.env_creators import (\n",
     "    create_default_brax_task_components,\n",
     "    get_aurora_scoring_fn,\n",
     ")\n",
-    "from qdax.environments.descriptor_extractors import (\n",
-    "    AuroraExtraInfoNormalization,\n",
+    "from qdax.tasks.brax.descriptor_extractors import (\n",
     "    get_aurora_encoding,\n",
     ")\n",
     "from qdax.core.neuroevolution.buffers.buffer import QDTransition\n",
     "from qdax.core.neuroevolution.networks.networks import MLP\n",
     "from qdax.core.emitters.mutation_operators import isoline_variation\n",
     "from qdax.core.emitters.standard_emitters import MixingEmitter\n",
     "\n",
-    "from qdax.custom_types import Observation\n",
+    "from qdax.custom_types import AuroraExtraInfoNormalization, Observation\n",
     "from qdax.utils import train_seq2seq\n",
     "\n",
     "\n",

examples/aurora.ipynb

@@ -59,7 +59,7 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
-    "from qdax import environments\n",
+    "import qdax.tasks.brax.v1 as environments\n",
     "from qdax.baselines.dads import DADS, DadsConfig, DadsTrainingState\n",
     "from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer\n",
     "from qdax.core.neuroevolution.sac_td3_utils import do_iteration_fn, warmstart_buffer\n",

examples/dads.ipynb

@@ -64,17 +64,17 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
-    "from qdax import environments\n",
+    "import qdax.tasks.brax.v1 as environments\n",
     "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids\n",
     "from qdax.core.emitters.dcrl_me_emitter import DCRLMEConfig, DCRLMEEmitter\n",
     "from qdax.core.emitters.mutation_operators import isoline_variation\n",
     "from qdax.core.map_elites import MAPElites\n",
     "from qdax.core.neuroevolution.buffers.buffer import DCRLTransition\n",
     "from qdax.core.neuroevolution.networks.networks import MLP, MLPDC\n",
     "from qdax.custom_types import EnvState, Params, RNGKey\n",
-    "from qdax.environments import descriptor_extractor\n",
-    "from qdax.environments.wrappers import OffsetRewardWrapper, ClipRewardWrapper\n",
-    "from qdax.tasks.brax_envs import scoring_function_brax_envs\n",
+    "from qdax.tasks.brax.v1 import descriptor_extractor\n",
+    "from qdax.tasks.brax.v1.wrappers.reward_wrappers import OffsetRewardWrapper, ClipRewardWrapper\n",
+    "from qdax.tasks.brax.v1.env_creators import scoring_function_brax_envs\n",
     "from qdax.utils.plotting import plot_map_elites_results\n",
     "\n",
     "from qdax.utils.metrics import CSVLogger, default_qd_metrics\n",
@@ -101,7 +101,7 @@
     "min_descriptor = -30.0 #@param {type:\"number\"}\n",
     "max_descriptor = 30.0 #@param {type:\"number\"}\n",
     "\n",
-    "num_iterations = 1000 #@param {type:\"integer\"}\n",
+    "num_iterations = 200 #@param {type:\"integer\"}\n",
     "batch_size = 256 #@param {type:\"integer\"}\n",
     "\n",
     "# Archive\n",

examples/dcrlme.ipynb

@@ -59,7 +59,7 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
-    "from qdax import environments\n",
+    "import qdax.tasks.brax.v1 as environments\n",
     "from qdax.baselines.diayn import DIAYN, DiaynConfig, DiaynTrainingState\n",
     "from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer\n",
     "from qdax.core.neuroevolution.sac_td3_utils import do_iteration_fn, warmstart_buffer\n",

examples/diayn.ipynb

@@ -70,8 +70,8 @@
     "\n",
     "from qdax.core.distributed_map_elites import DistributedMAPElites\n",
     "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids\n",
-    "from qdax import environments\n",
-    "from qdax.tasks.brax_envs import scoring_function_brax_envs as scoring_function\n",
+    "import qdax.tasks.brax.v1 as environments\n",
+    "from qdax.tasks.brax.v1.env_creators import scoring_function_brax_envs as scoring_function\n",
     "from qdax.core.neuroevolution.buffers.buffer import QDTransition\n",
     "from qdax.core.neuroevolution.networks.networks import MLP\n",
     "from qdax.core.emitters.mutation_operators import isoline_variation\n",

examples/distributed_mapelites.ipynb

@@ -65,8 +65,8 @@
     "\n",
     "from qdax.core.map_elites import MAPElites\n",
     "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids, MapElitesRepertoire\n",
-    "from qdax import environments\n",
-    "from qdax.tasks.brax_envs import scoring_function_brax_envs as scoring_function\n",
+    "import qdax.tasks.brax.v1 as environments\n",
+    "from qdax.tasks.brax.v1.env_creators import scoring_function_brax_envs as scoring_function\n",
     "from qdax.core.neuroevolution.buffers.buffer import QDTransition\n",
     "from qdax.core.neuroevolution.networks.networks import MLP\n",
     "from qdax.core.emitters.mutation_operators import isoline_variation\n",
@@ -99,7 +99,7 @@
     "batch_size = 100 #@param {type:\"number\"}\n",
     "env_name = 'walker2d_uni'#@param['ant_uni', 'hopper_uni', 'walker2d_uni', 'halfcheetah_uni', 'humanoid_uni', 'ant_omni', 'humanoid_omni']\n",
     "episode_length = 100 #@param {type:\"integer\"}\n",
-    "num_iterations = 1000 #@param {type:\"integer\"}\n",
+    "num_iterations = 500 #@param {type:\"integer\"}\n",
     "seed = 42 #@param {type:\"integer\"}\n",
     "policy_hidden_layer_sizes = (64, 64) #@param {type:\"raw\"}\n",
     "iso_sigma = 0.005 #@param {type:\"number\"}\n",

examples/mapelites.ipynb

@@ -64,9 +64,9 @@
     "import jax.numpy as jnp\n",
     "\n",
     "from qdax.core.map_elites import MAPElites\n",
-    "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids, MapElitesRepertoire\n",
-    "from qdax import environments\n",
-    "from qdax.tasks.brax_envs import scoring_function_brax_envs as scoring_function\n",
+    "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids\n",
+    "import qdax.tasks.brax.v1 as environments\n",
+    "from qdax.tasks.brax.v1.env_creators import scoring_function_brax_envs as scoring_function\n",
     "from qdax.core.neuroevolution.buffers.buffer import QDTransition\n",
     "from qdax.core.neuroevolution.networks.networks import MLP\n",
     "from qdax.core.emitters.mutation_operators import isoline_variation\n",
@@ -75,11 +75,6 @@
     "\n",
     "from qdax.utils.metrics import CSVLogger, default_qd_metrics\n",
     "\n",
-    "from jax.flatten_util import ravel_pytree\n",
-    "\n",
-    "from IPython.display import HTML\n",
-    "from brax.v1.io import html\n",
-    "\n",
     "\n",
     "if \"COLAB_TPU_ADDR\" in os.environ:\n",
     "  from jax.tools import colab_tpu\n",