Stub code for physical object dragging

Author: gasnica

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -21,6 +21,7 @@
 from genesis.utils.misc import tensor_to_array, ti_field_to_torch, ALLOCATE_TENSOR_WARNING
 
 from ..base_entity import Entity
+from .rigid_geom import RigidGeom
 from .rigid_joint import RigidJoint
 from .rigid_link import RigidLink
 from .rigid_equality import RigidEquality
@@ -3146,7 +3147,7 @@ def q_end(self):
         return self._q_start + self.n_qs
 
     @property
-    def geoms(self):
+    def geoms(self) -> list[RigidGeom]:
         """The list of collision geoms (`RigidGeom`) in the entity."""
         if self.is_built:
             return self._geoms
@@ -3168,7 +3169,7 @@ def vgeoms(self):
             return vgeoms
 
     @property
-    def links(self):
+    def links(self) -> list[RigidLink]:
         """The list of links (`RigidLink`) in the entity."""
         return self._links
 

genesis/engine/entities/rigid_entity/rigid_geom.py

@@ -30,7 +30,7 @@ class RigidGeom(RBC):
 
     def __init__(
         self,
-        link,
+        link: "RigidLink",
         idx,
         cell_start,
         vert_start,
@@ -494,14 +494,14 @@ def metadata(self):
         return self._metadata
 
     @property
-    def link(self):
+    def link(self) -> "RigidLink":
         """
         Get the link that the geom belongs to.
         """
         return self._link
 
     @property
-    def entity(self):
+    def entity(self) -> "RigidEntity":
         """
         Get the entity that the geom belongs to.
         """

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -14,6 +14,7 @@
 
 if TYPE_CHECKING:
     from .rigid_entity import RigidEntity
+    from genesis.engine.solvers.rigid.rigid_solver_decomp import RigidSolver
 
 
 @ti.data_oriented
@@ -51,7 +52,7 @@ def __init__(
     ):
         self._name: str = name
         self._entity: "RigidEntity" = entity
-        self._solver = entity.solver
+        self._solver: "RigidSolver" = entity.solver
         self._entity_idx_in_solver = entity.idx
 
         self._uid = gs.UID()
@@ -265,7 +266,7 @@ def get_quat(self, envs_idx=None):
         return self._solver.get_links_quat([self._idx], envs_idx).squeeze(-2)
 
     @gs.assert_built
-    def get_vel(self, envs_idx=None):
+    def get_vel(self, envs_idx=None) -> torch.Tensor:
         """
         Get the linear velocity of the link in the world frame.
 
@@ -277,7 +278,7 @@ def get_vel(self, envs_idx=None):
         return self._solver.get_links_vel([self._idx], envs_idx).squeeze(-2)
 
     @gs.assert_built
-    def get_ang(self, envs_idx=None):
+    def get_ang(self, envs_idx=None) -> torch.Tensor:
         """
         Get the angular velocity of the link in the world frame.
 

genesis/engine/scene.py

@@ -1201,7 +1201,7 @@ def active_solvers(self):
         return self._sim.active_solvers
 
     @property
-    def entities(self):
+    def entities(self) -> list[Entity]:
         """All the entities in the scene."""
         return self._sim.entities
 

genesis/ext/pyrender/interaction/mouse_spring.py

@@ -0,0 +1,82 @@
+
+from genesis.engine.entities.rigid_entity.rigid_entity import RigidEntity
+from genesis.engine.entities.rigid_entity.rigid_geom import RigidGeom
+
+from .ray import Plane, Ray, RayHit
+from .vec3 import Pose, Quat, Vec3, Color
+
+from genesis.engine.entities.rigid_entity.rigid_link import RigidLink
+
+def _ensure_torch_imported() -> None:
+    global torch
+    import torch
+
+class MouseSpring:
+    def __init__(self):
+        self.held_geom: RigidGeom | 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):
+        # 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_point_in_local = pose.inverse_transform_point(control_point)
+        self.prev_control_point = control_point
+
+    def detach(self):
+        self.held_geom = None
+
+    def apply_force(self, control_point: Vec3, delta_time: float):
+        _ensure_torch_imported()
+        
+        # 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())
+        lin_vel: Vec3 = Vec3.from_tensor(link.get_vel())
+        ang_vel: Vec3 = Vec3.from_tensor(link.get_ang())
+
+        pos_err_v: Vec3 = control_point - link_pos
+        vel_err_v: Vec3 = Vec3.zero() - lin_vel
+        inv_mass: float = float(1.0 / link.get_mass() if 0.0 < link.get_mass() else 0.0)
+
+        inv_dt: float = 1.0 / delta_time
+        tau: float = 1.0 / 2
+        damp: float = 1.0 * 2
+
+        total_impulse: Vec3 = Vec3.zero()
+
+        for i in range(3):
+            dir: Vec3 = Vec3.zero()
+            dir.v[i] = 1.0
+            pos_err: float = dir.dot(pos_err_v)
+            vel_err: float = dir.dot(vel_err_v)
+            err: float = tau * pos_err * inv_dt + damp * vel_err
+            vm: float = inv_mass
+            imp: float = err * vm
+
+            lin_vel += imp * dir * inv_mass
+            total_impulse.v[i] = imp
+
+        total_force = total_impulse * inv_dt
+        force_tensor: torch.Tensor = total_force.as_tensor().unsqueeze(0)
+        link.solver.apply_links_external_force(force_tensor, (link.idx,), ref='link_com', local=False)
+
+        # print("vel", lin_vel, "total_impulse", total_impulse)
+        # print("exp", lin_vel + total_force * inv_mass * link.solver.dt, "at dt", link.solver.dt, "inv dt", inv_dt)
+
+
+    @property
+    def is_attached(self) -> bool:
+        return self.held_geom is not None
+
+
+def _solve_mouse_spring():
+    pass

genesis/ext/pyrender/interaction/vec3.py

@@ -1,11 +1,14 @@
 from dataclasses import dataclass
-from typing import Union
+from typing import TYPE_CHECKING, Union
 
 import numpy as np
 from numpy.typing import NDArray
 
 from genesis.utils.misc import tensor_to_array
 
+if TYPE_CHECKING:
+    from genesis.engine.entities.rigid_entity.rigid_geom import RigidGeom
+
 # 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)]
@@ -209,6 +212,14 @@ def get_inverse(self) -> 'Pose':
         inv_pos = Vec3(-inv_pos.v[1:])
         return Pose(inv_pos, inv_rot)
 
+    @classmethod
+    def from_geom(cls, geom: 'RigidGeom') -> 'Pose':
+        assert geom._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(geom.get_pos())
+        quat = Quat.from_tensor(geom.get_quat())
+        return Pose(pos, quat)
+
 
 @dataclass
 class Color:

genesis/ext/pyrender/interaction/viewer_interaction.py

@@ -7,6 +7,7 @@
 from genesis.engine.entities.rigid_entity.rigid_entity import RigidEntity
 
 from .aabb import AABB
+from .mouse_spring import MouseSpring
 from .ray import Plane, Ray, RayHit
 from .vec3 import Pose, Quat, Vec3, Color
 from .viewer_interaction_base import ViewerInteractionBase, EVENT_HANDLE_STATE, EVENT_HANDLED
@@ -38,13 +39,15 @@ def __init__(self,
         self.camera_yfov: float = camera_yfov
 
         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.prev_mouse_pos: tuple[int, int] = (viewport_size[0] // 2, viewport_size[1] // 2)
 
         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
 
+        self.mouse_spring: MouseSpring = MouseSpring()
+
     @override
     def on_mouse_motion(self, x: int, y: int, dx: int, dy: int) -> EVENT_HANDLE_STATE:
         super().on_mouse_motion(x, y, dx, dy)
@@ -55,23 +58,7 @@ def on_mouse_drag(self, x: int, y: int, dx: int, dy: int, buttons: int, modifier
         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())
+            # actual processing done in update_on_sim_step()
 
             return EVENT_HANDLED
 
@@ -90,13 +77,44 @@ def on_mouse_press(self, x: int, y: int, button: int, modifiers: int) -> EVENT_H
                 pose: Pose = self.get_pose_of_first_geom(self.picked_entity)
                 self.picked_point_in_local = pose.inverse_transform_point(ray_hit.position)
 
+                self.mouse_spring.attach(self.picked_entity, 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
 
+            self.mouse_spring.detach()
+
+    @override
+    def on_resize(self, width: int, height: int) -> EVENT_HANDLE_STATE:
+        super().on_resize(width, height)
+        self.viewport_size = (width, height)
+        self.tan_half_fov = np.tan(0.5 * self.camera_yfov)
+
+    @override
+    def update_on_sim_step(self) -> None:
+        if self.picked_entity:
+            mouse_ray: Ray = self.screen_position_to_ray(*self.prev_mouse_pos)
+            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 = True
+                if use_force:
+                    # apply force
+                    self.mouse_spring.apply_force(new_mouse_3d_pos, self.scene.sim.dt)
+                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())
+
     @override
     def on_draw(self) -> None:
         super().on_draw()
@@ -202,11 +220,7 @@ def get_entities(self) -> list[RigidEntity]:
         return self.scene.sim.rigid_solver.entities
 
     def get_pose_of_first_geom(self, entity: RigidEntity) -> 'Pose':
-        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
-        return Pose(Vec3(gpos.cpu().numpy()), Quat(gquat.cpu().numpy()))
+        return Pose.from_geom(entity.geoms[0])
 
     def _draw_arrow(
         self, pos: Vec3, dir: Vec3, color: tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0),
@@ -215,6 +229,9 @@ def _draw_arrow(
         self.scene.draw_debug_line(pos.v, pos.v + dir.v, color=color)
 
     def _draw_entity_unrotated_obb(self, entity: RigidEntity) -> None:
+        if self.picked_entity:
+            return
+        
         if isinstance(entity.morph, gs.morphs.Plane):
             plane: gs.morphs.Plane = entity.morph
             pass

genesis/ext/pyrender/interaction/viewer_interaction_base.py

@@ -41,5 +41,12 @@ def on_key_release(self, symbol: int, modifiers: int) -> EVENT_HANDLE_STATE:
         if self.log_events:
             gs.logger.info(f"Key released: {chr(symbol)}")
 
+    def on_resize(self, width: int, height: int) -> EVENT_HANDLE_STATE:
+        if self.log_events:
+            gs.logger.info(f"Window resized to {width}x{height}")
+
+    def update_on_sim_step(self) -> None:
+        pass
+
     def on_draw(self) -> None:
         pass

genesis/ext/pyrender/viewer.py

@@ -751,7 +751,7 @@ def on_draw(self):
         if self._run_in_thread or not self.auto_start:
             self.render_lock.release()
 
-    def on_resize(self, width, height):
+    def on_resize(self, width: int, height: int) -> EVENT_HANDLE_STATE:
         """Resize the camera and trackball when the window is resized."""
         if self._renderer is None:
             return
@@ -763,6 +763,7 @@ def on_resize(self, width, height):
         self._trackball.resize(self._viewport_size)
         self._renderer.viewport_width = self._viewport_size[0]
         self._renderer.viewport_height = self._viewport_size[1]
+        self.viewer_interaction.on_resize(width, height)
         self.on_draw()
 
     def on_mouse_motion(self, x: int, y: int, dx: int, dy: int) -> EVENT_HANDLE_STATE:
@@ -1288,6 +1289,9 @@ def refresh(self):
         if self._is_active:
             self.flip()
 
+    def update_on_sim_step(self):
+        self.viewer_interaction.update_on_sim_step()
+
     def _compute_initial_camera_pose(self):
         centroid = self.scene.centroid
         if self.viewer_flags["view_center"] is not None:

genesis/vis/viewer.py

@@ -139,6 +139,8 @@ def update(self, auto_refresh=None):
         if self._followed_entity is not None:
             self.update_following()
 
+        self._pyrender_viewer.update_on_sim_step()
+
         with self.lock:
             self._pyrender_viewer.pending_buffer_updates |= self.context.update()