ruff

Signed-off-by: Alain Denzler <adenzler@nvidia.com>

Author: adenzler-nvidia

newton/_src/sim/builder.py

@@ -735,11 +735,15 @@ def _process_joint_custom_attributes(
 
         Joint attributes are processed based on their declared frequency:
         - JOINT frequency: Single value per joint
-        - JOINT_DOF frequency: List of values with length equal to joint DOF count
-        - JOINT_COORD frequency: List of values with length equal to joint coordinate count
+        - JOINT_DOF frequency: List or dict of values for each DOF
+        - JOINT_COORD frequency: List or dict of values for each coordinate
 
-        For DOF and COORD attributes, values must always be provided as lists with length
-        matching the joint's DOF or coordinate count.
+        For DOF and COORD attributes, values can be:
+        - A list with length matching the joint's DOF/coordinate count (all DOFs get values)
+        - A dict mapping DOF/coord indices to values (only specified indices get values, rest use defaults)
+        - For single-DOF joints with JOINT_DOF frequency: a single Warp vector/matrix value
+
+        When using dict format, unspecified indices will be filled with the attribute's default value during finalization.
 
         Args:
             joint_index: Index of the joint
@@ -764,7 +768,7 @@ def _process_joint_custom_attributes(
                 )
 
             elif custom_attr.frequency == ModelAttributeFrequency.JOINT_DOF:
-                # List of values, one per DOF
+                # Values per DOF - can be list or dict
                 dof_start = self.joint_qd_start[joint_index]
                 if joint_index + 1 < len(self.joint_qd_start):
                     dof_end = self.joint_qd_start[joint_index + 1]
@@ -773,27 +777,54 @@ def _process_joint_custom_attributes(
 
                 dof_count = dof_end - dof_start
 
-                if not isinstance(value, (list, tuple)):
-                    raise TypeError(
-                        f"JOINT_DOF attribute '{attr_key}' must be a list with length equal to joint DOF count ({dof_count})"
-                    )
-
-                if len(value) != dof_count:
-                    raise ValueError(
-                        f"JOINT_DOF attribute '{attr_key}' has {len(value)} values but joint has {dof_count} DOFs"
-                    )
+                # Check if value is a dict (mapping DOF index to value)
+                if isinstance(value, dict):
+                    # Dict format: only specified DOF indices have values, rest use defaults
+                    for dof_offset, dof_value in value.items():
+                        if not isinstance(dof_offset, int):
+                            raise TypeError(
+                                f"JOINT_DOF attribute '{attr_key}' dict keys must be integers (DOF indices), got {type(dof_offset)}"
+                            )
+                        if dof_offset < 0 or dof_offset >= dof_count:
+                            raise ValueError(
+                                f"JOINT_DOF attribute '{attr_key}' has invalid DOF index {dof_offset} (joint has {dof_count} DOFs)"
+                            )
+                        single_attr = {attr_key: dof_value}
+                        self._process_custom_attributes(
+                            entity_index=dof_start + dof_offset,
+                            custom_attrs=single_attr,
+                            expected_frequency=ModelAttributeFrequency.JOINT_DOF,
+                        )
+                else:
+                    # List format or single value for single-DOF joints
+                    value_sanitized = value
+                    if not isinstance(value_sanitized, (list, tuple)):
+                        # Check if it's a Warp vector/matrix type
+                        if wp.types.type_is_vector(type(value_sanitized)) or wp.types.type_is_matrix(
+                            type(value_sanitized)
+                        ):
+                            value_sanitized = [value_sanitized]
+                        else:
+                            raise TypeError(
+                                f"JOINT_DOF attribute '{attr_key}' must be a list with length equal to joint DOF count ({dof_count}), "
+                                f"a dict mapping DOF indices to values, or a single Warp vector/matrix value for single-DOF joints"
+                            )
 
-                # Apply each value to its corresponding DOF
-                for i, dof_value in enumerate(value):
-                    single_attr = {attr_key: dof_value}
-                    self._process_custom_attributes(
-                        entity_index=dof_start + i,
-                        custom_attrs=single_attr,
-                        expected_frequency=ModelAttributeFrequency.JOINT_DOF,
-                    )
+                    actual = len(value_sanitized)
+                    if actual != dof_count:
+                        raise ValueError(f"JOINT_DOF '{attr_key}': got {actual}, expected {dof_count}")
+
+                    # Apply each value to its corresponding DOF
+                    for i, dof_value in enumerate(value_sanitized):
+                        single_attr = {attr_key: dof_value}
+                        self._process_custom_attributes(
+                            entity_index=dof_start + i,
+                            custom_attrs=single_attr,
+                            expected_frequency=ModelAttributeFrequency.JOINT_DOF,
+                        )
 
             elif custom_attr.frequency == ModelAttributeFrequency.JOINT_COORD:
-                # List of values, one per coordinate
+                # Values per coordinate - can be list or dict
                 coord_start = self.joint_q_start[joint_index]
                 if joint_index + 1 < len(self.joint_q_start):
                     coord_end = self.joint_q_start[joint_index + 1]
@@ -802,24 +833,45 @@ def _process_joint_custom_attributes(
 
                 coord_count = coord_end - coord_start
 
-                if not isinstance(value, (list, tuple)):
-                    raise TypeError(
-                        f"JOINT_COORD attribute '{attr_key}' must be a list with length equal to joint coordinate count ({coord_count})"
-                    )
+                # Check if value is a dict (mapping coord index to value)
+                if isinstance(value, dict):
+                    # Dict format: only specified coord indices have values, rest use defaults
+                    for coord_offset, coord_value in value.items():
+                        if not isinstance(coord_offset, int):
+                            raise TypeError(
+                                f"JOINT_COORD attribute '{attr_key}' dict keys must be integers (coord indices), got {type(coord_offset)}"
+                            )
+                        if coord_offset < 0 or coord_offset >= coord_count:
+                            raise ValueError(
+                                f"JOINT_COORD attribute '{attr_key}' has invalid coord index {coord_offset} (joint has {coord_count} coordinates)"
+                            )
+                        single_attr = {attr_key: coord_value}
+                        self._process_custom_attributes(
+                            entity_index=coord_start + coord_offset,
+                            custom_attrs=single_attr,
+                            expected_frequency=ModelAttributeFrequency.JOINT_COORD,
+                        )
+                else:
+                    # List format
+                    if not isinstance(value, (list, tuple)):
+                        raise TypeError(
+                            f"JOINT_COORD attribute '{attr_key}' must be a list with length equal to joint coordinate count ({coord_count}) "
+                            f"or a dict mapping coordinate indices to values"
+                        )
 
-                if len(value) != coord_count:
-                    raise ValueError(
-                        f"JOINT_COORD attribute '{attr_key}' has {len(value)} values but joint has {coord_count} coordinates"
-                    )
+                    if len(value) != coord_count:
+                        raise ValueError(
+                            f"JOINT_COORD attribute '{attr_key}' has {len(value)} values but joint has {coord_count} coordinates"
+                        )
 
-                # Apply each value to its corresponding coordinate
-                for i, coord_value in enumerate(value):
-                    single_attr = {attr_key: coord_value}
-                    self._process_custom_attributes(
-                        entity_index=coord_start + i,
-                        custom_attrs=single_attr,
-                        expected_frequency=ModelAttributeFrequency.JOINT_COORD,
-                    )
+                    # Apply each value to its corresponding coordinate
+                    for i, coord_value in enumerate(value):
+                        single_attr = {attr_key: coord_value}
+                        self._process_custom_attributes(
+                            entity_index=coord_start + i,
+                            custom_attrs=single_attr,
+                            expected_frequency=ModelAttributeFrequency.JOINT_COORD,
+                        )
 
             else:
                 raise ValueError(

newton/_src/solvers/mujoco/kernels.py

@@ -903,24 +903,88 @@ def update_axis_properties_kernel(
 
 
 @wp.kernel
-def update_dof_properties_kernel(
+def update_joint_dof_properties_kernel(
+    joint_qd_start: wp.array(dtype=wp.int32),
+    joint_dof_dim: wp.array2d(dtype=wp.int32),
+    joint_mjc_dof_start: wp.array(dtype=wp.int32),
+    dof_to_mjc_joint: wp.array(dtype=wp.int32),
     joint_armature: wp.array(dtype=float),
     joint_friction: wp.array(dtype=float),
-    dofs_per_world: int,
+    joint_limit_ke: wp.array(dtype=float),
+    joint_limit_kd: wp.array(dtype=float),
+    solimplimit: wp.array(dtype=vec5),
+    joints_per_world: int,
     # outputs
     dof_armature: wp.array2d(dtype=float),
     dof_frictionloss: wp.array2d(dtype=float),
+    jnt_solimp: wp.array2d(dtype=vec5),
+    jnt_solref: wp.array2d(dtype=wp.vec2),
 ):
-    """Update DOF armature and friction loss values."""
+    """Update joint DOF properties including armature, friction loss, joint impedance limits, and solref.
+
+    This kernel properly maps Newton DOFs to MuJoCo DOFs using joint_mjc_dof_start.
+    For solimplimit and solref, we use dof_to_mjc_joint since jnt_solimp/jnt_solref are per-joint in MuJoCo.
+    If solimplimit is None, jnt_solimp won't be updated (MuJoCo defaults will be preserved).
+    """
     tid = wp.tid()
-    worldid = tid // dofs_per_world
-    dof_in_world = tid % dofs_per_world
+    worldid = tid // joints_per_world
+    joint_in_world = tid % joints_per_world
+
+    lin_axis_count = joint_dof_dim[tid, 0]
+    ang_axis_count = joint_dof_dim[tid, 1]
+
+    if lin_axis_count + ang_axis_count == 0:
+        return
+
+    newton_dof_start = joint_qd_start[tid]
+    mjc_dof_start = joint_mjc_dof_start[joint_in_world]
+
+    # Get the DOF start for the template joint (world 0)
+    # dof_to_mjc_joint is only populated for template DOFs (first world)
+    template_joint_idx = joint_in_world
+    template_dof_start = joint_qd_start[template_joint_idx]
+
+    # update linear dofs
+    for i in range(lin_axis_count):
+        newton_dof_index = newton_dof_start + i
+        template_dof_index = template_dof_start + i
+        mjc_dof_index = mjc_dof_start + i
+        mjc_joint_index = dof_to_mjc_joint[template_dof_index]
+
+        # Update armature and friction (per DOF)
+        dof_armature[worldid, mjc_dof_index] = joint_armature[newton_dof_index]
+        dof_frictionloss[worldid, mjc_dof_index] = joint_friction[newton_dof_index]
+
+        # Update joint limit solref using negative convention (per joint)
+        if joint_limit_ke[newton_dof_index] > 0.0:
+            jnt_solref[worldid, mjc_joint_index] = wp.vec2(
+                -joint_limit_ke[newton_dof_index], -joint_limit_kd[newton_dof_index]
+            )
+
+        # Update solimplimit (per joint)
+        if solimplimit:
+            jnt_solimp[worldid, mjc_joint_index] = solimplimit[newton_dof_index]
 
-    # update armature
-    dof_armature[worldid, dof_in_world] = joint_armature[tid]
+    # update angular dofs
+    for i in range(ang_axis_count):
+        newton_dof_index = newton_dof_start + lin_axis_count + i
+        template_dof_index = template_dof_start + lin_axis_count + i
+        mjc_dof_index = mjc_dof_start + lin_axis_count + i
+        mjc_joint_index = dof_to_mjc_joint[template_dof_index]
+
+        # Update armature and friction (per DOF)
+        dof_armature[worldid, mjc_dof_index] = joint_armature[newton_dof_index]
+        dof_frictionloss[worldid, mjc_dof_index] = joint_friction[newton_dof_index]
+
+        # Update joint limit solref using negative convention (per joint)
+        if joint_limit_ke[newton_dof_index] > 0.0:
+            jnt_solref[worldid, mjc_joint_index] = wp.vec2(
+                -joint_limit_ke[newton_dof_index], -joint_limit_kd[newton_dof_index]
+            )
 
-    # update friction loss
-    dof_frictionloss[worldid, dof_in_world] = joint_friction[tid]
+        # Update solimplimit (per joint)
+        if solimplimit:
+            jnt_solimp[worldid, mjc_joint_index] = solimplimit[newton_dof_index]
 
 
 @wp.kernel
@@ -932,16 +996,13 @@ def update_joint_transforms_kernel(
     joint_original_axis: wp.array(dtype=wp.vec3),
     joint_child: wp.array(dtype=wp.int32),
     joint_type: wp.array(dtype=wp.int32),
-    joint_limit_ke: wp.array(dtype=float),
-    joint_limit_kd: wp.array(dtype=float),
-    joint_mjc_dof_start: wp.array(dtype=wp.int32),
+    dof_to_mjc_joint: wp.array(dtype=wp.int32),
     body_mapping: wp.array(dtype=wp.int32),
     newton_body_to_mocap_index: wp.array(dtype=wp.int32),
     joints_per_world: int,
     # outputs
     joint_pos: wp.array2d(dtype=wp.vec3),
     joint_axis: wp.array2d(dtype=wp.vec3),
-    joint_solref: wp.array2d(dtype=wp.vec2),
     body_pos: wp.array2d(dtype=wp.vec3),
     body_quat: wp.array2d(dtype=wp.quat),
     mocap_pos: wp.array2d(dtype=wp.vec3),
@@ -958,14 +1019,6 @@ def update_joint_transforms_kernel(
 
     child_xform = joint_X_c[tid]
     parent_xform = joint_X_p[tid]
-    lin_axis_count = joint_dof_dim[tid, 0]
-    ang_axis_count = joint_dof_dim[tid, 1]
-    newton_dof_start = joint_dof_start[tid]
-    mjc_dof_start = joint_mjc_dof_start[joint_in_world]
-    if mjc_dof_start == -1:
-        # this should not happen
-        wp.printf("Joint %i has no MuJoCo DOF start index\n", joint_in_world)
-        return
 
     # update body pos and quat from parent joint transform
     child = joint_child[joint_in_world]  # Newton body id
@@ -984,27 +1037,31 @@ def update_joint_transforms_kernel(
         body_pos[worldid, body_id] = tf.p
         body_quat[worldid, body_id] = rotation
 
+    lin_axis_count = joint_dof_dim[tid, 0]
+    ang_axis_count = joint_dof_dim[tid, 1]
+
+    if lin_axis_count + ang_axis_count == 0:
+        return
+
+    newton_dof_start = joint_dof_start[tid]
+    template_dof_start = joint_dof_start[joint_in_world]
+    mjc_joint_index = dof_to_mjc_joint[template_dof_start]
+
     # update linear dofs
     for i in range(lin_axis_count):
         newton_dof_index = newton_dof_start + i
         axis = joint_original_axis[newton_dof_index]
-        ai = mjc_dof_start + i
+        ai = mjc_joint_index + i
         joint_axis[worldid, ai] = wp.quat_rotate(child_xform.q, axis)
         joint_pos[worldid, ai] = child_xform.p
-        # update joint limit solref using negative convention
-        if joint_limit_ke[newton_dof_index] > 0:
-            joint_solref[worldid, ai] = wp.vec2(-joint_limit_ke[newton_dof_index], -joint_limit_kd[newton_dof_index])
 
     # update angular dofs
     for i in range(ang_axis_count):
         newton_dof_index = newton_dof_start + lin_axis_count + i
         axis = joint_original_axis[newton_dof_index]
-        ai = mjc_dof_start + lin_axis_count + i
+        ai = mjc_joint_index + lin_axis_count + i
         joint_axis[worldid, ai] = wp.quat_rotate(child_xform.q, axis)
         joint_pos[worldid, ai] = child_xform.p
-        # update joint limit solref using negative convention
-        if joint_limit_ke[newton_dof_index] > 0:
-            joint_solref[worldid, ai] = wp.vec2(-joint_limit_ke[newton_dof_index], -joint_limit_kd[newton_dof_index])
 
 
 @wp.kernel(enable_backward=False)