Updating _build_networks_and_optimizer for full_rainbow agent.

PiperOrigin-RevId: 584053752
google · Nov 27, 2023 · 485ea99 · 485ea99
1 parent 8df4337
commit 485ea99
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/dopamine/jax/agents/full_rainbow/full_rainbow_agent.py b/dopamine/jax/agents/full_rainbow/full_rainbow_agent.py
@@ -259,7 +259,8 @@ def __init__(self,
 
   def _build_networks_and_optimizer(self):
     self._rng, rng = jax.random.split(self._rng)
-    self.online_params = self.network_def.init(rng, x=self.state,
+    state = self.preprocess_fn(self.state)
+    self.online_params = self.network_def.init(rng, x=state,
                                                support=self._support)
     self.optimizer = dqn_agent.create_optimizer(self._optimizer_name)
     self.optimizer_state = self.optimizer.init(self.online_params)