[Migration] Rigid init and forward kinematics (Genesis-Embodied-AI#1401)

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 -----------------------------------

genesis/engine/solvers/avatar_solver.py

@@ -58,8 +58,30 @@ def _kernel_step(self):
     @ti.kernel
     def _kernel_forward_kinematics_links_geoms(self, envs_idx: ti.types.ndarray()):
         for i_b in envs_idx:
-            self._func_forward_kinematics(i_b)
-            self._func_update_geoms(i_b)
+            self._func_forward_kinematics(
+                i_b,
+                self.links_state,
+                self.links_info,
+                self.joints_state,
+                self.joints_info,
+                self.dofs_state,
+                self.dofs_info,
+                self.entities_info,
+                self._rigid_global_info,
+                self._static_rigid_sim_config,
+            )
+            self._func_update_geoms(
+                i_b,
+                self.links_state,
+                self.links_info,
+                self.joints_state,
+                self.joints_info,
+                self.dofs_state,
+                self.dofs_info,
+                self.entities_info,
+                self._rigid_global_info,
+                self._static_rigid_sim_config,
+            )
 
     @ti.func
     def _func_detect_collision(self):

genesis/engine/solvers/rigid/array_class.py

@@ -11,19 +11,34 @@
 DofsInfo = ti.template()
 GeomsState = ti.template()
 GeomsInfo = ti.template()
-GeomsInitAABB = ti.template()
+GeomsInitAABB = ti.template()  # TODO: move to rigid global info
 LinksState = ti.template()
 LinksInfo = ti.template()
+JointsInfo = ti.template()
+JointsState = ti.template()
+VertsState = ti.template()
 VertsInfo = ti.template()
 EdgesInfo = ti.template()
+FacesInfo = ti.template()
+VVertsInfo = ti.template()
+VFacesInfo = ti.template()
+VGeomsInfo = ti.template()
+EntitiesState = ti.template()
+EntitiesInfo = ti.template()
+EqualitiesInfo = ti.template()
 
 
 @ti.data_oriented
 class RigidGlobalInfo:
-    def __init__(self, n_dofs: int, n_entities: int, n_geoms: int, _B: int, f_batch: Callable):
+    def __init__(self, solver, n_dofs: int, n_entities: int, n_geoms: int, _B: int, f_batch: Callable):
         self.n_awake_dofs = ti.field(dtype=gs.ti_int, shape=f_batch())
         self.awake_dofs = ti.field(dtype=gs.ti_int, shape=f_batch(n_dofs))
 
+        self.qpos0 = ti.field(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_qs_))
+        self.qpos = ti.field(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_qs_))
+
+        # self.links_T = ti.Matrix.field(n=4, m=4, dtype=gs.ti_float, shape=solver.n_links)
+
 
 # =========================================== Collider ===========================================