Merge pull request #41 from vakker/entity-types

Entity types

Author: gaorkl

.github/workflows/python-package.yml

@@ -37,7 +37,6 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        sudo apt-get update
         pip install flake8 pytest
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
     - name: Lint with flake8

simple_playgrounds/playground.py

@@ -70,6 +70,8 @@ def __init__(self, size):
         self._handle_interactions()
         self.sensor_collision_index = 2
 
+        self.entity_types_map = {}
+
     @staticmethod
     def _initialize_space():
         """ Method to initialize Pymunk empty space for 2D physics.
@@ -353,6 +355,16 @@ def _add_scene_element(self, new_scene_element, keep_position):
         if new_scene_element in self._disappeared_scene_elements:
             self._disappeared_scene_elements.remove(new_scene_element)
 
+    def create_entity_types_map(self,
+                                additional_types=[]):
+        entity_types = [type(e) for e in self.scene_elements]
+        entity_types.extend(additional_types)
+
+        self.entity_types_map = {}
+        for et in entity_types:
+            if et not in self.entity_types_map:
+                self.entity_types_map[et] = len(self.entity_types_map)
+
     def _entity_colliding(self, entity):
 
         collides = False

simple_playgrounds/playgrounds/collection/rl/foraging.py

@@ -5,8 +5,8 @@
                                                      Trajectory)
 
 
-@PlaygroundRegister.register('foraging', 'candy_collect')
-class CandyCollectEnv(SingleRoom):
+@PlaygroundRegister.register('foraging', 'candy_poison')
+class CandyPoisonEnv(SingleRoom):
     def __init__(self):
         super().__init__(size=(200, 200))
 
@@ -17,6 +17,7 @@ def __init__(self):
                                        width_length=(100, 100))
         self.agent_starting_area = area_start
 
+        additional_types = set()
         for loc in ["down-left", "up-right"]:
             area_center, size_area = self.get_area((0, 0), loc)
             area = CoordinateSampler(center=area_center,
@@ -27,16 +28,20 @@ def __init__(self):
                                          probability=0.01,
                                          limit=2)
             self.add_scene_element(field)
+            additional_types.add(field.entity_produced)
 
             field = scene_elements.Field(entity_produced=scene_elements.Poison,
                                          production_area=area,
                                          probability=0.01,
                                          limit=2)
             self.add_scene_element(field)
+            additional_types.add(field.entity_produced)
 
         self.time_limit = 2000
         self.time_limit_reached_reward = -1
 
+        self.create_entity_types_map(additional_types)
+
 
 @PlaygroundRegister.register('foraging', 'candy_fireballs')
 class CandyFireballs(SingleRoom):
@@ -49,6 +54,8 @@ def __init__(self, time_limit=100, probability_production=0.4):
             'size_tiles': 4
         }
 
+        additional_types = set()
+
         # First Fireball
         text_1 = {'color_min': [220, 0, 200], 'color_max': [255, 100, 220]}
         trajectory = Trajectory('waypoints',
@@ -60,6 +67,7 @@ def __init__(self, time_limit=100, probability_production=0.4):
             **text_1
         })
         self.add_scene_element(fireball, trajectory)
+        additional_types.add(type(fireball))
 
         # Second Fireball
         text_2 = {'color_min': [180, 0, 0], 'color_max': [220, 100, 0]}
@@ -92,5 +100,8 @@ def __init__(self, time_limit=100, probability_production=0.4):
                                      production_area=area_prod,
                                      probability=probability_production)
         self.add_scene_element(field)
+        additional_types.add(field.entity_produced)
 
         self.time_limit = time_limit
+
+        self.create_entity_types_map(additional_types)

simple_playgrounds/playgrounds/collection/rl/navigation.py

@@ -61,6 +61,8 @@ def __init__(self,
         self.time_limit = time_limit
         self.time_limit_reached_reward = reward_reached_time_limit
 
+        self.create_entity_types_map([Basic])
+
     def _set_goal(self):
 
         index_goal = random.randint(0, 3)
@@ -115,4 +117,4 @@ def __init__(self,
 
         self.time_limit = time_limit
 
-
+        self.create_entity_types_map()

simple_playgrounds/playgrounds/collection/rl/sequential.py

@@ -72,6 +72,8 @@ def _place_scene_elements(self):
             allow_overlapping=False)
         self.add_scene_element(self.dispenser, self.area_dispenser)
 
+        self.create_entity_types_map([self.dispenser.entity_produced])
+
     def reset(self):
         self.remove_scene_element(self.dispenser)
 
@@ -148,3 +150,5 @@ def __init__(
                                        area_shape='rectangle',
                                        width_length=area_start_shape)
         self.initial_agent_coordinates = area_start
+
+        self.create_entity_types_map([dispenser.entity_produced])