feat(schemas): joint-drive schema-fragment API (additive)

source/isaaclab/changelog.d/vidurv-schema-frag-jointdrive.minor.rst

@@ -0,0 +1,22 @@
+Added
+^^^^^
+
+* Added the joint-drive schema-fragment API: the
+  :class:`~isaaclab.sim.schemas.JointDriveFragment` marker and
+  :class:`~isaaclab.sim.schemas.UsdPhysicsDriveCfg` (writing the typed multi-instance
+  ``UsdPhysics.DriveAPI`` attributes). The drive fragment overrides its ``func`` with
+  :func:`~isaaclab.sim.schemas.apply_drive`, which selects the angular/linear instance, performs
+  the radian-to-degree conversion for angular drives, and skips tendon child prims.
+* Added :func:`~isaaclab.sim.schemas.apply_joint_drive_properties` (applies a list of joint-drive
+  fragments to all joint prims under a path; ``UsdPhysics.DriveAPI`` is presence-gated and applied
+  only when a :class:`~isaaclab.sim.schemas.UsdPhysicsDriveCfg` fragment is present).
+
+Changed
+^^^^^^^
+
+* Changed the spawner ``joint_drive_props`` slot
+  (:attr:`~isaaclab.sim.spawners.FileCfg.joint_drive_props`) to also accept a list of
+  :class:`~isaaclab.sim.schemas.JointDriveFragment` fragments. Legacy single cfgs continue to work
+  through a transition bridge at the from-files spawn site. Added the spawner-level
+  ``ensure_drives_exist`` flag to reproduce the legacy minimal-stiffness behaviour for the fragment
+  path.

source/isaaclab/isaaclab/sim/__init__.pyi

@@ -45,6 +45,7 @@ __all__ = [
     "DeformableBodyPropertiesCfg",
     "FixedTendonPropertiesCfg",
     "JointDriveBaseCfg",
+    "JointDriveFragment",
     "MassPropertiesCfg",
     "MeshCollisionPropertiesCfg",
     "MujocoJointDrivePropertiesCfg",
@@ -61,7 +62,10 @@ __all__ = [
     "RigidBodyBaseCfg",
     "RigidBodyFragment",
     "SchemaFragment",
+    "UsdPhysicsDriveCfg",
     "UsdPhysicsRigidBodyCfg",
+    "apply_drive",
+    "apply_joint_drive_properties",
     "apply_namespaced",
     "apply_rigid_body_properties",
     "SDFMeshPropertiesCfg",
@@ -220,6 +224,7 @@ from .schemas import (
     DeformableBodyPropertiesCfg,
     FixedTendonPropertiesCfg,
     JointDriveBaseCfg,
+    JointDriveFragment,
     MassPropertiesCfg,
     MeshCollisionPropertiesCfg,
     PhysxJointDrivePropertiesCfg,
@@ -231,8 +236,11 @@ from .schemas import (
     SpatialTendonPropertiesCfg,
     TriangleMeshPropertiesCfg,
     TriangleMeshSimplificationPropertiesCfg,
+    UsdPhysicsDriveCfg,
     UsdPhysicsRigidBodyCfg,
     activate_contact_sensors,
+    apply_drive,
+    apply_joint_drive_properties,
     apply_namespaced,
     apply_rigid_body_properties,
     define_articulation_root_properties,

source/isaaclab/isaaclab/sim/converters/mesh_converter.py

@@ -185,7 +185,7 @@ def _convert_asset(self, cfg: MeshConverterCfg):
         # apply mass properties
         if cfg.mass_props is not None:
             schemas.define_mass_properties(prim_path=xform_prim.GetPath(), cfg=cfg.mass_props, stage=stage)
-        # apply rigid body properties (transition routing: fragment list -> apply_*; legacy cfg -> define_*)
+        # apply rigid body properties (transition shim, remove later: fragment list -> apply_*; legacy cfg -> define_*)
         if cfg.rigid_props is not None:
             rigid_frags = cfg.rigid_props if isinstance(cfg.rigid_props, (list, tuple)) else [cfg.rigid_props]
             if rigid_frags and all(isinstance(f, schemas.SchemaFragment) for f in rigid_frags):

source/isaaclab/isaaclab/sim/schemas/__init__.pyi

@@ -8,6 +8,8 @@ __all__ = [
     "PHYSX_MESH_COLLISION_CFGS",
     "USD_MESH_COLLISION_CFGS",
     "activate_contact_sensors",
+    "apply_drive",
+    "apply_joint_drive_properties",
     "apply_namespaced",
     "apply_rigid_body_properties",
     "define_actuator_properties",
@@ -33,10 +35,12 @@ __all__ = [
     "DeformableBodyPropertiesBaseCfg",
     "DeformableBodyPropertiesCfg",
     "JointDriveBaseCfg",
+    "JointDriveFragment",
     "MassPropertiesCfg",
     "MeshCollisionBaseCfg",
     "RigidBodyFragment",
     "SchemaFragment",
+    "UsdPhysicsDriveCfg",
     "UsdPhysicsRigidBodyCfg",
     "MujocoJointDrivePropertiesCfg",
     "MujocoRigidBodyPropertiesCfg",
@@ -55,6 +59,8 @@ from .schemas import (
     PHYSX_MESH_COLLISION_CFGS,
     USD_MESH_COLLISION_CFGS,
     activate_contact_sensors,
+    apply_drive,
+    apply_joint_drive_properties,
     apply_namespaced,
     apply_rigid_body_properties,
     define_articulation_root_properties,
@@ -84,11 +90,13 @@ from .schemas_cfg import (
     DeformableBodyPropertiesBaseCfg,
     DeformableBodyPropertiesCfg,
     JointDriveBaseCfg,
+    JointDriveFragment,
     MassPropertiesCfg,
     MeshCollisionBaseCfg,
     RigidBodyBaseCfg,
     RigidBodyFragment,
     SchemaFragment,
+    UsdPhysicsDriveCfg,
     UsdPhysicsRigidBodyCfg,
 )
 

source/isaaclab/isaaclab/sim/schemas/schemas.py

@@ -792,6 +792,216 @@ def activate_contact_sensors(prim_path: str, threshold: float = 0.0, stage: Usd.
 """
 
 
+def apply_drive(cfg, prim_path: str, stage: Usd.Stage | None = None) -> bool:
+    """Apply a :class:`~isaaclab.sim.schemas.UsdPhysicsDriveCfg` fragment to a single joint prim.
+
+    This is the override ``func`` for the ``UsdPhysics.DriveAPI`` fragment: the drive attributes
+    live under a multi-instance schema, so the generic :func:`apply_namespaced` writer cannot be
+    used. The writer reproduces the solver-common drive logic of
+    :func:`modify_joint_drive_properties`:
+
+    * Selects the drive instance: ``"angular"`` for a revolute joint, ``"linear"`` for a prismatic
+      joint. For any other prim type, the function is a no-op and returns ``False``.
+    * Skips tendon child prims (prims carrying ``PhysxTendonAxisAPI`` without
+      ``PhysxTendonAxisRootAPI``), returning ``False``.
+    * Applies ``UsdPhysics.DriveAPI`` for the selected instance (presence-gated -- only applied when
+      this fragment is present).
+    * Converts angular-drive :attr:`stiffness` and :attr:`damping` from radians to degrees
+      (``N·m/rad`` -> ``N·m/deg`` and ``N·m·s/rad`` -> ``N·m·s/deg``); linear drives are written
+      as-is.
+    * Writes the typed ``drive:<inst>:physics:{type,maxForce,stiffness,damping}`` attributes,
+      mapping the :attr:`drive_type` field to the USD attribute named ``type``.
+
+    Args:
+        cfg: The :class:`~isaaclab.sim.schemas.UsdPhysicsDriveCfg` fragment to apply.
+        prim_path: The joint prim path to author on.
+        stage: The stage where to find the prim. Defaults to None, in which case the current
+            stage is used.
+
+    Returns:
+        True if the drive was applied to a joint prim, False if the prim is not a revolute or
+        prismatic joint (or is a tendon child).
+    """
+    if stage is None:
+        stage = get_current_stage()
+    prim = stage.GetPrimAtPath(prim_path)
+    if not prim.IsValid():
+        raise ValueError(f"Prim path '{prim_path}' is not valid.")
+
+    # select the drive instance based on the joint type
+    if prim.IsA(UsdPhysics.RevoluteJoint):
+        drive_api_name = "angular"
+    elif prim.IsA(UsdPhysics.PrismaticJoint):
+        drive_api_name = "linear"
+    else:
+        return False
+    # skip tendon child prims (carry PhysxTendonAxisAPI but not the root API)
+    applied_schemas_str = str(prim.GetAppliedSchemas())
+    if "PhysxTendonAxisAPI" in applied_schemas_str and "PhysxTendonAxisRootAPI" not in applied_schemas_str:
+        return False
+
+    # apply the multi-instance drive API (presence-gated anchor for the joint-drive family)
+    usd_drive_api = UsdPhysics.DriveAPI(prim, drive_api_name)
+    if not usd_drive_api:
+        usd_drive_api = UsdPhysics.DriveAPI.Apply(prim, drive_api_name)
+
+    # gather the solver-common drive fields
+    drive_type = cfg.drive_type
+    max_force = cfg.max_force
+    stiffness = cfg.stiffness
+    damping = cfg.damping
+
+    # angular drives use degree units in USD; convert stiffness/damping from radian units
+    is_linear_drive = prim.IsA(UsdPhysics.PrismaticJoint)
+    if not is_linear_drive:
+        if stiffness is not None:
+            # N-m/rad --> N-m/deg
+            stiffness = stiffness * math.pi / 180.0
+        if damping is not None:
+            # N-m-s/rad --> N-m-s/deg
+            damping = damping * math.pi / 180.0
+
+    # ``drive_type`` is a permanent inline carve-out: the USD attribute is named ``type``
+    # (a Python keyword-like name we cannot use as a cfg field). All other solver-common
+    # joint-drive fields follow the snake_case = camelCase convention.
+    for field_name, value in (
+        ("drive_type", drive_type),
+        ("max_force", max_force),
+        ("stiffness", stiffness),
+        ("damping", damping),
+    ):
+        if value is None:
+            continue
+        usd_attr_name = "type" if field_name == "drive_type" else field_name
+        safe_set_attribute_on_usd_schema(usd_drive_api, usd_attr_name, value, camel_case=True)
+
+    return True
+
+
+def apply_joint_drive_properties(
+    prim_path: str, fragments, stage: Usd.Stage | None = None, ensure_drives_exist: bool = False
+) -> bool:
+    """Apply a list of joint-drive fragments to all joint prims under a prim path.
+
+    Mirrors the recursion behaviour of :func:`modify_joint_drive_properties` (decorated with
+    :func:`~isaaclab.sim.utils.apply_nested`): the prim path and all its descendants are visited,
+    and the fragments are dispatched to every revolute/prismatic joint prim found. As soon as a
+    prim is successfully handled, its children are not descended into (nested joints are not
+    allowed).
+
+    Unlike :func:`apply_rigid_body_properties`, the joint-drive family has no implicit anchor:
+    ``UsdPhysics.DriveAPI`` is *presence-gated* and applied only by :func:`apply_drive` when a
+    :class:`~isaaclab.sim.schemas.UsdPhysicsDriveCfg` fragment is present in ``fragments``. Each
+    fragment is dispatched via its :attr:`~isaaclab.sim.schemas.SchemaFragment.func`, so backend
+    fragments carry backend-specific funcs and core never imports a backend.
+
+    Args:
+        prim_path: The root prim path to search for joint prims under.
+        fragments: An iterable of :class:`~isaaclab.sim.schemas.JointDriveFragment` instances.
+        stage: The stage where to find the prim. Defaults to None, in which case the current
+            stage is used.
+        ensure_drives_exist: If True, write a minimal stiffness (``1e-3``) to any drive whose
+            authored stiffness *and* damping are both zero, so that backends (e.g. Newton) treat
+            the drive as active. Reproduces the legacy
+            :attr:`~isaaclab.sim.schemas.JointDriveBaseCfg.ensure_drives_exist` behaviour. This is
+            a spawner-level flag, not a fragment field.
+
+    Returns:
+        True if the fragments were applied to at least one joint prim, False otherwise.
+    """
+    if stage is None:
+        stage = get_current_stage()
+
+    fragments = list(fragments)
+    # detect whether a UsdPhysics.DriveAPI fragment is present (governs presence-gating + the
+    # ensure_drives_exist behaviour, which only makes sense for the solver-common drive fragment)
+    drive_cfg = next((f for f in fragments if isinstance(f, schemas_cfg.UsdPhysicsDriveCfg)), None)
+
+    root_prim = stage.GetPrimAtPath(prim_path)
+    if not root_prim.IsValid():
+        raise ValueError(f"Prim path '{prim_path}' is not valid.")
+
+    count_success = 0
+    instanced_prim_paths = []
+    all_prims = [root_prim]
+    while len(all_prims) > 0:
+        child_prim = all_prims.pop(0)
+        child_prim_path = child_prim.GetPath().pathString
+        # skip instanced prims (cannot author on prototypes)
+        if child_prim.IsInstance():
+            instanced_prim_paths.append(child_prim_path)
+            continue
+        # a prim is a valid joint-drive target only if it is a revolute/prismatic joint
+        is_joint = child_prim.IsA(UsdPhysics.RevoluteJoint) or child_prim.IsA(UsdPhysics.PrismaticJoint)
+        if not is_joint:
+            all_prims += child_prim.GetChildren()
+            continue
+        # skip tendon-child joints (PhysxTendonAxisAPI without the root API) wholesale: no fragment
+        # may author on them, matching the legacy modify_joint_drive_properties writer. apply_drive
+        # guards itself too, but the physxJoint/mjc fragments would otherwise leak onto them.
+        applied_schemas_str = str(child_prim.GetAppliedSchemas())
+        if "PhysxTendonAxisAPI" in applied_schemas_str and "PhysxTendonAxisRootAPI" not in applied_schemas_str:
+            all_prims += child_prim.GetChildren()
+            continue
+        # dispatch each fragment via its func
+        success = False
+        for cfg in fragments:
+            func = cfg.func if callable(cfg.func) else string_to_callable(cfg.func)
+            if func(cfg, child_prim_path, stage):
+                success = True
+        # seed a minimal stiffness only when the drive fragment authored neither stiffness nor
+        # damping and the resulting drive is fully passive
+        if ensure_drives_exist and drive_cfg is not None and success:
+            _ensure_drive_exists(drive_cfg, child_prim)
+        if success:
+            count_success += 1
+        else:
+            all_prims += child_prim.GetChildren()
+
+    if count_success == 0:
+        logger.warning(
+            f"Could not apply joint-drive properties on any prims under: '{prim_path}'."
+            " This might be because no revolute/prismatic joint prims were found, or they are"
+            f" instanced. Discovered list of instanced prim paths: {instanced_prim_paths}"
+        )
+    return count_success > 0
+
+
+def _ensure_drive_exists(drive_cfg, prim) -> None:
+    """Seed a minimal stiffness on a fully-passive drive so backends treat it as active.
+
+    Reproduces the legacy ``ensure_drives_exist`` behaviour: if the drive fragment authored
+    neither :attr:`stiffness` nor :attr:`damping` and the authored drive currently has zero
+    (or unset) stiffness *and* damping, write a minimal stiffness of ``1e-3`` directly to the
+    drive API (converted to degree units for angular drives, matching :func:`apply_drive`). The
+    fragment is not mutated, so this is safe across multiple joint prims sharing one fragment.
+
+    Args:
+        drive_cfg: The :class:`~isaaclab.sim.schemas.UsdPhysicsDriveCfg` fragment.
+        prim: The joint prim being authored.
+    """
+    if drive_cfg.stiffness is not None or drive_cfg.damping is not None:
+        return
+    if prim.IsA(UsdPhysics.RevoluteJoint):
+        drive_api_name = "angular"
+    elif prim.IsA(UsdPhysics.PrismaticJoint):
+        drive_api_name = "linear"
+    else:
+        return
+    usd_drive_api = UsdPhysics.DriveAPI(prim, drive_api_name)
+    if not usd_drive_api:
+        usd_drive_api = UsdPhysics.DriveAPI.Apply(prim, drive_api_name)
+    cur_stiffness = usd_drive_api.GetStiffnessAttr().Get()
+    cur_damping = usd_drive_api.GetDampingAttr().Get()
+    if (cur_stiffness is None or cur_stiffness == 0.0) and (cur_damping is None or cur_damping == 0.0):
+        # mirror the legacy writer: 1e-3 is set before the rad->deg conversion, so an angular
+        # drive ends up with ``1e-3 * pi / 180``.
+        stiffness = 1e-3
+        if drive_api_name == "angular":
+            stiffness = stiffness * math.pi / 180.0
+        safe_set_attribute_on_usd_schema(usd_drive_api, "stiffness", stiffness, camel_case=True)
+
+
 @apply_nested
 def modify_joint_drive_properties(
     prim_path: str, cfg: schemas_cfg.JointDriveBaseCfg, stage: Usd.Stage | None = None