fix a few envs, mark pick clutter as to be fixed in the future.

mani_skill2/agents/controllers/base_controller.py

@@ -284,7 +284,7 @@ def set_action(self, action: np.ndarray):
 
         for uid, controller in self.controllers.items():
             start, end = self.action_mapping[uid]
-            controller.set_action(to_tensor(action[..., start:end]))
+            controller.set_action(action[..., start:end])
 
     def to_action_dict(self, action: np.ndarray):
         """Convert a flat action to a dict of actions."""

mani_skill2/agents/controllers/pd_ee_pose.py

@@ -4,6 +4,7 @@
 import numpy as np
 import sapien
 import sapien.physx as physx
+import torch
 from gymnasium import spaces
 from scipy.spatial.transform import Rotation
 
@@ -167,12 +168,13 @@ def _initialize_action_space(self):
         self.action_space = spaces.Box(low, high, dtype=np.float32)
 
     def _clip_and_scale_action(self, action):
+        # TODO (stao): support batched actions
         # NOTE(xiqiang): rotation should be clipped by norm.
         pos_action = clip_and_scale_action(
-            action[:3], self._action_space.low[:3], self._action_space.high[:3]
+            action[:3], self.action_space_low[:3], self.action_space_high[:3]
         )
         rot_action = action[3:]
-        rot_norm = np.linalg.norm(rot_action)
+        rot_norm = torch.linalg.norm(rot_action)
         if rot_norm > 1:
             rot_action = rot_action / rot_norm
         rot_action = rot_action * self.config.rot_bound

mani_skill2/agents/controllers/pd_joint_vel.py

@@ -29,8 +29,7 @@ def set_drive_property(self):
 
     def set_action(self, action: np.ndarray):
         action = self._preprocess_action(action)
-        for i, joint in enumerate(self.joints):
-            joint.set_drive_velocity_target(action[i])
+        self.articulation.set_joint_drive_velocity_targets(action, self.joints)
 
 
 @dataclass

mani_skill2/agents/robots/xmate3/xmate3.py

@@ -1,18 +1,20 @@
 import numpy as np
 import sapien
 import sapien.physx as physx
+import torch
 
 from mani_skill2 import ASSET_DIR
 from mani_skill2.agents.base_agent import BaseAgent
 from mani_skill2.agents.controllers import *
 from mani_skill2.sensors.camera import CameraConfig
-from mani_skill2.utils.common import np_compute_angle_between
+from mani_skill2.utils.common import compute_angle_between, np_compute_angle_between
 from mani_skill2.utils.sapien_utils import (
     compute_total_impulse,
     get_actor_contacts,
     get_obj_by_name,
     get_pairwise_contact_impulse,
 )
+from mani_skill2.utils.structs.actor import Actor
 
 
 class Xmate3Robotiq(BaseAgent):
@@ -172,38 +174,83 @@ def controller_configs(self):
         ),
     ]
 
-    def is_grasping(self, object: sapien.Entity = None, min_impulse=1e-6, max_angle=85):
-        contacts = self.scene.get_contacts()
-        if object is None:
-            finger1_contacts = get_actor_contacts(contacts, self.finger1_link)
-            finger2_contacts = get_actor_contacts(contacts, self.finger2_link)
-            return (
-                np.linalg.norm(compute_total_impulse(finger1_contacts)) >= min_impulse
-                and np.linalg.norm(compute_total_impulse(finger2_contacts))
-                >= min_impulse
+    def is_grasping(self, object: Actor = None, min_impulse=1e-6, max_angle=85):
+        # TODO (stao): is_grasping code needs to be updated for new GPU sim
+        if physx.is_gpu_enabled():
+            if object.name not in self.queries:
+                body_pairs = list(zip(self.finger1_link._bodies, object._bodies))
+                body_pairs += list(zip(self.finger2_link._bodies, object._bodies))
+                self.queries[object.name] = (
+                    self.scene.px.gpu_create_contact_query(body_pairs),
+                    (len(object._bodies), 3),
+                )
+                print(f"Create query for Panda grasp({object.name})")
+            query, contacts_shape = self.queries[object.name]
+            self.scene.px.gpu_query_contacts(query)
+            # query.cuda_contacts # (num_unique_pairs * num_envs, 3)
+            contacts = query.cuda_contacts.clone().reshape((-1, *contacts_shape))
+            lforce = torch.linalg.norm(contacts[0], axis=1)
+            rforce = torch.linalg.norm(contacts[1], axis=1)
+
+            # NOTE (stao): 0.5 * time_step is a decent value when tested on a pick cube task.
+            min_force = 0.5 * self.scene.px.timestep
+
+            # direction to open the gripper
+            ldirection = self.finger1_link.pose.to_transformation_matrix()[..., :3, 1]
+            rdirection = -self.finger2_link.pose.to_transformation_matrix()[..., :3, 1]
+            langle = compute_angle_between(ldirection, contacts[0])
+            rangle = compute_angle_between(rdirection, contacts[1])
+            lflag = torch.logical_and(
+                lforce >= min_force, torch.rad2deg(langle) <= max_angle
             )
+            rflag = torch.logical_and(
+                rforce >= min_force, torch.rad2deg(rangle) <= max_angle
+            )
+
+            return torch.logical_and(lflag, rflag)
         else:
-            limpulse = get_pairwise_contact_impulse(contacts, self.finger1_link, object)
-            rimpulse = get_pairwise_contact_impulse(contacts, self.finger2_link, object)
+            contacts = self.scene.get_contacts()
 
-            # direction to open the gripper
-            ldirection = self.finger1_link.pose.to_transformation_matrix()[:3, 1]
-            rdirection = -self.finger2_link.pose.to_transformation_matrix()[:3, 1]
+            if object is None:
+                finger1_contacts = get_actor_contacts(contacts, self.finger1_link)
+                finger2_contacts = get_actor_contacts(contacts, self.finger2_link)
+                return (
+                    np.linalg.norm(compute_total_impulse(finger1_contacts))
+                    >= min_impulse
+                    and np.linalg.norm(compute_total_impulse(finger2_contacts))
+                    >= min_impulse
+                )
+            else:
+                limpulse = get_pairwise_contact_impulse(
+                    contacts, self.finger1_link, object
+                )
+                rimpulse = get_pairwise_contact_impulse(
+                    contacts, self.finger2_link, object
+                )
 
-            # angle between impulse and open direction
-            langle = np_compute_angle_between(ldirection, limpulse)
-            rangle = np_compute_angle_between(rdirection, rimpulse)
+                # direction to open the gripper
+                ldirection = self.finger1_link.pose.to_transformation_matrix()[
+                    ..., :3, 1
+                ]
+                rdirection = -self.finger2_link.pose.to_transformation_matrix()[
+                    ..., :3, 1
+                ]
 
-            lflag = (
-                np.linalg.norm(limpulse) >= min_impulse
-                and np.rad2deg(langle) <= max_angle
-            )
-            rflag = (
-                np.linalg.norm(rimpulse) >= min_impulse
-                and np.rad2deg(rangle) <= max_angle
-            )
+                # TODO Convert this to batched code
+                # angle between impulse and open direction
+                langle = np_compute_angle_between(ldirection[0], limpulse)
+                rangle = np_compute_angle_between(rdirection[0], rimpulse)
+
+                lflag = (
+                    np.linalg.norm(limpulse) >= min_impulse
+                    and np.rad2deg(langle) <= max_angle
+                )
+                rflag = (
+                    np.linalg.norm(rimpulse) >= min_impulse
+                    and np.rad2deg(rangle) <= max_angle
+                )
 
-            return all([lflag, rflag])
+                return all([lflag, rflag])
 
     @staticmethod
     def build_grasp_pose(approaching, closing, center):

mani_skill2/envs/assembly/peg_insertion_side.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import sapien
+import torch
 from sapien import Pose
 from transforms3d.euler import euler2quat
 
@@ -58,7 +59,8 @@ def _build_box_with_hole(
         return builder.build_static(name)
 
     def _load_actors(self):
-        TableSceneBuilder().build(self._scene)
+        self.table_scene = TableSceneBuilder(env=self)
+        self.table_scene.build()
 
         # peg
         # length, radius = 0.1, 0.02
@@ -105,6 +107,7 @@ def _load_actors(self):
         self.box_hole_radius = inner_radius
 
     def _initialize_actors(self):
+        self.table_scene.initialize()
         xy = self._episode_rng.uniform([-0.1, -0.3], [0.1, 0])
         pos = np.hstack([xy, self.peg_half_size[2]])
         ori = np.pi / 2 + self._episode_rng.uniform(-np.pi / 3, np.pi / 3)
@@ -167,14 +170,17 @@ def has_peg_inserted(self):
         # Only head position is used in fact
         peg_head_pos_at_hole = (self.box_hole_pose.inv() * self.peg_head_pose).p
         # x-axis is hole direction
-        x_flag = -0.015 <= peg_head_pos_at_hole[0]
+        x_flag = -0.015 <= peg_head_pos_at_hole[:, 0]
         y_flag = (
-            -self.box_hole_radius <= peg_head_pos_at_hole[1] <= self.box_hole_radius
+            -self.box_hole_radius <= peg_head_pos_at_hole[:, 1] <= self.box_hole_radius
         )
         z_flag = (
-            -self.box_hole_radius <= peg_head_pos_at_hole[2] <= self.box_hole_radius
+            -self.box_hole_radius <= peg_head_pos_at_hole[:, 2] <= self.box_hole_radius
+        )
+        return (
+            torch.logical_and(torch.logical_and(x_flag, y_flag), z_flag),
+            peg_head_pos_at_hole,
         )
-        return (x_flag and y_flag and z_flag), peg_head_pos_at_hole
 
     def evaluate(self, **kwargs) -> dict:
         success, peg_head_pos_at_hole = self.has_peg_inserted()

mani_skill2/envs/pick_and_place/__init__.py

@@ -2,4 +2,3 @@
 from . import pick_cube
 from . import stack_cube
 from . import pick_single
-from . import pick_clutter

mani_skill2/envs/pick_and_place/base_env.py

@@ -12,7 +12,7 @@
 from mani_skill2.agents.robots.xmate3 import Xmate3Robotiq
 from mani_skill2.envs.sapien_env import BaseEnv
 from mani_skill2.sensors.camera import CameraConfig
-from mani_skill2.utils.sapien_utils import hide_entity, look_at
+from mani_skill2.utils.sapien_utils import hide_entity, look_at, to_numpy
 from mani_skill2.utils.structs.pose import vectorize_pose
 
 
@@ -34,7 +34,7 @@ def _build_cube(
         if render_material is None:
             render_material = self._renderer.create_material()
             render_material.base_color = (*color, 1.0)
-
+        half_size = to_numpy(half_size)
         builder = self._scene.create_actor_builder()
         builder.add_box_collision(half_size=half_size)
         builder.add_box_visual(half_size=half_size, material=render_material)
@@ -123,7 +123,7 @@ def _initialize_agent_v1(self):
             raise NotImplementedError(self.robot_uid)
 
     def _register_sensors(self):
-        pose = look_at([0.3, 0, 0.6], [-0.1, 0, 0.1])
+        pose = look_at([0.3, 0, 0.6], [-0.1, 0, 0.9])
         return CameraConfig(
             "base_camera", pose.p, pose.q, 128, 128, np.pi / 2, 0.01, 10
         )