forward kinematics in rigid entity

Author: YilingQiao

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1300,8 +1300,19 @@ def _kernel_inverse_kinematics(
             for i_sample in range(max_samples):
                 for _ in range(max_solver_iters):
                     # run FK to update link states using current q
-                    self._solver._func_forward_kinematics_entity(self._idx_in_solver, i_b)
-
+                    self._solver._func_forward_kinematics_entity(
+                        self._idx_in_solver,
+                        i_b,
+                        self._solver.links_state,
+                        self._solver.links_info,
+                        self._solver.joints_state,
+                        self._solver.joints_info,
+                        self._solver.dofs_state,
+                        self._solver.dofs_info,
+                        self._solver.entities_info,
+                        self._solver._rigid_global_info,
+                        self._solver._static_rigid_sim_config,
+                    )
                     # compute error
                     solved = True
                     for i_ee in range(n_links):
@@ -1386,7 +1397,19 @@ def _kernel_inverse_kinematics(
 
                 if not solved:
                     # re-compute final error if exited not due to solved
-                    self._solver._func_forward_kinematics_entity(self._idx_in_solver, i_b)
+                    self._solver._func_forward_kinematics_entity(
+                        self._idx_in_solver,
+                        i_b,
+                        self._solver.links_state,
+                        self._solver.links_info,
+                        self._solver.joints_state,
+                        self._solver.joints_info,
+                        self._solver.dofs_state,
+                        self._solver.dofs_info,
+                        self._solver.entities_info,
+                        self._solver._rigid_global_info,
+                        self._solver._static_rigid_sim_config,
+                    )
                     solved = True
                     for i_ee in range(n_links):
                         i_l_ee = links_idx[i_ee]
@@ -1480,7 +1503,19 @@ def _kernel_inverse_kinematics(
             # restore original qpos and link state
             for i_q in range(self.n_qs):
                 self._solver.qpos[i_q + self._q_start, i_b] = self._IK_qpos_orig[i_q, i_b]
-            self._solver._func_forward_kinematics_entity(self._idx_in_solver, i_b)
+            self._solver._func_forward_kinematics_entity(
+                self._idx_in_solver,
+                i_b,
+                self._solver.links_state,
+                self._solver.links_info,
+                self._solver.joints_state,
+                self._solver.joints_info,
+                self._solver.dofs_state,
+                self._solver.dofs_info,
+                self._solver.entities_info,
+                self._solver._rigid_global_info,
+                self._solver._static_rigid_sim_config,
+            )
 
     @gs.assert_built
     def forward_kinematics(self, qpos, qs_idx_local=None, links_idx_local=None, envs_idx=None):
@@ -1549,7 +1584,19 @@ def _kernel_forward_kinematics(
             # set new qpos
             self._solver.qpos[qs_idx[i_q_], envs_idx[i_b_]] = qpos[i_b_, i_q_]
             # run FK
-            self._solver._func_forward_kinematics_entity(self._idx_in_solver, envs_idx[i_b_])
+            self._solver._func_forward_kinematics_entity(
+                self._idx_in_solver,
+                envs_idx[i_b_],
+                self._solver.links_state,
+                self._solver.links_info,
+                self._solver.joints_state,
+                self._solver.joints_info,
+                self._solver.dofs_state,
+                self._solver.dofs_info,
+                self._solver.entities_info,
+                self._solver._rigid_global_info,
+                self._solver._static_rigid_sim_config,
+            )
 
         ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.PARTIAL)
         for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]):
@@ -1563,7 +1610,19 @@ def _kernel_forward_kinematics(
             # restore original qpos
             self._solver.qpos[qs_idx[i_q_], envs_idx[i_b_]] = self._IK_qpos_orig[qs_idx[i_q_], envs_idx[i_b_]]
             # run FK
-            self._solver._func_forward_kinematics_entity(self._idx_in_solver, envs_idx[i_b_])
+            self._solver._func_forward_kinematics_entity(
+                self._idx_in_solver,
+                envs_idx[i_b_],
+                self._solver.links_state,
+                self._solver.links_info,
+                self._solver.joints_state,
+                self._solver.joints_info,
+                self._solver.dofs_state,
+                self._solver.dofs_info,
+                self._solver.entities_info,
+                self._solver._rigid_global_info,
+                self._solver._static_rigid_sim_config,
+            )
 
     # ------------------------------------------------------------------------------------
     # --------------------------------- motion planing -----------------------------------