Greptile Overview

Greptile Summary

This PR fixes a PyTorch advanced indexing bug in EMAJointPositionToLimitsAction.reset() that caused shape mismatch errors when resetting environments with a subset of joints enabled. The issue occurred because advanced indexing with 2D env_ids tensors produced incompatible shapes during assignment operations. The fix introduces conditional code paths: when env_ids is None, it uses efficient slice indexing, and when env_ids is a tensor, it uses explicit 2D indexing with env_ids[:, None] combined with .view(len(env_ids), -1) to ensure proper shape compatibility. This action class is part of the MDP (Markov Decision Process) action space framework in IsaacLab, where it maps normalized actions to joint position limits with exponential moving average smoothing for robot control tasks.

Important Files Changed

Confidence score: 4/5

• This PR addresses a specific bug with a targeted fix that should resolve the shape mismatch issue without introducing side effects
• Score reflects that while the fix is well-reasoned and addresses the root cause, there are no accompanying tests to validate the fix works correctly for all edge cases (acknowledged by unchecked test checkbox in PR description)
• Pay close attention to source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py - ensure the conditional logic correctly handles both code paths and that the .view(len(env_ids), -1) reshape produces expected dimensions for various joint subset configurations

Sequence Diagram

sequenceDiagram
    participant User
    participant Environment
    participant EMAJointPositionToLimitsAction
    participant AssetData
    
    User->>Environment: reset(env_ids)
    Environment->>EMAJointPositionToLimitsAction: reset(env_ids)
    
    alt env_ids is None
        EMAJointPositionToLimitsAction->>EMAJointPositionToLimitsAction: super().reset(slice(None))
        EMAJointPositionToLimitsAction->>AssetData: Get joint_pos[:, joint_ids]
        AssetData-->>EMAJointPositionToLimitsAction: Return joint positions
        EMAJointPositionToLimitsAction->>EMAJointPositionToLimitsAction: Set _prev_applied_actions[:]
    else env_ids is tensor
        EMAJointPositionToLimitsAction->>EMAJointPositionToLimitsAction: super().reset(env_ids)
        EMAJointPositionToLimitsAction->>AssetData: Get joint_pos[env_ids[:, None], joint_ids]
        Note over EMAJointPositionToLimitsAction: Advanced indexing with 2D env_ids causes shape mismatch
        AssetData-->>EMAJointPositionToLimitsAction: Return joint positions (wrong shape)
        EMAJointPositionToLimitsAction->>EMAJointPositionToLimitsAction: view(len(env_ids), -1)
        EMAJointPositionToLimitsAction->>EMAJointPositionToLimitsAction: Set _prev_applied_actions[env_ids, :]
        Note over EMAJointPositionToLimitsAction: Left-hand side shape is ill-formed due to 2D env_ids
    end
    
    EMAJointPositionToLimitsAction-->>Environment: Reset complete
    Environment-->>User: Environment ready

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

Looks good.