adapt to changes in calvin_agent

mees · Nov 11, 2021 · accfb2e · accfb2e
1 parent e13cdcc
commit accfb2e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/calvin_env/envs/play_lmp_wrapper.py b/calvin_env/envs/play_lmp_wrapper.py
@@ -7,7 +7,7 @@
 try:
     from lfp.datasets.utils.episode_utils import process_actions, process_depth, process_rgb, process_state
 except ImportError:
-    from calvin_models.calvin_agent.datasets.utils.episode_utils import (
+    from calvin_agent.datasets.utils.episode_utils import (
         process_actions,
         process_depth,
         process_rgb,
@@ -63,9 +63,9 @@ def transform_observation(self, obs: Dict[str, Any]) -> Dict[str, Union[torch.Te
         rgb_obs = process_rgb(rgb_obs, self.observation_space_keys, self.transforms)
         depth_obs = process_depth(depth_obs, self.observation_space_keys, self.transforms)
 
-        state_obs["robot_obs"] = state_obs["robot_obs"].to(self.device)
-        rgb_obs.update({"rgb_obs": {k: v.to(self.device) for k, v in rgb_obs["rgb_obs"].items()}})
-        depth_obs.update({"depth_obs": {k: v.to(self.device) for k, v in depth_obs["depth_obs"].items()}})
+        state_obs["robot_obs"] = state_obs["robot_obs"].to(self.device).unsqueeze(0)
+        rgb_obs.update({"rgb_obs": {k: v.to(self.device).unsqueeze(0) for k, v in rgb_obs["rgb_obs"].items()}})
+        depth_obs.update({"depth_obs": {k: v.to(self.device).unsqueeze(0) for k, v in depth_obs["depth_obs"].items()}})
 
         return {**rgb_obs, **state_obs, **depth_obs}
 

diff --git a/calvin_env/envs/play_table_env.py b/calvin_env/envs/play_table_env.py
@@ -153,7 +153,7 @@ def render(self, mode="human"):
                 log.warning("Environment does not have camera")
                 return
             img = rgb_obs[0][:, :, ::-1]
-            cv2.imshow("simulation cam", img)
+            cv2.imshow("simulation cam", cv2.resize(img, (500, 500)))
             cv2.waitKey(1)
         elif mode == "rgb_array":
             assert len(rgb_obs) > 0, "Environment does not have camera"
@@ -261,21 +261,24 @@ def serialize(self):
         return data
 
 
-def get_env(dataset_path, obs_space, show_gui=True, **kwargs):
+def get_env(dataset_path, obs_space=None, show_gui=True, **kwargs):
     from pathlib import Path
 
     from omegaconf import OmegaConf
 
     render_conf = OmegaConf.load(Path(dataset_path) / ".hydra" / "merged_config.yaml")
 
-    exclude_keys = set(render_conf.cameras.keys()) - {
-        re.split("_", key)[1] for key in obs_space["rgb_obs"] + obs_space["depth_obs"]
-    }
-    for k in exclude_keys:
-        del render_conf.cameras[k]
+    if obs_space is not None:
+        exclude_keys = set(render_conf.cameras.keys()) - {
+            re.split("_", key)[1] for key in obs_space["rgb_obs"] + obs_space["depth_obs"]
+        }
+        for k in exclude_keys:
+            del render_conf.cameras[k]
     if "scene" in kwargs:
         scene_cfg = OmegaConf.load(Path(calvin_env.__file__).parents[1] / "conf/scene" / f"{kwargs['scene']}.yaml")
         OmegaConf.merge(render_conf, scene_cfg)
+    if not hydra.core.global_hydra.GlobalHydra.instance().is_initialized():
+        hydra.initialize(".")
     env = hydra.utils.instantiate(render_conf.env, show_gui=show_gui, use_vr=False, use_scene_info=True)
     return env