[MISC] Add viewer code to handle moving objects (Genesis-Embodied-AI#1378)

* Add simple box-moving interaction.
* Lazy load tensor in vec3.py, add Vec3/Quat.from_tensor/from_array.
* Rename oobb to obb.
* Vec3/Quat.as_tensor() is now a method.

Author: gasnica

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -55,24 +55,24 @@ def __init__(
         self._entity_idx_in_solver = entity.idx
 
         self._uid = gs.UID()
-        self._idx = idx
-        self._parent_idx = parent_idx
-        self._root_idx = root_idx
-        self._child_idxs = list()
-        self._invweight = invweight
-
-        self._joint_start = joint_start
-        self._n_joints = n_joints
-
-        self._geom_start = geom_start
-        self._cell_start = cell_start
-        self._vert_start = vert_start
-        self._face_start = face_start
-        self._edge_start = edge_start
-        self._verts_state_start = verts_state_start
-        self._vgeom_start = vgeom_start
-        self._vvert_start = vvert_start
-        self._vface_start = vface_start
+        self._idx: int = idx
+        self._parent_idx: int = parent_idx  # -1 if no parent
+        self._root_idx: int | None = root_idx  # None if no root
+        self._child_idxs: list[int] = list()
+        self._invweight: float | None = invweight
+
+        self._joint_start: int = joint_start
+        self._n_joints: int = n_joints
+
+        self._geom_start: int = geom_start
+        self._cell_start: int = cell_start
+        self._vert_start: int = vert_start
+        self._face_start: int = face_start
+        self._edge_start: int = edge_start
+        self._verts_state_start: int = verts_state_start
+        self._vgeom_start: int = vgeom_start
+        self._vvert_start: int = vvert_start
+        self._vface_start: int = vface_start
 
         # Link position & rotation at creation time:
         self._pos: ArrayLike = pos

genesis/ext/pyrender/interaction/aabb.py

@@ -54,7 +54,7 @@ def raycast(self, ray: Ray) -> RayHit:
         hit_pos = ray.origin + ray.direction * enter
         return RayHit(enter, hit_pos, normal)
 
-    def raycast_oobb(self, pose: Pose, ray: Ray) -> RayHit:
+    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)

genesis/ext/pyrender/interaction/vec3.py

@@ -4,11 +4,20 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from genesis.utils.misc import tensor_to_array
+
 # 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)]
 
 
+def _ensure_torch_imported() -> None:
+    global gs
+    import genesis as gs
+    global torch
+    import torch
+
+
 class Vec3:
     """
     Use this wrapper around np.array if you want to ensure adherence to float32 arithmethic
@@ -35,6 +44,9 @@ def __mul__(self, other: float) -> 'Vec3':
     def __rmul__(self, other: float) -> 'Vec3':
         return Vec3(self.v * np.float32(other))
 
+    def __neg__(self) -> 'Vec3':
+        return Vec3(-self.v)
+
     def dot(self, other: 'Vec3') -> float:
         return np.dot(self.v, other.v).item()
 
@@ -44,12 +56,19 @@ def cross(self, other: 'Vec3') -> 'Vec3':
     def normalized(self) -> 'Vec3':
         return Vec3(self.v / (np.linalg.norm(self.v) + 1e-24))
 
+    def magnitude(self) -> float:
+        return np.linalg.norm(self.v)
+
     def copy(self) -> 'Vec3':
         return Vec3(self.v.copy())
 
     def __repr__(self) -> str:
         return f"Vec3({self.v[0]}, {self.v[1]}, {self.v[2]})"
 
+    def as_tensor(self) -> 'torch.Tensor':
+        _ensure_torch_imported()
+        return torch.tensor(self.v, dtype=gs.tc_float) 
+
     @property
     def x(self) -> float:
         return self.v[0]
@@ -67,27 +86,17 @@ def from_xyz(cls, x: float, y: float, z: float) -> 'Vec3':
         return cls(np.array([x, y, z], dtype=np.float32))
 
     @classmethod
-    def from_int32(cls, v: NDArray[np.int32]) -> 'Vec3':
-        assert v.shape == (3,), f"Vec3 must be initialized with a 3-element array, got {v.shape}"
-        assert v.dtype == np.int32, f"from_int32 must be initialized with a int32 array, got {v.dtype}"
-        return cls.from_xyz(*v)
-
-    @classmethod
-    def from_int64(cls, v: NDArray[np.int64]) -> 'Vec3':
+    def from_array(cls, v: np.ndarray) -> 'Vec3':
         assert v.shape == (3,), f"Vec3 must be initialized with a 3-element array, got {v.shape}"
-        assert v.dtype == np.int64, f"from_int64 must be initialized with a int64 array, got {v.dtype}"
+        assert v.dtype == np.int32 or v.dtype == np.int64 or v.dtype == np.float32 or v.dtype == np.float64, \
+            f"from_array must be initialized with a array of ints/floats 32/64-bit, got {v.dtype}"
         return cls.from_xyz(*v)
 
     @classmethod
-    def from_float64(cls, v: NDArray[np.float64]) -> 'Vec3':
-        assert v.shape == (3,), f"Vec3 must be initialized with a 3-element array, got {v.shape}"
-        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)
+    def from_tensor(cls, v: 'torch.Tensor') -> 'Vec3':
+        _ensure_torch_imported()
+        array: np.ndarray = tensor_to_array(v)
+        return cls.from_array(array)
 
 
     @classmethod
@@ -136,6 +145,10 @@ def copy(self) -> 'Quat':
     def __repr__(self) -> str:
         return f"Quat({self.v[0]}, {self.v[1]}, {self.v[2]}, {self.v[3]})"
 
+    def as_tensor(self) -> 'torch.Tensor':
+        _ensure_torch_imported()
+        return torch.tensor(self.v, dtype=gs.tc_float) 
+
     @property
     def w(self) -> float:
         return self.v[0]
@@ -152,16 +165,21 @@ def y(self) -> float:
     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':
+    def from_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)
 
+    @classmethod
+    def from_tensor(cls, v: 'torch.Tensor') -> 'Quat':
+        _ensure_torch_imported()
+        array: np.ndarray = tensor_to_array(v)
+        return cls.from_array(array)
+
 
 @dataclass
 class Pose:

genesis/ext/pyrender/interaction/viewer_interaction.py

@@ -2,15 +2,14 @@
 from typing_extensions import override
 
 import numpy as np
-from numpy.typing import NDArray
 
 import genesis as gs
 from genesis.engine.entities.rigid_entity.rigid_entity import RigidEntity
 
 from .aabb import AABB
 from .ray import Plane, Ray, RayHit
 from .vec3 import Pose, Quat, Vec3, Color
-from .viewer_interaction_base import ViewerInteractionBase, EVENT_HANDLE_STATE
+from .viewer_interaction_base import ViewerInteractionBase, EVENT_HANDLE_STATE, EVENT_HANDLED
 
 if TYPE_CHECKING:
     from genesis.engine.entities.rigid_entity.rigid_geom import RigidGeom
@@ -24,14 +23,6 @@ class ViewerInteraction(ViewerInteractionBase):
     - mouse dragging
     """
 
-    camera: 'Node'
-    scene: 'Scene'
-    viewport_size: tuple[int, int]
-    camera_yfov: float
-
-    tan_half_fov: float
-    prev_mouse_pos: tuple[int, int]
-
     def __init__(self, 
         camera: 'Node', 
         scene: 'Scene', 
@@ -41,13 +32,18 @@ def __init__(self,
         camera_fov: float = 60.0,
     ):
         super().__init__(log_events)
-        self.camera = camera
-        self.scene = scene
-        self.viewport_size = viewport_size
-        self.camera_yfov = camera_yfov
+        self.camera: 'Node' = camera
+        self.scene: 'Scene' = scene
+        self.viewport_size: tuple[int, int] = viewport_size
+        self.camera_yfov: float = camera_yfov
 
-        self.tan_half_fov = np.tan(0.5 * self.camera_yfov)
-        self.prev_mouse_pos = tuple(np.array(viewport_size) / 2)
+        self.tan_half_fov: float = np.tan(0.5 * self.camera_yfov)
+        self.prev_mouse_pos: tuple[int, int] = (0.5 * viewport_size[0], 0.5 * viewport_size[1])
+
+        self.picked_entity: RigidEntity | None = None
+        self.picked_point_in_local: Vec3 | None = None
+        self.mouse_drag_plane: Plane | None = None
+        self.prev_mouse_3d_pos: Vec3 | None = None
 
     @override
     def on_mouse_motion(self, x: int, y: int, dx: int, dy: int) -> EVENT_HANDLE_STATE:
@@ -58,17 +54,55 @@ def on_mouse_motion(self, x: int, y: int, dx: int, dy: int) -> EVENT_HANDLE_STAT
     def on_mouse_drag(self, x: int, y: int, dx: int, dy: int, buttons: int, modifiers: int) -> EVENT_HANDLE_STATE:
         super().on_mouse_drag(x, y, dx, dy, buttons, modifiers)
         self.prev_mouse_pos = (x, y)
+        if self.picked_entity:
+            mouse_ray: Ray = self.screen_position_to_ray(x, y)
+            ray_hit: RayHit = self.mouse_drag_plane.raycast(mouse_ray)
+            assert ray_hit.is_hit
+            if ray_hit.is_hit:
+                new_mouse_3d_pos: Vec3 = ray_hit.position
+                delta_3d_pos: Vec3 = new_mouse_3d_pos - self.prev_mouse_3d_pos
+                self.prev_mouse_3d_pos = new_mouse_3d_pos
+
+                use_force: bool = False
+                if use_force:
+                    # apply force
+                    pass
+                else:
+                    #apply displacement
+                    pos = Vec3.from_tensor(self.picked_entity.get_pos())
+                    pos = pos + delta_3d_pos
+                    self.picked_entity.set_pos(pos.as_tensor())
+
+            return EVENT_HANDLED
 
     @override
     def on_mouse_press(self, x: int, y: int, button: int, modifiers: int) -> EVENT_HANDLE_STATE:
         super().on_mouse_press(x, y, button, modifiers)
+        if button == 1: # left mouse button
+            (ray_hit, self.picked_entity) = self.raycast_against_entities(self.screen_position_to_ray(x, y))
+            if self.picked_entity and ray_hit:
+                temp_fwd = self.get_camera_forward()
+                temp_back = -temp_fwd
+
+                self.mouse_drag_plane = Plane(temp_back, ray_hit.position)
+                self.prev_mouse_3d_pos = ray_hit.position
+
+                pose: Pose = self.get_pose_of_first_geom(self.picked_entity)
+                self.picked_point_in_local = pose.inverse_transform_point(ray_hit.position)
+
+    @override
+    def on_mouse_release(self, x: int, y: int, button: int, modifiers: int) -> EVENT_HANDLE_STATE:
+        super().on_mouse_release(x, y, button, modifiers)
+        if button == 1: # left mouse button
+            self.picked_entity = None
+            self.picked_point_in_local = None
 
     @override
     def on_draw(self) -> None:
         super().on_draw()
         if self.scene._visualizer is not None and self.scene._visualizer.viewer_lock is not None:
             self.scene.clear_debug_objects()
-            mouse_ray = self.screen_position_to_ray(*self.prev_mouse_pos)
+            mouse_ray: Ray = self.screen_position_to_ray(*self.prev_mouse_pos)
             closest_hit = None
             hit_entity: RigidEntity | None = None
 
@@ -77,7 +111,7 @@ def on_draw(self) -> None:
                 closest_hit = ray_hit
 
             for entity in self.get_entities():
-                ray_hit = self.raycast_against_entity_oobb(entity, mouse_ray)
+                ray_hit = self.raycast_against_entity_obb(entity, mouse_ray)
                 if ray_hit.is_hit:
                     if closest_hit is None or ray_hit.distance < closest_hit.distance:
                         closest_hit = ray_hit
@@ -87,7 +121,21 @@ def on_draw(self) -> None:
                 self.scene.draw_debug_sphere(closest_hit.position.v, 0.01, (0, 1, 0, 1))
                 self._draw_arrow(closest_hit.position, 0.25 * closest_hit.normal, (0, 1, 0, 1))
             if hit_entity:
-                self._draw_entity_unrotated_oobb(hit_entity)
+                self._draw_entity_unrotated_obb(hit_entity)
+
+            if self.picked_entity:
+                assert self.mouse_drag_plane is not None
+                assert self.picked_point_in_local is not None
+
+                # draw held point
+                pose: Pose = self.get_pose_of_first_geom(self.picked_entity)
+                held_point: Vec3 = pose.transform_point(self.picked_point_in_local)
+                self.scene.draw_debug_sphere(held_point.v, 0.02, Color.red().tuple())
+
+                plane_hit: RayHit = self.mouse_drag_plane.raycast(mouse_ray)
+                if plane_hit.is_hit:
+                    self.scene.draw_debug_sphere(plane_hit.position.v, 0.02, Color.red().tuple())
+                    self.scene.draw_debug_line(held_point.v, plane_hit.position.v, color=Color.red().tuple())
 
     def screen_position_to_ray(self, x: float, y: float) -> Ray:
         # convert screen position to ray
@@ -107,35 +155,49 @@ def screen_position_to_ray(self, x: float, y: float) -> Ray:
         # Note: ignoring pixel aspect ratio
 
         mtx = self.camera.matrix
-        position = Vec3.from_float64(mtx[:3, 3])
-        forward = Vec3.from_float64(-mtx[:3, 2])
-        right = Vec3.from_float64(mtx[:3, 0])
-        up = Vec3.from_float64(mtx[:3, 1])
+        position = Vec3.from_array(mtx[:3, 3])
+        forward = Vec3.from_array(-mtx[:3, 2])
+        right = Vec3.from_array(mtx[:3, 0])
+        up = Vec3.from_array(mtx[:3, 1])
 
         direction = forward + right * x + up * y
         return Ray(position, direction)
 
+    def get_camera_forward(self) -> Vec3:
+        mtx = self.camera.matrix
+        return Vec3.from_array(-mtx[:3, 2])
+
     def get_camera_ray(self) -> Ray:
         mtx = self.camera.matrix
-        position = Vec3.from_float64(mtx[:3, 3])
-        forward = Vec3.from_float64(-mtx[:3, 2])
+        position = Vec3.from_array(mtx[:3, 3])
+        forward = Vec3.from_array(-mtx[:3, 2])
         return Ray(position, forward)
 
     def _raycast_against_ground_plane(self, ray: Ray) -> RayHit:
         ground_plane = Plane(Vec3.from_xyz(0, 0, 1), Vec3.zero())
         return ground_plane.raycast(ray)
 
-    def raycast_against_entity_oobb(self, entity: RigidEntity, ray: Ray) -> RayHit:
+    def raycast_against_entity_obb(self, entity: RigidEntity, ray: Ray) -> RayHit:
         if isinstance(entity.morph, gs.morphs.Box):
             box: gs.morphs.Box = entity.morph
             size = Vec3.from_xyz(*box.size)
             pose = self.get_pose_of_first_geom(entity)
             aabb = AABB.from_center_and_size(Vec3.zero(), size)
-            ray_hit = aabb.raycast_oobb(pose, ray)
+            ray_hit = aabb.raycast_obb(pose, ray)
             return ray_hit
         else:
             return RayHit.no_hit()
 
+    def raycast_against_entities(self, ray: Ray) -> tuple[RayHit | None, RigidEntity | None]:
+        closest_hit = None
+        hit_entity: RigidEntity | None = None
+        for entity in self.get_entities():
+            ray_hit = self.raycast_against_entity_obb(entity, ray)
+            if ray_hit.is_hit and (closest_hit is None or ray_hit.distance < closest_hit.distance):
+                closest_hit = ray_hit
+                hit_entity = entity
+        return (closest_hit, hit_entity)
+
     def get_entities(self) -> list[RigidEntity]:
         return self.scene.sim.rigid_solver.entities
 
@@ -152,7 +214,7 @@ def _draw_arrow(
         self.scene.draw_debug_arrow(pos.v, dir.v, color=color)  # Only draws arrowhead -- bug?
         self.scene.draw_debug_line(pos.v, pos.v + dir.v, color=color)
 
-    def _draw_entity_unrotated_oobb(self, entity: RigidEntity) -> None:
+    def _draw_entity_unrotated_obb(self, entity: RigidEntity) -> None:
         if isinstance(entity.morph, gs.morphs.Plane):
             plane: gs.morphs.Plane = entity.morph
             pass
@@ -161,10 +223,9 @@ def _draw_entity_unrotated_oobb(self, entity: RigidEntity) -> None:
             size = Vec3.from_xyz(*box.size)
             geom: RigidGeom = entity.geoms[0]
             assert geom._solver.n_envs == 0, "ViewerInteraction only supports single-env for now"
-            gpos = geom.get_pos()  # squeezed if n_envs == 0
-            gquat = geom.get_quat()  # squeezed if n_envs == 0
-            pos = Vec3.from_any_array(gpos.cpu().numpy())
-            quat = Quat.from_any_array(gquat.cpu().numpy())
+            # geom.get_pos() and .get_quat() are squeezed if n_envs == 0
+            pos = Vec3.from_tensor(geom.get_pos())
+            quat = Quat.from_tensor(geom.get_quat())
             aabb = AABB.from_center_and_size(pos, size)
             aabb.expand(0.01)
             self.scene.draw_debug_box(aabb.v, color=Color.red().with_alpha(0.5).tuple(), wireframe=False)