[FEATURE] Get hibernation compiling and working

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1,6 +1,6 @@
 from copy import copy
 from itertools import chain
-from typing import Literal
+from typing import Literal, TYPE_CHECKING
 
 import numpy as np
 import taichi as ti
@@ -25,24 +25,31 @@
 from .rigid_joint import RigidJoint
 from .rigid_link import RigidLink
 
+if TYPE_CHECKING:
+    from genesis.engine.solvers.rigid.rigid_solver_decomp import RigidSolver
+    from genesis.engine.scene import Scene
+
 
 @ti.data_oriented
 class RigidEntity(Entity):
     """
     Entity class in rigid body systems. One rigid entity can be a robot, a terrain, a floating rigid body, etc.
     """
 
+    # override typing
+    _solver: "RigidSolver"
+
     def __init__(
         self,
         scene: "Scene",
-        solver: "Solver",
+        solver: "RigidSolver",
         material: Material,
         morph: Morph,
         surface: Surface,
         idx: int,
         idx_in_solver,
-        link_start=0,
-        joint_start=0,
+        link_start: int = 0,
+        joint_start: int = 0,
         q_start=0,
         dof_start=0,
         geom_start=0,
@@ -55,13 +62,13 @@ def __init__(
         vvert_start=0,
         vface_start=0,
         equality_start=0,
-        visualize_contact=False,
+        visualize_contact: bool = False,
     ):
         super().__init__(idx, scene, morph, solver, material, surface)
 
         self._idx_in_solver = idx_in_solver
-        self._link_start = link_start
-        self._joint_start = joint_start
+        self._link_start: int = link_start
+        self._joint_start: int = joint_start
         self._q_start = q_start
         self._dof_start = dof_start
         self._geom_start = geom_start
@@ -77,11 +84,11 @@ def __init__(
 
         self._base_links_idx = torch.tensor([self.base_link_idx], dtype=gs.tc_int, device=gs.device)
 
-        self._visualize_contact = visualize_contact
+        self._visualize_contact: bool = visualize_contact
 
-        self._is_free = morph.is_free
+        self._is_free: bool = morph.is_free
 
-        self._is_built = False
+        self._is_built: bool = False
 
         self._load_model()
 
@@ -2853,7 +2860,7 @@ def equalities(self):
         return self._equalities
 
     @property
-    def is_free(self):
+    def is_free(self) -> bool:
         """Whether the entity is free to move."""
         return self._is_free
 

genesis/engine/entities/rigid_entity/rigid_geom.py

@@ -4,6 +4,7 @@
 
 import igl
 import numpy as np
+from numpy.typing import NDArray
 import skimage
 import taichi as ti
 import torch
@@ -16,7 +17,9 @@
 from genesis.utils.misc import tensor_to_array
 
 if TYPE_CHECKING:
+    from genesis.engine.solvers.rigid.rigid_solver_decomp import RigidSolver
     from genesis.engine.materials.rigid import Rigid as RigidMaterial
+    from genesis.engine.mesh import Mesh
 
     from .rigid_entity import RigidEntity
     from .rigid_link import RigidLink
@@ -32,18 +35,18 @@ def __init__(
         self,
         link: "RigidLink",
         idx,
-        cell_start,
-        vert_start,
-        face_start,
-        edge_start,
-        verts_state_start,
-        mesh,
-        type,
-        friction,
+        cell_start: int,
+        vert_start: int,
+        face_start: int,
+        edge_start: int,
+        verts_state_start: int,
+        mesh: "Mesh",
+        type: gs.GEOM_TYPE,
+        friction: float,
         sol_params,
         init_pos,
         init_quat,
-        needs_coup,
+        needs_coup: bool,
         contype,
         conaffinity,
         center_init=None,
@@ -52,30 +55,30 @@ def __init__(
         self._link: "RigidLink" = link
         self._entity: "RigidEntity" = link.entity
         self._material: "RigidMaterial" = link.entity.material
-        self._solver = link.entity.solver
-        self._mesh = mesh
+        self._solver: "RigidSolver" = link.entity.solver
+        self._mesh: "Mesh" = mesh
 
         self._uid = gs.UID()
         self._idx = idx
-        self._type = type
-        self._friction = friction
+        self._type: gs.GEOM_TYPE = type
+        self._friction: float = friction
         self._sol_params = sol_params
-        self._needs_coup = needs_coup
+        self._needs_coup: bool = needs_coup
         self._contype = contype
         self._conaffinity = conaffinity
-        self._is_convex = mesh.is_convex
-        self._cell_start = cell_start
-        self._vert_start = vert_start
-        self._face_start = face_start
-        self._edge_start = edge_start
-        self._verts_state_start = verts_state_start
+        self._is_convex: bool = mesh.is_convex
+        self._cell_start: int = cell_start
+        self._vert_start: int = vert_start
+        self._face_start: int = face_start
+        self._edge_start: int = edge_start
+        self._verts_state_start: int = verts_state_start
 
-        self._coup_softness = self._material.coup_softness
-        self._coup_friction = self._material.coup_friction
-        self._coup_restitution = self._material.coup_restitution
+        self._coup_softness: float = self._material.coup_softness
+        self._coup_friction: float = self._material.coup_friction
+        self._coup_restitution: float = self._material.coup_restitution
 
-        self._init_pos = init_pos
-        self._init_quat = init_quat
+        self._init_pos: np.ndarray = init_pos
+        self._init_quat: np.ndarray = init_quat
 
         self._init_verts = mesh.verts
         self._init_faces = mesh.faces
@@ -455,14 +458,14 @@ def uid(self):
         return self._uid
 
     @property
-    def idx(self):
+    def idx(self) -> int:
         """
         Get the global index of the geom in RigidSolver.
         """
         return self._idx
 
     @property
-    def type(self):
+    def type(self) -> gs.GEOM_TYPE:
         """
         Get the type of the geom.
         """
@@ -504,25 +507,25 @@ def entity(self) -> "RigidEntity":
         return self._entity
 
     @property
-    def solver(self):
+    def solver(self) -> "RigidSolver":
         """
         Get the solver that the geom belongs to.s
         """
         return self._solver
 
     @property
-    def is_convex(self):
+    def is_convex(self) -> bool:
         """
         Get whether the geom is convex.
         """
         return self._is_convex
 
     @property
-    def mesh(self):
+    def mesh(self) -> "Mesh":
         return self._mesh
 
     @property
-    def needs_coup(self):
+    def needs_coup(self) -> bool:
         """
         Get whether the geom needs coupling with other non-rigid entities.
         """
@@ -550,35 +553,35 @@ def conaffinity(self):
         return self._conaffinity
 
     @property
-    def coup_softness(self):
+    def coup_softness(self) -> float:
         """
         Get the softness coefficient of the geom for coupling.
         """
         return self._coup_softness
 
     @property
-    def coup_friction(self):
+    def coup_friction(self) -> float:
         """
         Get the friction coefficient of the geom for coupling.
         """
         return self._coup_friction
 
     @property
-    def coup_restitution(self):
+    def coup_restitution(self) -> float:
         """
         Get the restitution coefficient of the geom for coupling.
         """
         return self._coup_restitution
 
     @property
-    def init_pos(self):
+    def init_pos(self) -> np.ndarray:
         """
         Get the initial position of the geom.
         """
         return self._init_pos
 
     @property
-    def init_quat(self):
+    def init_quat(self) -> np.ndarray:
         """
         Get the initial quaternion of the geom.
         """
@@ -725,14 +728,14 @@ def n_cells(self):
         return np.prod(self._sdf_res)
 
     @property
-    def n_verts(self):
+    def n_verts(self) -> int:
         """
         Number of vertices of the geom.
         """
         return len(self._init_verts)
 
     @property
-    def n_faces(self):
+    def n_faces(self) -> int:
         """
         Number of faces of the geom.
         """