Fix history vehicles key error + not recreate SMARTS instance

examples/history_vehicles_replacement_for_imitation_learning.py

@@ -20,6 +20,9 @@ def act(self, obs):
 
 def main(scenarios, headless, seed):
     scenarios_iterator = Scenario.scenario_variations(scenarios, [])
+    first_time = True
+    smarts = None
+
     for _ in scenarios:
         scenario = next(scenarios_iterator)
         agent_missions = scenario.discover_missions_of_traffic_histories()
@@ -33,14 +36,18 @@ def main(scenarios, headless, seed):
                 ),
                 agent_builder=KeepLaneAgent,
             )
-
             agent = agent_spec.build_agent()
 
-            smarts = SMARTS(
-                agent_interfaces={agent_id: agent_spec.interface},
-                traffic_sim=SumoTrafficSimulation(headless=True, auto_start=True),
-                envision=Envision(),
-            )
+            if first_time:
+                smarts = SMARTS(
+                    agent_interfaces={agent_id: agent_spec.interface},
+                    traffic_sim=SumoTrafficSimulation(headless=True, auto_start=True),
+                    envision=Envision(),
+                )
+                first_time = False
+            else:
+                smarts.switch_ego_agent({agent_id: agent_spec.interface})
+
             observations = smarts.reset(scenario)
 
             dones = {agent_id: False}
@@ -52,7 +59,7 @@ def main(scenarios, headless, seed):
                     {agent_id: agent_action}
                 )
 
-            smarts.destroy()
+    smarts.destroy()
 
 
 if __name__ == "__main__":

smarts/core/agent_manager.py

@@ -300,6 +300,9 @@ def send_observations_to_social_agents(self, observations):
             obs = observations[agent_id]
             self._remote_social_agents_action[agent_id] = remote_agent.act(obs)
 
+    def switch_initial_agent(self, agent_interface):
+        self._initial_interfaces = agent_interface
+
     def setup_agents(self, sim):
         self.init_ego_agents(sim)
         self.setup_social_agents(sim)

smarts/core/remote_agent_buffer.py

@@ -196,7 +196,9 @@ def get_manager_channel_stub(addr):
     try:
         # Wait until the grpc server is ready or timeout after 30 seconds
         grpc.channel_ready_future(channel).result(timeout=30)
-    except grpc.FutureTimeoutError:
-        raise RemoteAgentException("Timeout in connecting to remote zoo manager.")
+    except grpc.FutureTimeoutError as e:
+        raise RemoteAgentException(
+            "Timeout in connecting to remote zoo manager."
+        ) from e
     stub = manager_pb2_grpc.ManagerStub(channel)
     return channel, stub

smarts/core/smarts.py

@@ -369,6 +369,10 @@ def setup(self, scenario: Scenario):
     def add_provider(self, provider):
         self._providers.append(provider)
 
+    def switch_ego_agent(self, agent_interface):
+        self._agent_manager.switch_initial_agent(agent_interface)
+        self._is_setup = False
+
     def _setup_road_network(self):
         glb_path = self.scenario.map_glb_filepath
         if self._road_network_np:

smarts/core/trap_manager.py

@@ -145,6 +145,10 @@ def largest_vehicle_plane_dimension(vehicle):
                 ),
             )
             for v_id in sorted_vehicle_ids:
+                # Skip the capturing process if history traffic is used
+                if sim.scenario.traffic_history != {}:
+                    break
+
                 vehicle = vehicles[v_id]
                 point = Point(vehicle.position)