Source code for isaaclab.envs.mdp.commands.velocity_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 the velocity-based locomotion task."""
from __future__ import annotations
import logging
from collections.abc import Sequence
from typing import TYPE_CHECKING
import torch
import isaaclab.utils.math as math_utils
from isaaclab.assets import Articulation
from isaaclab.managers import CommandTerm
from isaaclab.markers import VisualizationMarkers
if TYPE_CHECKING:
from isaaclab.envs import ManagerBasedEnv
from .commands_cfg import NormalVelocityCommandCfg, UniformVelocityCommandCfg
# import logger
logger = logging.getLogger(__name__)
[docs]
class UniformVelocityCommand(CommandTerm):
r"""Command generator that generates a velocity command in SE(2) from uniform distribution.
The command comprises of a linear velocity in x and y direction and an angular velocity around
the z-axis. It is given in the robot's base frame.
If the :attr:`cfg.heading_command` flag is set to True, the angular velocity is computed from the heading
error similar to doing a proportional control on the heading error. The target heading is sampled uniformly
from the provided range. Otherwise, the angular velocity is sampled uniformly from the provided range.
Mathematically, the angular velocity is computed as follows from the heading command:
.. math::
\omega_z = \frac{1}{2} \text{wrap_to_pi}(\theta_{\text{target}} - \theta_{\text{current}})
"""
cfg: UniformVelocityCommandCfg
"""The configuration of the command generator."""
[docs]
def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedEnv):
"""Initialize the command generator.
Args:
cfg: The configuration of the command generator.
env: The environment.
Raises:
ValueError: If the heading command is active but the heading range is not provided.
"""
# initialize the base class
super().__init__(cfg, env)
# check configuration
if self.cfg.heading_command and self.cfg.ranges.heading is None:
raise ValueError(
"The velocity command has heading commands active (heading_command=True) but the `ranges.heading`"
" parameter is set to None."
)
if self.cfg.ranges.heading and not self.cfg.heading_command:
logger.warning(
f"The velocity command has the 'ranges.heading' attribute set to '{self.cfg.ranges.heading}'"
" but the heading command is not active. Consider setting the flag for the heading command to True."
)
# obtain the robot asset
# -- robot
self.robot: Articulation = env.scene[cfg.asset_name]
# crete buffers to store the command
# -- command: x vel, y vel, yaw vel, heading
self.vel_command_b = torch.zeros(self.num_envs, 3, device=self.device)
self.heading_target = torch.zeros(self.num_envs, device=self.device)
self.is_heading_env = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
self.is_standing_env = torch.zeros_like(self.is_heading_env)
# -- metrics: finalized per-episode means/rates, written at reset() and read by the base class
self.metrics["error_vel_xy"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["error_vel_yaw"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["success_rate"] = torch.zeros(self.num_envs, device=self.device)
# -- per-episode running sums (cleared at episode reset)
self._error_xy_sum = torch.zeros(self.num_envs, device=self.device)
self._error_yaw_sum = torch.zeros(self.num_envs, device=self.device)
self._step_count = torch.zeros(self.num_envs, 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/velocity"
self.cfg.element_names = self.cfg.element_names or ["lin_vel_x", "lin_vel_y", "ang_vel_z"]
def __str__(self) -> str:
"""Return a string representation of the command generator."""
msg = "UniformVelocityCommand:\n"
msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n"
msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n"
msg += f"\tHeading command: {self.cfg.heading_command}\n"
if self.cfg.heading_command:
msg += f"\tHeading probability: {self.cfg.rel_heading_envs}\n"
msg += f"\tStanding probability: {self.cfg.rel_standing_envs}"
return msg
"""
Properties
"""
@property
def command(self) -> torch.Tensor:
"""The desired base velocity command in the base frame. Shape is (num_envs, 3)."""
return self.vel_command_b
"""
Implementation specific functions.
"""
def _update_metrics(self):
# accumulate per-step tracking error sums; the per-episode mean is finalized in
# :meth:`reset` so the value is independent of episode length and
# ``resampling_time_range`` (no magic ``max_command_step`` divisor).
error_xy = torch.linalg.norm(self.vel_command_b[:, :2] - self.robot.data.root_lin_vel_b.torch[:, :2], dim=-1)
error_yaw = torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b.torch[:, 2])
self._error_xy_sum += error_xy
self._error_yaw_sum += error_yaw
self._step_count += 1.0
def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]:
# Finalize the just-ended episode's metrics into ``self.metrics`` BEFORE the base
# class reads them. ``success_rate`` is per-env binary: the *episode-mean* error
# is below both thresholds. Then super().reset() logs and zeros ``self.metrics``;
# we zero the running sums for ``env_ids`` afterwards so the next episode starts clean.
if env_ids is None:
env_ids = slice(None)
denom = self._step_count[env_ids].clamp_min(1.0)
mean_error_xy = self._error_xy_sum[env_ids] / denom
mean_error_yaw = self._error_yaw_sum[env_ids] / denom
self.metrics["error_vel_xy"][env_ids] = mean_error_xy
self.metrics["error_vel_yaw"][env_ids] = mean_error_yaw
self.metrics["success_rate"][env_ids] = (
(mean_error_xy < self.cfg.vel_xy_success_threshold) & (mean_error_yaw < self.cfg.vel_yaw_success_threshold)
).float()
extras = super().reset(env_ids)
# Route success_rate to the unified ``Metrics/success_rate`` path (shared TensorBoard
# / wandb card across tasks); pop it from the returned dict so CommandManager does
# not additionally log it under ``Metrics/<term_name>/success_rate``.
self._env.extras.setdefault("log", {})["Metrics/success_rate"] = extras.pop("success_rate")
self._error_xy_sum[env_ids] = 0.0
self._error_yaw_sum[env_ids] = 0.0
self._step_count[env_ids] = 0.0
return extras
def _resample_command(self, env_ids: Sequence[int]):
# sample velocity commands
r = torch.empty(len(env_ids), device=self.device)
# -- linear velocity - x direction
self.vel_command_b[env_ids, 0] = r.uniform_(*self.cfg.ranges.lin_vel_x)
# -- linear velocity - y direction
self.vel_command_b[env_ids, 1] = r.uniform_(*self.cfg.ranges.lin_vel_y)
# -- ang vel yaw - rotation around z
self.vel_command_b[env_ids, 2] = r.uniform_(*self.cfg.ranges.ang_vel_z)
# heading target
if self.cfg.heading_command:
self.heading_target[env_ids] = r.uniform_(*self.cfg.ranges.heading)
# update heading envs
self.is_heading_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_heading_envs
# update standing envs
self.is_standing_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_standing_envs
def _update_command(self):
"""Post-processes the velocity command.
This function sets velocity command to zero for standing environments and computes angular
velocity from heading direction if the heading_command flag is set.
"""
# Compute angular velocity from heading direction
if self.cfg.heading_command:
# resolve indices of heading envs
env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
# compute angular velocity
heading_error = math_utils.wrap_to_pi(
self.heading_target[env_ids] - self.robot.data.heading_w.torch[env_ids]
)
self.vel_command_b[env_ids, 2] = torch.clip(
self.cfg.heading_control_stiffness * heading_error,
min=self.cfg.ranges.ang_vel_z[0],
max=self.cfg.ranges.ang_vel_z[1],
)
# Enforce standing (i.e., zero velocity command) for standing envs
# TODO: check if conversion is needed
standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten()
self.vel_command_b[standing_env_ids, :] = 0.0
def _set_debug_vis_impl(self, debug_vis: bool):
# set visibility of markers
# note: parent only deals with callbacks. not their visibility
if debug_vis:
# create markers if necessary for the first time
if not hasattr(self, "goal_vel_visualizer"):
# -- goal
self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg)
# -- current
self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg)
# set their visibility to true
self.goal_vel_visualizer.set_visibility(True)
self.current_vel_visualizer.set_visibility(True)
else:
if hasattr(self, "goal_vel_visualizer"):
self.goal_vel_visualizer.set_visibility(False)
self.current_vel_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
# get marker location
# -- base state
base_pos_w = self.robot.data.root_pos_w.torch.clone()
offset = self.cfg.marker_pos_offset
base_pos_w[:, 0] += offset[0]
base_pos_w[:, 1] += offset[1]
base_pos_w[:, 2] += offset[2]
# -- resolve the scales and quaternions
vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2])
vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(
self.robot.data.root_lin_vel_b.torch[:, :2]
)
# display markers
environment_ids = self._env.scene._ALL_INDICES
self.goal_vel_visualizer.visualize(
base_pos_w,
vel_des_arrow_quat,
vel_des_arrow_scale,
environment_ids=environment_ids,
)
self.current_vel_visualizer.visualize(
base_pos_w,
vel_arrow_quat,
vel_arrow_scale,
environment_ids=environment_ids,
)
"""
Internal helpers.
"""
def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Converts the XY base velocity command to arrow direction rotation."""
# obtain default scale of the marker
default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale
# arrow-scale
arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1)
arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0
# arrow-direction
heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0])
zeros = torch.zeros_like(heading_angle)
arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle)
# convert everything back from base to world frame
base_quat_w = self.robot.data.root_quat_w.torch
arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat)
return arrow_scale, arrow_quat
[docs]
class NormalVelocityCommand(UniformVelocityCommand):
"""Command generator that generates a velocity command in SE(2) from a normal distribution.
The command comprises of a linear velocity in x and y direction and an angular velocity around
the z-axis. It is given in the robot's base frame.
The command is sampled from a normal distribution with mean and standard deviation specified in
the configuration. With equal probability, the sign of the individual components is flipped.
"""
cfg: NormalVelocityCommandCfg
"""The command generator configuration."""
[docs]
def __init__(self, cfg: NormalVelocityCommandCfg, env: ManagerBasedEnv):
"""Initializes the command generator.
Args:
cfg: The command generator configuration.
env: The environment.
"""
super().__init__(cfg, env)
# create buffers for zero commands envs
self.is_zero_vel_x_env = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
self.is_zero_vel_y_env = torch.zeros_like(self.is_zero_vel_x_env)
self.is_zero_vel_yaw_env = torch.zeros_like(self.is_zero_vel_x_env)
def __str__(self) -> str:
"""Return a string representation of the command generator."""
msg = "NormalVelocityCommand:\n"
msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n"
msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n"
msg += f"\tStanding probability: {self.cfg.rel_standing_envs}"
return msg
def _resample_command(self, env_ids):
# sample velocity commands
r = torch.empty(len(env_ids), device=self.device)
# -- linear velocity - x direction
self.vel_command_b[env_ids, 0] = r.normal_(mean=self.cfg.ranges.mean_vel[0], std=self.cfg.ranges.std_vel[0])
self.vel_command_b[env_ids, 0] *= torch.where(r.uniform_(0.0, 1.0) <= 0.5, 1.0, -1.0)
# -- linear velocity - y direction
self.vel_command_b[env_ids, 1] = r.normal_(mean=self.cfg.ranges.mean_vel[1], std=self.cfg.ranges.std_vel[1])
self.vel_command_b[env_ids, 1] *= torch.where(r.uniform_(0.0, 1.0) <= 0.5, 1.0, -1.0)
# -- angular velocity - yaw direction
self.vel_command_b[env_ids, 2] = r.normal_(mean=self.cfg.ranges.mean_vel[2], std=self.cfg.ranges.std_vel[2])
self.vel_command_b[env_ids, 2] *= torch.where(r.uniform_(0.0, 1.0) <= 0.5, 1.0, -1.0)
# update element wise zero velocity command
# TODO what is zero prob ?
self.is_zero_vel_x_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.ranges.zero_prob[0]
self.is_zero_vel_y_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.ranges.zero_prob[1]
self.is_zero_vel_yaw_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.ranges.zero_prob[2]
# update standing envs
self.is_standing_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_standing_envs
def _update_command(self):
"""Sets velocity command to zero for standing envs."""
# Enforce standing (i.e., zero velocity command) for standing envs
standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten() # TODO check if conversion is needed
self.vel_command_b[standing_env_ids, :] = 0.0
# Enforce zero velocity for individual elements
# TODO: check if conversion is needed
zero_vel_x_env_ids = self.is_zero_vel_x_env.nonzero(as_tuple=False).flatten()
zero_vel_y_env_ids = self.is_zero_vel_y_env.nonzero(as_tuple=False).flatten()
zero_vel_yaw_env_ids = self.is_zero_vel_yaw_env.nonzero(as_tuple=False).flatten()
self.vel_command_b[zero_vel_x_env_ids, 0] = 0.0
self.vel_command_b[zero_vel_y_env_ids, 1] = 0.0
self.vel_command_b[zero_vel_yaw_env_ids, 2] = 0.0