[MISC] Exploit the exactly-diagonal mass of primitive free bodies to tighten the sparse Hessian band

Author: duburcqa

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1197,6 +1197,14 @@ def _align_link(self, l_info, j_infos, cg_infos, vg_infos, morph, link_idx):
             offset_pos = gu.transform_by_trans_quat(global_com, offset_pos, offset_quat)
             offset_quat = gu.transform_quat_by_quat(principal_quat, offset_quat)
 
+        # 'aligned' marks a free body whose joint-space mass matrix is exactly diagonal, letting the solver drop the
+        # within-link off-diagonal mass coupling. Only primitives qualify: their inertia is analytic, principal-axis and
+        # COM-centered. Mesh/MJCF bodies are not anchored to their composite center of mass (that cannot be done per
+        # variant for a heterogeneous entity sharing one link frame), so their mass is not forced diagonal.
+        l_info["aligned"] = any(j_info["type"] == gs.JOINT_TYPE.FREE for j_info in j_infos) and isinstance(
+            morph, gs.options.morphs.Primitive
+        )
+
         # Refresh the free joint init_qpos to reflect the composed world pose.
         for j_info in j_infos:
             if j_info["type"] == gs.JOINT_TYPE.FREE:
@@ -2472,6 +2480,7 @@ def _add_by_info(self, l_info, j_infos, g_infos, morph, surface):
             invweight=l_info.get("invweight"),
             visualize_contact=self.visualize_contact,
             is_robot=l_info.get("is_robot", root_idx != -1),
+            aligned=l_info.get("aligned", False),
         )
         self._links.append(link)
 

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -607,6 +607,7 @@ def __init__(
         invweight: float | None,
         visualize_contact: bool,
         is_robot: bool,
+        aligned: bool = False,
     ):
         super().__init__(
             entity,
@@ -624,6 +625,9 @@ def __init__(
         )
 
         self._is_robot: bool = is_robot
+        # True when the link's inertia is diagonal in the joint frame, so a free body's joint-space mass matrix is
+        # exactly diagonal and only contacts couple its DOFs.
+        self._aligned: bool = aligned
 
         if self._is_fixed and not entity._batch_fixed_verts:
             verts_state_start = fixed_verts_state_start
@@ -1036,6 +1040,14 @@ def inertial_i(self) -> "np.typing.ArrayLike | None":
         """
         return self._inertial_i
 
+    @property
+    def aligned(self) -> bool:
+        """
+        Whether the link's inertia is diagonal in the joint frame (principal axes, anchored at the COM). When True for
+        a free body, its joint-space mass matrix is exactly diagonal, so only contacts couple its DOFs.
+        """
+        return self._aligned
+
     @property
     def geoms(self) -> list[RigidGeom]:
         """

genesis/engine/solvers/rigid/constraint/solver.py

@@ -2013,14 +2013,17 @@ def func_hessian_direct_batch(
                         * constraint_state.efc_D[i_c, i_b]
                         * constraint_state.active[i_c, i_b]
                     )
-        # H += M, restricted to the island's DOFs. Mass couples only DOFs within the same component (one island), so
-        # the block is dense over the island's gathered DOFs and never reaches another island's block. Iterating the
-        # gathered dof_id (rather than an entity's contiguous range) keeps the add inside this island even when a
-        # component's DOFs are non-contiguous (e.g. an entity whose free bodies interleave in DOF order).
+        # H += M, restricted to the island's DOFs. Mass couples only DOFs within the same kinematic-tree block, which
+        # is a contiguous global DOF range and so maps to a contiguous local range (dof_id is ascending). Bound the add
+        # by that block (dofs_mass_block_start, mapped to local via dof_local_pos) rather than the full constraint
+        # envelope: the envelope already includes mass coupling so block_start >= env_start, and entries below
+        # block_start are structurally zero mass. For an aligned free body the block is the diagonal, so only the
+        # diagonal mass is added; for articulated bodies the whole branch block is added.
         for i_d in range(n):
             i_dg = island_state.dof_id[dof_base + i_d, i_b]
-            env_i = island_state.dof_env_start_local[dof_base + i_d, i_b]
-            for j_d in range(env_i, i_d + 1):
+            mass_block_start = rigid_global_info.dofs_mass_block_start[i_dg]
+            mass_lo = island_state.dof_local_pos[mass_block_start, i_b]
+            for j_d in range(mass_lo, i_d + 1):
                 j_dg = island_state.dof_id[dof_base + j_d, i_b]
                 constraint_state.nt_H[i_b, i_dg, j_dg] = (
                     constraint_state.nt_H[i_b, i_dg, j_dg] + rigid_global_info.mass_mat[i_dg, j_dg, i_b]

genesis/engine/solvers/rigid/rigid_solver.py

@@ -883,7 +883,6 @@ def _init_mass_mat(self):
                 ):
                     mass_parent_mask[i_d, j_d] = 1.0
                 j_l = self.links[j_l].parent_idx
-        self._rigid_global_info.mass_parent_mask.from_numpy(mass_parent_mask)
 
         # Partition each entity's DOFs into contiguous, independently-factorable blocks. M is block-diagonal between
         # DOFs whose links are not kinematic ancestor/descendant of one another, so the per-block bounds let the
@@ -908,6 +907,20 @@ def _init_mass_mat(self):
         for i_d in range(self.n_dofs_):
             block_end[block_start[i_d]] = i_d + 1
         block_end = block_end[block_start]
+
+        # An aligned free body has a diagonal joint-space mass block, so zero its within-link off-diagonal mask to make
+        # the assembled mass exactly diagonal (else ~1e-6 round-off once it rotates) and the skyline envelope tighter.
+        for link in self.links:
+            if not link.aligned or link.n_dofs != 6 or block_start[link.dof_start] != link.dof_start:
+                continue
+            for i_d in range(link.dof_start, link.dof_end):
+                for j_d in range(link.dof_start, link.dof_end):
+                    if i_d != j_d:
+                        mass_parent_mask[i_d, j_d] = 0.0
+                block_start[i_d] = i_d
+                block_end[i_d] = i_d + 1
+
+        self._rigid_global_info.mass_parent_mask.from_numpy(mass_parent_mask)
         self._rigid_global_info.dofs_mass_block_start.from_numpy(block_start)
         self._rigid_global_info.dofs_mass_block_end.from_numpy(block_end)
 

tests/test_rigid_physics.py

@@ -7050,118 +7050,45 @@ def test_merge_entities(is_fixed, merge_fixed_links, show_viewer, tol, monkeypat
 @pytest.mark.slow  # ~450s
 @pytest.mark.required
 def test_heterogeneous_physics_parity(show_viewer, tol):
-    """Test heterogeneous simulation by comparing against independent homogeneous simulations.
+    """Verify a single heterogeneous entity is bit-identical to separate entities of its variants.
 
-    This test verifies that heterogeneous simulation produces identical physics results
-    to running separate homogeneous simulations for each variant, including per-variant
-    initial positions.
+    Uses the welded-child mesh objects from 'test_convexify' (each has an offset center of mass) so the parity check
+    exercises the inertia alignment, not just trivially-symmetric primitives.
     """
-    n_steps = 20
-    box_drop_height = 0.05
-    sphere_drop_height = 0.08
-
-    # Run homogeneous simulation with box only
-    scene_box = gs.Scene(
-        show_viewer=False,
-    )
-    scene_box.add_entity(gs.morphs.Plane())
-    box_obj = scene_box.add_entity(
-        gs.morphs.Box(
-            size=(0.04, 0.04, 0.04),
-            pos=(0.0, 0.0, box_drop_height),
-        )
-    )
-    scene_box.build()
-    for _ in range(n_steps):
-        scene_box.step()
-    box_pos = tensor_to_array(box_obj.get_pos())
-    box_vel = tensor_to_array(box_obj.get_vel())
-
-    # Run homogeneous simulation with sphere only
-    scene_sphere = gs.Scene(
-        show_viewer=False,
-    )
-    scene_sphere.add_entity(gs.morphs.Plane())
-    sphere_obj = scene_sphere.add_entity(
-        gs.morphs.Sphere(
-            radius=0.02,
-            pos=(0.1, 0.0, sphere_drop_height),
-        ),
-    )
-    scene_sphere.build()
-    for _ in range(n_steps):
-        scene_sphere.step()
-    sphere_pos = tensor_to_array(sphere_obj.get_pos())
-    sphere_vel = tensor_to_array(sphere_obj.get_vel())
+    n_steps = 100
+    drop_height = 0.2
+    variants = (("mug_1", "output.xml"), ("donut_0", "output.xml"), ("cup_2", "model.xml"), ("apple_15", "model.xml"))
+
+    def make_morph(name, xml):
+        asset_path = get_hf_dataset(pattern=f"{name}/*")
+        return gs.morphs.MJCF(file=f"{asset_path}/{name}/{xml}", pos=(0.0, 0.0, drop_height))
+
+    # Independent homogeneous reference for each variant.
+    homog_pos, homog_vel = [], []
+    for name, xml in variants:
+        scene = gs.Scene(show_viewer=False)
+        scene.add_entity(gs.morphs.Plane())
+        obj = scene.add_entity(make_morph(name, xml))
+        scene.build()
+        for _ in range(n_steps):
+            scene.step()
+        homog_pos.append(tensor_to_array(obj.get_pos()))
+        homog_vel.append(tensor_to_array(obj.get_vel()))
 
-    # Run heterogeneous simulation with both variants (different sizes AND positions)
-    # 4 envs with 2 variants: envs 0-1 get box, envs 2-3 get sphere
-    scene_het = gs.Scene(
-        show_viewer=show_viewer,
-    )
+    # Single heterogeneous entity with one variant per environment.
+    scene_het = gs.Scene(show_viewer=show_viewer)
     scene_het.add_entity(gs.morphs.Plane())
-    # Divergent per-variant yaw offsets, irrelevant to the dynamics of these symmetric primitives dropped flat (so
-    # the world references still match) but stripped per environment by the relative getters.
-    box_offset_euler = (0.0, 0.0, 30.0)
-    sphere_offset_euler = (0.0, 0.0, -45.0)
-    morphs_heterogeneous = (
-        gs.morphs.Box(
-            size=(0.04, 0.04, 0.04),
-            pos=(0.0, 0.0, box_drop_height),
-            offset_euler=box_offset_euler,
-        ),
-        gs.morphs.Sphere(
-            radius=0.02,
-            pos=(0.1, 0.0, sphere_drop_height),
-            offset_euler=sphere_offset_euler,
-        ),
-    )
-    het_obj = scene_het.add_entity(morph=morphs_heterogeneous)
-    scene_het.build(n_envs=4)
-
-    # Verify initial positions match per-variant morph.pos
-    het_pos_init = het_obj.get_pos()
-    assert_allclose(het_pos_init[0, 2], box_drop_height, tol=tol)
-    assert_allclose(het_pos_init[1, 2], box_drop_height, tol=tol)
-    assert_allclose(het_pos_init[2, 0], 0.1, tol=tol)
-    assert_allclose(het_pos_init[2, 2], sphere_drop_height, tol=tol)
-    assert_allclose(het_pos_init[3, 0], 0.1, tol=tol)
-    assert_allclose(het_pos_init[3, 2], sphere_drop_height, tol=tol)
-
-    # The relative getter strips each variant's own offset back to the user frame (identity), while the world frame
-    # carries the per-environment offset.
-    box_offset_quat = gu.xyz_to_quat(np.array(box_offset_euler), rpy=True, degrees=True)
-    sphere_offset_quat = gu.xyz_to_quat(np.array(sphere_offset_euler), rpy=True, degrees=True)
-    assert_allclose(het_obj.get_quat(relative=True), gu.identity_quat(), tol=tol)
-    het_quat_world = het_obj.get_quat(relative=False)
-    assert_allclose(het_quat_world[:2], box_offset_quat, tol=tol)
-    assert_allclose(het_quat_world[2:], sphere_offset_quat, tol=tol)
-
+    het_obj = scene_het.add_entity(morph=tuple(make_morph(name, xml) for name, xml in variants))
+    scene_het.build(n_envs=len(variants))
     for _ in range(n_steps):
         scene_het.step()
-    het_pos = het_obj.get_pos()
-    het_vel = het_obj.get_vel()
-
-    # Verify heterogeneous results match homogeneous results
-    # Envs 0-1 should match box simulation
-    assert_allclose(het_pos[0], box_pos, tol=tol)
-    assert_allclose(het_pos[1], box_pos, tol=tol)
-    assert_allclose(het_vel[0], box_vel, tol=tol)
-    assert_allclose(het_vel[1], box_vel, tol=tol)
-
-    # Envs 2-3 should match sphere simulation
-    assert_allclose(het_pos[2], sphere_pos, tol=tol)
-    assert_allclose(het_pos[3], sphere_pos, tol=tol)
-    assert_allclose(het_vel[2], sphere_vel, tol=tol)
-    assert_allclose(het_vel[3], sphere_vel, tol=tol)
-
-    # Box envs should have same mass, sphere envs should have same mass
-    mass = het_obj.get_mass()
-    assert_allclose(mass[0], mass[1], tol=tol)
-    assert_allclose(mass[2], mass[3], tol=tol)
-    # Box and sphere should have different masses
-    with pytest.raises(AssertionError):
-        assert_allclose(mass[0], mass[2], tol=tol)
+    het_pos = tensor_to_array(het_obj.get_pos())
+    het_vel = tensor_to_array(het_obj.get_vel())
+
+    # Each environment must match the standalone simulation of its variant.
+    for i in range(len(variants)):
+        assert_allclose(het_pos[i], homog_pos[i], tol=tol)
+        assert_allclose(het_vel[i], homog_vel[i], tol=tol)
 
 
 @pytest.mark.required