rotation task requires surfaces at the end

mees · Jan 10, 2022 · 141fc11 · 141fc11
1 parent 9b43fb9
commit 141fc11
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 7 deletions.
diff --git a/calvin_env/envs/play_lmp_wrapper.py b/calvin_env/envs/play_lmp_wrapper.py
@@ -38,7 +38,7 @@ def set_egl_device(device):
             logger.warning(
                 "Couldn't find correct EGL device. Setting EGL_VISIBLE_DEVICE=0. "
                 "When using DDP with many GPUs this can lead to OOM errors. "
-                "Did you install PyBullet correctly? Please refer to VREnv README"
+                "Did you install PyBullet correctly? Please refer to calvin env README"
             )
             egl_id = 0
         os.environ["EGL_VISIBLE_DEVICES"] = str(egl_id)

diff --git a/calvin_env/envs/play_table_env.py b/calvin_env/envs/play_table_env.py
@@ -70,7 +70,7 @@ def __init__(
             )
             for name in cameras
         ]
-        log.info(f"Using VREnv with commit {get_git_commit_hash(Path(calvin_env.__file__))}.")
+        log.info(f"Using calvin_env with commit {get_git_commit_hash(Path(calvin_env.__file__))}.")
 
     def __del__(self):
         self.close()

diff --git a/calvin_env/envs/tasks.py b/calvin_env/envs/tasks.py
@@ -54,24 +54,32 @@ def num_tasks(self):
 
     @staticmethod
     def rotate_object(
-        obj_name, z_degrees, x_y_threshold=30, z_threshold=180, movement_threshold=0.1, start_info=None, end_info=None
+            obj_name, z_degrees, x_y_threshold=30, z_threshold=180, movement_threshold=0.1, start_info=None,
+            end_info=None
     ):
         """
         Returns True if the object with obj_name was rotated more than z_degrees degrees around the z-axis while not
         being rotated more than x_y_threshold degrees around the x or y axis.
         z_degrees is negative for clockwise rotations and positive for counter-clockwise rotations.
         """
-        start_orn = R.from_quat(start_info["scene_info"]["movable_objects"][obj_name]["current_orn"])
-        end_orn = R.from_quat(end_info["scene_info"]["movable_objects"][obj_name]["current_orn"])
+        obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
+        obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
+        start_orn = R.from_quat(obj_start_info["current_orn"])
+        end_orn = R.from_quat(obj_end_info["current_orn"])
         rotation = end_orn * start_orn.inv()
         x, y, z = rotation.as_euler("xyz", degrees=True)
 
-        start_pos = np.array(start_info["scene_info"]["movable_objects"][obj_name]["current_pos"])
-        end_pos = np.array(end_info["scene_info"]["movable_objects"][obj_name]["current_pos"])
+        start_pos = np.array(obj_start_info["current_pos"])
+        end_pos = np.array(obj_end_info["current_pos"])
         pos_diff = end_pos - start_pos
         if np.linalg.norm(pos_diff) > movement_threshold:
             return False
 
+        end_contacts = set(c[2] for c in obj_end_info["contacts"])
+        robot_uid = {start_info["robot_info"]["uid"]}
+        if len(end_contacts - robot_uid) == 0:
+            return False
+
         if z_degrees > 0:
             return z_degrees < z < z_threshold and abs(x) < x_y_threshold and abs(y) < x_y_threshold
         else: