Merge pull request #24 from rk1a/fix-jumanji-spec-properties

Fix Jumanji environment spec access in examples
instadeepai · Apr 16, 2024 · 40a64e5 · 40a64e5
2 parents 2b8354d + 4ca6918
commit 40a64e5
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/examples/anakin_dqn_example.ipynb b/examples/anakin_dqn_example.ipynb
@@ -310,7 +310,7 @@
     "):\n",
     "    \"\"\"Sets up the experiment.\"\"\"\n",
     "    cores_count = len(jax.devices())  # get available TPU cores.\n",
-    "    network = get_network_fn(env.action_spec().num_values)  # define network.\n",
+    "    network = get_network_fn(env.action_spec.num_values)  # define network.\n",
     "    optim = optax.adam(step_size)  # define optimiser.\n",
     "\n",
     "    rng, rng_e, rng_p = random.split(random.PRNGKey(seed), num=3)  # prng keys.\n",

diff --git a/examples/anakin_ppo_example.ipynb b/examples/anakin_ppo_example.ipynb
@@ -253,7 +253,7 @@
     "    \"\"\"Sets up the experiment and returns the necessary information.\"\"\"\n",
     "\n",
     "    cores_count = len(jax.devices())  # get available TPU cores.\n",
-    "    network = get_network_fn(env.action_spec().num_values)  # define network.\n",
+    "    network = get_network_fn(env.action_spec.num_values)  # define network.\n",
     "    optim = optax.adam(step_size)  # define optimiser.\n",
     "\n",
     "    rng, rng_e, rng_p = random.split(random.PRNGKey(seed), num=3)  # prng keys.\n",

diff --git a/examples/anakin_prioritised_dqn_example.ipynb b/examples/anakin_prioritised_dqn_example.ipynb
@@ -333,7 +333,7 @@
     "):\n",
     "    \"\"\"Sets up the experiment.\"\"\"\n",
     "    cores_count = len(jax.devices())  # get available TPU cores.\n",
-    "    network = get_network_fn(env.action_spec().num_values)  # define network.\n",
+    "    network = get_network_fn(env.action_spec.num_values)  # define network.\n",
     "    optim = optax.adam(step_size)  # define optimiser.\n",
     "\n",
     "    rng, rng_e, rng_p = random.split(random.PRNGKey(seed), num=3)  # prng keys.\n",