diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 551023547..d45d61217 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -1,7 +1,7 @@ import inspect import os from itertools import chain -from typing import TYPE_CHECKING, Literal, Any, NamedTuple, Sequence +from typing import TYPE_CHECKING, Literal, Any, Sequence from functools import wraps import quadrants as qd @@ -11,6 +11,7 @@ import genesis as gs from genesis.engine.materials.base import Material +from genesis.engine.mesh import InertialProperties from genesis.options.morphs import Morph from genesis.options.surfaces import Surface from genesis.utils import array_class @@ -26,7 +27,15 @@ from .rigid_equality import RigidEquality from .rigid_geom import RigidGeom from .rigid_joint import RigidJoint -from .rigid_link import KinematicLink, RigidLink, compose_inertial_properties +from .rigid_link import ( + GeomInertialInfo, + KinematicLink, + LinkInertial, + RigidLink, + compose_inertial_properties, + finalize_inertial, + get_local_inertial_from_geom_info, +) if TYPE_CHECKING: from genesis.engine.scene import Scene @@ -50,87 +59,6 @@ def wrapper(self, *args, **kwargs): return wrapper -class InertialProperties(NamedTuple): - """A link's inertial properties: the 3x3 inertia matrix 'i' and the inertial frame pose 'pos'/'quat'.""" - - i: np.ndarray - pos: np.ndarray - quat: np.ndarray - - -def _get_original_inertial_properties(l_info, recompute_inertia): - """Return a link's original inertial properties for alignment, or None to recompute them from geometry. - - None is returned when the link has no valid declared inertia or 'recompute_inertia' is set. - """ - inertia_valid = ( - (l_info.get("inertial_mass") or 0.0) > gs.EPS - and l_info.get("inertial_i") is not None - and (np.diag(l_info["inertial_i"]) > 0.0).all() - and l_info.get("inertial_pos") is not None - ) - if not inertia_valid or recompute_inertia: - return None - return InertialProperties( - i=l_info["inertial_i"], - pos=np.array(l_info["inertial_pos"]), - quat=np.array(l_info["inertial_quat"]) if l_info.get("inertial_quat") is not None else gu.identity_quat(), - ) - - -def _align_geoms_to_inertia(cg_infos, vg_infos, file_inertial): - """Compute a root link's COM and principal-inertia axes, then re-express its geoms into that aligned frame. - - The frame is derived from 'file_inertial' (an 'InertialProperties') when provided, otherwise recomputed from the - (convexified) collision geometry. The collision and visual geom poses in 'cg_infos' / 'vg_infos' are re-expressed - in place so the link origin sits at the COM with the principal axes aligned. Returns '(global_com, principal_quat, - diagonal_inertial)', where 'diagonal_inertial' is the diagonalized 'InertialProperties' for the file path and None - when recomputed from geometry (the geom-based inertia is rebuilt from the re-expressed geoms at build time). - """ - if file_inertial is not None: - # Derive COM and principal axes from file-specified inertia, returning the diagonalized inertia in the new - # (aligned) link frame. - inertia_R = gu.quat_to_R(file_inertial.quat) - inertia_in_link = inertia_R @ file_inertial.i @ inertia_R.T - R_principal = uu.principal_axes_rot(inertia_in_link) - principal_quat = gu.R_to_quat(R_principal) - global_com = file_inertial.pos - diagonal_inertial = InertialProperties( - i=R_principal.T @ inertia_in_link @ R_principal, pos=gu.zero_pos(), quat=gu.identity_quat() - ) - else: - # Compute COM and principal axes from (convexified) collision geometry - geoms_inertial_info = [] - for cg_info in cg_infos: - if not (cg_info.get("contype", 0) or cg_info.get("conaffinity", 0)): - continue - mesh_inertial_info = cg_info["mesh"].get_inertial_info() - if mesh_inertial_info.volume > 0: - geoms_inertial_info.append( - ( - mesh_inertial_info.mass, - mesh_inertial_info.center_mass, - mesh_inertial_info.moment_inertia, - np.array(cg_info.get("pos", gu.zero_pos())), - np.array(cg_info.get("quat", gu.identity_quat())), - ) - ) - _global_mass, global_com, global_inertia = compose_inertial_properties(geoms_inertial_info) - R_principal = uu.principal_axes_rot(global_inertia) - principal_quat = gu.R_to_quat(R_principal) - diagonal_inertial = None - - # Re-express all geoms in the new link frame - for g_info in chain(cg_infos, vg_infos): - g_info["pos"], g_info["quat"] = gu.inv_transform_pos_quat_by_trans_quat( - np.array(g_info.get("pos", gu.zero_pos())), - np.array(g_info.get("quat", gu.identity_quat())), - global_com, - principal_quat, - ) - return global_com, principal_quat, diagonal_inertial - - @qd.data_oriented class KinematicEntity(Entity): """ @@ -194,6 +122,13 @@ def __init__( self._offset_pos = np.array(self._morph.offset_pos, dtype=gs.np_float) self._offset_quat = np.array(self._morph.offset_quat, dtype=gs.np_float) + # Transient per-link, per-variant (finalized inertia, mass-is-user-specified) computed at load from the (then- + # available) collision geometry, used only to compute the align anchor. Dropped after '_align_free_roots'; + # never persisted, since the collision geometry it derives from is irrelevant to a kinematic entity once built. + # Indexed by link-local position (link.idx - link_start), matching the link order and the '_links_offset_*' + # arrays: '_align_link' appends one entry per link in build order, the heterogeneous loop appends variants. + self._align_inertials: list[list[tuple[LinkInertial, bool]]] = [] + # Per-variant base-link offset for heterogeneous entities (None when homogeneous), primary first and aligned # with '_variant_init_qpos'. The primary carries the inertial alignment; the geometry-only variants carry their # morph offset alone, matching how their initial pose is built. @@ -299,38 +234,19 @@ def _load_heterogeneous_morphs(self): f"variant has {v_j_info['n_dofs']}." ) - # Post-process each link's geoms and align the floating base by its own geometry (mirrors '_align_link' - # for the primary), so an asymmetric variant behaves like its homogeneous equivalent. Diagonalized - # inertia from a file-specified base is written back so the per-variant inertial uses the aligned frame. + # Post-process each link's geoms. The COM/principal-axis anchoring of the floating base is deferred to + # '_align_free_roots' after build (where the finalized per-variant composite inertia is known); here + # only the morph pose offset is composed into the variant's init_qpos and offset. offset_pos = np.array(morph.offset_pos, dtype=gs.np_float) offset_quat = np.array(morph.offset_quat, dtype=gs.np_float) - global_com = principal_quat = None cg_vg_infos = [] for v_l_info, v_j_infos, v_g_infos in zip(v_l_infos, v_links_j_infos, v_links_g_infos): is_robot = v_l_info.get("is_robot", np.array(False, dtype=np.bool_)) cg_infos, vg_infos = self._postprocess_geoms_info(morph, v_g_infos, is_robot) - is_root = v_l_info["parent_idx"] == -1 - align = morph.align - if align is None: - align = ( - is_root - and not bool(is_robot) - and all(j_info["type"] == gs.JOINT_TYPE.FREE for j_info in v_j_infos) - ) - if align and is_root and any(j_info["type"] == gs.JOINT_TYPE.FREE for j_info in v_j_infos): - global_com, principal_quat, diagonal_inertial = _align_geoms_to_inertia( - cg_infos, vg_infos, _get_original_inertial_properties(v_l_info, morph.recompute_inertia) - ) - if diagonal_inertial is not None: - v_l_info["inertial_pos"] = diagonal_inertial.pos - v_l_info["inertial_quat"] = diagonal_inertial.quat - v_l_info["inertial_i"] = diagonal_inertial.i - offset_pos = gu.transform_by_trans_quat(global_com, offset_pos, offset_quat) - offset_quat = gu.transform_quat_by_quat(principal_quat, offset_quat) cg_vg_infos.append((cg_infos, vg_infos)) - # Extract variant's init_qpos from parsed joint infos, composing the morph offset and the floating - # base's alignment into the free joint so relative getters report the variant's user frame. + # Extract variant's init_qpos from parsed joint infos, composing the morph offset into the free joint so + # relative getters report the variant's user frame. variant_init_qpos_parts = [] for v_l_info, v_j_infos in zip(v_l_infos, v_links_j_infos): is_root = v_l_info["parent_idx"] == -1 @@ -343,9 +259,6 @@ def _load_heterogeneous_morphs(self): qpos[:3], qpos[3:7], ) - if global_com is not None: - init_pos = gu.transform_by_trans_quat(global_com, init_pos, init_quat) - init_quat = gu.transform_quat_by_quat(principal_quat, init_quat) qpos = np.concatenate([init_pos, init_quat]) variant_init_qpos_parts.append(qpos) if variant_init_qpos_parts: @@ -355,10 +268,22 @@ def _load_heterogeneous_morphs(self): self._variant_offset_pos.append(offset_pos) self._variant_offset_quat.append(offset_quat) - # Add geoms per link - for link, v_l_info, (cg_infos, vg_infos) in zip(self._links, v_l_infos, cg_vg_infos): + # Add geoms per link and stash this variant's finalized inertia for the align anchor. + recompute = morph.recompute_inertia + for i_link, (link, v_l_info, (cg_infos, vg_infos)) in enumerate( + zip(self._links, v_l_infos, cg_vg_infos) + ): self._add_heterogeneous_variant(link, cg_infos, vg_infos) self._on_heterogeneous_scene_variant_loaded(link, morph, v_l_info) + self._align_inertials[i_link].append( + self._finalize_inertial( + None if recompute else v_l_info.get("inertial_mass"), + None if recompute else v_l_info.get("inertial_pos"), + None if recompute else v_l_info.get("inertial_quat"), + None if recompute else v_l_info.get("inertial_i"), + cg_infos, + ) + ) elif isinstance(morph, (gs.morphs.Mesh, gs.morphs.Primitive)): if isinstance(morph, gs.morphs.Mesh): @@ -369,20 +294,14 @@ def _load_heterogeneous_morphs(self): gs.raise_exception("Mixing fixed and non-fixed morphs in heterogeneous entities is not supported.") cg_infos, vg_infos = self._postprocess_geoms_info(morph, g_infos, is_robot=False) - # Align this variant by its own geometry (mirrors '_align_link' for the primary), so an asymmetric - # variant behaves like its homogeneous equivalent. Primitives are never aligned, matching '_align_link'. + # The COM/principal-axis anchoring is deferred to '_align_free_roots' after build; compose only the + # morph pose offset here. offset_pos = np.array(morph.offset_pos, dtype=gs.np_float) offset_quat = np.array(morph.offset_quat, dtype=gs.np_float) - align = morph.align if isinstance(morph, gs.options.morphs.FileMorph) else False - if align is None: - align = not morph.fixed - global_com = principal_quat = None - if align and not morph.fixed: - global_com, principal_quat, _ = _align_geoms_to_inertia(cg_infos, vg_infos, None) - offset_pos = gu.transform_by_trans_quat(global_com, offset_pos, offset_quat) - offset_quat = gu.transform_quat_by_quat(principal_quat, offset_quat) self._add_heterogeneous_variant(self._links[0], cg_infos, vg_infos) + # Mesh/Primitive variants have no explicit inertial; the anchor inertia comes from their geometry. + self._align_inertials[0].append(self._finalize_inertial(None, None, None, None, cg_infos)) if morph.fixed: init_qpos = np.array((), dtype=gs.np_float) @@ -393,9 +312,6 @@ def _load_heterogeneous_morphs(self): np.array(morph.pos, dtype=gs.np_float), np.array(morph.quat, dtype=gs.np_float), ) - if global_com is not None: - init_pos = gu.transform_by_trans_quat(global_com, init_pos, init_quat) - init_quat = gu.transform_quat_by_quat(principal_quat, init_quat) init_qpos = np.concatenate([init_pos, init_quat]) self._variant_init_qpos.append(init_qpos) self._variant_offset_pos.append(offset_pos) @@ -986,6 +902,171 @@ def _build(self): self._vgeoms = self.vgeoms self._is_built = True + # The per-link inertia (and per heterogeneous variant) is now finalized, so anchor each aligned free root at + # its fixed subtree center of mass and principal axes. Must run before the solver reads the link poses and + # inertia. Defined here on the base class so kinematic and rigid entities anchor identically: a kinematic + # entity is commonly used to visualize the target reference motion a rigid entity tracks, so the two coexist + # and the same qpos must map to the same world pose for both. + self._align_free_roots() + + def _align_free_roots(self): + """Anchor each aligned free root at the center of mass and principal axes of its fixed subtree. + + Runs after build, so it composes the finalized inertia the solver's mass matrix uses, and applies the anchoring + per heterogeneous variant (each variant's own inertia and geoms). This delivers the 'align' option's promise (a + COM-centered, principal-axis frame, which also conditions the constraint solve) for mesh and fixed-child + bodies, not just primitives, while keeping a heterogeneous entity bit-identical to its variants' separate + entities. + """ + for root in self._links: + if root.parent_idx != -1 or not root.aligned: + continue + + # Gather the fixed subtree (root + transitive n_dofs == 0 descendants) and each link's pose in the root + # frame; links are in build order so a parent is always visited before its children. A DOF-bearing + # descendant makes the root an articulated chain rather than a single rigid body: its joint-space mass is + # not diagonal and the frames of its moving children are not re-expressed here, so such roots are skipped. + pose_in_root = {root.idx: (gu.zero_pos(), gu.identity_quat())} + subtree = [root] + articulated = False + for link in self._links: + if link is root or link.parent_idx not in pose_in_root: + continue + if link.n_dofs != 0: + articulated = True + break + pose_in_root[link.idx] = gu.transform_pos_quat_by_trans_quat( + np.asarray(link.pos, dtype=gs.np_float), + np.asarray(link.quat, dtype=gs.np_float), + *pose_in_root[link.parent_idx], + ) + subtree.append(link) + if articulated: + continue + + # Anchor each variant on its own finalized inertia (from the load-time stash, available to kinematic and + # rigid entities alike) and geoms. Composing the fixed subtree's inertia and shifting the body frame to its + # COM and principal axes re-expresses the variant's geoms (so the world geometry is unchanged), folds the + # composite into the variant's dynamics inertia (rigid only) and into its offset and init_qpos. + root_local = root.idx - self._link_start + is_heterogeneous = root._variant_vgeom_ranges is not None + for v in range(len(self._align_inertials[root_local])): + inertial_info = [] + mass_specified = set() + for link in subtree: + props, is_specified = self._align_inertials[link.idx - self._link_start][v] + if props.mass <= gs.EPS: + continue + mass_specified.add(is_specified) + rot = gu.quat_to_R(props.quat) + inertia_in_link = rot @ props.inertia @ rot.T + inertial_info.append( + GeomInertialInfo( + InertialProperties(props.mass, props.com, inertia_in_link), *pose_in_root[link.idx] + ) + ) + + # The composite center of mass and principal axes are density-independent only if every contributing + # link's mass comes from the same source. Mixing a user-specified mass with a density-estimated one + # would make the anchor density-dependent, and a kinematic entity has no density to fall back on, so + # alignment could differ from the rigid counterpart. Require all-or-none and raise otherwise. + if len(mass_specified) > 1: + gs.raise_exception( + f"Entity '{self.uid}': an aligned free body mixes user-specified and geometry-estimated link " + "masses. Specify the mass of all of its links or none of them." + ) + if not inertial_info: + continue + mass_total, com_root, inertia_root = compose_inertial_properties(inertial_info) + if mass_total <= gs.EPS: + continue + principal_R = uu.principal_axes_rot(inertia_root) + principal_quat = gu.R_to_quat(principal_R) + inertia_diag = principal_R.T @ inertia_root @ principal_R + + # Re-express the variant's geoms across the subtree so the body frame can move to (com_root, principal + # axes) while the world geometry stays fixed: bring each geom to the root frame, undo the anchoring + # there, bring it back to its link frame. The static link poses are not moved (they are shared across + # heterogeneous variants). A kinematic link has only visual geoms; a rigid link has both. + for link in subtree: + link_pos, link_quat = pose_in_root[link.idx] + if is_heterogeneous: + vgeom_start, vgeom_end = link._variant_vgeom_ranges[v] + geoms = [vg for vg in link.vgeoms if vgeom_start <= vg.idx < vgeom_end] + if isinstance(link, RigidLink): + geom_start, geom_end = link._variant_geom_ranges[v] + geoms += [g for g in link.geoms if geom_start <= g.idx < geom_end] + elif isinstance(link, RigidLink): + geoms = [*link.geoms, *link.vgeoms] + else: + geoms = list(link.vgeoms) + for geom in geoms: + pos, quat = gu.transform_pos_quat_by_trans_quat( + geom._init_pos, geom._init_quat, link_pos, link_quat + ) + pos, quat = gu.inv_transform_pos_quat_by_trans_quat(pos, quat, com_root, principal_quat) + pos, quat = gu.inv_transform_pos_quat_by_trans_quat(pos, quat, link_pos, link_quat) + geom._init_pos, geom._init_quat = pos, quat + geom._init_pos_tc = torch.from_numpy(pos).to(device=gs.device, dtype=gs.tc_float) + geom._init_quat_tc = torch.from_numpy(quat).to(device=gs.device, dtype=gs.tc_float) + + # Fold the composite (diagonal, COM-centered) into the dynamics inertia (rigid only; a kinematic entity + # has no dynamics inertia): the root carries the whole mass, the fixed children are subsumed. Rescale + # the unit-density estimate to the link masses RigidLink._build resolved (uniform per the mass gate). + if isinstance(root, RigidLink): + real_total = sum( + (link._variant_inertial[v][0] if is_heterogeneous else float(link.inertial_mass)) + for link in subtree + ) + scale = real_total / mass_total + zero_pos, identity_quat = gu.zero_pos(), gu.identity_quat() + for link in subtree: + if link is root: + props = LinkInertial(real_total, zero_pos, identity_quat, scale * inertia_diag) + else: + props = LinkInertial(gs.EPS, zero_pos, identity_quat, np.zeros((3, 3), dtype=gs.np_float)) + if is_heterogeneous: + link._variant_inertial[v] = props + else: + link._inertial_mass, link._inertial_pos, link._inertial_quat, link._inertial_i = props + + # Fold the anchoring into the offset (relative getters keep reporting the user frame) and the root + # free-joint init_qpos (the body frame moves to old_pose o (com_root, principal), keeping the + # re-expressed geoms in world). + if is_heterogeneous: + off_pos, off_quat = self._variant_offset_pos[v], self._variant_offset_quat[v] + self._variant_offset_pos[v] = gu.transform_by_trans_quat(com_root, off_pos, off_quat) + self._variant_offset_quat[v] = gu.transform_quat_by_quat(principal_quat, off_quat) + # The solver broadcasts '_links_offset_*' (not '_variant_offset_*') when all variants share an + # offset, so keep the base link offset in sync with the primary variant's aligned offset. + if v == 0: + self._links_offset_pos[root_local] = self._variant_offset_pos[0] + self._links_offset_quat[root_local] = self._variant_offset_quat[0] + if root_local == 0: + self._offset_pos = self._variant_offset_pos[0] + self._offset_quat = self._variant_offset_quat[0] + qpos = self._variant_init_qpos[v] + new_pos = gu.transform_by_trans_quat(com_root, qpos[:3], qpos[3:7]) + new_quat = gu.transform_quat_by_quat(principal_quat, qpos[3:7]) + qpos[:3], qpos[3:7] = new_pos, new_quat + else: + off_pos, off_quat = self._links_offset_pos[root_local], self._links_offset_quat[root_local] + self._links_offset_pos[root_local] = gu.transform_by_trans_quat(com_root, off_pos, off_quat) + self._links_offset_quat[root_local] = gu.transform_quat_by_quat(principal_quat, off_quat) + if root_local == 0: + self._offset_pos, self._offset_quat = self._links_offset_pos[0], self._links_offset_quat[0] + for joint in root.joints: + if joint.type == gs.JOINT_TYPE.FREE: + new_pos = gu.transform_by_trans_quat(com_root, joint._init_qpos[:3], joint._init_qpos[3:7]) + new_quat = gu.transform_quat_by_quat(principal_quat, joint._init_qpos[3:7]) + joint._init_qpos = np.concatenate([new_pos, new_quat]) + root._pos = gu.transform_by_trans_quat(com_root, np.asarray(root.pos, dtype=gs.np_float), root.quat) + root._quat = gu.transform_quat_by_quat(principal_quat, np.asarray(root.quat, dtype=gs.np_float)) + + # The collision-derived inertia was only needed to compute the anchors; drop it now that the alignment lives in + # the persistent geometry, offset and init_qpos (the collision geometry is irrelevant to a built entity). + self._align_inertials = [] + def _create_joints(self, j_infos, link_idx, joint_start): """Create RigidJoint objects from joint info dicts. @@ -1089,6 +1170,7 @@ def _add_by_info(self, l_info, j_infos, g_infos, morph, surface): quat=l_info["quat"], parent_idx=parent_idx, root_idx=root_idx, + aligned=l_info.get("aligned", False), ) self._links.append(link) @@ -1155,16 +1237,55 @@ def _postprocess_geoms_info(self, morph, g_infos, is_robot): return cg_infos, vg_infos + def _finalize_inertial(self, explicit_mass, explicit_com, explicit_quat, explicit_inertia, cg_infos): + """Finalize a link's local inertial at load for the align anchor, with whether its mass is user-specified. + + Uses unit density for the geometry estimate: the anchor (center of mass and principal axes) is density- + independent as long as a subtree's masses are all user-specified or all estimated (enforced by the gate in + '_align_free_roots'), so the material density - which a kinematic entity does not have - is never needed here. + """ + if cg_infos: + # Compose the unit-density geometry estimate from the collision-geom infos (the dict counterpart of + # RigidLink._build's geom-object path), shared so the align anchor matches the dynamics inertia. + geoms_inertial_info = tuple( + GeomInertialInfo( + get_local_inertial_from_geom_info(g_info), + np.asarray(g_info.get("pos", gu.zero_pos()), dtype=gs.np_float), + np.asarray(g_info.get("quat", gu.identity_quat()), dtype=gs.np_float), + ) + for g_info in cg_infos + ) + hint = compose_inertial_properties(geoms_inertial_info) + else: + hint = InertialProperties(0.0, np.zeros(3, dtype=gs.np_float), np.zeros((3, 3), dtype=gs.np_float)) + props = finalize_inertial(explicit_mass, explicit_com, explicit_quat, explicit_inertia, *hint) + return props, explicit_mass is not None and explicit_mass > 0.0 + def _align_link(self, l_info, j_infos, cg_infos, vg_infos, morph, link_idx): - """Carry the morph pose offset into a root link, align it to its COM/principal axes, and record its offset. + """Carry the morph pose offset into a root link and record its offset, and stash each link's finalized inertia. + + Only root (floating-base) links carry the morph 'offset_pos'/'offset_quat' and the 'aligned' flag; the resulting + body-frame offset is stored in '_links_offset_*' so the relative getters report the user's original pose. + Separately, every link's finalized inertia (explicit values, else from its collision geometry) is computed here + - the one point where the collision geometry is available to both kinematic and rigid entities - and stashed + transiently so '_align_free_roots' can derive the COM/principal anchor identically for both. + """ + # Stash the finalized local inertia of this link (primary variant). recompute_inertia discards explicit values + # for non-world-fixed links, exactly as RigidLink._build does; an aligned free body's subtree is never + # world-fixed, so it suffices to honor the morph flag here. + recompute = isinstance(morph, gs.options.morphs.FileMorph) and morph.recompute_inertia + self._align_inertials.append( + [ + self._finalize_inertial( + None if recompute else l_info.get("inertial_mass"), + None if recompute else l_info.get("inertial_pos"), + None if recompute else l_info.get("inertial_quat"), + None if recompute else l_info.get("inertial_i"), + cg_infos, + ) + ] + ) - Only root (floating-base) links are affected. The morph 'offset_pos'/'offset_quat' is composed into the link - world pose, then (for a free root that opts into alignment) the COM/principal-axis transform is applied on top; - l_info, j_infos, cg_infos and vg_infos are mutated in-place so kinematic and rigid entities share the same - aligned qpos and link frame. The resulting body-frame offset (morph offset then alignment) is stored in - '_links_offset_*' so the relative getters report the user's original pose; children carry no offset, and each - root is independent (multi-root entities keep a distinct offset per root). - """ if l_info["parent_idx"] != -1: return @@ -1182,20 +1303,15 @@ def _align_link(self, l_info, j_infos, cg_infos, vg_infos, morph, link_idx): align = not bool(l_info.get("is_robot", False)) and all( j_info["type"] == gs.JOINT_TYPE.FREE for j_info in j_infos ) - if align and any(j_info["type"] == gs.JOINT_TYPE.FREE for j_info in j_infos): - global_com, principal_quat, diagonal_inertial = _align_geoms_to_inertia( - cg_infos, vg_infos, _get_original_inertial_properties(l_info, morph.recompute_inertia) - ) - if diagonal_inertial is not None: - l_info["inertial_pos"] = diagonal_inertial.pos - l_info["inertial_quat"] = diagonal_inertial.quat - l_info["inertial_i"] = diagonal_inertial.i - # Shift the link frame to the COM and rotate to the principal axes, folding the alignment into the offset. - l_info["pos"] = gu.transform_by_trans_quat(global_com, l_info["pos"], l_info["quat"]) - l_info["quat"] = gu.transform_quat_by_quat(principal_quat, l_info["quat"]) - offset_pos = gu.transform_by_trans_quat(global_com, offset_pos, offset_quat) - offset_quat = gu.transform_quat_by_quat(principal_quat, offset_quat) + # A free body opting into alignment (or any primitive, which is inherently principal-axis and COM-centered) has + # an exactly-diagonal joint-space mass matrix once anchored. The COM/principal-axis anchoring itself is deferred + # to '_align_free_roots' after build, where the finalized composite inertia of the fixed subtree is known and + # can be applied per heterogeneous variant. Here only the morph pose offset is composed (child link poses are + # defined relative to it); the anchoring transform is folded in later. + l_info["aligned"] = any(j_info["type"] == gs.JOINT_TYPE.FREE for j_info in j_infos) and ( + align or isinstance(morph, gs.options.morphs.Primitive) + ) # Refresh the free joint init_qpos to reflect the composed world pose. for j_info in j_infos: @@ -2472,6 +2588,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) diff --git a/genesis/engine/entities/rigid_entity/rigid_link.py b/genesis/engine/entities/rigid_entity/rigid_link.py index 3394d110f..cb339556a 100644 --- a/genesis/engine/entities/rigid_entity/rigid_link.py +++ b/genesis/engine/entities/rigid_entity/rigid_link.py @@ -1,11 +1,11 @@ from itertools import starmap -from typing import TYPE_CHECKING, Sequence +from typing import TYPE_CHECKING, NamedTuple, Sequence import numpy as np import torch -import trimesh import genesis as gs +from genesis.engine.mesh import InertialProperties from genesis.repr_base import RBC from genesis.typing import LaxPositiveFArrayType, Matrix3x3Type, UnitVec4FType, Vec3FType from genesis.utils import geom as gu @@ -30,33 +30,36 @@ INERTIA_RATIO_MAX = 100.0 -def get_local_inertial_from_geom(geom: RigidGeom | RigidVisGeom, rho: float) -> tuple[float, Vec3FType, Matrix3x3Type]: - """ - Extract the local inertial properties (mass, center of mass, inertia tensor) of a given rigid geometry. - """ - geom_type = gs.GEOM_TYPE.MESH if isinstance(geom, RigidVisGeom) else geom.type +def get_local_inertial_from_geom_info(g_info: dict, rho: float = 1.0) -> InertialProperties: + """Local inertial properties (mass, center of mass, inertia tensor) of a parsed geometry-info dict. + Primitive types use the analytic formula on `data`; MESH defers to the mesh's cached unit-density mass properties + (`Mesh.get_inertial_info`) scaled by `rho`. This is the load-time computation available to kinematic and rigid + entities alike (it operates on the parsed info, not on a geom object), so the anchor it feeds matches the finalized + geom-derived inertia exactly. + """ + geom_type = g_info["type"] geom_com_local = np.zeros(3) if geom_type == gs.GEOM_TYPE.PLANE: geom_mass = 0.0 geom_inertia_local = np.zeros(3, dtype=gs.np_float) elif geom_type == gs.GEOM_TYPE.SPHERE: - radius = geom.data[0] + radius = g_info["data"][0] geom_mass = (4.0 / 3.0) * np.pi * radius**3 * rho I = (2.0 / 5.0) * geom_mass * radius**2 geom_inertia_local = np.diag([I, I, I]) elif geom_type == gs.GEOM_TYPE.ELLIPSOID: - hx, hy, hz = geom.data[:3] + hx, hy, hz = g_info["data"][:3] geom_mass = (4.0 / 3.0) * np.pi * hx * hy * hz * rho geom_inertia_local = (geom_mass / 5.0) * np.diag([hy**2 + hz**2, hx**2 + hz**2, hx**2 + hy**2]) elif geom_type == gs.GEOM_TYPE.CYLINDER: - radius, height = geom.data[:2] + radius, height = g_info["data"][:2] geom_mass = np.pi * radius**2 * height * rho I_r = (geom_mass / 12.0) * (3.0 * radius**2 + height**2) I_z = 0.5 * geom_mass * radius**2 geom_inertia_local = np.diag([I_r, I_r, I_z]) elif geom_type == gs.GEOM_TYPE.CAPSULE: - radius, height = geom.data[:2] + radius, height = g_info["data"][:2] m_cyl = np.pi * radius**2 * height * rho m_sph = (4.0 / 3.0) * np.pi * radius**3 * rho geom_mass = m_cyl + m_sph @@ -66,46 +69,46 @@ def get_local_inertial_from_geom(geom: RigidGeom | RigidVisGeom, rho: float) -> I_h = 0.5 * m_cyl * radius**2 + (2.0 / 5.0) * m_sph * radius**2 geom_inertia_local = np.diag([I_r, I_r, I_h]) elif geom_type == gs.GEOM_TYPE.BOX: - hx, hy, hz = geom.data[:3] + hx, hy, hz = g_info["data"][:3] geom_mass = (hx * hy * hz) * rho geom_inertia_local = (geom_mass / 12.0) * np.diag([hy**2 + hz**2, hx**2 + hz**2, hx**2 + hy**2]) else: - # MESH type - if isinstance(geom, RigidVisGeom): - inertia_mesh = trimesh.Trimesh(geom.init_vverts, geom.init_vfaces, process=False) + # MESH type: reuse the mesh's cached unit-density mass properties; mass and inertia scale linearly with density. + inertial = g_info["mesh"].get_inertial_info().inertial + if inertial is None: + geom_mass, geom_inertia_local = 0.0, np.zeros((3, 3), dtype=gs.np_float) else: - inertia_mesh = trimesh.Trimesh(geom.init_verts, geom.init_faces, process=False) + geom_mass = inertial.mass * rho + geom_com_local = inertial.com + geom_inertia_local = inertial.i * rho - if not inertia_mesh.is_watertight: - inertia_mesh = inertia_mesh.convex_hull + return InertialProperties(geom_mass, geom_com_local, geom_inertia_local) - # FIXME: without this check, some geom will have negative volume even after the above convex - # hull operation, e.g. 'tests/test_examples.py::test_example[rigid/terrain_from_mesh.py-None]' - if inertia_mesh.volume < 0.0: - inertia_mesh.invert() - inertia_mesh.density = rho - geom_mass = inertia_mesh.mass - geom_com_local = inertia_mesh.center_mass - geom_inertia_local = inertia_mesh.moment_inertia +class GeomInertialInfo(NamedTuple): + """A geom's intrinsic 'inertial' (in the geom's own frame) placed at pose 'pos'/'quat' in the parent link, as + consumed by 'compose_inertial_properties'.""" - return geom_mass, geom_com_local, geom_inertia_local + inertial: InertialProperties + pos: Vec3FType + quat: UnitVec4FType -def compose_inertial_properties( - geoms_inertial_info: Sequence[tuple[float, Vec3FType, Matrix3x3Type, Vec3FType, UnitVec4FType]], -) -> tuple[float, Vec3FType, Matrix3x3Type]: +def compose_inertial_properties(geoms_inertial_info: Sequence[GeomInertialInfo]) -> InertialProperties: """ Compose mass, center of mass, and inertia tensor from multiple geometries. """ global_mass = 0.0 if geoms_inertial_info: - geoms_mass, geoms_com_local, geoms_I_local, geoms_pos, geoms_quat = zip(*geoms_inertial_info) - geoms_mass = np.asarray(geoms_mass) + geoms_inertial, geoms_pos, geoms_quat = zip(*geoms_inertial_info) + geoms_mass = np.asarray([inertial.mass for inertial in geoms_inertial]) global_mass = geoms_mass.sum() if global_mass == 0.0: - return 0.0, np.zeros(3, dtype=np.float64), np.zeros((3, 3), dtype=np.float64) + return InertialProperties(0.0, np.zeros(3, dtype=np.float64), np.zeros((3, 3), dtype=np.float64)) + + geoms_com_local = [inertial.com for inertial in geoms_inertial] + geoms_I_local = [inertial.i for inertial in geoms_inertial] # Compute world COMs of each geom geoms_com_world = np.stack( @@ -126,17 +129,15 @@ def compose_inertial_properties( geom_I_world = gu.transform_inertia_by_T(geom_I_local, T_offset, geom_mass) global_inertia += geom_I_world - return global_mass, global_com, global_inertia + return InertialProperties(global_mass, global_com, global_inertia) -def compute_inertial_from_geoms( - geoms: Sequence[RigidGeom | RigidVisGeom], rho: float -) -> tuple[float, Vec3FType, Matrix3x3Type]: +def compute_inertial_from_geoms(geoms: Sequence[RigidGeom | RigidVisGeom], rho: float) -> InertialProperties: """ Compose inertial properties (mass, center of mass, inertia tensor) from multiple rigid geometries. - Handles all primitive collision geometry types analytically (SPHERE, ELLIPSOID, CYLINDER, CAPSULE, BOX) and falls - back to trimesh for MESH type. + Handles all primitive collision geometry types analytically (SPHERE, ELLIPSOID, CYLINDER, CAPSULE, BOX) and defers + to the mesh's cached unit-density mass properties for MESH type. Parameters ---------- @@ -144,20 +145,62 @@ def compute_inertial_from_geoms( List of geometry objects to compute inertial from. rho : float Material density (kg/m^3). + """ + geoms_inertial_info = [] + for geom in geoms: + if isinstance(geom, RigidVisGeom): + g_info = {"type": gs.GEOM_TYPE.MESH, "mesh": geom.vmesh} + elif geom.type == gs.GEOM_TYPE.MESH: + g_info = {"type": gs.GEOM_TYPE.MESH, "mesh": geom.mesh} + else: + g_info = {"type": geom.type, "data": geom.data} + geoms_inertial_info.append( + GeomInertialInfo(get_local_inertial_from_geom_info(g_info, rho), geom._init_pos, geom._init_quat) + ) + return compose_inertial_properties(geoms_inertial_info) + + +class LinkInertial(NamedTuple): + """A link's (or variant's) finalized inertial in the solver's representation: mass, center of mass 'com', and the + inertia tensor 'inertia' expressed in the principal frame 'quat' (identity when derived from geometry). Distinct + from 'InertialProperties', whose 'i' is a single tensor; the solver stores the principal tensor and its orientation + separately.""" + + mass: float + com: Vec3FType + quat: UnitVec4FType + inertia: Matrix3x3Type + + +def finalize_inertial( + explicit_mass, explicit_com, explicit_quat, explicit_inertia, hint_mass, hint_com, hint_inertia +) -> LinkInertial: + """Resolve a link's local inertial from its parsed explicit values and a geometry-derived estimate (hint). - Returns - ------- - tuple[float, np.ndarray, np.ndarray] - (total_mass, center_of_mass, inertia_tensor) + Explicit values are used when given; otherwise the geometry estimate is used, and an explicit mass rescales a + geometry-derived inertia. The hint can come from geom objects (`compute_inertial_from_geoms`, used by + `RigidLink._build`) or from parsed geometry-info (the load-time align stash) - the single resolution path keeps the + rigid dynamics inertia and the align anchor in lockstep. """ - # Extract inertia information - geoms_inertial_info = tuple( - (*get_local_inertial_from_geom(geom, rho), geom._init_pos, geom._init_quat) for geom in geoms + mass, com, quat, inertia = explicit_mass, explicit_com, explicit_quat, explicit_inertia + if (mass or hint_mass) > MASS_EPS and hint_mass > gs.EPS and mass is not None: + hint_inertia = hint_inertia * (mass / hint_mass) + hint_mass = mass + if mass is None: + mass = hint_mass + if com is None or inertia is None: + com, inertia, quat = hint_com, hint_inertia, gu.identity_quat() + if quat is None: + quat = gu.identity_quat() + # FIXME: Setting zero mass even for fixed links breaks physics for some reason... + # For non-fixed links, it must be non-zero in case for coupling with deformable body solvers. + return LinkInertial( + max(mass, gs.EPS), + np.asarray(com, dtype=gs.np_float), + np.asarray(quat, dtype=gs.np_float), + np.asarray(inertia, dtype=gs.np_float), ) - # Compose all inertia of all geometries in parent link frame - return compose_inertial_properties(geoms_inertial_info) - class KinematicLink(RBC): """ @@ -179,6 +222,7 @@ def __init__( quat: "np.typing.ArrayLike", parent_idx: int, root_idx: int | None, + aligned: bool = False, ): self._name: str = name self._entity: "KinematicEntity" = entity @@ -216,6 +260,10 @@ def __init__( self._pos: "np.typing.ArrayLike" = pos self._quat: "np.typing.ArrayLike" = quat + # True when the link's frame is reframed to its center of mass and principal axes (the 'align' option). Affects + # the persistent geometry of any entity (kinematic or rigid), so it lives on the base class. + self._aligned: bool = aligned + self._vgeoms: list[RigidVisGeom] = gs.List() # Heterogeneous variant tracking (None = not heterogeneous) @@ -462,6 +510,16 @@ def is_fixed(self): """ return self._is_fixed + @property + def aligned(self) -> bool: + """ + Whether the link opts into center-of-mass / principal-axis reframing (the 'align' option, set for a free body + that opts in or is a primitive). The reframing - and the resulting exactly-diagonal joint-space mass block it + enables - is applied only when the body is a single rigid body (a free root with no DOF-bearing descendant); + callers relying on the diagonal mass must check that condition too, as the solver does. + """ + return self._aligned + @property def invweight(self): """Inverse weight of the link. Always zero for KinematicLink (infinite mass).""" @@ -607,6 +665,7 @@ def __init__( invweight: float | None, visualize_contact: bool, is_robot: bool, + aligned: bool = False, ): super().__init__( entity, @@ -621,6 +680,7 @@ def __init__( quat, parent_idx, root_idx, + aligned, ) self._is_robot: bool = is_robot @@ -656,7 +716,7 @@ def __init__( self._geoms: list[RigidGeom] = gs.List() - # Heterogeneous variant tracking (None = not heterogeneous) + # Heterogeneous collision-geom variant tracking (None = not heterogeneous) self._variant_geom_ranges: list[tuple[int, int]] | None = None def _init_variant_tracking(self): @@ -719,6 +779,10 @@ def _build(self): aabb_min = np.minimum(aabb_min, verts.min(axis=0)) aabb_max = np.maximum(aabb_max, verts.max(axis=0)) + # The consistency-check block below rescales its working copy of the geometry estimate; keep the raw hint for + # the shared inertial resolution. + hint_mass_raw, hint_inertia_raw = hint_mass, np.array(hint_inertia) + # Make sure that provided spatial inertia is consistent with the estimate from the geometries if not fixed if (self._inertial_mass or hint_mass) > MASS_EPS and hint_mass > gs.EPS: if self._inertial_pos is not None: @@ -771,25 +835,29 @@ def _build(self): f"Mass is not specified and collision geoms can not be found for link '{self.name}'. " f"Using visual geoms to compute inertial properties." ) - if self._inertial_mass is None: - self._inertial_mass = hint_mass - if self._inertial_pos is None or self._inertial_i is None: - if self._inertial_pos is not None and self._inertial_i is None: - gs.logger.warning( - f"Ignoring center of mass of link '{self.name}' because inertia matrix is not specified." - ) - elif self._inertial_pos is None and self._inertial_i is not None: - gs.logger.warning( - f"Ignoring inertia matrix of link '{self.name}' because center of mass is not specified." - ) - self._inertial_pos = hint_com - self._inertial_i = hint_inertia - self._inertial_quat = gu.identity_quat() + if self._inertial_pos is not None and self._inertial_i is None: + gs.logger.warning( + f"Ignoring center of mass of link '{self.name}' because inertia matrix is not specified." + ) + elif self._inertial_pos is None and self._inertial_i is not None: + gs.logger.warning( + f"Ignoring inertia matrix of link '{self.name}' because center of mass is not specified." + ) self._invweight = None - # FIXME: Setting zero mass even for fixed links breaks physics for some reason... - # For non-fixed links, it must be non-zero in case for coupling with deformable body solvers. - self._inertial_mass = max(self._inertial_mass, gs.EPS) + # Resolve the final inertial from the explicit values and the geometry estimate, sharing finalize_inertial with + # the load-time align stash so the dynamics inertia and the align anchor stay in lockstep. + inertial = finalize_inertial( + self._inertial_mass, + self._inertial_pos, + self._inertial_quat, + self._inertial_i, + hint_mass_raw, + hint_com, + hint_inertia_raw, + ) + self._inertial_mass, self._inertial_pos = inertial.mass, inertial.com + self._inertial_quat, self._inertial_i = inertial.quat, inertial.inertia # Postpone computation of inverse weight if not specified if self._invweight is None: @@ -812,7 +880,7 @@ def _build(self): if v == 0: # Primary variant: use the link's own parsed/computed inertial self._variant_inertial.append( - ( + LinkInertial( self._inertial_mass, np.asarray(self._inertial_pos, dtype=gs.np_float), ( @@ -836,7 +904,7 @@ def _build(self): and v_i is not None ): self._variant_inertial.append( - ( + LinkInertial( v_mass, np.asarray(v_pos, dtype=gs.np_float), np.asarray(v_quat, dtype=gs.np_float) if v_quat is not None else gu.identity_quat(), @@ -859,7 +927,7 @@ def _build(self): mass = gs.EPS com = np.zeros(3, dtype=gs.np_float) inertia = np.zeros((3, 3), dtype=gs.np_float) - self._variant_inertial.append((mass, com, gu.identity_quat(), inertia)) + self._variant_inertial.append(LinkInertial(mass, com, gu.identity_quat(), inertia)) def _add_geom( self, diff --git a/genesis/engine/mesh.py b/genesis/engine/mesh.py index d7d78ea1d..e5040bd8f 100644 --- a/genesis/engine/mesh.py +++ b/genesis/engine/mesh.py @@ -15,16 +15,24 @@ import genesis.utils.point_cloud as pc from genesis.options.surfaces import Surface from genesis.repr_base import RBC +from genesis.typing import Matrix3x3Type, Vec3FType from genesis.utils.misc import redirect_libc_stderr +class InertialProperties(NamedTuple): + """A rigid body's intrinsic inertial: mass, center of mass 'com', and inertia tensor 'i' about that COM.""" + + mass: float + com: Vec3FType + i: Matrix3x3Type + + class MeshInertialInfo(NamedTuple): - """Mass properties of a mesh in its own frame; center_mass and moment_inertia are None for degenerate geometry.""" + """A mesh's mass properties in its own frame at unit density: the volume and the corresponding intrinsic + 'inertial', which is None for degenerate geometry that carries no well-defined mass distribution.""" volume: float - mass: float - center_mass: "np.ndarray | None" - moment_inertia: "np.ndarray | None" + inertial: "InertialProperties | None" class Mesh(RBC): @@ -306,13 +314,22 @@ def get_inertial_info(self): try: tmesh = self._mesh if self._mesh.is_watertight else self._mesh.convex_hull volume = float(tmesh.volume) + if volume < 0.0: + # Inward-facing winding gives a negative volume and inverted mass properties (e.g. a closed + # terrain block, which bypasses watertighten since it is already watertight). Flip the mesh so + # the inertia integral reflects the actual solid rather than estimating from an inverted one. + tmesh = tmesh.copy() + tmesh.invert() + volume = -volume center_mass = tmesh.center_mass except QhullError: volume, center_mass = 0.0, None if volume > 0.0 and np.all(np.isfinite(center_mass)): - self._inertial_info = MeshInertialInfo(volume, tmesh.mass, center_mass, tmesh.moment_inertia) + self._inertial_info = MeshInertialInfo( + volume, InertialProperties(tmesh.mass, center_mass, tmesh.moment_inertia) + ) else: - self._inertial_info = MeshInertialInfo(0.0, 0.0, None, None) + self._inertial_info = MeshInertialInfo(0.0, None) return self._inertial_info def get_vert_adjacency(self): diff --git a/genesis/engine/solvers/rigid/constraint/solver.py b/genesis/engine/solvers/rigid/constraint/solver.py index f8c5f9de4..09c6ddf78 100644 --- a/genesis/engine/solvers/rigid/constraint/solver.py +++ b/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] diff --git a/genesis/engine/solvers/rigid/rigid_solver.py b/genesis/engine/solvers/rigid/rigid_solver.py index 33dd0d95f..1a6004d63 100644 --- a/genesis/engine/solvers/rigid/rigid_solver.py +++ 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,28 @@ 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 whose only DOFs are its own free joint 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. A DOF-bearing (articulated) descendant adds off-diagonal base + # coupling, so the block must be exactly the root's own DOFs for the diagonalization to be valid. + for link in self.links: + # 'aligned' already implies a free joint; the block bounds must additionally be exactly the link's own DOFs + # (no DOF-bearing ancestor or descendant), otherwise the coupled block is not diagonal. + if ( + not link.aligned + or block_start[link.dof_start] != link.dof_start + or block_end[link.dof_start] != link.dof_end + ): + 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) diff --git a/tests/test_rigid_physics.py b/tests/test_rigid_physics.py index 5d98813ae..6be443295 100644 --- a/tests/test_rigid_physics.py +++ b/tests/test_rigid_physics.py @@ -1305,7 +1305,7 @@ def test_no_drift(gjk_collision, entity_kind, entity_type, ground_type, show_vie scene = gs.Scene( sim_options=gs.options.SimOptions( - dt=0.004, + dt=0.003, gravity=gravity_world, ), rigid_options=gs.options.RigidOptions( @@ -1476,7 +1476,7 @@ def test_no_drift(gjk_collision, entity_kind, entity_type, ground_type, show_vie if show_viewer: scene.visualizer.update() - for _ in range(300): + for _ in range(400): scene.step() pos_local = tensor_to_array(entity.get_pos()) @ R @@ -4144,10 +4144,11 @@ def test_convexify(euler, show_viewer, gjk_collision): # FIXME: The cup is falling on Windows OS because the convex decomposition provided by CoACD is different than # other platform, and much worst in practice, with the bottom of the tank that is not planar (even discontinuous). # cam.start_recording() - for i in range(1100): + n_settle = 1250 if euler == (74, 15, 90) else 1000 + for i in range(n_settle + 100): scene.step() # cam.render() - if i > 1000: + if i > n_settle: assert_allclose(gs_sim.rigid_solver.get_dofs_velocity(), 0.0, atol=1.0 if sys.platform == "win32" else 0.6) # cam.stop_recording(save_to_filename="video.mp4", fps=60) @@ -6449,7 +6450,7 @@ def test_ellipsoid(xml_path, show_viewer): @pytest.mark.slow # ~200s @pytest.mark.required -def test_mesh_align(show_viewer, tol): +def test_align_mesh(show_viewer, tol): INIT_POS = (0.0, 0.0, 0.1) mango_path = get_hf_dataset(pattern="glb/mango.glb") @@ -6576,7 +6577,7 @@ def test_mesh_align(show_viewer, tol): @pytest.mark.slow # ~200s @pytest.mark.required -def test_urdf_align(show_viewer, tol): +def test_align_urdf(show_viewer, tol): INIT_POS = (0.0, 0.0, 0.7) asset_path = get_hf_dataset(pattern="fork/*") @@ -6629,7 +6630,36 @@ def test_urdf_align(show_viewer, tol): @pytest.mark.required -def test_relative_offset_on_link_relative_geoms(show_viewer, tol): +def test_align_mixed_mass_raises(): + # Mixing a user-specified mass with a geometry-estimated one in an aligned free body makes the anchor density- + # dependent (so rigid and kinematic could align differently) and must raise. The fixed joint with + # merge_fixed_links=False keeps the child a distinct fixed link with unspecified mass while the base specifies one. + urdf = _build_two_link_revolute_urdf( + "mixed_mass_align", + "box", + {"size": "0.06 0.06 0.06"}, + links_inertial=[{"mass": 1.0, "ixx": 0.01, "iyy": 0.01, "izz": 0.01, "origin_xyz": "0 0 0"}, None], + joint_type="fixed", + ) + for material in (gs.materials.Rigid(), gs.materials.Kinematic()): + scene = gs.Scene( + show_viewer=False, + show_FPS=False, + ) + scene.add_entity( + gs.morphs.URDF( + file=urdf, + align=True, + merge_fixed_links=False, + ), + material=material, + ) + with pytest.raises(gs.GenesisException, match="mixes user-specified and geometry-estimated"): + scene.build() + + +@pytest.mark.required +def test_align_relative_offset_on_link_relative_geoms(show_viewer, tol): # To exercise the geom-frame offset strip the geoms MUST sit at non-identity poses relative to their link (explicit # collision/visual ) AND the morph offset MUST be a rotation that does not commute with them - otherwise the # conjugation degenerates to the plain morph offset and a naive (corrupted) strip would still pass. A convex- @@ -7050,118 +7080,62 @@ 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. - - This test verifies that heterogeneous simulation produces identical physics results - to running separate homogeneous simulations for each variant, including per-variant - initial positions. - """ - 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()) + # Uses the fixed-child mesh objects from 'test_convexify' (offset center of mass, distinct mass) so the per-env + # parity check exercises the inertia alignment, not just trivially-symmetric primitives. + n_steps = 100 + drop_height = 0.2 + variants = (("mug_1", "output.xml"), ("donut_0", "output.xml"), ("cup_2", "model.xml"), ("apple_15", "model.xml")) + # Divergent per-variant yaw offset, stripped to identity by the relative getter and carried by the world frame. + # Applied identically to the standalone reference so the dynamics still match. + offset_eulers = ((0.0, 0.0, 30.0), (0.0, 0.0, -45.0), (0.0, 0.0, 90.0), (0.0, 0.0, -120.0)) + # Distinct per-variant placement, dispatched per environment. + positions = ((0.0, 0.0, drop_height), (0.2, 0.0, drop_height), (0.0, 0.2, drop_height), (0.2, 0.2, drop_height)) + + def make_morph(name, xml, pos, offset_euler): + asset_path = get_hf_dataset(pattern=f"{name}/*") + return gs.morphs.MJCF(file=f"{asset_path}/{name}/{xml}", pos=pos, offset_euler=offset_euler) + + # Independent homogeneous reference per variant: world orientation, then post-drop pose, velocity and mass. + homog_quat_world, homog_pos, homog_vel, homog_mass = [], [], [], [] + for (name, xml), pos, offset_euler in zip(variants, positions, offset_eulers): + scene = gs.Scene(show_viewer=False) + scene.add_entity(gs.morphs.Plane()) + obj = scene.add_entity(make_morph(name, xml, pos, offset_euler)) + scene.build() + homog_quat_world.append(obj.get_quat(relative=False)) + for _ in range(n_steps): + scene.step() + homog_pos.append(obj.get_pos()) + homog_vel.append(obj.get_vel()) + homog_mass.append(obj.get_mass()) - # 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=tuple(make_morph(name, xml, pos, oe) for (name, xml), pos, oe in zip(variants, positions, offset_eulers)) ) - het_obj = scene_het.add_entity(morph=morphs_heterogeneous) - scene_het.build(n_envs=4) + scene_het.build(n_envs=len(variants)) - # 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) + # At init each variant sits at its own placement; the relative getter strips its offset (and inertial alignment) to + # identity in the user frame, while the world frame matches the standalone reference's world orientation. 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) + for i in range(len(variants)): + assert_allclose(het_obj.get_pos()[i], positions[i], tol=tol) + assert_allclose(het_obj.get_quat(relative=False)[i], homog_quat_world[i], tol=tol) 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 + + # After the drop each environment matches the standalone simulation of its variant in pose, velocity and mass. + for i in range(len(variants)): + assert_allclose(het_obj.get_pos()[i], homog_pos[i], tol=tol) + assert_allclose(het_obj.get_vel()[i], homog_vel[i], tol=tol) + assert_allclose(het_obj.get_mass()[i], homog_mass[i], tol=tol) + + # The variants are genuinely distinct: their masses are not all equal. with pytest.raises(AssertionError): - assert_allclose(mass[0], mass[2], tol=tol) + assert_allclose(het_obj.get_mass(), het_obj.get_mass()[0], tol=tol) @pytest.mark.required @@ -7448,8 +7422,10 @@ def test_pick_heterogenous_objects(show_viewer): assert np.all(lift_deltas > 0.05), f"All objects should be lifted (deltas={lift_deltas})" -def _build_two_link_revolute_urdf(name, geom_tag=None, geom_attribs=None, *, links_geoms=None, links_inertial=None): - """Build a 2-link prismatic URDF file and return its path. +def _build_two_link_revolute_urdf( + name, geom_tag=None, geom_attribs=None, *, links_geoms=None, links_inertial=None, joint_type="prismatic" +): + """Build a 2-link URDF file (prismatic joint by default) and return its path. Geometry can be specified either uniformly via (geom_tag, geom_attribs) — applied identically to all links — or per-link via links_geoms for full control. @@ -7458,9 +7434,13 @@ def _build_two_link_revolute_urdf(name, geom_tag=None, geom_attribs=None, *, lin ---------- links_geoms : list of list of (tag, attribs, origin_xyz) or None Per-link geometry specs. Each link gets a list of (tag, attribs, origin_xyz) tuples. - links_inertial : list of dict or None + links_inertial : list of (dict or None) or None Per-link inertial overrides. Each dict may contain 'mass', 'ixx', 'iyy', 'izz', - 'ixy', 'ixz', 'iyz', 'origin_xyz'. If None, zero mass/inertia is used (recomputed from geometry). + 'ixy', 'ixz', 'iyz', 'origin_xyz'. A None entry leaves that link's inertial unspecified + (recomputed from geometry); None for the whole list does so for every link. + joint_type : str + Type of the joint between the two links ('prismatic', 'revolute', 'fixed', ...). A fixed joint makes the + second link a fixed child of the first (a single rigid body). """ robot = ET.Element("robot", name=name) @@ -7478,7 +7458,7 @@ def _build_two_link_revolute_urdf(name, geom_tag=None, geom_attribs=None, *, lin ET.SubElement(geom_el, tag, **attribs) if origin_xyz: ET.SubElement(group, "origin", xyz=origin_xyz) - if links_inertial: + if links_inertial and links_inertial[i_link] is not None: inertial_props = links_inertial[i_link] inertial = ET.SubElement(link, "inertial") ET.SubElement(inertial, "mass", value=str(inertial_props["mass"])) @@ -7494,12 +7474,13 @@ def _build_two_link_revolute_urdf(name, geom_tag=None, geom_attribs=None, *, lin izz=str(inertial_props.get("izz", 0)), ) - joint = ET.SubElement(robot, "joint", name="joint1", type="prismatic") + joint = ET.SubElement(robot, "joint", name="joint1", type=joint_type) ET.SubElement(joint, "parent", link="base") ET.SubElement(joint, "child", link="moving") ET.SubElement(joint, "origin", xyz="0.1 0 0") - ET.SubElement(joint, "axis", xyz="1 0 0") - ET.SubElement(joint, "limit", lower="-1.0", upper="1.0", effort="100", velocity="1.0") + if joint_type != "fixed": + ET.SubElement(joint, "axis", xyz="1 0 0") + ET.SubElement(joint, "limit", lower="-1.0", upper="1.0", effort="100", velocity="1.0") return urdfpy.URDF._from_xml(robot, robot, get_assets_dir()) @@ -7521,13 +7502,7 @@ def _build_free_body_urdf(name, com_xyz): @pytest.mark.slow # ~250s @pytest.mark.required -def test_heterogeneous_inertial_alignment(show_viewer, tol): - """Test heterogeneous articulated simulation with vertex-based and primitive collision geometries. - - Variant A splits each box primitive into two half-height sub-boxes (top/bottom), - variant B uses sphere mesh collision geometry. Verifies dynamics, mass, CoM position, - inertia matrix, joint structure, and ground contact settling. - """ +def test_align_heterogeneous_inertial(show_viewer, tol): GRAVITY = -9.81 # Variant A: sphere mesh collision with explicit inertial properties per link @@ -7585,9 +7560,12 @@ def test_heterogeneous_inertial_alignment(show_viewer, tol): ) scene.add_entity(gs.morphs.Plane()) + # align=True is requested but must be ignored for these articulated robots: a free base with a DOF-bearing child is + # not a single rigid body, so its link frames and joint-space mass coupling must be left intact (aligning it would + # misplace the moving child and drop the base coupling). The link-spacing and settling assertions below verify this. het_morph = ( - gs.morphs.URDF(file=urdf_spheres, pos=(0.5, 0, 0.08)), - gs.morphs.URDF(file=urdf_boxes, pos=(0, 0, 0.02)), + gs.morphs.URDF(file=urdf_spheres, pos=(0.5, 0, 0.08), align=True), + gs.morphs.URDF(file=urdf_boxes, pos=(0, 0, 0.02), align=True), ) het_obj = scene.add_entity( morph=het_morph, @@ -7604,18 +7582,35 @@ def test_heterogeneous_inertial_alignment(show_viewer, tol): color=(0.0, 0.0, 1.0, 0.4), ), ) - # Free-floating single-link URDF objects with different off-center COMs. Unlike the articulated robots above - # (which stay unaligned), each variant is a basic rigid object, so its link frame is moved to its own COM. + # Free-floating single-link URDF objects with different off-center COMs. Unlike the articulated robots above (whose + # requested alignment is ignored), each variant is a basic rigid object, so its link frame is moved to its own COM. FREE_POS = (3.0, 0.0, 0.2) - free_het = scene.add_entity( - morph=( - gs.morphs.URDF(file=_build_free_body_urdf("free_body_a", "0.02 0 0"), pos=FREE_POS, align=True), - gs.morphs.URDF(file=_build_free_body_urdf("free_body_b", "0 0 0.03"), pos=FREE_POS, align=True), - ), - material=gs.materials.Rigid(rho=200.0), - ) + free_morph = ( + gs.morphs.URDF(file=_build_free_body_urdf("free_body_a", "0.02 0 0"), pos=FREE_POS, align=True), + gs.morphs.URDF(file=_build_free_body_urdf("free_body_b", "0 0 0.03"), pos=FREE_POS, align=True), + ) + free_het = scene.add_entity(morph=free_morph, material=gs.materials.Rigid(rho=200.0)) + # Kinematic counterpart of the aligned free bodies. The COM/principal anchoring is applied in the base entity, so + # for the same qpos a kinematic visualization and the rigid body it tracks must place identical world geometry. + free_kin = scene.add_entity(morph=free_morph, material=gs.materials.Kinematic()) + # A free entity whose two variants are identical, so the solver resolves a single shared offset and takes the + # broadcast path (the base link offset) rather than the per-env variant offset. That shared base offset must still + # carry the COM anchoring, else the relative getter cannot strip the alignment on this path. + dup_morph = ( + gs.morphs.URDF(file=_build_free_body_urdf("free_dup_a", "0.02 0 0"), pos=FREE_POS, align=True), + gs.morphs.URDF(file=_build_free_body_urdf("free_dup_b", "0.02 0 0"), pos=FREE_POS, align=True), + ) + free_dup = scene.add_entity(morph=dup_morph, material=gs.materials.Rigid()) scene.build(n_envs=4, env_spacing=(0.0, 0.5)) + # Same absolute qpos must map to the same world geometry for the aligned rigid body and its kinematic counterpart; + # an unanchored kinematic entity would interpret the qpos in a different link frame and diverge. + free_qpos = (1.0, 0.5, 0.8, 0.6, 0.5, 0.3, 0.0) + free_het.set_qpos(free_qpos) + free_kin.set_qpos(free_qpos) + assert_allclose(free_het.get_vAABB(), free_kin.get_vAABB(), tol=tol) + scene.reset() + # Each free-body variant (env 0 variant A, env 2 variant B) is aligned to its own COM: the link origin coincides # with the COM, and the relative getter strips the alignment back to the user pose. for i_env in (0, 2): @@ -7626,6 +7621,16 @@ def test_heterogeneous_inertial_alignment(show_viewer, tol): ) assert_allclose(free_het.get_pos(relative=True, envs_idx=i_env), FREE_POS, tol=tol) + # The duplicate-variant entity takes the broadcast offset path; its relative getter must still strip the shared COM + # anchoring back to the user pose in every env (a non-anchored base offset would leave it COM-shifted). + assert_allclose( + free_dup.get_links_pos(links_idx_local=[0], ref="link_com", relative=False), + free_dup.get_links_pos(links_idx_local=[0], ref="link_origin", relative=False), + tol=tol, + ) + assert_allclose(free_dup.get_pos(relative=True), FREE_POS, tol=tol) + assert_allclose(gu.quat_to_xyz(free_dup.get_quat(relative=True)), 0.0, tol=tol) + # Relative set_pos on a boolean-masked subset of envs: each selected env's relative getter must report its target # back, stripping its own per-variant offset. mask = torch.tensor([True, False, True, False], device=gs.device)