Cat change yuxin

metadrive/envs/scenario_env.py

@@ -379,6 +379,55 @@ def _reset_global_seed(self, force_seed=None):
             self.config["start_scenario_index"] + self.config["num_scenarios"])
         self.seed(current_seed)
 
+class ScenarioTrafficManagerWithCAT(ScenarioTrafficManager):
+    def __init__(self):
+        # Load CAT data
+        super().__init__()
+
+        import pickle
+        with open(f'{self.engine.global_config["CAT_dir"]}/save.pkl', 'rb') as f:
+            cat_res = pickle.load(f)
+
+
+        self.cat_traj = cat_res
+
+
+    def after_reset(self):
+        scenario_id = self.engine.data_manager.current_scenario['metadata']['scenario_id']
+        if scenario_id not in self.cat_traj:
+            return None
+        adv_agent = self.cat_traj[scenario_id]['adv_agent']
+
+        adv_traj = self.cat_traj[scenario_id]['adv_traj_raw']
+        # Change data
+        def func(scenario_dict):
+            assert scenario_dict['tracks'][adv_agent]['state']['position'][:, :2].shape == adv_traj[:, :2].shape
+
+            scenario_dict['tracks'][adv_agent]['state']['position'][:, :2] = adv_traj[:, :2]
+            scenario_dict['tracks'][adv_agent]['state']['velocity'][:] = adv_traj[:, 2:4]
+            scenario_dict['tracks'][adv_agent]['state']['heading'][:] = adv_traj[:, 4]
+            scenario_dict['tracks'][adv_agent]['state']['valid'] = np.ones(91)
+
+        self.engine.data_manager.overwrite_data(func, i=self.engine.global_random_seed)
+
+        ret = super().after_reset()
+        return ret
+
+class ScenarioEnvWithCAT(ScenarioEnv):
+    @classmethod
+    def default_config(cls):
+        config = super(ScenarioEnvWithCAT, cls).default_config()
+        config.update({
+            "CAT_dir": "",
+        })
+        return config
+
+    def setup_engine(self):
+        super(ScenarioEnvWithCAT, self).setup_engine()
+        if not self.config["no_traffic"]:
+            self.engine.update_manager("traffic_manager", ScenarioTrafficManagerWithCAT())
+
+
 
 if __name__ == "__main__":
     env = ScenarioEnv(

metadrive/manager/scenario_data_manager.py

@@ -69,6 +69,9 @@ def _get_scenario(self, i):
         assert isinstance(ret, SD)
         return ret
 
+    def overwrite_data(self, func, i):
+        func(self._scenarios[i])
+
     def before_reset(self):
         if not self.store_data:
             assert len(self._scenarios) <= 1, "It seems you access multiple scenarios in one episode"

metadrive/manager/scenario_traffic_manager.py

@@ -367,4 +367,4 @@ def reset_vehicle_type_count(np_random=None):
     if np_random is None:
         type_count = [0 for i in range(3)]
     else:
-        type_count = [np_random.randint(100) for i in range(3)]
+        type_count = [np_random.randint(100) for i in range(3)]