129 lines
5.9 KiB
Python
129 lines
5.9 KiB
Python
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}.")
|
|
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]
|
|
|
|
def _resample_command(self, env_ids: Sequence[int]):
|
|
"""Sample commands, with a configurable fraction reserved for straight walking."""
|
|
env_ids = torch.as_tensor(env_ids, device=self.device, dtype=torch.long)
|
|
super()._resample_command(env_ids)
|
|
|
|
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]
|
|
self.metrics["error_vel_xy"] += (
|
|
torch.norm(self.vel_command_b[:, :2] - body_lin_vel_heading[:, :2], dim=-1) / max_command_step
|
|
)
|
|
self.metrics["error_vel_yaw"] += (
|
|
torch.abs(self.vel_command_b[:, 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
|
|
|
|
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
|
|
|
|
|
|
@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."""
|
|
|
|
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."""
|