[BUG FIX] Fix link velocity computation for kinematic entity. (#2508)

* Fix FEM get pos.
* Fix link velocity computation for kinematic entity.
* Fix kinematic solver step not properly tracking dirty state.
* Remove unsupported link reference arg for kinematic entity.

Author: duburcqa

genesis/engine/entities/fem_entity.py

@@ -971,10 +971,10 @@ def remove_vertex_constraints(self, verts_idx_local=None, envs_idx=None):
     @qd.kernel
     def _kernel_get_verts_pos(self, f: qd.i32, pos: qd.types.ndarray(), verts_idx: qd.types.ndarray()):
         # get current position of vertices
-        for i_v, i_b in qd.ndrange(verts_idx.shape[0], verts_idx.shape[1]):
-            i_global = verts_idx[i_v, i_b] + self.v_start
+        for i_b, i_v_ in qd.ndrange(verts_idx.shape[0], verts_idx.shape[1]):
+            i_v = verts_idx[i_b, i_v_] + self.v_start
             for j in qd.static(range(3)):
-                pos[i_b, i_v, j] = self._solver.elements_v[f, i_global, i_b].pos[j]
+                pos[i_b, i_v_, j] = self._solver.elements_v[f, i_v, i_b].pos[j]
 
     def get_el2v(self):
         """

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1230,14 +1230,7 @@ def get_ang(self, envs_idx=None):
         return self._solver.get_links_ang(self.base_link_idx, envs_idx)[..., 0, :]
 
     @gs.assert_built
-    def get_links_pos(
-        self,
-        links_idx_local=None,
-        envs_idx=None,
-        *,
-        ref: Literal["link_origin", "link_com", "root_com"] = "link_origin",
-        unsafe=False,
-    ):
+    def get_links_pos(self, links_idx_local=None, envs_idx=None):
         """
         Returns the position of a given reference point for all the entity's links.
 
@@ -1247,19 +1240,14 @@ def get_links_pos(
             The indices of the links. Defaults to None.
         envs_idx : None | array_like, optional
             The indices of the environments. If None, all environments will be considered. Defaults to None.
-        ref: "link_origin" | "link_com" | "root_com"
-            The reference point being used to express the position of each link.
-            * "root_com": center of mass of the sub-entities to which the link belongs. As a reminder, a single
-              kinematic tree (aka. 'RigidEntity') may compromise multiple "physical" entities, i.e. a kinematic tree
-              that may have at most one free joint, at its root.
 
         Returns
         -------
         pos : torch.Tensor, shape (n_links, 3) or (n_envs, n_links, 3)
             The position of all the entity's links.
         """
         links_idx = self._get_global_idx(links_idx_local, self.n_links, self._link_start, unsafe=True)
-        return self._solver.get_links_pos(links_idx, envs_idx, ref=ref)
+        return self._solver.get_links_pos(links_idx, envs_idx)
 
     @gs.assert_built
     def get_links_quat(self, links_idx_local=None, envs_idx=None):
@@ -1318,14 +1306,7 @@ def get_vAABB(self, envs_idx=None):
         return torch.stack((aabbs[..., 0, :].min(dim=-2).values, aabbs[..., 1, :].max(dim=-2).values), dim=-2)
 
     @gs.assert_built
-    def get_links_vel(
-        self,
-        links_idx_local=None,
-        envs_idx=None,
-        *,
-        ref: Literal["link_origin", "link_com"] = "link_origin",
-        unsafe=False,
-    ):
+    def get_links_vel(self, links_idx_local=None, envs_idx=None):
         """
         Returns linear velocity of all the entity's links expressed at a given reference position in world coordinates.
 
@@ -1335,16 +1316,14 @@ def get_links_vel(
             The indices of the links. Defaults to None.
         envs_idx : None | array_like, optional
             The indices of the environments. If None, all environments will be considered. Defaults to None.
-        ref: "link_origin" | "link_com"
-            The reference point being used to expressed the velocity of each link.
 
         Returns
         -------
         vel : torch.Tensor, shape (n_links, 3) or (n_envs, n_links, 3)
             The linear velocity of all the entity's links.
         """
         links_idx = self._get_global_idx(links_idx_local, self.n_links, self._link_start, unsafe=True)
-        return self._solver.get_links_vel(links_idx, envs_idx, ref=ref)
+        return self._solver.get_links_vel(links_idx, envs_idx)
 
     @gs.assert_built
     def get_links_ang(self, links_idx_local=None, envs_idx=None):
@@ -1800,6 +1779,7 @@ class RigidEntity(KinematicEntity):
 
     if TYPE_CHECKING:
         material: gs.materials.Rigid
+        _solver: "RigidSolver"
 
     def __init__(
         self,
@@ -3176,6 +3156,61 @@ def get_AABB(self, envs_idx=None, *, allow_fast_approx: bool = False):
     def get_aabb(self):
         raise DeprecationError("This method has been removed. Please use 'get_AABB()' instead.")
 
+    @gs.assert_built
+    def get_links_pos(
+        self,
+        links_idx_local=None,
+        envs_idx=None,
+        *,
+        ref: Literal["link_origin", "link_com", "root_com"] = "link_origin",
+    ):
+        """
+        Returns the position of a given reference point for all the entity's links.
+
+        Parameters
+        ----------
+        links_idx_local : None | array_like
+            The indices of the links. Defaults to None.
+        envs_idx : None | array_like, optional
+            The indices of the environments. If None, all environments will be considered. Defaults to None.
+        ref: "link_origin" | "link_com" | "root_com"
+            The reference point being used to express the position of each link.
+            * "root_com": center of mass of the sub-entities to which the link belongs. As a reminder, a single
+              kinematic tree (aka. 'RigidEntity') may compromise multiple "physical" entities, i.e. a kinematic tree
+              that may have at most one free joint, at its root.
+
+        Returns
+        -------
+        pos : torch.Tensor, shape (n_links, 3) or (n_envs, n_links, 3)
+            The position of all the entity's links.
+        """
+        links_idx = self._get_global_idx(links_idx_local, self.n_links, self._link_start, unsafe=True)
+        return self._solver.get_links_pos(links_idx, envs_idx, ref=ref)
+
+    @gs.assert_built
+    def get_links_vel(
+        self, links_idx_local=None, envs_idx=None, *, ref: Literal["link_origin", "link_com"] = "link_origin"
+    ):
+        """
+        Returns linear velocity of all the entity's links expressed at a given reference position in world coordinates.
+
+        Parameters
+        ----------
+        links_idx_local : None | array_like
+            The indices of the links. Defaults to None.
+        envs_idx : None | array_like, optional
+            The indices of the environments. If None, all environments will be considered. Defaults to None.
+        ref: "link_origin" | "link_com"
+            The reference point being used to expressed the velocity of each link.
+
+        Returns
+        -------
+        vel : torch.Tensor, shape (n_links, 3) or (n_envs, n_links, 3)
+            The linear velocity of all the entity's links.
+        """
+        links_idx = self._get_global_idx(links_idx_local, self.n_links, self._link_start, unsafe=True)
+        return self._solver.get_links_vel(links_idx, envs_idx, ref=ref)
+
     @gs.assert_built
     def get_links_acc(self, links_idx_local=None, envs_idx=None):
         links_idx = self._get_global_idx(links_idx_local, self.n_links, self._link_start, unsafe=True)

genesis/engine/solvers/kinematic_solver.py

@@ -485,6 +485,8 @@ def substep_post_coupling(self, f):
                 rigid_global_info=self._rigid_global_info,
                 static_rigid_sim_config=self._static_rigid_sim_config,
             )
+            self._is_forward_pos_updated = True
+            self._is_forward_vel_updated = True
 
     def substep_post_coupling_grad(self, f):
         pass
@@ -942,72 +944,33 @@ def set_dofs_position(self, position, dofs_idx=None, envs_idx=None):
         self._is_forward_pos_updated = True
         self._is_forward_vel_updated = True
 
-    @staticmethod
-    def _convert_ref_to_idx(ref: Literal["link_origin", "link_com", "root_com"]):
-        if ref == "root_com":
-            return 0
-        elif ref == "link_com":
-            return 1
-        elif ref == "link_origin":
-            return 2
-        else:
-            gs.raise_exception("'ref' must be either 'link_origin', 'link_com', or 'root_com'.")
-
-    def get_links_pos(
-        self,
-        links_idx=None,
-        envs_idx=None,
-        *,
-        ref: Literal["link_origin", "link_com", "root_com"] = "link_origin",
-    ):
+    def get_links_pos(self, links_idx=None, envs_idx=None):
         if not gs.use_zerocopy:
             _, links_idx, envs_idx = self._sanitize_io_variables(
                 None, links_idx, self.n_links, "links_idx", envs_idx, (3,), skip_allocation=True
             )
-
-        ref = self._convert_ref_to_idx(ref)
-        if ref == 0:
-            tensor = qd_to_torch(self.links_state.root_COM, envs_idx, links_idx, transpose=True, copy=True)
-        elif ref == 1:
-            i_pos = qd_to_torch(self.links_state.i_pos, envs_idx, links_idx, transpose=True)
-            root_COM = qd_to_torch(self.links_state.root_COM, envs_idx, links_idx, transpose=True)
-            tensor = i_pos + root_COM
-        elif ref == 2:
-            tensor = qd_to_torch(self.links_state.pos, envs_idx, links_idx, transpose=True, copy=True)
-        else:
-            gs.raise_exception("'ref' must be either 'link_origin', 'link_com', or 'root_com'.")
-
+        tensor = qd_to_torch(self.links_state.pos, envs_idx, links_idx, transpose=True, copy=True)
         return tensor[0] if self.n_envs == 0 else tensor
 
     def get_links_quat(self, links_idx=None, envs_idx=None):
         tensor = qd_to_torch(self.links_state.quat, envs_idx, links_idx, transpose=True, copy=True)
         return tensor[0] if self.n_envs == 0 else tensor
 
-    def get_links_vel(
-        self, links_idx=None, envs_idx=None, *, ref: Literal["link_origin", "link_com", "root_com"] = "link_origin"
-    ):
+    def get_links_vel(self, links_idx=None, envs_idx=None):
         if gs.use_zerocopy:
             mask = (0, *indices_to_mask(links_idx)) if self.n_envs == 0 else indices_to_mask(envs_idx, links_idx)
             cd_vel = qd_to_torch(self.links_state.cd_vel, transpose=True)
-            if ref == "root_com":
-                return cd_vel[mask]
             cd_ang = qd_to_torch(self.links_state.cd_ang, transpose=True)
-            if ref == "link_com":
-                i_pos = qd_to_torch(self.links_state.i_pos, transpose=True)
-                delta = i_pos[mask]
-            else:
-                pos = qd_to_torch(self.links_state.pos, transpose=True)
-                root_COM = qd_to_torch(self.links_state.root_COM, transpose=True)
-                delta = pos[mask] - root_COM[mask]
-            return cd_vel[mask] + cd_ang[mask].cross(delta, dim=-1)
+            pos = qd_to_torch(self.links_state.pos, transpose=True)
+            root_COM = qd_to_torch(self.links_state.root_COM, transpose=True)
+            return cd_vel[mask] + cd_ang[mask].cross(pos[mask] - root_COM[mask], dim=-1)
 
         _tensor, links_idx, envs_idx = self._sanitize_io_variables(
             None, links_idx, self.n_links, "links_idx", envs_idx, (3,)
         )
         assert _tensor is not None
         tensor = _tensor[None] if self.n_envs == 0 else _tensor
-        ref = self._convert_ref_to_idx(ref)
-        kernel_get_links_vel(tensor, links_idx, envs_idx, ref, self.links_state, self._static_rigid_sim_config)
+        kernel_get_links_vel(tensor, links_idx, envs_idx, 2, self.links_state, self._static_rigid_sim_config)
         return _tensor
 
     def get_links_ang(self, links_idx=None, envs_idx=None):

genesis/engine/solvers/rigid/abd/forward_kinematics.py

@@ -144,6 +144,19 @@ def kernel_forward_kinematics(
             static_rigid_sim_config=static_rigid_sim_config,
             is_backward=False,
         )
+        func_COM_links(
+            i_b=i_b,
+            links_state=links_state,
+            links_info=links_info,
+            joints_state=joints_state,
+            joints_info=joints_info,
+            dofs_state=dofs_state,
+            dofs_info=dofs_info,
+            entities_info=entities_info,
+            rigid_global_info=rigid_global_info,
+            static_rigid_sim_config=static_rigid_sim_config,
+            is_backward=False,
+        )
         func_forward_velocity_batch(
             i_b=i_b,
             entities_info=entities_info,
@@ -185,6 +198,19 @@ def kernel_masked_forward_kinematics(
                 static_rigid_sim_config=static_rigid_sim_config,
                 is_backward=False,
             )
+            func_COM_links(
+                i_b=i_b,
+                links_state=links_state,
+                links_info=links_info,
+                joints_state=joints_state,
+                joints_info=joints_info,
+                dofs_state=dofs_state,
+                dofs_info=dofs_info,
+                entities_info=entities_info,
+                rigid_global_info=rigid_global_info,
+                static_rigid_sim_config=static_rigid_sim_config,
+                is_backward=False,
+            )
             func_forward_velocity_batch(
                 i_b=i_b,
                 entities_info=entities_info,
@@ -321,6 +347,7 @@ def func_COM_links_entity(
     static_rigid_sim_config: qd.template(),
     is_backward: qd.template(),
 ):
+    EPS = rigid_global_info.EPS[None]
     BW = qd.static(is_backward)
 
     # Becomes static loop in backward pass, because we assume this loop is an inner loop
@@ -372,7 +399,11 @@ def func_COM_links_entity(
 
             i_r = links_info.root_idx[I_l]
             if i_l == i_r:
-                links_state.root_COM[i_l, i_b] = links_state.root_COM_bw[i_l, i_b] / links_state.mass_sum[i_l, i_b]
+                mass_sum = links_state.mass_sum[i_l, i_b]
+                if mass_sum > EPS:
+                    links_state.root_COM[i_l, i_b] = links_state.root_COM_bw[i_l, i_b] / links_state.mass_sum[i_l, i_b]
+                else:
+                    links_state.root_COM[i_l, i_b] = links_state.i_pos_bw[i_r, i_b]
 
     for i_l_ in (
         range(entities_info.link_start[i_e], entities_info.link_end[i_e])
@@ -500,8 +531,6 @@ def func_COM_links_entity(
                     i_j = i_j_ if qd.static(not BW) else (i_j_ + links_info.joint_start[I_l])
 
                     if func_check_index_range(i_j, links_info.joint_start[I_l], links_info.joint_end[I_l], BW):
-                        EPS = rigid_global_info.EPS[None]
-
                         offset_pos = links_state.root_COM[i_l, i_b] - joints_state.xanchor[i_j, i_b]
                         I_j = [i_j, i_b] if qd.static(static_rigid_sim_config.batch_joints_info) else i_j
                         joint_type = joints_info.type[I_j]
@@ -571,7 +600,7 @@ def func_forward_kinematics_entity(
         if qd.static(not BW)
         else qd.static(range(static_rigid_sim_config.max_n_links_per_entity))
     ):
-        i_l = i_l_ if qd.static(not BW) else (i_l_ + entities_info.link_start[i_e])
+        i_l = gs.qd_int(i_l_ if qd.static(not BW) else (i_l_ + entities_info.link_start[i_e]))
 
         if func_check_index_range(i_l, entities_info.link_start[i_e], entities_info.link_end[i_e], BW):
             I_l = [i_l, i_b] if qd.static(static_rigid_sim_config.batch_links_info) else i_l
@@ -977,7 +1006,7 @@ def func_forward_velocity_entity(
         if qd.static(not BW)
         else qd.static(range(static_rigid_sim_config.max_n_links_per_entity))
     ):
-        i_l = i_l_ if qd.static(not BW) else (i_l_ + entities_info.link_start[i_e])
+        i_l = gs.qd_int(i_l_ if qd.static(not BW) else (i_l_ + entities_info.link_start[i_e]))
 
         if func_check_index_range(i_l, entities_info.link_start[i_e], entities_info.link_end[i_e], BW):
             I_l = [i_l, i_b] if qd.static(static_rigid_sim_config.batch_links_info) else i_l