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

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)

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
 

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)
 
             # ── 累计日志 ──

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]: