58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
from typing import TYPE_CHECKING
|
|
|
|
from isaaclab.managers import SceneEntityCfg
|
|
from isaaclab.utils.math import euler_xyz_from_quat, quat_apply_inverse, quat_from_euler_xyz
|
|
|
|
if TYPE_CHECKING:
|
|
from isaaclab.envs import ManagerBasedEnv
|
|
|
|
|
|
def _single_body_id(asset_cfg: SceneEntityCfg, term_name: str) -> int:
|
|
body_ids = asset_cfg.body_ids
|
|
if body_ids is None:
|
|
raise ValueError(f"{term_name} requires asset_cfg with exactly one body name.")
|
|
if isinstance(body_ids, int):
|
|
return body_ids
|
|
if isinstance(body_ids, slice) or len(body_ids) != 1:
|
|
raise ValueError(f"{term_name} requires exactly one body id, got {body_ids}.")
|
|
return body_ids[0]
|
|
|
|
|
|
def _heading_quat_with_offset(body_quat_w: torch.Tensor, heading_yaw_offset: float = 0.0) -> torch.Tensor:
|
|
roll, pitch, yaw = euler_xyz_from_quat(body_quat_w)
|
|
zeros = torch.zeros_like(yaw)
|
|
return quat_from_euler_xyz(zeros, zeros, yaw - heading_yaw_offset)
|
|
|
|
|
|
def robot_base_lin_vel_b(env: ManagerBasedEnv) -> torch.Tensor:
|
|
"""Base linear velocity expressed in the base frame."""
|
|
asset = env.scene["robot"]
|
|
# prefer direct base-frame velocity if available
|
|
if getattr(asset.data, "root_lin_vel_b", None) is not None:
|
|
return asset.data.root_lin_vel_b
|
|
# fallback: rotate world velocity into base frame
|
|
return quat_apply_inverse(asset.data.root_quat_w, asset.data.root_lin_vel_w)
|
|
|
|
|
|
def body_ang_vel_yaw_frame(
|
|
env: ManagerBasedEnv,
|
|
asset_cfg: SceneEntityCfg,
|
|
heading_yaw_offset: float = 0.0,
|
|
) -> torch.Tensor:
|
|
"""Body angular velocity expressed in a yaw-aligned frame for that body."""
|
|
body_id = _single_body_id(asset_cfg, "body_ang_vel_yaw_frame")
|
|
asset = env.scene[asset_cfg.name]
|
|
body_quat_w = asset.data.body_quat_w[:, body_id, :]
|
|
body_ang_vel_w = asset.data.body_ang_vel_w[:, body_id, :]
|
|
return quat_apply_inverse(_heading_quat_with_offset(body_quat_w, heading_yaw_offset), body_ang_vel_w)
|
|
|
|
|
|
def body_projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
|
|
"""Gravity projection in a configured body's local frame."""
|
|
body_id = _single_body_id(asset_cfg, "body_projected_gravity")
|
|
asset = env.scene[asset_cfg.name]
|
|
return quat_apply_inverse(asset.data.body_quat_w[:, body_id, :], asset.data.GRAVITY_VEC_W)
|