cmvr_ai_lab/source/engineai_lab/tasks/velocity/mdp/commands.py

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."""