[Bug Report] `find_joints` in `isaaclab_newton` returns incorrect mask/indices when `joint_subset` is provided

[fanzongsao]

Describe the bug

Affected Branches: dev/newton & feature/newton
Filepath: source/isaaclab_newton/isaaclab_newton/assets/utils/shared.py

The find_joints function calculates the returned mask and indices incorrectly when the joint_subset argument is not None.

The issue arises from a mismatch in index reference frames:

1. string_utils.resolve_matching_names returns indices relative to the provided joint_subset (local indices).
2. However, find_joints uses these local indices to populate a boolean mask that corresponds to the full joint_names list (global mask).

This results in the wrong joints being marked as True in the mask. Specifically, the mask simply marks the first indices as True (where is the number of matches), regardless of where those joints actually are in the robot's full joint list.

Steps to reproduce

Tested on AntWarpEnv in isaaclab_newton/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py.

Modified __init__ to test a subset query:

class AntWarpEnv(LocomotionWarpEnv):
    cfg: AntWarpEnvCfg

    def __init__(self, cfg: AntWarpEnvCfg, render_mode: str | None = None, **kwargs):
        super().__init__(cfg, render_mode, **kwargs)
        
        # Reproduce Bug: searching within a specific subset
        # Ant joint_names: ['front_left_leg', 'front_left_foot', 'front_right_leg', 'front_right_foot', 'left_back_leg', 'left_back_foot', 'right_back_leg', 'right_back_foot']
        target_subset = ['front_left_leg', 'front_left_foot', 'left_back_leg', 'left_back_foot']
        
        self._joint_dof_mask_with_subset, _, self._joint_dof_idx = self.robot.find_joints(
            ".*", 
            joint_subset=target_subset
        )

        dof_mask = self._joint_dof_mask_with_subset.numpy()
        print(f"Target Subset: {target_subset}")
        print(f"Result Mask:   {dof_mask}")

Output:

[ True  True  True  True False False False False]

Analysis:
The output mask corresponds to indices [0, 1, 2, 3].

• Actual (Bug): It selected front_left_leg, front_left_foot, front_right_leg, front_right_foot.
• Expected: It should select the joints specified in the subset: front_left_leg (0), front_left_foot (1), left_back_leg (4), left_back_foot (5).
• Correct Mask Should Be: [ True True False False True True False False]

Additional context

Proposed Solution

We need to map the matched names back to their global indices in joint_names when a subset is used.

Here is a zero-overhead fix for the default case (joint_subset=None):

def find_joints(
    joint_names: list[str],
    name_keys: str | Sequence[str],
    joint_subset: list[str] | None = None,
    preserve_order: bool = False,
    device: str = "cuda:0",
) -> tuple[wp.array, list[str], list[int]]:
    
    # Use subset if provided, otherwise default to full list
    search_target = joint_subset if joint_subset is not None else joint_names

    # 1. Search in the target (subset or full)
    indices, names = string_utils.resolve_matching_names(name_keys, search_target, preserve_order)

    # 2. FIX: If a subset was explicitly provided, map names back to global indices
    # This block is skipped if joint_subset is None, preserving zero overhead for default usage.
    if joint_subset is not None:
         indices = [joint_names.index(n) for n in names]

    # 3. Create global mask
    mask = np.zeros(len(joint_names), dtype=bool)
    mask[indices] = True
    mask = wp.array(mask, dtype=wp.bool, device=device)
    
    return mask, names, indices

Limitation of the Fix regarding preserve_order

There is a subtle side effect with this solution regarding the preserve_order parameter when a joint_subset is used:

When preserve_order=False (default), the function resolve_matching_names returns matches in the order they appear in the search target (joint_subset). If joint_subset is not sorted in the same order as the global joint_names, the returned global indices will follow the joint_subset order, not the joint_names order.

Essentially, preserve_order=False effectively becomes "preserve subset order" rather than "preserve global order" in this specific edge case.

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have checked that the issue is not in running Isaac Sim itself and is related to the repo

Acceptance Criteria

• find_joints should correctly set the mask bits corresponding to global joint_names indices even when joint_subset is provided.
• The fix should incur no performance penalty when joint_subset is None (the default path).

[RandomOakForest]

Thank you for posting this. The team will review it.

[AntoineRichard]

Thanks a lot for reporting this!

Here is a fix and some tests: #4456

@fanzongsao if you have the chance, let us know if it works for you! It comes on top of a much larger PR that should in theory bring back rigid objects and fixes a couple more bugs in the articulation.

[fanzongsao]

It works for me. Thank you👍 @AntoineRichard