fix key management in examples using pmap

adaptive-intelligent-robotics · Sep 20, 2024 · 5cac247 · 5cac247
1 parent 2040636
commit 5cac247
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 0 deletions.
diff --git a/examples/me_sac_pbt.ipynb b/examples/me_sac_pbt.ipynb
@@ -311,6 +311,10 @@
     "    observation_size=env.observation_size,\n",
     "    buffer_size=buffer_size,\n",
     ")\n",
+    "\n",
+    "# Need to convert to PRNGKey because of github.com/jax-ml/jax/issues/23647\n",
+    "keys = jax.random.key_data(keys)\n",
+    "\n",
     "keys, training_states, _ = jax.pmap(agent_init_fn, axis_name=\"p\", devices=devices)(keys)"
    ]
   },

diff --git a/examples/me_td3_pbt.ipynb b/examples/me_td3_pbt.ipynb
@@ -314,6 +314,10 @@
     "    observation_size=env.observation_size,\n",
     "    buffer_size=buffer_size,\n",
     ")\n",
+    "\n",
+    "# Need to convert to PRNGKey because of github.com/jax-ml/jax/issues/23647\n",
+    "keys = jax.random.key_data(keys)\n",
+    "\n",
     "keys, training_states, _ = jax.pmap(agent_init_fn, axis_name=\"p\", devices=devices)(keys)"
    ]
   },

diff --git a/examples/sac_pbt.ipynb b/examples/sac_pbt.ipynb
@@ -269,6 +269,10 @@
     "    observation_size=env.observation_size,\n",
     "    buffer_size=buffer_size,\n",
     ")\n",
+    "\n",
+    "# Need to convert to PRNGKey because of github.com/jax-ml/jax/issues/23647\n",
+    "keys = jax.random.key_data(keys)\n",
+    "\n",
     "keys, training_states, replay_buffers = jax.pmap(\n",
     "    agent_init_fn, axis_name=\"p\", devices=devices\n",
     ")(keys)"

diff --git a/examples/td3_pbt.ipynb b/examples/td3_pbt.ipynb
@@ -232,6 +232,10 @@
     "    observation_size=env.observation_size,\n",
     "    buffer_size=buffer_size,\n",
     ")\n",
+    "\n",
+    "# Need to convert to PRNGKey because of github.com/jax-ml/jax/issues/23647\n",
+    "keys = jax.random.key_data(keys)\n",
+    "\n",
     "keys, training_states, replay_buffers = jax.pmap(\n",
     "    agent_init_fn, axis_name=\"p\", devices=devices\n",
     ")(keys)"