From 011657a655974bf181974c74c047d51b88836ef9 Mon Sep 17 00:00:00 2001
From: thatguy11325 <148832074+thatguy11325@users.noreply.github.com>
Date: Thu, 27 Jun 2024 14:18:40 -0400
Subject: [PATCH] Log the map fewer times

---
 pokemonred_puffer/cleanrl_puffer.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/pokemonred_puffer/cleanrl_puffer.py b/pokemonred_puffer/cleanrl_puffer.py
index c8230a3..73c5e2f 100644
--- a/pokemonred_puffer/cleanrl_puffer.py
+++ b/pokemonred_puffer/cleanrl_puffer.py
@@ -303,11 +303,9 @@ def evaluate(self):
             for k, v in self.infos.items():
                 # Moves into models... maybe. Definitely moves.
                 # You could also just return infos and have it in demo
-                if "pokemon_exploration_map" in self.infos and self.config.save_overlay is True:
+                if "pokemon_exploration_map" in k and self.config.save_overlay is True:
                     if self.epoch % self.config.overlay_interval == 0:
-                        overlay = make_pokemon_red_overlay(
-                            np.stack(self.infos["pokemon_exploration_map"], axis=0)
-                        )
+                        overlay = make_pokemon_red_overlay(np.stack(self.infos[k], axis=0))
                         if self.wandb_client is not None:
                             self.stats["Media/aggregate_exploration_map"] = wandb.Image(
                                 overlay, file_type="jpg"