steffen/marl-factory-grid
1
0
Fork 0
mirror of https://github.com/illiumst/marl-factory-grid.git synced 2025-05-23 07:16:44 +02:00
Code Issues Packages Projects Releases Wiki Activity

updated simple factory to include dirt

Browse Source
This commit is contained in:
romue 2021-05-07 16:32:17 +02:00
parent 385e6868dd
commit ed10b3225d
1 changed files with 16 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
environments/factory/simple_factory.py
View File

@ -1,13 +1,26 @@
import numpy as np
from environments.factory.base_factory import BaseFactory from environments.factory.base_factory import BaseFactory
class SimpleFactory(BaseFactory): class SimpleFactory(BaseFactory):
def __init__(self, *args, max_dirt=5, **kwargs): def __init__(self, *args, max_dirt=5, **kwargs):
self.max_dirt = max_dirt
super(SimpleFactory, self).__init__(*args, **kwargs) 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__': if __name__ == '__main__':
factory = SimpleFactory(n_agents=1, max_dirt=2) factory = SimpleFactory(n_agents=1, max_dirt=8)
print(factory.state) #print(factory.state)
state, r, done, _ = factory.step(0) state, r, done, _ = factory.step(0)
print(state) #print(state)