247 lines
12 KiB
Python
247 lines
12 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}.")
|
|
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."""
|