added image_b to info dict

gym_donkeycar/envs/donkey_sim.py

@@ -87,7 +87,10 @@ def __init__(self, conf):
         # sensor size - height, width, depth
         self.camera_img_size = conf["cam_resolution"]
         self.image_array = np.zeros(self.camera_img_size)
-        self.last_obs = None
+        self.image_array_b = None
+        self.last_obs = self.image_array
+        self.time_received = time.time()
+        self.last_received = self.time_received
         self.hit = "none"
         self.cte = 0.0
         self.x = 0.0
@@ -208,6 +211,7 @@ def send_config(self, conf):
             )
             self.send_cam_config(**cam_config_b, msg_type="cam_config_b")
             logger.info(f"done sending cam config B. {cam_config_b}")
+            self.image_array_b = np.zeros(self.camera_img_size)
 
         if "lidar_config" in conf.keys():
             lidar_config = self.extract_keys(
@@ -320,7 +324,10 @@ def reset(self):
         self.timer.reset()
         time.sleep(1)
         self.image_array = np.zeros(self.camera_img_size)
+        self.image_array_b = None
         self.last_obs = self.image_array
+        self.time_received = time.time()
+        self.last_received = self.time_received
         self.hit = "none"
         self.cte = 0.0
         self.x = 0.0
@@ -353,15 +360,14 @@ def take_action(self, action):
         self.send_control(action[0], action[1])
 
     def observe(self):
-        while self.last_obs is self.image_array:
-            time.sleep(1.0 / 120.0)
+        while self.last_received == self.time_received:
+            time.sleep(0.001)
 
-        self.last_obs = self.image_array
+        self.last_received = self.time_received
         observation = self.image_array
         done = self.is_game_over()
         reward = self.calc_reward(done)
-        # info = {'pos': (self.x, self.y, self.z), 'cte': self.cte,
-        #        "speed": self.speed, "hit": self.hit}
+
         info = {
             "pos": (self.x, self.y, self.z),
             "cte": self.cte,
@@ -374,6 +380,10 @@ def observe(self):
             "car": (self.roll, self.pitch, self.yaw),
         }
 
+        # Add the second image to the dict
+        if self.image_array_b is not None:
+            info["image_b"] = self.image_array_b
+
         # self.timer.on_frame()
 
         return observation, reward, done, info
@@ -412,11 +422,20 @@ def on_telemetry(self, data):
 
         # always update the image_array as the observation loop will hang if not changing.
         self.image_array = np.asarray(image)
+        self.time_received = time.time()
+
+        if "image_b" in data:
+            imgString_b = data["image_b"]
+            image_b = Image.open(BytesIO(base64.b64decode(imgString_b)))
+            self.image_array_b = np.asarray(image_b)
+
+        if "pos_x" in data:
+            self.x = data["pos_x"]
+            self.y = data["pos_y"]
+            self.z = data["pos_z"]
 
-        self.x = data["pos_x"]
-        self.y = data["pos_y"]
-        self.z = data["pos_z"]
-        self.speed = data["speed"]
+        if "speed" in data:
+            self.speed = data["speed"]
 
         if "gyro_x" in data:
             self.gyro_x = data["gyro_x"]
@@ -449,7 +468,8 @@ def on_telemetry(self, data):
         if self.over:
             return
 
-        self.hit = data["hit"]
+        if "hit" in data:
+            self.hit = data["hit"]
 
         self.determine_episode_over()