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

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