fixing bugs from new data structure of RNG keys

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

Hide file tree

Showing 5 changed files with 14 additions and 2 deletions.
diff --git a/qdax/utils/sampling.py b/qdax/utils/sampling.py
@@ -159,8 +159,8 @@ def multi_sample_scoring_function(
         # vectorizing over axis 0 vectorizes over the num_samples random keys
         in_axes=(None, 0),
         # indicates that the vectorized axis will become axis 1, i.e., the final
-        # output is shape (batch_size, num_samples, ...)
-        out_axes=1,
+        # output is shape (batch_size, num_samples, ...) except for the random key
+        out_axes=(1, 1, 1, 0),
     )
     all_fitnesses, all_descriptors, all_extra_scores, _ = sample_scoring_fn(
         policies_params, keys

diff --git a/tests/baselines_test/me_pbt_sac_test.py b/tests/baselines_test/me_pbt_sac_test.py
@@ -163,6 +163,9 @@ def scoring_function(genotypes, random_key):  # type: ignore
         observation_size=env.observation_size,
         buffer_size=buffer_size,
     )
+
+    # Need to convert to PRNGKey because of github.com/jax-ml/jax/issues/23647
+    keys = jax.random.key_data(keys)
     keys, training_states, _ = jax.pmap(agent_init_fn, axis_name="p", devices=devices)(
         keys
     )

diff --git a/tests/baselines_test/me_pbt_td3_test.py b/tests/baselines_test/me_pbt_td3_test.py
@@ -161,6 +161,9 @@ def scoring_function(genotypes, random_key):  # type: ignore
         observation_size=env.observation_size,
         buffer_size=buffer_size,
     )
+
+    # Need to convert to PRNGKey because of github.com/jax-ml/jax/issues/23647
+    keys = jax.random.key_data(keys)
     keys, training_states, _ = jax.pmap(agent_init_fn, axis_name="p", devices=devices)(
         keys
     )

diff --git a/tests/baselines_test/pbt_sac_test.py b/tests/baselines_test/pbt_sac_test.py
@@ -104,6 +104,9 @@ def init_environments(random_key):  # type: ignore
         observation_size=env.observation_size,
         buffer_size=buffer_size,
     )
+
+    # Need to convert to PRNGKey because of github.com/jax-ml/jax/issues/23647
+    keys = jax.random.key_data(keys)
     keys, training_states, replay_buffers = jax.pmap(
         agent_init_fn, axis_name="p", devices=devices
     )(keys)

diff --git a/tests/baselines_test/pbt_td3_test.py b/tests/baselines_test/pbt_td3_test.py
@@ -100,6 +100,9 @@ def init_environments(random_key):  # type: ignore
         observation_size=env.observation_size,
         buffer_size=buffer_size,
     )
+
+    # Need to convert to PRNGKey because of github.com/jax-ml/jax/issues/23647
+    keys = jax.random.key_data(keys)
     keys, training_states, replay_buffers = jax.pmap(
         agent_init_fn, axis_name="p", devices=devices
     )(keys)