Repository: adaptive-intelligent-robotics/QDax (language: python)

Issue:
Bad evaluation of the behavior descriptor in MAP-Elites and PGAME
We observed strange results on several environments. It seems that the behavior descriptors are not properly evaluated in some cases.
More information in the related PR.


Discussion / hints:

Fix patch:
diff --git a/notebooks/mapelites_example.ipynb b/notebooks/mapelites_example.ipynb
index d7c8d5ec..98eb1dee 100644
--- a/notebooks/mapelites_example.ipynb
+++ b/notebooks/mapelites_example.ipynb
@@ -108,7 +108,7 @@
    "outputs": [],
    "source": [
     "# Init environment\n",
-    "env = environments.create(env_name)\n",
+    "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
     "random_key = jax.random.PRNGKey(seed)\n",
@@ -161,6 +161,8 @@
     "    \"\"\"\n",
     "\n",
     "    actions = policy_network.apply(policy_params, env_state.obs)\n",
+    "    \n",
+    "    state_desc = env_state.info[\"state_descriptor\"]\n",
     "    next_state = env.step(env_state, actions)\n",
     "\n",
     "    transition = QDTransition(\n",
@@ -170,7 +172,7 @@
     "        dones=next_state.done,\n",
     "        actions=actions,\n",
     "        truncations=next_state.info[\"truncation\"],\n",
-    "        state_desc=env_state.info[\"state_descriptor\"],\n",
+    "        state_desc=state_desc,\n",
     "        next_state_desc=next_state.info[\"state_descriptor\"],\n",
     "    )\n",
     "\n",
diff --git a/notebooks/pgame_example.ipynb b/notebooks/pgame_example.ipynb
index ddd316f3..e386e4e8 100644
--- a/notebooks/pgame_example.ipynb
+++ b/notebooks/pgame_example.ipynb
@@ -119,7 +119,7 @@
    "outputs": [],
    "source": [
     "# Init environment\n",
-    "env = environments.create(env_name)\n",
+    "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
     "random_key = jax.random.PRNGKey(seed)\n",
@@ -169,6 +169,8 @@
     "    \"\"\"\n",
     "\n",
     "    actions = policy_network.apply(policy_params, env_state.obs)\n",
+    "    \n",
+    "    state_desc = env_state.info[\"state_descriptor\"]\n",
     "    next_state = env.step(env_state, actions)\n",
     "\n",
     "    transition = QDTransition(\n",
@@ -178,7 +180,7 @@
     "        dones=next_state.done,\n",
     "        actions=actions,\n",
     "        truncations=next_state.info[\"truncation\"],\n",
-    "        state_desc=env_state.info[\"state_descriptor\"],\n",
+    "        state_desc=state_desc,\n",
     "        next_state_desc=next_state.info[\"state_descriptor\"],\n",
     "    )\n",
     "\n",


Test patch:
diff --git a/tests/core_test/map_elites_test.py b/tests/core_test/map_elites_test.py
index f0e9de4c..77eba0f5 100644
--- a/tests/core_test/map_elites_test.py
+++ b/tests/core_test/map_elites_test.py
@@ -29,7 +29,7 @@ def test_map_elites() -> None:
     max_bd = 1.0
 
     # Init environment
-    env = environments.create(env_name)
+    env = environments.create(env_name, episode_length=episode_length)
 
     # Init a random key
     random_key = jax.random.PRNGKey(seed)
@@ -65,6 +65,8 @@ def play_step_fn(
         """
 
         actions = policy_network.apply(policy_params, env_state.obs)
+
+        state_desc = env_state.info["state_descriptor"]
         next_state = env.step(env_state, actions)
 
         transition = QDTransition(
@@ -74,7 +76,7 @@ def play_step_fn(
             dones=next_state.done,
             actions=actions,
             truncations=next_state.info["truncation"],
-            state_desc=env_state.info["state_descriptor"],
+            state_desc=state_desc,
             next_state_desc=next_state.info["state_descriptor"],
         )
 
diff --git a/tests/core_test/pgame_test.py b/tests/core_test/pgame_test.py
index dfadcf49..809bdc8f 100644
--- a/tests/core_test/pgame_test.py
+++ b/tests/core_test/pgame_test.py
@@ -47,7 +47,7 @@ def test_pgame_elites() -> None:
     num_pg_training_steps = 5
 
     # Init environment
-    env = environments.create(env_name)
+    env = environments.create(env_name, episode_length=episode_length)
 
     # Init a random key
     random_key = jax.random.PRNGKey(seed)
@@ -77,6 +77,8 @@ def play_step_fn(
         """
 
         actions = policy_network.apply(policy_params, env_state.obs)
+
+        state_desc = env_state.info["state_descriptor"]
         next_state = env.step(env_state, actions)
 
         transition = QDTransition(
@@ -86,7 +88,7 @@ def play_step_fn(
             dones=next_state.done,
             actions=actions,
             truncations=next_state.info["truncation"],
-            state_desc=env_state.info["state_descriptor"],
+            state_desc=state_desc,
             next_state_desc=next_state.info["state_descriptor"],
         )