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

feat: DreamWaQ full replication — env, terrain, CENet, PPO

Browse Source
This commit is contained in:
8x54zj-m
2026-06-30 14:53:12 +08:00
parent c569f4d6d9
commit 422421d263
17 changed files with 2056 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

625
motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py Normal file
View File

@@ -0,0 +1,625 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""DreamWaQ environment for MotrixLab — 对齐 Manaro-Alpha/DreamWaQ。
架构:
CENet (VAE): history(225) → [128,64] → latent(16) + vel_est(3) = code(19)
Actor: code(19) + obs(45) = 64 → [512,256,128] → action(12)
Critic: privileged_obs(247) → [512,256,128] → value(1)
观测 (45-dim, 上游顺序):
ang_vel(3) + gravity(3) + commands(3) + joint_pos(12) + joint_vel(12) + actions(12)
特权观测 (247-dim):
obs(45) + base_vel(3) + 足部接触力(12) + heights(187)
注意：上游为 286（含全部 17 body × 3-axis 接触力），MotrixSim 仅用四足接触力
注册: "go1-dreamwaq-walk"
"""
import gymnasium as gym
import motrixsim as mtx
import numpy as np
from dataclasses import dataclass, field
import os
from motrix_envs import registry
from motrix_envs.locomotion.go1.cfg import Go1WalkNpEnvCfg
from motrix_envs.locomotion.go1.walk_np import Go1WalkTask
from motrix_envs.math import quaternion
_SCENE_PRINTED = False
def _scene_file():
"""选择地形场景。DREAMWAQ_TERRAIN=flat|pyramid|flat_stairs。默认 FLAT。"""
global _SCENE_PRINTED
terrain = os.environ.get("DREAMWAQ_TERRAIN", "flat").lower()
if terrain in ("flat_stairs", "flatstairs"):
fname = "scene_flat_stairs.xml"
elif terrain in ("pyramid", "terrain", "1"):
fname = "scene_dreamwaq_terrain.xml"
elif terrain in ("stairs", "stairs_terrain"):
fname = "scene_stairs_terrain.xml"
else:
fname = "scene_dreamwaq_flat.xml"
if not _SCENE_PRINTED:
print(f"[DreamWaQ] scene = {fname} [set DREAMWAQ_TERRAIN to switch]")
_SCENE_PRINTED = True
return os.path.join(os.path.dirname(__file__), "xmls", fname)
# ═══════════════════════════════════════════════════════════════════════
# Config
# ═══════════════════════════════════════════════════════════════════════
@registry.envcfg("go1-dreamwaq-walk")
@dataclass
class DreamWaQCfg(Go1WalkNpEnvCfg):
"""DreamWaQ 配置 — 10×20 地形网格，5 种类型，10 个难度级别。"""
model_file: str = field(default_factory=_scene_file)
# CENet 参数
num_latent: int = 16
cenet_out_dim: int = 19 # code = vel(3) + latent(16)
num_history: int = 5
# 特权观测维度：obs(45) + base_vel(3) + 足部接触力(12) + heights(187) = 247
num_privileged_obs: int = 247
max_episode_steps: int = 1000 # 与上游一致: 20s / 0.02s ctrl_dt
sim_dt: float = 0.005
ctrl_dt: float = 0.02
# 高度测量网格（与上游一致: 17×11）
height_points_x: tuple = tuple(np.linspace(-0.8, 0.8, 17).tolist())
height_points_y: tuple = tuple(np.linspace(-0.5, 0.5, 11).tolist())
# 地形网格
terrain_rows: int = 10
terrain_cols: int = 20 # 上游 10×20 网格：5 类型 × 10 难度
cell_size: float = 8.0
border_size: float = 5.0
# 命令范围（上游: [-1,1] 对称）
@dataclass
class Commands:
vel_limit = [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]
commands: Commands = field(default_factory=Commands)
def __init__(self):
super().__init__()
self.model_file = _scene_file()
self.num_latent = 16
self.cenet_out_dim = 19
self.num_history = 5
self.num_privileged_obs = 247
self.sim_dt = 0.005
self.ctrl_dt = 0.02
if "flat_stairs" in self.model_file:
self.terrain_rows = 2
else:
self.terrain_rows = 10
self.terrain_cols = 20
self.cell_size = 8.0
self.border_size = 5.0
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.control_config.stiffness = 28.0
self.control_config.damping = 0.7
self._apply_dreamwaq()
def _apply_dreamwaq(self):
"""DreamWaQ 奖励尺度 — 与上游 Go1RoughCfg 一致（上游会在 _prepare 中 × dt）。"""
r = self.reward_config.scales
r.clear()
r.update({
"tracking_lin_vel": 1.0,
"tracking_ang_vel": 0.5,
"lin_vel_z": -2.0,
"ang_vel_xy": -0.05,
"orientation": -0.2,
"dof_acc": -2.5e-7,
"base_height": -1.0,
"feet_air_time": 0.1,
"action_rate": -0.01,
"joint_power": -2e-5,
"smoothness": -0.01,
"power_distribution": -10e-6,
# 注意：stand_still 在上游被注释掉
})
self.reward_config.only_positive_rewards = True
self.reward_config.tracking_sigma = 0.25
# ═══════════════════════════════════════════════════════════════════════
# DreamWaQ 环境
# ═══════════════════════════════════════════════════════════════════════
@registry.env("go1-dreamwaq-walk", sim_backend="np")
class DreamWaQTask(Go1WalkTask):
"""DreamWaQ 环境：不对称观测、历史缓冲区、高度测量。"""
_cfg: DreamWaQCfg
def __init__(self, cfg: DreamWaQCfg = None, num_envs=1):
if cfg is None:
cfg = DreamWaQCfg()
super().__init__(cfg, num_envs)
# 覆盖观测空间（DreamWaQ 45-dim，比基类少 linvel 3-dim）
self._observation_space = gym.spaces.Box(
low=-np.inf, high=np.inf, shape=(45,), dtype=np.float32)
self._num_observation = 45
self._hx = np.array(cfg.height_points_x, dtype=np.float32)
self._hy = np.array(cfg.height_points_y, dtype=np.float32)
self._num_rows = cfg.terrain_rows
self._num_cols = cfg.terrain_cols
self._cell_size = cfg.cell_size
self._border = cfg.border_size
self._init_done = False
self._all_terrain_types = np.tile(
np.arange(self._num_cols), num_envs // self._num_cols + 1)[:num_envs]
np.random.shuffle(self._all_terrain_types)
self._hf_cache = None
self._hm_cache = None
# ── 命令：全范围，无课程（匹配上游）──
def resample_commands(self, num_envs: int) -> np.ndarray:
lim = np.array(self.cfg.commands.vel_limit, dtype=np.float32)
cmds = np.random.uniform(lim[0], lim[1], size=(num_envs, 3)).astype(np.float32)
small = np.linalg.norm(cmds[:, :2], axis=1) < 0.2
cmds[small, :2] = 0.0
return cmds
# ── 动作裁剪 + 力矩计算 ──
def apply_action(self, actions, state):
"""裁剪动作防止奖励计算溢出（上游 clip_actions=100）。"""
actions = np.clip(actions, -100.0, 100.0)
return super().apply_action(actions, state)
def _compute_torques(self, actions, data):
"""PD 控制器 + 域随机化。力矩裁剪防溢出。"""
state = getattr(self, '_state', None)
if state is not None:
motor_strength = state.info.get("motor_strength",
np.ones(self._num_envs, dtype=np.float32))
kp_factor = state.info.get("kp_factor",
np.ones(self._num_envs, dtype=np.float32))
kd_factor = state.info.get("kd_factor",
np.ones(self._num_envs, dtype=np.float32))
else:
motor_strength = np.ones(1, dtype=np.float32)
kp_factor = np.ones(1, dtype=np.float32)
kd_factor = np.ones(1, dtype=np.float32)
actions_scaled = actions * self.cfg.control_config.action_scale * motor_strength[:, np.newaxis]
torques = (self.kps * kp_factor[:, np.newaxis]) * (
actions_scaled + self.default_angles - self.get_dof_pos(data)
) - (self.kds * kd_factor[:, np.newaxis]) * self.get_dof_vel(data)
return np.clip(torques, -80.0, 80.0)
# ── 观测 ──
def _get_obs(self, data: mtx.SceneData, info: dict) -> np.ndarray:
"""45-dim: ang_vel + gravity + commands + joint_pos + joint_vel + actions。"""
gyro = self.get_gyro(data)
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
gravity = quaternion.rotate_inverse(base_quat, self.gravity_vec)
commands = info["commands"] * self.commands_scale
joint_pos = (self.get_dof_pos(data) - self.default_angles) * self.cfg.normalization.dof_pos
joint_vel = self.get_dof_vel(data) * self.cfg.normalization.dof_vel
noisy_gyro = gyro * self.cfg.normalization.ang_vel
actions = info["current_actions"]
obs = np.concatenate([
noisy_gyro, gravity, commands, joint_pos, joint_vel, actions
], axis=-1)
# 观测噪声
nc = self.cfg.noise_config
noise_scale_vec = np.zeros(45, dtype=np.float32)
noise_scale_vec[0:3] = nc.scale_gyro
noise_scale_vec[3:6] = nc.scale_gravity
noise_scale_vec[6:9] = 0.0
noise_scale_vec[9:21] = nc.scale_joint_angle
noise_scale_vec[21:33] = nc.scale_joint_vel
noise_scale_vec[33:45] = 0.0
noise = (2.0 * np.random.rand(*obs.shape).astype(np.float32) - 1.0) * nc.level
obs = obs + noise * noise_scale_vec[np.newaxis, :]
return obs
# ── 高度测量 ──
def _get_heights(self, data: mtx.SceneData) -> np.ndarray:
n = data.shape[0]
nx, ny = len(self._hx), len(self._hy)
pose = self._body.get_pose(data)
base_pos = pose[:, 0:3]
yaw = quaternion.get_yaw(pose[:, 3:7])
cos_yaw, sin_yaw = np.cos(yaw), np.sin(yaw)
try:
if self._hm_cache is None:
hf = self._model.get_hfield(0)
self._hf_cache = hf
hm = hf.height_matrix
self._hm_cache = hm
nr, nc = hm.shape
b = hf.bound
self._h_xmin, self._h_ymin = b[0], b[1]
self._h_xmax, self._h_ymax = b[3], b[4]
self._h_nr, self._h_nc = nr, nc
cfg_base = getattr(self.cfg, "hfield_z_base", None)
self._h_zbase = float(cfg_base) if cfg_base is not None else float(-hm[0, 0])
hm, nr, nc = self._hm_cache, self._h_nr, self._h_nc
xmin, ymin, xmax, ymax = self._h_xmin, self._h_ymin, self._h_xmax, self._h_ymax
z_base = self._h_zbase
w, h = xmax - xmin, ymax - ymin
except Exception:
return np.zeros((n, nx * ny), dtype=np.float32)
heights = np.zeros((n, nx * ny), dtype=np.float32)
idx = 0
for iy in range(ny):
gy = self._hy[iy]
for ix in range(nx):
gx = self._hx[ix]
wx = base_pos[:, 0] + cos_yaw * gx - sin_yaw * gy
wy = base_pos[:, 1] + sin_yaw * gx + cos_yaw * gy
col_f = (wx - xmin) / max(w, 1e-6) * (nc - 1)
row_f = (ymax - wy) / max(h, 1e-6) * (nr - 1)
col = np.clip(np.nan_to_num(col_f, nan=0).astype(np.int32), 0, nc - 1)
row = np.clip(np.nan_to_num(row_f, nan=0).astype(np.int32), 0, nr - 1)
heights[:, idx] = hm[row, col] + z_base
idx += 1
return heights
def _sample_terrain_height(self, xy: np.ndarray, radius: float = 0.0) -> np.ndarray:
n = xy.shape[0]
try:
hm, nr, nc = self._hm_cache, self._h_nr, self._h_nc
xmin, ymin, xmax, ymax = self._h_xmin, self._h_ymin, self._h_xmax, self._h_ymax
z_base = self._h_zbase
w, h = xmax - xmin, ymax - ymin
except Exception:
return np.zeros(n, dtype=np.float32)
col = np.clip(np.nan_to_num((xy[:, 0] - xmin) / max(w, 1e-6) * (nc - 1), nan=0).astype(np.int32), 0, nc - 1)
row = np.clip(np.nan_to_num((ymax - xy[:, 1]) / max(h, 1e-6) * (nr - 1), nan=0).astype(np.int32), 0, nr - 1)
if radius <= 0:
return (hm[row, col] + z_base).astype(np.float32)
dc = max(1, int(radius / max(w, 1e-6) * (nc - 1)))
dr = max(1, int(radius / max(h, 1e-6) * (nr - 1)))
out = np.empty(n, dtype=np.float32)
for i in range(n):
r, c = int(row[i]), int(col[i])
out[i] = hm[max(0, r - dr):r + dr + 1, max(0, c - dc):c + dc + 1].max()
return out + z_base
# ── 足部接触力 ──
def _read_contact_forces(self, data: mtx.SceneData) -> np.ndarray:
"""读取四足接触力（body frame）→ (N, 12)。"""
pose = self._body.get_pose(data)
base_quat = pose[:, 3:7]
forces = []
for foot in ["FR", "FL", "RR", "RL"]:
v = self._model.get_sensor_value(foot + "_foot_contact", data)
v_body = quaternion.rotate_inverse(base_quat, v)
forces.append(v_body)
return np.concatenate(forces, axis=1)
def _get_privileged_obs(self, data: mtx.SceneData, obs: np.ndarray) -> np.ndarray:
base_vel = self.get_local_linvel(data) * self.cfg.normalization.lin_vel
heights = self._get_heights(data)
contact_forces = self._read_contact_forces(data)
return np.concatenate([obs, base_vel, contact_forces, heights], axis=-1)
# ── update_observation ──
def update_observation(self, state):
data = state.data
obs = self._get_obs(data, state.info)
# 更新历史缓冲区
old_history = state.info.get("obs_history",
np.zeros((self._num_envs, self._cfg.num_history, 45), dtype=np.float32))
new_history = np.concatenate([old_history[:, 1:, :], obs[:, np.newaxis, :]], axis=1)
state.info["obs_history"] = new_history
current_step = state.info.get("steps", np.zeros(self._num_envs, dtype=np.int32))
# 中期 command 重采样（每 10 秒）
resample_steps = int(10.0 / self.cfg.ctrl_dt)
do_resample = (current_step % resample_steps) == 0
if do_resample.any():
n_resample = int(do_resample.sum())
state.info["commands"][do_resample] = self.resample_commands(n_resample)
# 周期性域随机化（每 4s）
rand_interval = int(4.0 / self.cfg.ctrl_dt)
last_rand = state.info.get("last_rand_step", np.zeros(self._num_envs, dtype=np.int32))
do_rand = (current_step - last_rand) >= rand_interval
if do_rand.any():
n_rand = int(do_rand.sum())
state.info["motor_strength"][do_rand] = np.random.uniform(0.9, 1.1, size=n_rand).astype(np.float32)
state.info["kp_factor"][do_rand] = np.random.uniform(0.9, 1.1, size=n_rand).astype(np.float32)
state.info["kd_factor"][do_rand] = np.random.uniform(0.9, 1.1, size=n_rand).astype(np.float32)
state.info["friction_coeff"][do_rand] = np.random.uniform(0.2, 1.25, size=n_rand).astype(np.float32)
state.info["added_mass"][do_rand] = np.random.uniform(-1.0, 2.0, size=n_rand).astype(np.float32)
state.info["com_displacement"][do_rand] = np.random.uniform(-0.05, 0.05, size=(n_rand, 3)).astype(np.float32)
state.info["last_rand_step"][do_rand] = current_step[do_rand]
# 特权观测
state.info["privileged_obs"] = self._get_privileged_obs(data, obs)
state.info["base_vel"] = self.get_local_linvel(data)
# 足部接触
cquerys = self._model.get_contact_query(data)
state.info["contacts"] = cquerys.is_colliding(self.foot_check).reshape(
(self._num_envs, self.foot_check_num))
state.info["feet_air_time"] = self.update_feet_air_time(state.info)
# 累计命令距离和跟踪
cmd_speed = np.linalg.norm(state.info["commands"][:, :2], axis=1)
state.info["ep_cmd_distance"] = state.info.get("ep_cmd_distance",
np.zeros(self._num_envs, dtype=np.float32)) + cmd_speed * self.cfg.ctrl_dt
state.info["ep_steps"] = state.info.get("ep_steps", np.zeros(self._num_envs, dtype=np.int32)) + 1
local_vel = self.get_local_linvel(data)[:, :2]
cmd_vel = state.info["commands"][:, :2]
vel_error = np.sum(np.square(cmd_vel - local_vel), axis=1)
tracking = np.exp(-vel_error / self.cfg.reward_config.tracking_sigma)
state.info["ep_tracking_sum"] = state.info.get("ep_tracking_sum",
np.zeros(self._num_envs, dtype=np.float32)) + tracking
return state.replace(obs=obs)
# ── 出生点 ──
def _make_origins(self, levels, indices):
n = len(levels)
half_x = self._border + self._num_cols * self._cell_size / 2.0
half_y = self._border + self._num_rows * self._cell_size / 2.0
origins = np.zeros((n, 2), dtype=np.float32)
for i in range(n):
row = levels[i]
col = self._all_terrain_types[int(indices[i]) % len(self._all_terrain_types)]
cx = -half_x + self._border + col * self._cell_size + self._cell_size / 2
cy = half_y - self._border - row * self._cell_size - self._cell_size / 2
origins[i, 0] = cx
origins[i, 1] = cy
return origins
# ── Reset ──
def reset(self, data) -> tuple[np.ndarray, dict]:
num_reset = data.shape[0]
state = getattr(self, '_state', None)
if state is not None and hasattr(state, 'done'):
done = state.done
done_idx = np.where(done)[0]
else:
done_idx = np.arange(num_reset)
# 游戏启发式地形课程（per-env，与上游一致）
if not hasattr(self, '_terrain_origins'):
all_levels = np.repeat(np.arange(self._num_rows), self._num_cols)
all_indices = np.tile(np.arange(self._num_cols), self._num_rows)
all_origins = self._make_origins(all_levels, all_indices)
self._terrain_origins = all_origins.reshape(self._num_rows, self._num_cols, 2)
self._max_init_level = 5
if num_reset > 0 and self._init_done and state is not None and hasattr(state, 'info'):
old_info = state.info
old_origins = old_info.get("env_origins", np.zeros((self._num_envs, 2), dtype=np.float32))[done]
old_levels = old_info.get("terrain_level", np.zeros(self._num_envs, dtype=np.int32))[done]
old_commands = old_info.get("commands", np.zeros((self._num_envs, 3), dtype=np.float32))[done]
base_pose = self._body.get_pose(state.data)
base_pos = base_pose[done, :2]
distance = np.linalg.norm(base_pos - old_origins, axis=1)
move_up = distance > (self._cell_size / 2.0)
cmd_speed = np.linalg.norm(old_commands[:, :2], axis=1)
required_dist = cmd_speed * (self.cfg.max_episode_steps * self.cfg.ctrl_dt) * 0.5
move_down = (distance < required_dist) & ~move_up
new_levels = np.where(move_up, old_levels + 1, old_levels)
new_levels = np.where(move_down, new_levels - 1, new_levels)
at_max = new_levels >= self._num_rows
if at_max.any():
new_levels[at_max] = np.random.randint(0, self._num_rows, size=int(at_max.sum()))
new_levels = np.clip(new_levels, 0, self._num_rows - 1)
else:
new_levels = np.random.randint(0, self._max_init_level + 1, size=num_reset, dtype=np.int32)
self._init_done = True
_force = getattr(self, "_force_level", None)
if _force is not None:
new_levels = np.full(num_reset, int(_force), dtype=np.int32)
if "scene_stairs_terrain" in self.cfg.model_file:
height_list = np.array([-1.0, 0.5, 1.5], dtype=np.float32)
offset_h = [[2, 0, 2, 1, 1], [2, 2, 1, 0, 0], [1, 1, 2, 1, 2],
[0, 1, 0, 2, 0], [0, 1, 1, 0, 2]]
offsets = []
for i in range(5):
for j in range(5):
offsets.append([(i-2)*8.0, (j-2)*8.0, height_list[offset_h[j][i]]])
offset_arr = np.array(offsets, dtype=np.float32)
idx = np.random.choice(len(offsets), size=num_reset)
new_origins = offset_arr[idx, :2]
self._stairs_z_offsets = offset_arr[idx, 2]
else:
indices = done_idx
if hasattr(self, '_terrain_origins'):
new_origins = np.zeros((num_reset, 2), dtype=np.float32)
for i in range(num_reset):
row = int(new_levels[i])
col = self._all_terrain_types[int(indices[i]) % len(self._all_terrain_types)]
new_origins[i] = self._terrain_origins[row, col]
else:
new_origins = self._make_origins(new_levels, indices)
self._stairs_z_offsets = np.zeros(num_reset, dtype=np.float32)
data.reset(self._model)
init_dof_pos = np.tile(self._init_dof_pos, (num_reset, 1))
init_dof_pos[:, 0] = new_origins[:, 0]
init_dof_pos[:, 1] = new_origins[:, 1]
if hasattr(self, '_stairs_z_offsets') and self._stairs_z_offsets.any():
init_dof_pos[:, 2] = self._stairs_z_offsets
else:
terrain_z = self._sample_terrain_height(new_origins, radius=0.35)
spawn_abs = getattr(self, "_spawn_absolute", None)
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)
yaw = np.random.uniform(-np.pi, np.pi, size=num_reset)
init_dof_pos[:, 3] = 0.0
init_dof_pos[:, 4] = 0.0
init_dof_pos[:, 5] = np.sin(yaw / 2)
init_dof_pos[:, 6] = np.cos(yaw / 2)
init_dof_vel = np.tile(self._init_dof_vel, (num_reset, 1))
data.set_dof_vel(init_dof_vel)
data.set_dof_pos(init_dof_pos, self._model)
self._model.forward_kinematic(data)
info = {
"current_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
"last_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
"last_last_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
"commands": self.resample_commands(num_reset),
"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),
"motor_strength": np.random.uniform(0.9, 1.1, size=num_reset).astype(np.float32),
"kp_factor": np.random.uniform(0.9, 1.1, size=num_reset).astype(np.float32),
"kd_factor": np.random.uniform(0.9, 1.1, size=num_reset).astype(np.float32),
"friction_coeff": np.random.uniform(0.2, 1.25, size=num_reset).astype(np.float32),
"added_mass": np.random.uniform(-1.0, 2.0, size=num_reset).astype(np.float32),
"com_displacement": np.random.uniform(-0.05, 0.05, size=(num_reset, 3)).astype(np.float32),
"last_rand_step": np.zeros(num_reset, dtype=np.int32),
"ep_cmd_distance": np.zeros(num_reset, dtype=np.float32),
"ep_steps": np.zeros(num_reset, dtype=np.int32),
"ep_tracking_sum": np.zeros(num_reset, dtype=np.float32),
}
obs = self._get_obs(data, info)
hist = np.zeros((num_reset, self._cfg.num_history, 45), dtype=np.float32)
hist[:, -1, :] = obs
info["obs_history"] = hist
info["privileged_obs"] = self._get_privileged_obs(data, obs)
info["base_vel"] = self.get_local_linvel(data)
info["terrain_level"] = new_levels
info["env_origins"] = new_origins
# 清除 ep_ 奖励累计器（跳过 ep_report 它是 dict）
if state is not None:
for k in list(state.info.keys()):
if k.startswith("ep_") and k != "ep_report" and k not in info:
info[k] = np.zeros(num_reset, dtype=np.float32)
return obs, info
# ── 奖励（与上游对齐，update_reward 中 × dt）──
def _get_reward(self, data: mtx.SceneData, info: dict) -> dict[str, np.ndarray]:
"""DreamWaQ 奖励项——与上游 Manaro-Alpha 对齐。"""
commands = info["commands"]
return {
"tracking_lin_vel": self._reward_tracking_lin_vel(data, commands),
"tracking_ang_vel": self._reward_tracking_ang_vel(data, commands),
"lin_vel_z": self._reward_lin_vel_z(data),
"ang_vel_xy": self._reward_ang_vel_xy(data),
"orientation": self._reward_orientation(data),
"dof_acc": self._reward_dof_acc(data, info),
"base_height": self._reward_base_height(data),
"feet_air_time": self._reward_feet_air_time(commands, info),
"action_rate": self._reward_action_rate(info),
"joint_power": self._reward_joint_power(data),
"smoothness": self._reward_smoothness(info),
"power_distribution": self._reward_power_distribution(data),
# stand_still 在上游被注释掉
}
def _reward_feet_air_time(self, commands, info):
"""足部腾空时间奖励——与上游 legged_robot.py 公式一致。"""
first_contact = info.get("first_contact")
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((air_time - 0.5) * first_contact, axis=1)
rew *= np.linalg.norm(commands[:, :2], axis=1) > 0.1
return rew
def update_reward(self, state):
"""存储各项奖励到 TensorBoard + 累计 episode 总和。"""
reward_dict = self._get_reward(state.data, state.info)
# 乘系数 + dt（与上游 _prepare_reward_function 对齐）
scales = self._cfg.reward_config.scales
dt = self._cfg.ctrl_dt
scaled_terms = {
k: v * scales.get(k, 0.0) * dt
for k, v in reward_dict.items()
}
state.info["reward_terms"] = {k: float(np.mean(v)) for k, v in scaled_terms.items()}
# 累计 episode 总和
for k, v in scaled_terms.items():
ek = f"ep_{k}"
state.info[ek] = state.info.get(ek, np.zeros(self._num_envs, dtype=np.float32)) + v
# 对即将结束的 env 存 ep_report
steps = state.info.get("steps", np.zeros(self._num_envs, dtype=np.int32))
will_end = state.terminated | (steps >= self._cfg.max_episode_steps - 1)
if will_end.any():
ep_report = state.info.get("ep_report", {})
for k in scaled_terms:
ek = f"ep_{k}"
vals = state.info.get(ek, np.zeros(self._num_envs, dtype=np.float32))[will_end]
ep_report[f"rew_{k}"] = float(np.mean(vals))
tl = state.info.get("terrain_level", np.zeros(self._num_envs, dtype=np.int32))[will_end]
ep_report["terrain_level"] = float(np.mean(tl))
state.info["ep_report"] = ep_report
state = super().update_reward(state)
if self._cfg.reward_config.only_positive_rewards:
state = state.replace(reward=np.maximum(state.reward, 0.0))
return state
# ── 额外奖励函数 ──
def _reward_base_height(self, data):
"""惩罚偏离目标基础高度——与上游 Go1RoughCfg.base_height_target=0.30 一致。"""
pose = self._body.get_pose(data)
base_z = pose[:, 2]
heights = self._get_heights(data)
ground_level = np.mean(heights, axis=1)
target = 0.30
return np.square(base_z - ground_level - target)
def _reward_joint_power(self, data):
torque = np.clip(data.actuator_ctrls, -100, 100)
vel = np.clip(self.get_dof_vel(data), -100, 100)
return np.sum(np.abs(torque * vel), axis=1)
def _reward_smoothness(self, info):
scale = self.cfg.control_config.action_scale
da = self.default_angles
a0 = info["current_actions"] * scale + da
a1 = info["last_actions"] * scale + da
a2 = info.get("last_last_actions", info["last_actions"]) * scale + da
diff = np.square(np.clip(a0 - 2.0 * a1 + a2, -1e4, 1e4))
diff = diff * (info["last_actions"] != 0)
diff = diff * (info.get("last_last_actions", info["last_actions"]) != 0)
return np.sum(diff, axis=1)
def _reward_power_distribution(self, data):
torque = np.clip(data.actuator_ctrls, -100, 100)
vel = np.clip(self.get_dof_vel(data), -100, 100)
power = torque * vel
return np.var(np.abs(power), axis=1)

3
motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/dreamwaq_terrain.png Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:5fe0435e385f736b46a910b53d30135d7f6280c0d8daa2bd2f64b7df7962d998
size 1439313

25
motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_flat.xml Normal file
View File

@@ -0,0 +1,25 @@
<mujoco model="go1 dreamwaq flat scene">
<include file="go1_motor_actuator.xml" />
<include file="materials.xml" />
<statistic center="0 0 0.1" extent="0.8" meansize="0.04" />
<visual>
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
<rgba haze="0.15 0.25 0.35 1" />
<global azimuth="120" elevation="-20" />
<map force="0.01" />
<scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2" />
<quality shadowsize="8192" />
</visual>
<worldbody>
<light pos="0 0 1.5" dir="0 0 -1" directional="true" />
<geom name="floor" pos="0 0 0" size="0 0 0.01" type="plane"
material="motphys-ground" contype="1" conaffinity="0" priority="1"
friction="0.6" condim="3" />
</worldbody>
<sensor>
<contact name="FR_foot_contact" geom2="FR_foot" geom1="floor" data="force" num="1" />
<contact name="FL_foot_contact" geom2="FL_foot" geom1="floor" data="force" num="1" />
<contact name="RR_foot_contact" geom2="RR_foot" geom1="floor" data="force" num="1" />
<contact name="RL_foot_contact" geom2="RL_foot" geom1="floor" data="force" num="1" />
</sensor>
</mujoco>

34
motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_terrain.xml Normal file
View File

@@ -0,0 +1,34 @@
<mujoco model="go1 dreamwaq terrain scene">
<include file="go1_motor_actuator.xml" />
<include file="materials.xml" />
<statistic center="0 0 0.2" extent="5" meansize="0.04" />
<visual>
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
<rgba haze="0.15 0.25 0.35 1" />
<global azimuth="120" elevation="-20" />
<map force="0.01" />
<scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2" />
<quality shadowsize="8192" />
</visual>
<asset>
<hfield name="dreamwaq_terrain"
file="assets/dreamwaq_terrain.png"
size="85.0 45.0 1.680 0.001" />
</asset>
<worldbody>
<light pos="0 0 4" dir="0 0 -1" directional="true" />
<geom name="floor" pos="0 0 0" type="hfield" hfield="dreamwaq_terrain"
material="motphys-ground" contype="1" conaffinity="0"
priority="1" friction="0.6" />
</worldbody>
<sensor>
<contact name="FR_foot_contact" geom2="FR_foot" geom1="floor" data="force" num="1" />
<contact name="FL_foot_contact" geom2="FL_foot" geom1="floor" data="force" num="1" />
<contact name="RR_foot_contact" geom2="RR_foot" geom1="floor" data="force" num="1" />
<contact name="RL_foot_contact" geom2="RL_foot" geom1="floor" data="force" num="1" />
</sensor>
</mujoco>