From af3ba4704f12f5bd608df36551eb1ac4926d8da7 Mon Sep 17 00:00:00 2001
From: ChenXi <chenxiliu2405@gmail.com>
Date: Thu, 19 Mar 2026 09:25:20 -0400
Subject: [PATCH] Add feet_airtime loss

---
 rl_game/get_up/config/t1_env_cfg.py | 56 ++++++++++++++++++++++++++---
 1 file changed, 52 insertions(+), 4 deletions(-)

diff --git a/rl_game/get_up/config/t1_env_cfg.py b/rl_game/get_up/config/t1_env_cfg.py
index 0bd0ed8..c71d7c5 100644
--- a/rl_game/get_up/config/t1_env_cfg.py
+++ b/rl_game/get_up/config/t1_env_cfg.py
@@ -104,6 +104,36 @@ def is_standing_still(
 
     return env.extras["stable_timer"] > standing_time
 
+def feet_airtime_penalty_local(
+        env: ManagerBasedRLEnv,
+        sensor_cfg: SceneEntityCfg,
+        threshold: float = 1.0
+) -> torch.Tensor:
+    """
+    自定义滞空惩罚逻辑：
+    如果脚部的垂直合力小于阈值，说明脚离地了。
+    返回一个 Tensor，离地时为 1.0，着地时为 0.0。
+    """
+    # 1. 获取传感器对象
+    contact_sensor = env.scene.sensors.get(sensor_cfg.name)
+
+    if contact_sensor is None:
+        # 如果没搜到传感器，返回全 0，防止程序崩溃
+        return torch.zeros(env.num_envs, device=env.device)
+
+    # 2. 获取触地力 (num_envs, num_bodies_in_sensor, 3)
+    # 我们取所有被监测 Body (左右脚) 的 Z 轴推力
+    # 如果所有脚的力都小于 threshold，判定为“完全腾空”
+    foot_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
+    is_in_air = torch.all(foot_forces_z < threshold, dim=-1)
+
+    return is_in_air.float()
+
+def root_vel_z_l2_local(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+    # 严厉惩罚 Z 轴正向速度（向上窜）
+    vel_z = env.scene[asset_cfg.name].data.root_lin_vel_w[:, 2]
+    return torch.square(torch.clamp(vel_z, min=0.0))
+
 
 # --- 2. 配置类 ---
 
@@ -171,7 +201,7 @@ class T1GetUpRewardCfg:
     # 2. 【条件高度奖】：双高度判定（头+盆骨），且必须脚踩地
     height_with_feet = RewTerm(
         func=standing_with_feet_reward,
-        weight=25.0,  # 作为核心引导，增加权重
+        weight=5.0,  # 作为核心引导，增加权重
         params={
             "min_head_height": 1.10,
             "min_pelvis_height": 0.65,
@@ -183,7 +213,7 @@ class T1GetUpRewardCfg:
     # 3. 手臂撑地奖：辅助脱离地面阶段
     arm_push_support = RewTerm(
         func=arm_push_up_reward,
-        weight=2.0,
+        weight=3.0,
         params={"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_hand_link"])}
     )
 
@@ -198,10 +228,28 @@ class T1GetUpRewardCfg:
         }
     )
 
-    # 5. 成功终极大奖
+    # 5. 抑制跳跃：严厉惩罚向上窜的速度
+    root_vel_z_penalty = RewTerm(
+        func=root_vel_z_l2_local,
+        weight=-10.0,  # 增大负权重
+        params={"asset_cfg": SceneEntityCfg("robot")}
+    )
+
+    # 6. 抑制滞空 (Airtime Penalty)
+    # 如果脚离开地面，按时间扣分
+    feet_airtime = RewTerm(
+        func=feet_airtime_penalty_local,
+        weight=-10.0,
+        params={
+            "sensor_cfg": SceneEntityCfg("feet_contact_sensor"),
+            "threshold": 0.1,  # 超过 0.2s 离地就开始扣分
+        }
+    )
+
+    # 7. 成功终极大奖
     is_success = RewTerm(
         func=lambda env, keys: env.termination_manager.get_term(keys),
-        weight=1000.0,
+        weight=500.0,
         params={"keys": "standing_success"}
     )