Occupation Map

environments/factory/base/base_factory.py

@@ -50,6 +50,11 @@ class BaseFactory(gym.Env):
     def movement_actions(self):
         return self._actions.movement_actions
 
+    @property
+    def params(self) -> dict:
+        d = {key: val for key, val in self.__dict__.items() if not key.startswith('_') and not key.startswith('__')}
+        return d
+
     def __enter__(self):
         return self if self.obs_prop.frames_to_stack == 0 else \
             FrameStack(self, self.obs_prop.frames_to_stack)
@@ -576,8 +581,7 @@ class BaseFactory(gym.Env):
 
     def save_params(self, filepath: Path):
         # noinspection PyProtectedMember
-        # d = {key: val._asdict() if hasattr(val, '_asdict') else val for key, val in self.__dict__.items()
-        d = {key: val for key, val in self.__dict__.items() if not key.startswith('_') and not key.startswith('__')}
+        d = self.params
         filepath.parent.mkdir(parents=True, exist_ok=True)
         with filepath.open('w') as f:
             simplejson.dump(d, f, indent=4, namedtuple_as_object=True)
@@ -587,6 +591,7 @@ class BaseFactory(gym.Env):
 
         for entity_group in self._entities:
             summary.update({f'{REC_TAC}{entity_group.name}': entity_group.summarize_states(n_steps=self._steps)})
+
         return summary
 
     def print(self, string):

environments/factory/factory_dirt.py

@@ -239,7 +239,12 @@ class DirtFactory(BaseFactory):
 
         if agent.temp_action == CLEAN_UP_ACTION:
             if agent.temp_valid:
+                # Reward if pickup succeds,
+                #  0.5 on every pickup
                 reward += 0.5
+                if self.dirt_prop.done_when_clean and (len(self[c.DIRT]) == 0):
+                    #  0.5 additional reward for the very last pickup
+                    reward += 0.5
                 self.print(f'{agent.name} did just clean up some dirt at {agent.pos}.')
                 info_dict.update(dirt_cleaned=1)
             else:

environments/logging/recorder.py

@@ -3,6 +3,8 @@ from collections import defaultdict
 from pathlib import Path
 from typing import Union
 
+import numpy as np
+import pandas as pd
 import simplejson
 from stable_baselines3.common.callbacks import BaseCallback
 
@@ -10,21 +12,39 @@ from environments.factory.base.base_factory import REC_TAC
 
 
 # noinspection PyAttributeOutsideInit
+from environments.helpers import Constants as c
+
+
 class RecorderCallback(BaseCallback):
 
-    def __init__(self, filepath: Union[str, Path], occupation_map: bool = False, trajectory_map: bool = False):
+    def __init__(self, filepath: Union[str, Path], occupation_map: bool = False, trajectory_map: bool = False,
+                 entities='all'):
         super(RecorderCallback, self).__init__()
         self.trajectory_map = trajectory_map
         self.occupation_map = occupation_map
         self.filepath = Path(filepath)
         self._recorder_dict = defaultdict(list)
         self._recorder_out_list = list()
+        self._env_params = None
         self.do_record: bool
+        if isinstance(entities, str):
+            if entities.lower() == 'all':
+                self._entities = None
+            else:
+                self._entities = [entities]
+        else:
+            self._entities = entities
         self.started = False
         self.closed = False
 
+    def read_params(self, params):
+        self._env_params = params
+
     def read_info(self, env_idx, info: dict):
         if info_dict := {key.replace(REC_TAC, ''): val for key, val in info.items() if key.startswith(f'{REC_TAC}')}:
+            if self._entities:
+                info_dict = {k: v for k, v in info_dict.items() if k in self._entities}
+
             info_dict.update(episode=(self.num_timesteps + env_idx))
             self._recorder_dict[env_idx].append(info_dict)
         else:
@@ -51,14 +71,27 @@ class RecorderCallback(BaseCallback):
         if self.do_record and self.started:
             # self.out_file.unlink(missing_ok=True)
             with self.filepath.open('w') as f:
-                out_dict = {'episodes': self._recorder_out_list}
+                out_dict = {'episodes': self._recorder_out_list, 'header': self._env_params}
                 try:
                     simplejson.dump(out_dict, f, indent=4)
                 except TypeError:
                     print('Shit')
 
             if self.occupation_map:
-                print('Recorder files were dumped to disk, now plotting the occupation map...')
+                a = np.zeros((15, 15))
+                for episode in  out_dict['episodes']:
+                    df = pd.DataFrame([y for x in episode['steps'] for y in x['Agents']])
+
+                    b = list(df[['x', 'y']].to_records(index=False))
+
+                    np.add.at(a, tuple(zip(*b)), 1)
+
+                # a = np.rot90(a)
+                import seaborn as sns
+                from matplotlib import pyplot as plt
+                hm = sns.heatmap(data=a)
+                hm.set_title('Very Nice Heatmap')
+                plt.show()
 
             if self.trajectory_map:
                 print('Recorder files were dumped to disk, now plotting the occupation map...')
@@ -75,7 +108,7 @@ class RecorderCallback(BaseCallback):
 
             for env_idx, done in list(
                     enumerate(self.locals.get('dones', []))) + list(
-                    enumerate(self.locals.get('done', []))):
+                enumerate(self.locals.get('done', []))):
                 self.read_done(env_idx, done)
         else:
             pass