Port boot_dqn from optix to optax

PiperOrigin-RevId: 330726981 Change-Id: Ib7fa620772c45cc4746627b271acc2dbb5455283
google-deepmind · Sep 9, 2020 · f2e6c21 · f2e6c21
1 parent 828f9bf
commit f2e6c21
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/bsuite/baselines/jax/boot_dqn/run.py b/bsuite/baselines/jax/boot_dqn/run.py
@@ -28,8 +28,8 @@
 
 import haiku as hk
 from jax import lax
-from jax.experimental import optix
 import jax.numpy as jnp
+import optax
 
 # Internal imports.
 
@@ -70,7 +70,7 @@ def network(inputs: jnp.ndarray) -> jnp.ndarray:
     x = hk.Flatten()(inputs)
     return net(x) + prior_scale * lax.stop_gradient(prior_net(x))
 
-  optimizer = optix.adam(learning_rate=1e-3)
+  optimizer = optax.adam(learning_rate=1e-3)
 
   agent = boot_dqn.BootstrappedDqn(
       obs_spec=env.observation_spec(),