Port boot_dqn from optix to optax

PiperOrigin-RevId: 330726981
Change-Id: Ib7fa620772c45cc4746627b271acc2dbb5455283

Author: mtthss

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(),