From ed10b3225ddc1280aa23dc360cf38694aa20f875 Mon Sep 17 00:00:00 2001 From: romue Date: Fri, 7 May 2021 16:32:17 +0200 Subject: [PATCH] updated simple factory to include dirt --- environments/factory/simple_factory.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/environments/factory/simple_factory.py b/environments/factory/simple_factory.py index d8be737..f809abc 100644 --- a/environments/factory/simple_factory.py +++ b/environments/factory/simple_factory.py @@ -1,13 +1,26 @@ +import numpy as np from environments.factory.base_factory import BaseFactory class SimpleFactory(BaseFactory): def __init__(self, *args, max_dirt=5, **kwargs): + self.max_dirt = max_dirt super(SimpleFactory, self).__init__(*args, **kwargs) + 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.state.sum(0) + free_for_dirt = np.argwhere(free_for_dirt == 0) + np.random.shuffle(free_for_dirt) + for x,y in free_for_dirt[:self.max_dirt]: + self.state[-1, x, y] = 1 + print(self.state) + if __name__ == '__main__': - factory = SimpleFactory(n_agents=1, max_dirt=2) - print(factory.state) + 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 + #print(state) \ No newline at end of file