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

fix: NaN chain break — rewards, std, obs, physics all protected

Browse Source
This commit is contained in:
8x54zj-m
2026-07-01 10:31:47 +08:00
parent 1eede03ef0
commit 35a6f66511
4 changed files with 26 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -419,7 +419,7 @@ class DreamWaQTask(Go1WalkTask):
cy = half_y - self._border - row * self._cell_size - self._cell_size / 2 cy = half_y - self._border - row * self._cell_size - self._cell_size / 2
all_origins[row, col] = [cx, cy] all_origins[row, col] = [cx, cy]
self._terrain_origins = all_origins 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'): if num_reset > 0 and self._init_done and state is not None and hasattr(state, 'info'):
old_info = state.info old_info = state.info
@@ -610,7 +610,8 @@ class DreamWaQTask(Go1WalkTask):
} }
# state.reward 直接用带 dt 的值（不调 super，super 不带 dt） # state.reward 直接用带 dt 的值（不调 super，super 不带 dt）
state = state.replace(reward=np.zeros(self._num_envs, dtype=np.float32)) 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 state.reward += v
if self._cfg.reward_config.only_positive_rewards: if self._cfg.reward_config.only_positive_rewards:
state.reward = np.maximum(state.reward, 0.0) state.reward = np.maximum(state.reward, 0.0)
@@ -650,9 +651,11 @@ class DreamWaQTask(Go1WalkTask):
return np.square(base_z - ground_level - target) return np.square(base_z - ground_level - target)
def _reward_joint_power(self, data): def _reward_joint_power(self, data):
torque = np.clip(data.actuator_ctrls, -100, 100) torque = np.nan_to_num(data.actuator_ctrls, nan=0.0)
vel = np.clip(self.get_dof_vel(data), -100, 100) vel = np.nan_to_num(self.get_dof_vel(data), nan=0.0)
return np.sum(np.abs(torque * vel), axis=1) 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): def _reward_smoothness(self, info):
scale = self.cfg.control_config.action_scale scale = self.cfg.control_config.action_scale
@@ -666,7 +669,9 @@ class DreamWaQTask(Go1WalkTask):
return np.sum(diff, axis=1) return np.sum(diff, axis=1)
def _reward_power_distribution(self, data): def _reward_power_distribution(self, data):
torque = np.clip(data.actuator_ctrls, -100, 100) torque = np.nan_to_num(data.actuator_ctrls, nan=0.0)
vel = np.clip(self.get_dof_vel(data), -100, 100) 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 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)

8
motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py
View File

@@ -121,8 +121,10 @@ class CENetActorModel(MLPModel):
if self.stochastic and not self.state_dependent_std: if self.stochastic and not self.state_dependent_std:
with torch.no_grad(): with torch.no_grad():
if self.noise_std_type == "scalar": 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)
elif self.noise_std_type == "log": 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) self.log_std.clamp_(min=-20.0, max=10.0)
super()._update_distribution(obs) super()._update_distribution(obs)
@@ -140,6 +142,12 @@ class CENetActorModel(MLPModel):
self._last_cenet_output = out self._last_cenet_output = out
code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent = 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) latent = torch.cat([code, policy_obs], dim=-1) # (N, 64)
return latent return latent

1
motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py
View File

@@ -133,6 +133,7 @@ class DreamWaQPPO(PPO):
# 每次更新后强制 std > 0，防止数值异常导致 NaN # 每次更新后强制 std > 0，防止数值异常导致 NaN
if hasattr(self.actor, 'std') and self.actor.stochastic: if hasattr(self.actor, 'std') and self.actor.stochastic:
with torch.no_grad(): 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) self.actor.std.clamp_(min=1e-6)
# ── 累计日志 ── # ── 累计日志 ──

8
motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py
View File

@@ -88,14 +88,14 @@ class RslrlNpEnvWrap(VecEnv):
支持 env 通过 state.info 传递 obs_history 和 privileged_obs。 支持 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: 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( obs_dict["obs_history"] = torch.from_numpy(hist).reshape(
self._num_envs, -1).to(self._device) self._num_envs, -1).to(self._device)
if "privileged_obs" in state.info: if "privileged_obs" in state.info:
obs_dict["privileged_obs"] = torch.from_numpy( priv = np.nan_to_num(state.info["privileged_obs"], nan=0.0)
state.info["privileged_obs"]).to(self._device) obs_dict["privileged_obs"] = torch.from_numpy(priv).to(self._device)
return obs_dict return obs_dict
def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: