[MISC] Picking boxes and Panda links with a MouseSpring (Genesis-Embodied-AI#1443)

Author: gasnica

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -15,6 +15,7 @@
 if TYPE_CHECKING:
     from .rigid_entity import RigidEntity
     from genesis.engine.solvers.rigid.rigid_solver_decomp import RigidSolver
+    from genesis.ext.pyrender.interaction.vec3 import Pose
 
 
 @ti.data_oriented
@@ -591,14 +592,14 @@ def invweight(self):
         return self._invweight
 
     @property
-    def pos(self):
+    def pos(self) -> ArrayLike:
         """
         The initial position of the link. For real-time position, use `link.get_pos()`.
         """
         return self._pos
 
     @property
-    def quat(self):
+    def quat(self) -> ArrayLike:
         """
         The initial quaternion of the link. For real-time quaternion, use `link.get_quat()`.
         """
@@ -744,6 +745,11 @@ def is_free(self):
         """
         return self.entity.is_free
 
+    @property
+    def pose(self) -> "Pose":
+        """Return the current pose of the link (note, this is not necessarily the same as the principal axes frame)."""
+        return Pose.from_link(self)
+
     # ------------------------------------------------------------------------------------
     # -------------------------------------- repr ----------------------------------------
     # ------------------------------------------------------------------------------------

genesis/engine/solvers/base_solver.py

@@ -152,7 +152,7 @@ def gravity(self):
         return self._gravity.to_numpy() if self._gravity is not None else None
 
     @property
-    def entities(self):
+    def entities(self) -> list[Entity]:
         return self._entities
 
     @property

genesis/ext/pyrender/interaction/aabb.py

@@ -1,3 +1,5 @@
+from dataclasses import dataclass
+
 import numpy as np
 from numpy.typing import NDArray
 
@@ -20,6 +22,10 @@ def min(self) -> Vec3:
     def max(self) -> Vec3:
         return Vec3(self.v[1])
 
+    @property
+    def extents(self) -> Vec3:
+        return self.max - self.min
+
     def expand(self, padding: float) -> None:
         self.v[0] -= padding
         self.v[1] += padding
@@ -54,28 +60,37 @@ def raycast(self, ray: Ray) -> RayHit:
         hit_pos = ray.origin + ray.direction * enter
         return RayHit(enter, hit_pos, normal)
 
-    def raycast_obb(self, pose: Pose, ray: Ray) -> RayHit:
-        inv_pose = pose.get_inverse()
-        origin2 = inv_pose.transform_point(ray.origin)
-        direction2 = inv_pose.transform_direction(ray.direction)
-        ray2 = Ray(origin2, direction2)
-        ray_hit = self.raycast(ray2)
-        if ray_hit.is_hit:
-            ray_hit.position = pose.transform_point(ray_hit.position)
-            ray_hit.normal = pose.transform_direction(ray_hit.normal)
-        return ray_hit
-
     def __repr__(self) -> str:
-        return f"Min({self.min.x}, {self.min.y}, {self.min.z}) Max({self.max.x}, {self.max.y}, {self.max.z})"
+        return f"AABB: Min({self.min.x}, {self.min.y}, {self.min.z}) Max({self.max.x}, {self.max.y}, {self.max.z})"
 
     @classmethod
     def from_min_max(cls, min: Vec3, max: Vec3) -> 'AABB':
-        bounds = np.stack((min.v, max.v))
+        bounds = np.stack((min.v, max.v), axis=0)
         return cls(bounds)
 
     @classmethod
-    def from_center_and_size(cls, center: Vec3, size: Vec3) -> 'AABB':
-        min = center - 0.5 * size
-        max = center + 0.5 * size
-        bounds = np.stack((min.v, max.v))
+    def from_center_and_half_extents(cls, center: Vec3, half_extents: Vec3) -> 'AABB':
+        min = center - half_extents
+        max = center + half_extents
+        bounds = np.stack((min.v, max.v), axis=0)
         return cls(bounds)
+
+
+@dataclass
+class OBB():
+    pose: Pose
+    half_extents: Vec3
+
+    def raycast(self, ray: Ray) -> RayHit:
+        origin2 = self.pose.inverse_transform_point(ray.origin)
+        direction2 = self.pose.inverse_transform_direction(ray.direction)
+        ray2 = Ray(origin2, direction2)
+        aabb = AABB.from_center_and_half_extents(Vec3.zero(), self.half_extents)
+        ray_hit = aabb.raycast(ray2)
+        if ray_hit.is_hit:
+            ray_hit.position = self.pose.transform_point(ray_hit.position)
+            ray_hit.normal = self.pose.transform_direction(ray_hit.normal)
+        return ray_hit
+
+    def __repr__(self) -> str:
+        return f"OBB(pose={self.pose}, half_extents={self.half_extents})"

genesis/ext/pyrender/interaction/mouse_spring.py

@@ -15,66 +15,84 @@ def _ensure_torch_imported() -> None:
     import torch
 
 class MouseSpring:
-    def __init__(self):
-        self.held_geom: RigidGeom | None = None
+    def __init__(self) -> None:
+        self.held_link: RigidLink | None = None
         self.held_point_in_local: Vec3 | None = None
         self.prev_control_point: Vec3 | None = None
 
-    def attach(self, picked_entity: RigidEntity, control_point: Vec3):
+    def attach(self, picked_link: RigidLink, control_point: Vec3) -> None:
         # for now, we just pick the first geometry
-        self.held_geom = picked_entity.geoms[0]
-        pose: Pose = Pose.from_geom(self.held_geom)
+        self.held_link = picked_link
+        pose: Pose = Pose.from_link(self.held_link)
         self.held_point_in_local = pose.inverse_transform_point(control_point)
         self.prev_control_point = control_point
 
-    def detach(self):
-        self.held_geom = None
+    def detach(self) -> None:
+        self.held_link = None
 
-    def apply_force(self, control_point: Vec3, delta_time: float):
+    def apply_force(self, control_point: Vec3, delta_time: float) -> None:
         _ensure_torch_imported()
+
+        # note when threaded: apply_force is called before attach!
+        # note2: that was before we added a lock to ViewerInteraction; this migth be fixed now
+        if not self.held_link:
+            return
 
-        # works ok:
-        # delta: Vec3 = control_point - self.prev_control_point
-        # pos = Vec3.from_tensor(self.held_geom.entity.get_pos())
-        # pos = pos + delta
-        # self.held_geom.entity.set_pos(pos.as_tensor())
         self.prev_control_point = control_point
 
         # do simple force on COM only:
-        link: RigidLink = self.held_geom.link
-        link_pos: Vec3 = Vec3.from_tensor(link.get_pos())
+        link: RigidLink = self.held_link
         lin_vel: Vec3 = Vec3.from_tensor(link.get_vel())
         ang_vel: Vec3 = Vec3.from_tensor(link.get_ang())
+        link_pose: Pose = Pose.from_link(link)
+        held_point_in_world: Vec3 = link_pose.transform_point(self.held_point_in_local)
+
+        # note: we should assert earlier that link inertial_pos/quat are not None
+        # todo: verify inertial_pos/quat are stored in local frame
+        link_T_principal: Pose = Pose(Vec3.from_arraylike(link.inertial_pos), Quat.from_arraylike(link.inertial_quat))
+        world_T_principal: Pose = link_pose * link_T_principal
+
+        arm_in_principal: Vec3 = link_T_principal.inverse_transform_point(self.held_point_in_local)   # for non-spherical inertia
+        arm_in_world: Vec3 = world_T_principal.rot * arm_in_principal  # for spherical inertia
 
-        pos_err_v: Vec3 = control_point - link_pos
-        vel_err_v: Vec3 = Vec3.zero() - lin_vel
+        pos_err_v: Vec3 = control_point - held_point_in_world
         inv_mass: float = float(1.0 / link.get_mass() if link.get_mass() > 0.0 else 0.0)
+        inv_spherical_inertia: float = float(1.0 / link.inertial_i[0, 0] if link.inertial_i[0, 0] > 0.0 else 0.0)
 
         inv_dt: float = 1.0 / delta_time
-        # these are temporary values, till we fix an issue with apply_links_external_force.
-        # after fixing it, use tau = damp = 1.0:
         tau: float = MOUSE_SPRING_POSITION_CORRECTION_FACTOR
         damp: float = MOUSE_SPRING_VELOCITY_CORRECTION_FACTOR
 
         total_impulse: Vec3 = Vec3.zero()
+        total_torque_impulse: Vec3 = Vec3.zero()
+
+        for i in range(3*4):
+            body_point_vel: Vec3 = lin_vel + ang_vel.cross(arm_in_world)
+            vel_err_v: Vec3 = Vec3.zero() - body_point_vel
 
-        for i in range(3):
             dir: Vec3 = Vec3.zero()
-            dir.v[i] = 1.0
+            dir.v[i % 3] = 1.0
             pos_err: float = dir.dot(pos_err_v)
             vel_err: float = dir.dot(vel_err_v)
             error: float = tau * pos_err * inv_dt + damp * vel_err
-            virtual_mass: float = 1.0 / (inv_mass + 1e-24)
+
+            arm_x_dir: Vec3 = arm_in_world.cross(dir)
+            virtual_mass: float = 1.0 / (inv_mass + arm_x_dir.sqr_magnitude() * inv_spherical_inertia + 1e-24)
             impulse: float = error * virtual_mass
 
-            lin_vel += impulse * dir * inv_mass
-            total_impulse.v[i] = impulse
+            lin_vel += impulse * inv_mass * dir
+            ang_vel += impulse * inv_spherical_inertia * arm_x_dir
+            total_impulse.v[i % 3] += impulse
+            total_torque_impulse += impulse * arm_x_dir
 
         # Apply the new force
         total_force = total_impulse * inv_dt
+        total_torque = total_torque_impulse * inv_dt
         force_tensor: torch.Tensor = total_force.as_tensor().unsqueeze(0)
+        torque_tensor: torch.Tensor = total_torque.as_tensor().unsqueeze(0)
         link.solver.apply_links_external_force(force_tensor, (link.idx,), ref='link_com', local=False)
+        link.solver.apply_links_external_torque(torque_tensor, (link.idx,), ref='link_com', local=False)
 
     @property
     def is_attached(self) -> bool:
-        return self.held_geom is not None
+        return self.held_link is not None

genesis/ext/pyrender/interaction/ray.py

@@ -1,11 +1,16 @@
 from dataclasses import dataclass
+import sys
+
+from genesis.engine.entities.rigid_entity.rigid_geom import RigidGeom
 
 from .vec3 import Vec3
 
 
 EPSILON = 1e-6
 EPSILON2 = EPSILON * EPSILON
 
+_MAX_RAY_DISTANCE = sys.float_info.max
+
 
 class Ray:
     origin: Vec3
@@ -24,14 +29,16 @@ class RayHit:
     distance: float
     position: Vec3
     normal: Vec3
+    geom: RigidGeom | None = None
 
     @property
     def is_hit(self) -> bool:
-        return 0 <= self.distance
+        assert 0.0 <= self.distance
+        return self.distance < _MAX_RAY_DISTANCE
 
     @classmethod
     def no_hit(cls) -> 'RayHit':
-        return RayHit(-1.0, Vec3.zero(), Vec3.zero())
+        return RayHit(_MAX_RAY_DISTANCE, Vec3.zero(), Vec3.zero(), None)
 
 
 class Plane:
@@ -50,4 +57,4 @@ def raycast(self, ray: Ray) -> RayHit:
             return RayHit.no_hit()
         else:
             dist_along_ray = dist / -dot
-            return RayHit(dist_along_ray, ray.origin + ray.direction * dist_along_ray, self.normal)
+            return RayHit(dist_along_ray, ray.origin + ray.direction * dist_along_ray, self.normal, None)

genesis/ext/pyrender/interaction/vec3.py

@@ -2,12 +2,13 @@
 from typing import TYPE_CHECKING, Union
 
 import numpy as np
-from numpy.typing import NDArray
+from numpy.typing import NDArray, ArrayLike
 
 from genesis.utils.misc import tensor_to_array
 
 if TYPE_CHECKING:
     from genesis.engine.entities.rigid_entity.rigid_geom import RigidGeom
+    from genesis.engine.entities.rigid_entity.rigid_link import RigidLink
 
 # If not needing runtime checks, we can just use annotated types:
 # Vec3 = Annotated[npt.NDArray[np.float32], (3,)]
@@ -62,6 +63,9 @@ def normalized(self) -> 'Vec3':
     def magnitude(self) -> float:
         return np.linalg.norm(self.v)
 
+    def sqr_magnitude(self) -> float:
+        return np.dot(self.v, self.v)
+
     def copy(self) -> 'Vec3':
         return Vec3(self.v.copy())
 
@@ -101,15 +105,30 @@ def from_tensor(cls, v: 'torch.Tensor') -> 'Vec3':
         array: np.ndarray = tensor_to_array(v)
         return cls.from_array(array)
 
+    @classmethod
+    def from_arraylike(cls, v: ArrayLike) -> 'Vec3':
+        if isinstance(v, np.ndarray):
+            return cls.from_array(v)
+        elif isinstance(v, torch.Tensor):
+            return cls.from_tensor(v)
+        elif isinstance(v, ArrayLike):
+            assert len(v) == 3, f"Vec3 must be initialized with a 3-element ArrayLike, got {len(v)}"
+            return cls.from_xyz(*v)
+        assert False
+
 
     @classmethod
-    def zero(cls):
+    def zero(cls) -> 'Vec3':
         return cls(np.array([0, 0, 0], dtype=np.float32))
 
     @classmethod
-    def one(cls):
+    def one(cls) -> 'Vec3':
         return cls(np.array([1, 1, 1], dtype=np.float32))
 
+    @classmethod
+    def full(cls, fill_value: float) -> 'Vec3':
+        return cls(np.full((3,), fill_value, dtype=np.float32))
+
 
 class Quat:
     v: NDArray[np.float32]
@@ -183,6 +202,17 @@ def from_tensor(cls, v: 'torch.Tensor') -> 'Quat':
         array: np.ndarray = tensor_to_array(v)
         return cls.from_array(array)
 
+    @classmethod
+    def from_arraylike(cls, v: ArrayLike) -> 'Quat':
+        if isinstance(v, np.ndarray):
+            return cls.from_array(v)
+        elif isinstance(v, torch.Tensor):
+            return cls.from_tensor(v)
+        elif isinstance(v, ArrayLike):
+            assert len(v) == 4, f"Quat must be initialized with a 4-element ArrayLike, got {len(v)}"
+            return cls.from_wxyz(*v)
+        assert False
+
 
 @dataclass
 class Pose:
@@ -212,6 +242,17 @@ def get_inverse(self) -> 'Pose':
         inv_pos = Vec3(-inv_pos.v[1:])
         return Pose(inv_pos, inv_rot)
 
+    def __mul__(self, other: Union['Pose', Vec3]) -> Union['Pose', Vec3]:
+        if isinstance(other, Pose):
+            return Pose(self.pos + self.rot * other.pos, self.rot * other.rot)
+        elif isinstance(other, Vec3):
+            return self.pos + self.rot * other
+        else:
+            return NotImplemented        
+
+    def __repr__(self) -> str:
+        return f"Pose(pos={self.pos}, rot={self.rot})"
+
     @classmethod
     def from_geom(cls, geom: 'RigidGeom') -> 'Pose':
         assert geom._solver.n_envs == 0, "ViewerInteraction only supports single-env for now"
@@ -220,6 +261,14 @@ def from_geom(cls, geom: 'RigidGeom') -> 'Pose':
         quat = Quat.from_tensor(geom.get_quat())
         return Pose(pos, quat)
 
+    @classmethod
+    def from_link(cls, link: 'RigidLink') -> 'Pose':
+        assert link._solver.n_envs == 0, "ViewerInteraction only supports single-env for now"
+        # geom.get_pos() and .get_quat() are squeezed if n_envs == 0
+        pos = Vec3.from_tensor(link.get_pos())
+        quat = Quat.from_tensor(link.get_quat())
+        return Pose(pos, quat)
+
 
 @dataclass
 class Color: