diff --git a/source/isaaclab_tasks/changelog.d/fix-anymal-symmetry-joint-order.fixed.rst b/source/isaaclab_tasks/changelog.d/fix-anymal-symmetry-joint-order.fixed.rst new file mode 100644 index 000000000000..21fd8748a9e0 --- /dev/null +++ b/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. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/symmetry/anymal.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/symmetry/anymal.py index 7c54f1b83bc2..5be817bf86fc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/symmetry/anymal.py +++ b/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 ``_``, 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]]) diff --git a/source/isaaclab_tasks/test/core/test_anymal_symmetry.py b/source/isaaclab_tasks/test/core/test_anymal_symmetry.py index 36c35ad3cbe3..fb2af10f0138 100644 --- a/source/isaaclab_tasks/test/core/test_anymal_symmetry.py +++ b/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)