diff --git a/exports_go1_dreamwaq/policy.onnx b/exports_go1_dreamwaq/policy.onnx index ccecef0..771f2a1 100644 Binary files a/exports_go1_dreamwaq/policy.onnx and b/exports_go1_dreamwaq/policy.onnx differ diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py b/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py index c1da422..4c4a6b2 100644 --- a/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py +++ b/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py @@ -121,18 +121,20 @@ class DreamWaQCfg(Go1WalkNpEnvCfg): "tracking_lin_vel": 1.5, "tracking_ang_vel": 1.0, "lin_vel_z": -2.0, - "ang_vel_xy": -0.05, - "orientation": -0.2, + "ang_vel_xy": -0.10, # 抑制晃动 + "orientation": -0.5, # 强制平稳姿态 "dof_acc": -2.5e-7, - "base_height": -1.0, + "base_height": -10.0, # 强惩罚身高偏差，避免蹲伏或踮脚 "feet_air_time": 0.1, "action_rate": -0.01, "joint_power": -2e-5, - "smoothness": -0.01, + "smoothness": -0.02, "power_distribution": -10e-6, - # 注意：stand_still 在上游被注释掉 + "stand_still": -0.5, + "dof_pos_limits": -5.0, # 关节限位软约束（参考 M20 修改版） + "collision": -1.0, # 惩罚身体碰撞 }) - self.reward_config.only_positive_rewards = True + self.reward_config.only_positive_rewards = False # 让坏行为负反馈直达策略 self.reward_config.tracking_sigma = 0.25 @@ -181,12 +183,25 @@ class DreamWaQTask(Go1WalkTask): # ── 动作裁剪 + 力矩计算 ── def apply_action(self, actions, state): - """裁剪动作防止奖励计算溢出（上游 clip_actions=100）。""" - actions = np.clip(actions, -100.0, 100.0) - return super().apply_action(actions, state) + """裁剪动作 + 随机延迟（模拟真实部署延迟，提高 sim-to-real 泛化）。""" + actions = np.clip(actions, -4.0, 4.0) + # 随机动作延迟：0-3 个控制步（0-60ms） + if not hasattr(self, '_action_buffer'): + self._action_buffer = np.zeros((self._num_envs, 3, self._num_action), dtype=np.float32) + self._latency_steps = np.random.randint(0, 4, size=self._num_envs) + # 滑动缓冲区，新动作插入末尾 + self._action_buffer = np.concatenate([ + self._action_buffer[:, 1:, :], actions[:, np.newaxis, :] + ], axis=1) + # 根据每个 env 的延迟取对应历史动作 + delayed_actions = np.zeros_like(actions) + for i in range(self._num_envs): + lat = self._latency_steps[i] + delayed_actions[i] = self._action_buffer[i, 2 - lat] # buffer[-1] 是最新 + return super().apply_action(delayed_actions, state) def _compute_torques(self, actions, data): - """PD 控制器 + 域随机化。力矩裁剪防溢出。""" + """PD 控制器 + 域随机化 + 关节限位。力矩裁剪防溢出。""" state = getattr(self, '_state', None) if state is not None: motor_strength = state.info.get("motor_strength", @@ -201,8 +216,13 @@ class DreamWaQTask(Go1WalkTask): kd_factor = np.ones(1, dtype=np.float32) actions_scaled = actions * self.cfg.control_config.action_scale * motor_strength[:, np.newaxis] + # 目标角度限制在关节范围内（与 MuJoCo 硬限位对齐） + target = actions_scaled + self.default_angles + lo = self._model.joint_limits[0][np.newaxis, :] # (1, 12) + hi = self._model.joint_limits[1][np.newaxis, :] + target = np.clip(target, lo, hi) torques = (self.kps * kp_factor[:, np.newaxis]) * ( - actions_scaled + self.default_angles - self.get_dof_pos(data) + target - self.get_dof_pos(data) ) - (self.kds * kd_factor[:, np.newaxis]) * self.get_dof_vel(data) return np.clip(torques, -80.0, 80.0) @@ -326,14 +346,31 @@ class DreamWaQTask(Go1WalkTask): def update_observation(self, state): data = state.data - # 清理物理崩溃残留的 NaN/Inf（仅在检测到时才修改，避免无谓开销） + # 清理物理崩溃残留的 NaN/Inf dv = data.dof_vel if np.any(~np.isfinite(dv)): - data.set_dof_vel(np.nan_to_num(np.array(dv), nan=0.0, posinf=0.0, neginf=0.0)) + dv_clean = np.array(dv) + dv_clean[:, 6:] = np.nan_to_num(dv_clean[:, 6:], nan=0.0, posinf=0.0, neginf=0.0) + data.set_dof_vel(dv_clean) dp = data.dof_pos if np.any(~np.isfinite(dp)): - data.set_dof_pos(np.nan_to_num(np.array(dp), nan=0.0, posinf=0.0, neginf=0.0), self._model) + dp_clean = np.array(dp) + # 四元数 NaN → 整行替换为单位四元数 [0,0,0,1] + quat_nan_row = np.any(~np.isfinite(dp_clean[:, 3:7]), axis=1) + if np.any(quat_nan_row): + dp_clean[quat_nan_row, 3:7] = [0.0, 0.0, 0.0, 1.0] + # 关节位置 NaN → 0 + dp_clean[:, 7:] = np.nan_to_num(dp_clean[:, 7:], nan=0.0, posinf=0.0, neginf=0.0) + data.set_dof_pos(dp_clean, self._model) + else: + quat_norm = np.linalg.norm(dp[:, 3:7], axis=1) + bad_quat = quat_norm < 1e-6 + if np.any(bad_quat): + dp_clean = np.array(dp) + dp_clean[bad_quat, 3:7] = [0.0, 0.0, 0.0, 1.0] + data.set_dof_pos(dp_clean, self._model) obs = self._get_obs(data, state.info) + obs = np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0) # 更新历史缓冲区：CENet 用不含当前帧的历史预测下一帧（论文 LVAE = MSE(õ_{t+1}, o_{t+1})） full_hist = state.info.get("obs_history_full", @@ -360,8 +397,10 @@ class DreamWaQTask(Go1WalkTask): 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["kp_factor"][do_rand] = np.random.uniform(0.85, 1.15, size=n_rand).astype(np.float32) + state.info["kd_factor"][do_rand] = np.random.uniform(0.85, 1.15, size=n_rand).astype(np.float32) + # 动作延迟随机化（每个 rand 周期重新分配延迟步数） + self._latency_steps[do_rand] = np.random.randint(0, 4, size=n_rand) 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) @@ -369,6 +408,8 @@ class DreamWaQTask(Go1WalkTask): # 特权观测 state.info["privileged_obs"] = self._get_privileged_obs(data, obs) + state.info["privileged_obs"] = np.nan_to_num( + state.info["privileged_obs"], nan=0.0, posinf=0.0, neginf=0.0) state.info["base_vel"] = self.get_local_linvel(data) # 足部接触 @@ -429,7 +470,9 @@ class DreamWaQTask(Go1WalkTask): cy = half_y - self._border - row * self._cell_size - self._cell_size / 2 all_origins[row, col] = [cx, cy] self._terrain_origins = all_origins - self._max_init_level = 0 # 上游原值 + # 跟踪课程进度：动态 max_init_level，避免将已学会的机器人拉回平地 + self._max_init_level = 0 + self._level_history = [] if num_reset > 0 and self._init_done and state is not None and hasattr(state, 'info'): old_info = state.info @@ -439,16 +482,117 @@ class DreamWaQTask(Go1WalkTask): 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 > 3.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 + # ── 多条件 Curriculum 升级 ── + # 1) 速度跟踪: exp(-L2²/σ), σ=0.25, >0.5 表示等效恒定 L2 误差约 0.42m/s + ep_tracking = old_info.get("ep_tracking_sum", np.zeros(self._num_envs, dtype=np.float32))[done] + ep_steps = old_info.get("ep_steps", np.zeros(self._num_envs, dtype=np.int32))[done] + mask = ep_steps > 0 + avg_tracking = np.zeros(num_reset, dtype=np.float32) + avg_tracking[mask] = ep_tracking[mask] / ep_steps[mask].astype(np.float32) + # 调试：打印 avg_tracking 分布（单次） + if not hasattr(self, '_dbg_tracking_printed'): + self._dbg_tracking_printed = True + valid = avg_tracking[mask] + if len(valid) > 0: + print(f"[Curric-DBG] avg_tracking: mean={np.mean(valid):.4f} " + f"min={np.min(valid):.4f} p10={np.percentile(valid,10):.4f} " + f"p50={np.percentile(valid,50):.4f} p90={np.percentile(valid,90):.4f} " + f"max={np.max(valid):.4f} >0.5={100*np.mean(valid>0.5):.0f}%") + # 2) 存活率: 必须自然结束(truncated), 而非摔倒终止(terminated) + term = state.terminated[done] if hasattr(state, 'terminated') else np.zeros(num_reset, dtype=bool) + not_fallen = ~term # True=活到 episode 结束 + # 3) 姿态: avg_orient 保留 dt，-0.0005 对应等效倾角约 13 deg + ep_orient = old_info.get("ep_orientation", np.zeros(self._num_envs, dtype=np.float32))[done] + avg_orient = np.zeros(num_reset, dtype=np.float32) + avg_orient[mask] = ep_orient[mask] / ep_steps[mask].astype(np.float32) + # 4) 身高: avg_base_height 保留 dt，-0.001 对应 RMS 高度误差约 7cm + ep_bh = old_info.get("ep_base_height", np.zeros(self._num_envs, dtype=np.float32))[done] + avg_bh = np.zeros(num_reset, dtype=np.float32) + avg_bh[mask] = ep_bh[mask] / ep_steps[mask].astype(np.float32) + # orient: -0.5*dt*sin²(theta); bh: -10*dt*z_err² + move_up_raw = ( + (avg_tracking > 0.5) + & not_fallen + & (avg_orient > -0.0005) + & (avg_bh > -0.001) + ) + # ── 渐进升级：连续通过 3 次 + 冷却期 + 每次只升 1 级 ── + if not hasattr(self, '_consecutive_pass_count'): + self._consecutive_pass_count = np.zeros(self._num_envs, dtype=np.int32) + self._pending_upgrade = np.zeros(self._num_envs, dtype=bool) + self._upgrade_cooldown = np.zeros(self._num_envs, dtype=np.int32) + + # 每 env: 连续通过 +1, 失败重置（仅更新被 reset 的 env） + self._consecutive_pass_count[done_idx] = np.where( + move_up_raw, + self._consecutive_pass_count[done_idx] + 1, 0) + # 冷却期以该 env 自己完成的 episode 为单位，而不是全局 reset 批次数。 + self._upgrade_cooldown[done_idx] = np.maximum(0, self._upgrade_cooldown[done_idx] - 1) + streak_ready = self._consecutive_pass_count[done_idx] >= 3 + cooldown_blocked = streak_ready & (self._upgrade_cooldown[done_idx] > 0) + # 连续 3 次通过 且 不在冷却期 → 升级（取被 reset env 的当前状态） + can_upgrade = streak_ready & (self._upgrade_cooldown[done_idx] == 0) + # 升级后设置 5 episode 冷却期，并从新等级重新累计连续通过次数。 + self._upgrade_cooldown[done_idx] = np.where( + can_upgrade, 5, self._upgrade_cooldown[done_idx]) + self._consecutive_pass_count[done_idx] = np.where( + can_upgrade, 0, self._consecutive_pass_count[done_idx]) + + move_up = can_upgrade + move_down_raw = ( + (~not_fallen) + | (avg_tracking < 0.35) + | (avg_orient <= -0.001) + | (avg_bh <= -0.002) + ) + move_down = move_down_raw & ~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) + # 跟踪 level 分布，动态更新 max_init_level（P90 - 2） + self._level_history.extend(new_levels.tolist()) + if len(self._level_history) > 500: + self._level_history = self._level_history[-500:] + if len(self._level_history) >= 10: + sorted_lv = sorted(self._level_history) + self._max_init_level = max(0, sorted_lv[int(len(sorted_lv) * 0.9)] - 2) + # ── 调试日志: curriculum 各条件通过率 ── + if not hasattr(self, '_curric_log_counter'): + self._curric_log_counter = 0 + self._curric_track_ok = 0; self._curric_survive_ok = 0 + self._curric_orient_ok = 0; self._curric_bh_ok = 0 + self._curric_raw_all_ok = 0; self._curric_streak_ok = 0 + self._curric_upgrade_ok = 0; self._curric_cooldown_block = 0 + self._curric_down_ok = 0; self._curric_down_track = 0 + self._curric_down_stability = 0; self._curric_total = 0 + self._curric_log_counter += 1 + self._curric_track_ok += int((avg_tracking > 0.5).sum()) + self._curric_survive_ok += int(not_fallen.sum()) + self._curric_orient_ok += int((avg_orient > -0.0005).sum()) + self._curric_bh_ok += int((avg_bh > -0.001).sum()) + self._curric_raw_all_ok += int(move_up_raw.sum()) + self._curric_streak_ok += int(streak_ready.sum()) + self._curric_upgrade_ok += int(move_up.sum()) + self._curric_cooldown_block += int(cooldown_blocked.sum()) + self._curric_down_ok += int(move_down.sum()) + self._curric_down_track += int((avg_tracking < 0.35).sum()) + self._curric_down_stability += int((~not_fallen | (avg_orient <= -0.001) | (avg_bh <= -0.002)).sum()) + self._curric_total += num_reset + if self._curric_log_counter % 50 == 0: + def pct(n): return 100*n/max(self._curric_total,1) + print(f"[Curric] n={self._curric_total:5d} " + f"track={pct(self._curric_track_ok):.0f}% " + f"survive={pct(self._curric_survive_ok):.0f}% " + f"orient={pct(self._curric_orient_ok):.0f}% " + f"bh={pct(self._curric_bh_ok):.0f}% " + f"rawALL={pct(self._curric_raw_all_ok):.0f}% " + f"streak3={pct(self._curric_streak_ok):.0f}% " + f"up={pct(self._curric_upgrade_ok):.0f}% " + f"cool={pct(self._curric_cooldown_block):.0f}% " + f"down={pct(self._curric_down_ok):.0f}% " + f"downT={pct(self._curric_down_track):.0f}% " + f"downS={pct(self._curric_down_stability):.0f}% " + f"max_init={self._max_init_level}") else: new_levels = np.random.randint(0, self._max_init_level + 1, size=num_reset, dtype=np.int32) self._init_done = True @@ -545,12 +689,32 @@ class DreamWaQTask(Go1WalkTask): return obs, info + # ── 新增奖励（参考 M20 修改版）── + + def _reward_dof_pos_limits(self, data): + """关节软限位惩罚——接近关节极限时扣分。""" + lo = self._model.joint_limits[0] + hi = self._model.joint_limits[1] + margin = 0.1 # rad 软边界 + pos = self.get_dof_pos(data) + lower = np.clip(lo + margin - pos, 0, None) + upper = np.clip(pos - (hi - margin), 0, None) + return np.sum(lower + upper, axis=1) + + def _reward_collision(self, data): + """身体碰撞惩罚——膝/肩触地扣分。""" + cquerys = self._model.get_contact_query(data) + penal = cquerys.is_colliding(self.termination_check) + return np.any(penal.reshape(self._num_envs, -1), axis=1).astype(np.float32) + def update_terminated(self, state): """基类接触终止 + hfield 边界终止（防止走出地形导致 NaN）。""" state = super().update_terminated(state) + if self._model.num_hfields == 0: + return state # flat 场景无 hfield + hf = self._model.get_hfield(0) pose = self._body.get_pose(state.data) base_xy = pose[:, :2] - hf = self._model.get_hfield(0) b = hf.bound out_of_bounds = ( (base_xy[:, 0] < b[0]) | (base_xy[:, 0] > b[3]) | @@ -577,7 +741,9 @@ class DreamWaQTask(Go1WalkTask): "joint_power": self._reward_joint_power(data), "smoothness": self._reward_smoothness(info), "power_distribution": self._reward_power_distribution(data), - # stand_still 在上游被注释掉 + "stand_still": self._reward_stand_still(data, commands), + "dof_pos_limits": self._reward_dof_pos_limits(data), + "collision": self._reward_collision(data), } def update_feet_air_time(self, info: dict): @@ -606,6 +772,13 @@ class DreamWaQTask(Go1WalkTask): 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) + def update_reward(self, state): """存储各项奖励到 TensorBoard + 更新 state.reward。 diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/xmls/go1_motor_actuator.xml b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/go1_motor_actuator.xml index 55744ca..0cc6797 100644 --- a/motrix_envs/src/motrix_envs/locomotion/go1/xmls/go1_motor_actuator.xml +++ b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/go1_motor_actuator.xml @@ -49,19 +49,10 @@ - - - - - - - + - - - - + @@ -97,20 +88,20 @@ - + - - - + + + - - + + @@ -120,20 +111,20 @@ - + - - - + + + - - + + @@ -143,20 +134,20 @@ - + - + + + - + + @@ -166,20 +157,20 @@ - + - + + + - + + diff --git a/motrix_envs/src/motrix_envs/np/env.py b/motrix_envs/src/motrix_envs/np/env.py index 8f2743c..050d472 100644 --- a/motrix_envs/src/motrix_envs/np/env.py +++ b/motrix_envs/src/motrix_envs/np/env.py @@ -182,17 +182,19 @@ class NpEnv(ABEnv): pass def physics_step(self): + self._physics_crashed_this_step = False # motrixsim.SceneModel.step only supports single step, so we loop try: for _ in range(self._cfg.sim_substeps): self._model.step(self._state.data) - except Exception as e: + except BaseException as e: # Rust panic / 物理崩溃 → 标记终止 + 清除腐蚀数据 n = self._state.data.shape[0] self._state.terminated[:] = True self._state.reward[:] = 0.0 # 强制重置物理状态，防止 NaN 传播到 observations/rewards self._state.data.reset(self._model) + self._physics_crashed_this_step = True if not hasattr(self, '_physics_crash_count'): self._physics_crash_count = 0 self._physics_crash_count += 1 @@ -213,6 +215,9 @@ class NpEnv(ABEnv): self._state = self.apply_action(actions, self._state) assert self._state is not None, "apply_action must return a valid NpEnvState" self.physics_step() + if getattr(self, "_physics_crashed_this_step", False): + self._reset_done_envs() + return self._state self._state = self.update_state(self._state) self._state.info["steps"] += 1 self._update_truncate() diff --git a/motrix_rl/src/motrix_rl/dwaq_rsl/actor_critic_dwaq.py b/motrix_rl/src/motrix_rl/dwaq_rsl/actor_critic_dwaq.py index 23105be..8bbb659 100644 --- a/motrix_rl/src/motrix_rl/dwaq_rsl/actor_critic_dwaq.py +++ b/motrix_rl/src/motrix_rl/dwaq_rsl/actor_critic_dwaq.py @@ -4,6 +4,43 @@ import torch import torch.nn as nn from torch.distributions import Normal + +class RunningStats(nn.Module): + """Running mean/variance tracker for actor observation normalization. + + Matches SKRL RunningStandardScaler behavior. Without it, actor sees + unscaled inputs at inference → garbage actions → robot collapses. + """ + def __init__(self, num_features: int, eps: float = 1e-8, clip: float = 5.0): + super().__init__() + self.eps = eps + self.clip = clip + self.register_buffer("count", torch.zeros(1, dtype=torch.int64)) + self.register_buffer("mean", torch.zeros(num_features)) + self.register_buffer("var", torch.ones(num_features)) + + @torch.no_grad() + def update(self, x): + if x.dim() != 2: + return + n = x.shape[0] + batch_mean = x.mean(dim=0) + batch_var = x.var(dim=0, unbiased=False) + batch_count = torch.tensor(n, dtype=torch.int64, device=x.device) + delta = batch_mean - self.mean + total = self.count + batch_count + self.mean.add_(delta * batch_count.float() / total.float()) + m_a = self.var * self.count.float() + m_b = batch_var * batch_count.float() + m2 = m_a + m_b + delta.pow(2) * self.count.float() * batch_count.float() / total.float() + self.var.copy_(m2 / total.float()) + self.count.add_(batch_count) + + def forward(self, x): + return torch.clamp( + (x - self.mean) / (self.var.sqrt() + self.eps), -self.clip, self.clip) + + class ActorCritic_DWAQ(nn.Module): def __init__(self, num_actor_obs, num_critic_obs, num_actions, cenet_in_dim, cenet_out_dim, activation="elu", init_noise_std=1.0,): super().__init__() @@ -52,6 +89,7 @@ class ActorCritic_DWAQ(nn.Module): ) self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + self.actor_normalizer = RunningStats(actor_input_dim, eps=1e-8, clip=5.0) self.distribution = None # disable args validation for speedup Normal.set_default_validate_args = False @@ -89,7 +127,7 @@ class ActorCritic_DWAQ(nn.Module): # code = mean_latent + var*code_temp # print("latent : ",code[0]) mean_vel = self.encode_mean_vel(distribution) - logvar_vel = self.encode_mean_vel(distribution) + logvar_vel = self.encode_logvar_vel(distribution) # FIXED: was encode_mean_vel code_latent = self.reparameterise(mean_latent,logvar_latent) code_vel = self.reparameterise(mean_vel,logvar_vel) code = torch.cat((code_vel,code_latent),dim=-1) @@ -117,17 +155,23 @@ class ActorCritic_DWAQ(nn.Module): self.distribution = Normal(mean, mean * 0.0 + self.std) def act(self, observations, obs_history, **kwargs): - code,_,decode,_,_,_,_ = self.cenet_forward(obs_history) - observations = torch.cat((code,observations),dim=-1) + code, _, decode, _, _, _, _ = self.cenet_forward(obs_history) + observations = torch.cat((code, observations), dim=-1) + if self.training: + self.actor_normalizer.update(observations) + if self.actor_normalizer.count > 10: + observations = self.actor_normalizer(observations) self.update_distribution(observations) return self.distribution.sample() def get_actions_log_prob(self, actions): return self.distribution.log_prob(actions).sum(dim=-1) - def act_inference(self, observations,obs_history): - code,_,decode,_,_,_,_ = self.cenet_forward(obs_history) - observations = torch.cat((code,observations),dim=-1) + def act_inference(self, observations, obs_history): + code, _, decode, _, _, _, _ = self.cenet_forward(obs_history) + observations = torch.cat((code, observations), dim=-1) + if self.actor_normalizer.count > 10: + observations = self.actor_normalizer(observations) actions_mean = self.actor(observations) return actions_mean diff --git a/motrix_rl/src/motrix_rl/dwaq_rsl/env_wrapper.py b/motrix_rl/src/motrix_rl/dwaq_rsl/env_wrapper.py index 93968df..4dac6b5 100644 --- a/motrix_rl/src/motrix_rl/dwaq_rsl/env_wrapper.py +++ b/motrix_rl/src/motrix_rl/dwaq_rsl/env_wrapper.py @@ -20,7 +20,7 @@ import torch class DwaqVecEnv: """Wraps a MotrixLab DreamWaQ NpEnv for the upstream DreamWaQ rsl_rl runner.""" - def __init__(self, env, device, num_obs=45, num_privileged_obs=235, + def __init__(self, env, device, num_obs=45, num_privileged_obs=247, num_obs_hist=5, num_actions=12, clip_actions=100.0): self._env = env self.device = device diff --git a/motrix_rl/src/motrix_rl/dwaq_rsl/runner.py b/motrix_rl/src/motrix_rl/dwaq_rsl/runner.py index f97cd14..57b586c 100644 --- a/motrix_rl/src/motrix_rl/dwaq_rsl/runner.py +++ b/motrix_rl/src/motrix_rl/dwaq_rsl/runner.py @@ -158,12 +158,19 @@ class OnPolicyRunner: f"rew={rew:.2f} eplen={elen:.0f}") def save(self, path, infos=None): - torch.save({ + save_dict = { "model_state_dict": self.alg.actor_critic.state_dict(), "optimizer_state_dict": self.alg.optimizer.state_dict(), "iter": self.current_learning_iteration, "infos": infos, - }, path) + } + # Persist RunningStats normalizer state for inference + if hasattr(self.alg.actor_critic, 'actor_normalizer'): + n = self.alg.actor_critic.actor_normalizer + save_dict["normalizer_count"] = n.count + save_dict["normalizer_mean"] = n.mean + save_dict["normalizer_var"] = n.var + torch.save(save_dict, path) def load(self, path, load_optimizer=True): loaded_dict = torch.load(path, map_location=self.device) @@ -171,6 +178,12 @@ class OnPolicyRunner: if load_optimizer: self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) self.current_learning_iteration = loaded_dict["iter"] + # Restore RunningStats normalizer state + if "normalizer_count" in loaded_dict and hasattr(self.alg.actor_critic, 'actor_normalizer'): + n = self.alg.actor_critic.actor_normalizer + n.count.copy_(loaded_dict["normalizer_count"]) + n.mean.copy_(loaded_dict["normalizer_mean"]) + n.var.copy_(loaded_dict["normalizer_var"]) return loaded_dict["infos"] def get_inference_policy(self, device=None): diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py b/motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py index c3bb562..5e8ae78 100644 --- a/motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py +++ b/motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py @@ -14,6 +14,7 @@ import torch import torch.nn as nn +from torch.distributions import Normal from tensordict import TensorDict from rsl_rl.models.mlp_model import MLPModel @@ -113,23 +114,43 @@ class CENetActorModel(MLPModel): # 在 nn.Module.__init__ 之后创建 VAE 子模块 self.vae = CENetVAE(cenet_in_dim, cenet_out_dim, activation) + self.action_clip = 4.0 + self.std_clip = 0.6 self._last_cenet_output = None # AdaBoot: 自适应 bootstrapping（论文 Section II-C） self._adaboot_cv_buffer = [] # 速度估计误差的 CV 历史 self._adaboot_prob = 1.0 # 当前 bootstrap 概率（1.0 = 完全信任 GT） def _update_distribution(self, obs: torch.Tensor) -> None: - """覆盖父类 — 强制 std > 0 再创建 Normal 分布（防止 NaN）。""" - # 先 clamp std，再调父类创建分布 + """覆盖父类 — 限制动作分布，避免采样动作进入 PD/关节限位饱和区。""" + mean = torch.clamp(self.mlp(obs), -self.action_clip, self.action_clip) if self.stochastic and not self.state_dependent_std: with torch.no_grad(): if self.noise_std_type == "scalar": self.std.nan_to_num_(nan=0.5, posinf=1.0, neginf=1.0) - self.std.clamp_(min=1e-6) + self.std.clamp_(min=1e-6, max=self.std_clip) elif self.noise_std_type == "log": self.log_std.nan_to_num_(nan=0.0, posinf=5.0, neginf=-5.0) - self.log_std.clamp_(min=-20.0, max=10.0) - super()._update_distribution(obs) + self.log_std.clamp_(min=-20.0, max=torch.log(torch.tensor(self.std_clip)).item()) + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}.") + else: + std = torch.full_like(mean, self.std_clip) + self.distribution = Normal(mean, std) + + def forward( + self, + obs: TensorDict, + masks: torch.Tensor | None = None, + hidden_state: HiddenState = None, + stochastic_output: bool = False, + ) -> torch.Tensor: + actions = super().forward(obs, masks, hidden_state, stochastic_output) + return torch.clamp(actions, -self.action_clip, self.action_clip) def _get_latent_dim(self) -> int: """Actor 实际输入：code(19) + policy(45) = 64。 @@ -164,10 +185,10 @@ class CENetActorModel(MLPModel): code = torch.cat([code_vel, code[:, 3:]], dim=-1) # 防止 VAE NaN 传播到下游 - if torch.isnan(code).any(): - code = torch.nan_to_num(code, nan=0.0) - if torch.isnan(policy_obs).any(): - policy_obs = torch.nan_to_num(policy_obs, nan=0.0) + if not torch.isfinite(code).all(): + code = torch.nan_to_num(code, nan=0.0, posinf=0.0, neginf=0.0) + if not torch.isfinite(policy_obs).all(): + policy_obs = torch.nan_to_num(policy_obs, nan=0.0, posinf=0.0, neginf=0.0) latent = torch.cat([code, policy_obs], dim=-1) # (N, 64) return latent diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py b/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py index 0a921ef..3c427f7 100644 --- a/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py +++ b/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py @@ -81,8 +81,17 @@ class DreamWaQPPO(PPO): returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, hid_states_batch, masks_batch, ) in generator: - # NaN 检测 — 数据有 NaN 就跳过这个 batch - if torch.isnan(obs_batch["policy"]).any() or torch.isnan(obs_batch["obs_history"]).any(): + # 多层有限值检测 — 数据异常就跳过这个 batch，避免污染网络参数。 + if (not torch.isfinite(obs_batch["policy"]).all() + or not torch.isfinite(obs_batch["obs_history"]).all() + or ("privileged_obs" in obs_batch and not torch.isfinite(obs_batch["privileged_obs"]).all()) + or not torch.isfinite(actions_batch).all() + or not torch.isfinite(target_values_batch).all() + or not torch.isfinite(advantages_batch).all() + or not torch.isfinite(returns_batch).all() + or not torch.isfinite(old_actions_log_prob_batch).all() + or not torch.isfinite(old_mu_batch).all() + or not torch.isfinite(old_sigma_batch).all()): continue # ── 标准 PPO 前向（各调用一次，复用结果）── self.actor(obs_batch, masks=masks_batch, stochastic_output=True) @@ -106,6 +115,11 @@ class DreamWaQPPO(PPO): actions_batch, actions_log_prob_batch, old_actions_log_prob_batch, advantages_batch) + # NaN guard: 损失值异常 → 跳过此 batch，保护模型权重 + if (not torch.isfinite(surrogate_loss) + or not torch.isfinite(value_loss)): + continue + # ── KL 自适应 schedule ── if self.desired_kl is not None and self.schedule == 'adaptive': with torch.inference_mode(): @@ -136,20 +150,41 @@ class DreamWaQPPO(PPO): - self.entropy_coef * entropy_loss + autoenc_loss ) + if not torch.isfinite(loss): + continue # ── 梯度更新 ── self.optimizer.zero_grad() loss.backward() + actor_grads_ok = all( + p.grad is None or torch.isfinite(p.grad).all() + for p in self.actor.parameters() + ) + critic_grads_ok = all( + p.grad is None or torch.isfinite(p.grad).all() + for p in self.critic.parameters() + ) + if not (actor_grads_ok and critic_grads_ok): + self.optimizer.zero_grad(set_to_none=True) + continue if self.max_grad_norm is not None: - nn.utils.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm) - nn.utils.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm) + try: + nn.utils.clip_grad_norm_( + self.actor.parameters(), self.max_grad_norm, error_if_nonfinite=True) + nn.utils.clip_grad_norm_( + self.critic.parameters(), self.max_grad_norm, error_if_nonfinite=True) + except RuntimeError: + self.optimizer.zero_grad(set_to_none=True) + continue nn.utils.clip_grad_value_(self.actor.parameters(), 10.0) self.optimizer.step() # 每次更新后强制 std > 0，防止数值异常导致 NaN if hasattr(self.actor, 'std') and self.actor.stochastic: with torch.no_grad(): - self.actor.std.nan_to_num_(nan=0.5, posinf=1.0, neginf=1.0) - self.actor.std.clamp_(min=1e-6) + if not torch.isfinite(self.actor.std).all(): + self.actor.std.data.fill_(0.5) + std_clip = getattr(self.actor, "std_clip", 0.6) + self.actor.std.data.clamp_(min=1e-6, max=std_clip) # ── 累计日志 ── mean_value_loss += value_loss.item() diff --git a/motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py b/motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py index 092ad38..c5037ce 100644 --- a/motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py +++ b/motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py @@ -4,7 +4,7 @@ CENet (VAE): history(225) → [128,64] → latent(16) + vel_est(3) = code(19) Decoder: code(19) → [64,128] → next_obs(45) Actor: code(19) + obs(45) = 64 → [512,256,128] → action(12) -Critic: privileged_obs(235) → [512,256,128] → value(1) +Critic: privileged_obs(247) → [512,256,128] → value(1) VAE Loss: reconstruction_MSE + velocity_MSE + beta * KL """ @@ -189,7 +189,7 @@ class DreamWaQWrapper: @property def privileged_obs(self): return self._env._state.info.get("privileged_obs", - np.zeros((self._num_envs, 235), dtype=np.float32)) + np.zeros((self._num_envs, 247), dtype=np.float32)) @property def obs_history(self): @@ -398,7 +398,7 @@ class DreamWaQTrainer: def __call__(self, inputs, role): kernel_init = nn.initializers.orthogonal(jnp.sqrt(2)) x = inputs["states"] - # Critic: obs(45) + base_vel(3) + heights(187) = 235 + # Critic: obs(45) + base_vel(3) + heights(187) = 247 # Layout: [code(19) | obs(45) | base_vel(3) | heights(187)] x_c = jnp.concatenate([x[:, 19:64], x[:, 64:254]], axis=-1) for d in value_cfg.hiddens: diff --git a/motrix_rl/src/motrix_rl/tasks/go1.py b/motrix_rl/src/motrix_rl/tasks/go1.py index fea4da5..908c11c 100644 --- a/motrix_rl/src/motrix_rl/tasks/go1.py +++ b/motrix_rl/src/motrix_rl/tasks/go1.py @@ -112,7 +112,7 @@ class rslrl: # Runner 设置（严格对齐上游 LeggedRobotCfgPPO + Go1RoughCfgPPO） runner.seed = 5 # 上游 seed=5 - runner.max_iterations = 3000 # 续训到 3000 轮 + runner.max_iterations = 5000 runner.num_steps_per_env = 24 runner.experiment_name = "go1_dreamwaq_walk" runner.save_interval = 50 diff --git a/scripts/diag_terrain_collision.py b/scripts/diag_terrain_collision.py new file mode 100644 index 0000000..581f09e --- /dev/null +++ b/scripts/diag_terrain_collision.py @@ -0,0 +1,458 @@ +#!/usr/bin/env python3 +"""地形碰撞诊断工具 — 系统性检测 hfield PNG 地形的物理交互质量。 + +检测项： + 1. 脚部穿透（foot z < terrain z） + 2. 接触力异常（为 0 但应该接触 / 过大） + 3. NaN/Inf 传播 + 4. 物理爆炸（关节速度飙升） + 5. 机器人卡住（速度接近 0 但有命令） + 6. 躯干沉入地形（base z < terrain z） + +用法: + uv run scripts/diag_terrain_collision.py # 全量检测 + uv run scripts/diag_terrain_collision.py --level 5 # 仅 level 5 + uv run scripts/diag_terrain_collision.py --type stairs # 仅楼梯地形 + uv run scripts/diag_terrain_collision.py --render # 渲染可视化 +""" +import argparse, os, sys, time +import numpy as np + +os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") +os.environ.setdefault("JAX_PLATFORMS", "cpu") +os.environ.setdefault("DREAMWAQ_TERRAIN", "pyramid") + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import motrix_envs.locomotion.go1.dreamwaq # noqa: F401 +from motrix_envs import registry as env_registry + +# ═══════════════════════════════════════════════════════════════════════ +# 配置 +# ═══════════════════════════════════════════════════════════════════════ + +NUM_ROWS = 10 +NUM_COLS = 20 +CELL_M = 8.0 +BORDER_M = 5.0 +SETTLE_STEPS = 100 # 自由落体稳定步数（200Hz） +TEST_STEPS = 400 # 测试步数 +STEPS_PER_CELL = 5 # 每个 cell 测试的 spawn 位置数 + +# 地形类型名（与 gen_dreamwaq_terrain.py PROPORTIONS 对应） +TYPE_NAMES = ["平滑斜坡", "粗糙斜坡", "下行楼梯", "上行楼梯", "离散障碍"] +CUM = [0.1, 0.2, 0.55, 0.9, 1.0] + + +def _cell_origin(row, col): + half_x = BORDER_M + NUM_COLS * CELL_M / 2.0 + half_y = BORDER_M + NUM_ROWS * CELL_M / 2.0 + cx = -half_x + BORDER_M + col * CELL_M + CELL_M / 2 + cy = half_y - BORDER_M - row * CELL_M - CELL_M / 2 + return cx, cy + + +def _terrain_type(col): + choice = col / NUM_COLS + 0.001 + for i, c in enumerate(CUM): + if choice < c: + return i + return len(TYPE_NAMES) - 1 + + +# ═══════════════════════════════════════════════════════════════════════ +# 诊断核心 +# ═══════════════════════════════════════════════════════════════════════ + +class CollisionDiag: + def __init__(self, env, render=False): + self.env = env + self.render = render + self.renderer = None + if render: + from motrix_envs.np.renderer import NpRenderer + self.renderer = NpRenderer(env) + + # 统计数据 + self.stats = { + "penetrations": [], # (row, col, ttype, foot_z - terrain_z) + "contact_anomalies": [], # (row, col, ttype, max_contact_force) + "nan_events": [], # (row, col, ttype, step) + "explosions": [], # (row, col, ttype, max_dof_vel) + "stuck_events": [], # (row, col, ttype, speed) + "base_sink": [], # (row, col, ttype, base_z - terrain_z) + "ok_count": 0, + "total_tested": 0, + } + + def _get_terrain_z(self, x, y): + """查询 (x, y) 处的地形高度。""" + hm = self.env._hm_cache + if hm is None: + return 0.0 + nr, nc = hm.shape + xmin = self.env._h_xmin + ymax = self.env._h_ymax + w = self.env._h_xmax - xmin + h = ymax - self.env._h_ymin + z_base = self.env._h_zbase + col = int(np.clip((x - xmin) / max(w, 1e-6) * (nc - 1), 0, nc - 1)) + row = int(np.clip((ymax - y) / max(h, 1e-6) * (nr - 1), 0, nr - 1)) + return float(hm[row, col]) + z_base + + def _get_foot_positions(self, state): + """获取所有足部在世界坐标系中的位置 (N, 4, 3)。""" + data = state.data + n = data.shape[0] + feet = np.zeros((n, 4, 3), dtype=np.float32) + foot_names = ["FR_foot", "FL_foot", "RR_foot", "RL_foot"] + for i, name in enumerate(foot_names): + gidx = self.env._model.get_geom_index(name) + if gidx is not None: + pose = self.env._model.geoms[gidx].get_pose(data) + feet[:, i, :] = pose[:, :3] + return feet + + def test_cell(self, row, col, level): + """在单个 cell 上测试多个 spawn 位置。""" + cx, cy = _cell_origin(row, col) + ttype = _terrain_type(col) + tname = TYPE_NAMES[ttype] + + # 在 cell 内取几个 spawn 点（中央 + 四角附近） + offsets = [ + (0, 0), + (CELL_M * 0.3, 0), (-CELL_M * 0.3, 0), + (0, CELL_M * 0.3), (0, -CELL_M * 0.3), + ][:STEPS_PER_CELL] + + for ox, oy in offsets: + sx, sy = cx + ox, cy + oy + terrain_z = float(self.env._sample_terrain_height( + np.array([[sx, sy]], dtype=np.float32), radius=0.35)[0]) + + self.stats["total_tested"] += 1 + self._test_single(sx, sy, terrain_z, row, col, level, ttype, tname) + + def _test_single(self, sx, sy, terrain_z, row, col, level, ttype, tname): + """单个 spawn 点的完整测试。""" + env = self.env + data = env._state.data + n = 1 + + # Spawn + init_pos = env._init_dof_pos.copy().reshape(1, -1) + init_pos[0, 0] = sx + init_pos[0, 1] = sy + spawn_z = terrain_z + 0.45 # 标准 clearance + init_pos[0, 2] = spawn_z + yaw = np.random.uniform(-np.pi, np.pi) + init_pos[0, 3:7] = [0, 0, np.sin(yaw / 2), np.cos(yaw / 2)] + + data.reset(env._model) + data.set_dof_pos(init_pos, env._model) + env._model.forward_kinematic(data) + + # 给一个随机动作防止完全静止（小幅度 PD 控制） + action = 0.1 * (np.random.randn(n, 12).astype(np.float32)) + + has_nan = False + max_dof_vel = 0.0 + max_contact_force = 0.0 + min_foot_z = spawn_z + min_base_z = spawn_z + total_speed = 0.0 + + for step in range(SETTLE_STEPS + TEST_STEPS): + # Apply PD + actions_scaled = action * env.cfg.control_config.action_scale + target = actions_scaled + env.default_angles.reshape(1, -1) + lo = env._model.joint_limits[0].reshape(1, -1) + hi = env._model.joint_limits[1].reshape(1, -1) + target = np.clip(target, lo, hi) + torques = env.kps * (target - env.get_dof_pos(data)) - env.kds * env.get_dof_vel(data) + data.actuator_ctrls = np.clip(torques, -80.0, 80.0) + + env._model.step(data) + + # ── 检测 ── + + # 1. NaN 检测 + dv = np.array(data.dof_vel) + dp = np.array(data.dof_pos) + if np.any(~np.isfinite(dv)) or np.any(~np.isfinite(dp)): + self.stats["nan_events"].append((row, col, ttype, step)) + has_nan = True + break + + # 2. 速度爆炸 + cur_max_vel = float(np.max(np.abs(dv[:, 6:]))) + max_dof_vel = max(max_dof_vel, cur_max_vel) + if cur_max_vel > 100.0: + self.stats["explosions"].append((row, col, ttype, cur_max_vel)) + break + + # 3. 接触力 + # Read contact forces via MotrixSim sensor API + cf_vals = [] + for foot in ["FR", "FL", "RR", "RL"]: + try: + v = env._model.get_sensor_value(f"{foot}_foot_contact", data) + cf_vals.append(v[0]) + except Exception: + cf_vals.append(np.zeros(3, dtype=np.float32)) + cf_raw = np.concatenate(cf_vals) + cur_cf = float(np.max(np.abs(cf_raw))) + max_contact_force = max(max_contact_force, cur_cf) + + # 4. 足部穿透（settle 之后测） + if step >= SETTLE_STEPS: + feet = self._get_foot_positions(env._state) + for fi in range(4): + fz = float(feet[0, fi, 2]) + tz = self._get_terrain_z(float(feet[0, fi, 0]), float(feet[0, fi, 1])) + penetration = tz - fz # 正 = 穿透 + min_foot_z = min(min_foot_z, fz - tz) # 相对地形的高度 + if penetration > 0.02: # > 2cm 穿透 + self.stats["penetrations"].append( + (row, col, ttype, float(penetration))) + + # 5. 躯干沉入 + base_pose = env._body.get_pose(data) + bz = float(base_pose[0, 2]) + tz = self._get_terrain_z(float(base_pose[0, 0]), float(base_pose[0, 1])) + min_base_z = min(min_base_z, bz - tz) + if bz < tz + 0.10 and step >= SETTLE_STEPS: # 躯干离地面 < 10cm + self.stats["base_sink"].append((row, col, ttype, float(bz - tz))) + + # 6. 速度（卡住检测） + if step >= SETTLE_STEPS: + linvel = env.get_local_linvel(data)[0] + total_speed += float(np.linalg.norm(linvel[:2])) + + if self.renderer and step % 20 == 0: + self.renderer.render() + time.sleep(0.005) + + # ── 统计 ── + if not has_nan and max_dof_vel < 100.0: + avg_speed = total_speed / max(TEST_STEPS, 1) + if avg_speed < 0.01 and step == SETTLE_STEPS + TEST_STEPS - 1: + self.stats["stuck_events"].append((row, col, ttype, avg_speed)) + + if not has_nan and max_dof_vel < 50.0 and min_base_z > 0.05: + self.stats["ok_count"] += 1 + + # 记录接触力异常（完全无接触力但足部应该着地） + if max_contact_force < 1e-3 and step == SETTLE_STEPS + TEST_STEPS - 1: + feet = self._get_foot_positions(env._state) + for fi in range(4): + fz = float(feet[0, fi, 2]) + tz = self._get_terrain_z(float(feet[0, fi, 0]), float(feet[0, fi, 1])) + if abs(fz - tz) < 0.05: # 足部接近地面 + self.stats["contact_anomalies"].append( + (row, col, ttype, 0.0)) # 0 = 无接触力 + break + + def run(self, target_level=None, target_type=None): + """遍历所有地形 cell 测试。""" + total = 0 + start_t = time.time() + + for row in range(NUM_ROWS): + if target_level is not None and row != target_level: + continue + for col in range(NUM_COLS): + ttype = _terrain_type(col) + if target_type is not None and ttype != target_type: + continue + total += 1 + + print(f"[Diag] 测试 {NUM_ROWS}×{NUM_COLS} 网格, " + f"每 cell {STEPS_PER_CELL} 个 spawn 点 = {total * STEPS_PER_CELL} 次") + print(f"[Diag] Settle={SETTLE_STEPS}步 Test={TEST_STEPS}步") + + count = 0 + for row in range(NUM_ROWS): + if target_level is not None and row != target_level: + continue + for col in range(NUM_COLS): + ttype = _terrain_type(col) + if target_type is not None and ttype != target_type: + continue + self.test_cell(row, col, row) + count += 1 + if count % 10 == 0: + elapsed = time.time() - start_t + print(f" [{count}/{total}] {elapsed:.1f}s " + f"OK={self.stats['ok_count']}/{self.stats['total_tested']} " + f"穿透={len(self.stats['penetrations'])} " + f"NaN={len(self.stats['nan_events'])} " + f"爆炸={len(self.stats['explosions'])}") + + elapsed = time.time() - start_t + self._report(elapsed) + + def _report(self, elapsed): + s = self.stats + t = s["total_tested"] + print("\n" + "=" * 70) + print(f"地形碰撞诊断报告 ({elapsed:.1f}s, {t} 次测试)") + print("=" * 70) + + # 总体统计 + ok_rate = s["ok_count"] / max(t, 1) * 100 + print(f"\n ✅ 通过: {s['ok_count']}/{t} ({ok_rate:.1f}%)") + + # 穿透 + if s["penetrations"]: + pens = s["penetrations"] + pvals = [p[3] for p in pens] + print(f"\n 🔴 足部穿透: {len(pens)} 次") + print(f" 最大穿透: {max(pvals):.3f}m 平均: {np.mean(pvals):.3f}m") + + # 按地形类型分组 + by_type = {} + for p in pens: + tn = TYPE_NAMES[p[2]] + by_type.setdefault(tn, []).append(p[3]) + for tn, vals in sorted(by_type.items()): + print(f" {tn}: {len(vals)}次 max={max(vals):.3f}m avg={np.mean(vals):.3f}m") + + # 按 level 分组 + by_level = {} + for p in pens: + by_level.setdefault(p[0], []).append(p[3]) + print(f" 按 level: ", end="") + for lv in sorted(by_level.keys()): + vals = by_level[lv] + print(f"L{lv}({len(vals)}x max={max(vals):.3f}) ", end="") + print() + else: + print(f"\n 🟢 足部穿透: 0 次 ✓") + + # NaN + if s["nan_events"]: + print(f"\n 🔴 NaN/Inf: {len(s['nan_events'])} 次") + for n in s["nan_events"][:5]: + print(f" L{n[0]} C{n[1]:02d} {TYPE_NAMES[n[2]]} @step {n[3]}") + else: + print(f" 🟢 NaN/Inf: 0 次 ✓") + + # 物理爆炸 + if s["explosions"]: + print(f"\n 🟡 速度爆炸 (>{100}rad/s): {len(s['explosions'])} 次") + for e in s["explosions"][:5]: + print(f" L{e[0]} C{e[1]:02d} {TYPE_NAMES[e[2]]} max_vel={e[3]:.1f}") + else: + print(f" 🟢 速度爆炸: 0 次 ✓") + + # 接触力异常 + if s["contact_anomalies"]: + print(f"\n 🟡 接触力异常 (无接触但足部近地): {len(s['contact_anomalies'])} 次") + else: + print(f" 🟢 接触力异常: 0 次 ✓") + + # 卡住 + if s["stuck_events"]: + print(f"\n 🟡 卡住 (速度<0.01m/s): {len(s['stuck_events'])} 次") + by_type = {} + for e in s["stuck_events"]: + tn = TYPE_NAMES[e[2]] + by_type.setdefault(tn, []).append(1) + for tn, vals in sorted(by_type.items()): + print(f" {tn}: {len(vals)}次") + else: + print(f" 🟢 卡住: 0 次 ✓") + + # 躯干沉入 + if s["base_sink"]: + sinks = s["base_sink"] + print(f"\n 🟡 躯干贴地 (<10cm): {len(sinks)} 次") + vals = [s[3] for s in sinks] + print(f" 最差: {min(vals):.3f}m 平均: {np.mean(vals):.3f}m") + else: + print(f" 🟢 躯干贴地: 0 次 ✓") + + # 按地形类型汇总 + print(f"\n── 按地形类型汇总 ──") + for ti, tn in enumerate(TYPE_NAMES): + pens_t = sum(1 for p in s["penetrations"] if p[2] == ti) + nans_t = sum(1 for n in s["nan_events"] if n[2] == ti) + exps_t = sum(1 for e in s["explosions"] if e[2] == ti) + icon = "🔴" if (pens_t + nans_t + exps_t) > 0 else "🟢" + print(f" {icon} {tn}: 穿透{pens_t} NaN{nans_t} 爆炸{exps_t}") + + # 结论 + print(f"\n── 全局结论 ──") + if ok_rate > 95 and len(s["nan_events"]) == 0 and len(s["penetrations"]) < 5: + print(f" ✅ 仿真环境正常 — hfield 地形碰撞基本可靠") + elif ok_rate > 80: + print(f" ⚠ 仿真环境存在一些问题 — 部分地形类型/level 有问题") + else: + print(f" 🔴 仿真环境问题严重 — 需要排查 MotrixSim hfield 实现") + + if len(s["penetrations"]) > 10: + print(f" 💡 穿透较多 → hfield 碰撞面可能偏软或 contact margin 过大") + print(f" 可尝试增大 geom margin 或检查 MotrixSim 碰撞参数") + + +# ═══════════════════════════════════════════════════════════════════════ +# Main +# ═══════════════════════════════════════════════════════════════════════ + +def main(): + p = argparse.ArgumentParser(description="DreamWaQ 地形碰撞诊断") + p.add_argument("--level", type=int, default=None, help="仅测试指定 level (0-9)") + p.add_argument("--type", dest="ttype", default=None, + choices=["slope", "rough", "stairs_down", "stairs_up", "obstacles"], + help="仅测试指定地形类型") + p.add_argument("--render", action="store_true", help="渲染可视化") + p.add_argument("--settle", type=int, default=SETTLE_STEPS) + p.add_argument("--test", type=int, default=TEST_STEPS) + args = p.parse_args() + + type_map = {"slope": 0, "rough": 1, "stairs_down": 2, "stairs_up": 3, "obstacles": 4} + target_type = type_map.get(args.ttype) + + print("[Diag] 创建 DreamWaQ 环境...") + env = env_registry.make("go1-dreamwaq-walk", num_envs=1) + + # 如果指定 level，强制执行 + if args.level is not None: + env._force_level = args.level + print(f"[Diag] 强制 level={args.level}") + + env.init_state() + + # 确保 terrain 初始化（触发 hfield cache） + if env._model.num_hfields > 0: + hf = env._model.get_hfield(0) + env._hm_cache = hf.height_matrix + env._h_xmin, env._h_ymin = hf.bound[0], hf.bound[1] + env._h_xmax, env._h_ymax = hf.bound[3], hf.bound[4] + env._h_nr, env._h_nc = hf.height_matrix.shape + cfg_base = getattr(env.cfg, "hfield_z_base", None) + env._h_zbase = float(cfg_base) if cfg_base is not None else float(-hf.height_matrix[0, 0]) + print(f"[Diag] HField: {hf.height_matrix.shape} bound={hf.bound}") + print(f"[Diag] z_base={env._h_zbase:.3f}") + else: + print("[Diag] WARNING: 无 hfield（flat 场景），跳过诊断") + return + + diag = CollisionDiag(env, render=args.render) + diag.SETTLE_STEPS = args.settle + diag.TEST_STEPS = args.test + + try: + diag.run(target_level=args.level, target_type=target_type) + except KeyboardInterrupt: + print("\n[Diag] 中断") + finally: + if diag.renderer: + diag.renderer.close() + print("[Diag] 完成") + + +if __name__ == "__main__": + main() diff --git a/scripts/dreamwaq_sim2sim_mujoco.py b/scripts/dreamwaq_sim2sim_mujoco.py index ca9ea4e..3565408 100644 --- a/scripts/dreamwaq_sim2sim_mujoco.py +++ b/scripts/dreamwaq_sim2sim_mujoco.py @@ -31,7 +31,8 @@ HISTORY_LEN = 5 ACTION_SCALE = 0.25 KP = 28.0 KD = 0.7 -CLIP_ACTIONS = 23.7 +CLIP_ACTIONS = 4.0 +CLIP_TORQUES = 80.0 CLIP_OBS = 100.0 MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0 @@ -122,7 +123,7 @@ def main(): # Select XML scene terrain_map = { - "flat": "scene_motor_actuator.xml", + "flat": "scene_dreamwaq_flat.xml", "rough": "scene_rough_terrain.xml", "stairs": "scene_stairs_terrain.xml", "dreamwaq": "scene_dreamwaq_terrain.xml", @@ -243,10 +244,22 @@ def main(): action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS) last_action = action.copy() - # PD control + # PD control（对齐训练：目标限位 + 力矩裁剪） target = DEFAULT_ANGLES + action * ACTION_SCALE + # 关节目标限位（与训练一致） + jnt_lo = model.jnt_range[:, 0].copy() if hasattr(model, 'jnt_range') else None + jnt_hi = model.jnt_range[:, 1].copy() if hasattr(model, 'jnt_range') else None + # MuJoCo model.actuator_trnid 可能不直接暴露，改用 model.jnt_range + try: + trnid = model.actuator_trnid[:, 0] # transmission joint indices + lo = model.jnt_range[trnid, 0] + hi = model.jnt_range[trnid, 1] + except Exception: + lo = np.full(12, -12.0) + hi = np.full(12, 12.0) + target = np.clip(target, lo, hi) torques = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18] - data.ctrl[:] = np.clip(torques, -CLIP_ACTIONS, CLIP_ACTIONS) + data.ctrl[:] = np.clip(torques, -CLIP_TORQUES, CLIP_TORQUES) mujoco.mj_step(model, data) view.sync() diff --git a/scripts/export_dreamwaq_onnx_new.py b/scripts/export_dreamwaq_onnx_new.py index 4250dcc..b0cb07d 100644 --- a/scripts/export_dreamwaq_onnx_new.py +++ b/scripts/export_dreamwaq_onnx_new.py @@ -28,7 +28,7 @@ class DwaqInfer(torch.nn.Module): """obs: (N,45), obs_history: (N,225) → action: (N,12)""" code = self.vae.deterministic_code(obs_history) # (N,19) latent = torch.cat([code, obs], dim=-1) # (N,64) - return self.actor_mlp(latent) + return torch.clamp(self.actor_mlp(latent), -4.0, 4.0) def main(): diff --git a/scripts/export_dreamwaq_rsl_onnx.py b/scripts/export_dreamwaq_rsl_onnx.py index bfa863a..1f19256 100644 --- a/scripts/export_dreamwaq_rsl_onnx.py +++ b/scripts/export_dreamwaq_rsl_onnx.py @@ -22,13 +22,20 @@ PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) class DwaqInfer(tnn.Module): - """Deterministic inference: CENet mean code + actor.""" + """Deterministic inference: CENet mean code + RunningStats normalize + actor.""" def __init__(self, ac: ActorCritic_DWAQ): super().__init__() self.encoder = ac.encoder self.encode_mean_vel = ac.encode_mean_vel self.encode_mean_latent = ac.encode_mean_latent self.actor = ac.actor + # Bake RunningStats into the exported model + if hasattr(ac, 'actor_normalizer') and ac.actor_normalizer.count > 10: + self.register_buffer("norm_mean", ac.actor_normalizer.mean.clone()) + self.register_buffer("norm_std", ac.actor_normalizer.var.sqrt().clone()) + self.has_norm = True + else: + self.has_norm = False def forward(self, obs, obs_history): h = self.encoder(obs_history.reshape(obs_history.shape[0], -1)) # (B,225)->(B,64) @@ -36,6 +43,8 @@ class DwaqInfer(tnn.Module): latent = self.encode_mean_latent(h) # (B,16) mean latent code = torch.cat([vel, latent], dim=-1) # (B,19) = [vel, latent] x = torch.cat([code, obs], dim=-1) # (B,64) = [code, obs] + if self.has_norm: + x = torch.clamp((x - self.norm_mean) / (self.norm_std + 1e-8), -5.0, 5.0) return self.actor(x) # (B,12) diff --git a/scripts/play_amp_mujoco.py b/scripts/play_amp_mujoco.py new file mode 100644 index 0000000..e18e4a6 --- /dev/null +++ b/scripts/play_amp_mujoco.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Deploy Go1 AMP model in MuJoCo.""" +import mujoco, numpy as np, os, torch, time, sys +from mujoco import viewer + +# Use AMP's rsl_rl (NOT MotrixLab's) +sys.path.insert(0, '/home/8x54zj-m/amp_go2/rsl_rl') +sys.path.insert(0, '/home/8x54zj-m/amp_go2/legged_gym') +# Force reload of rsl_rl to pick up AMP version +for mod in list(sys.modules.keys()): + if 'rsl_rl' in mod: + del sys.modules[mod] + +XML_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + 'motrix_envs/src/motrix_envs/locomotion/go1/xmls') +CKPT = '/home/8x54zj-m/amp_go2/legged_gym/logs/go1_amp/Jul05_01-15-41_flat/model_6000.pt' + +import argparse +p = argparse.ArgumentParser() +p.add_argument("--terrain", default="flat", choices=["flat", "flat_stairs"]) +args = p.parse_args() + +scene = "scene_flat_stairs.xml" if args.terrain == "flat_stairs" else "go1_motor_actuator.xml" +os.chdir(XML_DIR) +m = mujoco.MjModel.from_xml_string(open(scene).read()) +d = mujoco.MjData(m) + +# Spawn on level 1 platform if using flat_stairs +spawn_x, spawn_y, spawn_z = 0.0, 0.0, 0.45 +if args.terrain == "flat_stairs": + col = np.random.randint(0, 4) + spawn_x = -12.0 + col * 8.0 + spawn_y = -4.0 # level 1 center + # Sample terrain height at spawn + floor_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "floor") + if floor_id >= 0 and m.geom_type[floor_id] == mujoco.mjtGeom.mjGEOM_HFIELD: + hf = m.geom_dataid[floor_id] + nr, nc = int(m.hfield_nrow[hf]), int(m.hfield_ncol[hf]) + sx, sy, zt, sb = m.hfield_size[hf] + hd = m.hfield_data[m.hfield_adr[hf]:m.hfield_adr[hf]+nr*nc].reshape(nr, nc) + gp = m.geom_pos[floor_id] + col_idx = int(np.clip(((spawn_x-gp[0])/sx*0.5+0.5)*(nc-1), 0, nc-1)) + row_idx = int(np.clip(((spawn_y-gp[1])/sy*0.5+0.5)*(nr-1), 0, nr-1)) + spawn_z = float(gp[2] + sb + hd[row_idx, col_idx] * zt) + 0.45 + print(f"[Terrain] flat_stairs level=1 spawn=({spawn_x:.1f},{spawn_y:.1f},{spawn_z:.2f})") + +ckpt = torch.load(CKPT, map_location='cpu', weights_only=False) +normalizer = ckpt['amp_normalizer'] + +actor = torch.nn.Sequential( + torch.nn.Linear(45, 512), torch.nn.ELU(), + torch.nn.Linear(512, 256), torch.nn.ELU(), + torch.nn.Linear(256, 128), torch.nn.ELU(), + torch.nn.Linear(128, 12) +) +sd = {k.replace('actor.',''): v for k,v in ckpt['model_state_dict'].items() if k.startswith('actor.')} +actor.load_state_dict(sd) +actor.eval() + +DEFAULT = np.array([-0.1,0.8,-1.5, 0.1,0.8,-1.5, -0.1,1.0,-1.5, 0.1,1.0,-1.5], dtype=np.float32) +LIN_VEL_SCALE, ANG_VEL_SCALE = 2.0, 0.25 +DOF_POS_SCALE, DOF_VEL_SCALE = 1.0, 0.05 +CMD_SCALE = np.array([2.0, 2.0, 0.25], dtype=np.float32) +KP, KD, ACT_SCALE = 20.0, 0.5, 0.25 + +def get_gravity(q): + qw,qx,qy,qz = q[3],q[0],q[1],q[2] + return np.array([2*(-qz*qx+qw*qy), -2*(qz*qy+qw*qx), 1-2*(qw*qw+qz*qz)]) + +def compute_obs(d, last_action, cmd): + obs = np.zeros(45, dtype=np.float32) + obs[0:3] = d.qvel[3:6] * ANG_VEL_SCALE + obs[3:6] = get_gravity(d.qpos[3:7]) + obs[6:9] = cmd * CMD_SCALE + obs[9:21] = (d.qpos[7:19] - DEFAULT) * DOF_POS_SCALE + obs[21:33] = d.qvel[6:18] * DOF_VEL_SCALE + obs[33:45] = last_action + return obs + +d.qpos[7:19] = DEFAULT; d.qpos[0:3] = [spawn_x, spawn_y, spawn_z]; d.qvel[:] = 0 +mujoco.mj_forward(m, d) + +view = viewer.launch_passive(m, d) +cmd = np.array([0.8, 0.0, 0.0], dtype=np.float32) +last_action = np.zeros(12, dtype=np.float32) + +print("AMP Go1 | W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出") + +while view.is_running(): + if view.is_running(): + obs = compute_obs(d, last_action, cmd) + obs_t = torch.from_numpy(obs).unsqueeze(0).float() + obs_norm = obs_t.clone() + obs_norm[:, :43] = normalizer.normalize(obs_t[:, :43]) + with torch.no_grad(): + act = actor(obs_norm)[0].numpy() + act = np.clip(act, -18, 18) + last_action = act.copy() + target = DEFAULT + act * ACT_SCALE + d.ctrl[:] = KP * (target - d.qpos[7:19]) - KD * d.qvel[6:18] + mujoco.mj_step(m, d) + view.sync() + time.sleep(0.001) diff --git a/scripts/test_abrupt_stop.py b/scripts/test_abrupt_stop.py new file mode 100644 index 0000000..a95cadf --- /dev/null +++ b/scripts/test_abrupt_stop.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +"""无头 MuJoCo 测试：前进 → 突然停止 → 观察姿态变化。""" +import numpy as np +import mujoco +import onnxruntime as ort +import os, sys, time + +_PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +XML_DIR = os.path.join(_PROJECT, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls") +ONNX = os.path.join(_PROJECT, "runs/go1-dreamwaq-walk/rslrl/26-07-02_20-03-23-_36054_PPO/policy.onnx") + +NUM_OBS = 45 +NUM_ACTIONS = 12 +HISTORY_LEN = 5 +ACTION_SCALE = 0.25 +KP = 28.0 +KD = 0.7 +DECIMATION = 4 +DEFAULT_ANGLES = np.array([0.0, 0.9, -1.8, 0.0, 0.9, -1.8, 0.0, 0.9, -1.8, 0.0, 0.9, -1.8], dtype=np.float32) + + +def compute_obs(model, data, commands, last_action): + obs = np.zeros(NUM_OBS, dtype=np.float32) + g = None # no gyro sensor in headless + obs[0:3] = (g if g is not None else data.qvel[3:6]) * 0.25 + grav_world = model.opt.gravity.copy() + grav_world = grav_world / np.linalg.norm(grav_world) + R = data.xmat[1].reshape(3, 3) + obs[3:6] = (R.T @ grav_world).astype(np.float32) + obs[6:9] = commands * np.array([2.0, 2.0, 0.25], dtype=np.float32) + obs[9:21] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0 + obs[21:33] = data.qvel[6:18] * 0.05 + obs[33:45] = last_action + return obs + + +def get_base_pose(data): + """返回 base 的 z 高度和 roll/pitch（度）。""" + quat = data.xquat[1] # base body quaternion + w, x, y, z = quat[0], quat[1], quat[2], quat[3] + # roll, pitch from quaternion + sinr_cosp = 2 * (w * x + y * z) + cosr_cosp = 1 - 2 * (x * x + y * y) + roll = np.arctan2(sinr_cosp, cosr_cosp) + sinp = 2 * (w * y - z * x) + pitch = np.arcsin(np.clip(sinp, -1, 1)) + return float(data.xpos[1, 2]), np.degrees(roll), np.degrees(pitch) + + +def main(): + session = ort.InferenceSession(ONNX, providers=['CPUExecutionProvider']) + xml_file = os.path.join(XML_DIR, "scene_dreamwaq_flat.xml") + model = mujoco.MjModel.from_xml_path(xml_file) + data = mujoco.MjData(model) + + # 初始化 + data.qpos[7:19] = DEFAULT_ANGLES + data.qpos[2] = 0.35 + mujoco.mj_forward(model, data) + + history = np.zeros((1, HISTORY_LEN, NUM_OBS), dtype=np.float32) + last_action = np.zeros(NUM_ACTIONS, dtype=np.float32) + + print(f"{'Step':>5s} {'Cmd_vx':>7s} {'base_z':>8s} {'roll':>8s} {'pitch':>8s} {'speed_xy':>8s}") + print("-" * 60) + + for step in range(3000): + # 命令：前 1500 步前进，之后突然停止 + if step < 1500: + vx, vy, wz = 0.5, 0.0, 0.0 # 前进 + else: + vx, vy, wz = 0.0, 0.0, 0.0 # 突然停止！ + + for _ in range(DECIMATION): + mujoco.mj_step(model, data) + + if step % DECIMATION == 0: + cmd = np.array([vx, vy, wz], dtype=np.float32) + obs = compute_obs(model, data, cmd, last_action) + history = np.concatenate([history[:, 1:, :], obs.reshape(1, 1, -1)], axis=1) + outputs = session.run(None, { + 'obs': obs.reshape(1, -1).astype(np.float32), + 'obs_history': history.reshape(1, -1).astype(np.float32), + }) + action = outputs[0][0] + last_action = action + + target = DEFAULT_ANGLES + action * ACTION_SCALE + data.ctrl[:] = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18] + + # 每 100 步打印 + if step % 100 == 0: + bz, roll, pitch = get_base_pose(data) + speed_xy = np.linalg.norm(data.qvel[0:2]) + print(f"{step:5d} {vx:7.1f} {bz:8.3f} {roll:+8.1f} {pitch:+8.1f} {speed_xy:8.3f}") + + # 最终状态 + bz, roll, pitch = get_base_pose(data) + print(f"\n最终: base_z={bz:.3f} roll={roll:.1f}° pitch={pitch:.1f}°") + if abs(roll) > 30 or abs(pitch) > 30: + print("⚠ 机器人倾覆！突然停止导致翻跟头") + else: + print("✅ 机器人在突然停止后保持稳定") + + +if __name__ == "__main__": + main()