Align DreamWaQ Go1 environment rewards and initialization
This commit is contained in:
@@ -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 --
|
||||
|
||||
@@ -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。
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -112,7 +112,7 @@ class rslrl:
|
||||
|
||||
# Runner 设置(严格对齐上游 LeggedRobotCfgPPO + Go1RoughCfgPPO)
|
||||
runner.seed = 5 # 上游 seed=5
|
||||
runner.max_iterations = 5000
|
||||
runner.max_iterations = 3000
|
||||
runner.num_steps_per_env = 24
|
||||
runner.experiment_name = "go1_dreamwaq_walk"
|
||||
runner.save_interval = 50
|
||||
|
||||
Reference in New Issue
Block a user