diff --git a/environments/factory/factory_cleaning.py b/environments/factory/base_factory.py similarity index 86% rename from environments/factory/factory_cleaning.py rename to environments/factory/base_factory.py index 1f035fe..8cc652e 100644 --- a/environments/factory/factory_cleaning.py +++ b/environments/factory/base_factory.py @@ -3,7 +3,7 @@ from pathlib import Path from environments import helpers as h -class Factory(object): +class BaseFactory(object): LEVELS_DIR = 'levels' def __init__(self, level='simple', n_agents=1, max_steps=1e3): @@ -16,6 +16,7 @@ class Factory(object): def reset(self): self.done = False + self.steps = 0 self.agents = np.zeros((self.n_agents, *self.level.shape), dtype=np.int8) free_cells = np.argwhere(self.level == 0) np.random.shuffle(free_cells) @@ -30,6 +31,7 @@ class Factory(object): assert type(actions) in [int, list] if type(actions) == int: actions = [actions] + self.steps += 1 r = 0 # level, agent 1,..., agent n, for i, a in enumerate(actions): @@ -44,20 +46,26 @@ class Factory(object): collisions_vec[i+1] = 0 # no self-collisions collision_vecs.append(collisions_vec) self.handle_collisions(collisions_vec) + r += self.step_core(collisions_vec, actions, r) + if self.steps >= self.max_steps: + self.done = True return self.state, r, self.done, {} def make_move(self, agent_i, old_pos, new_pos): (x, y), (x_new, y_new) = old_pos, new_pos self.state[agent_i+1, x, y] = 0 self.state[agent_i+1, x_new, y_new] = 1 - print( old_pos, new_pos) + print(old_pos, new_pos) def handle_collisions(self, vecs): pass + def step_core(self, collisions_vec, actions, r): + return 0 + if __name__ == '__main__': - factory = Factory(n_agents=1) + factory = BaseFactory(n_agents=1) print(factory.state) state, r, done, _ = factory.step(0) print(state) \ No newline at end of file