Merge branch 'main' into yiling/250703_migrate_test

Author: YilingQiao

.github/workflows/generic-cpu.yml

@@ -5,6 +5,10 @@ on:
     branches:
       - main
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   generic-cpu:
     strategy:

.github/workflows/linux-gpu.yml

@@ -7,6 +7,12 @@ on:
       - main
   pull_request:
 
+concurrency:
+  # Cancel all workflows that are stil running if any when updating branches associated with PRs,
+  # BUT don't do anything for workflows that are not triggered by PRs.
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   linux-gpu:
     runs-on: [self-hosted, coreweave]

examples/drone/hover_env.py

@@ -90,6 +90,9 @@ def __init__(self, num_envs, env_cfg, obs_cfg, reward_cfg, command_cfg, show_vie
         # build scene
         self.scene.build(n_envs=num_envs)
 
+        # Set damping to a small value to avoid numerical instability
+        self.drone.set_dofs_damping(torch.tensor([0.0, 0.0, 0.0, 1e-4, 1e-4, 1e-4]))
+
         # prepare reward functions and multiply reward scales by dt
         self.reward_functions, self.episode_sums = dict(), dict()
         for name in self.reward_scales.keys():

genesis/__init__.py

@@ -1,23 +1,19 @@
-# import taichi while suppressing its output
+import io
 import os
 import sys
 import site
 import atexit
 import logging as _logging
 import traceback
 from platform import system
-from unittest.mock import patch
+from contextlib import redirect_stdout
 
+# Import taichi while collecting its output without printing directly
+_ti_outputs = io.StringIO()
 
-_ti_outputs = []
+os.environ.setdefault("TI_ENABLE_PYBUF", "0" if sys.stdout is sys.__stdout__ else "1")
 
-
-def fake_print(*args, **kwargs):
-    output = "".join(args)
-    _ti_outputs.append(output)
-
-
-with patch("builtins.print", fake_print):
+with redirect_stdout(_ti_outputs):
     import taichi as ti
 
 try:
@@ -52,11 +48,9 @@ def init(
     logger_verbose_time=False,
     performance_mode: bool = False,  # True: compilation up to 6x slower (GJK), but runs ~1-5% faster
 ):
-    # Consider Genesis as initialized right away
     global _initialized
     if _initialized:
         raise_exception("Genesis already initialized.")
-    _initialized = True
 
     # genesis._theme
     global _theme
@@ -193,7 +187,7 @@ def init(
         ti_arch = TI_ARCH[platform][gs_backend.cpu]
 
     # init taichi
-    with patch("builtins.print", fake_print):
+    with redirect_stdout(_ti_outputs):
         ti.init(
             arch=ti_arch,
             # debug is causing segfault on some machines
@@ -227,9 +221,10 @@ def init(
         f"Running on ~~<[{device_name}]>~~ with backend ~~<{backend}>~~. Device memory: ~~<{total_mem:.2f}>~~ GB."
     )
 
-    for ti_output in _ti_outputs:
+    for ti_output in _ti_outputs.getvalue().splitlines():
         logger.debug(ti_output)
-    _ti_outputs.clear()
+    _ti_outputs.truncate(0)
+    _ti_outputs.seek(0)
 
     global exit_callbacks
     exit_callbacks = []
@@ -238,6 +233,8 @@ def init(
         f"🚀 Genesis initialized. 🔖 version: ~~<{__version__}>~~, 🌱 seed: ~~<{seed}>~~, 📏 precision: '~~<{precision}>~~', 🐛 debug: ~~<{debug}>~~, 🎨 theme: '~~<{theme}>~~'."
     )
 
+    _initialized = True
+
 
 ########################## init ##########################
 def destroy():

genesis/engine/bvh.py

@@ -343,38 +343,38 @@ def compute_bounds(self):
     @ti.kernel
     def _kernel_compute_bounds_init(self):
         self.updated[None] = True
-        self.internal_node_active.fill(0)
-        self.internal_node_ready.fill(0)
+        self.internal_node_active.fill(False)
+        self.internal_node_ready.fill(False)
 
         for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs):
             idx = ti.i32(self.morton_codes[i_b, i])
             self.nodes[i_b, i + self.n_aabbs - 1].bound.min = self.aabbs[i_b, idx].min
             self.nodes[i_b, i + self.n_aabbs - 1].bound.max = self.aabbs[i_b, idx].max
             parent_idx = self.nodes[i_b, i + self.n_aabbs - 1].parent
             if parent_idx != -1:
-                self.internal_node_active[i_b, parent_idx] = 1
+                self.internal_node_active[i_b, parent_idx] = True
 
     @ti.kernel
     def _kernel_compute_bounds_one_layer(self):
         self.updated[None] = False
         for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs - 1):
-            if self.internal_node_active[i_b, i] == 0:
+            if not self.internal_node_active[i_b, i]:
                 continue
             left_bound = self.nodes[i_b, self.nodes[i_b, i].left].bound
             right_bound = self.nodes[i_b, self.nodes[i_b, i].right].bound
             self.nodes[i_b, i].bound.min = ti.min(left_bound.min, right_bound.min)
             self.nodes[i_b, i].bound.max = ti.max(left_bound.max, right_bound.max)
             parent_idx = self.nodes[i_b, i].parent
             if parent_idx != -1:
-                self.internal_node_ready[i_b, parent_idx] = 1
-            self.internal_node_active[i_b, i] = 0
+                self.internal_node_ready[i_b, parent_idx] = True
+            self.internal_node_active[i_b, i] = False
             self.updated[None] = True
 
         for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs - 1):
-            if self.internal_node_ready[i_b, i] == 0:
+            if not self.internal_node_ready[i_b, i]:
                 continue
-            self.internal_node_active[i_b, i] = 1
-            self.internal_node_ready[i_b, i] = 0
+            self.internal_node_active[i_b, i] = True
+            self.internal_node_ready[i_b, i] = False
 
     @ti.kernel
     def query(self, aabbs: ti.template()):

genesis/engine/entities/rigid_entity/rigid_geom.py

@@ -1,5 +1,6 @@
 import os
 import pickle as pkl
+from typing import TYPE_CHECKING
 
 import igl
 import numpy as np
@@ -14,6 +15,12 @@
 from genesis.repr_base import RBC
 from genesis.utils.misc import tensor_to_array
 
+if TYPE_CHECKING:
+    from genesis.engine.materials.rigid import Rigid as RigidMaterial
+
+    from .rigid_entity import RigidEntity
+    from .rigid_link import RigidLink
+
 
 @ti.data_oriented
 class RigidGeom(RBC):
@@ -42,9 +49,9 @@ def __init__(
         center_init=None,
         data=None,
     ):
-        self._link = link
-        self._entity = link.entity
-        self._material = link.entity.material
+        self._link: "RigidLink" = link
+        self._entity: "RigidEntity" = link.entity
+        self._material: "RigidMaterial" = link.entity.material
         self._solver = link.entity.solver
         self._mesh = mesh
 

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -1,4 +1,7 @@
+from typing import TYPE_CHECKING
+
 import numpy as np
+from numpy.typing import ArrayLike
 import taichi as ti
 import torch
 
@@ -9,6 +12,9 @@
 
 from .rigid_geom import RigidGeom, RigidVisGeom
 
+if TYPE_CHECKING:
+    from .rigid_entity import RigidEntity
+
 
 @ti.data_oriented
 class RigidLink(RBC):
@@ -43,8 +49,8 @@ def __init__(
         invweight,
         visualize_contact,
     ):
-        self._name = name
-        self._entity = entity
+        self._name: str = name
+        self._entity: "RigidEntity" = entity
         self._solver = entity.solver
         self._entity_idx_in_solver = entity.idx
 
@@ -68,16 +74,18 @@ def __init__(
         self._vvert_start = vvert_start
         self._vface_start = vface_start
 
-        self._pos = pos
-        self._quat = quat
-        self._inertial_pos = inertial_pos
-        self._inertial_quat = inertial_quat
+        # Link position & rotation at creation time:
+        self._pos: ArrayLike = pos
+        self._quat: ArrayLike = quat
+        # Link's center-of-mass position & principal axes frame rotation at creation time:
+        self._inertial_pos: ArrayLike = inertial_pos
+        self._inertial_quat: ArrayLike = inertial_quat
         self._inertial_mass = inertial_mass
         self._inertial_i = inertial_i
 
         self._visualize_contact = visualize_contact
 
-        self._geoms = gs.List()
+        self._geoms: list[RigidGeom] = gs.List()
         self._vgeoms = gs.List()
 
     def _build(self):
@@ -99,7 +107,7 @@ def _build(self):
                 is_fixed = False
         if self._root_idx is None:
             self._root_idx = gs.np_int(link.idx)
-        self.is_fixed = gs.np_int(is_fixed)
+        self.is_fixed: np.int32 = gs.np_int(is_fixed)  # note: type inconsistent with is_built, is_free, is_leaf: bool
 
         # inertial_mass and inertia_i
         if self._inertial_mass is None:

genesis/engine/solvers/rigid/collider_decomp.py

@@ -153,6 +153,8 @@ def _init_collision_fields(self) -> None:
         self._max_collision_pairs = min(n_possible_pairs, self._solver._max_collision_pairs)
         self._max_contact_pairs = self._max_collision_pairs * self._n_contacts_per_pair
 
+        # FIXME: 'ti.static_print' cannot be used as it will be printed systematically, completely ignoring guard
+        # condition, while 'print' is slowing down the kernel even if every called in practice...
         self._warn_msg_max_collision_pairs = (
             f"{colors.YELLOW}[Genesis] [00:00:00] [WARNING] Ignoring contact pair to avoid exceeding max "
             f"({self._max_contact_pairs}). Please increase the value of RigidSolver's option "
@@ -777,7 +779,7 @@ def _func_broad_phase(self):
                                 continue
 
                             if self.n_broad_pairs[i_b] == self._max_collision_pairs:
-                                ti.static_print(self._warn_msg_max_collision_pairs)
+                                # print(self._warn_msg_max_collision_pairs)
                                 break
                             self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][0] = i_ga
                             self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][1] = i_gb
@@ -887,22 +889,27 @@ def _func_narrow_phase_convex_vs_convex(self):
                 i_ga = self.broad_collision_pairs[i_pair, i_b][0]
                 i_gb = self.broad_collision_pairs[i_pair, i_b][1]
 
+                if self._solver.geoms_info[i_ga].type > self._solver.geoms_info[i_gb].type:
+                    i_ga, i_gb = i_gb, i_ga
+
                 if (
                     self._solver.geoms_info[i_ga].is_convex
                     and self._solver.geoms_info[i_gb].is_convex
                     and not self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.TERRAIN
-                    and not (
-                        self._solver._enable_multi_contact
-                        and self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
-                        and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
-                    )
                     and not (
                         self._solver._box_box_detection
                         and self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.BOX
                         and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
                     )
                 ):
-                    self._func_convex_convex_contact(i_ga, i_gb, i_b)
+                    if ti.static(sys.platform == "darwin"):
+                        self._func_convex_convex_contact(i_ga, i_gb, i_b)
+                    else:
+                        if not (
+                            self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
+                            and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
+                        ):
+                            self._func_convex_convex_contact(i_ga, i_gb, i_b)
 
     @ti.kernel
     def _func_narrow_phase_convex_specializations(self):
@@ -915,15 +922,11 @@ def _func_narrow_phase_convex_specializations(self):
                 if self._solver.geoms_info[i_ga].type > self._solver.geoms_info[i_gb].type:
                     i_ga, i_gb = i_gb, i_ga
 
-                if (
-                    self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
-                    and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
-                ):
-                    if ti.static(sys.platform == "darwin"):
-                        # FIXME: It seems redundant, why don't we just call _func_plane_box_contact directly?
-                        # Anyway in this function, we will call _func_plane_box_contact.
-                        self._func_convex_convex_contact(i_ga, i_gb, i_b)
-                    else:
+                if ti.static(sys.platform != "darwin"):
+                    if (
+                        self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
+                        and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX
+                    ):
                         self._func_plane_box_contact(i_ga, i_gb, i_b)
 
                 if ti.static(self._solver._box_box_detection):
@@ -1110,7 +1113,10 @@ def _func_add_contact(self, i_ga, i_gb, normal, contact_pos, penetration, i_b):
         i_col = self.n_contacts[i_b]
 
         if i_col == self._max_contact_pairs:
-            ti.static_print(self._warn_msg_max_collision_pairs)
+            # FIXME: 'ti.static_print' cannot be used as it will be printed systematically, completely ignoring guard
+            # condition, while 'print' is slowing down the kernel even if every called in practice...
+            # print(self._warn_msg_max_collision_pairs)
+            pass
         else:
             ga_info = self._solver.geoms_info[i_ga]
             gb_info = self._solver.geoms_info[i_gb]
@@ -1194,9 +1200,6 @@ def _func_contact_orthogonals(self, i_ga, i_gb, normal, i_b):
 
     @ti.func
     def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
-        if self._solver.geoms_info[i_ga].type > self._solver.geoms_info[i_gb].type:
-            i_ga, i_gb = i_gb, i_ga
-
         if (
             self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE
             and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX