Merge branch 'new-style-jax-rng-keys' of github.com:miltonllera/QDax …

…into miltonllera-new-style-jax-rng-keys
adaptive-intelligent-robotics · Sep 13, 2024 · 7ba39e3 · 7ba39e3
2 parents 85b6ba6 + b1d0685
commit 7ba39e3
Show file tree

Hide file tree

Showing 60 changed files with 73 additions and 73 deletions.
diff --git a/README.md b/README.md
@@ -64,7 +64,7 @@ min_bd = 0.0
 max_bd = 1.0
 
 # Init a random key
-random_key = jax.random.PRNGKey(seed)
+random_key = jax.random.key(seed)
 
 # Init population of controllers
 random_key, subkey = jax.random.split(random_key)

diff --git a/examples/aurora.ipynb b/examples/aurora.ipynb
@@ -146,7 +146,7 @@
     "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
-    "random_key = jax.random.PRNGKey(seed)\n",
+    "random_key = jax.random.key(seed)\n",
     "\n",
     "# Init policy network\n",
     "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",

diff --git a/examples/cmaes.ipynb b/examples/cmaes.ipynb
@@ -178,7 +178,7 @@
    "outputs": [],
    "source": [
     "state = cmaes.init()\n",
-    "random_key = jax.random.PRNGKey(0)"
+    "random_key = jax.random.key(0)"
    ]
   },
   {

diff --git a/examples/cmame.ipynb b/examples/cmame.ipynb
@@ -205,7 +205,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "random_key = jax.random.PRNGKey(0)\n",
+    "random_key = jax.random.key(0)\n",
     "# in CMA-ME settings (from the paper), there is no init population\n",
     "# we multipy by zero to reproduce this setting\n",
     "initial_population = jax.random.uniform(random_key, shape=(batch_size, num_dimensions)) * 0.\n",

diff --git a/examples/cmamega.ipynb b/examples/cmamega.ipynb
@@ -198,7 +198,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "random_key = jax.random.PRNGKey(0)\n",
+    "random_key = jax.key(0)\n",
     "# no initial population - give all the same value as emitter init value\n",
     "initial_population = jax.random.uniform(random_key, shape=(batch_size, num_dimensions)) * 0.\n",
     "\n",

diff --git a/examples/dads.ipynb b/examples/dads.ipynb
@@ -163,7 +163,7 @@
     "    eval_metrics=True,\n",
     ")\n",
     "\n",
-    "key = jax.random.PRNGKey(seed)\n",
+    "key = jax.random.key(seed)\n",
     "env_state = jax.jit(env.reset)(rng=key)\n",
     "eval_env_first_state = jax.jit(eval_env.reset)(rng=key)\n",
     "\n",
@@ -499,7 +499,7 @@
    "outputs": [],
    "source": [
     "rollout = []\n",
-    "random_key = jax.random.PRNGKey(seed=1)\n",
+    "random_key = jax.random.key(seed=1)\n",
     "state = jit_env_reset(rng=random_key)\n",
     "while not state.done:\n",
     "    rollout.append(state)\n",

diff --git a/examples/diayn.ipynb b/examples/diayn.ipynb
@@ -162,7 +162,7 @@
     "    eval_metrics=True,\n",
     ")\n",
     "\n",
-    "key = jax.random.PRNGKey(seed)\n",
+    "key = jax.random.key(seed)\n",
     "env_state = jax.jit(env.reset)(rng=key)\n",
     "eval_env_first_state = jax.jit(eval_env.reset)(rng=key)\n",
     "\n",
@@ -490,7 +490,7 @@
    "outputs": [],
    "source": [
     "rollout = []\n",
-    "random_key = jax.random.PRNGKey(seed=1)\n",
+    "random_key = jax.random.key(seed=1)\n",
     "state = jit_env_reset(rng=random_key)\n",
     "while not state.done:\n",
     "    rollout.append(state)\n",

diff --git a/examples/distributed_mapelites.ipynb b/examples/distributed_mapelites.ipynb
@@ -176,7 +176,7 @@
     "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
-    "random_key = jax.random.PRNGKey(seed)\n",
+    "random_key = jax.random.key(seed)\n",
     "\n",
     "# Init policy network\n",
     "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",

diff --git a/examples/jumanji_snake.ipynb b/examples/jumanji_snake.ipynb
@@ -113,7 +113,7 @@
     "env = jumanji.make('Snake-v1')\n",
     "\n",
     "# Reset your (jit-able) environment\n",
-    "key = jax.random.PRNGKey(0)\n",
+    "key = jax.random.key(0)\n",
     "state, timestep = jax.jit(env.reset)(key)\n",
     "\n",
     "# Interact with the (jit-able) environment\n",
@@ -137,7 +137,7 @@
    "outputs": [],
    "source": [
     "# Init a random key\n",
-    "random_key = jax.random.PRNGKey(seed)\n",
+    "random_key = jax.random.key(seed)\n",
     "\n",
     "# get number of actions\n",
     "num_actions = env.action_spec().maximum + 1\n",

diff --git a/examples/mapelites.ipynb b/examples/mapelites.ipynb
@@ -133,7 +133,7 @@
     "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
-    "random_key = jax.random.PRNGKey(seed)\n",
+    "random_key = jax.Key(seed)\n",
     "\n",
     "# Init policy network\n",
     "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",
@@ -494,7 +494,7 @@
    "outputs": [],
    "source": [
     "rollout = []\n",
-    "rng = jax.random.PRNGKey(seed=1)\n",
+    "rng = jax.random.key(seed=1)\n",
     "state = jit_env_reset(rng=rng)\n",
     "while not state.done:\n",
     "    rollout.append(state)\n",

diff --git a/examples/me_sac_pbt.ipynb b/examples/me_sac_pbt.ipynb
@@ -150,7 +150,7 @@
    "outputs": [],
    "source": [
     "%%time\n",
-    "key = jax.random.PRNGKey(seed)\n",
+    "key = jax.random.key(seed)\n",
     "key, subkey = jax.random.split(key)\n",
     "env_states = jax.jit(env.reset)(rng=subkey)\n",
     "eval_env_first_states = jax.jit(eval_env.reset)(rng=subkey)"
@@ -504,7 +504,7 @@
     "%%time\n",
     "rollout = []\n",
     "\n",
-    "rng = jax.random.PRNGKey(seed=1)\n",
+    "rng = jax.random.key(seed=1)\n",
     "env_state = jax.jit(env.reset)(rng=rng)\n",
     "\n",
     "training_state, env_state = jax.tree_map(\n",

diff --git a/examples/me_td3_pbt.ipynb b/examples/me_td3_pbt.ipynb
@@ -155,7 +155,7 @@
    "outputs": [],
    "source": [
     "%%time\n",
-    "key = jax.random.PRNGKey(seed)\n",
+    "key = jax.random.key(seed)\n",
     "key, subkey = jax.random.split(key)\n",
     "env_states = jax.jit(env.reset)(rng=subkey)\n",
     "eval_env_first_states = jax.jit(eval_env.reset)(rng=subkey)"

diff --git a/examples/mees.ipynb b/examples/mees.ipynb
@@ -151,7 +151,7 @@
     "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
-    "random_key = jax.random.PRNGKey(seed)\n",
+    "random_key = jax.random.key(seed)\n",
     "\n",
     "# Init policy network\n",
     "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",

diff --git a/examples/mels.ipynb b/examples/mels.ipynb
@@ -140,7 +140,7 @@
     "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
-    "random_key = jax.random.PRNGKey(seed)\n",
+    "random_key = jax.random.key(seed)\n",
     "\n",
     "# Init policy network\n",
     "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",
@@ -509,7 +509,7 @@
    "outputs": [],
    "source": [
     "rollout = []\n",
-    "rng = jax.random.PRNGKey(seed=1)\n",
+    "rng = jax.random.key(seed=1)\n",
     "state = jit_env_reset(rng=rng)\n",
     "while not state.done:\n",
     "    rollout.append(state)\n",

diff --git a/examples/mome.ipynb b/examples/mome.ipynb
@@ -205,7 +205,7 @@
    "outputs": [],
    "source": [
     "# initial population\n",
-    "random_key = jax.random.PRNGKey(42)\n",
+    "random_key = jax.random.key(42)\n",
     "random_key, subkey = jax.random.split(random_key)\n",
     "genotypes = jax.random.uniform(\n",
     "    random_key, (batch_size, num_variables), minval=minval, maxval=maxval, dtype=jnp.float32\n",

diff --git a/examples/nsga2_spea2.ipynb b/examples/nsga2_spea2.ipynb
@@ -181,7 +181,7 @@
    "outputs": [],
    "source": [
     "# Initial population\n",
-    "random_key = jax.random.PRNGKey(0)\n",
+    "random_key = jax.random.key(0)\n",
     "random_key, subkey = jax.random.split(random_key)\n",
     "genotypes = jax.random.uniform(\n",
     "    subkey, (batch_size, genotype_dim), minval=minval, maxval=maxval, dtype=jnp.float32\n",

diff --git a/examples/omgmega.ipynb b/examples/omgmega.ipynb
@@ -184,7 +184,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "random_key = jax.random.PRNGKey(0)\n",
+    "random_key = jax.random.key(0)\n",
     "\n",
     "# defines the population\n",
     "random_key, subkey = jax.random.split(random_key)\n",

diff --git a/examples/pga_aurora.ipynb b/examples/pga_aurora.ipynb
@@ -164,7 +164,7 @@
     "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
-    "random_key = jax.random.PRNGKey(seed)\n",
+    "random_key = jax.random.key(seed)\n",
     "\n",
     "# Init policy network\n",
     "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",

diff --git a/examples/pgame.ipynb b/examples/pgame.ipynb
@@ -143,7 +143,7 @@
     "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
-    "random_key = jax.random.PRNGKey(seed)\n",
+    "random_key = jax.random.key(seed)\n",
     "\n",
     "# Init policy network\n",
     "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",

diff --git a/examples/qdpg.ipynb b/examples/qdpg.ipynb
@@ -157,7 +157,7 @@
     "env = environments.create(env_name, episode_length=episode_length)\n",
     "\n",
     "# Init a random key\n",
-    "random_key = jax.random.PRNGKey(seed)\n",
+    "random_key = jax.random.key(seed)\n",
     "\n",
     "# Init policy network\n",
     "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",

diff --git a/examples/sac_pbt.ipynb b/examples/sac_pbt.ipynb
@@ -210,7 +210,7 @@
    "outputs": [],
    "source": [
     "# %%time\n",
-    "key = jax.random.PRNGKey(seed)\n",
+    "key = jax.random.key(seed)\n",
     "key, *keys = jax.random.split(key, num=1 + num_devices)\n",
     "keys = jnp.stack(keys)\n",
     "env_states, eval_env_first_states = jax.pmap(\n",
@@ -518,7 +518,7 @@
     "%%time\n",
     "rollout = []\n",
     "\n",
-    "rng = jax.random.PRNGKey(seed=1)\n",
+    "rng = jax.random.key(seed=1)\n",
     "env_state = jax.jit(env.reset)(rng=rng)\n",
     "\n",
     "training_state, env_state = jax.tree_util.tree_map(\n",

diff --git a/examples/scripts/me_example.py b/examples/scripts/me_example.py
@@ -26,7 +26,7 @@ def run_me() -> None:
     max_bd = 1.0
 
     # Init a random key
-    random_key = jax.random.PRNGKey(seed)
+    random_key = jax.random.key(seed)
 
     # Init population of controllers
     random_key, subkey = jax.random.split(random_key)

diff --git a/examples/smerl.ipynb b/examples/smerl.ipynb
@@ -168,7 +168,7 @@
     "    eval_metrics=True,\n",
     ")\n",
     "\n",
-    "key = jax.random.PRNGKey(seed)\n",
+    "key = jax.Key(seed)\n",
     "env_state = jax.jit(env.reset)(rng=key)\n",
     "eval_env_first_state = jax.jit(eval_env.reset)(rng=key)\n",
     "\n",
@@ -504,7 +504,7 @@
    "outputs": [],
    "source": [
     "rollout = []\n",
-    "random_key = jax.random.PRNGKey(seed=1)\n",
+    "random_key = jax.random.key(seed=1)\n",
     "state = jit_env_reset(rng=random_key)\n",
     "while not state.done:\n",
     "    rollout.append(state)\n",

diff --git a/examples/td3_pbt.ipynb b/examples/td3_pbt.ipynb
@@ -181,7 +181,7 @@
    "outputs": [],
    "source": [
     "%%time\n",
-    "key = jax.random.PRNGKey(seed)\n",
+    "key = jax.random.key(seed)\n",
     "key, *keys = jax.random.split(key, num=1 + num_devices)\n",
     "keys = jnp.stack(keys)\n",
     "env_states, eval_env_first_states = jax.pmap(\n",

diff --git a/qdax/core/containers/mapelites_repertoire.py b/qdax/core/containers/mapelites_repertoire.py
@@ -63,7 +63,7 @@ def compute_cvt_centroids(
         init="k-means++",
         n_clusters=num_centroids,
         n_init=1,
-        random_state=RandomState(subkey),
+        random_state=RandomState(jax.random.key_data(subkey)),
     )
     k_means.fit(x)
     centroids = k_means.cluster_centers_

diff --git a/qdax/core/neuroevolution/networks/seq2seq_networks.py b/qdax/core/neuroevolution/networks/seq2seq_networks.py
@@ -16,7 +16,7 @@
 from flax import linen as nn
 
 Array = Any
-PRNGKey = Any
+PRNGKey = jax.Array
 
 
 class EncoderLSTM(nn.Module):
@@ -52,7 +52,7 @@ def select_carried_state(new_state: Array, old_state: Array) -> Array:
     def initialize_carry(batch_size: int, hidden_size: int) -> Tuple[Array, Array]:
         # Use a dummy key since the default state init fn is just zeros.
         return nn.LSTMCell(hidden_size, parent=None).initialize_carry(  # type: ignore
-            jax.random.PRNGKey(0), (batch_size, hidden_size)
+            jax.random.key(0), (batch_size, hidden_size)
         )
 
 

diff --git a/qdax/tasks/README.md b/qdax/tasks/README.md
@@ -19,7 +19,7 @@ Notes:
 import jax
 from qdax.tasks.arm import arm_scoring_function
 
-random_key = jax.random.PRNGKey(0)
+random_key = jax.random.key(0)
 
 # Get scoring function
 scoring_fn = arm_scoring_function
@@ -56,7 +56,7 @@ desc_size = 2
 import jax
 from qdax.tasks.standard_functions import sphere_scoring_function
 
-random_key = jax.random.PRNGKey(0)
+random_key = jax.random.key(0)
 
 # Get scoring function
 scoring_fn = sphere_scoring_function
@@ -98,7 +98,7 @@ desc_size = 2
 import jax
 from qdax.tasks.hypervolume_functions import square_scoring_function
 
-random_key = jax.random.PRNGKey(0)
+random_key = jax.random.key(0)
 
 # Get scoring function
 scoring_fn = square_scoring_function

diff --git a/tests/baselines_test/cmame_test.py b/tests/baselines_test/cmame_test.py
@@ -81,7 +81,7 @@ def metrics_fn(repertoire: MapElitesRepertoire) -> Dict[str, jnp.ndarray]:
         max_fitness = jnp.max(adjusted_fitness)
         return {"qd_score": qd_score, "max_fitness": max_fitness, "coverage": coverage}
 
-    random_key = jax.random.PRNGKey(0)
+    random_key = jax.random.key(0)
     initial_population = (
         jax.random.uniform(random_key, shape=(batch_size, num_dimensions)) * 0.0
     )

diff --git a/tests/baselines_test/cmamega_test.py b/tests/baselines_test/cmamega_test.py
@@ -95,7 +95,7 @@ def metrics_fn(repertoire: MapElitesRepertoire) -> Dict[str, jnp.ndarray]:
         max_fitness = jnp.max(adjusted_fitness)
         return {"qd_score": qd_score, "max_fitness": max_fitness, "coverage": coverage}
 
-    random_key = jax.random.PRNGKey(0)
+    random_key = jax.random.key(0)
     initial_population = jax.random.uniform(
         random_key, shape=(batch_size, num_dimensions)
     )

diff --git a/tests/baselines_test/dads_smerl_test.py b/tests/baselines_test/dads_smerl_test.py
@@ -72,7 +72,7 @@ def test_dads_smerl() -> None:
         eval_metrics=True,
     )
 
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     env_state = jax.jit(env.reset)(rng=key)
     eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
 

diff --git a/tests/baselines_test/dads_test.py b/tests/baselines_test/dads_test.py
@@ -66,7 +66,7 @@ def test_dads() -> None:
         eval_metrics=True,
     )
 
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     env_state = jax.jit(env.reset)(rng=key)
     eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
 

diff --git a/tests/baselines_test/diayn_smerl_test.py b/tests/baselines_test/diayn_smerl_test.py
@@ -69,7 +69,7 @@ def test_diayn_smerl() -> None:
         eval_metrics=True,
     )
 
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     env_state = jax.jit(env.reset)(rng=key)
     eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
 

diff --git a/tests/baselines_test/diayn_test.py b/tests/baselines_test/diayn_test.py
@@ -62,7 +62,7 @@ def test_diayn() -> None:
         eval_metrics=True,
     )
 
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     env_state = jax.jit(env.reset)(rng=key)
     eval_env_first_state = jax.jit(eval_env.reset)(rng=key)