Fix input-processing when preprocess_fn is explicitly passed.

PiperOrigin-RevId: 408423178
SuReLI · Nov 9, 2021 · 95308e7 · 95308e7
1 parent ff80e75
commit 95308e7
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/dopamine/jax/agents/implicit_quantile/implicit_quantile_agent.py b/dopamine/jax/agents/implicit_quantile/implicit_quantile_agent.py
@@ -386,7 +386,7 @@ def step(self, reward, observation):
 
     self._rng, self.action = select_action(self.network_def,
                                            self.online_params,
-                                           self.state,
+                                           self.preprocess_fn(self.state),
                                            self._rng,
                                            self.num_quantile_samples,
                                            self.num_actions,
@@ -419,9 +419,9 @@ def _train_step(self):
             self.target_network_params,
             self.optimizer,
             self.optimizer_state,
-            self.replay_elements['state'],
+            self.preprocess_fn(self.replay_elements['state']),
             self.replay_elements['action'],
-            self.replay_elements['next_state'],
+            self.preprocess_fn(self.replay_elements['next_state']),
             self.replay_elements['reward'],
             self.replay_elements['terminal'],
             self.num_tau_samples,

diff --git a/dopamine/jax/agents/quantile/quantile_agent.py b/dopamine/jax/agents/quantile/quantile_agent.py
@@ -240,9 +240,9 @@ def _train_step(self):
             self.target_network_params,
             self.optimizer,
             self.optimizer_state,
-            self.replay_elements['state'],
+            self.preprocess_fn(self.replay_elements['state']),
             self.replay_elements['action'],
-            self.replay_elements['next_state'],
+            self.preprocess_fn(self.replay_elements['next_state']),
             self.replay_elements['reward'],
             self.replay_elements['terminal'],
             self._kappa,

diff --git a/dopamine/jax/agents/rainbow/rainbow_agent.py b/dopamine/jax/agents/rainbow/rainbow_agent.py
@@ -352,7 +352,7 @@ def step(self, reward, observation):
 
     self._rng, self.action = select_action(self.network_def,
                                            self.online_params,
-                                           self.state,
+                                           self.preprocess_fn(self.state),
                                            self._rng,
                                            self.num_actions,
                                            self.eval_mode,
@@ -398,9 +398,9 @@ def _train_step(self):
             self.target_network_params,
             self.optimizer,
             self.optimizer_state,
-            self.replay_elements['state'],
+            self.preprocess_fn(self.replay_elements['state']),
             self.replay_elements['action'],
-            self.replay_elements['next_state'],
+            self.preprocess_fn(self.replay_elements['next_state']),
             self.replay_elements['reward'],
             self.replay_elements['terminal'],
             loss_weights,