cyy/Motrixlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Align DreamWaQ Go1 environment rewards and initialization

Browse Source
This commit is contained in:
8x54zj-m
2026-07-22 14:08:10 +08:00
parent 0555776c11
commit 82f4997deb
5 changed files with 85 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
motrix_envs/src/motrix_envs/locomotion/go1/cfg.py
View File

@@ -118,6 +118,7 @@ class RewardConfig:
tracking_sigma: float = 0.25
max_foot_height: float = 0.1
only_positive_rewards: bool = True
# -- docs-tag-end: go1-reward-config --

54
motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py
View File

@@ -120,6 +120,13 @@ class DreamWaQCfg(Go1WalkNpEnvCfg):
self.height_points_x = tuple(np.linspace(-0.8, 0.8, 17).tolist())
self.height_points_y = tuple(np.linspace(-0.5, 0.5, 11).tolist())
self.commands = DreamWaQCfg.Commands()
self.init_state.pos = [0.0, 0.0, 0.34]
self.spawn_clearance = 0.34
self.init_state.default_joint_angles = {
"FL_hip": 0.1, "RL_hip": 0.1, "FR_hip": -0.1, "RR_hip": -0.1,
"FL_thigh": 0.8, "RL_thigh": 1.0, "FR_thigh": 0.8, "RR_thigh": 1.0,
"FL_calf": -1.5, "RL_calf": -1.5, "FR_calf": -1.5, "RR_calf": -1.5,
}
self.control_config.stiffness = 28.0
self.control_config.damping = 0.7
self.control_config.action_scale = 0.25 # 上游原值
@@ -131,23 +138,23 @@ class DreamWaQCfg(Go1WalkNpEnvCfg):
r = self.reward_config.scales
r.clear()
r.update({
"tracking_lin_vel": 1.5,
"tracking_ang_vel": 1.0,
"tracking_lin_vel": 1.0,
"tracking_ang_vel": 0.5,
"lin_vel_z": -2.0,
"ang_vel_xy": -0.10, # 抑制晃动
"orientation": -0.5, # 强制平稳姿态
"ang_vel_xy": -0.05,
"orientation": -0.2,
"dof_acc": -2.5e-7,
"base_height": -5.0,
"feet_air_time": 0.03,
"base_height": -10.0,
"feet_air_time": 0.1,
"action_rate": -0.01,
"joint_power": -2e-5,
"smoothness": -0.02,
"smoothness": -0.01,
"power_distribution": -10e-6,
"stand_still": -0.5,
"dof_pos_limits": -5.0, # 关节限位软约束（参考 M20 修改版）
"collision": -1.0, # 惩罚身体碰撞
"dof_pos_limits": -5.0,
"collision": -1.0,
})
self.reward_config.only_positive_rewards = False # 让坏行为负反馈直达策略
self.reward_config.only_positive_rewards = True
self.reward_config.tracking_sigma = 0.25
@@ -645,7 +652,9 @@ class DreamWaQTask(Go1WalkTask):
if spawn_abs is not None:
init_dof_pos[:, 2] = float(spawn_abs)
else:
init_dof_pos[:, 2] = terrain_z + getattr(self, "_spawn_clearance", 0.45)
init_dof_pos[:, 2] = terrain_z + getattr(
self, "_spawn_clearance", self.cfg.spawn_clearance
)
yaw = np.random.uniform(-np.pi, np.pi, size=num_reset)
init_dof_pos[:, 3] = 0.0
init_dof_pos[:, 4] = 0.0
@@ -711,7 +720,7 @@ class DreamWaQTask(Go1WalkTask):
def _reward_collision(self, data):
"""身体碰撞惩罚——膝/肩触地扣分。"""
cquerys = self._model.get_contact_query(data)
penal = cquerys.is_colliding(self.termination_check)
penal = cquerys.is_colliding(self.penalized_contact)
return np.any(penal.reshape(self._num_envs, -1), axis=1).astype(np.float32)
def update_terminated(self, state):
@@ -757,12 +766,12 @@ class DreamWaQTask(Go1WalkTask):
"""覆盖基类 — 在清零前保存 first_contact 供奖励函数使用。"""
feet_air_time = info["feet_air_time"]
contacts = info["contacts"]
# 保存着地瞬间的状态（脚刚触地，且之前腾空时间 > 0）
info["first_contact"] = (feet_air_time > 0.0) & contacts
info["air_time_at_contact"] = feet_air_time.copy()
# 基类逻辑：累加 dt，着地清零
contact_filt = contacts | info.get("last_contacts", False)
info["first_contact"] = (feet_air_time > 0.0) & contact_filt
feet_air_time = feet_air_time + self.cfg.ctrl_dt
feet_air_time = feet_air_time * (~contacts)
info["air_time_at_contact"] = feet_air_time.copy()
info["last_contacts"] = contacts.copy()
feet_air_time = feet_air_time * (~contact_filt)
info["feet_air_time"] = feet_air_time
return feet_air_time
@@ -775,16 +784,15 @@ class DreamWaQTask(Go1WalkTask):
air_time = info.get("air_time_at_contact")
if first_contact is None or air_time is None:
return np.zeros(self._num_envs, dtype=np.float32)
rew = np.sum(np.maximum(air_time - 0.25, 0.0) * first_contact, axis=1)
rew = np.sum((air_time - 0.5) * first_contact, axis=1)
rew *= np.linalg.norm(commands[:, :2], axis=1) > 0.1
return rew
def _reward_stand_still(self, data, commands):
"""惩罚有命令但速度接近零的'卡住'行为。"""
cmd_norm = np.linalg.norm(commands[:, :2], axis=1)
vel_norm = np.linalg.norm(self.get_local_linvel(data)[:, :2], axis=1)
stifled = (cmd_norm > 0.2) & (vel_norm < 0.1) # 有命令但基本不动
return stifled.astype(np.float32)
"""惩罚零速度命令下偏离默认姿态的关节。"""
return np.sum(np.abs(self.get_dof_pos(data) - self.default_angles), axis=1) * (
np.linalg.norm(commands[:, :2], axis=1) < 0.1
)
def update_reward(self, state):
"""存储各项奖励到 TensorBoard + 更新 state.reward。

19
motrix_envs/src/motrix_envs/locomotion/go1/walk_np.py
View File

@@ -150,6 +150,24 @@ class Go1WalkTask(NpEnv):
)
self.num_check = self.termination_contact.shape[0]
self.penalized_contact = None
for ground_index in self.ground:
for name in cfg.asset.penalize_contacts_on:
matches = [g for g in self._model.geom_names
if g is not None and name in g]
for geom_name in matches:
pair = np.array(
[[self._model.get_geom_index(geom_name), ground_index]],
dtype=np.uint32,
)
if self.penalized_contact is None:
self.penalized_contact = pair
else:
self.penalized_contact = np.append(
self.penalized_contact, pair, axis=0)
if self.penalized_contact is None:
self.penalized_contact = self.termination_contact.copy()
self.foot = None
for gournd_index in self.ground:
for i in self._model.geom_names:
@@ -293,6 +311,7 @@ class Go1WalkTask(NpEnv):
"last_dof_vel": np.zeros((num_reset, self._num_action), dtype=np.float32),
"feet_air_time": np.zeros((num_reset, self.foot_check_num), dtype=np.float32),
"contacts": np.zeros((num_reset, self.foot_check_num), dtype=np.bool),
"last_contacts": np.zeros((num_reset, self.foot_check_num), dtype=np.bool),
}
obs = self._get_obs(data, info)
return obs, info

34
motrix_envs/tests/test_dreamwaq_state_safety.py
View File

@@ -2,7 +2,7 @@ from types import SimpleNamespace
import numpy as np
from motrix_envs.locomotion.go1.dreamwaq import _sanitize_dof_pos
from motrix_envs.locomotion.go1.dreamwaq import DreamWaQCfg, DreamWaQTask, _sanitize_dof_pos
from motrix_envs.locomotion.go1.walk_np import Go1WalkTask
@@ -42,3 +42,35 @@ def test_apply_action_advances_three_frame_action_history():
np.testing.assert_array_equal(state.info["last_last_actions"], 1.0)
np.testing.assert_array_equal(state.info["last_actions"], 2.0)
np.testing.assert_array_equal(state.info["current_actions"], 3.0)
def test_dreamwaq_feet_air_time_matches_upstream_contact_filtering():
task = DreamWaQTask.__new__(DreamWaQTask)
task._cfg = SimpleNamespace(ctrl_dt=0.02)
task._num_envs = 1
info = {
"feet_air_time": np.array([[0.6, 0.0]], dtype=np.float32),
"contacts": np.array([[True, False]]),
"last_contacts": np.array([[False, False]]),
}
task.update_feet_air_time(info)
reward = task._reward_feet_air_time(
np.array([[1.0, 0.0, 0.0]], dtype=np.float32), info
)
np.testing.assert_allclose(reward, [0.12])
np.testing.assert_allclose(info["feet_air_time"], [[0.0, 0.02]])
np.testing.assert_array_equal(info["last_contacts"], info["contacts"])
def test_dreamwaq_config_matches_upstream_go1_defaults():
cfg = DreamWaQCfg()
assert cfg.reward_config.only_positive_rewards
assert cfg.reward_config.scales["tracking_lin_vel"] == 1.0
assert cfg.reward_config.scales["tracking_ang_vel"] == 0.5
assert cfg.reward_config.scales["feet_air_time"] == 0.1
assert cfg.reward_config.scales["base_height"] == -10.0
assert cfg.init_state.default_joint_angles["FL_hip"] == 0.1
assert cfg.init_state.default_joint_angles["RR_hip"] == -0.1