Propagating observation dtype to the replay buffer in RainbowAgent si…

…milarly to DQNAgent. PiperOrigin-RevId: 243108947
ribeirojose · Apr 16, 2019 · 2ce82ea · 2ce82ea
1 parent 2435916
commit 2ce82ea
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/dopamine/agents/rainbow/rainbow_agent.py b/dopamine/agents/rainbow/rainbow_agent.py
@@ -194,7 +194,8 @@ def _build_replay_buffer(self, use_staging):
         stack_size=self.stack_size,
         use_staging=use_staging,
         update_horizon=self.update_horizon,
-        gamma=self.gamma)
+        gamma=self.gamma,
+        observation_dtype=self.observation_dtype.as_numpy_dtype)
 
   def _build_target_distribution(self):
     """Builds the C51 target distribution as per Bellemare et al. (2017).