1015 lines
44 KiB
Python
1015 lines
44 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
from typing import TYPE_CHECKING
|
|
|
|
from isaaclab.envs import mdp
|
|
from isaaclab.managers import SceneEntityCfg
|
|
from isaaclab.sensors import ContactSensor
|
|
from isaaclab.utils.math import (
|
|
euler_xyz_from_quat,
|
|
quat_apply_inverse,
|
|
quat_from_euler_xyz,
|
|
quat_mul,
|
|
quat_rotate_inverse,
|
|
wrap_to_pi,
|
|
yaw_quat,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from isaaclab.envs import ManagerBasedRLEnv
|
|
|
|
|
|
def _to_env_device(env: ManagerBasedRLEnv, tensor: torch.Tensor) -> torch.Tensor:
|
|
"""Move sensor tensors back to the RL environment device when Isaac uses another CUDA device."""
|
|
return tensor.to(env.device)
|
|
|
|
|
|
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 _heading_yaw_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)
|
|
return wrap_to_pi(yaw - heading_yaw_offset)
|
|
|
|
|
|
def _command_is_moving(
|
|
env: ManagerBasedRLEnv,
|
|
command_name: str,
|
|
linear_threshold: float = 0.1,
|
|
angular_threshold: float = 0.1,
|
|
) -> torch.Tensor:
|
|
commands = env.command_manager.get_command(command_name)
|
|
return (torch.norm(commands[:, :2], dim=1) > linear_threshold) | (
|
|
torch.abs(commands[:, 2]) > angular_threshold
|
|
)
|
|
|
|
|
|
def action_smoothness(env: ManagerBasedRLEnv) -> torch.Tensor:
|
|
"""Penalize action second-order differences to encourage smooth control."""
|
|
action_manager = env.action_manager
|
|
prev_prev_action = getattr(action_manager, "_prev_prev_action", None)
|
|
if prev_prev_action is None:
|
|
prev_prev_action = torch.zeros_like(action_manager.action)
|
|
action_manager._prev_prev_action = prev_prev_action
|
|
|
|
second_diff = action_manager.action + prev_prev_action - 2.0 * action_manager.prev_action
|
|
reward = torch.sum(torch.square(second_diff), dim=1)
|
|
|
|
# Update action history for the next step and clear reset environments.
|
|
prev_prev_action.copy_(action_manager.prev_action)
|
|
reset_buf = getattr(env, "reset_buf", None)
|
|
if reset_buf is not None:
|
|
reset_env_ids = reset_buf.nonzero(as_tuple=False).squeeze(-1)
|
|
if reset_env_ids.numel() > 0:
|
|
prev_prev_action[reset_env_ids] = 0.0
|
|
|
|
return reward
|
|
|
|
|
|
def _epoch_curriculum_scale(env: ManagerBasedRLEnv, start_scale: float, power: float, interval_epochs: int) -> float:
|
|
"""Compute epoch-based scale by exponentiating by power every interval_epochs."""
|
|
num_step = env.common_step_counter
|
|
interval = max(int(interval_epochs), 1)
|
|
updates = num_step // interval
|
|
return float(start_scale) ** (float(power) ** updates)
|
|
|
|
|
|
def action_smoothness_with_curriculum(
|
|
env: ManagerBasedRLEnv, start_scale: float, power: float, interval_epochs: int
|
|
) -> torch.Tensor:
|
|
"""Action smoothness penalty with epoch-based curriculum scaling."""
|
|
reward = action_smoothness(env)
|
|
return reward * _epoch_curriculum_scale(env, start_scale, power, interval_epochs)
|
|
|
|
|
|
def action_rate_with_curriculum(env: ManagerBasedRLEnv, start_scale: float, power: float, interval_epochs: int) -> torch.Tensor:
|
|
"""Action rate penalty with epoch-based curriculum scaling."""
|
|
reward = mdp.action_rate_l2(env)
|
|
return reward * _epoch_curriculum_scale(env, start_scale, power, interval_epochs)
|
|
|
|
|
|
def energy_cost_with_curriculum(
|
|
env: ManagerBasedRLEnv,
|
|
start_scale: float,
|
|
power: float,
|
|
interval_epochs: int,
|
|
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
|
) -> torch.Tensor:
|
|
"""Energy cost penalty with epoch-based curriculum scaling."""
|
|
reward = energy_cost(env, asset_cfg=asset_cfg)
|
|
return reward * _epoch_curriculum_scale(env, start_scale, power, interval_epochs)
|
|
|
|
|
|
def feet_air_time_similarity(
|
|
env: ManagerBasedRLEnv,
|
|
sensor_cfg: SceneEntityCfg,
|
|
scale: float = 4.0,
|
|
min_air_time: float = 0.0,
|
|
) -> torch.Tensor:
|
|
"""Reward similar air time between two feet.
|
|
|
|
The reward is computed when either foot makes a new contact. It compares the last completed air times
|
|
for the two feet and rewards small differences.
|
|
"""
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
body_ids = sensor_cfg.body_ids
|
|
if body_ids is None or len(body_ids) != 2:
|
|
raise ValueError("feet_air_time_similarity expects exactly two foot body ids in sensor_cfg.body_ids.")
|
|
|
|
first_contact = _to_env_device(env, contact_sensor.compute_first_contact(env.step_dt)[:, body_ids])
|
|
last_air_time = _to_env_device(env, contact_sensor.data.last_air_time[:, body_ids])
|
|
|
|
recent_contact = torch.any(first_contact > 0.0, dim=1)
|
|
valid = torch.all(last_air_time > min_air_time, dim=1)
|
|
diff = torch.abs(last_air_time[:, 0] - last_air_time[:, 1])
|
|
reward = torch.exp(-diff * scale)
|
|
return reward * (recent_contact & valid)
|
|
|
|
|
|
def feet_air_time(
|
|
env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float
|
|
) -> torch.Tensor:
|
|
"""Reward long steps while keeping contact-sensor tensors on the env device."""
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
first_contact = _to_env_device(env, contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids])
|
|
last_air_time = _to_env_device(env, contact_sensor.data.last_air_time[:, sensor_cfg.body_ids])
|
|
reward = torch.sum((last_air_time - threshold) * first_contact, dim=1)
|
|
reward *= torch.norm(env.command_manager.get_command(command_name)[:, :2], dim=1) > 0.1
|
|
return reward
|
|
|
|
|
|
def feet_air_time_positive_on_contact(
|
|
env: ManagerBasedRLEnv,
|
|
command_name: str,
|
|
sensor_cfg: SceneEntityCfg,
|
|
min_air_time: float = 0.05,
|
|
max_air_time: float = 0.25,
|
|
linear_threshold: float = 0.1,
|
|
angular_threshold: float = 0.1,
|
|
) -> torch.Tensor:
|
|
"""Reward completed swing times without penalizing short exploratory steps."""
|
|
if max_air_time <= min_air_time:
|
|
raise ValueError("max_air_time must be greater than min_air_time.")
|
|
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
first_contact = _to_env_device(env, contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids])
|
|
last_air_time = _to_env_device(env, contact_sensor.data.last_air_time[:, sensor_cfg.body_ids])
|
|
completed_swing = torch.clamp(last_air_time - min_air_time, min=0.0, max=max_air_time - min_air_time)
|
|
reward = torch.sum(completed_swing * first_contact, dim=1)
|
|
moving = _command_is_moving(env, command_name, linear_threshold, angular_threshold)
|
|
return reward * moving
|
|
|
|
|
|
def feet_air_time_positive_biped(
|
|
env: ManagerBasedRLEnv,
|
|
command_name: str,
|
|
threshold: float,
|
|
sensor_cfg: SceneEntityCfg,
|
|
linear_threshold: float = 0.1,
|
|
angular_threshold: float = 0.1,
|
|
) -> torch.Tensor:
|
|
"""Dense biped air-time reward with contact-sensor tensors on the env device."""
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
air_time = _to_env_device(env, contact_sensor.data.current_air_time[:, sensor_cfg.body_ids])
|
|
contact_time = _to_env_device(env, contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids])
|
|
in_contact = contact_time > 0.0
|
|
in_mode_time = torch.where(in_contact, contact_time, air_time)
|
|
single_stance = torch.sum(in_contact.int(), dim=1) == 1
|
|
reward = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0]
|
|
reward = torch.clamp(reward, max=threshold)
|
|
moving = _command_is_moving(env, command_name, linear_threshold, angular_threshold)
|
|
return reward * moving
|
|
|
|
|
|
def track_lin_vel_xy_yaw_frame_exp(
|
|
env, sigma: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), stand_threshold: float = 0.06
|
|
) -> torch.Tensor:
|
|
"""Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame using exponential kernel."""
|
|
commands = env.command_manager.get_command(command_name)
|
|
stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & (
|
|
torch.abs(commands[:, 2]) < stand_threshold
|
|
)
|
|
asset = env.scene[asset_cfg.name]
|
|
vel_yaw = quat_apply_inverse(yaw_quat(asset.data.root_quat_w), asset.data.root_lin_vel_w[:, :3])
|
|
lin_vel_error_square = torch.sum(torch.square(commands[:, :2] - vel_yaw[:, :2]), dim=1)
|
|
lin_vel_error_abs = torch.sum(torch.abs(commands[:, :2] - vel_yaw[:, :2]), dim=1)
|
|
rew_square = torch.exp(-lin_vel_error_square * sigma)
|
|
rew_abs = torch.exp(-lin_vel_error_abs * sigma)
|
|
return torch.where(stand_command, rew_abs, rew_square)
|
|
|
|
|
|
def track_ang_vel_z_world_exp(
|
|
env, command_name: str, sigma: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), stand_threshold: float = 0.06
|
|
) -> torch.Tensor:
|
|
"""Reward tracking of angular velocity commands (yaw) in world frame using exponential kernel."""
|
|
commands = env.command_manager.get_command(command_name)
|
|
stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & (
|
|
torch.abs(commands[:, 2]) < stand_threshold
|
|
)
|
|
asset = env.scene[asset_cfg.name]
|
|
ang_vel_error_square = torch.square(commands[:, 2] - asset.data.root_ang_vel_w[:, 2])
|
|
ang_vel_error_abs = torch.abs(commands[:, 2] - asset.data.root_ang_vel_w[:, 2])
|
|
rew_square = torch.exp(-ang_vel_error_square * sigma)
|
|
rew_abs = torch.exp(-ang_vel_error_abs * sigma)
|
|
return torch.where(stand_command, rew_abs, rew_square)
|
|
|
|
|
|
def track_lin_vel_xy_yaw_frame_exp_body(
|
|
env,
|
|
sigma: float,
|
|
command_name: str,
|
|
asset_cfg: SceneEntityCfg,
|
|
stand_threshold: float = 0.06,
|
|
heading_yaw_offset: float = 0.0,
|
|
) -> torch.Tensor:
|
|
"""Track planar velocity of a configured body in its yaw-aligned frame."""
|
|
body_id = _single_body_id(asset_cfg, "track_lin_vel_xy_yaw_frame_exp_body")
|
|
commands = env.command_manager.get_command(command_name)
|
|
stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & (
|
|
torch.abs(commands[:, 2]) < stand_threshold
|
|
)
|
|
asset = env.scene[asset_cfg.name]
|
|
body_quat_w = asset.data.body_quat_w[:, body_id, :]
|
|
body_lin_vel_w = asset.data.body_lin_vel_w[:, body_id, :]
|
|
vel_yaw = quat_apply_inverse(_heading_quat_with_offset(body_quat_w, heading_yaw_offset), body_lin_vel_w)
|
|
lin_vel_error_square = torch.sum(torch.square(commands[:, :2] - vel_yaw[:, :2]), dim=1)
|
|
lin_vel_error_abs = torch.sum(torch.abs(commands[:, :2] - vel_yaw[:, :2]), dim=1)
|
|
rew_square = torch.exp(-lin_vel_error_square * sigma)
|
|
rew_abs = torch.exp(-lin_vel_error_abs * sigma)
|
|
return torch.where(stand_command, rew_abs, rew_square)
|
|
|
|
|
|
def track_ang_vel_z_world_exp_body(
|
|
env,
|
|
command_name: str,
|
|
sigma: float,
|
|
asset_cfg: SceneEntityCfg,
|
|
stand_threshold: float = 0.06,
|
|
) -> torch.Tensor:
|
|
"""Track world-frame yaw angular velocity of a configured body."""
|
|
body_id = _single_body_id(asset_cfg, "track_ang_vel_z_world_exp_body")
|
|
commands = env.command_manager.get_command(command_name)
|
|
stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & (
|
|
torch.abs(commands[:, 2]) < stand_threshold
|
|
)
|
|
asset = env.scene[asset_cfg.name]
|
|
ang_vel_error_square = torch.square(commands[:, 2] - asset.data.body_ang_vel_w[:, body_id, 2])
|
|
ang_vel_error_abs = torch.abs(commands[:, 2] - asset.data.body_ang_vel_w[:, body_id, 2])
|
|
rew_square = torch.exp(-ang_vel_error_square * sigma)
|
|
rew_abs = torch.exp(-ang_vel_error_abs * sigma)
|
|
return torch.where(stand_command, rew_abs, rew_square)
|
|
|
|
|
|
def track_ang_vel_z_world_exp_bodies(
|
|
env,
|
|
command_name: str,
|
|
sigma: float,
|
|
asset_cfg: SceneEntityCfg,
|
|
stand_threshold: float = 0.06,
|
|
) -> torch.Tensor:
|
|
"""Track yaw rate across all configured bodies so internal waist motion cannot satisfy the command alone."""
|
|
body_ids = asset_cfg.body_ids
|
|
if body_ids is None or isinstance(body_ids, slice):
|
|
raise ValueError("track_ang_vel_z_world_exp_bodies requires one or more explicitly resolved body ids.")
|
|
|
|
commands = env.command_manager.get_command(command_name)
|
|
stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & (
|
|
torch.abs(commands[:, 2]) < stand_threshold
|
|
)
|
|
asset = env.scene[asset_cfg.name]
|
|
body_yaw_rates = asset.data.body_ang_vel_w[:, body_ids, 2]
|
|
command_yaw_rate = commands[:, 2].unsqueeze(1)
|
|
ang_vel_error_square = torch.mean(torch.square(command_yaw_rate - body_yaw_rates), dim=1)
|
|
ang_vel_error_abs = torch.mean(torch.abs(command_yaw_rate - body_yaw_rates), dim=1)
|
|
rew_square = torch.exp(-ang_vel_error_square * sigma)
|
|
rew_abs = torch.exp(-ang_vel_error_abs * sigma)
|
|
return torch.where(stand_command, rew_abs, rew_square)
|
|
|
|
|
|
def feet_stumble(
|
|
env, sensor_cfg: SceneEntityCfg, tangential_threshold: float = 2.0, normal_threshold: float = 1.0
|
|
) -> torch.Tensor:
|
|
"""Penalize feet hitting vertical surfaces using contact forces.
|
|
|
|
Flags a stumble when tangential force exceeds ``tangential_threshold`` while the normal force stays
|
|
below ``normal_threshold``. Returns the count of stumbling feet per environment.
|
|
"""
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
forces = _to_env_device(env, contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :])
|
|
tangential = torch.norm(forces[..., :2], dim=-1) > tangential_threshold
|
|
small_normal = torch.abs(forces[..., 2]) < normal_threshold
|
|
stumble = tangential & small_normal
|
|
return stumble.sum(dim=1)
|
|
|
|
def feet_contact(
|
|
env, sensor_cfg: SceneEntityCfg, command_name: str, stand_threshold: float = 0.06, force_threshold: float = 5.0
|
|
) -> torch.Tensor:
|
|
"""Reward valid foot contacts during walking and standing.
|
|
|
|
- When the command is effectively zero (stand), reward 1 only if both feet are in contact.
|
|
- Otherwise, reward 1 if any recent timestep had exactly one foot in contact.
|
|
Uses contact force history if available, else falls back to the latest forces.
|
|
"""
|
|
commands = env.command_manager.get_command(command_name)
|
|
stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & (
|
|
torch.abs(commands[:, 2]) < stand_threshold
|
|
)
|
|
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
contact_history = contact_sensor.data.net_forces_w_history
|
|
if contact_history is None:
|
|
contact_history = contact_sensor.data.net_forces_w.unsqueeze(1)
|
|
contact_history = _to_env_device(env, contact_history)
|
|
|
|
contacts = contact_history[:, :, sensor_cfg.body_ids, 2] > force_threshold
|
|
contact_num_buf = torch.sum(contacts, dim=-1)
|
|
|
|
reward = stand_command.float()
|
|
contact_mask = (~stand_command) & torch.any(contact_num_buf == 1, dim=1)
|
|
reward[contact_mask] = 1.0
|
|
|
|
|
|
return reward
|
|
|
|
def feet_contact_fixed(
|
|
env, sensor_cfg: SceneEntityCfg, command_name: str, stand_threshold: float = 0.06, force_threshold: float = 5.0
|
|
) -> torch.Tensor:
|
|
"""Reward valid foot contacts during walking and standing.
|
|
|
|
- When the command is effectively zero (stand), reward 1.
|
|
- Otherwise, reward 1 if any recent timestep had exactly one foot in contact.
|
|
Uses contact force history if available, else falls back to the latest forces.
|
|
"""
|
|
commands = env.command_manager.get_command(command_name)
|
|
stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & (
|
|
torch.abs(commands[:, 2]) < stand_threshold
|
|
)
|
|
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
contact_history = contact_sensor.data.net_forces_w_history
|
|
if contact_history is None:
|
|
contact_history = contact_sensor.data.net_forces_w.unsqueeze(1)
|
|
contact_history = _to_env_device(env, contact_history)
|
|
|
|
contacts = contact_history[:, :, sensor_cfg.body_ids, 2] > force_threshold
|
|
contact_num_buf = torch.sum(contacts, dim=-1)
|
|
|
|
# For stand, require both feet in contact at the latest timestep.
|
|
stand_contact = contact_num_buf[:, -1] == 2
|
|
reward = (stand_command & stand_contact).float()
|
|
contact_mask = (~stand_command) & torch.any(contact_num_buf == 1, dim=1)
|
|
reward[contact_mask] = 1.0
|
|
|
|
return reward
|
|
|
|
|
|
def biped_contact_mode_reward(
|
|
env: ManagerBasedRLEnv,
|
|
sensor_cfg: SceneEntityCfg,
|
|
command_name: str,
|
|
force_threshold: float = 5.0,
|
|
linear_threshold: float = 0.1,
|
|
angular_threshold: float = 0.1,
|
|
) -> torch.Tensor:
|
|
"""Reward double support while standing and exactly one supporting foot while moving."""
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
forces = _to_env_device(env, contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :])
|
|
contacts = torch.norm(forces, dim=-1) > force_threshold
|
|
contact_count = torch.sum(contacts.int(), dim=1)
|
|
moving = _command_is_moving(env, command_name, linear_threshold, angular_threshold)
|
|
return torch.where(moving, contact_count == 1, contact_count == 2).float()
|
|
|
|
|
|
def swing_foot_clearance_reward(
|
|
env: ManagerBasedRLEnv,
|
|
asset_cfg: SceneEntityCfg,
|
|
sensor_cfg: SceneEntityCfg,
|
|
command_name: str,
|
|
target_height: float,
|
|
std: float,
|
|
force_threshold: float = 5.0,
|
|
linear_threshold: float = 0.1,
|
|
angular_threshold: float = 0.1,
|
|
) -> torch.Tensor:
|
|
"""Reward swing-foot height relative to the supporting foot."""
|
|
if std <= 0.0:
|
|
raise ValueError("std must be positive.")
|
|
|
|
asset = env.scene[asset_cfg.name]
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
forces = _to_env_device(env, contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :])
|
|
contacts = torch.norm(forces, dim=-1) > force_threshold
|
|
single_stance = torch.sum(contacts.int(), dim=1) == 1
|
|
swing_feet = ~contacts
|
|
|
|
foot_height = asset.data.body_pos_w[:, asset_cfg.body_ids, 2]
|
|
if foot_height.shape[1] != contacts.shape[1]:
|
|
raise ValueError("asset_cfg and sensor_cfg must resolve the same number of feet.")
|
|
stance_height = torch.sum(foot_height * contacts, dim=1, keepdim=True)
|
|
swing_clearance = foot_height - stance_height
|
|
clearance_reward = torch.exp(-torch.square(swing_clearance - target_height) / (std * std))
|
|
clearance_reward = torch.sum(clearance_reward * swing_feet, dim=1)
|
|
|
|
moving = _command_is_moving(env, command_name, linear_threshold, angular_threshold)
|
|
return clearance_reward * single_stance * moving
|
|
|
|
|
|
def feet_slide(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
|
"""Penalize foot sliding while keeping contact tensors on the env device."""
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
contacts = _to_env_device(
|
|
env, contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0]
|
|
) > 1.0
|
|
asset = env.scene[asset_cfg.name]
|
|
body_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2]
|
|
return torch.sum(body_vel.norm(dim=-1) * contacts, dim=1)
|
|
|
|
|
|
|
|
def feet_position(env,
|
|
asset_cfg: SceneEntityCfg,
|
|
command_name: str,
|
|
stand_threshold: float = 0.06,
|
|
ankle_distance: float = 0.22,
|
|
base_height_target: float = 0.82,
|
|
) -> torch.Tensor:
|
|
"""Reward keeping feet near a desired stance when standing; otherwise return 1."""
|
|
commands = env.command_manager.get_command(command_name)
|
|
stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & (
|
|
torch.abs(commands[:, 2]) < stand_threshold
|
|
)
|
|
asset = env.scene[asset_cfg.name]
|
|
|
|
|
|
feet_pos_w = asset.data.body_pos_w[:, asset_cfg.body_ids, :]
|
|
base_pos_w = asset.data.root_pos_w
|
|
base_quat_w = asset.data.root_quat_w
|
|
|
|
# isolate yaw heading; zero roll/pitch
|
|
r, p, y = euler_xyz_from_quat(base_quat_w)
|
|
heading_quat = quat_from_euler_xyz(torch.zeros_like(r), torch.zeros_like(p), y)
|
|
feet_pos_rel = feet_pos_w - base_pos_w.unsqueeze(1)
|
|
# Expand heading quaternions per foot to satisfy broadcasting expected by quat_apply_inverse.
|
|
num_envs, num_feet, _ = feet_pos_rel.shape
|
|
heading_quat_per_foot = heading_quat.unsqueeze(1).expand(-1, num_feet, -1).reshape(-1, 4)
|
|
feet_pos_rel_flat = feet_pos_rel.reshape(-1, 3)
|
|
feet_pos_heading = quat_apply_inverse(heading_quat_per_foot, feet_pos_rel_flat).reshape(num_envs, num_feet, 3)
|
|
|
|
desired_x = torch.zeros((num_envs, num_feet), device=feet_pos_heading.device)
|
|
desired_y = torch.cat(
|
|
(
|
|
(ankle_distance * 0.5) * torch.ones((num_envs, num_feet // 2), device=feet_pos_heading.device),
|
|
(-ankle_distance * 0.5) * torch.ones((num_envs, num_feet - num_feet // 2), device=feet_pos_heading.device),
|
|
),
|
|
dim=1,
|
|
)
|
|
desired_z = -(base_height_target - 0.045) * torch.ones((num_envs, num_feet), device=feet_pos_heading.device)
|
|
desired = torch.stack((desired_x, desired_y, desired_z), dim=-1)
|
|
|
|
position_error = torch.sum(torch.abs(feet_pos_heading - desired), dim=(1, 2))
|
|
reward_stand = torch.exp(-position_error * 3.0)
|
|
return torch.where(stand_command, reward_stand, torch.ones_like(reward_stand))
|
|
|
|
|
|
def feet_position_relative_to_body(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
reference_body_cfg: SceneEntityCfg,
|
|
command_name: str,
|
|
desired_foot_positions: tuple[tuple[float, float, float], ...],
|
|
stand_threshold: float = 0.06,
|
|
heading_yaw_offset: float = 0.0,
|
|
scale: float = 3.0,
|
|
) -> torch.Tensor:
|
|
"""Reward standing foot positions relative to a configured reference body."""
|
|
reference_body_id = _single_body_id(reference_body_cfg, "feet_position_relative_to_body")
|
|
commands = env.command_manager.get_command(command_name)
|
|
stand_command = (torch.norm(commands[:, :2], dim=1) < stand_threshold) & (
|
|
torch.abs(commands[:, 2]) < stand_threshold
|
|
)
|
|
asset = env.scene[asset_cfg.name]
|
|
reference_asset = env.scene[reference_body_cfg.name]
|
|
|
|
feet_pos_w = asset.data.body_pos_w[:, asset_cfg.body_ids, :]
|
|
reference_pos_w = reference_asset.data.body_pos_w[:, reference_body_id, :]
|
|
reference_quat_w = reference_asset.data.body_quat_w[:, reference_body_id, :]
|
|
|
|
num_envs, num_feet, _ = feet_pos_w.shape
|
|
feet_pos_rel = feet_pos_w - reference_pos_w.unsqueeze(1)
|
|
heading_quat = _heading_quat_with_offset(reference_quat_w, heading_yaw_offset)
|
|
heading_quat_per_foot = heading_quat.unsqueeze(1).expand(-1, num_feet, -1).reshape(-1, 4)
|
|
feet_pos_heading = quat_apply_inverse(heading_quat_per_foot, feet_pos_rel.reshape(-1, 3)).reshape(num_envs, num_feet, 3)
|
|
|
|
desired = torch.tensor(desired_foot_positions, dtype=feet_pos_heading.dtype, device=feet_pos_heading.device)
|
|
if desired.shape != (num_feet, 3):
|
|
raise ValueError(f"desired_foot_positions must have shape ({num_feet}, 3), got {tuple(desired.shape)}.")
|
|
position_error = torch.sum(torch.abs(feet_pos_heading - desired.unsqueeze(0)), dim=(1, 2))
|
|
reward_stand = torch.exp(-position_error * scale)
|
|
return torch.where(stand_command, reward_stand, torch.ones_like(reward_stand))
|
|
|
|
|
|
def feet_regulation(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
sensor_cfg: SceneEntityCfg,
|
|
ankle_height: float = 0.045,
|
|
base_height_target: float = 0.82,
|
|
clearance_clip: float = 1.0,
|
|
scale: float = 40.0,
|
|
) -> torch.Tensor:
|
|
"""Penalize feet moving fast while too close to terrain."""
|
|
asset = env.scene[asset_cfg.name]
|
|
raycaster = env.scene.sensors[sensor_cfg.name]
|
|
|
|
# terrain height estimate from rays
|
|
ray_hits = raycaster.data.ray_hits_w
|
|
if ray_hits is None:
|
|
terrain_height = torch.zeros(asset.data.body_pos_w.shape[0], device=asset.data.body_pos_w.device)
|
|
else:
|
|
# use the highest hit point in world frame as terrain height reference
|
|
terrain_height = torch.max(ray_hits[..., 2], dim=-1).values
|
|
|
|
foot_pos_w = asset.data.body_pos_w[:, asset_cfg.body_ids, 2]
|
|
foot_vel_xy = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2]
|
|
|
|
foot_clearance = torch.clamp(foot_pos_w - ankle_height - terrain_height.unsqueeze(1), min=0.0, max=clearance_clip)
|
|
speed_term = torch.square(torch.norm(foot_vel_xy, dim=-1))
|
|
height_term = torch.exp(-foot_clearance / max(base_height_target, 1e-6) * scale)
|
|
reward = torch.sum(height_term * speed_term, dim=1)
|
|
return reward
|
|
|
|
|
|
def feet_landing_velocity(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
sensor_cfg: SceneEntityCfg,
|
|
velocity_threshold: float = 0.25,
|
|
power: float = 2.0,
|
|
) -> torch.Tensor:
|
|
"""Penalize high downward landing speed at first contact to reduce impact noise."""
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
first_contact = _to_env_device(env, contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids])
|
|
|
|
asset = env.scene[asset_cfg.name]
|
|
foot_vel_z = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, 2]
|
|
landing_speed = torch.clamp(-foot_vel_z - velocity_threshold, min=0.0)
|
|
penalty = torch.sum(torch.pow(landing_speed, power) * first_contact, dim=1)
|
|
return penalty
|
|
|
|
def feet_z_velocity(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
sensor_cfg: SceneEntityCfg,
|
|
power: float = 2.0,
|
|
) -> torch.Tensor:
|
|
asset = env.scene[asset_cfg.name]
|
|
foot_vel_z = asset.data.body_lin_vel_w[:, sensor_cfg.body_ids, 2]
|
|
z_speed = torch.clamp(-foot_vel_z, min=0.0)
|
|
penalty = torch.sum(torch.pow(z_speed, power), dim=1)
|
|
return penalty
|
|
|
|
|
|
def foot_sound_suppression(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
sensor_cfg: SceneEntityCfg,
|
|
max_delta: float = 0.5,
|
|
) -> torch.Tensor:
|
|
"""Penalize frame-to-frame changes in foot vertical velocity to suppress landing sounds."""
|
|
asset = env.scene[asset_cfg.name]
|
|
foot_vel_z = asset.data.body_lin_vel_w[:, sensor_cfg.body_ids, 2]
|
|
|
|
try:
|
|
prev_vel = env._prev_foot_z_vel
|
|
except AttributeError:
|
|
env._prev_foot_z_vel = foot_vel_z.clone()
|
|
return foot_vel_z.new_zeros(foot_vel_z.shape[0])
|
|
|
|
delta = foot_vel_z - prev_vel
|
|
delta = torch.clamp(delta, min=-max_delta, max=max_delta)
|
|
penalty = torch.sum(delta * delta, dim=1)
|
|
|
|
prev_vel.copy_(foot_vel_z)
|
|
reset_env_ids = env.reset_buf.nonzero(as_tuple=False).squeeze(-1)
|
|
if reset_env_ids.numel() > 0:
|
|
prev_vel[reset_env_ids] = foot_vel_z[reset_env_ids]
|
|
penalty[reset_env_ids] = 0.0
|
|
|
|
return penalty
|
|
|
|
|
|
|
|
def base_height_tracking(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), target_height: float = 0.82) -> torch.Tensor:
|
|
"""Reward keeping the base height near a target height."""
|
|
asset = env.scene[asset_cfg.name]
|
|
height_error = torch.abs(asset.data.root_pos_w[:, 2] - target_height)
|
|
return torch.exp(-height_error * 30.0)
|
|
|
|
|
|
def body_height_tracking(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
target_height: float,
|
|
scale: float = 30.0,
|
|
) -> torch.Tensor:
|
|
"""Reward keeping a configured body height near a target height."""
|
|
body_id = _single_body_id(asset_cfg, "body_height_tracking")
|
|
asset = env.scene[asset_cfg.name]
|
|
height_error = torch.abs(asset.data.body_pos_w[:, body_id, 2] - target_height)
|
|
return torch.exp(-height_error * scale)
|
|
|
|
|
|
def body_vertical_velocity_l2(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
deadband: float = 0.0,
|
|
) -> torch.Tensor:
|
|
"""Penalize vertical body velocity outside a small natural-motion deadband."""
|
|
body_id = _single_body_id(asset_cfg, "body_vertical_velocity_l2")
|
|
asset = env.scene[asset_cfg.name]
|
|
vertical_speed = torch.abs(asset.data.body_lin_vel_w[:, body_id, 2])
|
|
return torch.square(torch.clamp(vertical_speed - deadband, min=0.0))
|
|
|
|
|
|
def body_roll_pitch_ang_vel_l2(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
deadband: float = 0.0,
|
|
) -> torch.Tensor:
|
|
"""Penalize horizontal angular speed while allowing normal gait oscillation."""
|
|
body_id = _single_body_id(asset_cfg, "body_roll_pitch_ang_vel_l2")
|
|
asset = env.scene[asset_cfg.name]
|
|
horizontal_ang_speed = torch.linalg.norm(asset.data.body_ang_vel_w[:, body_id, :2], dim=1)
|
|
return torch.square(torch.clamp(horizontal_ang_speed - deadband, min=0.0))
|
|
|
|
|
|
def joint_deviation_l1_with_deadband(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
deadband: float = 0.0,
|
|
) -> torch.Tensor:
|
|
"""Penalize joint deviation from defaults while allowing a small neutral range."""
|
|
if deadband < 0.0:
|
|
raise ValueError("deadband must be non-negative.")
|
|
asset = env.scene[asset_cfg.name]
|
|
joint_error = torch.abs(
|
|
asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
|
|
)
|
|
return torch.sum(torch.clamp(joint_error - deadband, min=0.0), dim=1)
|
|
|
|
|
|
def _forward_swing_speed_scale(
|
|
commands: torch.Tensor,
|
|
min_forward_speed: float,
|
|
full_swing_speed: float,
|
|
min_swing_scale: float,
|
|
) -> torch.Tensor:
|
|
"""Scale arm swing with commanded forward speed while keeping slow walking visible."""
|
|
if full_swing_speed <= min_forward_speed:
|
|
raise ValueError("full_swing_speed must be greater than min_forward_speed.")
|
|
if not 0.0 <= min_swing_scale <= 1.0:
|
|
raise ValueError("min_swing_scale must be in [0, 1].")
|
|
|
|
forward_speed = torch.abs(commands[:, 0])
|
|
speed_progress = torch.clamp(
|
|
(forward_speed - min_forward_speed) / (full_swing_speed - min_forward_speed),
|
|
min=0.0,
|
|
max=1.0,
|
|
)
|
|
moving = forward_speed >= min_forward_speed
|
|
moving_scale = min_swing_scale + (1.0 - min_swing_scale) * speed_progress
|
|
return torch.where(moving, moving_scale, torch.zeros_like(moving_scale))
|
|
|
|
|
|
def cross_body_arm_swing_reward(
|
|
env,
|
|
arm_asset_cfg: SceneEntityCfg,
|
|
feet_asset_cfg: SceneEntityCfg,
|
|
reference_body_cfg: SceneEntityCfg,
|
|
command_name: str,
|
|
heading_yaw_offset: float = 0.0,
|
|
min_forward_speed: float = 0.1,
|
|
full_swing_speed: float = 0.6,
|
|
min_swing_scale: float = 0.0,
|
|
phase_distance: float = 0.25,
|
|
shoulder_amplitude: float = 0.18,
|
|
elbow_flexion: float = 0.15,
|
|
shoulder_phase_signs: tuple[float, float] = (-1.0, -1.0),
|
|
elbow_flexion_signs: tuple[float, float] = (1.0, -1.0),
|
|
std: float = 0.2,
|
|
) -> torch.Tensor:
|
|
"""Track speed-scaled shoulder and elbow targets that oppose the leg phase.
|
|
|
|
Arm joints must be ordered as left/right shoulder pitch followed by left/right elbow flexion.
|
|
Feet must be ordered left then right.
|
|
"""
|
|
if phase_distance <= 0.0 or std <= 0.0:
|
|
raise ValueError("phase_distance and std must be positive.")
|
|
|
|
joint_ids = arm_asset_cfg.joint_ids
|
|
foot_ids = feet_asset_cfg.body_ids
|
|
if joint_ids is None or isinstance(joint_ids, slice) or len(joint_ids) != 4:
|
|
raise ValueError("cross_body_arm_swing_reward requires four ordered arm joint ids.")
|
|
if foot_ids is None or isinstance(foot_ids, slice) or len(foot_ids) != 2:
|
|
raise ValueError("cross_body_arm_swing_reward requires two ordered foot body ids.")
|
|
|
|
reference_body_id = _single_body_id(reference_body_cfg, "cross_body_arm_swing_reward")
|
|
asset = env.scene[arm_asset_cfg.name]
|
|
feet_asset = env.scene[feet_asset_cfg.name]
|
|
reference_asset = env.scene[reference_body_cfg.name]
|
|
|
|
feet_pos_w = feet_asset.data.body_pos_w[:, foot_ids, :]
|
|
reference_pos_w = reference_asset.data.body_pos_w[:, reference_body_id, :]
|
|
reference_quat_w = reference_asset.data.body_quat_w[:, reference_body_id, :]
|
|
feet_pos_rel = feet_pos_w - reference_pos_w.unsqueeze(1)
|
|
|
|
num_envs = feet_pos_w.shape[0]
|
|
heading_quat = _heading_quat_with_offset(reference_quat_w, heading_yaw_offset)
|
|
heading_quat = heading_quat.unsqueeze(1).expand(-1, 2, -1).reshape(-1, 4)
|
|
feet_pos_heading = quat_apply_inverse(heading_quat, feet_pos_rel.reshape(-1, 3)).reshape(num_envs, 2, 3)
|
|
leg_phase = torch.clamp(
|
|
(feet_pos_heading[:, 0, 0] - feet_pos_heading[:, 1, 0]) / phase_distance,
|
|
min=-1.0,
|
|
max=1.0,
|
|
)
|
|
|
|
commands = env.command_manager.get_command(command_name)
|
|
speed_scale = _forward_swing_speed_scale(
|
|
commands,
|
|
min_forward_speed=min_forward_speed,
|
|
full_swing_speed=full_swing_speed,
|
|
min_swing_scale=min_swing_scale,
|
|
)
|
|
|
|
joint_pos = asset.data.joint_pos[:, joint_ids]
|
|
target_pos = asset.data.default_joint_pos[:, joint_ids].clone()
|
|
shoulder_signs = joint_pos.new_tensor(shoulder_phase_signs)
|
|
elbow_signs = joint_pos.new_tensor(elbow_flexion_signs)
|
|
target_pos[:, :2] += (
|
|
shoulder_amplitude * speed_scale * leg_phase
|
|
).unsqueeze(1) * shoulder_signs.unsqueeze(0)
|
|
target_pos[:, 2:] += (elbow_flexion * speed_scale).unsqueeze(1) * elbow_signs.unsqueeze(0)
|
|
|
|
mean_square_error = torch.mean(torch.square(joint_pos - target_pos), dim=1)
|
|
return torch.exp(-mean_square_error / (std * std))
|
|
|
|
|
|
def contralateral_arm_phase_reward(
|
|
env,
|
|
arm_body_cfg: SceneEntityCfg,
|
|
feet_asset_cfg: SceneEntityCfg,
|
|
reference_body_cfg: SceneEntityCfg,
|
|
command_name: str,
|
|
heading_yaw_offset: float = 0.0,
|
|
min_forward_speed: float = 0.1,
|
|
full_swing_speed: float = 0.6,
|
|
min_swing_scale: float = 0.0,
|
|
phase_distance: float = 0.25,
|
|
arm_phase_amplitude: float = 0.1,
|
|
std: float = 0.05,
|
|
) -> torch.Tensor:
|
|
"""Reward physical left/right arm motion that is opposite to the same-side leg.
|
|
|
|
Arm bodies and feet must both be ordered left then right. When the left foot is ahead,
|
|
the right arm body is expected to be ahead of the left arm body, and vice versa.
|
|
Using upper-arm/elbow bodies prevents wrist motion from satisfying the phase target.
|
|
"""
|
|
if phase_distance <= 0.0 or arm_phase_amplitude <= 0.0 or std <= 0.0:
|
|
raise ValueError("phase_distance, arm_phase_amplitude, and std must be positive.")
|
|
|
|
arm_body_ids = arm_body_cfg.body_ids
|
|
foot_ids = feet_asset_cfg.body_ids
|
|
if arm_body_ids is None or isinstance(arm_body_ids, slice) or len(arm_body_ids) != 2:
|
|
raise ValueError("contralateral_arm_phase_reward requires two ordered arm body ids.")
|
|
if foot_ids is None or isinstance(foot_ids, slice) or len(foot_ids) != 2:
|
|
raise ValueError("contralateral_arm_phase_reward requires two ordered foot body ids.")
|
|
|
|
reference_body_id = _single_body_id(reference_body_cfg, "contralateral_arm_phase_reward")
|
|
arm_asset = env.scene[arm_body_cfg.name]
|
|
feet_asset = env.scene[feet_asset_cfg.name]
|
|
reference_asset = env.scene[reference_body_cfg.name]
|
|
|
|
reference_pos_w = reference_asset.data.body_pos_w[:, reference_body_id, :]
|
|
reference_quat_w = reference_asset.data.body_quat_w[:, reference_body_id, :]
|
|
heading_quat = _heading_quat_with_offset(reference_quat_w, heading_yaw_offset)
|
|
|
|
arm_pos_rel = arm_asset.data.body_pos_w[:, arm_body_ids, :] - reference_pos_w.unsqueeze(1)
|
|
feet_pos_rel = feet_asset.data.body_pos_w[:, foot_ids, :] - reference_pos_w.unsqueeze(1)
|
|
num_envs = reference_pos_w.shape[0]
|
|
heading_quat_pairs = heading_quat.unsqueeze(1).expand(-1, 2, -1).reshape(-1, 4)
|
|
arm_pos_heading = quat_apply_inverse(heading_quat_pairs, arm_pos_rel.reshape(-1, 3)).reshape(num_envs, 2, 3)
|
|
feet_pos_heading = quat_apply_inverse(heading_quat_pairs, feet_pos_rel.reshape(-1, 3)).reshape(num_envs, 2, 3)
|
|
|
|
leg_phase = torch.clamp(
|
|
(feet_pos_heading[:, 0, 0] - feet_pos_heading[:, 1, 0]) / phase_distance,
|
|
min=-1.0,
|
|
max=1.0,
|
|
)
|
|
commands = env.command_manager.get_command(command_name)
|
|
speed_scale = _forward_swing_speed_scale(
|
|
commands,
|
|
min_forward_speed=min_forward_speed,
|
|
full_swing_speed=full_swing_speed,
|
|
min_swing_scale=min_swing_scale,
|
|
)
|
|
|
|
physical_arm_phase = arm_pos_heading[:, 1, 0] - arm_pos_heading[:, 0, 0]
|
|
target_arm_phase = arm_phase_amplitude * speed_scale * leg_phase
|
|
phase_error = physical_arm_phase - target_arm_phase
|
|
return torch.exp(-torch.square(phase_error) / (std * std))
|
|
|
|
|
|
def energy_cost(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
|
"""Penalize energy consumption approximated by the sum of squared joint torques."""
|
|
asset = env.scene[asset_cfg.name]
|
|
joint_torques = asset.data.applied_torque[:, :]
|
|
joint_vel = asset.data.joint_vel[:, :]
|
|
power = joint_torques * joint_vel
|
|
energy = torch.sum(torch.abs(power), dim=1)
|
|
return energy
|
|
|
|
def feet_orientation(env, asset_cfg: SceneEntityCfg, command_name: str, stand_threshold: float = 0.06) -> torch.Tensor:
|
|
"""Reward aligning feet orientation; ignore yaw error while turning."""
|
|
commands = env.command_manager.get_command(command_name)
|
|
yaw_command = torch.abs(commands[:, 2]) > stand_threshold
|
|
|
|
asset = env.scene[asset_cfg.name]
|
|
feet_quat = asset.data.body_quat_w[:, asset_cfg.body_ids, :]
|
|
base_quat = asset.data.root_quat_w
|
|
|
|
num_envs, num_feet, _ = feet_quat.shape
|
|
feet_flat = feet_quat.reshape(-1, 4)
|
|
roll, pitch, yaw = euler_xyz_from_quat(feet_flat)
|
|
roll = roll.reshape(num_envs, num_feet)
|
|
pitch = pitch.reshape(num_envs, num_feet)
|
|
yaw = yaw.reshape(num_envs, num_feet)
|
|
|
|
_, _, base_yaw = euler_xyz_from_quat(base_quat)
|
|
|
|
feet_roll_pitch_error = torch.sum(torch.abs(torch.stack((roll, pitch), dim=-1)), dim=-1)
|
|
feet_yaw_error = torch.abs(wrap_to_pi(yaw - base_yaw.unsqueeze(1)))
|
|
|
|
rew = torch.sum(feet_roll_pitch_error + feet_yaw_error, dim=1)
|
|
rew[yaw_command] = torch.sum(feet_roll_pitch_error[yaw_command], dim=1)
|
|
return torch.exp(-rew * 2.0)
|
|
|
|
|
|
def feet_orientation_relative_to_body(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
reference_body_cfg: SceneEntityCfg,
|
|
command_name: str,
|
|
stand_threshold: float = 0.06,
|
|
heading_yaw_offset: float = 0.0,
|
|
foot_frame_offsets_rpy: tuple[tuple[float, float, float], ...] | None = None,
|
|
scale: float = 2.0,
|
|
) -> torch.Tensor:
|
|
"""Reward physical sole orientation relative to a configured reference body's heading."""
|
|
reference_body_id = _single_body_id(reference_body_cfg, "feet_orientation_relative_to_body")
|
|
commands = env.command_manager.get_command(command_name)
|
|
yaw_command = torch.abs(commands[:, 2]) > stand_threshold
|
|
|
|
asset = env.scene[asset_cfg.name]
|
|
reference_asset = env.scene[reference_body_cfg.name]
|
|
feet_quat = asset.data.body_quat_w[:, asset_cfg.body_ids, :]
|
|
reference_quat = reference_asset.data.body_quat_w[:, reference_body_id, :]
|
|
|
|
num_envs, num_feet, _ = feet_quat.shape
|
|
feet_flat = feet_quat.reshape(-1, 4)
|
|
if foot_frame_offsets_rpy is not None:
|
|
offsets = torch.tensor(foot_frame_offsets_rpy, dtype=feet_quat.dtype, device=feet_quat.device)
|
|
if offsets.shape != (num_feet, 3):
|
|
raise ValueError(f"foot_frame_offsets_rpy must have shape ({num_feet}, 3), got {tuple(offsets.shape)}.")
|
|
offset_quat = quat_from_euler_xyz(offsets[:, 0], offsets[:, 1], offsets[:, 2])
|
|
offset_quat = offset_quat.unsqueeze(0).expand(num_envs, -1, -1).reshape(-1, 4)
|
|
feet_flat = quat_mul(feet_flat, offset_quat)
|
|
|
|
roll, pitch, yaw = euler_xyz_from_quat(feet_flat)
|
|
roll = roll.reshape(num_envs, num_feet)
|
|
pitch = pitch.reshape(num_envs, num_feet)
|
|
yaw = yaw.reshape(num_envs, num_feet)
|
|
|
|
reference_yaw = _heading_yaw_with_offset(reference_quat, heading_yaw_offset)
|
|
feet_roll_pitch_error = torch.sum(torch.abs(torch.stack((roll, pitch), dim=-1)), dim=-1)
|
|
feet_yaw_error = torch.abs(wrap_to_pi(yaw - reference_yaw.unsqueeze(1)))
|
|
|
|
rew = torch.sum(feet_roll_pitch_error + feet_yaw_error, dim=1)
|
|
rew[yaw_command] = torch.sum(feet_roll_pitch_error[yaw_command], dim=1)
|
|
return torch.exp(-rew * scale)
|
|
|
|
|
|
def base_orientation(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
|
"""Reward keeping the base roll/pitch near zero."""
|
|
asset = env.scene[asset_cfg.name]
|
|
roll, pitch, yaw = euler_xyz_from_quat(asset.data.root_quat_w)
|
|
base_euler = torch.stack((roll, pitch, yaw), dim=-1)
|
|
return torch.exp(-torch.sum(torch.abs(base_euler[:, :2]), dim=-1) * 10.0)
|
|
|
|
|
|
def body_orientation(env, asset_cfg: SceneEntityCfg, scale: float = 10.0) -> torch.Tensor:
|
|
"""Reward keeping a configured body's roll/pitch near zero."""
|
|
body_id = _single_body_id(asset_cfg, "body_orientation")
|
|
asset = env.scene[asset_cfg.name]
|
|
roll, pitch, yaw = euler_xyz_from_quat(asset.data.body_quat_w[:, body_id, :])
|
|
return torch.exp(-torch.sum(torch.abs(torch.stack((roll, pitch), dim=-1)), dim=-1) * scale)
|
|
|
|
|
|
def body_yaw_alignment(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
reference_body_cfg: SceneEntityCfg,
|
|
heading_yaw_offset: float = 0.0,
|
|
reference_heading_yaw_offset: float = 0.0,
|
|
scale: float = 4.0,
|
|
) -> torch.Tensor:
|
|
"""Reward keeping one body's heading aligned with another body's heading."""
|
|
body_id = _single_body_id(asset_cfg, "body_yaw_alignment")
|
|
reference_body_id = _single_body_id(reference_body_cfg, "body_yaw_alignment")
|
|
asset = env.scene[asset_cfg.name]
|
|
reference_asset = env.scene[reference_body_cfg.name]
|
|
body_yaw = _heading_yaw_with_offset(asset.data.body_quat_w[:, body_id, :], heading_yaw_offset)
|
|
reference_yaw = _heading_yaw_with_offset(
|
|
reference_asset.data.body_quat_w[:, reference_body_id, :], reference_heading_yaw_offset
|
|
)
|
|
return torch.exp(-torch.abs(wrap_to_pi(body_yaw - reference_yaw)) * scale)
|
|
|
|
|
|
def body_yaw_rate_difference_l2(
|
|
env,
|
|
asset_cfg: SceneEntityCfg,
|
|
reference_body_cfg: SceneEntityCfg,
|
|
) -> torch.Tensor:
|
|
"""Penalize relative world-frame yaw rate between two configured bodies."""
|
|
body_id = _single_body_id(asset_cfg, "body_yaw_rate_difference_l2")
|
|
reference_body_id = _single_body_id(reference_body_cfg, "body_yaw_rate_difference_l2")
|
|
asset = env.scene[asset_cfg.name]
|
|
reference_asset = env.scene[reference_body_cfg.name]
|
|
yaw_rate = asset.data.body_ang_vel_w[:, body_id, 2]
|
|
reference_yaw_rate = reference_asset.data.body_ang_vel_w[:, reference_body_id, 2]
|
|
return torch.square(yaw_rate - reference_yaw_rate)
|
|
|
|
|
|
def reward_waist_pos(
|
|
env,
|
|
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
|
joint_names: list[str] | str | None = None,
|
|
scale: float = 5.0,
|
|
) -> torch.Tensor:
|
|
"""Penalize motion on specific waist joints; joint_names provided by user."""
|
|
asset = env.scene[asset_cfg.name]
|
|
if joint_names is not None:
|
|
joint_ids = asset.find_joints(joint_names, preserve_order=True)[0]
|
|
elif asset_cfg.joint_ids is not None:
|
|
joint_ids = asset_cfg.joint_ids
|
|
else:
|
|
raise ValueError("penalize_waist_joint_motion requires joint_names or asset_cfg.joint_ids.")
|
|
|
|
pos_err = torch.abs(asset.data.joint_pos[:, joint_ids])
|
|
return torch.exp(-torch.sum(pos_err, dim=1) * scale)
|
|
|
|
def penalize_foot_stumble(env, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
|
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
|
contacts = _to_env_device(
|
|
env, contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0]
|
|
) > 1.0
|
|
asset = env.scene[asset_cfg.name]
|
|
body_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2]
|
|
return torch.sum(body_vel.norm(dim=-1) * contacts, dim=1)
|
|
|
|
def joint_deviation_exp(
|
|
env: ManagerBasedRLEnv,
|
|
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
|
tolerance: float = 0.1,
|
|
scale: float = 3.0,
|
|
max_err: float = 50.0,
|
|
) -> torch.Tensor:
|
|
"""Penalize joint positions deviating from defaults, in an exponential, configurable way."""
|
|
asset: Articulation = env.scene[asset_cfg.name]
|
|
joint_ids = asset_cfg.joint_ids if asset_cfg.joint_ids is not None else slice(None)
|
|
joint_pos = asset.data.joint_pos[:, joint_ids]
|
|
default_pos = getattr(asset.data, "default_joint_pos", None)
|
|
if default_pos is not None:
|
|
default_pos = default_pos[:, joint_ids]
|
|
else:
|
|
default_pos = torch.zeros_like(joint_pos)
|
|
print("Warning: joint_deviation_exp reward called but default_joint_pos not set in asset data; assuming zeros.")
|
|
|
|
joint_error = torch.norm(joint_pos - default_pos, dim=1)
|
|
joint_error = torch.clamp(joint_error - tolerance, min=0.0, max=max_err)
|
|
return torch.exp(-joint_error * scale)
|