support query of vel and pos for a point on a link

Author: YilingQiao

examples/rigid/point_position_velocity.py

@@ -0,0 +1,114 @@
+import argparse
+
+import numpy as np
+import torch
+
+import genesis as gs
+from genesis.utils import geom as gu
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--vis", action="store_true", default=False)
+    args = parser.parse_args()
+
+    gs.init(backend=gs.cpu)
+
+    viewer_options = gs.options.ViewerOptions(
+        camera_pos=(0, -3.5, 2.5),
+        camera_lookat=(0.0, 0.0, 0.5),
+        camera_fov=40,
+        max_FPS=60,
+    )
+
+    scene = gs.Scene(
+        viewer_options=viewer_options,
+        show_viewer=args.vis,
+    )
+
+    scene.add_entity(morph=gs.morphs.Plane())
+
+    robot = scene.add_entity(
+        morph=gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),
+        material=gs.materials.Rigid(),
+    )
+
+    scene.build()
+
+    link4 = robot.get_link("link4")
+    right_finger = robot.get_link("right_finger")
+    hand_link = robot.get_link("hand")
+
+    local_point_offsets = torch.tensor(
+        [[0.05, 0.0, 0.1], [0.02, 0.0, 0.01], [0.08, 0.0, 0.05], [0.0, 0.03, 0.08]], dtype=gs.tc_float, device=gs.device
+    )
+
+    point_links_idx_local = torch.tensor(
+        [link4.idx_local, right_finger.idx_local, hand_link.idx_local, hand_link.idx_local],
+        dtype=gs.tc_int,
+        device=gs.device,
+    )
+
+    all_points_info = [
+        {"link_name": "link4", "point_idx": 0},
+        {"link_name": "right_finger", "point_idx": 1},
+        {"link_name": "hand", "point_idx": 2},
+        {"link_name": "hand", "point_idx": 3},
+    ]
+
+    center = torch.tensor([0.5, 0.0, 0.4])
+    radius = 0.15
+    frequency = 0.3
+
+    point_colors = [
+        (0, 0, 1, 0.8),
+        (1, 0, 0, 0.6),
+        (0, 1, 0, 0.6),
+        (1, 0, 1, 0.6),
+    ]
+
+    n_steps = 1000
+    for step in range(n_steps):
+        t = step * scene.dt
+
+        target_pos = center + torch.tensor(
+            [radius * np.cos(2 * np.pi * frequency * t), radius * np.sin(2 * np.pi * frequency * t), 0.0]
+        )
+
+        qpos_target = robot.inverse_kinematics(
+            link=hand_link,
+            pos=target_pos,
+            max_solver_iters=20,
+            damping=0.1,
+            pos_tol=1e-3,
+        )
+        robot.control_dofs_position(qpos_target)
+
+        scene.step()
+
+        world_positions, world_velocities = robot.get_link_point_position_velocity(
+            local_offsets=local_point_offsets, links_idx_local=point_links_idx_local
+        )
+
+        scene.clear_debug_objects()
+
+        for i, point_info in enumerate(all_points_info):
+            point_idx = point_info["point_idx"]
+
+            point_pos = world_positions[point_idx].numpy()
+
+            scene.draw_debug_sphere(pos=point_pos, radius=0.012, color=point_colors[i % len(point_colors)])
+
+        for i, point_info in enumerate(all_points_info):
+            point_idx = point_info["point_idx"]
+
+            point_vel = world_velocities[point_idx]
+            point_pos = world_positions[point_idx].numpy()
+
+            scene.draw_debug_arrow(
+                pos=point_pos, vec=point_vel.numpy() * 0.3, radius=0.005, color=point_colors[i % len(point_colors)]
+            )
+
+
+if __name__ == "__main__":
+    main()

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -821,7 +821,7 @@ def get_jacobian(self, link, local_point=None):
         link : RigidLink
             The target link.
         local_point : torch.Tensor or None, shape (3,)
-            Coordinates of the point in the link’s *local* frame.
+            Coordinates of the point in the link's *local* frame.
             If None, the link origin is used (back-compat).
 
         Returns
@@ -2882,3 +2882,57 @@ def is_free(self):
     def is_local_collision_mask(self):
         """Whether the contype and conaffinity bitmasks of this entity only applies to self-collision."""
         return self._is_local_collision_mask
+
+    @gs.assert_built
+    def get_link_point_position_velocity(self, local_offsets, links_idx_local, envs_idx=None, *, unsafe=False):
+        """
+        Compute the world position and velocity of points given their local offsets from link origins.
+
+        Parameters
+        ----------
+        local_offsets : torch.Tensor or array_like
+            Local offsets from link origins. Shape should be:
+            - (n_points, 3) for multiple points
+            - (n_envs, n_points, 3) for batched environments
+        links_idx_local : torch.Tensor or array_like
+            The link indices for each point. Shape should be:
+            - (n_points,) for multiple points
+            - (n_envs, n_points) for batched environments
+        envs_idx : None | array_like, optional
+            The indices of the environments. If None, all environments will be considered. Defaults to None.
+        unsafe : bool, optional
+            Whether to skip safety checks. Defaults to False.
+
+        Returns
+        -------
+        positions : torch.Tensor
+            World positions of the points. Shape: (n_points, 3) or (n_envs, n_points, 3)
+        velocities : torch.Tensor
+            World velocities of the points. Shape matches positions.
+
+        Examples
+        --------
+        # Multiple points on different links
+        local_offsets = torch.tensor([
+            [0.1, 0.0, 0.05],  # Point 1 offset
+            [0.02, 0.0, 0.01], # Point 2 offset
+            [0.08, 0.0, 0.05], # Point 3 offset
+            [0.0, 0.03, 0.08]  # Point 4 offset
+        ])
+        links_idx_local = torch.tensor([0, 1, 2, 2])  # Link indices for each point
+
+        pos, vel = entity.get_link_point_position_velocity(
+            local_offsets=local_offsets,
+            links_idx_local=links_idx_local
+        )
+        """
+        links_pos = self.get_links_pos(links_idx_local, envs_idx, unsafe=unsafe)
+        links_quat = self.get_links_quat(links_idx_local, envs_idx, unsafe=unsafe)
+        links_vel = self.get_links_vel(links_idx_local, envs_idx, ref="link_origin", unsafe=unsafe)
+        links_ang = self.get_links_ang(links_idx_local, envs_idx, unsafe=unsafe)
+
+        world_positions = gu.transform_by_trans_quat(local_offsets, links_pos, links_quat)
+        arm_vectors = world_positions - links_pos
+        rotational_vel = torch.cross(links_ang, arm_vectors, dim=-1)
+        world_velocities = links_vel + rotational_vel
+        return world_positions, world_velocities

tests/test_rigid_physics.py

@@ -274,8 +274,18 @@ def _build_multi_pendulum(n):
 
 
 @pytest.fixture(scope="session")
-def pendulum(asset_tmp_path):
-    return _build_multi_pendulum(n=1)
+def one_pendulum(asset_tmp_path):
+
+    mjcf = ET.Element("mujoco", model="one_stick_robot")
+    # ET.SubElement(mjcf, "option", timestep="0.05")
+    default = ET.SubElement(mjcf, "default")
+    ET.SubElement(default, "geom", contype="1", conaffinity="1", condim="3")
+    worldbody = ET.SubElement(mjcf, "worldbody")
+    link1 = ET.SubElement(worldbody, "body", name="body1", pos="0.0 0.0 0.0")
+    ET.SubElement(link1, "geom", type="capsule", fromto="0 0 0 0.0 0 1.0", size="0.05", rgba="0 0 1 0.3")
+    ET.SubElement(link1, "joint", type="hinge", name="joint1", axis="0 1 0")
+
+    return mjcf
 
 
 @pytest.fixture(scope="session")
@@ -2877,3 +2887,59 @@ def test_contype_conaffinity(show_viewer, tol):
     assert_allclose(box1.get_pos(), np.array([0.0, 0.0, 0.25]), atol=1e-3)
     assert_allclose(box2.get_pos(), np.array([0.0, 0.0, 0.75]), atol=1e-3)
     assert_allclose(box3.get_pos(), np.array([0.0, 0.0, 0.75]), atol=1e-3)
+
+
+@pytest.mark.required
+@pytest.mark.merge_fixed_links(False)
+@pytest.mark.parametrize("model_name", ["one_pendulum"])
+@pytest.mark.parametrize("backend", [gs.cpu])
+@pytest.mark.parametrize("gs_solver", [gs.constraint_solver.CG])
+@pytest.mark.parametrize("gs_integrator", [gs.integrator.Euler])
+def test_point_position_velocity(gs_sim, backend, tol):
+    """Test get_link_point_position_velocity against analytical pendulum motion."""
+    scene = gs_sim
+    pendulum = scene.entities[0]
+
+    arm_link = pendulum.get_link("body1")
+    r = 1.0
+    local_point_offsets = torch.tensor([[0.0, 0.0, 1.0]], dtype=gs.tc_float, device=gs.device)
+    point_links_idx_local = torch.tensor([arm_link.idx_local], dtype=gs.tc_int, device=gs.device)
+
+    pendulum.set_dofs_position(torch.tensor([np.pi / 6], dtype=gs.tc_float, device=gs.device))
+    pendulum.set_dofs_velocity(torch.tensor([0.0], dtype=gs.tc_float, device=gs.device))
+
+    n_steps = 50
+
+    for step in range(n_steps):
+        scene.step()
+
+        theta = pendulum.get_dofs_position()[0]
+        theta_dot = pendulum.get_dofs_velocity()[0]
+
+        world_positions, world_velocities = pendulum.get_link_point_position_velocity(
+            local_offsets=local_point_offsets, links_idx_local=point_links_idx_local
+        )
+
+        pos_analytical = torch.tensor([r * np.sin(theta), 0.0, r * np.cos(theta)], dtype=gs.tc_float, device=gs.device)
+
+        # Velocity: v = r * θ̇ * tangent_direction
+        # Tangent direction = [cos(θ), 0, sin(θ)]
+        vel_analytical = torch.tensor(
+            [r * theta_dot * np.cos(theta), 0.0, -r * theta_dot * np.sin(theta)], dtype=gs.tc_float, device=gs.device
+        )
+
+        # Compare with API results
+        pos_queried = world_positions[0]
+        vel_queried = world_velocities[0]
+
+        # Calculate errors
+        pos_error = torch.norm(pos_queried - pos_analytical).item()
+        vel_error = torch.norm(vel_queried - vel_analytical).item()
+        print(step, "pos", pos_error, pos_queried, pos_analytical)
+        print(step, "vel", vel_error, vel_queried, vel_analytical)
+
+        assert_allclose(pos_error, 0.0, atol=tol, err_msg="position error too large")
+        # why cannot be tol?
+        assert_allclose(vel_error, 0.0, atol=0.02, err_msg="velocity error too large")
+
+    # Use appropriate tolerances