diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py b/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py index 544ff61..ce09d4a 100644 --- a/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py +++ b/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py @@ -419,7 +419,7 @@ 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 = 5 # 上游原值,随机 0-5 起步 + self._max_init_level = 3 # 先 0-3,适应后再提到 5 if num_reset > 0 and self._init_done and state is not None and hasattr(state, 'info'): old_info = state.info @@ -610,7 +610,8 @@ class DreamWaQTask(Go1WalkTask): } # state.reward 直接用带 dt 的值(不调 super,super 不带 dt) state = state.replace(reward=np.zeros(self._num_envs, dtype=np.float32)) - for v in scaled_terms.values(): + for k, v in scaled_terms.items(): + v = np.nan_to_num(v, nan=0.0, posinf=0.0, neginf=0.0) state.reward += v if self._cfg.reward_config.only_positive_rewards: state.reward = np.maximum(state.reward, 0.0) @@ -650,9 +651,11 @@ class DreamWaQTask(Go1WalkTask): 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) + torque = np.nan_to_num(data.actuator_ctrls, nan=0.0) + vel = np.nan_to_num(self.get_dof_vel(data), nan=0.0) + torque = np.clip(torque, -100, 100) + vel = np.clip(vel, -100, 100) + return np.clip(np.sum(np.abs(torque * vel), axis=1), 0, 1e6) def _reward_smoothness(self, info): scale = self.cfg.control_config.action_scale @@ -666,7 +669,9 @@ class DreamWaQTask(Go1WalkTask): 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) + torque = np.nan_to_num(data.actuator_ctrls, nan=0.0) + vel = np.nan_to_num(self.get_dof_vel(data), nan=0.0) + torque = np.clip(torque, -100, 100) + vel = np.clip(vel, -100, 100) power = torque * vel - return np.var(np.abs(power), axis=1) + return np.nan_to_num(np.var(np.abs(power), axis=1), nan=0.0) 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 cbbb1bf..cfb7317 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 @@ -121,8 +121,10 @@ class CENetActorModel(MLPModel): 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) 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) @@ -140,6 +142,12 @@ class CENetActorModel(MLPModel): self._last_cenet_output = out code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent = out + # 防止 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) + 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 5a5afa9..c7d453d 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 @@ -133,6 +133,7 @@ class DreamWaQPPO(PPO): # 每次更新后强制 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) # ── 累计日志 ── diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py b/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py index dbb7a3d..5c7ebb7 100644 --- a/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py +++ b/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py @@ -88,14 +88,14 @@ class RslrlNpEnvWrap(VecEnv): 支持 env 通过 state.info 传递 obs_history 和 privileged_obs。 """ - obs_dict = {"policy": torch.from_numpy(state.obs).to(self._device)} + obs_dict = {"policy": torch.from_numpy(np.nan_to_num(state.obs, nan=0.0)).to(self._device)} if "obs_history" in state.info: - hist = state.info["obs_history"] # (N, num_history, obs_dim) + hist = np.nan_to_num(state.info["obs_history"], nan=0.0) obs_dict["obs_history"] = torch.from_numpy(hist).reshape( self._num_envs, -1).to(self._device) if "privileged_obs" in state.info: - obs_dict["privileged_obs"] = torch.from_numpy( - state.info["privileged_obs"]).to(self._device) + priv = np.nan_to_num(state.info["privileged_obs"], nan=0.0) + obs_dict["privileged_obs"] = torch.from_numpy(priv).to(self._device) return obs_dict def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: