[MISC] Make collision bitmasks filter local to entity for MJCF morphs. (Genesis-Embodied-AI#1499)

* Make collision bitmask filter local to entity for MJCF morphs.
* Fix broken 'RigidSolver.links_info.entity_idx'.

Author: duburcqa

examples/collision/contype.py

@@ -37,8 +37,8 @@ def main():
             pos=(0.025, 0, 0.5),
             quat=(0, 0, 0, 1),
             size=(0.1, 0.1, 0.1),
-            contype=1,
-            conaffinity=1,
+            contype=0b001,
+            conaffinity=0b001,
         ),
         surface=gs.surfaces.Default(
             color=(1.0, 0.0, 0.0, 1.0),
@@ -49,8 +49,8 @@ def main():
             pos=(-0.025, 0, 1.0),
             quat=(0, 0, 0, 1),
             size=(0.1, 0.1, 0.1),
-            contype=2,
-            conaffinity=2,
+            contype=0b010,
+            conaffinity=0b010,
         ),
         surface=gs.surfaces.Default(
             color=(0.0, 1.0, 0.0, 1.0),
@@ -61,8 +61,8 @@ def main():
             pos=(0.0, 0, 1.5),
             quat=(0, 0, 0, 1),
             size=(0.1, 0.1, 0.1),
-            contype=3,
-            conaffinity=3,
+            contype=0b011,
+            conaffinity=0b011,
         ),
         surface=gs.surfaces.Default(
             color=(0.0, 0.0, 1.0, 1.0),
@@ -74,8 +74,8 @@ def main():
             scale=0.004,
             euler=(0, 0, 90),
             pos=(-0.1, 0.0, 1.0),
-            contype=4,
-            conaffinity=4,
+            contype=0b100,
+            conaffinity=0b100,
         ),
     )
 

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -102,6 +102,7 @@ def _load_model(self):
             gs.raise_exception(f"Unsupported morph: {self._morph}.")
 
         self._requires_jac_and_IK = self._morph.requires_jac_and_IK
+        self._is_local_collision_mask = isinstance(self._morph, gs.morphs.MJCF)
 
         self._update_child_idxs()
 
@@ -2874,3 +2875,8 @@ def equalities(self):
     def is_free(self):
         """Whether the entity is free to move."""
         return self._is_free
+
+    @property
+    def is_local_collision_mask(self):
+        """Whether the contype and conaffinity bitmasks of this entity only applies to self-collision."""
+        return self._is_local_collision_mask

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -54,7 +54,7 @@ def __init__(
         self._name: str = name
         self._entity: "RigidEntity" = entity
         self._solver: "RigidSolver" = entity.solver
-        self._entity_idx_in_solver = entity.idx
+        self._entity_idx_in_solver = entity._idx_in_solver
 
         self._uid = gs.UID()
         self._idx: int = idx

genesis/engine/scene.py

@@ -294,9 +294,8 @@ def add_entity(
             material = gs.materials.Rigid()
 
         if surface is None:
-            surface = (
-                gs.surfaces.Default()
-            )  # assign a local surface, otherwise modification will apply on global default surface
+            # assign a local surface, otherwise modification will apply on global default surface
+            surface = gs.surfaces.Default()
 
         if isinstance(material, gs.materials.Rigid):
             # small sdf res is sufficient for primitives regardless of size

genesis/engine/solvers/rigid/collider_decomp.py

@@ -158,13 +158,16 @@ def _compute_collision_pair_validity(self):
             links_root_idx = links_root_idx[:, 0]
             links_parent_idx = links_parent_idx[:, 0]
             links_is_fixed = links_is_fixed[:, 0]
+        entities_is_local_collision_mask = solver.entities_info.is_local_collision_mask.to_numpy()
 
         n_possible_pairs = 0
         collision_pair_validity = np.zeros((n_geoms, n_geoms), dtype=gs.np_int)
         for i_ga in range(n_geoms):
             for i_gb in range(i_ga + 1, n_geoms):
                 i_la = geoms_link_idx[i_ga]
                 i_lb = geoms_link_idx[i_gb]
+                i_ea = links_entity_idx[i_la]
+                i_eb = links_entity_idx[i_lb]
 
                 # geoms in the same link
                 if i_la == i_lb:
@@ -182,7 +185,10 @@ def _compute_collision_pair_validity(self):
                         continue
 
                 # contype and conaffinity
-                if not (
+                if (
+                    (i_ea == i_eb)
+                    or not (entities_is_local_collision_mask[i_ea] or entities_is_local_collision_mask[i_eb])
+                ) and not (
                     (geoms_contype[i_ga] & geoms_conaffinity[i_gb]) or (geoms_contype[i_gb] & geoms_conaffinity[i_ga])
                 ):
                     continue

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -822,6 +822,9 @@ def _init_entity_fields(self):
             entities_gravity_compensation=np.array(
                 [entity.gravity_compensation for entity in entities], dtype=gs.np_float
             ),
+            entities_is_local_collision_mask=np.array(
+                [entity.is_local_collision_mask for entity in entities], dtype=gs.np_bool
+            ),
             # taichi variables
             entities_info=self.entities_info,
             entities_state=self.entities_state,
@@ -2917,6 +2920,7 @@ def kernel_init_entity_fields(
     entities_geom_start: ti.types.ndarray(),
     entities_geom_end: ti.types.ndarray(),
     entities_gravity_compensation: ti.types.ndarray(),
+    entities_is_local_collision_mask: ti.types.ndarray(),
     # taichi variables
     entities_info: array_class.EntitiesInfo,
     entities_state: array_class.EntitiesState,
@@ -2942,6 +2946,7 @@ def kernel_init_entity_fields(
         entities_info.n_geoms[i] = entities_geom_end[i] - entities_geom_start[i]
 
         entities_info.gravity_compensation[i] = entities_gravity_compensation[i]
+        entities_info.is_local_collision_mask[i] = entities_is_local_collision_mask[i]
 
         if ti.static(static_rigid_sim_config.batch_dofs_info):
             for i_d, i_b in ti.ndrange((entities_dof_start[i], entities_dof_end[i]), _B):

genesis/options/morphs.py

@@ -135,7 +135,6 @@ class Primitive(Morph):
     ----
     This class should *not* be instantiated directly.
 
-
     Parameters
     ----------
     pos : tuple, shape (3,), optional
@@ -685,13 +684,25 @@ class MJCF(FileMorph):
 
     Note
     ----
-    MJCF file always contains a worldbody, which we will skip during loading.
-    The robots/objects in MJCF come with their own baselink pose.
-    If `pos`, `euler`, or `quat` is specified, it will override the baselink pose in the MJCF file.
+    MJCF file always contains a 'world' body. Although this body is added to the kinematic tree, it is used to define
+    the initial pose of the root link. If `pos`, `euler`, or `quat` is specified, it will override the root pose that
+    was originally specified in the MJCF file.
+
+    Note
+    ----
+    Genesis currently processes MJCF as if it describing a single entity instead of an actual scene. This means that
+    there is a single gigantic kinematic chain comprising multiple physical kinematic chains connected together using
+    fee joints. The definition of kinematic chain has been stretched a bit to allow us. In particular, there must be
+    multiple root links instead of a single one. One other related limitation is global / world options defined in MJCF
+    but must be set at the scene-level in Genesis are completely ignored at the moment, e.g. the simulation timestep,
+    integrator or constraint solver. Building an actual scene hierarchy with multiple independent entities may be
+    supported in the future.
 
-    The current version of Genesis asumes there's only one child of the worldbody.
-    However, it's possible that a MJCF file contains a scene, not just a single robot, in which case the worldbody will
-    have multiple kinematic trees. We will support such cases in the future.
+    Note
+    ----
+    Collision filters defined in MJCF are considered "local", i.e. they only apply to collision pairs for which both
+    geometries along to that specific entity. This means that there is no way to filter out collision pairs between
+    primitive and MJCF entity at the moment.
 
     Parameters
     ----------

genesis/utils/array_class.py

@@ -1862,6 +1862,7 @@ def get_entities_info(solver):
         "geom_end": V(dtype=gs.ti_int, shape=shape),
         "n_geoms": V(dtype=gs.ti_int, shape=shape),
         "gravity_compensation": V(dtype=gs.ti_float, shape=shape),
+        "is_local_collision_mask": V(dtype=gs.ti_bool, shape=shape),
     }
 
     if use_ndarray: