diff --git a/environments/factory/base_factory.py b/environments/factory/base_factory.py index cb317ad..0fbb306 100644 --- a/environments/factory/base_factory.py +++ b/environments/factory/base_factory.py @@ -61,10 +61,3 @@ class BaseFactory(object): def step_core(self, collisions_vec, actions, r): return 0 - - -if __name__ == '__main__': - factory = BaseFactory(n_agents=1) - print(factory.state) - state, r, done, _ = factory.step(0) - print(state) \ No newline at end of file diff --git a/environments/factory/simple_factory.py b/environments/factory/simple_factory.py new file mode 100644 index 0000000..d8be737 --- /dev/null +++ b/environments/factory/simple_factory.py @@ -0,0 +1,13 @@ +from environments.factory.base_factory import BaseFactory + + +class SimpleFactory(BaseFactory): + def __init__(self, *args, max_dirt=5, **kwargs): + super(SimpleFactory, self).__init__(*args, **kwargs) + + +if __name__ == '__main__': + factory = SimpleFactory(n_agents=1, max_dirt=2) + print(factory.state) + state, r, done, _ = factory.step(0) + print(state) \ No newline at end of file