Merge 05e95d8 into cd3d378

BerkeleyAutomation · Sep 13, 2018 · 4a41c6d · 4a41c6d
2 parents cd3d378 + 05e95d8
commit 4a41c6d
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 9 deletions.
diff --git a/examples/fluids_example.py b/examples/fluids_example.py
@@ -4,7 +4,7 @@
 
 simulator = fluids.FluidSim(visualization_level=1,        # How much debug visualization you want to enable. Set to 0 for no vis
                             fps=0,                        # If set to non 0, caps the FPS. Target is 30
-                            obs_space=fluids.OBS_BIRDSEYE,# OBS_BIRDSEYE, OBS_GRID, or OBS_NONE
+                            obs_space=fluids.OBS_GRID,# OBS_BIRDSEYE, OBS_GRID, or OBS_NONE
                             background_control=fluids.BACKGROUND_CSP) # BACKGROUND_CSP or BACKGROUND_NULL
 
 state = fluids.State(

diff --git a/fluids/assets/__init__.py b/fluids/assets/__init__.py
@@ -18,7 +18,9 @@
             PedCrossing,
             Pedestrian,
             Car,
-            TrafficLight,
+            "TrafficLight-Red",
+            "TrafficLight-Green",
+            "TrafficLight-Yellow",
             CrossWalkLight,
             CrossWalk,
             Waypoint]

diff --git a/fluids/assets/traffic_light.py b/fluids/assets/traffic_light.py
@@ -1,8 +1,6 @@
 from fluids.assets.shape import Shape
+from fluids.consts import *
 
-RED = (0xf6, 0x11, 0x46)
-YELLOW = (0xfc, 0xef, 0x5e),
-GREEN = (0, 0xc6, 0x44)
 
 
 

diff --git a/fluids/consts.py b/fluids/consts.py
@@ -10,3 +10,7 @@
 
 REWARD_PATH = "fluids_reward_path"
 REWARD_NONE = "fluids_reward_none"
+
+RED = (0xf6, 0x11, 0x46)
+YELLOW = (0xfc, 0xef, 0x5e),
+GREEN = (0, 0xc6, 0x44)
diff --git a/fluids/obs/grid.py b/fluids/obs/grid.py
@@ -5,6 +5,8 @@
 from fluids.obs.obs import FluidsObs
 from fluids.utils import rotation_array
 from scipy.misc import imresize
+from fluids.consts import *
+
 class GridObservation(FluidsObs):
     """
     Grid observation type. 
@@ -30,6 +32,13 @@ def __init__(self, car, obs_dim=500, shape=(500,500)):
         for k, obj in iteritems(state.objects):
             if (car.can_collide(obj) or type(obj) in {TrafficLight, Lane, Street}) and self.grid_square.intersects(obj):
                 typ = type(obj)
+                if typ == TrafficLight:
+                    if obj.color == RED:
+                        typ = "TrafficLight-Red"
+                    elif obj.color == GREEN:
+                        typ = "TrafficLight-Green"
+                    elif obj.color == YELLOW:
+                        typ = "TrafficLight-Yellow"
                 if typ not in collideable_map:
                     collideable_map[typ] = []
                 collideable_map[typ].append(obj)
@@ -43,7 +52,9 @@ def __init__(self, car, obs_dim=500, shape=(500,500)):
         undrivable_window = pygame.Surface((self.grid_dim, self.grid_dim))
         car_window        = pygame.Surface((self.grid_dim, self.grid_dim))
         ped_window        = pygame.Surface((self.grid_dim, self.grid_dim))
-        light_window      = pygame.Surface((self.grid_dim, self.grid_dim))
+        light_window_red  = pygame.Surface((self.grid_dim, self.grid_dim))
+        light_window_green= pygame.Surface((self.grid_dim, self.grid_dim))
+        light_window_yellow=pygame.Surface((self.grid_dim, self.grid_dim))
         direction_window  = pygame.Surface((self.grid_dim, self.grid_dim))
         direction_pixel_window \
                           = pygame.Surface((self.grid_dim, self.grid_dim))
@@ -79,9 +90,18 @@ def __init__(self, car, obs_dim=500, shape=(500,500)):
         for obj in collideable_map[Pedestrian]:
             rel_obj = obj.get_relative(rel)
             rel_obj.render(ped_window, border=None)
-        for obj in collideable_map[TrafficLight]:
+
+        for obj in collideable_map["TrafficLight-Red"]:
+            rel_obj = obj.get_relative(rel)
+            rel_obj.render(light_window_red, border=None)
+
+        for obj in collideable_map["TrafficLight-Green"]:
+            rel_obj = obj.get_relative(rel)
+            rel_obj.render(light_window_green, border=None)
+
+        for obj in collideable_map["TrafficLight-Yellow"]:
             rel_obj = obj.get_relative(rel)
-            rel_obj.render(light_window, border=None)
+            rel_obj.render(light_window_green, border=None)
 
         point = (int(gd/6), int(gd/2))
         edge_point = None
@@ -119,7 +139,9 @@ def is_on_screen(point, gd):
                                                                               undrivable_window,
                                                                               car_window,
                                                                               ped_window,
-                                                                              light_window,
+                                                                              light_window_red,
+                                                                              light_window_green,
+                                                                              light_window_yellow,
                                                                               direction_window,
                                                                               direction_pixel_window,
                                                                               direction_edge_window