Source code for isaaclab.envs.mdp.commands.pose_command

# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Sub-module containing command generators for pose tracking."""

from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING

import torch

from isaaclab.assets import Articulation
from isaaclab.managers import CommandTerm
from isaaclab.markers import VisualizationMarkers
from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES
from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique

if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedEnv

    from .commands_cfg import UniformPoseCommandCfg




class UniformPoseCommand(CommandTerm):
    """Command generator for generating pose commands uniformly.

    The command generator generates poses by sampling positions uniformly within specified
    regions in cartesian space. For orientation, it samples uniformly the euler angles
    (roll-pitch-yaw) and converts them into quaternion representation (x, y, z, w).

    The position and orientation commands are generated in the base frame of the robot, and not the
    simulation world frame. This means that users need to handle the transformation from the
    base frame to the simulation world frame themselves.

    .. caution::

        Sampling orientations uniformly is not strictly the same as sampling euler angles uniformly.
        This is because rotations are defined by 3D non-Euclidean space, and the mapping
        from euler angles to rotations is not one-to-one.

    """

    cfg: UniformPoseCommandCfg
    """Configuration for the command generator."""



    def __init__(self, cfg: UniformPoseCommandCfg, env: ManagerBasedEnv):
        """Initialize the command generator class.

        Args:
            cfg: The configuration parameters for the command generator.
            env: The environment object.
        """
        # initialize the base class
        super().__init__(cfg, env)

        # extract the robot and body index for which the command is generated
        self.robot: Articulation = env.scene[cfg.asset_name]
        self.body_idx = self.robot.find_bodies(cfg.body_name)[0][0]

        # create buffers
        # -- commands: (x, y, z, qx, qy, qz, qw) in root frame
        self.pose_command_b = torch.zeros(self.num_envs, 7, device=self.device)
        self.pose_command_b[:, 3] = 1.0
        self.pose_command_w = torch.zeros_like(self.pose_command_b)
        # -- metrics
        self.metrics["position_error"] = torch.zeros(self.num_envs, device=self.device)
        self.metrics["orientation_error"] = torch.zeros(self.num_envs, device=self.device)
        # -- per-episode sticky success bit (only used when at least one success threshold is set)
        self._track_success = (
            cfg.position_success_threshold is not None or cfg.orientation_success_threshold is not None
        )
        if self._track_success:
            self._succeeded = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)

        # adds (optional) cmd kind and element names for leapp export
        # during export, semantic data about this command will be used to annotate the command input
        self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/pose"
        self.cfg.element_names = self.cfg.element_names or POSE7_ELEMENT_NAMES


    def __str__(self) -> str:
        msg = "UniformPoseCommand:\n"
        msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n"
        msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n"
        return msg

    """
    Properties
    """

    @property
    def command(self) -> torch.Tensor:
        """The desired pose command. Shape is (num_envs, 7).

        The first three elements correspond to the position, followed by the quaternion orientation in (x, y, z, w).
        """
        return self.pose_command_b

    def compute_success(self) -> torch.Tensor:
        """Compute whether the current body pose satisfies the configured success thresholds.

        Successful entries are also recorded by the episode-level success tracker when enabled.

        Returns:
            A boolean tensor indicating which environments satisfy every configured threshold.
            If no success threshold is configured, all entries are false.
        """
        position_error, orientation_error = self._compute_error()
        success = self._compute_success(position_error, orientation_error)
        if self._track_success:
            self._succeeded |= success
        return success

    """
    Implementation specific functions.
    """

    def _update_metrics(self):
        position_error, orientation_error = self._compute_error()
        self.metrics["position_error"] = position_error
        self.metrics["orientation_error"] = orientation_error
        if self._track_success:
            self._succeeded |= self._compute_success(position_error, orientation_error)

    def _compute_error(self) -> tuple[torch.Tensor, torch.Tensor]:
        # transform command from base frame to simulation world frame
        self.pose_command_w[:, :3], self.pose_command_w[:, 3:] = combine_frame_transforms(
            self.robot.data.root_pos_w.torch,
            self.robot.data.root_quat_w.torch,
            self.pose_command_b[:, :3],
            self.pose_command_b[:, 3:],
        )
        # compute the error
        pos_error, rot_error = compute_pose_error(
            self.pose_command_w[:, :3],
            self.pose_command_w[:, 3:],
            self.robot.data.body_pos_w.torch[:, self.body_idx],
            self.robot.data.body_quat_w.torch[:, self.body_idx],
        )
        return torch.linalg.norm(pos_error, dim=-1), torch.linalg.norm(rot_error, dim=-1)

    def _compute_success(self, position_error: torch.Tensor, orientation_error: torch.Tensor) -> torch.Tensor:
        success = torch.ones(self.num_envs, dtype=torch.bool, device=self.device)
        if self.cfg.position_success_threshold is not None:
            success &= position_error < self.cfg.position_success_threshold
        if self.cfg.orientation_success_threshold is not None:
            success &= orientation_error < self.cfg.orientation_success_threshold
        if not self._track_success:
            success[:] = False
        return success

    def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]:
        extras = super().reset(env_ids)
        if self._track_success:
            if env_ids is None:
                env_ids = slice(None)
            # Write the unified ``Metrics/success_rate`` directly to env extras so it shares
            # a TensorBoard card with the same metric from other tasks.
            self._env.extras.setdefault("log", {})["Metrics/success_rate"] = (
                self._succeeded[env_ids].float().mean().item()
            )
            self._succeeded[env_ids] = False
        return extras

    def _resample_command(self, env_ids: Sequence[int]):
        # sample new pose targets
        # -- position
        r = torch.empty(len(env_ids), device=self.device)
        self.pose_command_b[env_ids, 0] = r.uniform_(*self.cfg.ranges.pos_x)
        self.pose_command_b[env_ids, 1] = r.uniform_(*self.cfg.ranges.pos_y)
        self.pose_command_b[env_ids, 2] = r.uniform_(*self.cfg.ranges.pos_z)
        # -- orientation
        euler_angles = torch.zeros_like(self.pose_command_b[env_ids, :3])
        euler_angles[:, 0].uniform_(*self.cfg.ranges.roll)
        euler_angles[:, 1].uniform_(*self.cfg.ranges.pitch)
        euler_angles[:, 2].uniform_(*self.cfg.ranges.yaw)
        quat = quat_from_euler_xyz(euler_angles[:, 0], euler_angles[:, 1], euler_angles[:, 2])
        # make sure the quaternion has real part as positive
        self.pose_command_b[env_ids, 3:] = quat_unique(quat) if self.cfg.make_quat_unique else quat

    def _update_command(self):
        pass

    def _set_debug_vis_impl(self, debug_vis: bool):
        # create markers if necessary for the first time
        if debug_vis:
            if not hasattr(self, "goal_pose_visualizer"):
                # -- goal pose
                self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg)
                # -- current body pose
                self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg)
            # set their visibility to true
            self.goal_pose_visualizer.set_visibility(True)
            self.current_pose_visualizer.set_visibility(True)
        else:
            if hasattr(self, "goal_pose_visualizer"):
                self.goal_pose_visualizer.set_visibility(False)
                self.current_pose_visualizer.set_visibility(False)

    def _debug_vis_callback(self, event):
        # check if robot is initialized
        # note: this is needed in-case the robot is de-initialized. we can't access the data
        if not self.robot.is_initialized:
            return
        # update the markers
        environment_ids = self._env.scene._ALL_INDICES
        # -- goal pose
        self.goal_pose_visualizer.visualize(
            self.pose_command_w[:, :3],
            self.pose_command_w[:, 3:],
            environment_ids=environment_ids,
        )
        # -- current body pose
        body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx]
        self.current_pose_visualizer.visualize(
            body_link_pose_w[:, :3],
            body_link_pose_w[:, 3:7],
            environment_ids=environment_ids,
        )