[MISC] Add raycasting against box-entities in the scene (Genesis-Embodied-AI#1346)

* Add raycasting against Aabb and Oobb. 
* Add helper classes: Quat, Pose, Aabb, Color
* Highlight hit object
* Update ViewerInteraction to run this logic every frame
* Add partial information on member-field types to RigidGeom, and RigidLink

Author: gasnica

genesis/engine/entities/rigid_entity/rigid_geom.py

@@ -1,5 +1,6 @@
 import os
 import pickle as pkl
+from typing import TYPE_CHECKING
 
 import igl
 import numpy as np
@@ -14,6 +15,12 @@
 from genesis.repr_base import RBC
 from genesis.utils.misc import tensor_to_array
 
+if TYPE_CHECKING:
+    from genesis.engine.materials.rigid import Rigid as RigidMaterial
+
+    from .rigid_entity import RigidEntity
+    from .rigid_link import RigidLink
+
 
 @ti.data_oriented
 class RigidGeom(RBC):
@@ -42,9 +49,9 @@ def __init__(
         center_init=None,
         data=None,
     ):
-        self._link = link
-        self._entity = link.entity
-        self._material = link.entity.material
+        self._link: "RigidLink" = link
+        self._entity: "RigidEntity" = link.entity
+        self._material: "RigidMaterial" = link.entity.material
         self._solver = link.entity.solver
         self._mesh = mesh
 

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -1,4 +1,7 @@
+from typing import TYPE_CHECKING
+
 import numpy as np
+from numpy.typing import ArrayLike
 import taichi as ti
 import torch
 
@@ -9,6 +12,9 @@
 
 from .rigid_geom import RigidGeom, RigidVisGeom
 
+if TYPE_CHECKING:
+    from .rigid_entity import RigidEntity
+
 
 @ti.data_oriented
 class RigidLink(RBC):
@@ -43,8 +49,8 @@ def __init__(
         invweight,
         visualize_contact,
     ):
-        self._name = name
-        self._entity = entity
+        self._name: str = name
+        self._entity: "RigidEntity" = entity
         self._solver = entity.solver
         self._entity_idx_in_solver = entity.idx
 
@@ -68,16 +74,18 @@ def __init__(
         self._vvert_start = vvert_start
         self._vface_start = vface_start
 
-        self._pos = pos
-        self._quat = quat
-        self._inertial_pos = inertial_pos
-        self._inertial_quat = inertial_quat
+        # Link position & rotation at creation time:
+        self._pos: ArrayLike = pos
+        self._quat: ArrayLike = quat
+        # Link's center-of-mass position & principal axes frame rotation at creation time:
+        self._inertial_pos: ArrayLike = inertial_pos
+        self._inertial_quat: ArrayLike = inertial_quat
         self._inertial_mass = inertial_mass
         self._inertial_i = inertial_i
 
         self._visualize_contact = visualize_contact
 
-        self._geoms = gs.List()
+        self._geoms: list[RigidGeom] = gs.List()
         self._vgeoms = gs.List()
 
     def _build(self):
@@ -99,7 +107,7 @@ def _build(self):
                 is_fixed = False
         if self._root_idx is None:
             self._root_idx = gs.np_int(link.idx)
-        self.is_fixed = gs.np_int(is_fixed)
+        self.is_fixed: np.int32 = gs.np_int(is_fixed)  # note: type inconsistent with is_built, is_free, is_leaf: bool
 
         # inertial_mass and inertia_i
         if self._inertial_mass is None:

genesis/ext/pyrender/interaction/aabb.py

@@ -0,0 +1,81 @@
+import numpy as np
+from numpy.typing import NDArray
+
+from .ray import Ray, RayHit, EPSILON
+from .vec3 import Pose, Vec3
+
+class AABB:
+    v: NDArray[np.float32]
+
+    def __init__(self, v: NDArray[np.float32]):
+        assert v.shape == (2, 3,), f"Aabb must be initialized with a (2,3)-element array, got {v.shape}"
+        assert v.dtype == np.float32, f"Aabb must be initialized with a float32 array, got {v.dtype}"
+        self.v = v
+
+    @property
+    def min(self) -> Vec3:
+        return Vec3(self.v[0])
+
+    @property
+    def max(self) -> Vec3:
+        return Vec3(self.v[1])
+
+    def expand(self, padding: float) -> None:
+        self.v[0] -= padding
+        self.v[1] += padding
+
+    def raycast(self, ray: Ray) -> RayHit:
+        """
+        Standard AABB slab implementation. Early-exits and returns no-hit for rays withing the XY, XZ, or YZ planes.
+        Ignores hits for rays originating inside the AABB.
+        """
+        if (np.abs(ray.direction.v) < EPSILON).any():
+            # unhandled ray case: early-exit
+            return RayHit.no_hit()
+
+        tmin = (self.v[0] - ray.origin.v) / ray.direction.v
+        tmax = (self.v[1] - ray.origin.v) / ray.direction.v
+        mmin = np.minimum(tmin, tmax)
+        mmax = np.maximum(tmin, tmax)
+        min_idx = np.argmax(mmin)
+        max_idx = np.argmin(mmax)
+        tnear = mmin[min_idx]
+        tfar = mmax[max_idx]
+
+        # Drop hits coming from inside
+        if tfar < tnear or tnear < 0:  # tfar < 0
+            return RayHit.no_hit()
+
+        # Calculate enter point and normal
+        enter = tnear  # if 0 <= tnear else tfar
+        normal = Vec3.zero()
+        normal.v[min_idx] = -np.sign(ray.direction.v[min_idx])
+
+        hit_pos = ray.origin + ray.direction * enter
+        return RayHit(enter, hit_pos, normal)
+
+    def raycast_oobb(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})"
+
+    @classmethod
+    def from_min_max(cls, min: Vec3, max: Vec3) -> 'AABB':
+        bounds = np.stack((min.v, max.v))
+        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))
+        return cls(bounds)

genesis/ext/pyrender/interaction/ray.py

@@ -4,6 +4,7 @@
 
 
 EPSILON = 1e-6
+EPSILON2 = EPSILON * EPSILON
 
 
 class Ray:
@@ -20,11 +21,17 @@ def __repr__(self) -> str:
 
 @dataclass
 class RayHit:
-    is_hit: bool
     distance: float
-    normal: Vec3
     position: Vec3
-    object_idx: int
+    normal: Vec3
+
+    @property
+    def is_hit(self) -> bool:
+        return 0 <= self.distance
+
+    @classmethod
+    def no_hit(cls) -> 'RayHit':
+        return RayHit(-1.0, Vec3.zero(), Vec3.zero())
 
 
 class Plane:
@@ -40,15 +47,7 @@ def raycast(self, ray: Ray) -> RayHit:
         dist = ray.origin.dot(self.normal) + self.distance
 
         if -EPSILON < dot or dist < EPSILON:
-            return RayHit(is_hit=False, distance=0, normal=Vec3.zero(), position=Vec3.zero(), object_idx=-1)
-
-        dist_along_ray = dist / -dot
-
-        return RayHit(
-            is_hit=True,
-            distance=dist_along_ray,
-            normal=self.normal,
-            position=ray.origin + ray.direction * dist_along_ray,
-            object_idx=0
-        )
-
+            return RayHit.no_hit()
+        else:
+            dist_along_ray = dist / -dot
+            return RayHit(dist_along_ray, ray.origin + ray.direction * dist_along_ray, self.normal)

genesis/ext/pyrender/interaction/vec3.py

@@ -1,6 +1,13 @@
+from dataclasses import dataclass
+from typing import Union
+
 import numpy as np
 from numpy.typing import NDArray
 
+# If not needing runtime checks, we can just use annotated types:
+# Vec3 = Annotated[npt.NDArray[np.float32], (3,)]
+# Aabb = Annotated[npt.NDArray[np.float32], (2, 3)]
+
 
 class Vec3:
     """
@@ -43,6 +50,17 @@ def copy(self) -> 'Vec3':
     def __repr__(self) -> str:
         return f"Vec3({self.v[0]}, {self.v[1]}, {self.v[2]})"
 
+    @property
+    def x(self) -> float:
+        return self.v[0]
+
+    @property
+    def y(self) -> float:
+        return self.v[1]
+
+    @property
+    def z(self) -> float:
+        return self.v[2]
 
     @classmethod
     def from_xyz(cls, x: float, y: float, z: float) -> 'Vec3':
@@ -66,6 +84,11 @@ def from_float64(cls, v: NDArray[np.float64]) -> 'Vec3':
         assert v.dtype == np.float64, f"from_float64 must be initialized with a float64 array, got {v.dtype}"
         return cls.from_xyz(*v)
 
+    @classmethod
+    def from_any_array(cls, v: np.ndarray) -> 'Vec3':
+        assert v.shape == (3,), f"Vec3 must be initialized with a 3-element array, got {v.shape}"
+        return cls.from_xyz(*v)
+
 
     @classmethod
     def zero(cls):
@@ -74,3 +97,126 @@ def zero(cls):
     @classmethod
     def one(cls):
         return cls(np.array([1, 1, 1], dtype=np.float32))
+
+
+class Quat:
+    v: NDArray[np.float32]
+    def __init__(self, v: NDArray[np.float32]):
+        assert v.shape == (4,), f"Quat must be initialized with a 4-element array, got {v.shape}"
+        assert v.dtype == np.float32, f"Quat must be initialized with a float32 array, got {v.dtype}"
+        self.v = v
+
+    def get_inverse(self) -> 'Quat':
+        quat_inv = self.v.copy()
+        quat_inv[1:] *= -1
+        return Quat(quat_inv)
+
+    def __mul__(self, other: Union['Quat', Vec3]) -> Union['Quat', Vec3]:
+        if isinstance(other, Quat):
+            # Quaternion * Quaternion
+            w1, x1, y1, z1 = self.w, self.x, self.y, self.z
+            w2, x2, y2, z2 = other.w, other.x, other.y, other.z
+            return Quat.from_wxyz(
+                w1*w2 - x1*x2 - y1*y2 - z1*z2,
+                w1*x2 + x1*w2 + y1*z2 - z1*y2,
+                w1*y2 - x1*z2 + y1*w2 + z1*x2,
+                w1*z2 + x1*y2 - y1*x2 + z1*w2
+            )
+        elif isinstance(other, Vec3):  # (other, np.ndarray) and other.shape == (3,):
+            # Quaternion * Vector3 -> rotate vector
+            v_quat = Quat.from_wxyz(0, *other.v)
+            result = self * v_quat * self.get_inverse()
+            return Vec3(result.v[1:])
+        else:
+            return NotImplemented
+
+    def copy(self) -> 'Quat':
+        return Quat(self.v.copy())
+
+    def __repr__(self) -> str:
+        return f"Quat({self.v[0]}, {self.v[1]}, {self.v[2]}, {self.v[3]})"
+
+    @property
+    def w(self) -> float:
+        return self.v[0]
+
+    @property
+    def x(self) -> float:
+        return self.v[1]
+
+    @property
+    def y(self) -> float:
+        return self.v[2]
+
+    @property
+    def z(self) -> float:
+        return self.v[3]
+
+
+    @classmethod
+    def from_wxyz(cls, w: float, x: float, y: float, z: float) -> 'Quat':
+        return cls(np.array([w, x, y, z], dtype=np.float32))
+
+    @classmethod
+    def from_any_array(cls, v: np.ndarray) -> 'Quat':
+        assert v.shape == (4,), f"Quat must be initialized with a 4-element array, got {v.shape}"
+        return cls.from_wxyz(*v)
+
+
+@dataclass
+class Pose:
+    pos: Vec3
+    rot: Quat
+
+    # todo: consider using a single np.array with views
+
+    def transform_point(self, point: Vec3) -> Vec3:
+        return self.pos + self.rot * point
+
+    def inverse_transform_point(self, point: Vec3) -> Vec3:
+        return self.rot.get_inverse() * (point - self.pos)
+
+    def transform_direction(self, direction: Vec3) -> Vec3:
+        return self.rot * direction
+
+    def inverse_transform_direction(self, direction: Vec3) -> Vec3:
+        return self.rot.get_inverse() * direction
+
+    def get_inverse(self) -> 'Pose':
+        inv_rot = self.rot.get_inverse()
+        # inv_pos = -1.0 * (inv_rot * self.pos)
+        # faster -- avoid repeated quat inversion:
+        pos_quat = Quat.from_wxyz(0, *self.pos.v)
+        inv_pos = inv_rot * pos_quat * self.rot
+        inv_pos = Vec3(-inv_pos.v[1:])
+        return Pose(inv_pos, inv_rot)
+
+
+@dataclass
+class Color:
+    r: float
+    g: float
+    b: float
+    a: float
+
+    def tuple(self) -> tuple[float, float, float, float]:
+        return (self.r, self.g, self.b, self.a)
+
+    def with_alpha(self, alpha: float) -> 'Color':
+        return Color(self.r, self.g, self.b, alpha)
+
+    @classmethod
+    def red(cls) -> 'Color':
+        return cls(1.0, 0.0, 0.0, 1.0)
+
+    @classmethod
+    def green(cls) -> 'Color':
+        return cls(0.0, 1.0, 0.0, 1.0)
+
+    @classmethod
+    def blue(cls) -> 'Color':
+        return cls(0.0, 0.0, 1.0, 1.0)
+
+    @classmethod
+    def yellow(cls) -> 'Color':
+        return cls(1.0, 1.0, 0.0, 1.0)