Adding functionality to plotting code to be able to reproduce figures…

… from Bellemare, Dabney, Munos (2017). This includes: proper Atari image dimensions (4:3), and a x axis on the return distribution histogram that indicates return values. PiperOrigin-RevId: 400192045
SuReLI · Oct 6, 2021 · cc3cfcd · cc3cfcd
1 parent 0d155c1
commit cc3cfcd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/dopamine/utils/atari_plotter.py b/dopamine/utils/atari_plotter.py
@@ -32,6 +32,8 @@ class AtariPlotter(plotter.Plotter):
   _defaults = {
       'x': 0,
       'y': 0,
+      'input_width': 160,
+      'input_height': 210,
       'width': 160,
       'height': 210,
   }
@@ -46,8 +48,8 @@ def __init__(self, parameter_dict=None):
     """
     super(AtariPlotter, self).__init__(parameter_dict)
     assert 'environment' in self.parameters
-    self.game_surface = pygame.Surface((self.parameters['width'],
-                                        self.parameters['height']))
+    self.game_surface = pygame.Surface((self.parameters['input_width'],
+                                        self.parameters['input_height']))
 
   def draw(self):
     """Render the Atari 2600 frame.

diff --git a/dopamine/utils/plotter.py b/dopamine/utils/plotter.py
@@ -66,6 +66,11 @@ def _setup_plot(self):
     if 'title' in self.parameters:
       self.plot.set_title(self.parameters['title'],
                           fontsize=self.parameters['fontsize'] + 2)
+    if 'xticks' in self.parameters:
+      self.plot.set_xticks(self.parameters['xticks'])
+    if 'xticklabels' in self.parameters:
+      self.plot.set_xticklabels(self.parameters['xticklabels'])
+
     self.plot.tick_params(labelsize=self.parameters['fontsize'])
 
   @abc.abstractmethod