pokemon logging demo for leanke

clean_pufferl.py

@@ -318,10 +318,12 @@ def evaluate(data):
             continue
         if 'pokemon_exploration_map' in k:
             import cv2
+            from pokemon_red_eval import make_pokemon_red_overlay
             bg = cv2.imread('kanto_map_dsv.png')
             overlay = make_pokemon_red_overlay(bg, sum(v))
             if data.wandb is not None:
                 data.stats['Media/exploration_map'] = data.wandb.Image(overlay)
+            # @Leanke: Add your infos['learner']['x'] etc
         try: # TODO: Better checks on log data types
             data.stats[k] = np.mean(v)
         except:

demo.py

@@ -230,7 +230,7 @@ def train(args, env_module, make_env):
             agent_creator=make_policy,
             agent_kwargs={'env_module': env_module, 'args': args},
             model_path=args.eval_model_path,
-            device=args.args.device
+            device=args.train.device
         )
     elif args.mode == 'evaluate' and pkg == 'pokemon_red':
         import pokemon_red_eval

pokemon_red_eval.py

@@ -10,7 +10,7 @@ def make_pokemon_red_overlay(bg, counts):
 
     # Convert counts to hue map
     hsv = np.zeros((*counts.shape, 3))
-    hsv[..., 0] = scaled*(240.0/360.0)
+    hsv[..., 0] = 2*(1-scaled)/3
     hsv[..., 1] = nonzero
     hsv[..., 2] = nonzero