BUGFIX: derived keys were getting dropped as a result there were no e…

…e_pose in the roboset data
vikashplus · Mar 15, 2023 · 65b63a3 · 65b63a3
1 parent a8bf20f
commit 65b63a3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mj_envs/logger/roboset_logger.py b/mj_envs/logger/roboset_logger.py
@@ -41,8 +41,8 @@ def path2dataset(self, path:dict, config_path=None)->dict:
 
         # Derived =====
         pose_ee = []
-        if 'pos_ee' in path_keys or 'rot_ee' in path_keys:
-            assert ('pos_ee' in path_keys and 'rot_ee' in path_keys), "Both pose_ee and rot_ee are required"
+        if 'env_infos/obs_dict/pos_ee' in path_keys or 'env_infos/obs_dict/rot_ee' in path_keys:
+            assert ('env_infos/obs_dict/pos_ee' in path_keys and 'env_infos/obs_dict/rot_ee' in path_keys), "Both pose_ee and rot_ee are required"
             dataset['derived/pose_ee'] = np.hstack([path['env_infos/obs_dict/pos_ee'], path['env_infos/obs_dict/rot_ee']])
 
         # Config =====