from __future__ import annotations from collections.abc import Sequence from dataclasses import MISSING import torch import isaaclab.utils.math as math_utils from isaaclab.envs.mdp.commands import UniformVelocityCommand, UniformVelocityCommandCfg from isaaclab.utils import configclass class BodyVelocityCommand(UniformVelocityCommand): """Velocity command whose metrics and markers use a configured robot body.""" cfg: BodyVelocityCommandCfg def __init__(self, cfg: BodyVelocityCommandCfg, env): if not 0.0 <= cfg.rel_straight_envs <= 1.0: raise ValueError(f"rel_straight_envs must be in [0, 1], got {cfg.rel_straight_envs}.") if not 0.0 <= cfg.rel_high_speed_envs <= 1.0: raise ValueError(f"rel_high_speed_envs must be in [0, 1], got {cfg.rel_high_speed_envs}.") if not 0.0 <= cfg.rel_replay_speed_envs <= 1.0: raise ValueError(f"rel_replay_speed_envs must be in [0, 1], got {cfg.rel_replay_speed_envs}.") if cfg.rel_high_speed_envs + cfg.rel_replay_speed_envs > 1.0: raise ValueError("rel_high_speed_envs + rel_replay_speed_envs must not exceed 1.") if cfg.rel_high_speed_envs > 0.0 and cfg.high_speed_bandwidth <= 0.0: raise ValueError("high_speed_bandwidth must be positive when high-speed sampling is enabled.") if cfg.rel_replay_speed_envs > 0.0: if cfg.replay_speed_range is None: raise ValueError("replay_speed_range is required when replay-speed sampling is enabled.") if cfg.replay_speed_range[0] > cfg.replay_speed_range[1]: raise ValueError("replay_speed_range must be ordered from low to high.") if cfg.command_ramp_rates is not None: if len(cfg.command_ramp_rates) != 3: raise ValueError("command_ramp_rates must contain vx, vy, and yaw rates.") if any(rate <= 0.0 for rate in cfg.command_ramp_rates): raise ValueError("All command_ramp_rates must be positive.") super().__init__(cfg, env) body_ids, body_names = self.robot.find_bodies(cfg.body_name, preserve_order=True) if len(body_ids) != 1: raise ValueError( f"BodyVelocityCommand requires exactly one body matching {cfg.body_name!r}; " f"found {body_names}." ) self.body_id = body_ids[0] self._command_ramp_rates = None self._ramped_vel_command_b = None if cfg.command_ramp_rates is not None: self._command_ramp_rates = torch.tensor( cfg.command_ramp_rates, device=self.device, dtype=self.vel_command_b.dtype ).view(1, 3) self._ramped_vel_command_b = torch.zeros_like(self.vel_command_b) @property def command(self) -> torch.Tensor: """Velocity command exposed to observations and rewards.""" ramped_command = getattr(self, "_ramped_vel_command_b", None) return self.vel_command_b if ramped_command is None else ramped_command def _resolve_env_ids(self, env_ids: Sequence[int] | slice | None) -> torch.Tensor: if env_ids is None: return torch.arange(self.num_envs, device=self.device) if isinstance(env_ids, slice): return torch.arange(self.num_envs, device=self.device)[env_ids] return torch.as_tensor(env_ids, device=self.device, dtype=torch.long) def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: """Reset raw targets and restart enabled command ramps from zero.""" resolved_env_ids = self._resolve_env_ids(env_ids) extras = super().reset(resolved_env_ids) if self._ramped_vel_command_b is not None: self._ramped_vel_command_b[resolved_env_ids] = 0.0 return extras def _resample_command(self, env_ids: Sequence[int]): """Sample uniform, high-speed, and old-speed replay command buckets.""" env_ids = self._resolve_env_ids(env_ids) super()._resample_command(env_ids) speed_bucket_probability = self.cfg.rel_high_speed_envs + self.cfg.rel_replay_speed_envs if speed_bucket_probability > 0.0 and env_ids.numel() > 0: moving_env_ids = env_ids[~self.is_standing_env[env_ids]] if moving_env_ids.numel() > 0: selector = torch.rand(moving_env_ids.numel(), device=self.device) high_mask = selector < self.cfg.rel_high_speed_envs high_env_ids = moving_env_ids[high_mask] if high_env_ids.numel() > 0: current_low, current_high = self.cfg.ranges.lin_vel_x high_low = max(current_low, current_high - self.cfg.high_speed_bandwidth) high_samples = torch.empty(high_env_ids.numel(), device=self.device) self.vel_command_b[high_env_ids, 0] = high_samples.uniform_(high_low, current_high) replay_mask = (selector >= self.cfg.rel_high_speed_envs) & ( selector < speed_bucket_probability ) replay_env_ids = moving_env_ids[replay_mask] if replay_env_ids.numel() > 0: replay_low, configured_replay_high = self.cfg.replay_speed_range replay_high = min(configured_replay_high, self.cfg.ranges.lin_vel_x[1]) if replay_low <= replay_high: replay_samples = torch.empty(replay_env_ids.numel(), device=self.device) self.vel_command_b[replay_env_ids, 0] = replay_samples.uniform_(replay_low, replay_high) if self.cfg.rel_straight_envs <= 0.0 or env_ids.numel() == 0: return straight_mask = torch.rand(env_ids.numel(), device=self.device) <= self.cfg.rel_straight_envs straight_env_ids = env_ids[straight_mask] self.vel_command_b[straight_env_ids, 1:] = 0.0 self.is_heading_env[straight_env_ids] = False def _body_heading_quat_w(self) -> torch.Tensor: body_quat_w = self.robot.data.body_quat_w[:, self.body_id, :] _, _, yaw = math_utils.euler_xyz_from_quat(body_quat_w) zeros = torch.zeros_like(yaw) return math_utils.quat_from_euler_xyz(zeros, zeros, yaw - self.cfg.heading_yaw_offset) def _body_lin_vel_heading(self) -> torch.Tensor: body_lin_vel_w = self.robot.data.body_lin_vel_w[:, self.body_id, :] return math_utils.quat_apply_inverse(self._body_heading_quat_w(), body_lin_vel_w) def _update_metrics(self): max_command_time = self.cfg.resampling_time_range[1] max_command_step = max_command_time / self._env.step_dt body_lin_vel_heading = self._body_lin_vel_heading() body_yaw_rate = self.robot.data.body_ang_vel_w[:, self.body_id, 2] command = self.command self.metrics["error_vel_xy"] += ( torch.norm(command[:, :2] - body_lin_vel_heading[:, :2], dim=-1) / max_command_step ) self.metrics["error_vel_yaw"] += ( torch.abs(command[:, 2] - body_yaw_rate) / max_command_step ) def _update_command(self): if self.cfg.heading_command: env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten() if len(env_ids) > 0: body_quat_w = self.robot.data.body_quat_w[env_ids, self.body_id, :] _, _, body_yaw = math_utils.euler_xyz_from_quat(body_quat_w) body_heading = body_yaw - self.cfg.heading_yaw_offset heading_error = math_utils.wrap_to_pi(self.heading_target[env_ids] - body_heading) 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], ) standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten() self.vel_command_b[standing_env_ids, :] = 0.0 if self._ramped_vel_command_b is not None: max_delta = self._command_ramp_rates * self._env.step_dt delta = self.vel_command_b - self._ramped_vel_command_b self._ramped_vel_command_b += torch.maximum(torch.minimum(delta, max_delta), -max_delta) def _debug_vis_callback(self, _event): if not self.robot.is_initialized: return marker_pos_w = self.robot.data.body_pos_w[:, self.body_id, :].clone() marker_pos_w[:, 2] += self.cfg.marker_height_offset desired_pos_w = marker_pos_w.clone() actual_pos_w = marker_pos_w.clone() desired_pos_w[:, 2] += 0.5 * self.cfg.marker_vertical_separation actual_pos_w[:, 2] -= 0.5 * self.cfg.marker_vertical_separation desired_scale, desired_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) actual_scale, actual_quat = self._resolve_xy_velocity_to_arrow(self._body_lin_vel_heading()[:, :2]) self.goal_vel_visualizer.visualize(desired_pos_w, desired_quat, desired_scale) self.current_vel_visualizer.visualize(actual_pos_w, actual_quat, actual_scale) def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: default_scale = self.goal_vel_visualizer.cfg.markers["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 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) arrow_quat = math_utils.quat_mul(self._body_heading_quat_w(), arrow_quat) return arrow_scale, arrow_quat def forward_speed_range_curriculum( env, env_ids: Sequence[int], command_name: str, min_forward_speed: float, initial_max_speed: float, final_max_speed: float, speed_increment: float, stage_steps: int, ) -> dict[str, float]: """Increase the sampled forward-speed ceiling in fixed training stages.""" del env_ids if not min_forward_speed <= initial_max_speed <= final_max_speed: raise ValueError("Expected min_forward_speed <= initial_max_speed <= final_max_speed.") if speed_increment <= 0.0 or stage_steps <= 0: raise ValueError("speed_increment and stage_steps must be positive.") stage = env.common_step_counter // stage_steps max_forward_speed = min(initial_max_speed + stage * speed_increment, final_max_speed) command_term = env.command_manager.get_term(command_name) command_term.cfg.ranges.lin_vel_x = (min_forward_speed, max_forward_speed) return {"stage": float(stage), "max_forward_speed": float(max_forward_speed)} @configclass class BodyVelocityCommandCfg(UniformVelocityCommandCfg): """Configuration for body-referenced planar velocity commands.""" class_type: type = BodyVelocityCommand body_name: str = MISSING """Body used for velocity metrics and command visualization.""" heading_yaw_offset: float = 0.0 """Yaw offset from the body's URDF frame to the locomotion heading frame.""" rel_straight_envs: float = 0.0 """Fraction of sampled environments with zero lateral and yaw commands.""" rel_high_speed_envs: float = 0.0 """Fraction of moving environments sampled from the upper forward-speed band.""" high_speed_bandwidth: float = 0.3 """Width of the upper forward-speed sampling band in m/s.""" rel_replay_speed_envs: float = 0.0 """Fraction of moving environments sampled from the old-speed replay range.""" replay_speed_range: tuple[float, float] | None = None """Independent forward-speed range retained to prevent low-speed forgetting.""" command_ramp_rates: tuple[float, float, float] | None = None """Per-axis vx, vy, and yaw slew rates in m/s^2, m/s^2, and rad/s^2.""" marker_height_offset: float = 0.35 """Vertical marker offset from the configured body origin in meters.""" marker_vertical_separation: float = 0.08 """Vertical separation between desired and measured velocity arrows in meters."""