Merge pull request #19 from AIasd/sim_scene_patch

Update sim_scene's get_obs function to accommodate updates on vec_map and state/obs format.

src/trajdata/simulation/sim_scene.py

@@ -121,12 +121,13 @@ def get_obs(
         self, collate: bool = True, get_map: bool = True
     ) -> Union[AgentBatch, Dict[str, Any]]:
         agent_data_list: List[AgentBatchElement] = list()
+        self.cache.set_obs_format(self.dataset.obs_format)
+
         for agent in self.agents:
             scene_time_agent = SceneTimeAgent(
                 self.scene, self.scene_ts, self.agents, agent, self.cache
             )
-            agent_data_list.append(
-                AgentBatchElement(
+            batch_element: AgentBatchElement = AgentBatchElement(
                     self.cache,
                     -1,  # Not used
                     scene_time_agent,
@@ -136,11 +137,23 @@ def get_obs(
                     incl_robot_future=False,
                     incl_raster_map=get_map and self.dataset.incl_raster_map,
                     raster_map_params=self.dataset.raster_map_params,
+                    map_api=self.dataset._map_api,
+                    vector_map_params=self.dataset.vector_map_params,
+                    state_format=self.dataset.state_format,
                     standardize_data=self.dataset.standardize_data,
                     standardize_derivatives=self.dataset.standardize_derivatives,
                     max_neighbor_num=self.dataset.max_neighbor_num,
                 )
-            )
+            agent_data_list.append(batch_element)
+
+            for key, extra_fn in self.dataset.extras.items():
+                batch_element.extras[key] = extra_fn(batch_element)
+
+            for transform_fn in self.dataset.transforms:
+                batch_element = transform_fn(batch_element)
+
+            if not self.dataset.vector_map_params.get("collate", False):
+                batch_element.vec_map = None
 
             # Need to reset transformations for each agent since each
             # AgentBatchElement transforms (standardizes) the cache.