核心奖励函数

• 理解核心奖励函数，包含跟踪奖励，能量奖励，姿态奖励，足部奖励和其他奖励等

概述

核心奖励函数定义在 playground/common/rewards.py 中，包含 18 个独立的奖励和代价函数。这些函数是纯 JAX 函数，接受张量输入并返回标量输出，设计为可组合的构建块。

跟踪奖励

reward_tracking_lin_vel

作用：引导机器人按照命令速度前进或横向移动。

def reward_tracking_lin_vel(commands, local_vel, tracking_sigma):
    y_tol = 0.1  # 横向容差
    error_x = jp.square(commands[0] - local_vel[0])
    error_y = jp.clip(jp.abs(local_vel[1] - commands[1]) - y_tol, 0.0, None)
    lin_vel_error = error_x + jp.square(error_y)
    return jp.nan_to_num(jp.exp(-lin_vel_error / tracking_sigma))

设计思路：

• 前进方向（x）：严格跟踪，误差平方
• 横向方向（y）：允许 ±0.1 m/s 的容差（避免不必要的侧滑补偿）
• 使用高斯核 exp(-error/sigma) 将范围映射到 [0, 1]

reward_tracking_ang_vel

作用：引导机器人按命令旋转。

def reward_tracking_ang_vel(commands, ang_vel, tracking_sigma):
    ang_vel_error = jp.square(commands[2] - ang_vel[2])
    return jp.nan_to_num(jp.exp(-ang_vel_error / tracking_sigma))

只关心 Z 轴（偏航）角速度，忽略 X/Y 轴角速度。

能量奖励

cost_torques

作用：最小化关节力矩消耗，提高能效。

def cost_torques(torques):
    return jp.nan_to_num(jp.sum(jp.square(torques)))

力矩平方和意味着：

• 较小的力矩消耗较小 → 无需精确为零，只需保持较低
• 平方放大大力矩的惩罚 → 避免过大的控制力

cost_energy

作用：最小化机械功（力矩 × 速度）。

def cost_energy(qvel, qfrc_actuator):
    return jp.nan_to_num(jp.sum(jp.abs(qvel) * jp.abs(qfrc_actuator)))

与 cost_torques 不同，此函数考虑了关节运动速度，更接近真实功耗。

cost_action_rate

作用：鼓励平滑的动作序列，防止抖动。

def cost_action_rate(act, last_act):
    return jp.nan_to_num(jp.sum(jp.square(act - last_act)))

相邻动作差的平方和：

• 动作变化越小 → 代价越低
• 防止策略输出高频抖动
• 类似于 L2 正则化在动作空间上的效果

姿态奖励

cost_orientation

作用：保持机器人直立姿态。

def cost_orientation(torso_zaxis):
    return jp.nan_to_num(jp.sum(jp.square(torso_zaxis[:2])))

原理：

• torso_zaxis 是躯干 Z 轴在世界帧中的投影
• 完全直立时：Z 轴 = [0, 0, 1]，前两维为 0，代价 = 0
• 倾斜时：Z 轴在 XY 平面有分量，产生代价
• 平方使小倾斜几乎无代价，大倾斜代价显著

cost_base_height

作用：保持机器人基座高度在目标值附近。

def cost_base_height(base_height, base_height_target):
    return jp.nan_to_num(jp.square(base_height - base_height_target))

cost_stand_still

作用：在命令速度为零时，保持默认姿态不动。

def cost_stand_still(commands, qpos, qvel, default_pose, ignore_head=False):
    cmd_norm = jp.linalg.norm(commands[:3])
    
    if not ignore_head:
        pose_cost = jp.sum(jp.abs(qpos - default_pose))
        vel_cost = jp.sum(jp.abs(qvel))
    else:
        # 忽略头部，只考虑腿部
        left_leg = qpos[:5]; right_leg = qpos[9:]
        left_def = default_pose[:5]; right_def = default_pose[9:]
        pose_cost = jp.sum(jp.abs(left_leg - left_def)) + \
                    jp.sum(jp.abs(right_leg - right_def))
        vel_cost = jp.sum(jp.abs(left_leg_vel)) + jp.sum(jp.abs(right_leg_vel))
    
    return jp.nan_to_num(pose_cost + vel_cost) * (cmd_norm < 0.01)

关键设计：(cmd_norm < 0.01) 条件确保仅当命令接近零时才施加此代价。运动时允许偏离默认姿态。

cost_pose

作用：加权关节位置偏差惩罚。

def cost_pose(qpos, default_pose, weights):
    return jp.nan_to_num(jp.sum(jp.square(qpos - default_pose) * weights))

cost_head_pos

作用：控制头部位置跟踪命令。

def cost_head_pos(joints_qpos, joints_qvel, cmd):
    cmd_norm = jp.linalg.norm(cmd[:3])
    head_cmd = cmd[3:]  # [neck_pitch, head_pitch, head_yaw, head_roll]
    head_pos = joints_qpos[5:9]
    head_pos_error = jp.sum(jp.square(head_pos - head_cmd))
    return jp.nan_to_num(head_pos_error) * (cmd_norm > 0.01)

头部仅在机器人实际运动时跟随命令。

足部奖励

reward_feet_air_time

作用：鼓励足部腾空时间（即每一步离地足够久）。

def reward_feet_air_time(air_time, first_contact, commands, 
                          threshold_min=0.1, threshold_max=0.5):
    cmd_norm = jp.linalg.norm(commands[:3])
    air_time = (air_time - threshold_min) * first_contact
    air_time = jp.clip(air_time, max=threshold_max - threshold_min)
    reward = jp.sum(air_time)
    reward *= cmd_norm > 0.01
    return jp.nan_to_num(reward)

• 仅在着地瞬间 (first_contact) 计算
• 腾空时间低于 0.1s → 不计奖励
• 腾空时间超过 0.5s → 不再增加奖励
• 足部合理腾空有利于稳定步态

cost_feet_height

作用：控制足部摆动最高点在合理范围。

def cost_feet_height(swing_peak, first_contact, max_foot_height):
    error = swing_peak / max_foot_height - 1.0
    return jp.nan_to_num(jp.sum(jp.square(error) * first_contact))

cost_feet_slip

作用：惩罚足部打滑。

def cost_feet_slip(contact, global_linvel):
    body_vel = global_linvel[:2]
    reward = jp.sum(jp.linalg.norm(body_vel, axis=-1) * contact)
    return jp.nan_to_num(reward)

当足部接触地面时，足部速度应接近零。

cost_feet_clearance

作用：控制足部摆动离地高度。

def cost_feet_clearance(feet_vel, foot_pos, max_foot_height):
    vel_norm = jp.sqrt(jp.linalg.norm(vel_xy, axis=-1))
    foot_z = foot_pos[..., -1]
    delta = jp.abs(foot_z - max_foot_height)
    return jp.nan_to_num(jp.sum(delta * vel_norm))

其他奖励

reward_base_y_swing

作用：鼓励基座横向摆动（模拟侧向体重转移）。

def reward_base_y_swing(base_y_speed, freq, amplitude, t, tracking_sigma):
    target_y_speed = amplitude * jp.sin(2 * jp.pi * freq * t)
    y_speed_error = jp.square(target_y_speed - base_y_speed)
    return jp.nan_to_num(jp.exp(-y_speed_error / tracking_sigma))

cost_joint_pos_limits

作用：惩罚关节超出软限制。

def cost_joint_pos_limits(qpos, soft_lowers, soft_uppers):
    out_of_limits = -jp.clip(qpos - soft_lowers, None, 0.0)
    out_of_limits += jp.clip(qpos - soft_uppers, 0.0, None)
    return jp.nan_to_num(jp.sum(out_of_limits))

cost_termination

作用：终止时施加惩罚。

def cost_termination(done):
    return done  # 0 或 1

reward_alive

作用：每步存活奖励，防止策略过早让机器人倒下。

def reward_alive():
    return jp.array(1.0)