From a28820e8c345334625e7b57848fcff093ecbe5da Mon Sep 17 00:00:00 2001
From: romue <robert.mueller@campus.lmu.de>
Date: Mon, 10 May 2021 11:56:09 +0200
Subject: [PATCH] updated simple factory

---
 environments/factory/base_factory.py   | 17 +++++++++++++----
 environments/factory/simple_factory.py | 23 +++++++++++++----------
 2 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/environments/factory/base_factory.py b/environments/factory/base_factory.py
index 8880dbc..df4827f 100644
--- a/environments/factory/base_factory.py
+++ b/environments/factory/base_factory.py
@@ -12,6 +12,7 @@ class BaseFactory(object):
         self.level = h.one_hot_level(
             h.parse_level(Path(__file__).parent / self.LEVELS_DIR / f'{level}.txt')
         )#[np.newaxis, ...]
+        self.slice_strings = {0: 'level', **{i: f'agent#{i}' for i in range(1, self.n_agents+1)}}
         self.reset()
 
     def reset(self):
@@ -23,7 +24,6 @@ class BaseFactory(object):
         for i in range(self.n_agents):
             r, c = free_cells[i]
             self.agents[i, r, c] = 1
-        free_cells = free_cells[self.n_agents:]
         self.state = np.concatenate((self.level[np.newaxis, ...], self.agents), 0)
         return self.state, 0, self.done, {}
 
@@ -34,18 +34,21 @@ class BaseFactory(object):
         self.steps += 1
         r = 0
         # level, agent 1,..., agent n,
+        collision_vecs = np.zeros((self.n_agents, self.state.shape[0]))  # n_agents x n_slices
         for i, a in enumerate(actions):
             old_pos, new_pos, valid = h.check_agent_move(state=self.state, dim=i+1, action=a)
             if valid:
                 self.make_move(i, old_pos, new_pos)
-        collision_vecs = []
+            else:  # trying to leave the level
+                collision_vecs[i, 0] = 1
         for i in range(self.n_agents):  # might as well save the positions (redundant)
             agent_slice = self.state[i+1]
             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)
-        reward, info = self.step_core(collisions_vec, actions, r)
+            #collision_vecs.append(collisions_vec)
+            collision_vecs[i] += collisions_vec
+        reward, info = self.step_core(np.array(collision_vecs), actions, r)
         r += reward
         if self.steps >= self.max_steps:
             self.done = True
@@ -56,5 +59,11 @@ class BaseFactory(object):
         self.state[agent_i+1, x, y] = 0
         self.state[agent_i+1, x_new, y_new] = 1
 
+    def free_cells(self):
+        free_cells = self.state.sum(0)
+        free_cells = np.argwhere(free_cells == 0)
+        np.random.shuffle(free_cells)
+        return free_cells
+
     def step_core(self, collisions_vec, actions, r):
         return 0, {}
diff --git a/environments/factory/simple_factory.py b/environments/factory/simple_factory.py
index ba4c77f..f5b1604 100644
--- a/environments/factory/simple_factory.py
+++ b/environments/factory/simple_factory.py
@@ -6,25 +6,28 @@ class SimpleFactory(BaseFactory):
     def __init__(self, *args, max_dirt=5, **kwargs):
         self.max_dirt = max_dirt
         super(SimpleFactory, self).__init__(*args, **kwargs)
+        self.slice_strings.update({self.state.shape[0]-1: 'dirt'})
 
     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_for_dirt()
+        free_for_dirt = self.free_cells()
         for x, y in free_for_dirt[:self.max_dirt]:
             self.state[-1, x, y] = 1
-        print(self.state)
 
-    def free_for_dirt(self):
-        free_for_dirt = self.state.sum(0)
-        free_for_dirt = np.argwhere(free_for_dirt == 0)
-        np.random.shuffle(free_for_dirt)
-        return free_for_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 '
+                  f'{[self.slice_strings[entity] for entity in cols]}')
+        return 0, {}
+
 
 
 if __name__ == '__main__':
+    import random
     factory = SimpleFactory(n_agents=1, max_dirt=8)
-    #print(factory.state)
-    state, r, done, _ = factory.step(0)
-    #print(state)
\ No newline at end of file
+    random_actions = [random.randint(0,8) for _ in range(200)]
+    for action in random_actions:
+        state, r, done, _ = factory.step(action)
\ No newline at end of file