updated simple factory

environments/factory/base_factory.py

@@ -46,7 +46,6 @@ class BaseFactory(object):
             x, y = np.argwhere(agent_slice == 1)[0]
             collisions_vec = self.state[:, x, y].copy()  # otherwise you overwrite the grid/state
             collisions_vec[i+1] = 0  # no self-collisions
-            #collision_vecs.append(collisions_vec)
             collision_vecs[i] += collisions_vec
         reward, info = self.step_core(np.array(collision_vecs), actions, r)
         r += reward

environments/factory/simple_factory.py

@@ -8,18 +8,21 @@ class SimpleFactory(BaseFactory):
         super(SimpleFactory, self).__init__(*args, **kwargs)
         self.slice_strings.update({self.state.shape[0]-1: 'dirt'})
 
+    def spawn_dirt(self):
+        free_for_dirt = self.free_cells()
+        for x, y in free_for_dirt[:self.max_dirt]:  # randomly distribute dirt across the grid
+            self.state[-1, x, y] = 1
+
     def reset(self):
         super().reset()
         dirt_slice = np.zeros((1, *self.state.shape[1:]))
         self.state = np.concatenate((self.state, dirt_slice))  # dirt is now the last slice
-        free_for_dirt = self.free_cells()
-        for x, y in free_for_dirt[:self.max_dirt]:
-            self.state[-1, x, y] = 1
+        self.spawn_dirt()
 
     def step_core(self, collisions_vecs, actions, r):
         for agent_i, cols in enumerate(collisions_vecs):
             cols = np.argwhere(cols != 0).flatten()
-            print(f'Agent #{agent_i} has collisions with '
+            print(f't = {self.steps}\tAgent {agent_i} has collisions with '
                   f'{[self.slice_strings[entity] for entity in cols]}')
         return 0, {}