Fix ANYmal symmetry augmentation for non-PhysX joint orders

source/isaaclab_tasks/changelog.d/fix-anymal-symmetry-joint-order.fixed.rst

@@ -0,0 +1,3 @@
+* Fixed the ANYmal symmetry augmentation in :mod:`isaaclab_tasks.core.velocity.mdp.symmetry.anymal` mirroring
+  the wrong joints when the articulation's joint order is not the PhysX order, such as on the default Newton
+  backend. The left-right and front-back joint permutations are now resolved from the joint names.

source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/symmetry/anymal.py

@@ -41,6 +41,11 @@ def compute_symmetric_states(
     Returns:
         Augmented observations and actions tensors, or None if the respective input was None.
     """
+    if obs is None and actions is None:
+        return None, None
+
+    # resolve the mirrored joints from the robot's joint names (the order depends on the backend)
+    joint_names = tuple(env.unwrapped.scene["robot"].joint_names)
 
     # observations
     if obs is not None:
@@ -52,14 +57,16 @@ def compute_symmetric_states(
         # -- original
         obs_aug["policy"][:batch_size] = obs["policy"][:]
         # -- left-right
-        obs_aug["policy"][batch_size : 2 * batch_size] = _transform_policy_obs_left_right(env.unwrapped, obs["policy"])
+        obs_aug["policy"][batch_size : 2 * batch_size] = _transform_policy_obs_left_right(
+            env.unwrapped, obs["policy"], joint_names
+        )
         # -- front-back
         obs_aug["policy"][2 * batch_size : 3 * batch_size] = _transform_policy_obs_front_back(
-            env.unwrapped, obs["policy"]
+            env.unwrapped, obs["policy"], joint_names
         )
         # -- diagonal
         obs_aug["policy"][3 * batch_size :] = _transform_policy_obs_front_back(
-            env.unwrapped, obs_aug["policy"][batch_size : 2 * batch_size]
+            env.unwrapped, obs_aug["policy"][batch_size : 2 * batch_size], joint_names
         )
     else:
         obs_aug = None
@@ -72,11 +79,13 @@ def compute_symmetric_states(
         # -- original
         actions_aug[:batch_size] = actions[:]
         # -- left-right
-        actions_aug[batch_size : 2 * batch_size] = _transform_actions_left_right(actions)
+        actions_aug[batch_size : 2 * batch_size] = _transform_actions_left_right(actions, joint_names)
         # -- front-back
-        actions_aug[2 * batch_size : 3 * batch_size] = _transform_actions_front_back(actions)
+        actions_aug[2 * batch_size : 3 * batch_size] = _transform_actions_front_back(actions, joint_names)
         # -- diagonal
-        actions_aug[3 * batch_size :] = _transform_actions_front_back(actions_aug[batch_size : 2 * batch_size])
+        actions_aug[3 * batch_size :] = _transform_actions_front_back(
+            actions_aug[batch_size : 2 * batch_size], joint_names
+        )
     else:
         actions_aug = None
 
@@ -88,7 +97,9 @@ def compute_symmetric_states(
 """
 
 
-def _transform_policy_obs_left_right(env: ManagerBasedRLEnv, obs: torch.Tensor) -> torch.Tensor:
+def _transform_policy_obs_left_right(
+    env: ManagerBasedRLEnv, obs: torch.Tensor, joint_names: tuple[str, ...]
+) -> torch.Tensor:
     """Apply a left-right symmetry transformation to the observation tensor.
 
     This function modifies the given observation tensor by applying transformations
@@ -101,15 +112,18 @@ def _transform_policy_obs_left_right(env: ManagerBasedRLEnv, obs: torch.Tensor)
     Args:
         env: The environment instance from which the observation is obtained.
         obs: The observation tensor to be transformed.
+        joint_names: The robot's joint names, in the order of the joint observations.
 
     Returns:
         The transformed observation tensor with left-right symmetry applied.
     """
-    perm, sign = _policy_obs_symmetry("left_right", obs.shape[1], _has_height_scan(env), str(obs.device))
+    perm, sign = _policy_obs_symmetry("left_right", obs.shape[1], _has_height_scan(env), joint_names, str(obs.device))
     return obs[:, perm] * sign
 
 
-def _transform_policy_obs_front_back(env: ManagerBasedRLEnv, obs: torch.Tensor) -> torch.Tensor:
+def _transform_policy_obs_front_back(
+    env: ManagerBasedRLEnv, obs: torch.Tensor, joint_names: tuple[str, ...]
+) -> torch.Tensor:
     """Applies a front-back symmetry transformation to the observation tensor.
 
     This function modifies the given observation tensor by applying transformations
@@ -121,11 +135,12 @@ def _transform_policy_obs_front_back(env: ManagerBasedRLEnv, obs: torch.Tensor)
     Args:
         env: The environment instance from which the observation is obtained.
         obs: The observation tensor to be transformed.
+        joint_names: The robot's joint names, in the order of the joint observations.
 
     Returns:
         The transformed observation tensor with front-back symmetry applied.
     """
-    perm, sign = _policy_obs_symmetry("front_back", obs.shape[1], _has_height_scan(env), str(obs.device))
+    perm, sign = _policy_obs_symmetry("front_back", obs.shape[1], _has_height_scan(env), joint_names, str(obs.device))
     return obs[:, perm] * sign
 
 
@@ -134,7 +149,7 @@ def _transform_policy_obs_front_back(env: ManagerBasedRLEnv, obs: torch.Tensor)
 """
 
 
-def _transform_actions_left_right(actions: torch.Tensor) -> torch.Tensor:
+def _transform_actions_left_right(actions: torch.Tensor, joint_names: tuple[str, ...]) -> torch.Tensor:
     """Applies a left-right symmetry transformation to the actions tensor.
 
     This function modifies the given actions tensor by applying transformations
@@ -144,15 +159,16 @@ def _transform_actions_left_right(actions: torch.Tensor) -> torch.Tensor:
 
     Args:
         actions: The actions tensor to be transformed.
+        joint_names: The robot's joint names, in the order of the actions.
 
     Returns:
         The transformed actions tensor with left-right symmetry applied.
     """
-    perm, sign = _joint_symmetry("left_right", str(actions.device))
+    perm, sign = _joint_symmetry("left_right", joint_names, str(actions.device))
     return actions[:, perm] * sign
 
 
-def _transform_actions_front_back(actions: torch.Tensor) -> torch.Tensor:
+def _transform_actions_front_back(actions: torch.Tensor, joint_names: tuple[str, ...]) -> torch.Tensor:
     """Applies a front-back symmetry transformation to the actions tensor.
 
     This function modifies the given actions tensor by applying transformations
@@ -162,30 +178,25 @@ def _transform_actions_front_back(actions: torch.Tensor) -> torch.Tensor:
 
     Args:
         actions: The actions tensor to be transformed.
+        joint_names: The robot's joint names, in the order of the actions.
 
     Returns:
         The transformed actions tensor with front-back symmetry applied.
     """
-    perm, sign = _joint_symmetry("front_back", str(actions.device))
+    perm, sign = _joint_symmetry("front_back", joint_names, str(actions.device))
     return actions[:, perm] * sign
 
 
 """
 Helper functions for symmetry.
 
-In Isaac Sim, the joint ordering is as follows:
-[
-    'LF_HAA', 'LH_HAA', 'RF_HAA', 'RH_HAA',
-    'LF_HFE', 'LH_HFE', 'RF_HFE', 'RH_HFE',
-    'LF_KFE', 'LH_KFE', 'RF_KFE', 'RH_KFE'
-]
+The ANYmal joints are named ``<side><end>_<joint>``, where the side is ``L`` (left) or ``R`` (right), the end is
+``F`` (front) or ``H`` (hind), and the joint is ``HAA`` (hip abduction/adduction), ``HFE`` (hip flexion/extension)
+or ``KFE`` (knee flexion/extension). The joint order of the articulation depends on the physics backend and on
+:attr:`~isaaclab.assets.ArticulationCfg.joint_ordering`, so the mirrored joints are resolved by name:
 
-Correspondingly, the joint ordering for the ANYmal robot is:
-
-* LF = left front --> [0, 4, 8]
-* LH = left hind --> [1, 5, 9]
-* RF = right front --> [2, 6, 10]
-* RH = right hind --> [3, 7, 11]
+* left-right: swap ``L`` and ``R`` and negate the ``HAA`` joints.
+* front-back: swap ``F`` and ``H`` and negate the ``HFE`` and ``KFE`` joints.
 
 Each transform is a column permutation and a sign per column: ``out[:, i] = sign[i] * x[:, perm[i]]``.
 """
@@ -195,16 +206,9 @@ def _transform_actions_front_back(actions: torch.Tensor) -> torch.Tensor:
     "left_right": (1, -1, 1, -1, 1, -1, 1, -1, 1, 1, -1, -1),
     "front_back": (-1, 1, 1, 1, -1, -1, -1, 1, 1, -1, 1, -1),
 }
-# source joint of each switched joint: left <-> right swaps LF/RF and LH/RH, front <-> back swaps LF/LH and RF/RH
-_JOINT_PERMS = {
-    "left_right": (2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9),
-    "front_back": (1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10),
-}
-# left <-> right flips the HAA joints, front <-> back flips the HFE and KFE joints
-_JOINT_SIGNS = {
-    "left_right": (-1,) * 4 + (1,) * 8,
-    "front_back": (1,) * 4 + (-1,) * 8,
-}
+# side and end letters swapped, and joint types negated, per transform
+_JOINT_SWAPS = {"left_right": str.maketrans("LR", "RL"), "front_back": str.maketrans("FH", "HF")}
+_NEGATED_JOINTS = {"left_right": ("HAA",), "front_back": ("HFE", "KFE")}
 # height-scan grid dimension flipped per transform, for the (11, 17) grid
 _HEIGHT_SCAN_FLIP_DIM = {"left_right": 0, "front_back": 1}
 
@@ -214,25 +218,33 @@ def _has_height_scan(env: ManagerBasedRLEnv) -> bool:
     return "height_scan" in env.observation_manager.active_terms["policy"]
 
 
+def _joint_permutation(kind: str, joint_names: tuple[str, ...]) -> tuple[list[int], list[float]]:
+    """Source joint index and sign of each joint for a transform, resolved from the joint names."""
+    perm = [joint_names.index(name[:2].translate(_JOINT_SWAPS[kind]) + name[2:]) for name in joint_names]
+    sign = [-1.0 if name[3:] in _NEGATED_JOINTS[kind] else 1.0 for name in joint_names]
+    return perm, sign
+
+
 @functools.cache
-def _joint_symmetry(kind: str, device: str) -> tuple[torch.Tensor, torch.Tensor]:
-    """Column permutation and signs of a transform over the 12 joints, cached per device."""
-    return (
-        torch.tensor(_JOINT_PERMS[kind], dtype=torch.long, device=device),
-        torch.tensor(_JOINT_SIGNS[kind], device=device),
-    )
+def _joint_symmetry(kind: str, joint_names: tuple[str, ...], device: str) -> tuple[torch.Tensor, torch.Tensor]:
+    """Column permutation and signs of a transform over the joints, cached per joint order and device."""
+    perm, sign = _joint_permutation(kind, joint_names)
+    return torch.tensor(perm, dtype=torch.long, device=device), torch.tensor(sign, device=device)
 
 
 @functools.cache
-def _policy_obs_symmetry(kind: str, num_obs: int, height_scan: bool, device: str) -> tuple[torch.Tensor, torch.Tensor]:
-    """Column permutation and signs of a transform over the policy observation, cached per device."""
+def _policy_obs_symmetry(
+    kind: str, num_obs: int, height_scan: bool, joint_names: tuple[str, ...], device: str
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Column permutation and signs of a transform over the policy observation, cached per joint order and device."""
     perm = list(range(num_obs))
     sign = [1] * num_obs
     sign[:12] = _BASE_SIGNS[kind]
     # joint positions, joint velocities, and last actions
+    joint_perm, joint_sign = _joint_permutation(kind, joint_names)
     for start in (12, 24, 36):
-        perm[start : start + 12] = [start + joint for joint in _JOINT_PERMS[kind]]
-        sign[start : start + 12] = _JOINT_SIGNS[kind]
+        perm[start : start + 12] = [start + joint for joint in joint_perm]
+        sign[start : start + 12] = joint_sign
     # note: this is hard-coded for grid-pattern of ordering "xy" and size (1.6, 1.0)
     if height_scan:
         grid = torch.arange(48, 235).view(11, 17).flip(dims=[_HEIGHT_SCAN_FLIP_DIM[kind]])

source/isaaclab_tasks/test/core/test_anymal_symmetry.py

@@ -13,6 +13,20 @@
 
 from isaaclab_tasks.core.velocity.mdp.symmetry import anymal
 
+# Native joint orders of ANYmal-D on the two physics backends.
+PHYSX_ORDER = [
+    "LF_HAA", "LH_HAA", "RF_HAA", "RH_HAA",
+    "LF_HFE", "LH_HFE", "RF_HFE", "RH_HFE",
+    "LF_KFE", "LH_KFE", "RF_KFE", "RH_KFE",
+]  # fmt: skip
+NEWTON_ORDER = [
+    "LF_HAA", "LF_HFE", "LF_KFE",
+    "LH_HAA", "LH_HFE", "LH_KFE",
+    "RF_HAA", "RF_HFE", "RF_KFE",
+    "RH_HAA", "RH_HFE", "RH_KFE",
+]  # fmt: skip
+
+
 # Reference: the index-assignment formulation the permutation tables must reproduce.
 _LEFT = [0, 4, 8, 1, 5, 9]
 _RIGHT = [2, 6, 10, 3, 7, 11]
@@ -72,13 +86,16 @@ def _reference_augment(obs: torch.Tensor, actions: torch.Tensor, height_scan: bo
 
 @pytest.mark.parametrize("height_scan", [False, True])
 def test_compute_symmetric_states_matches_reference(height_scan: bool):
-    """Verify the cached permutation and sign tables reproduce the index-assignment formulation."""
+    """Verify the cached permutation and sign tables reproduce the index-assignment formulation in the PhysX order."""
     num_envs = 5
     num_obs = 235 if height_scan else 48
     obs = torch.randn(num_envs, num_obs)
     actions = torch.randn(num_envs, 12)
     active_terms = ["base_lin_vel", "height_scan"] if height_scan else ["base_lin_vel"]
-    unwrapped = SimpleNamespace(observation_manager=SimpleNamespace(active_terms={"policy": active_terms}))
+    unwrapped = SimpleNamespace(
+        scene={"robot": SimpleNamespace(joint_names=PHYSX_ORDER)},
+        observation_manager=SimpleNamespace(active_terms={"policy": active_terms}),
+    )
     env = SimpleNamespace(unwrapped=unwrapped)
 
     obs_aug, actions_aug = anymal.compute_symmetric_states(
@@ -88,3 +105,33 @@ def test_compute_symmetric_states_matches_reference(height_scan: bool):
 
     assert torch.equal(obs_aug["policy"], expected_obs)
     assert torch.equal(actions_aug, expected_actions)
+
+
+def test_compute_symmetric_states_follows_joint_names():
+    """Verify the Newton joint order is mirrored like the PhysX order, with the joint columns reordered."""
+    joint_cols = [PHYSX_ORDER.index(name) for name in NEWTON_ORDER]
+    obs_cols = list(range(12)) + [start + joint for start in (12, 24, 36) for joint in joint_cols]
+    obs = torch.randn(5, 48)
+    actions = torch.randn(5, 12)
+
+    outputs = []
+    for joint_names, obs_idx, action_idx in (
+        (PHYSX_ORDER, list(range(48)), list(range(12))),
+        (NEWTON_ORDER, obs_cols, joint_cols),
+    ):
+        env = SimpleNamespace(
+            scene={"robot": SimpleNamespace(joint_names=joint_names)},
+            observation_manager=SimpleNamespace(active_terms={"policy": []}),
+        )
+        env.unwrapped = env
+        policy = TensorDict({"policy": obs[:, obs_idx]}, batch_size=[5])
+        outputs.append(anymal.compute_symmetric_states(env, policy, actions[:, action_idx]))
+    (physx_obs, physx_actions), (newton_obs, newton_actions) = outputs
+
+    assert torch.equal(newton_obs["policy"], physx_obs["policy"][:, obs_cols])
+    assert torch.equal(newton_actions, physx_actions[:, joint_cols])
+
+
+def test_compute_symmetric_states_without_inputs():
+    """Return ``(None, None)`` without reading the scene when there is nothing to augment."""
+    assert anymal.compute_symmetric_states(SimpleNamespace()) == (None, None)